fix: flowcontrol - finalize requests during shutdown (#2156)

fixes #2099

Signed-off-by: nicole-lihui <nicole.li@daocloud.io>

Author: nicole-lihui

pkg/epp/flowcontrol/controller/controller.go

@@ -378,8 +378,15 @@ func (fc *FlowController) tryDistribution(
 	return item, finalErr
 }
 
+func finalizeOnControllerShutdown(item *internal.FlowItem) (types.QueueOutcome, error) {
+	item.Finalize(types.ErrFlowControllerNotRunning)
+
+	finalState := item.FinalState()
+	return finalState.Outcome, finalState.Err
+}
+
 // awaitFinalization blocks until an item is finalized, either by the processor (synchronously) or by the controller
-// itself due to context expiry (asynchronously).
+// itself due to context expiry or shutdown (asynchronously).
 func (fc *FlowController) awaitFinalization(
 	reqCtx context.Context,
 	item *internal.FlowItem,
@@ -388,13 +395,20 @@ func (fc *FlowController) awaitFinalization(
 	case <-reqCtx.Done():
 		// Asynchronous Finalization (Controller-initiated):
 		// The request Context expired (Cancellation/TTL) while the item was being processed.
+		if fc.parentCtx.Err() != nil {
+			return finalizeOnControllerShutdown(item)
+		}
+
 		cause := context.Cause(reqCtx)
 		item.Finalize(cause)
 
 		// The processor will eventually discard this "zombie" item during its cleanup sweep.
 		finalState := item.FinalState()
 		return finalState.Outcome, finalState.Err
 
+	case <-fc.parentCtx.Done():
+		return finalizeOnControllerShutdown(item)
+
 	case finalState := <-item.Done():
 		// Synchronous Finalization (Processor-initiated):
 		// The processor finalized the item (Dispatch, Reject, Shutdown).

pkg/epp/flowcontrol/controller/controller_test.go

@@ -367,6 +367,71 @@ func TestFlowController_EnqueueAndWait(t *testing.T) {
 			assert.Equal(t, types.QueueOutcomeRejectedOther, outcome,
 				"outcome should be QueueOutcomeRejectedOther on shutdown")
 		})
+		t.Run("OnControllerShutdownDuringFinalization", func(t *testing.T) {
+			t.Parallel()
+			ctx, cancel := context.WithCancel(t.Context())
+			h := newUnitHarness(ctx, t, &Config{}, nil, nil)
+			item := internal.NewItem(newTestRequest(defaultFlowKey), 0, time.Now())
+
+			result := make(chan struct {
+				outcome types.QueueOutcome
+				err     error
+			}, 1)
+			go func() {
+				outcome, err := h.fc.awaitFinalization(context.Background(), item)
+				result <- struct {
+					outcome types.QueueOutcome
+					err     error
+				}{outcome: outcome, err: err}
+			}()
+
+			cancel()
+			select {
+			case r := <-result:
+				require.Error(t, r.err, "awaitFinalization must fail when controller shuts down")
+				assert.ErrorIs(t, r.err, types.ErrFlowControllerNotRunning,
+					"error should wrap ErrFlowControllerNotRunning")
+				assert.Equal(t, types.QueueOutcomeRejectedOther, r.outcome,
+					"outcome should be QueueOutcomeRejectedOther on shutdown")
+			case <-time.After(time.Second):
+				t.Fatal("awaitFinalization did not return after controller shutdown")
+			}
+		})
+		t.Run("OnControllerShutdownTakesPrecedenceOverRequestCancellation", func(t *testing.T) {
+			t.Parallel()
+			ctx, cancel := context.WithCancel(t.Context())
+			h := newUnitHarness(ctx, t, &Config{}, nil, nil)
+			reqCtx, reqCancel := context.WithCancel(context.Background())
+			item := internal.NewItem(newTestRequest(defaultFlowKey), 0, time.Now())
+
+			reqCancel()
+			cancel()
+
+			outcome, err := h.fc.awaitFinalization(reqCtx, item)
+			require.Error(t, err, "awaitFinalization must fail when controller shuts down")
+			assert.ErrorIs(t, err, types.ErrFlowControllerNotRunning,
+				"controller shutdown should take precedence over request cancellation")
+			assert.Equal(t, types.QueueOutcomeRejectedOther, outcome,
+				"shutdown should return the rejected outcome")
+		})
+		t.Run("OnControllerShutdownPreservesQueuedOutcome", func(t *testing.T) {
+			t.Parallel()
+			ctx, cancel := context.WithCancel(t.Context())
+			h := newUnitHarness(ctx, t, &Config{}, nil, nil)
+			item := internal.NewItem(newTestRequest(defaultFlowKey), 0, time.Now())
+			item.SetHandle(&fwkfcmocks.MockQueueItemHandle{})
+
+			cancel()
+
+			outcome, err := h.fc.awaitFinalization(context.Background(), item)
+			require.Error(t, err, "awaitFinalization must fail when controller shuts down")
+			assert.ErrorIs(t, err, types.ErrEvicted,
+				"a queued item should be evicted, not rejected, during shutdown")
+			assert.ErrorIs(t, err, types.ErrFlowControllerNotRunning,
+				"queued shutdown should preserve the shutdown cause")
+			assert.Equal(t, types.QueueOutcomeEvictedOther, outcome,
+				"a queued item should return the evicted outcome")
+		})
 
 		t.Run("OnRegistryConnectionError", func(t *testing.T) {
 			t.Parallel()

pkg/epp/flowcontrol/integration_test.go

@@ -782,6 +782,11 @@ func TestConcurrentEnqueueDuringShutdown(t *testing.T) {
 
 	h := newHarness(t, harnessOpts{
 		detector: detector,
+		controllerCfg: &controller.Config{
+			DefaultRequestTTL:        0,
+			ExpiryCleanupInterval:    10 * time.Millisecond,
+			EnqueueChannelBufferSize: 100,
+		},
 	})
 
 	key := flowcontrol.FlowKey{ID: "flow-a", Priority: 0}
@@ -791,8 +796,8 @@ func TestConcurrentEnqueueDuringShutdown(t *testing.T) {
 	for i := 0; i < numRequests; i++ {
 		id := fmt.Sprintf("req-%d", i)
 		go func() {
-			req := &testRequest{id: id, key: key, byteSize: 100, ttl: 5 * time.Minute}
-			outcome, err := h.fc.EnqueueAndWait(h.ctx, req)
+			req := &testRequest{id: id, key: key, byteSize: 100}
+			outcome, err := h.fc.EnqueueAndWait(context.Background(), req)
 			results <- dispatchResult{id: id, outcome: outcome, err: err}
 		}()
 	}
@@ -809,6 +814,8 @@ func TestConcurrentEnqueueDuringShutdown(t *testing.T) {
 				"no request should dispatch (detector is blocked and controller is shutting down)")
 			require.Error(t, r.err,
 				"every request should receive an error during shutdown")
+			require.ErrorIs(t, r.err, fcTypes.ErrFlowControllerNotRunning,
+				"every request should report the shutdown cause")
 		case <-time.After(5 * time.Second):
 			t.Fatalf("request %d hung during concurrent shutdown", i)
 		}