Fix executor coordinator graceful shutdown. (#747)

* Add test.

* Graceful shutdown for executor coordinator.

Author: winder

executor/executor_coordinator.go

@@ -94,12 +94,10 @@ func (ec *Coordinator) Start(_ context.Context) error {
 		})
 
 		// Start worker goroutines
-		ec.wg.Add(ec.workerCount)
 		for i := 0; i < ec.workerCount; i++ {
-			go func() {
-				defer ec.wg.Done()
+			ec.wg.Go(func() {
 				ec.handleMessage(c)
-			}()
+			})
 		}
 
 		ec.lggr.Infow("Coordinator started")
@@ -116,12 +114,12 @@ func (ec *Coordinator) Close() error {
 			ec.cancel()
 		}
 
-		// Close channel to signal workers to stop
-		close(ec.workerPoolTasks)
-
 		// Wait for all goroutines to finish
 		ec.wg.Wait()
 
+		// It is safe to close the channel once all goroutines have stopped.
+		close(ec.workerPoolTasks)
+
 		// Update running state to reflect in healthcheck and readiness
 		ec.running.Store(false)
 
@@ -204,6 +202,7 @@ func (ec *Coordinator) runProcessingLoop(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
+			ec.lggr.Infow("Processing loop exiting")
 			return
 		case <-ticker.C:
 			currentTime := ec.timeProvider.GetTime()
@@ -217,8 +216,13 @@ func (ec *Coordinator) runProcessingLoop(ctx context.Context) {
 			)
 			for _, payload := range readyMessages {
 				ec.inFlightAdd(payload.MessageID)
-				// If the channel is full, we will block here, but messages will continue to be accumulate in the heap.
-				ec.workerPoolTasks <- payload
+				// If the channel is full, we will block here, but messages will continue to accumulate in the heap.
+				select {
+				case ec.workerPoolTasks <- payload:
+				case <-ctx.Done():
+					ec.lggr.Infow("Processing loop dropping payload to exit")
+					continue
+				}
 			}
 		case <-reportingTicker.C:
 			ec.monitoring.Metrics().RecordMessageHeapSize(ctx, int64(ec.delayedMessageHeap.Len()))
@@ -230,12 +234,12 @@ func (ec *Coordinator) handleMessage(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
+			ec.lggr.Infow("Message handler exiting")
 			return
 		case payload, ok := <-ec.workerPoolTasks:
-			if !ok {
-				return
+			if ok {
+				ec.processPayload(ctx, payload)
 			}
-			ec.processPayload(ctx, payload)
 		}
 	}
 }

executor/executor_coordinator_test.go

@@ -559,3 +559,83 @@ func TestDuplicateMessageIDFromStreamWhenAlreadyInHeap_IsSkippedByHeapAndHandleM
 	time.Sleep(2 * time.Second)
 	require.True(t, mock.AssertExpectationsForObjects(t, mockExecutor))
 }
+
+// TestGracefulShutdown tests that when Close() is called while a message is being processed, the processing loop will
+// shut down gracefully. This state is simulated by blocking the HandleMessage() call until Close() is called, and then
+// asserting that we logged the message about dropping a payload to exit.
+func TestGracefulShutdown(t *testing.T) {
+	lggr, hook := logger.TestObserved(t, zapcore.InfoLevel)
+	currentTime := time.Now().UTC()
+	mockTimeProvider := mocks.NewMockTimeProvider(t)
+	mockTimeProvider.EXPECT().GetTime().Return(currentTime).Maybe()
+
+	seqNum := uint64(0)
+	messageGenerator := func() common.MessageWithMetadata {
+		seqNum++
+		return common.MessageWithMetadata{
+			Message: protocol.Message{
+				DestChainSelector:   1,
+				SourceChainSelector: 2,
+				SequenceNumber:      protocol.SequenceNumber(seqNum),
+			},
+			Metadata: common.MessageMetadata{
+				IngestionTimestamp: currentTime,
+			},
+		}
+	}
+
+	results := make(chan common.MessageWithMetadata, 1)
+	results <- messageGenerator()
+	messageSubscriber := mocks.NewMockMessageSubscriber(t)
+	messageSubscriber.EXPECT().Start(mock.Anything).Return(results, nil, nil)
+
+	unblockHandle := make(chan struct{})
+	mockExecutor := mocks.NewMockExecutor(t)
+	mockExecutor.EXPECT().Start(mock.Anything).Return(nil)
+	mockExecutor.EXPECT().CheckValidMessage(mock.Anything, mock.Anything).Return(nil).Maybe()
+	mockExecutor.EXPECT().HandleMessage(mock.Anything, mock.Anything).Run(func(context.Context, protocol.Message) {
+		<-unblockHandle
+	}).Return(false, nil).Maybe()
+
+	leaderElector := mocks.NewMockLeaderElector(t)
+	leaderElector.EXPECT().GetReadyTimestamp(mock.Anything, mock.Anything, mock.Anything).Return(currentTime).Maybe()
+	leaderElector.EXPECT().GetRetryDelay(mock.Anything).Return(time.Second).Maybe()
+
+	ec, err := executor.NewCoordinator(
+		lggr,
+		mockExecutor,
+		messageSubscriber,
+		leaderElector,
+		monitoring.NewNoopExecutorMonitoring(),
+		8*time.Hour,
+		mockTimeProvider,
+		1,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, ec)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	require.NoError(t, ec.Start(ctx))
+	time.Sleep(1 * time.Second)
+
+	// Block close for 1 second to ensure the processing loop is forced to drop a payload.
+	go func() {
+		time.Sleep(1 * time.Second)
+		close(unblockHandle)
+	}()
+	require.NoError(t, ec.Close())
+
+	// Assert that we logged the message about dropping a payload.
+	found := func() bool {
+		for _, entry := range hook.All() {
+			entryStr := fmt.Sprintf("%+v", entry)
+			if strings.Contains(entryStr, "Processing loop dropping payload to exit") {
+				return true
+			}
+		}
+		return false
+	}
+	require.Eventuallyf(t, found, 2*time.Second, 100*time.Millisecond, "executor coordinator did not stop in time")
+}