github.com/koko1123/flow-go-1@v0.29.6/module/jobqueue/workerpool.go (about)

     1  package jobqueue
     2  
     3  import (
     4  	"github.com/koko1123/flow-go-1/module"
     5  	"github.com/koko1123/flow-go-1/module/component"
     6  	"github.com/koko1123/flow-go-1/module/irrecoverable"
     7  	"github.com/koko1123/flow-go-1/module/util"
     8  )
     9  
    10  // WorkerPool implements the jobqueue.Worker interface, and wraps the processing to make it
    11  // compatible with the Component interface.
    12  type WorkerPool struct {
    13  	component.Component
    14  
    15  	cm        *component.ComponentManager
    16  	processor JobProcessor
    17  	notify    NotifyDone
    18  	ch        chan module.Job
    19  }
    20  
    21  // JobProcessor is called by the worker to execute each job. It should only return when the job has
    22  // completed, either successfully or after performing any failure handling.
    23  // It takes 3 arguments:
    24  //   - irrecoverable.SignalerContext: this is used to signal shutdown to the worker and throw any
    25  //     irrecoverable errors back to the parent. The signaller context is passed in from consumer's
    26  //     Start method
    27  //   - module.Job: the job to be processed. The processor is responsible for decoding into the
    28  //     expected format.
    29  //   - func(): Call this closure after the job is considered complete. This is a convenience method
    30  //     that avoid needing to a separate ProcessingNotifier for simple usecases. If a different method
    31  //     is used to signal jobs are done to the consumer, this function can be ignored.
    32  type JobProcessor func(irrecoverable.SignalerContext, module.Job, func())
    33  
    34  // NotifyDone should be the consumer's NotifyJobIsDone method, or a wrapper for that method. It is
    35  // wrapped in a closure and added as an argument to the JobProcessor to notify the consumer that
    36  // the job is done.
    37  type NotifyDone func(module.JobID)
    38  
    39  // NewWorkerPool returns a new WorkerPool
    40  func NewWorkerPool(processor JobProcessor, notify NotifyDone, workers uint64) *WorkerPool {
    41  	w := &WorkerPool{
    42  		processor: processor,
    43  		notify:    notify,
    44  		ch:        make(chan module.Job),
    45  	}
    46  
    47  	builder := component.NewComponentManagerBuilder()
    48  
    49  	for i := uint64(0); i < workers; i++ {
    50  		builder.AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    51  			ready()
    52  			w.workerLoop(ctx)
    53  		})
    54  	}
    55  
    56  	w.cm = builder.Build()
    57  	w.Component = w.cm
    58  
    59  	return w
    60  }
    61  
    62  // Run executes the worker's JobProcessor for the provided job.
    63  // Run is non-blocking.
    64  func (w *WorkerPool) Run(job module.Job) error {
    65  	// don't accept new jobs after shutdown is signalled
    66  	if util.CheckClosed(w.cm.ShutdownSignal()) {
    67  		return nil
    68  	}
    69  
    70  	select {
    71  	case <-w.cm.ShutdownSignal():
    72  		return nil
    73  	case w.ch <- job:
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  // workerLoop processes incoming jobs passed via the Run method. The job execution is wrapped in a
    80  // goroutine to support passing the worker's irrecoverable.SignalerContext into the processor.
    81  func (w *WorkerPool) workerLoop(ctx irrecoverable.SignalerContext) {
    82  	for {
    83  		select {
    84  		case <-ctx.Done():
    85  			return
    86  		case job := <-w.ch:
    87  			w.processor(ctx, job, func() {
    88  				w.notify(job.ID())
    89  			})
    90  		}
    91  	}
    92  }