fix(messagequeue): bound a subscription to the caller's context

Makefile

@@ -148,7 +148,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
@@ -628,6 +628,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

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{})

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.

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