fix(processor): export the last partial batch on shutdown

pkg/cannon/cannon.go

@@ -55,7 +55,10 @@ type Cannon struct {
 
 	scheduler gocron.Scheduler
 
+	// deriversMu guards eventDerivers and stopping, which the beacon's OnReady callback and Shutdown both touch.
+	deriversMu    sync.Mutex
 	eventDerivers []deriver.EventDeriver
+	stopping      bool
 
 	coordinatorClient *coordinator.Client
 
@@ -178,49 +181,94 @@ func (c *Cannon) Start(ctx context.Context) error {
 		c.log.WithField("network", c.Config.Ethereum.OverrideNetworkName).WithContext(ctx).Info("Overriding network name")
 	}
 
-	if err := c.beacon.Start(ctx); err != nil {
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
+
+	// The beacon node runs until it fails, so watch the signal alongside it; cancelling it stops the derivers.
+	beaconCtx, stopBeacon := context.WithCancel(ctx)
+	defer stopBeacon()
+
+	beaconErr := make(chan error, 1)
+
+	go func() {
+		beaconErr <- c.beacon.Start(beaconCtx)
+	}()
+
+	select {
+	case err := <-beaconErr:
 		return err
+	case sig := <-signals:
+		c.log.WithContext(ctx).Printf("Caught signal: %v", sig)
 	}
 
-	cancel := make(chan os.Signal, 1)
-	signal.Notify(cancel, syscall.SIGTERM, syscall.SIGINT)
+	stopBeacon()
 
-	sig := <-cancel
-	c.log.WithContext(ctx).Printf("Caught signal: %v", sig)
+	var errs []error
+
+	// Join the beacon node before the derivers and sinks it feeds are torn down.
+	select {
+	case err := <-beaconErr:
+		if err != nil && !errors.Is(err, context.Canceled) {
+			errs = append(errs, fmt.Errorf("beacon node: %w", err))
+		}
+	case <-time.After(beaconStopTimeout):
+		c.log.WithContext(ctx).Warn("Beacon node did not stop in time")
+	}
 
 	if err := c.Shutdown(ctx); err != nil {
-		return err
+		errs = append(errs, err)
 	}
 
-	return nil
+	return errors.Join(errs...)
 }
 
+const (
+	// beaconStopTimeout and shutdownTimeout together stay below the chart's 90s terminationGracePeriodSeconds.
+	beaconStopTimeout = 10 * time.Second
+	shutdownTimeout   = 75 * time.Second
+)
+
+// Shutdown stops the producers first and then flushes the sinks, logging failures so one cannot skip the rest.
 func (c *Cannon) Shutdown(ctx context.Context) error {
 	c.log.WithContext(ctx).Printf("Shutting down")
 
-	for _, sink := range c.sinks {
-		if err := sink.Stop(ctx); err != nil {
-			return err
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), shutdownTimeout)
+	defer cancel()
+
+	var errs []error
+
+	if err := c.scheduler.Shutdown(); err != nil {
+		errs = append(errs, fmt.Errorf("scheduler: %w", err))
+	}
+
+	c.deriversMu.Lock()
+	c.stopping = true
+	derivers := c.eventDerivers
+	c.deriversMu.Unlock()
+
+	for _, deriver := range derivers {
+		if err := deriver.Stop(ctx); err != nil {
+			errs = append(errs, fmt.Errorf("deriver %s: %w", deriver.Name(), err))
 		}
 	}
 
 	for _, fun := range c.shutdownFuncs {
 		if err := fun(ctx); err != nil {
-			return err
+			errs = append(errs, err)
 		}
 	}
 
-	if err := c.scheduler.Shutdown(); err != nil {
-		return err
+	for _, sink := range c.sinks {
+		if err := sink.Stop(ctx); err != nil {
+			errs = append(errs, fmt.Errorf("sink %s: %w", sink.Name(), err))
+		}
 	}
 
-	for _, deriver := range c.eventDerivers {
-		if err := deriver.Stop(ctx); err != nil {
-			return err
-		}
+	for _, err := range errs {
+		c.log.WithError(err).WithContext(ctx).Error("Shutdown step failed")
 	}
 
-	return nil
+	return errors.Join(errs...)
 }
 
 func (c *Cannon) ServeMetrics(ctx context.Context) error {
@@ -1176,6 +1224,14 @@ func (c *Cannon) startBeaconBlockProcessor(ctx context.Context) error {
 			)
 		}
 
+		c.deriversMu.Lock()
+		defer c.deriversMu.Unlock()
+
+		// Derivers started after Shutdown would feed sinks that are already stopping.
+		if c.stopping {
+			return nil
+		}
+
 		c.eventDerivers = eventDerivers
 
 		// Refresh the spec every epoch
@@ -1195,6 +1251,10 @@ func (c *Cannon) startBeaconBlockProcessor(ctx context.Context) error {
 
 			go func() {
 				if err := c.startDeriverWhenReady(ctx, d); err != nil {
+					if ctx.Err() != nil {
+						return
+					}
+
 					c.log.
 						WithField("deriver", d.Name()).
 						WithError(err).WithContext(ctx).Fatal("Failed to start deriver")

pkg/output/http/exporter_test.go

@@ -207,6 +207,23 @@ func TestHTTPSinkProcessorAttemptsSparseGroupsAfterEarlierGroupFailure(t *testin
 	assert.Equal(t, blockSpan.SpanContext().TraceID(), spanContextFromHTTPTraceparent(requests[2].traceparent).TraceID())
 }
 
+func TestHTTPSinkStopExportsQueuedEvents(t *testing.T) {
+	recorder := startRecordingHTTPServer(t)
+	sink := newTestHTTPSink(t, recorder.server.URL, 512)
+
+	require.NoError(t, sink.Start(context.Background()))
+
+	for i := range 3 {
+		require.NoError(t, sink.HandleNewDecoratedEvent(context.Background(), newHTTPTestEvent(xatu.Event_BEACON_API_ETH_V1_EVENTS_BLOCK, fmt.Sprintf("id-%d", i))))
+	}
+
+	require.NoError(t, sink.Stop(context.Background()))
+
+	requests := recorder.snapshotRequests()
+	require.Len(t, requests, 1)
+	assert.Equal(t, []string{"id-0", "id-1", "id-2"}, requests[0].eventIDs)
+}
+
 func startRecordingHTTPServer(t *testing.T) *recordingHTTPServer {
 	t.Helper()
 

pkg/processor/batch.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"runtime"
 	"sort"
 	"strings"
 	"sync"
@@ -217,11 +216,13 @@ type BatchItemProcessor[T any] struct {
 	batchCh chan []*TraceableItem[T]
 	name    string
 
-	timer         *time.Timer
-	stopWait      sync.WaitGroup
-	stopOnce      sync.Once
-	stopCh        chan struct{}
-	stopWorkersCh chan struct{}
+	timer    *time.Timer
+	stopWait sync.WaitGroup
+	stopOnce sync.Once
+	stopCh   chan struct{}
+
+	// queueMu is held for reading while an item is sent to the queue, so Shutdown can close it safely.
+	queueMu sync.RWMutex
 
 	// Metrics
 	metrics *Metrics
@@ -285,16 +286,15 @@ func NewBatchItemProcessor[T any](exporter ItemExporter[T], name string, log obs
 	}
 
 	bvp := BatchItemProcessor[T]{
-		e:             exporter,
-		o:             o,
-		log:           log,
-		name:          name,
-		metrics:       metrics,
-		timer:         time.NewTimer(o.BatchTimeout),
-		queue:         make(chan *TraceableItem[T], o.MaxQueueSize),
-		batchCh:       make(chan []*TraceableItem[T], o.Workers),
-		stopCh:        make(chan struct{}),
-		stopWorkersCh: make(chan struct{}),
+		e:       exporter,
+		o:       o,
+		log:     log,
+		name:    name,
+		metrics: metrics,
+		timer:   time.NewTimer(o.BatchTimeout),
+		queue:   make(chan *TraceableItem[T], o.MaxQueueSize),
+		batchCh: make(chan []*TraceableItem[T], o.Workers),
+		stopCh:  make(chan struct{}),
 	}
 
 	bvp.log.WithFields(
@@ -311,21 +311,27 @@ func NewBatchItemProcessor[T any](exporter ItemExporter[T], name string, log obs
 	return &bvp, nil
 }
 
+// Start launches the batch builder and workers; exports ignore ctx's cancellation so Shutdown can still flush.
 func (bvp *BatchItemProcessor[T]) Start(ctx context.Context) {
-	bvp.stopWait.Add(bvp.o.Workers)
+	bvp.stopWait.Add(bvp.o.Workers + 1)
 
 	bvp.metrics.SetWorkerCount(bvp.name, float64(bvp.o.Workers))
 
 	bvp.log.WithContext(ctx).Infof("Starting %d workers for %s", bvp.o.Workers, bvp.name)
 
+	exportCtx := context.WithoutCancel(ctx)
+
 	for i := 0; i < bvp.o.Workers; i++ {
 		go func(num int) {
 			defer bvp.stopWait.Done()
-			bvp.worker(ctx, num)
+
+			bvp.worker(exportCtx, num)
 		}(i)
 	}
 
 	go func() {
+		defer bvp.stopWait.Done()
+
 		bvp.batchBuilder(ctx)
 		bvp.log.WithContext(ctx).Info("Batch builder exited")
 	}()
@@ -453,7 +459,7 @@ func (bvp *BatchItemProcessor[T]) exportWithTimeout(ctx context.Context, itemsBa
 	return nil
 }
 
-// Shutdown shuts down the batch item processor.
+// Shutdown rejects new writes, exports every queued item and shuts the exporter down, or returns ctx's error first.
 func (bvp *BatchItemProcessor[T]) Shutdown(ctx context.Context) error {
 	var err error
 
@@ -464,14 +470,15 @@ func (bvp *BatchItemProcessor[T]) Shutdown(ctx context.Context) error {
 
 			close(bvp.stopCh)
 
-			bvp.timer.Stop()
-
-			bvp.drainQueue()
-
-			close(bvp.stopWorkersCh)
+			// Closing the queue makes the builder send its last batch and close batchCh; workers exit once it drains.
+			bvp.queueMu.Lock()
+			close(bvp.queue)
+			bvp.queueMu.Unlock()
 
 			bvp.stopWait.Wait()
 
+			bvp.timer.Stop()
+
 			if bvp.e != nil {
 				if err = bvp.e.Shutdown(ctx); err != nil {
 					bvp.log.WithError(err).WithContext(ctx).Error("failed to shutdown processor")
@@ -552,15 +559,14 @@ func (bvp *BatchItemProcessor[T]) batchBuilder(ctx context.Context) {
 
 	var batch []*TraceableItem[T]
 
+	defer close(bvp.batchCh)
+
 	for {
 		select {
-		case <-bvp.stopWorkersCh:
-			log.WithContext(ctx).Info("Stopping batch builder")
-			return
 		case item, ok := <-bvp.queue:
 			if !ok {
-				// Channel is closed, send any remaining items in the batch for processing
-				// before shutting down.
+				log.WithContext(ctx).Info("Stopping batch builder")
+
 				if len(batch) > 0 {
 					bvp.sendBatch(batch, "shutdown")
 				}
@@ -602,64 +608,24 @@ func (bvp *BatchItemProcessor[T]) sendBatch(batch []*TraceableItem[T], reason st
 	log.Tracef("Batch sent to batch channel")
 }
 
+// worker exports batches until the batch builder closes batchCh.
 func (bvp *BatchItemProcessor[T]) worker(ctx context.Context, number int) {
-	for {
-		select {
-		case <-bvp.stopWorkersCh:
-			bvp.log.WithContext(ctx).Infof("Stopping worker %d", number)
-
-			return
-		case batch := <-bvp.batchCh:
-			bvp.timer.Reset(bvp.o.BatchTimeout)
-
-			if err := bvp.exportWithTimeout(ctx, batch); err != nil {
-				bvp.log.WithError(err).WithContext(ctx).Error("failed to export items")
-			}
+	for batch := range bvp.batchCh {
+		bvp.timer.Reset(bvp.o.BatchTimeout)
 
-			bvp.metrics.SetItemsQueued(bvp.name, float64(len(bvp.queue)))
+		if err := bvp.exportWithTimeout(ctx, batch); err != nil {
+			bvp.log.WithError(err).WithContext(ctx).Error("failed to export items")
 		}
-	}
-}
-
-func (bvp *BatchItemProcessor[T]) drainQueue() {
-	bvp.log.Info("Draining queue: waiting for the batch builder to process remaining items")
 
-	// First wait for queue to be processed
-	for len(bvp.queue) > 0 {
-		time.Sleep(10 * time.Millisecond)
-	}
-
-	bvp.log.Info("Draining queue: waiting for workers to finish processing batches")
-
-	// Then wait for any in-flight batches
-	for len(bvp.batchCh) > 0 {
-		time.Sleep(10 * time.Millisecond)
-	}
-
-	bvp.log.Info("Draining queue: all items processed")
-
-	close(bvp.queue)
-}
-
-func recoverSendOnClosedChan() {
-	x := recover()
-
-	switch err := x.(type) {
-	case nil:
-		return
-	case runtime.Error:
-		if err.Error() == "send on closed channel" {
-			return
-		}
+		bvp.metrics.SetItemsQueued(bvp.name, float64(len(bvp.queue)))
 	}
 
-	panic(x)
+	bvp.log.WithContext(ctx).Infof("Stopping worker %d", number)
 }
 
 func (bvp *BatchItemProcessor[T]) enqueueOrDrop(ctx context.Context, item *TraceableItem[T]) error {
-	// This ensures the bvp.queue<- below does not panic as the
-	// processor shuts down.
-	defer recoverSendOnClosedChan()
+	bvp.queueMu.RLock()
+	defer bvp.queueMu.RUnlock()
 
 	select {
 	case <-bvp.stopCh: