diff --git a/Makefile b/Makefile index d80bd17e..831ed58b 100644 --- a/Makefile +++ b/Makefile @@ -364,7 +364,7 @@ local-stovepipe-stop: ## Stop the Stovepipe service mocks: ## Generate mock files using mockgen @echo "Generating mocks..." - @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./platform/extension/counter/... ./platform/extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/extension/speculation/enumerator/... ./submitqueue/extension/speculation/dependencylimit/... ./submitqueue/extension/speculation/selector/... ./submitqueue/extension/speculation/selectionlimit/... ./submitqueue/extension/speculation/prioritizer/... ./submitqueue/extension/speculation/prioritizationlimit/... ./submitqueue/extension/validator/... ./submitqueue/extension/speculation/pathscorer/... ./platform/consumer/... ./stovepipe/extension/storage/... ./stovepipe/extension/sourcecontrol/... + @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./platform/extension/counter/... ./platform/extension/consumergate/... ./platform/extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/extension/speculation/enumerator/... ./submitqueue/extension/speculation/dependencylimit/... ./submitqueue/extension/speculation/selector/... ./submitqueue/extension/speculation/selectionlimit/... ./submitqueue/extension/speculation/prioritizer/... ./submitqueue/extension/speculation/prioritizationlimit/... ./submitqueue/extension/validator/... ./submitqueue/extension/speculation/pathscorer/... ./platform/consumer/... ./stovepipe/extension/storage/... ./stovepipe/extension/sourcecontrol/... @echo "Mocks generated successfully!" proto: ## Generate protobuf files from .proto definitions diff --git a/platform/consumer/BUILD.bazel b/platform/consumer/BUILD.bazel index c0e54aa7..44c26168 100644 --- a/platform/consumer/BUILD.bazel +++ b/platform/consumer/BUILD.bazel @@ -5,6 +5,7 @@ go_library( srcs = [ "consumer.go", "controller.go", + "gate.go", "registry.go", ], importpath = "github.com/uber/submitqueue/platform/consumer", @@ -12,6 +13,7 @@ go_library( deps = [ "//platform/base/messagequeue:go_default_library", "//platform/errs:go_default_library", + "//platform/extension/consumergate:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/metrics:go_default_library", "@com_github_uber_go_tally//:go_default_library", @@ -23,13 +25,17 @@ go_test( name = "go_default_test", srcs = [ "consumer_test.go", + "gate_internal_test.go", + "gate_test.go", "registry_test.go", ], + embed = [":go_default_library"], deps = [ - ":go_default_library", "//platform/base/messagequeue:go_default_library", "//platform/consumer/mock:go_default_library", "//platform/errs:go_default_library", + "//platform/extension/consumergate:go_default_library", + "//platform/extension/consumergate/noop:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mock:go_default_library", "//submitqueue/core/topickey:go_default_library", diff --git a/platform/consumer/consumer.go b/platform/consumer/consumer.go index 93c19ffe..c83482f1 100644 --- a/platform/consumer/consumer.go +++ b/platform/consumer/consumer.go @@ -23,6 +23,7 @@ import ( "github.com/uber-go/tally" "github.com/uber/submitqueue/platform/errs" + "github.com/uber/submitqueue/platform/extension/consumergate" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" "github.com/uber/submitqueue/platform/metrics" "go.uber.org/zap" @@ -62,6 +63,7 @@ type consumer struct { metricsScope tally.Scope registry TopicRegistry processor errs.ErrorProcessor + gate consumergate.Gate mu sync.Mutex stopped bool @@ -86,12 +88,18 @@ type activeSubscription struct { // consumers such as DLQ reconciliation that must redeliver on any failure. // processor must not be nil; callers that genuinely want no transformation // can pass errs.NewClassifierProcessor() with no classifiers. -func New(logger *zap.SugaredLogger, scope tally.Scope, registry TopicRegistry, processor errs.ErrorProcessor) Consumer { +// +// gate is the consumer-gate implementation consulted before each delivery +// reaches its controller. Pass noop.New() (from +// platform/extension/consumergate/noop) for services that do not need runtime +// gating. gate must not be nil. +func New(logger *zap.SugaredLogger, scope tally.Scope, registry TopicRegistry, processor errs.ErrorProcessor, gate consumergate.Gate) Consumer { return &consumer{ logger: logger, metricsScope: scope.SubScope("consumer"), registry: registry, processor: processor, + gate: gate, subscriptions: make(map[TopicKey]*activeSubscription), } } @@ -342,6 +350,14 @@ func (m *consumer) processPartition(ctx context.Context, controller Controller, func (m *consumer) processDelivery(ctx context.Context, controller Controller, delivery extqueue.Delivery, controllerScope tally.Scope) { const opName = "process" + // Consumer gate: block the delivery while the controller's gate is closed. + // A false return means the consumer is shutting down while blocked — leave + // the delivery in-flight (no process, no ack/nack) so its visibility lapses + // into a normal redelivery. Gate errors fail open inside waitGate. + if !m.waitGate(ctx, controller, delivery, controllerScope) { + return + } + start := time.Now() metrics.NamedCounter(controllerScope, opName, "messages_received", 1) diff --git a/platform/consumer/consumer_test.go b/platform/consumer/consumer_test.go index 12e3a285..cf03cda4 100644 --- a/platform/consumer/consumer_test.go +++ b/platform/consumer/consumer_test.go @@ -30,6 +30,7 @@ import ( "github.com/uber/submitqueue/platform/consumer" consumermock "github.com/uber/submitqueue/platform/consumer/mock" "github.com/uber/submitqueue/platform/errs" + consumergatenoop "github.com/uber/submitqueue/platform/extension/consumergate/noop" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" "github.com/uber/submitqueue/submitqueue/core/topickey" @@ -92,7 +93,7 @@ func TestNew(t *testing.T) { reg, err := consumer.NewTopicRegistry(nil) require.NoError(t, err) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) require.NotNil(t, c) } @@ -101,7 +102,7 @@ func TestConsumer_Register(t *testing.T) { logger := zaptest.NewLogger(t).Sugar() reg, _ := consumer.NewTopicRegistry(nil) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler1 := consumermock.NewMockController(ctrl) setupController(handler1, "handler1", topickey.TopicKeyStart, "group1", nil) @@ -121,7 +122,7 @@ func TestConsumer_Register_DuplicateTopic(t *testing.T) { logger := zaptest.NewLogger(t).Sugar() reg, _ := consumer.NewTopicRegistry(nil) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler1 := consumermock.NewMockController(ctrl) setupController(handler1, "handler1", topickey.TopicKeyStart, "group1", nil) @@ -141,7 +142,7 @@ func TestConsumer_Register_AfterStop(t *testing.T) { logger := zaptest.NewLogger(t).Sugar() reg, _ := consumer.NewTopicRegistry(nil) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) err := c.Stop(1000) require.NoError(t, err) @@ -157,7 +158,7 @@ func TestConsumer_Start_NoHandlers(t *testing.T) { logger := zaptest.NewLogger(t).Sugar() reg, _ := consumer.NewTopicRegistry(nil) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) err := c.Start(context.Background()) assert.Error(t, err) @@ -168,7 +169,7 @@ func TestConsumer_Start_AfterStop(t *testing.T) { logger := zaptest.NewLogger(t).Sugar() reg, _ := consumer.NewTopicRegistry(nil) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "handler1", topickey.TopicKeyStart, "group1", nil) @@ -194,7 +195,7 @@ func TestConsumer_Start_MissingSubscriptionConfig(t *testing.T) { ) require.NoError(t, err) - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "handler", topickey.TopicKeyStart, "group", nil) @@ -220,7 +221,7 @@ func TestConsumer_Start_SubscribeFailure(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "handler", topickey.TopicKeyStart, "group", nil) @@ -246,7 +247,7 @@ func TestConsumer_ProcessDelivery_Success(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handledMsg := "" handler := consumermock.NewMockController(ctrl) @@ -292,7 +293,7 @@ func TestConsumer_ProcessDelivery_Error(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", @@ -334,7 +335,7 @@ func TestConsumer_ProcessDelivery_NonRetryableError(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", @@ -385,7 +386,7 @@ func TestConsumer_Stop(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", nil) @@ -443,7 +444,7 @@ func TestConsumer_ObservabilityTags(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - testC := consumer.New(logger, testScope, reg, errs.NewClassifierProcessor()) + testC := consumer.New(logger, testScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", @@ -518,7 +519,7 @@ func TestConsumer_AckNackLatencyTracking(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, scope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, scope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", @@ -563,7 +564,7 @@ func TestConsumer_ErrorMetrics(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, scope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, scope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) handler := consumermock.NewMockController(ctrl) setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", @@ -619,7 +620,7 @@ func TestConsumer_PerPartitionProcessing(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) // Track processing by partition partBDone := make(chan struct{}) @@ -704,7 +705,7 @@ func TestConsumer_PartitionOrdering(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) // Mutex + shared slice captures processing order for assertion; // a channel would only signal completion, not record the sequence. @@ -773,7 +774,7 @@ func TestConsumer_PartitionWorkerCleanup(t *testing.T) { reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") - c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor()) + c := consumer.New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) processedCount := int64(0) diff --git a/platform/consumer/gate.go b/platform/consumer/gate.go new file mode 100644 index 00000000..25098e41 --- /dev/null +++ b/platform/consumer/gate.go @@ -0,0 +1,125 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package consumer + +import ( + "context" + "sync" + "time" + + "github.com/uber-go/tally" + "github.com/uber/submitqueue/platform/extension/consumergate" + extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" + "github.com/uber/submitqueue/platform/metrics" +) + +const ( + // parkExtensionMs is the visibility extension applied to a parked delivery + // on each keep-in-flight tick, keeping it in-flight without burning retry + // budget (milliseconds). Must comfortably exceed the extension cadence. + parkExtensionMs = int64(30000) + + // keepInFlightInterval is how often a blocked delivery's visibility is + // extended while the gate holds it. + keepInFlightInterval = 10 * time.Second +) + +// waitGate consults the consumer gate before letting a delivery reach the +// controller. It returns true when the delivery may proceed, false when the +// consumer is shutting down while blocked (the delivery is left in-flight for +// redelivery). Gate errors fail open: the delivery proceeds with a logged +// warning. +func (m *consumer) waitGate(ctx context.Context, controller Controller, delivery extqueue.Delivery, scope tally.Scope) bool { + const opName = "gate" + + msg := delivery.Message() + parked := consumergate.Parked{ + ConsumerGroup: controller.ConsumerGroup(), + Topic: controller.TopicKey().String(), + MessageID: msg.ID, + PartitionKey: msg.PartitionKey, + Payload: msg.Payload, + Attempt: delivery.Attempt(), + } + + stopKeepAlive := m.keepInFlight(ctx, delivery, keepInFlightInterval) + start := time.Now() + err := m.gate.Wait(ctx, parked) + stopKeepAlive() + metrics.NamedTimer(scope, opName, "wait_latency", time.Since(start)) + + if err == nil { + return true + } + + if ctx.Err() != nil { + metrics.NamedCounter(scope, opName, "shutdown_while_blocked", 1) + m.logger.Infow("consumer shutdown while delivery was blocked by gate", + "consumer_group", parked.ConsumerGroup, + "topic", parked.Topic, + "message_id", msg.ID, + ) + return false + } + + // Any other error: fail open — let the delivery through. + metrics.NamedCounter(scope, opName, "wait_errors", 1) + m.logger.Errorw("gate wait failed, failing open", + "consumer_group", parked.ConsumerGroup, + "topic", parked.Topic, + "message_id", msg.ID, + "error", err, + ) + return true +} + +// keepInFlight starts a goroutine that periodically extends a delivery's +// visibility timeout while it is blocked by the gate, so the queue does not +// redeliver a parked message. The returned stop function joins the goroutine +// before returning, guaranteeing no extension races the controller. +func (m *consumer) keepInFlight(ctx context.Context, delivery extqueue.Delivery, interval time.Duration) (stop func()) { + var wg sync.WaitGroup + done := make(chan struct{}) + + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ctx.Done(): + return + case <-ticker.C: + if err := delivery.ExtendVisibilityTimeout(ctx, parkExtensionMs); err != nil { + m.logger.Errorw("failed to extend visibility of parked delivery", + "message_id", delivery.Message().ID, + "error", err, + ) + } + } + } + }() + + var once sync.Once + return func() { + once.Do(func() { + close(done) + wg.Wait() + }) + } +} diff --git a/platform/consumer/gate_internal_test.go b/platform/consumer/gate_internal_test.go new file mode 100644 index 00000000..0fe19a30 --- /dev/null +++ b/platform/consumer/gate_internal_test.go @@ -0,0 +1,77 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package consumer + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/errs" + consumergatenoop "github.com/uber/submitqueue/platform/extension/consumergate/noop" + "go.uber.org/zap/zaptest" +) + +// fakeDeliveryForExtend is a hand-rolled fake delivery that records +// ExtendVisibilityTimeout calls on a channel. +type fakeDeliveryForExtend struct { + msg entityqueue.Message + extended chan int64 +} + +func (d *fakeDeliveryForExtend) Message() entityqueue.Message { return d.msg } +func (d *fakeDeliveryForExtend) Attempt() int { return 1 } +func (d *fakeDeliveryForExtend) ReceivedAt() int64 { return 0 } +func (d *fakeDeliveryForExtend) Metadata() map[string]string { return nil } +func (d *fakeDeliveryForExtend) DeliveryID() string { return d.msg.ID } +func (d *fakeDeliveryForExtend) Ack(_ context.Context) error { return nil } +func (d *fakeDeliveryForExtend) Nack(_ context.Context, _ int64) error { return nil } +func (d *fakeDeliveryForExtend) Reject(_ context.Context, _ string) error { return nil } +func (d *fakeDeliveryForExtend) ExtendVisibilityTimeout(_ context.Context, ms int64) error { + d.extended <- ms + return nil +} + +func TestKeepInFlight_ExtendsAndStopJoins(t *testing.T) { + logger := zaptest.NewLogger(t).Sugar() + reg, err := NewTopicRegistry(nil) + require.NoError(t, err) + + c := New(logger, tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()).(*consumer) + + ctx := context.Background() + del := &fakeDeliveryForExtend{ + msg: entityqueue.NewMessage("msg-1", []byte("p"), "part", nil), + extended: make(chan int64, 16), + } + + stop := c.keepInFlight(ctx, del, 5*time.Millisecond) + + // Wait for at least one extension call. + select { + case ms := <-del.extended: + assert.Equal(t, parkExtensionMs, ms) + case <-time.After(200 * time.Millisecond): + t.Fatal("timed out waiting for ExtendVisibilityTimeout call") + } + + // stop() joins the goroutine before returning — its return is the + // guarantee that no extension can race the controller. + stop() +} diff --git a/platform/consumer/gate_test.go b/platform/consumer/gate_test.go new file mode 100644 index 00000000..5f124c23 --- /dev/null +++ b/platform/consumer/gate_test.go @@ -0,0 +1,291 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package consumer_test + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + consumermock "github.com/uber/submitqueue/platform/consumer/mock" + "github.com/uber/submitqueue/platform/errs" + "github.com/uber/submitqueue/platform/extension/consumergate" + extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" + queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "go.uber.org/mock/gomock" + "go.uber.org/zap/zaptest" +) + +// fakeGate is a channel-instrumented consumergate.Gate so tests can await the +// park/release transitions instead of sleeping. +type fakeGate struct { + mu sync.Mutex + closed map[consumergate.Key]bool + err error + + parked chan consumergate.Parked + released chan string // message IDs +} + +func newFakeGate() *fakeGate { + return &fakeGate{ + closed: make(map[consumergate.Key]bool), + parked: make(chan consumergate.Parked, 16), + released: make(chan string, 16), + } +} + +func (f *fakeGate) setClosed(key consumergate.Key, closed bool) { + f.mu.Lock() + defer f.mu.Unlock() + f.closed[key] = closed +} + +func (f *fakeGate) setErr(err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.err = err +} + +func (f *fakeGate) isClosed(consumerGroup, partitionKey string) bool { + if f.closed[consumergate.Key{ConsumerGroup: consumerGroup}] { + return true + } + return f.closed[consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey}] +} + +// Wait implements consumergate.Gate. It checks the err field first. If open, +// returns nil immediately. If closed, sends the Parked on the parked channel +// and polls on a 2ms timer until the gate opens or ctx is cancelled. On open, +// sends the message ID on the released channel and returns nil. +func (f *fakeGate) Wait(ctx context.Context, parked consumergate.Parked) error { + f.mu.Lock() + if f.err != nil { + err := f.err + f.mu.Unlock() + return err + } + closed := f.isClosed(parked.ConsumerGroup, parked.PartitionKey) + f.mu.Unlock() + + if !closed { + return nil + } + + f.parked <- parked + + ticker := time.NewTicker(2 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + f.mu.Lock() + closed = f.isClosed(parked.ConsumerGroup, parked.PartitionKey) + f.mu.Unlock() + if !closed { + f.released <- parked.MessageID + return nil + } + } + } +} + +// startGatedConsumer builds a consumer with the fake gate directly as the 5th +// New arg, one registered mock controller, and a live subscription fed by the +// returned delivery channel. +func startGatedConsumer(t *testing.T, ctrl *gomock.Controller, gate consumergate.Gate, processFunc func(context.Context, consumer.Delivery) error) (consumer.Consumer, chan extqueue.Delivery) { + t.Helper() + + deliveryChan := make(chan extqueue.Delivery, 4) + mockSub := queuemock.NewMockSubscriber(ctrl) + mockSub.EXPECT().Subscribe(gomock.Any(), gomock.Any(), gomock.Any()).Return(deliveryChan, nil) + + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Subscriber().Return(mockSub) + + reg := newRegistry(t, mockQ, topickey.TopicKeyStart, "test-group") + + c := consumer.New(zaptest.NewLogger(t).Sugar(), tally.NoopScope, reg, errs.NewClassifierProcessor(), gate) + + handler := consumermock.NewMockController(ctrl) + setupController(handler, "test-handler", topickey.TopicKeyStart, "test-group", processFunc) + require.NoError(t, c.Register(handler)) + + require.NoError(t, c.Start(context.Background())) + return c, deliveryChan +} + +// gatedDelivery builds a MockDelivery that also tolerates visibility +// extensions while parked. +func gatedDelivery(ctrl *gomock.Controller, msg entityqueue.Message) (*queuemock.MockDelivery, chan struct{}) { + mockDel := queuemock.NewMockDelivery(ctrl) + done := setupDelivery(mockDel, msg, nil, nil) + mockDel.EXPECT().ExtendVisibilityTimeout(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return mockDel, done +} + +func TestConsumer_Gate_OpenGatePassesThrough(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + + handledMsg := "" + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(_ context.Context, delivery consumer.Delivery) error { + handledMsg = delivery.Message().ID + return nil + }) + + msg := entityqueue.NewMessage("msg-1", []byte("payload"), "partition1", nil) + mockDel, done := gatedDelivery(ctrl, msg) + + deliveryChan <- mockDel + <-done + + assert.Equal(t, "msg-1", handledMsg) + assert.Empty(t, gate.parked, "an open gate must not park deliveries") + + require.NoError(t, c.Stop(30000)) +} + +func TestConsumer_Gate_ParksThenReleases(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group"}, true) + + var processed atomic.Bool + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(context.Context, consumer.Delivery) error { + processed.Store(true) + return nil + }) + + msg := entityqueue.NewMessage("msg-1", []byte("payload"), "partition1", nil) + mockDel, done := gatedDelivery(ctrl, msg) + deliveryChan <- mockDel + + // The parked record is written before the gate blocks, so awaiting it + // proves the gate caught the message before the controller saw it. + parked := <-gate.parked + assert.Equal(t, "test-group", parked.ConsumerGroup) + assert.Equal(t, topickey.TopicKeyStart.String(), parked.Topic) + assert.Equal(t, "msg-1", parked.MessageID) + assert.Equal(t, "partition1", parked.PartitionKey) + assert.Equal(t, []byte("payload"), parked.Payload) + assert.Equal(t, 1, parked.Attempt) + assert.False(t, processed.Load(), "controller must not run while its gate is closed") + + // Open the gate: the parked delivery proceeds, the release is recorded, + // and the message is acked. + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group"}, false) + assert.Equal(t, "msg-1", <-gate.released) + <-done + assert.True(t, processed.Load()) + + require.NoError(t, c.Stop(30000)) +} + +func TestConsumer_Gate_PartitionScoped(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group", PartitionKey: "gated-partition"}, true) + + var handled sync.Map + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(_ context.Context, delivery consumer.Delivery) error { + handled.Store(delivery.Message().ID, true) + return nil + }) + + gatedMsg := entityqueue.NewMessage("gated-msg", []byte("p"), "gated-partition", nil) + gatedDel, gatedDone := gatedDelivery(ctrl, gatedMsg) + openMsg := entityqueue.NewMessage("open-msg", []byte("p"), "open-partition", nil) + openDel, openDone := gatedDelivery(ctrl, openMsg) + + deliveryChan <- gatedDel + parked := <-gate.parked + assert.Equal(t, "gated-msg", parked.MessageID) + + // Unrelated traffic keeps flowing through the same controller while one + // partition is parked. + deliveryChan <- openDel + <-openDone + _, ok := handled.Load("open-msg") + assert.True(t, ok) + _, ok = handled.Load("gated-msg") + assert.False(t, ok) + + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group", PartitionKey: "gated-partition"}, false) + <-gatedDone + _, ok = handled.Load("gated-msg") + assert.True(t, ok) + + require.NoError(t, c.Stop(30000)) +} + +func TestConsumer_Gate_ShutdownWhileParked(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group"}, true) + + var processed atomic.Bool + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(context.Context, consumer.Delivery) error { + processed.Store(true) + return nil + }) + + msg := entityqueue.NewMessage("msg-1", []byte("payload"), "partition1", nil) + mockDel, _ := gatedDelivery(ctrl, msg) + deliveryChan <- mockDel + <-gate.parked + + // Stopping while parked must not stall shutdown, must not invoke the + // controller, and must not ack/nack — the delivery is left in-flight for + // redelivery after its visibility lapses. + require.NoError(t, c.Stop(30000)) + assert.False(t, processed.Load()) + assert.Empty(t, gate.released, "a delivery dropped at shutdown is not released") +} + +func TestConsumer_Gate_FailsOpenOnReadError(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + gate.setClosed(consumergate.Key{ConsumerGroup: "test-group"}, true) + gate.setErr(fmt.Errorf("gate medium unavailable")) + + handledMsg := "" + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(_ context.Context, delivery consumer.Delivery) error { + handledMsg = delivery.Message().ID + return nil + }) + + msg := entityqueue.NewMessage("msg-1", []byte("payload"), "partition1", nil) + mockDel, done := gatedDelivery(ctrl, msg) + + deliveryChan <- mockDel + <-done + + assert.Equal(t, "msg-1", handledMsg, "a broken gate medium must not stall the pipeline") + assert.Empty(t, gate.parked) + + require.NoError(t, c.Stop(30000)) +} diff --git a/platform/extension/consumergate/BUILD.bazel b/platform/extension/consumergate/BUILD.bazel new file mode 100644 index 00000000..aa818783 --- /dev/null +++ b/platform/extension/consumergate/BUILD.bazel @@ -0,0 +1,8 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["consumergate.go"], + importpath = "github.com/uber/submitqueue/platform/extension/consumergate", + visibility = ["//visibility:public"], +) diff --git a/platform/extension/consumergate/README.md b/platform/extension/consumergate/README.md new file mode 100644 index 00000000..f78bd8c7 --- /dev/null +++ b/platform/extension/consumergate/README.md @@ -0,0 +1,23 @@ +# Consumer Gate Extension + +Runtime stop/start of individual queue controllers — for deterministic e2e scenario control and for operational pause of a consuming stage — without stopping the service that hosts them. Design: [doc/rfc/consumer-gate.md](../../../doc/rfc/consumer-gate.md). + +## Contract + +A gate is identified by a consumer group (every controller subscribes with a unique one, so it is the controller's stable runtime name), optionally narrowed to a single partition. The gate owns both the blocking mechanism and the parked-delivery observation records. While a gate is closed, its `Wait` method blocks the delivery, stamping `ParkedAtMs` when it actually blocks and `ReleasedAtMs` when the gate opens. The consumer owns queue mechanics (visibility extension via `keepInFlight`, keeping the delivery in-flight without burning retry budget) and the fail-open posture (a `Wait` error lets the delivery through). Stopping is a barrier, not preemption — a message already inside the controller runs to completion. + +The package defines two interfaces plus the `Config`: + +- `Gate` exposes a single `Wait(ctx, Parked) error` method. Polling implementations (see `file/`) re-check gate state on a timer; notification-capable implementations can release the instant the gate opens. The consumer passes `Gate` as a required argument to `consumer.New`; callers that do not need gating wire the `noop/` implementation. +- `Admin` is the write surface tests and tooling use: close a gate, open it, list what a stopped controller is holding. + +Parked records are the "observe" half of stop/observe/start: awaiting one is the only way to *know* a stop caught a specific message (as opposed to the message not having arrived yet), and the `ReleasedAtMs` stamp proves the delivery later proceeded. + +## Failure posture + +Gating is auxiliary: if `Wait` returns a non-ctx error, the consumer logs, counts the error, and lets the delivery through (fail open). A broken gate medium must not become a pipeline stall. + +## Implementations + +- [file/](file/) — gate state as plain files in a shared directory. Pausing a controller is writing a small file, resuming is `rm`, inspecting a paused stage is `ls` and `cat`. In the e2e stack the directory is bind-mounted into every service container, so the test process manipulates gates and reads parked records as local files. Wait polls on a configurable interval with a per-partition cache that avoids a filesystem stat per message. +- [noop/](noop/) — a no-op gate whose Wait returns nil immediately. Wire it in services that do not need runtime gating. diff --git a/platform/extension/consumergate/consumergate.go b/platform/extension/consumergate/consumergate.go new file mode 100644 index 00000000..3c9ff366 --- /dev/null +++ b/platform/extension/consumergate/consumergate.go @@ -0,0 +1,148 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package consumergate defines the consumer-gate extension: runtime stop/start +// of individual queue controllers without stopping the service that hosts them. +// +// A gate is keyed by consumer group (the controller's stable runtime name), +// optionally narrowed to a single partition. While a gate is closed, the gate +// implementation blocks the delivery inside its Wait method — the delivery +// stays in-flight (visibility periodically extended by the consumer, retry +// count untouched) until the gate opens or the consumer shuts down. The gate +// implementation owns both the wait mechanism and the parked-delivery +// observation records: it stamps ParkedAtMs when it actually blocks a delivery +// and ReleasedAtMs when the gate opens. The consumer owns queue mechanics +// (visibility extension) and the fail-open posture (a Wait error lets the +// delivery through). +// +// The package holds the contract only: Gate (the blocking interface consulted +// by the consumer), Admin (the write surface used by tests and tooling), +// Config, and the Factory interface. Implementations live in subdirectories +// (see file/, noop/). See doc/rfc/consumer-gate.md for the design. +package consumergate + +//go:generate mockgen -source=consumergate.go -destination=mock/consumergate_mock.go -package=mock + +import "context" + +// Key identifies a gate: a consumer group, optionally narrowed to one partition. +type Key struct { + // ConsumerGroup is the gated controller's consumer group — its stable runtime name. + ConsumerGroup string + + // PartitionKey optionally narrows the gate to a single partition. + // Empty gates every partition of the consumer group. + PartitionKey string +} + +// Metadata records why a gate was closed, for the operator who finds it later. +type Metadata struct { + // Reason is a human-readable explanation for the closure. + Reason string + + // CreatedBy identifies who or what closed the gate. + CreatedBy string + + // CreatedAtMs is when the gate was closed (Unix milliseconds). + CreatedAtMs int64 +} + +// Parked is the descriptor of one delivery presented to a gate. The consumer +// fills the identity and payload fields; the gate implementation stamps +// ParkedAtMs and ReleasedAtMs (callers leave them zero). +type Parked struct { + // ConsumerGroup is the consumer group whose gate is consulted. + ConsumerGroup string + + // Topic is the topic key (the stable logical name) the delivery was + // consumed from. + Topic string + + // MessageID is the queue message ID of the delivery. + MessageID string + + // PartitionKey is the partition the delivery belongs to. + PartitionKey string + + // Payload is the message payload, recorded so an observer can assert on it. + Payload []byte + + // Attempt is the delivery attempt the message is on. + Attempt int + + // ParkedAtMs is when the delivery was parked (Unix milliseconds). Stamped + // by the gate implementation when it actually blocks; callers leave it zero. + ParkedAtMs int64 + + // ReleasedAtMs is when the parked delivery proceeded into the controller + // (Unix milliseconds). Stamped by the gate implementation when the gate + // opens; zero while still parked. Callers leave it zero. + ReleasedAtMs int64 +} + +// Gate is the blocking interface consulted by the consumer for each delivery. +// Implementations must be safe for concurrent use. +type Gate interface { + // Wait blocks while the gate for the delivery's consumer group and + // partition is closed; returns nil when the delivery may proceed + // (immediately when the gate is open). Implementations own the wait + // mechanism (a polling implementation re-reads gate state on a timer; a + // notification-capable implementation can release the instant the gate + // opens) AND the parked-delivery observation records: an implementation + // records the delivery when it actually blocks and stamps its release + // when it returns. + // + // Returns ctx.Err() when ctx is cancelled while blocked (the delivery + // must not proceed). Any other error means gate state could not be read + // or recorded — the consumer fails open. + Wait(ctx context.Context, parked Parked) error +} + +// Admin is the write surface used by tests and tooling to operate gates and +// inspect what a stopped controller is holding. +type Admin interface { + // Close closes the gate for the key. Closing an already-closed gate + // overwrites its metadata. + Close(ctx context.Context, key Key, meta Metadata) error + + // Open opens the gate for the key. Opening an already-open gate is a no-op. + Open(ctx context.Context, key Key) error + + // ListParked returns every parked-delivery record for the consumer group, + // including records already stamped released. Callers filter by topic, + // message ID, or release state. + ListParked(ctx context.Context, consumerGroup string) ([]Parked, error) +} + +// Config holds the knobs for polling-based gate implementations. +type Config struct { + // PollIntervalMs is the cadence at which polling implementations re-read + // gate state (milliseconds). Notification-capable implementations may + // ignore it. + PollIntervalMs int64 +} + +// DefaultConfig returns the default gate configuration: 1s poll interval. +func DefaultConfig() Config { + return Config{ + PollIntervalMs: 1000, + } +} + +// Factory creates Gate instances for dependency injection. Factory +// implementations live in the wiring layer, not in this package. +type Factory interface { + // For returns a Gate for the given configuration. + For(cfg Config) (Gate, error) +} diff --git a/platform/extension/consumergate/file/BUILD.bazel b/platform/extension/consumergate/file/BUILD.bazel new file mode 100644 index 00000000..a37f3011 --- /dev/null +++ b/platform/extension/consumergate/file/BUILD.bazel @@ -0,0 +1,20 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["store.go"], + importpath = "github.com/uber/submitqueue/platform/extension/consumergate/file", + visibility = ["//visibility:public"], + deps = ["//platform/extension/consumergate:go_default_library"], +) + +go_test( + name = "go_default_test", + srcs = ["store_test.go"], + embed = [":go_default_library"], + deps = [ + "//platform/extension/consumergate:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/platform/extension/consumergate/file/README.md b/platform/extension/consumergate/file/README.md new file mode 100644 index 00000000..31626803 --- /dev/null +++ b/platform/extension/consumergate/file/README.md @@ -0,0 +1,21 @@ +# File-Backed Consumer Gate + +Stores gate state as plain files under a configured root directory. Presence of a gate file means the gate is closed; deleting it opens the gate. See the [extension README](../README.md) and [doc/rfc/consumer-gate.md](../../../../doc/rfc/consumer-gate.md) for the contract and design rationale. + +## Layout + +``` +{dir}/gates/{consumer_group}/all # gates every partition of the controller +{dir}/gates/{consumer_group}/p-{urlenc(partition)} # gates one partition +{dir}/parked/{consumer_group}/{topic}/{urlenc(id)}.json # one parked delivery record +``` + +Partition keys and message IDs may contain `/` (request IDs like `queue/1`), so they are URL-encoded in file names. Gate files carry human-readable JSON metadata (`reason`, `created_by`, `created_at_ms`); parked records carry the payload, attempt, `parked_at_ms`, and a `released_at_ms` stamped when the delivery proceeds. All writes go through temp-file-plus-rename so readers never see partial JSON. + +## Operating it by hand + +Pause a controller: write any JSON to `{dir}/gates/{group}/all`. Resume: `rm` the file. Inspect what a paused stage is holding: `ls`/`cat` under `{dir}/parked/{group}/`. + +## Reach and limits + +The directory is trivially shareable out of process — the e2e stack bind-mounts a host directory into every service container, and the test manipulates gates and reads parked records as local files. A file gates only the replicas that see the directory; fleet-wide pause needs the deployment platform to distribute the file, or a store-backed implementation of the same contract. diff --git a/platform/extension/consumergate/file/store.go b/platform/extension/consumergate/file/store.go new file mode 100644 index 00000000..d5e41955 --- /dev/null +++ b/platform/extension/consumergate/file/store.go @@ -0,0 +1,386 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package file implements the consumergate contract with plain files in a +// shared directory. Presence of a gate file means the gate is closed; deleting +// the file opens it. Layout under the configured root: +// +// gates/{consumer_group}/all gates every partition +// gates/{consumer_group}/p-{urlenc(partition)} gates one partition +// parked/{consumer_group}/{topic}/{urlenc(id)}.json one parked delivery record +// +// Consumer groups and topics are filesystem-safe by the repo's naming rules; +// partition keys and message IDs may contain "/" (request IDs like "queue/1"), +// so they are URL-encoded in file names. Gate files hold human-readable JSON +// metadata so an operator finding a paused controller can tell why. All writes +// go through temp-file-plus-rename so readers never see partial JSON. +// +// Wait consults a per-(consumerGroup, partitionKey) cached verdict with a TTL +// equal to the poll interval. An open gate costs at most one directory stat per +// partition per interval, not per message. When the gate is closed, Wait writes +// the parked record, then polls on a ticker at the poll interval, re-reading +// gate state directly and updating the cache on each tick. When the gate opens, +// it stamps ReleasedAtMs on the record and returns nil. +// +// The medium is deliberately the simplest that satisfies the contract: pausing +// a controller is writing a small file, resuming is rm, inspecting a paused +// stage is ls and cat, and a bind mount makes all of it reachable from outside +// the service process. A file gates only the replicas that see the directory — +// fleet-wide coordination belongs to a future store-backed implementation of +// the same contract. +package file + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/uber/submitqueue/platform/extension/consumergate" +) + +// Store implements consumergate.Gate and consumergate.Admin over a directory. +type Store struct { + dir string + pollInterval time.Duration + + // cache holds the last gate verdict per (consumer group, partition), + // refreshed at most once per pollInterval. Wait therefore does not hit + // the filesystem per message when the gate is open. + mu sync.Mutex + cache map[cacheKey]cacheEntry +} + +// cacheKey identifies one cached gate verdict. +type cacheKey struct { + consumerGroup string + partitionKey string +} + +// cacheEntry is a cached gate verdict and when it was read. +type cacheEntry struct { + gated bool + at time.Time +} + +// Verify interface compliance at compile time. +var ( + _ consumergate.Gate = (*Store)(nil) + _ consumergate.Admin = (*Store)(nil) +) + +// New returns a file-backed consumergate store rooted at dir. The directory +// does not need to exist yet — reads treat a missing tree as "no gates, no +// parked records", and writes create what they need. cfg.PollIntervalMs +// controls how often a blocked delivery re-checks gate state; values <= 0 +// fall back to the default (1s). +func New(dir string, cfg consumergate.Config) *Store { + if cfg.PollIntervalMs <= 0 { + cfg.PollIntervalMs = consumergate.DefaultConfig().PollIntervalMs + } + return &Store{ + dir: dir, + pollInterval: time.Duration(cfg.PollIntervalMs) * time.Millisecond, + cache: make(map[cacheKey]cacheEntry), + } +} + +// gatePath returns the gate file path for a key: the "all" marker when the key +// has no partition, or the partition-scoped "p-..." marker otherwise. +func (s *Store) gatePath(key consumergate.Key) string { + name := "all" + if key.PartitionKey != "" { + name = "p-" + url.QueryEscape(key.PartitionKey) + } + return filepath.Join(s.dir, "gates", key.ConsumerGroup, name) +} + +// parkedPath returns the parked-record file path for one delivery. +func (s *Store) parkedPath(consumerGroup, topic, messageID string) string { + return filepath.Join(s.dir, "parked", consumerGroup, topic, url.QueryEscape(messageID)+".json") +} + +// isGated reports whether deliveries for the consumer group and partition are +// currently gated, either by an all-partitions gate or by a gate scoped to +// exactly this partition. +func (s *Store) isGated(consumerGroup, partitionKey string) (bool, error) { + paths := []string{s.gatePath(consumergate.Key{ConsumerGroup: consumerGroup})} + if partitionKey != "" { + paths = append(paths, s.gatePath(consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey})) + } + for _, p := range paths { + switch _, err := os.Stat(p); { + case err == nil: + return true, nil + case os.IsNotExist(err): + // Not gated by this marker; check the next. + default: + return false, fmt.Errorf("failed to stat gate file %s: %w", p, err) + } + } + return false, nil +} + +// cachedIsGated returns the cached gate verdict for the group/partition, +// refreshing it from the filesystem when older than the poll interval. +func (s *Store) cachedIsGated(consumerGroup, partitionKey string) (bool, error) { + key := cacheKey{consumerGroup: consumerGroup, partitionKey: partitionKey} + now := time.Now() + + s.mu.Lock() + entry, ok := s.cache[key] + s.mu.Unlock() + if ok && now.Sub(entry.at) < s.pollInterval { + return entry.gated, nil + } + + gated, err := s.isGated(consumerGroup, partitionKey) + if err != nil { + return false, err + } + + s.mu.Lock() + s.cache[key] = cacheEntry{gated: gated, at: now} + s.mu.Unlock() + + return gated, nil +} + +// updateCache writes a fresh verdict into the cache. +func (s *Store) updateCache(consumerGroup, partitionKey string, gated bool) { + key := cacheKey{consumerGroup: consumerGroup, partitionKey: partitionKey} + s.mu.Lock() + s.cache[key] = cacheEntry{gated: gated, at: time.Now()} + s.mu.Unlock() +} + +// Wait implements consumergate.Gate. It returns nil immediately when the gate +// is open. When the gate is closed, it writes the parked record (stamping +// ParkedAtMs) and polls on a ticker at the configured poll interval until the +// gate opens (stamping ReleasedAtMs) or ctx is cancelled (returning ctx.Err()). +func (s *Store) Wait(ctx context.Context, parked consumergate.Parked) error { + // Fast path: consult the cache. + gated, err := s.cachedIsGated(parked.ConsumerGroup, parked.PartitionKey) + if err != nil { + return err + } + if !gated { + return nil + } + + // Gate is closed — write the parked record. + parked.ParkedAtMs = time.Now().UnixMilli() + if err := s.recordParked(parked); err != nil { + return err + } + + // Poll until the gate opens or ctx is cancelled. + ticker := time.NewTicker(s.pollInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + + gated, err := s.isGated(parked.ConsumerGroup, parked.PartitionKey) + if err != nil { + return err + } + s.updateCache(parked.ConsumerGroup, parked.PartitionKey, gated) + if !gated { + if err := s.recordReleased(parked.ConsumerGroup, parked.Topic, parked.MessageID, time.Now().UnixMilli()); err != nil { + return err + } + return nil + } + } +} + +// recordParked writes a parked-delivery record. Re-recording the same delivery +// (e.g. after a redelivery) overwrites the previous record. +func (s *Store) recordParked(parked consumergate.Parked) error { + path := s.parkedPath(parked.ConsumerGroup, parked.Topic, parked.MessageID) + if err := writeJSON(path, parkedRecord(parked)); err != nil { + return fmt.Errorf("failed to write parked record %s: %w", path, err) + } + return nil +} + +// recordReleased stamps ReleasedAtMs on the existing parked record. +func (s *Store) recordReleased(consumerGroup, topic, messageID string, releasedAtMs int64) error { + path := s.parkedPath(consumerGroup, topic, messageID) + rec, err := readParked(path) + if err != nil { + return fmt.Errorf("failed to read parked record %s: %w", path, err) + } + rec.ReleasedAtMs = releasedAtMs + if err := writeJSON(path, rec); err != nil { + return fmt.Errorf("failed to write parked record %s: %w", path, err) + } + return nil +} + +// Close implements consumergate.Admin by writing the gate file for the key. +func (s *Store) Close(_ context.Context, key consumergate.Key, meta consumergate.Metadata) error { + if key.ConsumerGroup == "" { + return fmt.Errorf("gate key requires a consumer group") + } + path := s.gatePath(key) + if err := writeJSON(path, gateRecord{ + Reason: meta.Reason, + CreatedBy: meta.CreatedBy, + CreatedAtMs: meta.CreatedAtMs, + }); err != nil { + return fmt.Errorf("failed to write gate file %s: %w", path, err) + } + return nil +} + +// Open implements consumergate.Admin by removing the gate file for the key. +// Opening an already-open gate is a no-op. +func (s *Store) Open(_ context.Context, key consumergate.Key) error { + path := s.gatePath(key) + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to remove gate file %s: %w", path, err) + } + return nil +} + +// ListParked implements consumergate.Admin. It returns every parked record for +// the consumer group across all topics; a missing tree yields an empty list. +func (s *Store) ListParked(_ context.Context, consumerGroup string) ([]consumergate.Parked, error) { + groupDir := filepath.Join(s.dir, "parked", consumerGroup) + topics, err := os.ReadDir(groupDir) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("failed to read parked dir %s: %w", groupDir, err) + } + + var out []consumergate.Parked + for _, topic := range topics { + if !topic.IsDir() { + continue + } + topicDir := filepath.Join(groupDir, topic.Name()) + entries, err := os.ReadDir(topicDir) + if err != nil { + return nil, fmt.Errorf("failed to read parked dir %s: %w", topicDir, err) + } + for _, entry := range entries { + // Skip anything that is not a finished record (e.g. temp files + // awaiting rename). + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { + continue + } + rec, err := readParked(filepath.Join(topicDir, entry.Name())) + if err != nil { + return nil, err + } + out = append(out, consumergate.Parked(rec)) + } + } + return out, nil +} + +// gateRecord is the JSON layout of a gate file. +type gateRecord struct { + // Reason is a human-readable explanation for the closure. + Reason string `json:"reason"` + // CreatedBy identifies who or what closed the gate. + CreatedBy string `json:"created_by"` + // CreatedAtMs is when the gate was closed (Unix milliseconds). + CreatedAtMs int64 `json:"created_at_ms"` +} + +// parkedRecord is the JSON layout of a parked-delivery record. It mirrors +// consumergate.Parked field-for-field; the named type only pins the wire tags. +type parkedRecord struct { + // ConsumerGroup is the consumer group whose gate parked the delivery. + ConsumerGroup string `json:"consumer_group"` + // Topic is the topic name the delivery was consumed from. + Topic string `json:"topic"` + // MessageID is the queue message ID of the parked delivery. + MessageID string `json:"message_id"` + // PartitionKey is the partition the delivery belongs to. + PartitionKey string `json:"partition_key"` + // Payload is the message payload (base64 in the JSON encoding). + Payload []byte `json:"payload"` + // Attempt is the delivery attempt the parked message is on. + Attempt int `json:"attempt"` + // ParkedAtMs is when the delivery was parked (Unix milliseconds). + ParkedAtMs int64 `json:"parked_at_ms"` + // ReleasedAtMs is when the delivery proceeded; zero while parked. + ReleasedAtMs int64 `json:"released_at_ms"` +} + +// readParked loads and decodes one parked-record file. +func readParked(path string) (parkedRecord, error) { + data, err := os.ReadFile(path) + if err != nil { + return parkedRecord{}, fmt.Errorf("failed to read parked record %s: %w", path, err) + } + var rec parkedRecord + if err := json.Unmarshal(data, &rec); err != nil { + return parkedRecord{}, fmt.Errorf("failed to decode parked record %s: %w", path, err) + } + return rec, nil +} + +// writeJSON writes v as indented JSON via temp-file-plus-rename in the target +// directory, so concurrent readers never observe partial content. On any +// failure after the temp file is created, the temp file is removed in a single +// deferred cleanup; removal errors are joined with the causal error. +func writeJSON(path string, v any) (retErr error) { + data, err := json.MarshalIndent(v, "", " ") + if err != nil { + return fmt.Errorf("failed to encode %s: %w", path, err) + } + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("failed to create dir %s: %w", dir, err) + } + tmp, err := os.CreateTemp(dir, filepath.Base(path)+".tmp") + if err != nil { + return fmt.Errorf("failed to create temp file in %s: %w", dir, err) + } + tmpName := tmp.Name() + defer func() { + if retErr != nil { + if rmErr := os.Remove(tmpName); rmErr != nil && !os.IsNotExist(rmErr) { + retErr = errors.Join(retErr, rmErr) + } + } + }() + if _, err := tmp.Write(data); err != nil { + closeErr := tmp.Close() + return errors.Join(fmt.Errorf("failed to write temp file %s: %w", tmpName, err), closeErr) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("failed to close temp file %s: %w", tmpName, err) + } + if err := os.Rename(tmpName, path); err != nil { + return fmt.Errorf("failed to rename %s to %s: %w", tmpName, path, err) + } + return nil +} diff --git a/platform/extension/consumergate/file/store_test.go b/platform/extension/consumergate/file/store_test.go new file mode 100644 index 00000000..c5d0e753 --- /dev/null +++ b/platform/extension/consumergate/file/store_test.go @@ -0,0 +1,332 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package file + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/platform/extension/consumergate" +) + +// testCfg keeps Wait tests fast: 5ms poll interval. +var testCfg = consumergate.Config{PollIntervalMs: 5} + +func TestIsGated(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + close []consumergate.Key + group string + partition string + want bool + }{ + { + name: "no gates", + group: "orchestrator-batch", + partition: "queue-a", + want: false, + }, + { + name: "all-partitions gate matches any partition", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch"}}, + group: "orchestrator-batch", + partition: "queue-a", + want: true, + }, + { + name: "all-partitions gate matches empty partition", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch"}}, + group: "orchestrator-batch", + partition: "", + want: true, + }, + { + name: "partition gate matches its partition", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"}}, + group: "orchestrator-batch", + partition: "queue-a", + want: true, + }, + { + name: "partition gate leaves other partitions open", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"}}, + group: "orchestrator-batch", + partition: "queue-b", + want: false, + }, + { + name: "gate on one group leaves other groups open", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch"}}, + group: "runway-merge", + partition: "queue-a", + want: false, + }, + { + name: "partition key with slash is encoded and matched", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue/1"}}, + group: "orchestrator-batch", + partition: "queue/1", + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store := New(t.TempDir(), testCfg) + for _, key := range tt.close { + require.NoError(t, store.Close(ctx, key, consumergate.Metadata{Reason: "test", CreatedBy: "unit", CreatedAtMs: 1})) + } + got, err := store.isGated(tt.group, tt.partition) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestOpenClosesGate(t *testing.T) { + ctx := context.Background() + store := New(t.TempDir(), testCfg) + key := consumergate.Key{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"} + + require.NoError(t, store.Close(ctx, key, consumergate.Metadata{Reason: "pause", CreatedBy: "unit", CreatedAtMs: 1})) + gated, err := store.isGated(key.ConsumerGroup, key.PartitionKey) + require.NoError(t, err) + require.True(t, gated) + + require.NoError(t, store.Open(ctx, key)) + gated, err = store.isGated(key.ConsumerGroup, key.PartitionKey) + require.NoError(t, err) + assert.False(t, gated) + + // Opening an already-open gate is a no-op. + require.NoError(t, store.Open(ctx, key)) +} + +func TestCloseRequiresConsumerGroup(t *testing.T) { + store := New(t.TempDir(), testCfg) + err := store.Close(context.Background(), consumergate.Key{}, consumergate.Metadata{}) + require.Error(t, err) +} + +func TestParkedRecordLifecycle(t *testing.T) { + ctx := context.Background() + store := New(t.TempDir(), testCfg) + + parked := consumergate.Parked{ + ConsumerGroup: "runway-mergeconflictcheck", + Topic: "merge-conflict-check", + MessageID: "e2e-queue/42", + PartitionKey: "e2e-queue", + Payload: []byte(`{"id":"e2e-queue/42"}`), + Attempt: 1, + ParkedAtMs: 1111, + } + require.NoError(t, store.recordParked(parked)) + + records, err := store.ListParked(ctx, parked.ConsumerGroup) + require.NoError(t, err) + require.Len(t, records, 1) + assert.Equal(t, parked, records[0]) + assert.Zero(t, records[0].ReleasedAtMs) + + // Re-recording the same delivery (redelivery) overwrites, not duplicates. + parked.Attempt = 2 + require.NoError(t, store.recordParked(parked)) + records, err = store.ListParked(ctx, parked.ConsumerGroup) + require.NoError(t, err) + require.Len(t, records, 1) + assert.Equal(t, 2, records[0].Attempt) + + require.NoError(t, store.recordReleased(parked.ConsumerGroup, parked.Topic, parked.MessageID, 2222)) + records, err = store.ListParked(ctx, parked.ConsumerGroup) + require.NoError(t, err) + require.Len(t, records, 1) + assert.Equal(t, int64(2222), records[0].ReleasedAtMs) + assert.Equal(t, parked.Payload, records[0].Payload, "release stamp must preserve the recorded payload") +} + +func TestListParkedEmpty(t *testing.T) { + store := New(t.TempDir(), testCfg) + records, err := store.ListParked(context.Background(), "no-such-group") + require.NoError(t, err) + assert.Empty(t, records) +} + +func TestListParkedSkipsTempFiles(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store := New(dir, testCfg) + + parked := consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "id", + PartitionKey: "part", + ParkedAtMs: 1, + } + require.NoError(t, store.recordParked(parked)) + + // Simulate an in-flight temp file awaiting rename alongside the record. + tmpPath := filepath.Join(dir, "parked", "group", "topic", "id.json.tmp123") + require.NoError(t, os.WriteFile(tmpPath, []byte("partial"), 0o644)) + + records, err := store.ListParked(ctx, "group") + require.NoError(t, err) + assert.Len(t, records, 1) +} + +func TestMissingDirIsNotGated(t *testing.T) { + store := New(filepath.Join(t.TempDir(), "does-not-exist"), testCfg) + gated, err := store.isGated("group", "part") + require.NoError(t, err) + assert.False(t, gated) +} + +func TestWait_OpenGateReturnsImmediately(t *testing.T) { + ctx := context.Background() + store := New(t.TempDir(), testCfg) + + err := store.Wait(ctx, consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "msg-1", + PartitionKey: "part", + Payload: []byte("hello"), + Attempt: 1, + }) + require.NoError(t, err) + + // No parked record should exist — the gate was open. + records, err := store.ListParked(ctx, "group") + require.NoError(t, err) + assert.Empty(t, records) +} + +func TestWait_ClosedGateParksThenReleases(t *testing.T) { + ctx := context.Background() + store := New(t.TempDir(), testCfg) + key := consumergate.Key{ConsumerGroup: "group"} + + require.NoError(t, store.Close(ctx, key, consumergate.Metadata{Reason: "test", CreatedBy: "unit", CreatedAtMs: 1})) + + parked := consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "msg-1", + PartitionKey: "part", + Payload: []byte("hello"), + Attempt: 1, + } + + // Run Wait in a goroutine. + waitDone := make(chan error, 1) + go func() { + waitDone <- store.Wait(ctx, parked) + }() + + // Wait until the parked record appears. + require.Eventually(t, func() bool { + records, err := store.ListParked(ctx, "group") + return err == nil && len(records) == 1 + }, 2*store.pollInterval, store.pollInterval/2) + + // Assert the parked record fields. + records, err := store.ListParked(ctx, "group") + require.NoError(t, err) + require.Len(t, records, 1) + assert.Equal(t, "msg-1", records[0].MessageID) + assert.Equal(t, "topic", records[0].Topic) + assert.Equal(t, "part", records[0].PartitionKey) + assert.Equal(t, []byte("hello"), records[0].Payload) + assert.Equal(t, 1, records[0].Attempt) + assert.NotZero(t, records[0].ParkedAtMs) + assert.Zero(t, records[0].ReleasedAtMs) + + // Assert Wait has not returned yet. + select { + case <-waitDone: + t.Fatal("Wait returned before the gate was opened") + default: + } + + // Open the gate — Wait should return nil and the record should have + // ReleasedAtMs stamped. + require.NoError(t, store.Open(ctx, key)) + require.NoError(t, <-waitDone) + + records, err = store.ListParked(ctx, "group") + require.NoError(t, err) + require.Len(t, records, 1) + assert.NotZero(t, records[0].ReleasedAtMs) +} + +func TestWait_ClosedGateCtxCancel(t *testing.T) { + store := New(t.TempDir(), testCfg) + key := consumergate.Key{ConsumerGroup: "group"} + + ctx, cancel := context.WithCancel(context.Background()) + require.NoError(t, store.Close(ctx, key, consumergate.Metadata{Reason: "test", CreatedBy: "unit", CreatedAtMs: 1})) + + parked := consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "msg-1", + PartitionKey: "part", + Payload: []byte("hello"), + Attempt: 1, + } + + waitDone := make(chan error, 1) + go func() { + waitDone <- store.Wait(ctx, parked) + }() + + // Wait until the parked record appears, then cancel. + require.Eventually(t, func() bool { + records, err := store.ListParked(ctx, "group") + return err == nil && len(records) == 1 + }, 2*store.pollInterval, store.pollInterval/2) + + cancel() + err := <-waitDone + require.ErrorIs(t, err, context.Canceled) + + // The record should remain unreleased. + records, err2 := store.ListParked(context.Background(), "group") + require.NoError(t, err2) + require.Len(t, records, 1) + assert.Zero(t, records[0].ReleasedAtMs) +} + +func TestWait_MediumError(t *testing.T) { + // Make the store root a regular file so stat fails with ENOTDIR. + dir := filepath.Join(t.TempDir(), "not-a-dir") + require.NoError(t, os.WriteFile(dir, []byte("x"), 0o644)) + + store := New(dir, testCfg) + err := store.Wait(context.Background(), consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "msg-1", + PartitionKey: "part", + }) + require.Error(t, err) +} diff --git a/platform/extension/consumergate/mock/BUILD.bazel b/platform/extension/consumergate/mock/BUILD.bazel new file mode 100644 index 00000000..f0b94bf2 --- /dev/null +++ b/platform/extension/consumergate/mock/BUILD.bazel @@ -0,0 +1,12 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["consumergate_mock.go"], + importpath = "github.com/uber/submitqueue/platform/extension/consumergate/mock", + visibility = ["//visibility:public"], + deps = [ + "//platform/extension/consumergate:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/platform/extension/consumergate/mock/consumergate_mock.go b/platform/extension/consumergate/mock/consumergate_mock.go new file mode 100644 index 00000000..3c18f7fa --- /dev/null +++ b/platform/extension/consumergate/mock/consumergate_mock.go @@ -0,0 +1,162 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: consumergate.go +// +// Generated by this command: +// +// mockgen -source=consumergate.go -destination=mock/consumergate_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + consumergate "github.com/uber/submitqueue/platform/extension/consumergate" + gomock "go.uber.org/mock/gomock" +) + +// MockGate is a mock of Gate interface. +type MockGate struct { + ctrl *gomock.Controller + recorder *MockGateMockRecorder + isgomock struct{} +} + +// MockGateMockRecorder is the mock recorder for MockGate. +type MockGateMockRecorder struct { + mock *MockGate +} + +// NewMockGate creates a new mock instance. +func NewMockGate(ctrl *gomock.Controller) *MockGate { + mock := &MockGate{ctrl: ctrl} + mock.recorder = &MockGateMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockGate) EXPECT() *MockGateMockRecorder { + return m.recorder +} + +// Wait mocks base method. +func (m *MockGate) Wait(ctx context.Context, parked consumergate.Parked) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Wait", ctx, parked) + ret0, _ := ret[0].(error) + return ret0 +} + +// Wait indicates an expected call of Wait. +func (mr *MockGateMockRecorder) Wait(ctx, parked any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Wait", reflect.TypeOf((*MockGate)(nil).Wait), ctx, parked) +} + +// MockAdmin is a mock of Admin interface. +type MockAdmin struct { + ctrl *gomock.Controller + recorder *MockAdminMockRecorder + isgomock struct{} +} + +// MockAdminMockRecorder is the mock recorder for MockAdmin. +type MockAdminMockRecorder struct { + mock *MockAdmin +} + +// NewMockAdmin creates a new mock instance. +func NewMockAdmin(ctrl *gomock.Controller) *MockAdmin { + mock := &MockAdmin{ctrl: ctrl} + mock.recorder = &MockAdminMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockAdmin) EXPECT() *MockAdminMockRecorder { + return m.recorder +} + +// Close mocks base method. +func (m *MockAdmin) Close(ctx context.Context, key consumergate.Key, meta consumergate.Metadata) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close", ctx, key, meta) + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockAdminMockRecorder) Close(ctx, key, meta any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockAdmin)(nil).Close), ctx, key, meta) +} + +// ListParked mocks base method. +func (m *MockAdmin) ListParked(ctx context.Context, consumerGroup string) ([]consumergate.Parked, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListParked", ctx, consumerGroup) + ret0, _ := ret[0].([]consumergate.Parked) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListParked indicates an expected call of ListParked. +func (mr *MockAdminMockRecorder) ListParked(ctx, consumerGroup any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListParked", reflect.TypeOf((*MockAdmin)(nil).ListParked), ctx, consumerGroup) +} + +// Open mocks base method. +func (m *MockAdmin) Open(ctx context.Context, key consumergate.Key) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Open", ctx, key) + ret0, _ := ret[0].(error) + return ret0 +} + +// Open indicates an expected call of Open. +func (mr *MockAdminMockRecorder) Open(ctx, key any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockAdmin)(nil).Open), ctx, key) +} + +// MockFactory is a mock of Factory interface. +type MockFactory struct { + ctrl *gomock.Controller + recorder *MockFactoryMockRecorder + isgomock struct{} +} + +// MockFactoryMockRecorder is the mock recorder for MockFactory. +type MockFactoryMockRecorder struct { + mock *MockFactory +} + +// NewMockFactory creates a new mock instance. +func NewMockFactory(ctrl *gomock.Controller) *MockFactory { + mock := &MockFactory{ctrl: ctrl} + mock.recorder = &MockFactoryMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockFactory) EXPECT() *MockFactoryMockRecorder { + return m.recorder +} + +// For mocks base method. +func (m *MockFactory) For(cfg consumergate.Config) (consumergate.Gate, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "For", cfg) + ret0, _ := ret[0].(consumergate.Gate) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// For indicates an expected call of For. +func (mr *MockFactoryMockRecorder) For(cfg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "For", reflect.TypeOf((*MockFactory)(nil).For), cfg) +} diff --git a/platform/extension/consumergate/noop/BUILD.bazel b/platform/extension/consumergate/noop/BUILD.bazel new file mode 100644 index 00000000..8456374a --- /dev/null +++ b/platform/extension/consumergate/noop/BUILD.bazel @@ -0,0 +1,19 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["gate.go"], + importpath = "github.com/uber/submitqueue/platform/extension/consumergate/noop", + visibility = ["//visibility:public"], + deps = ["//platform/extension/consumergate:go_default_library"], +) + +go_test( + name = "go_default_test", + srcs = ["gate_test.go"], + embed = [":go_default_library"], + deps = [ + "//platform/extension/consumergate:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/platform/extension/consumergate/noop/README.md b/platform/extension/consumergate/noop/README.md new file mode 100644 index 00000000..e680ae98 --- /dev/null +++ b/platform/extension/consumergate/noop/README.md @@ -0,0 +1,3 @@ +# No-op Consumer Gate + +A consumergate.Gate whose Wait returns nil immediately — every delivery flows straight to its controller. Wire it in services and tests that do not need runtime gating. diff --git a/platform/extension/consumergate/noop/gate.go b/platform/extension/consumergate/noop/gate.go new file mode 100644 index 00000000..b2a87a35 --- /dev/null +++ b/platform/extension/consumergate/noop/gate.go @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package noop provides a no-op consumergate.Gate whose Wait returns nil +// immediately. Wire it in services and tests that do not need runtime gating. +package noop + +import ( + "context" + + "github.com/uber/submitqueue/platform/extension/consumergate" +) + +// Verify interface compliance at compile time. +var _ consumergate.Gate = Gate{} + +// Gate is a no-op consumer gate. Wait always returns nil immediately. +type Gate struct{} + +// New returns a no-op Gate. +func New() Gate { + return Gate{} +} + +// Wait implements consumergate.Gate. It returns nil immediately — the delivery +// is never gated. +func (Gate) Wait(_ context.Context, _ consumergate.Parked) error { + return nil +} diff --git a/platform/extension/consumergate/noop/gate_test.go b/platform/extension/consumergate/noop/gate_test.go new file mode 100644 index 00000000..77349699 --- /dev/null +++ b/platform/extension/consumergate/noop/gate_test.go @@ -0,0 +1,36 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package noop + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/platform/extension/consumergate" +) + +func TestGate_WaitReturnsNil(t *testing.T) { + g := New() + err := g.Wait(context.Background(), consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "msg-1", + PartitionKey: "part", + Payload: []byte("hello"), + Attempt: 1, + }) + require.NoError(t, err) +} diff --git a/service/runway/server/BUILD.bazel b/service/runway/server/BUILD.bazel index 49067b0f..f4efcf0f 100644 --- a/service/runway/server/BUILD.bazel +++ b/service/runway/server/BUILD.bazel @@ -12,6 +12,8 @@ go_library( "//platform/errs:go_default_library", "//platform/errs/generic:go_default_library", "//platform/errs/mysql:go_default_library", + "//platform/extension/consumergate:go_default_library", + "//platform/extension/consumergate/file:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//runway/controller:go_default_library", diff --git a/service/runway/server/main.go b/service/runway/server/main.go index ee6a0936..6aa3c036 100644 --- a/service/runway/server/main.go +++ b/service/runway/server/main.go @@ -34,6 +34,8 @@ import ( "github.com/uber/submitqueue/platform/errs" genericerrs "github.com/uber/submitqueue/platform/errs/generic" mysqlerrs "github.com/uber/submitqueue/platform/errs/mysql" + "github.com/uber/submitqueue/platform/extension/consumergate" + consumergatefile "github.com/uber/submitqueue/platform/extension/consumergate/file" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" "github.com/uber/submitqueue/runway/controller" @@ -152,6 +154,7 @@ func run() error { genericerrs.Classifier, mysqlerrs.Classifier, ), + newConsumerGate(logger), ) mergerFactory := newMergerFactory() @@ -288,3 +291,23 @@ func newTopicRegistry(q extqueue.Queue, subscriberName string) (consumer.TopicRe }, }) } + +// defaultConsumerGateDir is the path the compose stack bind-mounts for gate +// state files. A missing directory simply means every gate is open. +const defaultConsumerGateDir = "/var/submitqueue/consumergate" + +// getEnv returns environment variable value or default if not set. +func getEnv(key, defaultVal string) string { + if val := os.Getenv(key); val != "" { + return val + } + return defaultVal +} + +// newConsumerGate returns a file-backed consumer gate rooted at the directory +// from CONSUMER_GATE_DIR (defaulting to defaultConsumerGateDir). +func newConsumerGate(logger *zap.Logger) consumergate.Gate { + dir := getEnv("CONSUMER_GATE_DIR", defaultConsumerGateDir) + logger.Info("consumer gate configured", zap.String("dir", dir)) + return consumergatefile.New(dir, consumergate.DefaultConfig()) +} diff --git a/service/stovepipe/server/BUILD.bazel b/service/stovepipe/server/BUILD.bazel index 559e3222..7fd566c4 100644 --- a/service/stovepipe/server/BUILD.bazel +++ b/service/stovepipe/server/BUILD.bazel @@ -11,6 +11,7 @@ go_library( "//platform/errs:go_default_library", "//platform/errs/generic:go_default_library", "//platform/errs/mysql:go_default_library", + "//platform/extension/consumergate/noop:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//service/stovepipe/server/mapper:go_default_library", diff --git a/service/stovepipe/server/main.go b/service/stovepipe/server/main.go index 3c4b5c4c..8b1b1b5a 100644 --- a/service/stovepipe/server/main.go +++ b/service/stovepipe/server/main.go @@ -33,6 +33,7 @@ import ( "github.com/uber/submitqueue/platform/errs" genericerrs "github.com/uber/submitqueue/platform/errs/generic" mysqlerrs "github.com/uber/submitqueue/platform/errs/mysql" + consumergatenoop "github.com/uber/submitqueue/platform/extension/consumergate/noop" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" "github.com/uber/submitqueue/service/stovepipe/server/mapper" @@ -208,12 +209,14 @@ func run() error { return fmt.Errorf("failed to create topic registry: %w", err) } - // Consumer running the process stage. + // Consumer running the process stage. Stovepipe has no gated deployment + // yet, so it uses the no-op gate. primaryConsumer := consumer.New(logger.Sugar(), scope.SubScope("consumer"), registry, errs.NewClassifierProcessor( genericerrs.Classifier, mysqlerrs.Classifier, ), + consumergatenoop.New(), ) processController := process.NewController( diff --git a/service/submitqueue/docker-compose.yml b/service/submitqueue/docker-compose.yml index e8d2c4bd..97a80a31 100644 --- a/service/submitqueue/docker-compose.yml +++ b/service/submitqueue/docker-compose.yml @@ -7,6 +7,12 @@ # # Quick start: # make e2e-test +# +# Consumer gate: every service shares one host directory (bind-mounted at +# /var/submitqueue/consumergate) so tests and operators can stop/start +# individual queue controllers by writing/removing gate files from the host. +# Override the host side with SQ_CONSUMER_GATE_DIR (the e2e suite points it at +# a per-run temp dir); it defaults to /tmp/sq-consumergate for local runs. services: # Application Database - Stores business data (requests, counters, etc.) @@ -57,6 +63,10 @@ services: - QUEUE_CONFIG_PATH=/root/queues.yaml # Stable subscriber name for the request-log consumer - HOSTNAME=gateway-dev + # Consumer-gate state shared with the host (see header comment) + - CONSUMER_GATE_DIR=/var/submitqueue/consumergate + volumes: + - ${SQ_CONSUMER_GATE_DIR:-/tmp/sq-consumergate}:/var/submitqueue/consumergate depends_on: mysql-app: condition: service_healthy @@ -76,6 +86,10 @@ services: # Queue infrastructure connection (separate database) - QUEUE_MYSQL_DSN=root:root@tcp(mysql-queue:3306)/submitqueue?parseTime=true - HOSTNAME=orchestrator-dev + # Consumer-gate state shared with the host (see header comment) + - CONSUMER_GATE_DIR=/var/submitqueue/consumergate + volumes: + - ${SQ_CONSUMER_GATE_DIR:-/tmp/sq-consumergate}:/var/submitqueue/consumergate depends_on: mysql-app: condition: service_healthy @@ -98,6 +112,10 @@ services: # Queue infrastructure connection (shared with the orchestrator) - QUEUE_MYSQL_DSN=root:root@tcp(mysql-queue:3306)/submitqueue?parseTime=true - HOSTNAME=runway-dev + # Consumer-gate state shared with the host (see header comment) + - CONSUMER_GATE_DIR=/var/submitqueue/consumergate + volumes: + - ${SQ_CONSUMER_GATE_DIR:-/tmp/sq-consumergate}:/var/submitqueue/consumergate depends_on: mysql-queue: condition: service_healthy diff --git a/service/submitqueue/gateway/server/BUILD.bazel b/service/submitqueue/gateway/server/BUILD.bazel index 9d24713c..47ac6e93 100644 --- a/service/submitqueue/gateway/server/BUILD.bazel +++ b/service/submitqueue/gateway/server/BUILD.bazel @@ -16,6 +16,8 @@ go_library( "//platform/errs:go_default_library", "//platform/errs/generic:go_default_library", "//platform/errs/mysql:go_default_library", + "//platform/extension/consumergate:go_default_library", + "//platform/extension/consumergate/file:go_default_library", "//platform/extension/counter/mysql:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", diff --git a/service/submitqueue/gateway/server/main.go b/service/submitqueue/gateway/server/main.go index 73cd122d..4d2bb453 100644 --- a/service/submitqueue/gateway/server/main.go +++ b/service/submitqueue/gateway/server/main.go @@ -33,6 +33,8 @@ import ( "github.com/uber/submitqueue/platform/errs" genericerrs "github.com/uber/submitqueue/platform/errs/generic" mysqlerrs "github.com/uber/submitqueue/platform/errs/mysql" + "github.com/uber/submitqueue/platform/extension/consumergate" + consumergatefile "github.com/uber/submitqueue/platform/extension/consumergate/file" mysqlcounter "github.com/uber/submitqueue/platform/extension/counter/mysql" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" @@ -296,6 +298,7 @@ func run() error { genericerrs.Classifier, mysqlerrs.Classifier, ), + newConsumerGate(logger), ) logController := logctrl.NewController(logger.Sugar(), scope, store, topickey.TopicKeyLog, "gateway-log") @@ -374,3 +377,23 @@ func run() error { return err } + +// defaultConsumerGateDir is the path the compose stack bind-mounts for gate +// state files. A missing directory simply means every gate is open. +const defaultConsumerGateDir = "/var/submitqueue/consumergate" + +// getEnv returns environment variable value or default if not set. +func getEnv(key, defaultVal string) string { + if val := os.Getenv(key); val != "" { + return val + } + return defaultVal +} + +// newConsumerGate returns a file-backed consumer gate rooted at the directory +// from CONSUMER_GATE_DIR (defaulting to defaultConsumerGateDir). +func newConsumerGate(logger *zap.Logger) consumergate.Gate { + dir := getEnv("CONSUMER_GATE_DIR", defaultConsumerGateDir) + logger.Info("consumer gate configured", zap.String("dir", dir)) + return consumergatefile.New(dir, consumergate.DefaultConfig()) +} diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 2ccb3493..8e6921e2 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -17,6 +17,8 @@ go_library( "//platform/errs:go_default_library", "//platform/errs/generic:go_default_library", "//platform/errs/mysql:go_default_library", + "//platform/extension/consumergate:go_default_library", + "//platform/extension/consumergate/file:go_default_library", "//platform/extension/counter:go_default_library", "//platform/extension/counter/mysql:go_default_library", "//platform/extension/messagequeue:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 6ed707ca..f4a36501 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -37,6 +37,8 @@ import ( "github.com/uber/submitqueue/platform/errs" genericerrs "github.com/uber/submitqueue/platform/errs/generic" mysqlerrs "github.com/uber/submitqueue/platform/errs/mysql" + "github.com/uber/submitqueue/platform/extension/consumergate" + consumergatefile "github.com/uber/submitqueue/platform/extension/consumergate/file" "github.com/uber/submitqueue/platform/extension/counter" mysqlcounter "github.com/uber/submitqueue/platform/extension/counter/mysql" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" @@ -220,6 +222,11 @@ func run() error { // so every non-nil error from a DLQ controller is forced retryable — // reconciliation must redeliver on any failure because the DLQ // subscriptions are final destinations (there is no further DLQ). + // Consumer gate: both consumers are gated uniformly — the gate keys on + // consumer group, so a DLQ stage is paused by its own group name just + // like a primary stage. + gate := newConsumerGate(logger) + primaryConsumer := consumer.New(logger.Sugar(), scope.SubScope("consumer"), registry, errs.NewClassifierProcessor( genericerrs.Classifier, @@ -228,9 +235,11 @@ func run() error { // errors surfaced from either backend. mysqlerrs.Classifier, ), + gate, ) dlqConsumer := consumer.New(logger.Sugar(), scope.SubScope("consumer-dlq"), registry, errs.AlwaysRetryableProcessor, + gate, ) // Build the per-queue extension registry: each queue resolves to its own @@ -785,6 +794,18 @@ func getEnv(key, defaultVal string) string { return defaultVal } +// defaultConsumerGateDir is the path the compose stack bind-mounts for gate +// state files. A missing directory simply means every gate is open. +const defaultConsumerGateDir = "/var/submitqueue/consumergate" + +// newConsumerGate returns a file-backed consumer gate rooted at the directory +// from CONSUMER_GATE_DIR (defaulting to defaultConsumerGateDir). +func newConsumerGate(logger *zap.Logger) consumergate.Gate { + dir := getEnv("CONSUMER_GATE_DIR", defaultConsumerGateDir) + logger.Info("consumer gate configured", zap.String("dir", dir)) + return consumergatefile.New(dir, consumergate.DefaultConfig()) +} + // parseTimeout parses a duration from environment variable with fallback to default. // Returns defaultVal if envVal is empty or cannot be parsed. func parseTimeout(envVal string, defaultVal time.Duration) time.Duration { diff --git a/test/e2e/submitqueue/BUILD.bazel b/test/e2e/submitqueue/BUILD.bazel index eda42174..bb4cf79f 100644 --- a/test/e2e/submitqueue/BUILD.bazel +++ b/test/e2e/submitqueue/BUILD.bazel @@ -18,12 +18,19 @@ go_test( "e2e", "external", "integration", + # The suite bind-mounts a host /tmp directory (consumer-gate state) into + # the compose services. Bazel's hermetic sandbox /tmp is invisible to + # the Docker daemon, so this test must run unsandboxed. + "no-sandbox", ], deps = [ "//api/base/change/protopb:go_default_library", "//api/base/mergestrategy/protopb:go_default_library", + "//api/runway/messagequeue:go_default_library", "//api/submitqueue/gateway/protopb:go_default_library", "//api/submitqueue/orchestrator/protopb:go_default_library", + "//platform/extension/consumergate:go_default_library", + "//platform/extension/consumergate/file:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mysql:go_default_library", diff --git a/test/e2e/submitqueue/harness_test.go b/test/e2e/submitqueue/harness_test.go index b8d18eef..6f9dbdeb 100644 --- a/test/e2e/submitqueue/harness_test.go +++ b/test/e2e/submitqueue/harness_test.go @@ -28,11 +28,14 @@ package e2e_test // stage is genuinely stuck, not a timing race. import ( + "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" changepb "github.com/uber/submitqueue/api/base/change/protopb" mergestrategypb "github.com/uber/submitqueue/api/base/mergestrategy/protopb" gatewaypb "github.com/uber/submitqueue/api/submitqueue/gateway/protopb" + "github.com/uber/submitqueue/platform/extension/consumergate" "github.com/uber/submitqueue/submitqueue/entity" ) @@ -130,6 +133,80 @@ func (s *E2EIntegrationSuite) assertStatusesInOrder(sqid string, want ...entity. sqid, want, got) } +// assertStatusesNever asserts that none of the banned statuses ever appeared +// in the request_log status timeline. +func (s *E2EIntegrationSuite) assertStatusesNever(sqid string, banned ...entity.RequestStatus) { + t := s.T() + got := s.timeline(sqid) + for _, b := range banned { + assert.NotContainsf(t, got, b, + "request_log for %s must never contain %q; got %v", sqid, b, got) + } +} + +// closeGate closes the consumer gate for the consumer group, scoped to one +// partition (the queue name for pipeline topics). The gate must be closed +// before the message that must be caught is published — that makes the stop +// exact by construction rather than a timing race. +func (s *E2EIntegrationSuite) closeGate(consumerGroup, partitionKey, reason string) { + t := s.T() + key := consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey} + require.NoError(t, s.gate.Close(s.ctx, key, consumergate.Metadata{ + Reason: reason, + CreatedBy: "e2e-suite", + CreatedAtMs: time.Now().UnixMilli(), + }), "failed to close gate %+v", key) + s.log.Logf("Closed consumer gate %s (partition %q)", consumerGroup, partitionKey) +} + +// openGate opens the consumer gate for the consumer group and partition. +// Opening an already-open gate is a no-op, so it is safe to call from a defer +// after an explicit open. +func (s *E2EIntegrationSuite) openGate(consumerGroup, partitionKey string) { + t := s.T() + key := consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey} + require.NoError(t, s.gate.Open(s.ctx, key), "failed to open gate %+v", key) + s.log.Logf("Opened consumer gate %s (partition %q)", consumerGroup, partitionKey) +} + +// awaitParked polls the shared gate directory until the delivery identified by +// (consumer group, topic key, message ID) has a parked record, and returns it. +// The record is written by the gated service before it blocks, so observing it +// proves the stopped controller is holding exactly this message — as opposed +// to the message simply not having arrived yet. +func (s *E2EIntegrationSuite) awaitParked(consumerGroup, topic, messageID string) consumergate.Parked { + return s.awaitParkedRecord(consumerGroup, topic, messageID, false) +} + +// awaitReleased polls until the parked record for the delivery carries a +// released stamp — proof the parked delivery proceeded into the controller +// after the gate opened. +func (s *E2EIntegrationSuite) awaitReleased(consumerGroup, topic, messageID string) consumergate.Parked { + return s.awaitParkedRecord(consumerGroup, topic, messageID, true) +} + +// awaitParkedRecord is the shared poll behind awaitParked/awaitReleased. +func (s *E2EIntegrationSuite) awaitParkedRecord(consumerGroup, topic, messageID string, released bool) consumergate.Parked { + t := s.T() + var found consumergate.Parked + require.Eventually(t, func() bool { + records, err := s.gate.ListParked(s.ctx, consumerGroup) + if err != nil { + s.log.Logf("ListParked(%s) failed, retrying: %v", consumerGroup, err) + return false + } + for _, r := range records { + if r.Topic == topic && r.MessageID == messageID && (!released || r.ReleasedAtMs != 0) { + found = r + return true + } + } + return false + }, persistTimeout, persistPollInterval, + "delivery %s on topic %s should be parked (released=%v) by gate %s", messageID, topic, released, consumerGroup) + return found +} + // terminalState reads the request's current internal RequestState from the // operating store (mysql-app). Unlike the status timeline, RequestState is // point-in-time — the Request entity is updated in place under optimistic diff --git a/test/e2e/submitqueue/suite_test.go b/test/e2e/submitqueue/suite_test.go index 19d27bc3..74cd394c 100644 --- a/test/e2e/submitqueue/suite_test.go +++ b/test/e2e/submitqueue/suite_test.go @@ -25,6 +25,7 @@ package e2e_test import ( "context" "database/sql" + "os" "path/filepath" "testing" "time" @@ -33,8 +34,11 @@ import ( "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "github.com/uber-go/tally" + runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" gatewaypb "github.com/uber/submitqueue/api/submitqueue/gateway/protopb" orchestratorpb "github.com/uber/submitqueue/api/submitqueue/orchestrator/protopb" + "github.com/uber/submitqueue/platform/extension/consumergate" + consumergatefile "github.com/uber/submitqueue/platform/extension/consumergate/file" "github.com/uber/submitqueue/submitqueue/entity" "github.com/uber/submitqueue/submitqueue/extension/storage" storagemysql "github.com/uber/submitqueue/submitqueue/extension/storage/mysql" @@ -55,6 +59,7 @@ type E2EIntegrationSuite struct { queueDB *sql.DB // Queue database requestLog storage.RequestLogStore // White-box view of the request_log status timeline (app DB) requestStore storage.RequestStore // White-box view of the internal RequestState (app DB) + gate *consumergatefile.Store // Consumer-gate control plane (shared dir bind-mounted into services) } func TestE2EIntegration(t *testing.T) { @@ -83,13 +88,31 @@ func (s *E2EIntegrationSuite) SetupSuite() { repoRoot := testutil.FindRepoRoot(t) t.Setenv("REPO_ROOT", repoRoot) + // Consumer-gate directory, bind-mounted into every service container by the + // compose file (SQ_CONSUMER_GATE_DIR → /var/submitqueue/consumergate). The + // suite closes/opens gates and reads parked records through the same file + // implementation the services use. Created directly under /tmp — not + // t.TempDir() — because it must be a host path the Docker daemon can bind + // mount, and because the containers write parked records into it as root, + // which would make t.TempDir()'s mandatory cleanup fail. Removal is + // therefore best-effort. + gateDir, err := os.MkdirTemp("/tmp", "sq-consumergate-") + require.NoError(t, err, "failed to create consumer-gate dir") + t.Cleanup(func() { + if rmErr := os.RemoveAll(gateDir); rmErr != nil { + s.log.Logf("best-effort consumer-gate dir cleanup failed (root-owned parked records): %v", rmErr) + } + }) + t.Setenv("SQ_CONSUMER_GATE_DIR", gateDir) + s.gate = consumergatefile.New(gateDir, consumergate.DefaultConfig()) + // Use docker-compose from service/submitqueue (full stack) // NOTE: Assumes Linux binaries are pre-built via make target composeFile := filepath.Join(repoRoot, "service/submitqueue/docker-compose.yml") s.stack = testutil.NewComposeStack(t, s.log, s.ctx, composeFile, "e2e-submitqueue") // Start the compose stack (Gateway + Orchestrator + 2 MySQL DBs) - err := s.stack.Up() + err = s.stack.Up() require.NoError(t, err, "failed to start compose stack") s.log.Logf("Compose stack started successfully") @@ -230,33 +253,76 @@ func (s *E2EIntegrationSuite) TestCancelRequest_InvalidSqid() { "empty sqid should map to InvalidArgument; got %s", st.Code()) } -// TestCancel_RecordsIntent verifies the deterministic half of the cancel flow: -// Cancel returns OK and the gateway synchronously records a "cancelling" intent -// entry in the request_log (written directly to the app DB before the RPC -// returns, right after the Land "accepted" entry). +// TestCancel_CaughtPreBatch_NeverLands drives the deterministic cancel +// scenario from doc/rfc/consumer-gate.md as stop → observe → start: the +// consumer gate stops runway's merge-conflict-check controller before the +// request's check message can be answered, so the request is provably held +// pre-batch while the cancel lands. The change must never reach the repo. +// +// 1. Stop: close the gate for runway-mergeconflictcheck, scoped to this +// queue's partition, before landing — exact by construction, no timing. +// 2. Land: the orchestrator runs the request to the merge-conflict-check +// hand-off; runway's subscriber delivers the check and the gate parks it. +// 3. Observe: awaiting the parked record proves the controller is stopped and +// holding exactly this request's check (there is otherwise no signal +// distinguishing "gated and parked" from "not arrived yet"). +// 4. Act while stopped: cancel the request. It is pre-batch by construction, +// so the cancel controller drives it terminal Cancelled directly. +// 5. Start: open the gate. The parked check proceeds as the same attempt, +// runway answers the now-stale check, and the orchestrator drops the +// signal for the halted request. // -// It deliberately does NOT assert the terminal "cancelled" outcome. Cancellation -// is best-effort and races the pipeline: on the hermetic stack the happy path -// reaches "landed" in ~2s, and a cancel published before the orchestrator's -// start controller has created the request is rejected to the DLQ and reconciled -// to "error". Asserting a terminal "cancelled" deterministically needs a -// pipeline-pause lever (e.g. a runway "park" marker that withholds the -// merge-conflict-check signal so the request is caught pre-batch) — that is the -// next incremental, per-stage addition on top of this harness. -func (s *E2EIntegrationSuite) TestCancel_RecordsIntent() { +// The drop in step 5 is asserted without sleeping: a sentinel request landed +// on the same queue after the gate opens shares the check and signal +// partitions with the stale message, so the sentinel reaching "landed" proves +// the stale signal was already consumed — at which point the cancelled +// request must still be terminal Cancelled, never batched, never landed. +func (s *E2EIntegrationSuite) TestCancel_CaughtPreBatch_NeverLands() { t := s.T() - sqid := s.land("e2e-cancel-queue", "github://github.example.com/uber/e2e-cancel/pull/9999/abcdef0123456789abcdef0123456789abcdef01") - s.log.Logf("Land (cancel path) succeeded: sqid=%s; cancelling", sqid) + const queue = "e2e-cancel-queue" + const gateGroup = "runway-mergeconflictcheck" + gateTopic := runwaymq.TopicKeyMergeConflictCheck.String() + + s.closeGate(gateGroup, queue, "e2e: hold merge-conflict check to catch cancel pre-batch") + // Reopen even if an assertion below fails, so teardown does not stop the + // stack with a delivery still parked. Opening twice is a no-op. + defer s.openGate(gateGroup, queue) + + sqid := s.land(queue, "github://github.example.com/uber/e2e-cancel/pull/9999/abcdef0123456789abcdef0123456789abcdef01") + s.log.Logf("Land (cancel path) succeeded: sqid=%s; awaiting parked check", sqid) + + parked := s.awaitParked(gateGroup, gateTopic, sqid) + assert.Equal(t, queue, parked.PartitionKey, "check message should be partitioned by queue") + assert.NotEmpty(t, parked.Payload, "parked record should carry the check payload") + // The controller is provably stopped and holding this request's check; + // cancel now. The request cannot be batched until the check is answered, + // so the cancel controller takes the not-batched path to terminal + // Cancelled. _, err := s.gatewayClient.Cancel(s.ctx, &gatewaypb.CancelRequest{Sqid: sqid, Reason: "e2e cancel test"}) require.NoError(t, err, "Cancel failed") - // The gateway writes "accepted" (on Land) and "cancelling" (on Cancel) - // synchronously to the same store, so both are present the moment Cancel - // returns — no polling needed. + s.awaitStatus(sqid, entity.RequestStatusCancelled) s.assertStatusesInOrder(sqid, entity.RequestStatusAccepted, entity.RequestStatusCancelling, + entity.RequestStatusCancelled, ) + assert.Equal(t, entity.RequestStateCancelled, s.terminalState(sqid), + "operating store should show request %s terminal cancelled while its check is parked", sqid) + + // Start the controller again and prove the parked delivery proceeded. + s.openGate(gateGroup, queue) + s.awaitReleased(gateGroup, gateTopic, sqid) + + // Sentinel on the same queue: its landing proves the stale signal ahead of + // it on the same partitions was consumed. + sentinel := s.land(queue, "github://github.example.com/uber/e2e-cancel/pull/10000/1234567890abcdef1234567890abcdef12345678") + s.awaitStatus(sentinel, entity.RequestStatusLanded) + + // The stale check answer was dropped: the cancelled request never advanced. + assert.Equal(t, entity.RequestStateCancelled, s.terminalState(sqid), + "request %s must stay terminal cancelled after its stale check signal is processed", sqid) + s.assertStatusesNever(sqid, entity.RequestStatusBatched, entity.RequestStatusLanded) } diff --git a/test/integration/submitqueue/core/consumer/BUILD.bazel b/test/integration/submitqueue/core/consumer/BUILD.bazel index 3afe7e9e..1924a69d 100644 --- a/test/integration/submitqueue/core/consumer/BUILD.bazel +++ b/test/integration/submitqueue/core/consumer/BUILD.bazel @@ -15,6 +15,7 @@ go_test( "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/errs:go_default_library", + "//platform/extension/consumergate/noop:go_default_library", "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//test/testutil:go_default_library", diff --git a/test/integration/submitqueue/core/consumer/consumer_test.go b/test/integration/submitqueue/core/consumer/consumer_test.go index 98709155..fc558d9f 100644 --- a/test/integration/submitqueue/core/consumer/consumer_test.go +++ b/test/integration/submitqueue/core/consumer/consumer_test.go @@ -17,6 +17,7 @@ import ( entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" + consumergatenoop "github.com/uber/submitqueue/platform/extension/consumergate/noop" extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" "github.com/uber/submitqueue/test/testutil" @@ -147,7 +148,7 @@ func (s *ConsumerIntegrationSuite) newConsumer(t *testing.T, q extqueue.Queue, t }) require.NoError(t, err) - return consumer.New(logger, tally.NoopScope, registry, errs.NewClassifierProcessor()) + return consumer.New(logger, tally.NoopScope, registry, errs.NewClassifierProcessor(), consumergatenoop.New()) } func (s *ConsumerIntegrationSuite) TestConsumerPerPartitionIsolation() {