github.com/onflow/flow-go@v0.33.17/module/jobqueue/component_consumer.go (about)

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