From 58b5493ab4f415aa23e695bd6b6a426f7776f134 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 21 Aug 2026 10:46:54 -0700 Subject: [PATCH] fix(messagequeue): bound a subscription to the caller's context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? The `Subscriber` interface documents that the delivery channel closes when the subscriber is closed *or the context is cancelled*. The MySQL subscriber — the only implementation — honoured just the first half: `Subscribe`'s `ctx` appeared solely in the signature and was never read, and the subscription instead ran on an independent `context.Background()` that only `Close` could cancel. That inert parameter is what let the package's own test leak. `TestSubscriber_Subscribe` called `Subscribe` with no paired `Close`, leaving a `managePartitions` supervisor per topic running on a one-second discovery ticker. Under the race detector the suite runs slowly enough for a tick to land after the subtest has returned, calling `GetLeasedPartitions` against a finished gomock controller while `zaptest` logs to a completed `*testing.T`. Without `-race` the suite finishes inside a second and the tick never fires, so the failure surfaces only under `--@rules_go//go/config:race` — which CI does not run. ### What? `Subscribe` derives the subscription context from the caller's `ctx`. `Close` still cancels that context directly, so a caller whose `ctx` never completes is unaffected. Honouring `ctx` opens a hazard that could not previously exist: cancellation ends the supervisor and closes `deliveryCh`, but leaves the entry in the subscriptions map, so the next `Subscribe` on that key would hand a new caller an already-closed channel. Each subscription now carries a `done` channel, closed by the supervisor on its way out; `Subscribe` treats an entry whose `done` is closed as absent and replaces it. That signal has to be a channel rather than a lock-guarded flag, because `Close` holds `subMu` across `cancelFunc` and `wg.Wait` for every subscription — acquiring the same mutex on the shutdown path would deadlock against the very `Close` that triggered it. The consumer now hands `Subscribe` a detached context. Every service starts its consumers with a SIGTERM-cancelled context, and feeding that straight through would close the delivery channel the moment shutdown began, cutting the consume loop's drain short. The consume loop already ran detached for exactly that reason, so the two now share one context. Production shutdown behaviour is unchanged. `TestSubscriber_Subscribe` now closes its subscriber, and two new tests cover the semantics: cancelling the context closes the delivery channel, and a stale entry is replaced rather than reused. `make test-race` runs the unit suite under the detector. It needs `--build_tests_only` — the `//...` pattern otherwise pulls in the cross-compiled `*_linux` binaries, which are built without cgo, and race instrumentation requires cgo. ## Test Plan Reproduced first: three of five runs of the filtered test failed with `WARNING: DATA RACE`, the reader being the `managePartitions` goroutine reaching `GetLeasedPartitions` while the subtest goroutine had already finished. ✅ 10/10 runs clean under `--@rules_go//go/config:race` for both `//platform/extension/messagequeue/mysql` and `//platform/consumer`; confirmed with `-test.v` that the new tests actually ran rather than silently matching nothing. ✅ `make test-race` — 109/109 targets pass repo-wide, so no other package carried a latent leak. ✅ `make test` (109/109), gazelle produced no BUILD changes, license headers clean, mocks unchanged. ✅ Integration, uncached, covering the real subscribe/close and shutdown-ordering paths: `//test/integration/submitqueue/core/consumer` and `//test/integration/extension/messagequeue/mysql`. Checked that the new guard is meaningful: with the stale-entry eviction removed, `TestSubscriber_SubscribeReplacesStaleSubscription` fails on both assertions and passes again once restored. --- Makefile | 8 ++- platform/consumer/consumer.go | 11 ++- .../messagequeue/mysql/subscriber.go | 43 ++++++++--- .../messagequeue/mysql/subscriber_test.go | 72 +++++++++++++++++++ 4 files changed, 122 insertions(+), 12 deletions(-) diff --git a/Makefile b/Makefile index 8f3a57c83..0a0b0161b 100644 --- a/Makefile +++ b/Makefile @@ -145,7 +145,7 @@ define assert_clean fi endef -.PHONY: build build-all-linux build-runway-linux build-submitqueue-gateway-client build-submitqueue-gateway-linux build-submitqueue-gateway-server build-submitqueue-orchestrator-linux build-stovepipe-linux build-stovepipe-linux-debug check-gazelle check-mocks check-tidy clean clean-proto demo-requests deps e2e-test fmt gazelle integration-test integration-test-submitqueue-consumer integration-test-extensions integration-test-submitqueue-gateway integration-test-submitqueue-orchestrator license-fix lint lint-binary lint-fmt lint-license local-init-runway-queue-schema local-init-stovepipe-schemas local-runway-start local-runway-stop local-submitqueue-stop local-submitqueue-clean local-submitqueue-gateway-start local-submitqueue-gateway-stop local-init-submitqueue-schemas local-submitqueue-logs local-submitqueue-orchestrator-start local-submitqueue-orchestrator-stop local-submitqueue-ps local-submitqueue-restart local-submitqueue-start local-stop local-stovepipe-debug-start local-stovepipe-logs local-stovepipe-start local-stovepipe-stop mocks proto query-deps query-targets run-client-runway run-client-submitqueue-gateway run-client-submitqueue-orchestrator run-client-stovepipe run-queue-admin test test-no-cache tidy tidy-bazel tidy-go help +.PHONY: build build-all-linux build-runway-linux build-submitqueue-gateway-client build-submitqueue-gateway-linux build-submitqueue-gateway-server build-submitqueue-orchestrator-linux build-stovepipe-linux build-stovepipe-linux-debug check-gazelle check-mocks check-tidy clean clean-proto demo-requests deps e2e-test fmt gazelle integration-test integration-test-submitqueue-consumer integration-test-extensions integration-test-submitqueue-gateway integration-test-submitqueue-orchestrator license-fix lint lint-binary lint-fmt lint-license local-init-runway-queue-schema local-init-stovepipe-schemas local-runway-start local-runway-stop local-submitqueue-stop local-submitqueue-clean local-submitqueue-gateway-start local-submitqueue-gateway-stop local-init-submitqueue-schemas local-submitqueue-logs local-submitqueue-orchestrator-start local-submitqueue-orchestrator-stop local-submitqueue-ps local-submitqueue-restart local-submitqueue-start local-stop local-stovepipe-debug-start local-stovepipe-logs local-stovepipe-start local-stovepipe-stop mocks proto query-deps query-targets run-client-runway run-client-submitqueue-gateway run-client-submitqueue-orchestrator run-client-stovepipe run-queue-admin test test-no-cache test-race tidy tidy-bazel tidy-go help build: ## Build all services and examples @@ -624,6 +624,12 @@ test-no-cache: ## Run unit tests without cache (force re-run) @echo "Running unit tests (no cache)..." @$(BAZEL) test //... --test_tag_filters=-manual,-integration --nocache_test_results +test-race: ## Run unit tests with the race detector (not run in CI) + @echo "Running unit tests with race detector..." + @# build_tests_only keeps the cross-compiled *_linux binaries out of the + @# build: they are cgo-free, and race instrumentation requires cgo. + @$(BAZEL) test //... --test_tag_filters=-manual,-integration --build_tests_only --@rules_go//go/config:race + tidy: tidy-go tidy-bazel ## Run go mod tidy and bazel mod tidy tidy-bazel: ## Run bazel mod tidy diff --git a/platform/consumer/consumer.go b/platform/consumer/consumer.go index c7888d264..753620db5 100644 --- a/platform/consumer/consumer.go +++ b/platform/consumer/consumer.go @@ -207,13 +207,20 @@ func (m *consumer) subscribe(ctx context.Context, controller Controller) error { } subscriber := q.Subscriber() - deliveryChan, err := subscriber.Subscribe(ctx, topicName, config) + + // Subscribe ties the subscription's lifetime to the ctx it receives, so the + // raw ctx must not reach it: that ctx is cancelled on SIGTERM, which would + // close deliveryChan the moment shutdown begins and cut consumeLoop's drain + // short. Detaching keeps teardown ordered by Stop and subscriber.Close. + detachedCtx := context.WithoutCancel(ctx) + + deliveryChan, err := subscriber.Subscribe(detachedCtx, topicName, config) if err != nil { return fmt.Errorf("subscribe failed: %w", err) } // Manage the controller lifecycle independently of the caller's context. - controllerCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + controllerCtx, cancel := context.WithCancel(detachedCtx) // Track active subscription done := make(chan struct{}) diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index 8140ca57a..3b732c581 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -126,6 +126,14 @@ type subscription struct { // Close() waits on this to know the entire subscription is shut down. wg sync.WaitGroup + // done is closed once managePartitions has exited, which implies + // deliveryCh is already closed. Subscribe consults it to tell a live + // subscription from one whose ctx was cancelled independently of Close, + // so it can replace the stale entry instead of handing a new caller a + // closed channel. A channel rather than a flag: signalling it must not + // require a lock (see managePartitions). + done chan struct{} + // workerWg tracks all partition worker goroutines independently of wg. // During shutdown, managePartitions waits on workerWg before closing // deliveryCh to guarantee no worker can send on a closed channel. @@ -472,7 +480,10 @@ func (s *subscriber) advanceWatermark(ctx context.Context, consumerGroup, topic, return nil } -// Subscribe starts consuming messages from the specified topic +// Subscribe starts consuming messages from the specified topic. The returned +// channel is closed when ctx is cancelled or Close is called, whichever comes +// first. Consumption that must outlive a request-scoped ctx and be torn down +// only by Close therefore needs a detached context (context.WithoutCancel). func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueue.SubscriptionConfig) (_ <-chan extqueue.Delivery, retErr error) { op := metrics.Begin(s.scope, "subscribe", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -491,10 +502,18 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu s.subMu.Lock() defer s.subMu.Unlock() - // Check if already subscribed + // A subscription whose supervisor already exited (its ctx was cancelled + // without Close, which would have cleared the map) has a closed + // deliveryCh. Evicting it here and falling through to build a fresh one + // keeps that closed channel from being handed to the next caller. if sub, exists := s.subscriptions[subKey]; exists { - s.logger.Debugw("reusing existing subscription", "topic", topic, "consumer_group", config.ConsumerGroup) - return sub.deliveryCh, nil + select { + case <-sub.done: + delete(s.subscriptions, subKey) + default: + s.logger.Debugw("reusing existing subscription", "topic", topic, "consumer_group", config.ConsumerGroup) + return sub.deliveryCh, nil + } } s.logger.Infow("creating new subscription", @@ -505,15 +524,16 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu "batch_size", config.BatchSize, ) - // Create new subscription - // Use a cancellable context for managing the subscription lifecycle - // and close will cancel the context to signal goroutines to stop - subCtx, cancel := context.WithCancel(context.Background()) + // Derived from the caller's ctx, so cancelling ctx tears this subscription + // down exactly as Close does. Close cancels subCtx directly and so remains + // effective for a caller whose ctx never completes. + subCtx, cancel := context.WithCancel(ctx) sub := &subscription{ topic: topic, config: config, deliveryCh: make(chan extqueue.Delivery, config.BatchSize*2), cancelFunc: cancel, + done: make(chan struct{}), workers: make(map[string]*partitionWorker), } @@ -549,9 +569,14 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu // 3. workerWg.Wait(): blocks until all workers have fully exited -- this ensures // no worker can send on deliveryCh after step 4 // 4. close(deliveryCh): safe because step 3 guarantees no senders remain -// 5. managePartitions returns -> wg.Done() fires -> Close() unblocks +// 5. managePartitions returns -> done and wg.Done() fire -> Close() unblocks func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { defer sub.wg.Done() + // Deferred so every exit path marks the subscription stale for Subscribe. + // Must not take s.subMu: Close holds it across cancelFunc()+wg.Wait() for + // every subscription, so locking here would deadlock against the very + // Close that triggered this shutdown. + defer close(sub.done) cfg := sub.config // Common log fields for all operations in this subscription's lifecycle. diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index b7d08e212..14777d00b 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -114,7 +114,18 @@ func TestSubscriber_Subscribe(t *testing.T) { mockOffsetStore := NewMockoffsetStore(ctrl) mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + // Reached via releaseAllLeases on the shutdown path, and by the + // discovery ticker if it fires before teardown. + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) + // Close waits for managePartitions to exit; a bare cancel would only + // signal it. Registered after the ctrl.Finish defer above so LIFO + // runs it first — gomock fails a mock call made after Finish. + defer func() { + require.NoError(t, sub.Close()) + }() + ctx := context.Background() cfg := testSubscriptionConfig() @@ -133,6 +144,67 @@ func TestSubscriber_Subscribe(t *testing.T) { } } +func TestSubscriber_SubscribeContextCancellation(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockMessageStore := NewMockmessageStore(ctrl) + mockOffsetStore := NewMockoffsetStore(ctrl) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + + sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) + defer func() { + require.NoError(t, sub.Close()) + }() + + ctx, cancel := context.WithCancel(context.Background()) + ch, err := sub.Subscribe(ctx, "test_topic", testSubscriptionConfig()) + require.NoError(t, err) + + cancel() + + // Closing deliveryCh is the last step of the shutdown path, so a blocking + // receive is the synchronization point (test timeout handles a hang). + _, ok := <-ch + assert.False(t, ok, "delivery channel should close when the subscribe context is cancelled") +} + +func TestSubscriber_SubscribeReplacesStaleSubscription(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockMessageStore := NewMockmessageStore(ctrl) + mockOffsetStore := NewMockoffsetStore(ctrl) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + + sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) + defer func() { + require.NoError(t, sub.Close()) + }() + + cfg := testSubscriptionConfig() + + ctx, cancel := context.WithCancel(context.Background()) + stale, err := sub.Subscribe(ctx, "test_topic", cfg) + require.NoError(t, err) + + cancel() + _, ok := <-stale + require.False(t, ok) + + fresh, err := sub.Subscribe(context.Background(), "test_topic", cfg) + require.NoError(t, err) + assert.NotEqual(t, stale, fresh, "a new subscription must not reuse the stale closed channel") + + select { + case _, ok := <-fresh: + t.Fatalf("new subscription channel should be open, receive returned ok=%v", ok) + default: + } +} + func TestSQLDelivery_Ack(t *testing.T) { tests := []struct { name string