diff --git a/Makefile b/Makefile index f06d0aea6..7efa18381 100644 --- a/Makefile +++ b/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 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