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

     1  package jobqueue
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/rs/zerolog"
     7  
     8  	"github.com/koko1123/flow-go-1/module"
     9  	"github.com/koko1123/flow-go-1/module/component"
    10  	"github.com/koko1123/flow-go-1/module/irrecoverable"
    11  	"github.com/koko1123/flow-go-1/module/util"
    12  	"github.com/koko1123/flow-go-1/storage"
    13  )
    14  
    15  type ComponentConsumer struct {
    16  	component.Component
    17  
    18  	cm           *component.ComponentManager
    19  	consumer     module.JobConsumer
    20  	jobs         module.Jobs
    21  	workSignal   <-chan struct{}
    22  	preNotifier  NotifyDone
    23  	postNotifier NotifyDone
    24  	log          zerolog.Logger
    25  }
    26  
    27  // NewComponentConsumer creates a new ComponentConsumer consumer
    28  func NewComponentConsumer(
    29  	log zerolog.Logger,
    30  	workSignal <-chan struct{},
    31  	progress storage.ConsumerProgress,
    32  	jobs module.Jobs,
    33  	defaultIndex uint64,
    34  	processor JobProcessor, // method used to process jobs
    35  	maxProcessing uint64,
    36  	maxSearchAhead uint64,
    37  ) *ComponentConsumer {
    38  
    39  	c := &ComponentConsumer{
    40  		workSignal: workSignal,
    41  		jobs:       jobs,
    42  		log:        log,
    43  	}
    44  
    45  	// create a worker pool with maxProcessing workers to process jobs
    46  	worker := NewWorkerPool(
    47  		processor,
    48  		func(id module.JobID) { c.NotifyJobIsDone(id) },
    49  		maxProcessing,
    50  	)
    51  	c.consumer = NewConsumer(c.log, c.jobs, progress, worker, maxProcessing, maxSearchAhead)
    52  
    53  	builder := component.NewComponentManagerBuilder().
    54  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    55  			worker.Start(ctx)
    56  			if err := util.WaitClosed(ctx, worker.Ready()); err != nil {
    57  				c.log.Info().Msg("job consumer startup aborted")
    58  				<-worker.Done()
    59  				c.log.Info().Msg("job consumer shutdown complete")
    60  				return
    61  			}
    62  
    63  			c.log.Info().Msg("job consumer starting")
    64  			err := c.consumer.Start(defaultIndex)
    65  			if err != nil {
    66  				ctx.Throw(fmt.Errorf("could not start consumer: %w", err))
    67  			}
    68  
    69  			ready()
    70  
    71  			<-ctx.Done()
    72  			c.log.Info().Msg("job consumer shutdown started")
    73  
    74  			// blocks until all running jobs have stopped
    75  			c.consumer.Stop()
    76  
    77  			<-worker.Done()
    78  			c.log.Info().Msg("job consumer shutdown complete")
    79  		}).
    80  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    81  			ready()
    82  			c.processingLoop(ctx)
    83  		})
    84  
    85  	cm := builder.Build()
    86  	c.cm = cm
    87  	c.Component = cm
    88  
    89  	return c
    90  }
    91  
    92  // SetPreNotifier sets a notification function that is invoked before marking a job as done in the
    93  // consumer.
    94  //
    95  // Note: This guarantees that the function is called at least once for each job, but may be executed
    96  // before consumer updates the last processed index.
    97  func (c *ComponentConsumer) SetPreNotifier(fn NotifyDone) {
    98  	c.preNotifier = fn
    99  }
   100  
   101  // SetPostNotifier sets a notification function that is invoked after marking a job as done in the
   102  // consumer.
   103  //
   104  // Note: This guarantees that the function is executed after consumer updates the last processed index,
   105  // but notifications may be missed in the event of a crash.
   106  func (c *ComponentConsumer) SetPostNotifier(fn NotifyDone) {
   107  	c.postNotifier = fn
   108  }
   109  
   110  // NotifyJobIsDone is invoked by the worker to let the consumer know that it is done
   111  // processing a (block) job.
   112  func (c *ComponentConsumer) NotifyJobIsDone(jobID module.JobID) uint64 {
   113  	if c.preNotifier != nil {
   114  		c.preNotifier(jobID)
   115  	}
   116  
   117  	// notify wrapped consumer that job is complete
   118  	processedIndex := c.consumer.NotifyJobIsDone(jobID)
   119  
   120  	if c.postNotifier != nil {
   121  		c.postNotifier(jobID)
   122  	}
   123  
   124  	return processedIndex
   125  }
   126  
   127  // Size returns number of in-memory block jobs that block consumer is processing.
   128  func (c *ComponentConsumer) Size() uint {
   129  	return c.consumer.Size()
   130  }
   131  
   132  // Head returns the highest job index available
   133  func (c *ComponentConsumer) Head() (uint64, error) {
   134  	return c.jobs.Head()
   135  }
   136  
   137  // LastProcessedIndex returns the last processed job index
   138  func (c *ComponentConsumer) LastProcessedIndex() uint64 {
   139  	return c.consumer.LastProcessedIndex()
   140  }
   141  
   142  func (c *ComponentConsumer) processingLoop(ctx irrecoverable.SignalerContext) {
   143  	c.log.Debug().Msg("listening for new jobs")
   144  	for {
   145  		select {
   146  		case <-ctx.Done():
   147  			return
   148  		case <-c.workSignal:
   149  			c.consumer.Check()
   150  		}
   151  	}
   152  }