github.com/benchkram/bob@v0.0.0-20240314204020-b7a57f2f9be9/bob/playbook/workers.go (about)

     1  package playbook
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/benchkram/bob/bobtask/processed"
    10  )
    11  
    12  type workerManager struct {
    13  	// workerWG wait for all workers to be stopped.
    14  	workerWG sync.WaitGroup
    15  
    16  	// workloadQueues to send workload to a worker
    17  	workloadQueues []chan *Status
    18  
    19  	// workers report themself as idle
    20  	// by putting their id on this channel.
    21  	idleChan chan int
    22  
    23  	errorsMutex sync.Mutex
    24  	errors      []error
    25  
    26  	processedMutex sync.Mutex
    27  	processed      []*processed.Task
    28  
    29  	shutdownMutext sync.Mutex
    30  	shutdown       bool
    31  }
    32  
    33  func newWorkerManager() *workerManager {
    34  	s := &workerManager{
    35  		workloadQueues: []chan *Status{},
    36  		idleChan:       make(chan int, 1000),
    37  
    38  		errors:    []error{},
    39  		processed: []*processed.Task{},
    40  	}
    41  	return s
    42  }
    43  
    44  func (wm *workerManager) stopWorkers() {
    45  	wm.shutdownMutext.Lock()
    46  	wm.shutdown = true
    47  	wm.shutdownMutext.Unlock()
    48  }
    49  
    50  func (wm *workerManager) closeWorkloadQueues() {
    51  	for _, q := range wm.workloadQueues {
    52  		close(q)
    53  	}
    54  }
    55  
    56  func (wm *workerManager) canShutdown() bool {
    57  	wm.shutdownMutext.Lock()
    58  	defer wm.shutdownMutext.Unlock()
    59  	return wm.shutdown
    60  
    61  }
    62  
    63  func (wm *workerManager) addError(err error) {
    64  	wm.errorsMutex.Lock()
    65  	wm.errors = append(wm.errors, err)
    66  	wm.errorsMutex.Unlock()
    67  }
    68  
    69  func (wm *workerManager) addProcessedTask(t *processed.Task) {
    70  	wm.processedMutex.Lock()
    71  	wm.processed = append(wm.processed, t)
    72  	wm.processedMutex.Unlock()
    73  }
    74  
    75  // startWorkers and return a state to interact with the workers.
    76  func (p *Playbook) startWorkers(ctx context.Context, workers int) *workerManager {
    77  
    78  	wm := newWorkerManager()
    79  
    80  	// Start the workers which listen on task queue
    81  	for i := 0; i < workers; i++ {
    82  
    83  		// create workload queue for this worker
    84  		// A non-blocking queue with a too big size allows to shutdown gracefully in
    85  		// case of error. The workload pump might pump more tasks to the queues
    86  		// and to avoid blocking behaviour the queue size is not nil.
    87  		queue := make(chan *Status, 1000)
    88  		wm.workloadQueues = append(wm.workloadQueues, queue)
    89  
    90  		wm.workerWG.Add(1)
    91  		go func(workerID int) {
    92  			// signal availability to receive workload
    93  			wm.idleChan <- workerID
    94  
    95  			for t := range queue {
    96  				t.SetStart(time.Now())
    97  				_ = p.setTaskState(t.Task.TaskID, StateRunning, nil)
    98  
    99  				// check if a shutdown is required.
   100  				if wm.canShutdown() {
   101  					break
   102  				}
   103  
   104  				processedTask, err := p.build(ctx, t.Task)
   105  				if err != nil {
   106  					wm.addError(fmt.Errorf("(worker) [task: %s], %w", t.Name(), err))
   107  
   108  					// stopp workers asap.
   109  					wm.stopWorkers()
   110  				}
   111  				wm.addProcessedTask(processedTask)
   112  
   113  				// done with processing. signal availability.
   114  				wm.idleChan <- workerID
   115  			}
   116  			wm.workerWG.Done()
   117  		}(i)
   118  	}
   119  
   120  	return wm
   121  
   122  }