Skip to content

Commit 5063ddd

Browse files
authored
fix(consumer): derive controllerCtx from Background, not caller context (#368)
## Summary - `consumer.subscribe` created `controllerCtx` from the caller's `ctx`, which inherits any deadline (e.g. Fx `OnStart`'s 15-second timeout). The consume loop selects on `controllerCtx.Done()` and exits when the deadline fires, silently dropping all subsequent messages. - The subscriber already uses `context.Background()`; this aligns the consumer to match. `Stop()` is the shutdown mechanism, not caller context cancellation. ## Test plan - [x] New test `TestConsumer_ConsumeLoopSurvivesCallerDeadline`: starts the consumer with a 50ms deadline context, waits for expiry, delivers a message, asserts it is processed - [x] All existing consumer tests pass
1 parent 75e20ce commit 5063ddd

2 files changed

Lines changed: 55 additions & 7 deletions

File tree

platform/consumer/consumer.go

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,8 @@ type Consumer interface {
4343
Register(controller Controller) error
4444

4545
// Start subscribes to all registered controllers' topics and begins consuming messages.
46-
// Context is cancelled when the consumer is stopped, the implementation should propagate it to the controllers
47-
// running message processing. The implementation can react immediately to the context cancellation by returning `ctx.Err()` instead of starting the message processing,
48-
// but can also opt out to defer the cancellation after the message processing routine is set up.
46+
// ctx governs only the synchronous subscribe calls; consume loops run independently
47+
// and must be terminated by calling Stop().
4948
// Start() will only be called once at the application startup, so it does not need to be idempotent.
5049
Start(ctx context.Context) error
5150

@@ -191,8 +190,8 @@ func (m *consumer) subscribe(ctx context.Context, controller Controller) error {
191190
return fmt.Errorf("subscribe failed: %w", err)
192191
}
193192

194-
// Create cancellable context for this controller
195-
controllerCtx, cancel := context.WithCancel(ctx)
193+
// Manage the controller lifecycle independently of the caller's context.
194+
controllerCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
196195

197196
// Track active subscription
198197
done := make(chan struct{})
@@ -227,7 +226,7 @@ func (m *consumer) subscribe(ctx context.Context, controller Controller) error {
227226
// └── processPartition("part-N")
228227
//
229228
// Shutdown sequence:
230-
// 1. ctx is cancelled (by Stop or parent context)
229+
// 1. ctx is cancelled (by Stop)
231230
// 2. consumeLoop exits the select loop and runs the deferred cleanup
232231
// 3. All partition channels are closed, causing processPartition goroutines to
233232
// drain remaining buffered messages and return (range loop ends)
@@ -322,7 +321,7 @@ func (m *consumer) shutdownPartitions(partitionChs map[string]chan extqueue.Deli
322321
//
323322
// The loop exits when either:
324323
// - deliveryCh is closed (consumeLoop cleanup)
325-
// - ctx is cancelled (graceful shutdown)
324+
// - ctx is cancelled (by Stop)
326325
//
327326
// On context cancellation, the current delivery being read from the channel is
328327
// dropped without processing. This is safe because the queue's visibility timeout

platform/consumer/consumer_test.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -811,3 +811,52 @@ func TestConsumer_PartitionWorkerCleanup(t *testing.T) {
811811
err = c.Stop(30000)
812812
require.NoError(t, err)
813813
}
814+
815+
func TestConsumer_ConsumeLoopSurvivesCallerDeadline(t *testing.T) {
816+
ctrl := gomock.NewController(t)
817+
logger := zaptest.NewLogger(t).Sugar()
818+
819+
deliveryChan := make(chan extqueue.Delivery, 1)
820+
mockSub := queuemock.NewMockSubscriber(ctrl)
821+
mockSub.EXPECT().Subscribe(gomock.Any(), gomock.Any(), gomock.Any()).Return(deliveryChan, nil)
822+
823+
mockQ := queuemock.NewMockQueue(ctrl)
824+
mockQ.EXPECT().Subscriber().Return(mockSub)
825+
826+
reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group")
827+
828+
c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor())
829+
830+
processed := make(chan string, 1)
831+
handler := consumermock.NewMockController(ctrl)
832+
setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group",
833+
func(ctx context.Context, delivery consumer.Delivery) error {
834+
processed <- delivery.Message().ID
835+
return nil
836+
},
837+
)
838+
839+
err := c.Register(handler)
840+
require.NoError(t, err)
841+
842+
// Start with a context that expires quickly, simulating an Fx OnStart hook.
843+
startCtx, startCancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
844+
defer startCancel()
845+
846+
err = c.Start(startCtx)
847+
require.NoError(t, err)
848+
849+
<-startCtx.Done()
850+
851+
msg := entityqueue.NewMessage("after-deadline", []byte("payload"), "partition1", nil)
852+
mockDel := queuemock.NewMockDelivery(ctrl)
853+
done := setupDelivery(mockDel, msg, nil, nil)
854+
855+
deliveryChan <- mockDel
856+
<-done
857+
858+
assert.Equal(t, "after-deadline", <-processed)
859+
860+
err = c.Stop(30000)
861+
require.NoError(t, err)
862+
}

0 commit comments

Comments
 (0)