From b436d8f50fcec262e754877c06d1ccc4b549adac Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Mon, 7 Sep 2026 22:34:47 -0700 Subject: [PATCH 1/2] [2/4 messagequeue] Shard MySQL queues by tenant identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Vitess needs a stable vindex that is not the Kafka-style partition key. Isolation has to be an explicit tenant column, with identity carried on publish and consume so two tenants sharing a partition key cannot collide. ### What? - Prefix every queue table and store query with tenant; set queue_offsets PK to (tenant, topic, partition_key, consumer_group). - Use typed (tenant, partitionKey) identities in consumers, gates, and MySQL workers; stop workers when lease discovery cannot confirm ownership. - Require tenant on publish.Message and validate payload queue against message tenant in pipeline controllers (the publish signature change is not compilable without those call sites). - Add ParseRequiredTenants and wire Stovepipe ingest to the configured tenant list so NewIngestController still builds. - Require tenant xor all-tenants on MQ admin list commands; message inspect/delete/requeue take the full (tenant, topic, partition, id) identity. ## Test Plan ✅ `go test` on mysql, publish, consumer, messagequeue identity, service/messagequeue, and start controller Co-authored-by: Cursor --- platform/base/messagequeue/BUILD.bazel | 8 +- platform/base/messagequeue/identity.go | 45 ++ platform/base/messagequeue/identity_test.go | 70 ++ platform/base/messagequeue/message.go | 9 +- platform/base/messagequeue/partition.go | 32 + platform/consumer/consumer.go | 55 +- platform/consumer/consumer_test.go | 194 ++++- platform/extension/consumergate/BUILD.bazel | 1 + .../extension/consumergate/consumergate.go | 15 +- .../extension/consumergate/file/BUILD.bazel | 6 +- platform/extension/consumergate/file/store.go | 89 ++- .../extension/consumergate/file/store_test.go | 98 ++- .../extension/consumergate/noop/BUILD.bazel | 1 + .../extension/consumergate/noop/gate_test.go | 6 +- .../extension/messagequeue/mysql/BUILD.bazel | 1 + .../extension/messagequeue/mysql/constants.go | 1 + .../messagequeue/mysql/ctl/BUILD.bazel | 13 +- .../messagequeue/mysql/ctl/README.md | 76 +- .../messagequeue/mysql/ctl/commands.go | 129 +++- .../messagequeue/mysql/ctl/commands_test.go | 75 ++ .../messagequeue/mysql/ctl/lib/admin.go | 211 +++-- .../messagequeue/mysql/ctl/lib/admin_test.go | 228 ++++-- .../mysql/delivery_state_store.go | 86 ++- .../mysql/delivery_state_store_test.go | 52 +- .../messagequeue/mysql/identifier.go | 68 ++ .../messagequeue/mysql/message_store.go | 119 +-- .../messagequeue/mysql/message_store_test.go | 63 +- .../messagequeue/mysql/mock_stores.go | 232 +++--- .../messagequeue/mysql/offset_store.go | 50 +- .../messagequeue/mysql/offset_store_test.go | 29 +- .../mysql/partition_lease_store.go | 96 +-- .../mysql/partition_lease_store_test.go | 56 +- .../extension/messagequeue/mysql/publisher.go | 31 +- .../messagequeue/mysql/publisher_test.go | 96 ++- .../mysql/schema/queue_delivery_state.sql | 16 +- .../mysql/schema/queue_messages.sql | 44 +- .../mysql/schema/queue_offsets.sql | 27 +- .../mysql/schema/queue_partition_leases.sql | 26 +- .../schema/queue_subscriber_heartbeats.sql | 22 +- platform/extension/messagequeue/mysql/sql.go | 10 + .../extension/messagequeue/mysql/sql_test.go | 47 ++ .../extension/messagequeue/mysql/stores.go | 134 ++-- .../messagequeue/mysql/subscriber.go | 731 +++++++++++------- .../mysql/subscriber_heartbeat_store.go | 41 +- .../mysql/subscriber_heartbeat_store_test.go | 52 +- .../messagequeue/mysql/subscriber_test.go | 728 ++++++++++++++--- platform/hook/publisher.go | 16 +- platform/hook/publisher_test.go | 11 +- platform/publish/publish.go | 60 +- platform/publish/publish_test.go | 75 +- runway/controller/dlq/BUILD.bazel | 1 + runway/controller/dlq/dlq.go | 14 +- runway/controller/dlq/dlq_test.go | 20 + runway/controller/merge/BUILD.bazel | 1 + runway/controller/merge/merge.go | 11 +- runway/controller/merge/merge_test.go | 18 + .../controller/mergeconflictcheck/BUILD.bazel | 1 + .../mergeconflictcheck/mergeconflictcheck.go | 11 +- .../mergeconflictcheck_test.go | 18 + service/messagequeue/BUILD.bazel | 18 + service/messagequeue/tenant.go | 104 +++ service/messagequeue/tenant_test.go | 110 +++ service/stovepipe/server/BUILD.bazel | 1 + service/stovepipe/server/main.go | 8 + stovepipe/controller/BUILD.bazel | 1 - stovepipe/controller/build/BUILD.bazel | 1 + stovepipe/controller/build/build.go | 11 +- stovepipe/controller/build/build_test.go | 18 +- stovepipe/controller/buildsignal/BUILD.bazel | 1 + .../controller/buildsignal/buildsignal.go | 11 +- .../buildsignal/buildsignal_test.go | 18 +- stovepipe/controller/dlq/BUILD.bazel | 1 + stovepipe/controller/dlq/build.go | 8 +- stovepipe/controller/dlq/build_test.go | 9 +- stovepipe/controller/dlq/buildsignal.go | 5 + stovepipe/controller/dlq/buildsignal_test.go | 10 +- stovepipe/controller/dlq/dlq_test.go | 20 +- stovepipe/controller/dlq/request.go | 5 + stovepipe/controller/ingest.go | 47 +- stovepipe/controller/ingest_test.go | 10 +- stovepipe/controller/process/BUILD.bazel | 1 + stovepipe/controller/process/process.go | 13 +- stovepipe/controller/process/process_test.go | 20 +- stovepipe/controller/record/BUILD.bazel | 1 + stovepipe/controller/record/record.go | 6 +- stovepipe/controller/record/record_test.go | 24 +- submitqueue/core/request/log.go | 8 +- submitqueue/core/request/log_test.go | 3 + submitqueue/core/request/terminate_test.go | 6 +- submitqueue/gateway/controller/cancel.go | 8 +- submitqueue/gateway/controller/cancel_test.go | 1 + submitqueue/gateway/controller/land.go | 8 +- submitqueue/gateway/controller/land_test.go | 1 + .../gateway/controller/log/BUILD.bazel | 1 + submitqueue/gateway/controller/log/log.go | 4 + .../gateway/controller/log/log_test.go | 17 + .../orchestrator/controller/batch/BUILD.bazel | 1 + .../orchestrator/controller/batch/batch.go | 12 +- .../controller/batch/batch_test.go | 22 +- .../orchestrator/controller/build/BUILD.bazel | 1 + .../orchestrator/controller/build/build.go | 11 +- .../controller/build/build_test.go | 15 +- .../controller/buildsignal/BUILD.bazel | 1 + .../controller/buildsignal/buildsignal.go | 11 +- .../buildsignal/buildsignal_test.go | 25 +- .../controller/cancel/BUILD.bazel | 1 + .../orchestrator/controller/cancel/cancel.go | 11 +- .../controller/cancel/cancel_test.go | 14 +- .../controller/conclude/BUILD.bazel | 1 + .../controller/conclude/conclude.go | 4 + .../controller/conclude/conclude_test.go | 30 +- .../controller/dependencyanalysis/BUILD.bazel | 1 + .../dependencyanalysis/dependencyanalysis.go | 11 +- .../dependencyanalysis_test.go | 12 + .../orchestrator/controller/dlq/BUILD.bazel | 1 + .../orchestrator/controller/dlq/batch.go | 5 + .../orchestrator/controller/dlq/batch_test.go | 17 +- .../controller/dlq/buildsignal.go | 5 + .../controller/dlq/buildsignal_test.go | 19 +- .../orchestrator/controller/dlq/dlq_test.go | 16 +- .../controller/dlq/mergeconflictsignal.go | 5 + .../dlq/mergeconflictsignal_test.go | 15 +- .../controller/dlq/mergesignal.go | 5 + .../controller/dlq/mergesignal_test.go | 15 +- .../controller/dlq/publisher_test.go | 1 + .../orchestrator/controller/dlq/request.go | 5 + .../controller/dlq/request_test.go | 44 +- .../orchestrator/controller/dlq/speculate.go | 12 +- .../controller/dlq/speculate_test.go | 15 +- .../orchestrator/controller/merge/BUILD.bazel | 1 + .../orchestrator/controller/merge/merge.go | 11 +- .../controller/merge/merge_test.go | 21 +- .../mergeconflictsignal/BUILD.bazel | 1 + .../mergeconflictsignal.go | 11 +- .../mergeconflictsignal_test.go | 18 + .../controller/mergesignal/BUILD.bazel | 1 + .../controller/mergesignal/mergesignal.go | 12 +- .../mergesignal/mergesignal_test.go | 19 +- .../controller/speculate/BUILD.bazel | 1 + .../controller/speculate/run_test.go | 1 + .../controller/speculate/speculate.go | 12 +- .../controller/speculate/speculate_test.go | 14 +- .../orchestrator/controller/start/BUILD.bazel | 1 + .../orchestrator/controller/start/start.go | 11 +- .../controller/start/start_test.go | 18 + .../controller/validate/BUILD.bazel | 1 + .../controller/validate/validate.go | 11 +- .../controller/validate/validate_test.go | 28 +- test/e2e/submitqueue/BUILD.bazel | 1 + test/e2e/submitqueue/harness_test.go | 30 +- test/e2e/submitqueue/suite_test.go | 42 +- .../extension/messagequeue/mysql/BUILD.bazel | 5 +- .../messagequeue/mysql/queue_test.go | 640 +++++++-------- .../mysql/tenant_isolation_test.go | 168 ++++ tool/linter/queueshard/main.go | 51 +- tool/linter/queueshard/main_test.go | 40 +- 156 files changed, 4653 insertions(+), 1896 deletions(-) create mode 100644 platform/base/messagequeue/identity.go create mode 100644 platform/base/messagequeue/identity_test.go create mode 100644 platform/base/messagequeue/partition.go create mode 100644 platform/extension/messagequeue/mysql/ctl/commands_test.go create mode 100644 platform/extension/messagequeue/mysql/identifier.go create mode 100644 service/messagequeue/BUILD.bazel create mode 100644 service/messagequeue/tenant.go create mode 100644 service/messagequeue/tenant_test.go create mode 100644 test/integration/extension/messagequeue/mysql/tenant_isolation_test.go diff --git a/platform/base/messagequeue/BUILD.bazel b/platform/base/messagequeue/BUILD.bazel index 97b4f74aa..4e8dfb148 100644 --- a/platform/base/messagequeue/BUILD.bazel +++ b/platform/base/messagequeue/BUILD.bazel @@ -4,7 +4,9 @@ go_library( name = "go_default_library", srcs = [ "context.go", + "identity.go", "message.go", + "partition.go", ], importpath = "github.com/uber/submitqueue/platform/base/messagequeue", visibility = ["//visibility:public"], @@ -14,8 +16,12 @@ go_test( name = "go_default_test", srcs = [ "context_test.go", + "identity_test.go", "message_test.go", ], embed = [":go_default_library"], - deps = ["@com_github_stretchr_testify//assert:go_default_library"], + deps = [ + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], ) diff --git a/platform/base/messagequeue/identity.go b/platform/base/messagequeue/identity.go new file mode 100644 index 000000000..40d250dab --- /dev/null +++ b/platform/base/messagequeue/identity.go @@ -0,0 +1,45 @@ +// 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 messagequeue + +import "fmt" + +// ValidateTenantMetadata verifies that the persisted shard identity agrees +// with queue metadata when metadata is present. +func ValidateTenantMetadata(message Message) error { + if message.Tenant == "" { + return fmt.Errorf("message tenant is required") + } + queueName := message.Metadata[MetadataKeyQueueName] + if queueName != "" && queueName != message.Tenant { + return fmt.Errorf("message tenant %q does not match queue metadata %q", message.Tenant, queueName) + } + return nil +} + +// ValidatePayloadQueue verifies that a queue-bearing payload belongs to the +// persisted message tenant. +func ValidatePayloadQueue(message Message, payloadQueue string) error { + if err := ValidateTenantMetadata(message); err != nil { + return err + } + if payloadQueue == "" { + return fmt.Errorf("payload queue is required") + } + if payloadQueue != message.Tenant { + return fmt.Errorf("message tenant %q does not match payload queue %q", message.Tenant, payloadQueue) + } + return nil +} diff --git a/platform/base/messagequeue/identity_test.go b/platform/base/messagequeue/identity_test.go new file mode 100644 index 000000000..bece4ce36 --- /dev/null +++ b/platform/base/messagequeue/identity_test.go @@ -0,0 +1,70 @@ +// 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 messagequeue + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestValidateTenantMetadata(t *testing.T) { + tests := []struct { + name string + message Message + wantErr bool + }{ + {name: "tenant without metadata", message: Message{Tenant: "queue-a"}}, + {name: "matching metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-a"}}}, + {name: "empty tenant", message: Message{}, wantErr: true}, + {name: "conflicting metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-b"}}, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateTenantMetadata(tt.message) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + }) + } +} + +func TestValidatePayloadQueue(t *testing.T) { + tests := []struct { + name string + message Message + payloadQueue string + wantErr bool + }{ + {name: "matching queue", message: Message{Tenant: "queue-a"}, payloadQueue: "queue-a"}, + {name: "empty payload queue", message: Message{Tenant: "queue-a"}, wantErr: true}, + {name: "conflicting payload queue", message: Message{Tenant: "queue-a"}, payloadQueue: "queue-b", wantErr: true}, + {name: "conflicting metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-b"}}, payloadQueue: "queue-a", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidatePayloadQueue(tt.message, tt.payloadQueue) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + }) + } +} diff --git a/platform/base/messagequeue/message.go b/platform/base/messagequeue/message.go index 77ef1f00e..d3fa781a8 100644 --- a/platform/base/messagequeue/message.go +++ b/platform/base/messagequeue/message.go @@ -43,11 +43,15 @@ type Message struct { // Use for trace IDs, request IDs, and cross-service metadata. Metadata map[string]string - // PartitionKey determines which partition/shard this message goes to. - // Messages with the same PartitionKey are guaranteed ordered delivery. + // PartitionKey determines the ordering partition within Tenant. + // Messages with the same Tenant and PartitionKey are delivered in order. // Optional - if empty, backend may use round-robin distribution. PartitionKey string + // Tenant is the shard isolation identity persisted by the backend. + // Domains map their queue name onto this field at publish time. + Tenant string + // PublishedAt is when the message was published (Unix milliseconds). PublishedAt int64 } @@ -79,6 +83,7 @@ func (m Message) Copy() Message { Payload: payloadCopy, Metadata: maps.Clone(m.Metadata), PartitionKey: m.PartitionKey, + Tenant: m.Tenant, PublishedAt: m.PublishedAt, } } diff --git a/platform/base/messagequeue/partition.go b/platform/base/messagequeue/partition.go new file mode 100644 index 000000000..52394cc5d --- /dev/null +++ b/platform/base/messagequeue/partition.go @@ -0,0 +1,32 @@ +// 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 messagequeue + +// PartitionIdentity uniquely identifies a partition within a tenant. +// It is comparable and may be used as a map key. +type PartitionIdentity struct { + // Tenant is the shard isolation identity. + Tenant string + // PartitionKey is the partition key within the tenant. + PartitionKey string +} + +// PartitionIdentity returns the message's tenant-scoped partition identity. +func (m Message) PartitionIdentity() PartitionIdentity { + return PartitionIdentity{ + Tenant: m.Tenant, + PartitionKey: m.PartitionKey, + } +} diff --git a/platform/consumer/consumer.go b/platform/consumer/consumer.go index 687d15663..d94d20a9d 100644 --- a/platform/consumer/consumer.go +++ b/platform/consumer/consumer.go @@ -22,6 +22,7 @@ import ( "time" "github.com/uber-go/tally" + "github.com/uber/submitqueue/platform/base/failure" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/extension/consumergate" @@ -282,10 +283,10 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic "topic_key", topicKey, ) - // partitionChs maps partition keys to per-partition delivery channels. + // partitionChs maps tenant-scoped partitions to delivery channels. // Each channel is created lazily on the first message for that partition // and is never removed — partitions are stable for the lifetime of a subscription. - partitionChs := make(map[string]chan extqueue.Delivery) + partitionChs := make(map[entityqueue.PartitionIdentity]chan extqueue.Delivery) var wg sync.WaitGroup for { @@ -307,15 +308,18 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic m.shutdownPartitions(partitionChs, &wg) return } + if m.rejectDeliveryWithInvalidTenant(ctx, controller, delivery) { + continue + } // Route delivery to its partition's channel, creating the channel // and spawning a processPartition goroutine if this is the first // message for that partition. - partitionKey := delivery.Message().PartitionKey - ch, exists := partitionChs[partitionKey] + partition := delivery.Message().PartitionIdentity() + ch, exists := partitionChs[partition] if !exists { ch = make(chan extqueue.Delivery, batchSize) - partitionChs[partitionKey] = ch + partitionChs[partition] = ch wg.Add(1) go func(pCh <-chan extqueue.Delivery) { defer wg.Done() @@ -336,9 +340,32 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic } } +func (m *consumer) rejectDeliveryWithInvalidTenant(ctx context.Context, controller Controller, delivery extqueue.Delivery) bool { + msg := delivery.Message() + if err := entityqueue.ValidateTenantMetadata(msg); err != nil { + m.logger.Errorw("rejecting message with inconsistent queue identity", + "controller", controller.Name(), + "topic_key", controller.TopicKey(), + "message_id", msg.ID, + "tenant", msg.Tenant, + "error", err, + ) + if rejectErr := delivery.Reject(ctx, failure.New(err.Error())); rejectErr != nil { + m.logger.Errorw("failed to reject message with inconsistent queue identity", + "controller", controller.Name(), + "topic_key", controller.TopicKey(), + "message_id", msg.ID, + "error", rejectErr, + ) + } + return true + } + return false +} + // shutdownPartitions closes all partition channels to signal processPartition // goroutines to exit, then waits for them to finish draining. -func (m *consumer) shutdownPartitions(partitionChs map[string]chan extqueue.Delivery, wg *sync.WaitGroup) { +func (m *consumer) shutdownPartitions(partitionChs map[entityqueue.PartitionIdentity]chan extqueue.Delivery, wg *sync.WaitGroup) { for _, ch := range partitionChs { close(ch) } @@ -373,11 +400,9 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d const opName = "process" msg := delivery.Message() - queueName := msg.Metadata[entityqueue.MetadataKeyQueueName] + queueName := msg.Tenant ctx = entityqueue.WithQueueName(ctx, queueName) - if queueName != "" { - ctx = metrics.WithContextTags(ctx, metrics.NewTag("queue", queueName)) - } + ctx = metrics.WithContextTags(ctx, metrics.NewTag("queue", queueName)) // Consumer gate: a delivery whose gate is closed is recorded as parked and // postponed (barrier + re-check on redelivery); a false return also covers @@ -575,7 +600,10 @@ func (m *consumer) checkGate(ctx context.Context, controller Controller, deliver consumerGroup := controller.ConsumerGroup() topic := controller.TopicKey().String() - entry, err := m.gate.Enter(ctx, consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: msg.PartitionKey}) + entry, err := m.gate.Enter(ctx, consumergate.Key{ + ConsumerGroup: consumerGroup, + Partition: msg.PartitionIdentity(), + }) if err != nil { if errors.Is(err, context.Canceled) { // Shutting down: leave the delivery in flight; visibility lapses @@ -708,10 +736,11 @@ func (m *consumer) unsubscribeAll(timeoutMs int64) error { var timedOutControllers []string for topicKey, sub := range m.subscriptions { start := time.Now() + timer := time.NewTimer(remaining) select { case <-sub.done: - // Controller stopped gracefully - case <-time.After(remaining): + timer.Stop() + case <-timer.C: m.logger.Errorw("timeout waiting for controller to stop", "controller", sub.controller.Name(), "topic_key", topicKey, diff --git a/platform/consumer/consumer_test.go b/platform/consumer/consumer_test.go index ceb557026..f7dd078ff 100644 --- a/platform/consumer/consumer_test.go +++ b/platform/consumer/consumer_test.go @@ -41,6 +41,7 @@ import ( const ( testTopicKeyStart TopicKey = "start" testTopicKeyValidate TopicKey = "validate" + testTenant = "test-tenant" ) // testController is a configurable Controller used by consumer tests. @@ -108,6 +109,9 @@ func newRegistry(t *testing.T, q extqueue.Queue, topicKey TopicKey, consumerGrou // that closes when Ack or Nack is called. func setupDelivery(del *queuemock.MockDelivery, msg entityqueue.Message, ackErr, nackErr error) chan struct{} { done := make(chan struct{}) + if msg.Tenant == "" { + msg.Tenant = testTenant + } del.EXPECT().Message().Return(msg).AnyTimes() del.EXPECT().Attempt().Return(1).AnyTimes() del.EXPECT().ReceivedAt().Return(time.Now().UnixMilli()).AnyTimes() @@ -311,6 +315,7 @@ func TestConsumer_ProcessDelivery_Success(t *testing.T) { msg := entityqueue.NewMessage("test-msg-1", []byte("payload"), "partition1", map[string]string{ entityqueue.MetadataKeyQueueName: "monorepo/main", }) + msg.Tenant = "monorepo/main" mockDel := queuemock.NewMockDelivery(ctrl) done := setupDelivery(mockDel, msg, nil, nil) @@ -325,6 +330,89 @@ func TestConsumer_ProcessDelivery_Success(t *testing.T) { require.NoError(t, err) } +func TestConsumer_ProcessDelivery_RejectsConflictingTenantMetadata(t *testing.T) { + ctrl := gomock.NewController(t) + deliveryChan := make(chan extqueue.Delivery, 1) + 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) + + c := New( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + newRegistry(t, mockQ, testTopicKeyStart, "test-group"), + errs.NewClassifierProcessor(), + consumergatenoop.New(), + ) + handler := &testController{} + setupController(handler, "test-handler", testTopicKeyStart, "test-group", + func(context.Context, Delivery) error { + assert.Fail(t, "controller must not receive inconsistent delivery identity") + return nil + }, + ) + require.NoError(t, c.Register(handler)) + require.NoError(t, c.Start(context.Background())) + + msg := entityqueue.NewMessage("mismatch", []byte("payload"), "partition1", map[string]string{ + entityqueue.MetadataKeyQueueName: "tenant-b", + }) + msg.Tenant = "tenant-a" + rejected := make(chan struct{}) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Reject(gomock.Any(), gomock.Any()).DoAndReturn(func(context.Context, failure.Failure) error { + close(rejected) + return nil + }) + + deliveryChan <- delivery + <-rejected + require.NoError(t, c.Stop(30000)) +} + +func TestConsumer_ProcessDelivery_RejectsEmptyTenant(t *testing.T) { + ctrl := gomock.NewController(t) + deliveryChan := make(chan extqueue.Delivery, 1) + 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) + + c := New( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + newRegistry(t, mockQ, testTopicKeyStart, "test-group"), + errs.NewClassifierProcessor(), + consumergatenoop.New(), + ) + handler := &testController{} + setupController(handler, "test-handler", testTopicKeyStart, "test-group", + func(context.Context, Delivery) error { + assert.Fail(t, "controller must not receive a delivery without a tenant") + return nil + }, + ) + require.NoError(t, c.Register(handler)) + require.NoError(t, c.Start(context.Background())) + + msg := entityqueue.NewMessage("missing-tenant", []byte("payload"), "partition1", map[string]string{ + entityqueue.MetadataKeyQueueName: testTenant, + }) + rejected := make(chan struct{}) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Reject(gomock.Any(), gomock.Any()).DoAndReturn(func(context.Context, failure.Failure) error { + close(rejected) + return nil + }) + + deliveryChan <- delivery + <-rejected + require.NoError(t, c.Stop(30000)) +} + func TestConsumer_ProcessDelivery_Error(t *testing.T) { ctrl := gomock.NewController(t) logger := zaptest.NewLogger(t).Sugar() @@ -451,6 +539,7 @@ func TestConsumer_ProcessDelivery_Hold(t *testing.T) { require.NoError(t, c.Start(ctx)) msg := entityqueue.NewMessage("held-msg", []byte("payload"), "partition1", nil) + msg.Tenant = testTenant done := make(chan struct{}) var gotDelayMs int64 mockDel := queuemock.NewMockDelivery(ctrl) @@ -518,6 +607,7 @@ func TestConsumer_ProcessDelivery_NonRetryableError(t *testing.T) { require.NoError(t, err) msg := entityqueue.NewMessage("poison-msg", []byte("bad"), "partition1", nil) + msg.Tenant = testTenant done := make(chan struct{}) mockDel := queuemock.NewMockDelivery(ctrl) mockDel.EXPECT().Message().Return(msg).AnyTimes() @@ -592,6 +682,7 @@ func TestConsumer_ProcessDelivery_FailureFromControllerError(t *testing.T) { require.NoError(t, c.Start(ctx)) msg := entityqueue.NewMessage("msg-1", []byte("bad"), "partition1", nil) + msg.Tenant = testTenant done := make(chan struct{}) var got failure.Failure @@ -1018,6 +1109,7 @@ func TestConsumer_PerPartitionProcessing(t *testing.T) { // Send message to partition A (will block in controller) msgA := entityqueue.NewMessage("msg-a", []byte("payload-a"), "partition-a", nil) + msgA.Tenant = testTenant mockDelA := queuemock.NewMockDelivery(ctrl) mockDelA.EXPECT().Message().Return(msgA).AnyTimes() mockDelA.EXPECT().Attempt().Return(1).AnyTimes() @@ -1034,6 +1126,7 @@ func TestConsumer_PerPartitionProcessing(t *testing.T) { // Send message to partition B (should process despite A being blocked) msgB := entityqueue.NewMessage("msg-b", []byte("payload-b"), "partition-b", nil) + msgB.Tenant = testTenant mockDelB := queuemock.NewMockDelivery(ctrl) mockDelB.EXPECT().Message().Return(msgB).AnyTimes() mockDelB.EXPECT().Attempt().Return(1).AnyTimes() @@ -1052,6 +1145,53 @@ func TestConsumer_PerPartitionProcessing(t *testing.T) { require.NoError(t, err) } +func TestConsumer_SamePartitionKeyAcrossTenantsProcessesIndependently(t *testing.T) { + ctrl := gomock.NewController(t) + deliveryChan := make(chan extqueue.Delivery, 2) + 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, testTopicKeyStart, "test-group") + c := New(zaptest.NewLogger(t).Sugar(), tally.NoopScope, reg, errs.NewClassifierProcessor(), consumergatenoop.New()) + + tenantABlocked := make(chan struct{}) + tenantBProcessed := make(chan struct{}) + handler := &testController{} + setupController(handler, "test-handler", testTopicKeyStart, "test-group", func(ctx context.Context, delivery Delivery) error { + if delivery.Message().Tenant == "tenant-a" { + close(tenantABlocked) + <-ctx.Done() + return nil + } + close(tenantBProcessed) + return nil + }) + require.NoError(t, c.Register(handler)) + require.NoError(t, c.Start(context.Background())) + + msgA := entityqueue.NewMessage("msg-a", []byte("a"), "shared", nil) + msgA.Tenant = "tenant-a" + delA := queuemock.NewMockDelivery(ctrl) + delA.EXPECT().Message().Return(msgA).AnyTimes() + delA.EXPECT().Attempt().Return(1).AnyTimes() + delA.EXPECT().Ack(gomock.Any()).Return(nil).MaxTimes(1) + deliveryChan <- delA + <-tenantABlocked + + msgB := entityqueue.NewMessage("msg-b", []byte("b"), "shared", nil) + msgB.Tenant = "tenant-b" + delB := queuemock.NewMockDelivery(ctrl) + delB.EXPECT().Message().Return(msgB).AnyTimes() + delB.EXPECT().Attempt().Return(1).AnyTimes() + delB.EXPECT().Ack(gomock.Any()).Return(nil).MaxTimes(1) + deliveryChan <- delB + <-tenantBProcessed + + require.NoError(t, c.Stop(30000)) +} + // TestConsumer_PartitionOrdering verifies that messages within a single partition // are processed in order. func TestConsumer_PartitionOrdering(t *testing.T) { @@ -1100,6 +1240,7 @@ func TestConsumer_PartitionOrdering(t *testing.T) { // Send 3 messages to the same partition for i, id := range []string{"msg-1", "msg-2", "msg-3"} { msg := entityqueue.NewMessage(id, []byte("payload"), "same-partition", nil) + msg.Tenant = testTenant mockDel := queuemock.NewMockDelivery(ctrl) mockDel.EXPECT().Message().Return(msg).AnyTimes() mockDel.EXPECT().Attempt().Return(1).AnyTimes() @@ -1261,11 +1402,11 @@ func (f *fakeGate) setErr(err error) { f.err = err } -func (f *fakeGate) isClosed(consumerGroup, partitionKey string) bool { +func (f *fakeGate) isClosed(consumerGroup string, partition entityqueue.PartitionIdentity) bool { if f.closed[consumergate.Key{ConsumerGroup: consumerGroup}] { return true } - return f.closed[consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey}] + return f.closed[consumergate.Key{ConsumerGroup: consumerGroup, Partition: partition}] } // Enter implements consumergate.Gate. It checks the err field first, then @@ -1278,7 +1419,7 @@ func (f *fakeGate) Enter(_ context.Context, key consumergate.Key) (consumergate. if f.err != nil { return nil, f.err } - return &fakeEntry{gate: f, key: key, blocked: f.isClosed(key.ConsumerGroup, key.PartitionKey)}, nil + return &fakeEntry{gate: f, key: key, blocked: f.isClosed(key.ConsumerGroup, key.Partition)}, nil } // fakeEntry is the entry handed out by fakeGate.Enter. @@ -1295,7 +1436,8 @@ func (e *fakeEntry) Park(_ context.Context, descriptor consumergate.DeliveryDesc ConsumerGroup: e.key.ConsumerGroup, Topic: descriptor.Topic, MessageID: descriptor.MessageID, - PartitionKey: e.key.PartitionKey, + Tenant: e.key.Partition.Tenant, + PartitionKey: e.key.Partition.PartitionKey, Payload: descriptor.Payload, Attempt: descriptor.Attempt, } @@ -1338,6 +1480,9 @@ func startGatedConsumer(t *testing.T, ctrl *gomock.Controller, gate consumergate // those calls fails the test. func gatedDelivery(t *testing.T, ctrl *gomock.Controller, msg entityqueue.Message) (*queuemock.MockDelivery, chan struct{}) { t.Helper() + if msg.Tenant == "" { + msg.Tenant = testTenant + } mockDel := queuemock.NewMockDelivery(ctrl) mockDel.EXPECT().Message().Return(msg).AnyTimes() mockDel.EXPECT().Attempt().Return(1).AnyTimes() @@ -1422,7 +1567,10 @@ func TestConsumer_Gate_BlockedParksAndPostpones(t *testing.T) { func TestConsumer_Gate_PartitionScoped(t *testing.T) { ctrl := gomock.NewController(t) gate := newFakeGate() - gate.close(consumergate.Key{ConsumerGroup: "test-group", PartitionKey: "gated-partition"}) + gate.close(consumergate.Key{ + ConsumerGroup: "test-group", + Partition: entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "gated-partition"}, + }) var handled sync.Map c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(_ context.Context, delivery Delivery) error { @@ -1451,7 +1599,10 @@ func TestConsumer_Gate_PartitionScoped(t *testing.T) { assert.False(t, ok) // Open the gate; the redelivery of the gated message processes. - gate.open(consumergate.Key{ConsumerGroup: "test-group", PartitionKey: "gated-partition"}) + gate.open(consumergate.Key{ + ConsumerGroup: "test-group", + Partition: entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "gated-partition"}, + }) redelivery := queuemock.NewMockDelivery(ctrl) gatedDone := setupDelivery(redelivery, gatedMsg, nil, nil) deliveryChan <- redelivery @@ -1462,6 +1613,37 @@ func TestConsumer_Gate_PartitionScoped(t *testing.T) { require.NoError(t, c.Stop(30000)) } +func TestConsumer_Gate_SamePartitionKeyAcrossTenantsIsIndependent(t *testing.T) { + ctrl := gomock.NewController(t) + gate := newFakeGate() + gatedPartition := entityqueue.PartitionIdentity{Tenant: "tenant-a", PartitionKey: "shared"} + gate.close(consumergate.Key{ConsumerGroup: "test-group", Partition: gatedPartition}) + + handled := make(chan string, 1) + c, deliveryChan := startGatedConsumer(t, ctrl, gate, func(_ context.Context, delivery Delivery) error { + handled <- delivery.Message().Tenant + return nil + }) + + gatedMsg := entityqueue.NewMessage("gated-msg", []byte("a"), "shared", nil) + gatedMsg.Tenant = "tenant-a" + gatedDel, postponed := gatedDelivery(t, ctrl, gatedMsg) + deliveryChan <- gatedDel + parked := <-gate.parked + assert.Equal(t, "tenant-a", parked.Tenant) + <-postponed + + openMsg := entityqueue.NewMessage("open-msg", []byte("b"), "shared", nil) + openMsg.Tenant = "tenant-b" + openDel := queuemock.NewMockDelivery(ctrl) + openDone := setupDelivery(openDel, openMsg, nil, nil) + deliveryChan <- openDel + <-openDone + assert.Equal(t, "tenant-b", <-handled) + + require.NoError(t, c.Stop(30000)) +} + func TestConsumer_Gate_StopWhileGated(t *testing.T) { ctrl := gomock.NewController(t) gate := newFakeGate() diff --git a/platform/extension/consumergate/BUILD.bazel b/platform/extension/consumergate/BUILD.bazel index aa8187837..896f6127f 100644 --- a/platform/extension/consumergate/BUILD.bazel +++ b/platform/extension/consumergate/BUILD.bazel @@ -5,4 +5,5 @@ go_library( srcs = ["consumergate.go"], importpath = "github.com/uber/submitqueue/platform/extension/consumergate", visibility = ["//visibility:public"], + deps = ["//platform/base/messagequeue:go_default_library"], ) diff --git a/platform/extension/consumergate/consumergate.go b/platform/extension/consumergate/consumergate.go index 581bf7552..5173aa768 100644 --- a/platform/extension/consumergate/consumergate.go +++ b/platform/extension/consumergate/consumergate.go @@ -34,16 +34,20 @@ package consumergate //go:generate mockgen -source=consumergate.go -destination=mock/consumergate_mock.go -package=mock -import "context" +import ( + "context" + + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" +) // 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 + // Partition optionally narrows the gate to a tenant-scoped partition. + // The zero value gates every partition of the consumer group. + Partition entityqueue.PartitionIdentity } // Metadata records why a gate was closed, for the operator who finds it later. @@ -89,6 +93,9 @@ type Parked struct { // MessageID is the queue message ID of the delivery. MessageID string + // Tenant is the shard isolation identity. + Tenant string + // PartitionKey is the partition the delivery belongs to. PartitionKey string diff --git a/platform/extension/consumergate/file/BUILD.bazel b/platform/extension/consumergate/file/BUILD.bazel index a37f30118..cf7afc925 100644 --- a/platform/extension/consumergate/file/BUILD.bazel +++ b/platform/extension/consumergate/file/BUILD.bazel @@ -5,7 +5,10 @@ go_library( srcs = ["store.go"], importpath = "github.com/uber/submitqueue/platform/extension/consumergate/file", visibility = ["//visibility:public"], - deps = ["//platform/extension/consumergate:go_default_library"], + deps = [ + "//platform/base/messagequeue:go_default_library", + "//platform/extension/consumergate:go_default_library", + ], ) go_test( @@ -13,6 +16,7 @@ go_test( srcs = ["store_test.go"], embed = [":go_default_library"], deps = [ + "//platform/base/messagequeue:go_default_library", "//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/store.go b/platform/extension/consumergate/file/store.go index cbe8c596f..464b3e7cc 100644 --- a/platform/extension/consumergate/file/store.go +++ b/platform/extension/consumergate/file/store.go @@ -16,9 +16,9 @@ // 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 +// gates/{consumer_group}/all gates every partition +// gates/{consumer_group}/partitions/t-{tenant}/p-{partition} gates one partition +// parked/{consumer_group}/{topic}/t-{tenant}/p-{partition}/{urlenc(id)}.json // // 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"), @@ -51,6 +51,7 @@ import ( "strings" "time" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/extension/consumergate" ) @@ -75,25 +76,39 @@ func New(dir string) *Store { // 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) + if key.Partition == (entityqueue.PartitionIdentity{}) { + return filepath.Join(s.dir, "gates", key.ConsumerGroup, "all") } - return filepath.Join(s.dir, "gates", key.ConsumerGroup, name) + return filepath.Join( + s.dir, + "gates", + key.ConsumerGroup, + "partitions", + "t-"+url.QueryEscape(key.Partition.Tenant), + "p-"+url.QueryEscape(key.Partition.PartitionKey), + ) } // 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") +func (s *Store) parkedPath(consumerGroup, topic string, partition entityqueue.PartitionIdentity, messageID string) string { + return filepath.Join( + s.dir, + "parked", + consumerGroup, + topic, + "t-"+url.QueryEscape(partition.Tenant), + "p-"+url.QueryEscape(partition.PartitionKey), + 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) { +func (s *Store) isGated(consumerGroup string, partition entityqueue.PartitionIdentity) (bool, error) { paths := []string{s.gatePath(consumergate.Key{ConsumerGroup: consumerGroup})} - if partitionKey != "" { - paths = append(paths, s.gatePath(consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey})) + if partition != (entityqueue.PartitionIdentity{}) { + paths = append(paths, s.gatePath(consumergate.Key{ConsumerGroup: consumerGroup, Partition: partition})) } for _, p := range paths { switch _, err := os.Stat(p); { @@ -111,7 +126,7 @@ func (s *Store) isGated(consumerGroup, partitionKey string) (bool, error) { // Enter implements consumergate.Gate. It returns an unblocked Entry when the // gate identified by key is open, and a blocked Entry when it is closed. func (s *Store) Enter(_ context.Context, key consumergate.Key) (consumergate.Entry, error) { - gated, err := s.isGated(key.ConsumerGroup, key.PartitionKey) + gated, err := s.isGated(key.ConsumerGroup, key.Partition) if err != nil { return nil, err } @@ -140,7 +155,8 @@ func (e entry) Park(_ context.Context, descriptor consumergate.DeliveryDescripto ConsumerGroup: e.key.ConsumerGroup, Topic: descriptor.Topic, MessageID: descriptor.MessageID, - PartitionKey: e.key.PartitionKey, + Tenant: e.key.Partition.Tenant, + PartitionKey: e.key.Partition.PartitionKey, Payload: descriptor.Payload, Attempt: descriptor.Attempt, ParkedAtMs: time.Now().UnixMilli(), @@ -150,13 +166,14 @@ func (e entry) Park(_ context.Context, descriptor consumergate.DeliveryDescripto // Unpark implements consumergate.Entry. Removing an absent record is a no-op, // so callers may invoke it unconditionally on the admit path. func (e entry) Unpark(_ context.Context, descriptor consumergate.DeliveryDescriptor) error { - return e.store.removeParked(e.key.ConsumerGroup, descriptor.Topic, descriptor.MessageID) + return e.store.removeParked(e.key.ConsumerGroup, descriptor.Topic, e.key.Partition, descriptor.MessageID) } // 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) + partition := entityqueue.PartitionIdentity{Tenant: parked.Tenant, PartitionKey: parked.PartitionKey} + path := s.parkedPath(parked.ConsumerGroup, parked.Topic, partition, parked.MessageID) if err := writeJSON(path, parkedRecord(parked)); err != nil { return fmt.Errorf("failed to write parked record %s: %w", path, err) } @@ -165,8 +182,8 @@ func (s *Store) recordParked(parked consumergate.Parked) error { // removeParked removes a parked-delivery record. Removing an already-absent // record is a no-op. -func (s *Store) removeParked(consumerGroup, topic, messageID string) error { - path := s.parkedPath(consumerGroup, topic, messageID) +func (s *Store) removeParked(consumerGroup, topic string, partition entityqueue.PartitionIdentity, messageID string) error { + path := s.parkedPath(consumerGroup, topic, partition, messageID) if err := os.Remove(path); err != nil && !os.IsNotExist(err) { return fmt.Errorf("failed to remove parked record %s: %w", path, err) } @@ -217,21 +234,39 @@ func (s *Store) ListParked(_ context.Context, consumerGroup string) ([]consumerg continue } topicDir := filepath.Join(groupDir, topic.Name()) - entries, err := os.ReadDir(topicDir) + tenantDirs, 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") { + for _, tenantDir := range tenantDirs { + if !tenantDir.IsDir() { continue } - rec, err := readParked(filepath.Join(topicDir, entry.Name())) + tenantPath := filepath.Join(topicDir, tenantDir.Name()) + partitionDirs, err := os.ReadDir(tenantPath) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to read parked dir %s: %w", tenantPath, err) + } + for _, partitionDir := range partitionDirs { + if !partitionDir.IsDir() { + continue + } + partitionPath := filepath.Join(tenantPath, partitionDir.Name()) + entries, err := os.ReadDir(partitionPath) + if err != nil { + return nil, fmt.Errorf("failed to read parked dir %s: %w", partitionPath, err) + } + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { + continue + } + rec, err := readParked(filepath.Join(partitionPath, entry.Name())) + if err != nil { + return nil, err + } + out = append(out, consumergate.Parked(rec)) + } } - out = append(out, consumergate.Parked(rec)) } } return out, nil @@ -256,6 +291,8 @@ type parkedRecord struct { Topic string `json:"topic"` // MessageID is the queue message ID of the parked delivery. MessageID string `json:"message_id"` + // Tenant is the shard isolation identity. + Tenant string `json:"tenant"` // PartitionKey is the partition the delivery belongs to. PartitionKey string `json:"partition_key"` // Payload is the message payload (base64 in the JSON encoding). diff --git a/platform/extension/consumergate/file/store_test.go b/platform/extension/consumergate/file/store_test.go index 1ab0abf07..c6317914f 100644 --- a/platform/extension/consumergate/file/store_test.go +++ b/platform/extension/consumergate/file/store_test.go @@ -22,9 +22,14 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/extension/consumergate" ) +func partitionIdentity(tenant, partitionKey string) entityqueue.PartitionIdentity { + return entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: partitionKey} +} + func TestIsGated(t *testing.T) { ctx := context.Background() @@ -32,55 +37,69 @@ func TestIsGated(t *testing.T) { name string close []consumergate.Key group string - partition string + partition entityqueue.PartitionIdentity want bool }{ { name: "no gates", group: "orchestrator-batch", - partition: "queue-a", + partition: partitionIdentity("", "queue-a"), want: false, }, { name: "all-partitions gate matches any partition", close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch"}}, group: "orchestrator-batch", - partition: "queue-a", + partition: partitionIdentity("", "queue-a"), want: true, }, { name: "all-partitions gate matches empty partition", close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch"}}, group: "orchestrator-batch", - partition: "", + partition: entityqueue.PartitionIdentity{}, want: true, }, { name: "partition gate matches its partition", - close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"}}, + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("", "queue-a")}}, group: "orchestrator-batch", - partition: "queue-a", + partition: partitionIdentity("", "queue-a"), want: true, }, { name: "partition gate leaves other partitions open", - close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"}}, + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("", "queue-a")}}, group: "orchestrator-batch", - partition: "queue-b", + partition: partitionIdentity("", "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", + partition: partitionIdentity("", "queue-a"), want: false, }, { name: "partition key with slash is encoded and matched", - close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue/1"}}, + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("", "queue/1")}}, + group: "orchestrator-batch", + partition: partitionIdentity("", "queue/1"), + want: true, + }, + { + name: "partition gate leaves same key in another tenant open", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("tenant-a", "shared")}}, + group: "orchestrator-batch", + partition: partitionIdentity("tenant-b", "shared"), + want: false, + }, + { + name: "partition gate matches tenant and partition", + close: []consumergate.Key{{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("tenant-a", "shared")}}, group: "orchestrator-batch", - partition: "queue/1", + partition: partitionIdentity("tenant-a", "shared"), want: true, }, } @@ -101,15 +120,15 @@ func TestIsGated(t *testing.T) { func TestOpenClosesGate(t *testing.T) { ctx := context.Background() store := New(t.TempDir()) - key := consumergate.Key{ConsumerGroup: "orchestrator-batch", PartitionKey: "queue-a"} + key := consumergate.Key{ConsumerGroup: "orchestrator-batch", Partition: partitionIdentity("", "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) + gated, err := store.isGated(key.ConsumerGroup, key.Partition) require.NoError(t, err) require.True(t, gated) require.NoError(t, store.Open(ctx, key)) - gated, err = store.isGated(key.ConsumerGroup, key.PartitionKey) + gated, err = store.isGated(key.ConsumerGroup, key.Partition) require.NoError(t, err) assert.False(t, gated) @@ -131,6 +150,7 @@ func TestParkedRecordLifecycle(t *testing.T) { ConsumerGroup: "runway-mergeconflictcheck", Topic: "merge-conflict-check", MessageID: "e2e-queue/42", + Tenant: "e2e-queue", PartitionKey: "e2e-queue", Payload: []byte(`{"id":"e2e-queue/42"}`), Attempt: 1, @@ -151,13 +171,46 @@ func TestParkedRecordLifecycle(t *testing.T) { require.Len(t, records, 1) assert.Equal(t, 2, records[0].Attempt) - require.NoError(t, store.removeParked(parked.ConsumerGroup, parked.Topic, parked.MessageID)) + partition := partitionIdentity(parked.Tenant, parked.PartitionKey) + require.NoError(t, store.removeParked(parked.ConsumerGroup, parked.Topic, partition, parked.MessageID)) records, err = store.ListParked(ctx, parked.ConsumerGroup) require.NoError(t, err) assert.Empty(t, records) // Removing an already-absent record is a no-op. - require.NoError(t, store.removeParked(parked.ConsumerGroup, parked.Topic, parked.MessageID)) + require.NoError(t, store.removeParked(parked.ConsumerGroup, parked.Topic, partition, parked.MessageID)) +} + +func TestParkedRecordsWithSameMessageIDAcrossTenantsAreIndependent(t *testing.T) { + ctx := context.Background() + store := New(t.TempDir()) + base := consumergate.Parked{ + ConsumerGroup: "group", + Topic: "topic", + MessageID: "shared-id", + PartitionKey: "shared", + ParkedAtMs: 1, + } + tenantA := base + tenantA.Tenant = "tenant-a" + tenantB := base + tenantB.Tenant = "tenant-b" + + require.NoError(t, store.recordParked(tenantA)) + require.NoError(t, store.recordParked(tenantB)) + records, err := store.ListParked(ctx, base.ConsumerGroup) + require.NoError(t, err) + assert.ElementsMatch(t, []consumergate.Parked{tenantA, tenantB}, records) + + require.NoError(t, store.removeParked( + base.ConsumerGroup, + base.Topic, + partitionIdentity(tenantA.Tenant, tenantA.PartitionKey), + base.MessageID, + )) + records, err = store.ListParked(ctx, base.ConsumerGroup) + require.NoError(t, err) + assert.Equal(t, []consumergate.Parked{tenantB}, records) } func TestListParkedEmpty(t *testing.T) { @@ -176,13 +229,14 @@ func TestListParkedSkipsTempFiles(t *testing.T) { ConsumerGroup: "group", Topic: "topic", MessageID: "id", + Tenant: "tenant", 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") + tmpPath := filepath.Join(dir, "parked", "group", "topic", "t-tenant", "p-part", "id.json.tmp123") require.NoError(t, os.WriteFile(tmpPath, []byte("partial"), 0o644)) records, err := store.ListParked(ctx, "group") @@ -192,7 +246,7 @@ func TestListParkedSkipsTempFiles(t *testing.T) { func TestMissingDirIsNotGated(t *testing.T) { store := New(filepath.Join(t.TempDir(), "does-not-exist")) - gated, err := store.isGated("group", "part") + gated, err := store.isGated("group", partitionIdentity("", "part")) require.NoError(t, err) assert.False(t, gated) } @@ -208,7 +262,7 @@ func TestEnter_OpenGateUnblocked(t *testing.T) { Attempt: 1, } - entry, err := store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", PartitionKey: "part"}) + entry, err := store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", Partition: partitionIdentity("", "part")}) require.NoError(t, err) assert.False(t, entry.Blocked()) @@ -234,7 +288,7 @@ func TestEnter_ClosedGateParkAndRelease(t *testing.T) { Attempt: 1, } - entry, err := store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", PartitionKey: "part"}) + entry, err := store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", Partition: partitionIdentity("", "part")}) require.NoError(t, err) require.True(t, entry.Blocked()) @@ -262,7 +316,7 @@ func TestEnter_ClosedGateParkAndRelease(t *testing.T) { // Open the gate; the next Enter is unblocked and Unpark removes the record. require.NoError(t, store.Open(ctx, key)) - entry, err = store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", PartitionKey: "part"}) + entry, err = store.Enter(ctx, consumergate.Key{ConsumerGroup: "group", Partition: partitionIdentity("", "part")}) require.NoError(t, err) require.False(t, entry.Blocked()) require.NoError(t, entry.Unpark(ctx, descriptor)) @@ -278,6 +332,6 @@ func TestEnter_MediumError(t *testing.T) { require.NoError(t, os.WriteFile(dir, []byte("x"), 0o644)) store := New(dir) - _, err := store.Enter(context.Background(), consumergate.Key{ConsumerGroup: "group", PartitionKey: "part"}) + _, err := store.Enter(context.Background(), consumergate.Key{ConsumerGroup: "group", Partition: partitionIdentity("", "part")}) require.Error(t, err) } diff --git a/platform/extension/consumergate/noop/BUILD.bazel b/platform/extension/consumergate/noop/BUILD.bazel index d4c9c2a0c..83b889e6d 100644 --- a/platform/extension/consumergate/noop/BUILD.bazel +++ b/platform/extension/consumergate/noop/BUILD.bazel @@ -13,6 +13,7 @@ go_test( srcs = ["gate_test.go"], embed = [":go_default_library"], deps = [ + "//platform/base/messagequeue:go_default_library", "//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/noop/gate_test.go b/platform/extension/consumergate/noop/gate_test.go index 6a8c7df03..36ccaf7e2 100644 --- a/platform/extension/consumergate/noop/gate_test.go +++ b/platform/extension/consumergate/noop/gate_test.go @@ -20,12 +20,16 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/extension/consumergate" ) func TestGate_EnterNeverBlocks(t *testing.T) { g := New() - entry, err := g.Enter(context.Background(), consumergate.Key{ConsumerGroup: "group", PartitionKey: "part"}) + entry, err := g.Enter(context.Background(), consumergate.Key{ + ConsumerGroup: "group", + Partition: entityqueue.PartitionIdentity{PartitionKey: "part"}, + }) require.NoError(t, err) assert.False(t, entry.Blocked()) diff --git a/platform/extension/messagequeue/mysql/BUILD.bazel b/platform/extension/messagequeue/mysql/BUILD.bazel index 12b2f7a95..fbc1f83b9 100644 --- a/platform/extension/messagequeue/mysql/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/BUILD.bazel @@ -6,6 +6,7 @@ go_library( "constants.go", "delivery_state_store.go", "errors.go", + "identifier.go", "message_store.go", "mock_stores.go", "offset_store.go", diff --git a/platform/extension/messagequeue/mysql/constants.go b/platform/extension/messagequeue/mysql/constants.go index 1d5c7fc9a..ca9da3adc 100644 --- a/platform/extension/messagequeue/mysql/constants.go +++ b/platform/extension/messagequeue/mysql/constants.go @@ -18,6 +18,7 @@ package mysql const ( // Common log field names (used extensively across all stores) + logTenant = "tenant" logTopic = "topic" logPartitionKey = "partition_key" logMessageID = "message_id" diff --git a/platform/extension/messagequeue/mysql/ctl/BUILD.bazel b/platform/extension/messagequeue/mysql/ctl/BUILD.bazel index 7804469f4..b7dd40e04 100644 --- a/platform/extension/messagequeue/mysql/ctl/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/ctl/BUILD.bazel @@ -1,4 +1,4 @@ -load("@rules_go//go:def.bzl", "go_binary", "go_library") +load("@rules_go//go:def.bzl", "go_binary", "go_library", "go_test") go_library( name = "go_default_library", @@ -20,3 +20,14 @@ go_binary( embed = [":go_default_library"], visibility = ["//visibility:public"], ) + +go_test( + name = "go_default_test", + srcs = ["commands_test.go"], + embed = [":go_default_library"], + deps = [ + "//platform/extension/messagequeue/mysql/ctl/lib:go_default_library", + "@com_github_spf13_cobra//:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/platform/extension/messagequeue/mysql/ctl/README.md b/platform/extension/messagequeue/mysql/ctl/README.md index e6cb64c2c..44e06d23c 100644 --- a/platform/extension/messagequeue/mysql/ctl/README.md +++ b/platform/extension/messagequeue/mysql/ctl/README.md @@ -24,40 +24,45 @@ Or pass `--dsn` on every command. Via Make (uses Bazel): ```bash -make run-queue-admin ARGS="list-topics" -make run-queue-admin ARGS="topic-stats --topic merge_queue" +make run-queue-admin ARGS="list-topics --tenant monorepo/main" +make run-queue-admin ARGS="topic-stats --tenant monorepo/main --topic merge_queue" ``` Via Bazel directly: ```bash -bazel run //platform/extension/messagequeue/mysql/ctl -- list-topics -bazel run //platform/extension/messagequeue/mysql/ctl -- topic-stats --topic merge_queue +bazel run //platform/extension/messagequeue/mysql/ctl -- list-topics --tenant monorepo/main +bazel run //platform/extension/messagequeue/mysql/ctl -- topic-stats --tenant monorepo/main --topic merge_queue ``` ## Commands +Commands operate on one tenant unless they explicitly accept `--all-tenants`. The `list-topics`, `list-offsets`, `list-leases`, and `stale-leases` commands require exactly one of `--tenant` or `--all-tenants`; fleet-wide queries are never implicit. + ### Inspect Topics ```bash -# List all topics with message counts -queue-admin list-topics +# List one tenant's topics with message counts +queue-admin list-topics --tenant monorepo/main + +# Explicitly list topics across every tenant +queue-admin list-topics --all-tenants # Detailed stats for a topic (total messages, DLQ count, partitions, consumer groups) -queue-admin topic-stats --topic merge_queue +queue-admin topic-stats --tenant monorepo/main --topic merge_queue ``` ### Inspect Messages ```bash # List messages (default limit 50) -queue-admin list-messages --topic merge_queue +queue-admin list-messages --tenant monorepo/main --topic merge_queue # Filter by partition, custom limit -queue-admin list-messages --topic merge_queue --partition uber/cadence --limit 10 +queue-admin list-messages --tenant monorepo/main --topic merge_queue --partition uber/cadence --limit 10 # Full message details including payload and metadata -queue-admin inspect-message --topic merge_queue --message-id msg-123 +queue-admin inspect-message --tenant monorepo/main --topic merge_queue --partition uber/cadence --message-id msg-123 ``` ### Manage Messages @@ -66,13 +71,13 @@ Destructive commands prompt for confirmation by default. Use `--no-interactive` ```bash # Delete a single message -queue-admin delete-message --topic merge_queue --message-id msg-123 +queue-admin delete-message --tenant monorepo/main --topic merge_queue --partition uber/cadence --message-id msg-123 # Purge all messages from a topic -queue-admin purge-topic --topic merge_queue +queue-admin purge-topic --tenant monorepo/main --topic merge_queue # Skip confirmation prompt (for scripting) -queue-admin purge-topic --topic merge_queue --no-interactive +queue-admin purge-topic --tenant monorepo/main --topic merge_queue --no-interactive ``` ### Dead Letter Queue (DLQ) @@ -81,26 +86,26 @@ DLQ messages live in the same `queue_messages` table under `topic + "_dlq"` (def ```bash # List DLQ messages -queue-admin list-dlq --topic merge_queue +queue-admin list-dlq --tenant monorepo/main --topic merge_queue # Inspect a DLQ message (use the DLQ topic name) -queue-admin inspect-message --topic merge_queue_dlq --message-id msg-456 +queue-admin inspect-message --tenant monorepo/main --topic merge_queue_dlq --partition uber/cadence --message-id msg-456 # Move a DLQ message back to the original topic -queue-admin requeue-dlq --topic merge_queue --message-id msg-456 +queue-admin requeue-dlq --tenant monorepo/main --topic merge_queue --partition uber/cadence --message-id msg-456 # Purge all DLQ messages -queue-admin purge-dlq --topic merge_queue +queue-admin purge-dlq --tenant monorepo/main --topic merge_queue # Custom DLQ suffix (if not using default "_dlq") -queue-admin list-dlq --topic merge_queue --dlq-suffix _dead +queue-admin list-dlq --tenant monorepo/main --topic merge_queue --dlq-suffix _dead ``` ### Consumer Lag ```bash # Per-partition lag for all consumer groups on a topic -queue-admin consumer-lag --topic merge_queue +queue-admin consumer-lag --tenant monorepo/main --topic merge_queue ``` Output shows `ACKED` (last processed offset), `LATEST` (newest message offset), and `LAG` (unprocessed count) per partition per consumer group. @@ -108,31 +113,38 @@ Output shows `ACKED` (last processed offset), `LATEST` (newest message offset), ### Consumer Offsets ```bash -# List all consumer group offsets -queue-admin list-offsets +# List one tenant's consumer group offsets +queue-admin list-offsets --tenant monorepo/main + +# Explicitly list offsets across every tenant +queue-admin list-offsets --all-tenants # Filter by consumer group -queue-admin list-offsets --consumer-group orchestrator +queue-admin list-offsets --tenant monorepo/main --consumer-group orchestrator # Reset offset to 0 (reprocess all messages) -queue-admin reset-offset --consumer-group orchestrator --topic merge_queue --partition uber/cadence +queue-admin reset-offset --tenant monorepo/main --consumer-group orchestrator --topic merge_queue --partition uber/cadence # Reset to a specific offset -queue-admin reset-offset --consumer-group orchestrator --topic merge_queue --partition uber/cadence --offset 42 +queue-admin reset-offset --tenant monorepo/main --consumer-group orchestrator --topic merge_queue --partition uber/cadence --offset 42 ``` ### Partition Leases ```bash -# List all active partition leases (who owns what) -queue-admin list-leases +# List one tenant's active partition leases +queue-admin list-leases --tenant monorepo/main + +# Explicitly list active leases across every tenant +queue-admin list-leases --all-tenants # Find stale leases (not renewed within threshold, likely dead workers) -queue-admin stale-leases # default 60s threshold -queue-admin stale-leases --threshold 30000 # 30s threshold +queue-admin stale-leases --tenant monorepo/main # default 60s threshold +queue-admin stale-leases --tenant monorepo/main --threshold 30000 # 30s threshold +queue-admin stale-leases --all-tenants # explicit fleet-wide query # Force-release a stuck lease -queue-admin release-lease --consumer-group orchestrator --topic merge_queue --partition uber/cadence +queue-admin release-lease --tenant monorepo/main --consumer-group orchestrator --topic merge_queue --partition uber/cadence ``` ### JSON Output @@ -140,7 +152,7 @@ queue-admin release-lease --consumer-group orchestrator --topic merge_queue --pa Add `--json` to any read command for machine-readable output: ```bash -queue-admin list-topics --json -queue-admin consumer-lag --topic merge_queue --json -queue-admin list-messages --topic merge_queue --json | jq '.[] | .ID' +queue-admin list-topics --tenant monorepo/main --json +queue-admin consumer-lag --tenant monorepo/main --topic merge_queue --json +queue-admin list-messages --tenant monorepo/main --topic merge_queue --json | jq '.[] | .ID' ``` diff --git a/platform/extension/messagequeue/mysql/ctl/commands.go b/platform/extension/messagequeue/mysql/ctl/commands.go index afb351265..c381b4502 100644 --- a/platform/extension/messagequeue/mysql/ctl/commands.go +++ b/platform/extension/messagequeue/mysql/ctl/commands.go @@ -111,36 +111,47 @@ func newRootCmd() *cobra.Command { return rootCmd } +func addTenantScopeFlags(cmd *cobra.Command, tenant *string, allTenants *bool) { + cmd.Flags().StringVar(tenant, "tenant", "", "Tenant name") + cmd.Flags().BoolVar(allTenants, "all-tenants", false, "Query all tenants") + cmd.MarkFlagsOneRequired("tenant", "all-tenants") + cmd.MarkFlagsMutuallyExclusive("tenant", "all-tenants") +} + func newListTopicsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - return &cobra.Command{ + var tenant string + var allTenants bool + cmd := &cobra.Command{ Use: "list-topics", Short: "List all topics with message counts", RunE: func(cmd *cobra.Command, args []string) error { - topics, err := (*store).ListTopics(cmd.Context()) + topics, err := (*store).ListTopics(cmd.Context(), lib.TenantScope{Tenant: tenant, AllTenants: allTenants}) if err != nil { return err } if *jsonOut { return lib.FormatJSON(os.Stdout, topics) } - headers := []string{"TOPIC", "MESSAGES"} + headers := []string{"TENANT", "TOPIC", "MESSAGES"} var rows [][]string for _, t := range topics { - rows = append(rows, []string{t.Topic, strconv.FormatInt(t.MessageCount, 10)}) + rows = append(rows, []string{t.Tenant, t.Topic, strconv.FormatInt(t.MessageCount, 10)}) } lib.FormatTable(os.Stdout, headers, rows) return nil }, } + addTenantScopeFlags(cmd, &tenant, &allTenants) + return cmd } func newTopicStatsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var topic, dlqSuffix string + var tenant, topic, dlqSuffix string cmd := &cobra.Command{ Use: "topic-stats", Short: "Show detailed statistics for a topic", RunE: func(cmd *cobra.Command, args []string) error { - stats, err := (*store).GetTopicStats(cmd.Context(), topic, dlqSuffix) + stats, err := (*store).GetTopicStats(cmd.Context(), tenant, topic, dlqSuffix) if err != nil { return err } @@ -149,6 +160,7 @@ func newTopicStatsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { } headers := []string{"FIELD", "VALUE"} rows := [][]string{ + {"Tenant", stats.Tenant}, {"Topic", stats.Topic}, {"Total Messages", strconv.FormatInt(stats.TotalMessages, 10)}, {"DLQ Count", strconv.FormatInt(stats.DLQCount, 10)}, @@ -159,20 +171,22 @@ func newTopicStatsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") cmd.Flags().StringVar(&dlqSuffix, "dlq-suffix", "_dlq", "DLQ topic suffix") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newListMessagesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var topic, partition string + var tenant, topic, partition string var limit int cmd := &cobra.Command{ Use: "list-messages", Short: "List messages for a topic", RunE: func(cmd *cobra.Command, args []string) error { - messages, err := (*store).ListMessages(cmd.Context(), topic, partition, limit) + messages, err := (*store).ListMessages(cmd.Context(), tenant, topic, partition, limit) if err != nil { return err } @@ -194,31 +208,34 @@ func newListMessagesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") cmd.Flags().StringVar(&partition, "partition", "", "Filter by partition key") cmd.Flags().IntVar(&limit, "limit", 50, "Maximum number of messages to show") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newInspectMessageCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var topic, messageID string + var tenant, topic, partition, messageID string cmd := &cobra.Command{ Use: "inspect-message", Short: "Show full message details including payload and metadata", RunE: func(cmd *cobra.Command, args []string) error { - detail, found, err := (*store).InspectMessage(cmd.Context(), topic, messageID) + detail, found, err := (*store).InspectMessage(cmd.Context(), tenant, topic, partition, messageID) if err != nil { return err } if !found { - return fmt.Errorf("message %q not found in topic %q", messageID, topic) + return fmt.Errorf("message %q not found in topic %q partition %q", messageID, topic, partition) } if *jsonOut { return lib.FormatJSON(os.Stdout, detail) } headers := []string{"FIELD", "VALUE"} rows := [][]string{ + {"Tenant", detail.Tenant}, {"Offset", strconv.FormatInt(detail.Offset, 10)}, {"ID", detail.ID}, {"Topic", detail.Topic}, @@ -244,43 +261,51 @@ func newInspectMessageCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") + cmd.Flags().StringVar(&partition, "partition", "", "Partition key (required)") cmd.Flags().StringVar(&messageID, "message-id", "", "Message ID (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") + cmd.MarkFlagRequired("partition") cmd.MarkFlagRequired("message-id") return cmd } func newDeleteMessageCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var topic, messageID string + var tenant, topic, partition, messageID string cmd := &cobra.Command{ Use: "delete-message", Short: "Delete a specific message", RunE: func(cmd *cobra.Command, args []string) error { - if err := confirmAction(*noInteractive, fmt.Sprintf("Delete message %q from topic %q?", messageID, topic)); err != nil { + if err := confirmAction(*noInteractive, fmt.Sprintf("Delete message %q from topic %q partition %q?", messageID, topic, partition)); err != nil { return err } - affected, err := (*store).DeleteMessage(cmd.Context(), topic, messageID) + affected, err := (*store).DeleteMessage(cmd.Context(), tenant, topic, partition, messageID) if err != nil { return err } if affected == 0 { - fmt.Fprintf(os.Stderr, "No message found with ID %q in topic %q\n", messageID, topic) + fmt.Fprintf(os.Stderr, "No message found with ID %q in topic %q partition %q\n", messageID, topic, partition) return nil } - fmt.Printf("Deleted message %q from topic %q\n", messageID, topic) + fmt.Printf("Deleted message %q from topic %q partition %q\n", messageID, topic, partition) return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") + cmd.Flags().StringVar(&partition, "partition", "", "Partition key (required)") cmd.Flags().StringVar(&messageID, "message-id", "", "Message ID (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") + cmd.MarkFlagRequired("partition") cmd.MarkFlagRequired("message-id") return cmd } func newPurgeTopicCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var topic string + var tenant, topic string cmd := &cobra.Command{ Use: "purge-topic", Short: "Delete all messages for a topic", @@ -288,7 +313,7 @@ func newPurgeTopicCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comman if err := confirmAction(*noInteractive, fmt.Sprintf("Purge ALL messages from topic %q?", topic)); err != nil { return err } - affected, err := (*store).PurgeTopic(cmd.Context(), topic) + affected, err := (*store).PurgeTopic(cmd.Context(), tenant, topic) if err != nil { return err } @@ -296,20 +321,22 @@ func newPurgeTopicCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comman return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newListDLQCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var topic, dlqSuffix string + var tenant, topic, dlqSuffix string var limit int cmd := &cobra.Command{ Use: "list-dlq", Short: "List dead-letter queue messages for a topic", RunE: func(cmd *cobra.Command, args []string) error { dlqTopic := topic + dlqSuffix - messages, err := (*store).ListMessages(cmd.Context(), dlqTopic, "", limit) + messages, err := (*store).ListMessages(cmd.Context(), tenant, dlqTopic, "", limit) if err != nil { return err } @@ -330,36 +357,42 @@ func newListDLQCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Original topic name (required)") cmd.Flags().StringVar(&dlqSuffix, "dlq-suffix", "_dlq", "DLQ topic suffix") cmd.Flags().IntVar(&limit, "limit", 50, "Maximum number of messages to show") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newRequeueDLQCmd(store **lib.AdminStore) *cobra.Command { - var topic, messageID, dlqSuffix string + var tenant, topic, partition, messageID, dlqSuffix string cmd := &cobra.Command{ Use: "requeue-dlq", Short: "Move a DLQ message back to its original topic", RunE: func(cmd *cobra.Command, args []string) error { - if err := (*store).RequeueDLQ(cmd.Context(), topic, messageID, dlqSuffix); err != nil { + if err := (*store).RequeueDLQ(cmd.Context(), tenant, topic, partition, messageID, dlqSuffix); err != nil { return err } fmt.Printf("Requeued message %q from DLQ back to topic %q\n", messageID, topic) return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Original topic name (required)") + cmd.Flags().StringVar(&partition, "partition", "", "Partition key (required)") cmd.Flags().StringVar(&messageID, "message-id", "", "Message ID (required)") cmd.Flags().StringVar(&dlqSuffix, "dlq-suffix", "_dlq", "DLQ topic suffix") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") + cmd.MarkFlagRequired("partition") cmd.MarkFlagRequired("message-id") return cmd } func newPurgeDLQCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var topic, dlqSuffix string + var tenant, topic, dlqSuffix string cmd := &cobra.Command{ Use: "purge-dlq", Short: "Delete all DLQ messages for a topic", @@ -368,7 +401,7 @@ func newPurgeDLQCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command if err := confirmAction(*noInteractive, fmt.Sprintf("Purge ALL messages from DLQ topic %q?", dlqTopic)); err != nil { return err } - affected, err := (*store).PurgeTopic(cmd.Context(), dlqTopic) + affected, err := (*store).PurgeTopic(cmd.Context(), tenant, dlqTopic) if err != nil { return err } @@ -376,29 +409,33 @@ func newPurgeDLQCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Original topic name (required)") cmd.Flags().StringVar(&dlqSuffix, "dlq-suffix", "_dlq", "DLQ topic suffix") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newListOffsetsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var consumerGroup string + var tenant, consumerGroup string + var allTenants bool cmd := &cobra.Command{ Use: "list-offsets", Short: "Show consumer group offsets", RunE: func(cmd *cobra.Command, args []string) error { - offsets, err := (*store).ListOffsets(cmd.Context(), consumerGroup) + offsets, err := (*store).ListOffsets(cmd.Context(), lib.TenantScope{Tenant: tenant, AllTenants: allTenants}, consumerGroup) if err != nil { return err } if *jsonOut { return lib.FormatJSON(os.Stdout, offsets) } - headers := []string{"CONSUMER_GROUP", "TOPIC", "PARTITION", "OFFSET_ACKED", "UPDATED_AT"} + headers := []string{"TENANT", "CONSUMER_GROUP", "TOPIC", "PARTITION", "OFFSET_ACKED", "UPDATED_AT"} var rows [][]string for _, o := range offsets { rows = append(rows, []string{ + o.Tenant, o.ConsumerGroup, o.Topic, o.PartitionKey, @@ -410,12 +447,13 @@ func newListOffsetsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + addTenantScopeFlags(cmd, &tenant, &allTenants) cmd.Flags().StringVar(&consumerGroup, "consumer-group", "", "Filter by consumer group") return cmd } func newResetOffsetCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var consumerGroup, topic, partition string + var tenant, consumerGroup, topic, partition string var offset int64 cmd := &cobra.Command{ Use: "reset-offset", @@ -424,7 +462,7 @@ func newResetOffsetCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comma if err := confirmAction(*noInteractive, fmt.Sprintf("Reset offset to %d for consumer-group=%q topic=%q partition=%q?", offset, consumerGroup, topic, partition)); err != nil { return err } - affected, err := (*store).ResetOffset(cmd.Context(), consumerGroup, topic, partition, offset) + affected, err := (*store).ResetOffset(cmd.Context(), tenant, consumerGroup, topic, partition, offset) if err != nil { return err } @@ -436,10 +474,12 @@ func newResetOffsetCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comma return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&consumerGroup, "consumer-group", "", "Consumer group name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") cmd.Flags().StringVar(&partition, "partition", "", "Partition key (required)") cmd.Flags().Int64Var(&offset, "offset", 0, "New offset value (default 0)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("consumer-group") cmd.MarkFlagRequired("topic") cmd.MarkFlagRequired("partition") @@ -447,21 +487,24 @@ func newResetOffsetCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comma } func newListLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - return &cobra.Command{ + var tenant string + var allTenants bool + cmd := &cobra.Command{ Use: "list-leases", Short: "Show all active partition leases", RunE: func(cmd *cobra.Command, args []string) error { - leases, err := (*store).ListLeases(cmd.Context()) + leases, err := (*store).ListLeases(cmd.Context(), lib.TenantScope{Tenant: tenant, AllTenants: allTenants}) if err != nil { return err } if *jsonOut { return lib.FormatJSON(os.Stdout, leases) } - headers := []string{"CONSUMER_GROUP", "TOPIC", "PARTITION", "LEASED_BY", "LEASED_AT", "RENEWED_AT"} + headers := []string{"TENANT", "CONSUMER_GROUP", "TOPIC", "PARTITION", "LEASED_BY", "LEASED_AT", "RENEWED_AT"} var rows [][]string for _, l := range leases { rows = append(rows, []string{ + l.Tenant, l.ConsumerGroup, l.Topic, l.PartitionKey, @@ -474,15 +517,17 @@ func newListLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + addTenantScopeFlags(cmd, &tenant, &allTenants) + return cmd } func newConsumerLagCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - var topic string + var tenant, topic string cmd := &cobra.Command{ Use: "consumer-lag", Short: "Show per-partition consumer lag for a topic", RunE: func(cmd *cobra.Command, args []string) error { - lags, err := (*store).ConsumerLag(cmd.Context(), topic) + lags, err := (*store).ConsumerLag(cmd.Context(), tenant, topic) if err != nil { return err } @@ -505,18 +550,22 @@ func newConsumerLagCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") return cmd } func newStaleLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { + var tenant string + var allTenants bool var thresholdMs int64 cmd := &cobra.Command{ Use: "stale-leases", Short: "Show partition leases not renewed within a threshold", RunE: func(cmd *cobra.Command, args []string) error { - leases, err := (*store).StaleLeases(cmd.Context(), thresholdMs) + leases, err := (*store).StaleLeases(cmd.Context(), lib.TenantScope{Tenant: tenant, AllTenants: allTenants}, thresholdMs) if err != nil { return err } @@ -527,10 +576,11 @@ func newStaleLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { if *jsonOut { return lib.FormatJSON(os.Stdout, leases) } - headers := []string{"CONSUMER_GROUP", "TOPIC", "PARTITION", "LEASED_BY", "LEASED_AT", "RENEWED_AT"} + headers := []string{"TENANT", "CONSUMER_GROUP", "TOPIC", "PARTITION", "LEASED_BY", "LEASED_AT", "RENEWED_AT"} var rows [][]string for _, l := range leases { rows = append(rows, []string{ + l.Tenant, l.ConsumerGroup, l.Topic, l.PartitionKey, @@ -543,12 +593,13 @@ func newStaleLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + addTenantScopeFlags(cmd, &tenant, &allTenants) cmd.Flags().Int64Var(&thresholdMs, "threshold", 60000, "Staleness threshold in milliseconds (default 60s)") return cmd } func newReleaseLeaseCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var consumerGroup, topic, partition string + var tenant, consumerGroup, topic, partition string cmd := &cobra.Command{ Use: "release-lease", Short: "Force-release a partition lease", @@ -556,7 +607,7 @@ func newReleaseLeaseCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comm if err := confirmAction(*noInteractive, fmt.Sprintf("Release lease for consumer-group=%q topic=%q partition=%q?", consumerGroup, topic, partition)); err != nil { return err } - affected, err := (*store).ReleaseLease(cmd.Context(), consumerGroup, topic, partition) + affected, err := (*store).ReleaseLease(cmd.Context(), tenant, consumerGroup, topic, partition) if err != nil { return err } @@ -568,9 +619,11 @@ func newReleaseLeaseCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comm return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&consumerGroup, "consumer-group", "", "Consumer group name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") cmd.Flags().StringVar(&partition, "partition", "", "Partition key (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("consumer-group") cmd.MarkFlagRequired("topic") cmd.MarkFlagRequired("partition") diff --git a/platform/extension/messagequeue/mysql/ctl/commands_test.go b/platform/extension/messagequeue/mysql/ctl/commands_test.go new file mode 100644 index 000000000..ce8550429 --- /dev/null +++ b/platform/extension/messagequeue/mysql/ctl/commands_test.go @@ -0,0 +1,75 @@ +// 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 main + +import ( + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/platform/extension/messagequeue/mysql/ctl/lib" +) + +func TestListCommandsRequireExplicitTenantScope(t *testing.T) { + var store *lib.AdminStore + jsonOut := false + commandFactories := map[string]func() *cobra.Command{ + "list topics": func() *cobra.Command { return newListTopicsCmd(&store, &jsonOut) }, + "list offsets": func() *cobra.Command { return newListOffsetsCmd(&store, &jsonOut) }, + "list leases": func() *cobra.Command { return newListLeasesCmd(&store, &jsonOut) }, + "stale leases": func() *cobra.Command { return newStaleLeasesCmd(&store, &jsonOut) }, + } + + for name, commandFactory := range commandFactories { + t.Run(name, func(t *testing.T) { + t.Run("missing scope", func(t *testing.T) { + cmd := commandFactory() + cmd.SetArgs(nil) + cmd.SilenceErrors = true + cmd.SilenceUsage = true + require.Error(t, cmd.Execute()) + }) + + t.Run("conflicting scope", func(t *testing.T) { + cmd := commandFactory() + cmd.SetArgs([]string{"--tenant", "acme", "--all-tenants"}) + cmd.SilenceErrors = true + cmd.SilenceUsage = true + require.Error(t, cmd.Execute()) + }) + }) + } +} + +func TestMessageCommandsRequirePartition(t *testing.T) { + var store *lib.AdminStore + jsonOut := false + noInteractive := true + commandFactories := map[string]func() *cobra.Command{ + "inspect": func() *cobra.Command { return newInspectMessageCmd(&store, &jsonOut) }, + "delete": func() *cobra.Command { return newDeleteMessageCmd(&store, &noInteractive) }, + "requeue": func() *cobra.Command { return newRequeueDLQCmd(&store) }, + } + + for name, commandFactory := range commandFactories { + t.Run(name, func(t *testing.T) { + cmd := commandFactory() + cmd.SetArgs([]string{"--tenant", "acme", "--topic", "orders", "--message-id", "msg-1"}) + cmd.SilenceErrors = true + cmd.SilenceUsage = true + require.Error(t, cmd.Execute()) + }) + } +} diff --git a/platform/extension/messagequeue/mysql/ctl/lib/admin.go b/platform/extension/messagequeue/mysql/ctl/lib/admin.go index 8ccdc77a2..c7693ea8b 100644 --- a/platform/extension/messagequeue/mysql/ctl/lib/admin.go +++ b/platform/extension/messagequeue/mysql/ctl/lib/admin.go @@ -30,6 +30,21 @@ type AdminStore struct { db *sql.DB } +// TenantScope selects either one tenant or every tenant. +type TenantScope struct { + // Tenant is the single tenant to select. + Tenant string + // AllTenants explicitly selects every tenant. + AllTenants bool +} + +func (s TenantScope) validate() error { + if (s.Tenant != "") == s.AllTenants { + return fmt.Errorf("exactly one of tenant or all tenants is required") + } + return nil +} + // NewAdminStore creates a new AdminStore backed by the given database connection. func NewAdminStore(db *sql.DB) *AdminStore { return &AdminStore{db: db} @@ -37,6 +52,8 @@ func NewAdminStore(db *sql.DB) *AdminStore { // MessageSummary contains a subset of message fields for listing. type MessageSummary struct { + // Tenant identifies the queue tenant. + Tenant string // Offset is the auto-incrementing sequence number Offset int64 // ID is the unique message identifier @@ -70,6 +87,8 @@ type MessageDetail struct { // OffsetInfo contains consumer group offset information. type OffsetInfo struct { + // Tenant identifies the queue tenant. + Tenant string // ConsumerGroup is the consumer group name ConsumerGroup string // Topic is the topic being consumed @@ -84,6 +103,8 @@ type OffsetInfo struct { // LeaseInfo contains partition lease information. type LeaseInfo struct { + // Tenant identifies the queue tenant. + Tenant string // ConsumerGroup is the consumer group name ConsumerGroup string // Topic is the topic being consumed @@ -100,6 +121,8 @@ type LeaseInfo struct { // TopicInfo contains a topic name and its message count. type TopicInfo struct { + // Tenant identifies the queue tenant. + Tenant string // Topic is the queue topic name Topic string // MessageCount is the number of messages in this topic @@ -108,6 +131,8 @@ type TopicInfo struct { // TopicStats contains detailed statistics for a topic. type TopicStats struct { + // Tenant identifies the queue tenant. + Tenant string // Topic is the queue topic name Topic string // TotalMessages is the total number of messages @@ -120,13 +145,18 @@ type TopicStats struct { ConsumerGroupCount int64 } -// ListTopics returns all topics with their message counts. -func (s *AdminStore) ListTopics(ctx context.Context) ([]TopicInfo, error) { - query := fmt.Sprintf( - "SELECT topic, COUNT(*) FROM %s GROUP BY topic ORDER BY topic", - mysql.MessagesTableName, - ) - rows, err := s.db.QueryContext(ctx, query) +// ListTopics returns topics and their message counts within scope. +func (s *AdminStore) ListTopics(ctx context.Context, scope TenantScope) ([]TopicInfo, error) { + if err := scope.validate(); err != nil { + return nil, err + } + query := fmt.Sprintf("SELECT tenant, topic, COUNT(*) FROM %s GROUP BY tenant, topic ORDER BY tenant, topic", mysql.MessagesTableName) + var args []any + if !scope.AllTenants { + query = fmt.Sprintf("SELECT tenant, topic, COUNT(*) FROM %s WHERE tenant = ? GROUP BY tenant, topic ORDER BY tenant, topic", mysql.MessagesTableName) + args = append(args, scope.Tenant) + } + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("list topics: %w", err) } @@ -135,7 +165,7 @@ func (s *AdminStore) ListTopics(ctx context.Context) ([]TopicInfo, error) { var topics []TopicInfo for rows.Next() { var t TopicInfo - if err := rows.Scan(&t.Topic, &t.MessageCount); err != nil { + if err := rows.Scan(&t.Tenant, &t.Topic, &t.MessageCount); err != nil { return nil, fmt.Errorf("scan topic row: %w", err) } topics = append(topics, t) @@ -144,13 +174,13 @@ func (s *AdminStore) ListTopics(ctx context.Context) ([]TopicInfo, error) { } // GetTopicStats returns detailed statistics for a topic. -func (s *AdminStore) GetTopicStats(ctx context.Context, topic string, dlqSuffix string) (TopicStats, error) { - stats := TopicStats{Topic: topic} +func (s *AdminStore) GetTopicStats(ctx context.Context, tenant string, topic string, dlqSuffix string) (TopicStats, error) { + stats := TopicStats{Tenant: tenant, Topic: topic} // Total messages err := s.db.QueryRowContext(ctx, - fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE topic = ?", mysql.MessagesTableName), - topic, + fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE tenant = ? AND topic = ?", mysql.MessagesTableName), + tenant, topic, ).Scan(&stats.TotalMessages) if err != nil { return stats, fmt.Errorf("count total: %w", err) @@ -159,8 +189,8 @@ func (s *AdminStore) GetTopicStats(ctx context.Context, topic string, dlqSuffix // DLQ count dlqTopic := topic + dlqSuffix err = s.db.QueryRowContext(ctx, - fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE topic = ?", mysql.MessagesTableName), - dlqTopic, + fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE tenant = ? AND topic = ?", mysql.MessagesTableName), + tenant, dlqTopic, ).Scan(&stats.DLQCount) if err != nil { return stats, fmt.Errorf("count dlq: %w", err) @@ -168,8 +198,8 @@ func (s *AdminStore) GetTopicStats(ctx context.Context, topic string, dlqSuffix // Distinct partitions err = s.db.QueryRowContext(ctx, - fmt.Sprintf("SELECT COUNT(DISTINCT partition_key) FROM %s WHERE topic = ?", mysql.MessagesTableName), - topic, + fmt.Sprintf("SELECT COUNT(DISTINCT partition_key) FROM %s WHERE tenant = ? AND topic = ?", mysql.MessagesTableName), + tenant, topic, ).Scan(&stats.PartitionCount) if err != nil { return stats, fmt.Errorf("count partitions: %w", err) @@ -177,8 +207,8 @@ func (s *AdminStore) GetTopicStats(ctx context.Context, topic string, dlqSuffix // Consumer groups from offsets err = s.db.QueryRowContext(ctx, - fmt.Sprintf("SELECT COUNT(DISTINCT consumer_group) FROM %s WHERE topic = ?", mysql.OffsetsTableName), - topic, + fmt.Sprintf("SELECT COUNT(DISTINCT consumer_group) FROM %s WHERE tenant = ? AND topic = ?", mysql.OffsetsTableName), + tenant, topic, ).Scan(&stats.ConsumerGroupCount) if err != nil { return stats, fmt.Errorf("count consumer groups: %w", err) @@ -188,19 +218,19 @@ func (s *AdminStore) GetTopicStats(ctx context.Context, topic string, dlqSuffix } // ListMessages returns messages for a topic, optionally filtered by partition. -func (s *AdminStore) ListMessages(ctx context.Context, topic string, partition string, limit int) ([]MessageSummary, error) { +func (s *AdminStore) ListMessages(ctx context.Context, tenant string, topic string, partition string, limit int) ([]MessageSummary, error) { var rows *sql.Rows var err error if partition != "" { rows, err = s.db.QueryContext(ctx, - fmt.Sprintf("SELECT `offset`, id, topic, partition_key, created_at, published_at FROM %s WHERE topic = ? AND partition_key = ? ORDER BY `offset` LIMIT ?", mysql.MessagesTableName), - topic, partition, limit, + fmt.Sprintf("SELECT tenant, `offset`, id, topic, partition_key, created_at, published_at FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? ORDER BY `offset` LIMIT ?", mysql.MessagesTableName), + tenant, topic, partition, limit, ) } else { rows, err = s.db.QueryContext(ctx, - fmt.Sprintf("SELECT `offset`, id, topic, partition_key, created_at, published_at FROM %s WHERE topic = ? ORDER BY `offset` LIMIT ?", mysql.MessagesTableName), - topic, limit, + fmt.Sprintf("SELECT tenant, `offset`, id, topic, partition_key, created_at, published_at FROM %s WHERE tenant = ? AND topic = ? ORDER BY `offset` LIMIT ?", mysql.MessagesTableName), + tenant, topic, limit, ) } if err != nil { @@ -211,7 +241,7 @@ func (s *AdminStore) ListMessages(ctx context.Context, topic string, partition s var messages []MessageSummary for rows.Next() { var m MessageSummary - if err := rows.Scan(&m.Offset, &m.ID, &m.Topic, &m.PartitionKey, &m.CreatedAt, &m.PublishedAt); err != nil { + if err := rows.Scan(&m.Tenant, &m.Offset, &m.ID, &m.Topic, &m.PartitionKey, &m.CreatedAt, &m.PublishedAt); err != nil { return nil, fmt.Errorf("scan message row: %w", err) } messages = append(messages, m) @@ -220,14 +250,14 @@ func (s *AdminStore) ListMessages(ctx context.Context, topic string, partition s } // InspectMessage returns full message details including payload and DLQ fields. -func (s *AdminStore) InspectMessage(ctx context.Context, topic string, messageID string) (MessageDetail, bool, error) { +func (s *AdminStore) InspectMessage(ctx context.Context, tenant string, topic string, partition string, messageID string) (MessageDetail, bool, error) { var d MessageDetail var metadataJSON []byte err := s.db.QueryRowContext(ctx, - fmt.Sprintf("SELECT `offset`, id, topic, partition_key, created_at, published_at, payload, metadata, failed_at, failure_count, last_error, original_topic FROM %s WHERE topic = ? AND id = ?", mysql.MessagesTableName), - topic, messageID, - ).Scan(&d.Offset, &d.ID, &d.Topic, &d.PartitionKey, &d.CreatedAt, &d.PublishedAt, &d.Payload, &metadataJSON, &d.FailedAt, &d.FailureCount, &d.LastError, &d.OriginalTopic) + fmt.Sprintf("SELECT tenant, `offset`, id, topic, partition_key, created_at, published_at, payload, metadata, failed_at, failure_count, last_error, original_topic FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ?", mysql.MessagesTableName), + tenant, topic, partition, messageID, + ).Scan(&d.Tenant, &d.Offset, &d.ID, &d.Topic, &d.PartitionKey, &d.CreatedAt, &d.PublishedAt, &d.Payload, &metadataJSON, &d.FailedAt, &d.FailureCount, &d.LastError, &d.OriginalTopic) if err == sql.ErrNoRows { return d, false, nil } @@ -247,11 +277,11 @@ func (s *AdminStore) InspectMessage(ctx context.Context, topic string, messageID return d, true, nil } -// DeleteMessage deletes a specific message by topic and ID. -func (s *AdminStore) DeleteMessage(ctx context.Context, topic string, messageID string) (int64, error) { +// DeleteMessage deletes a specific message by its unique identity. +func (s *AdminStore) DeleteMessage(ctx context.Context, tenant string, topic string, partition string, messageID string) (int64, error) { result, err := s.db.ExecContext(ctx, - fmt.Sprintf("DELETE FROM %s WHERE topic = ? AND id = ?", mysql.MessagesTableName), - topic, messageID, + fmt.Sprintf("DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ?", mysql.MessagesTableName), + tenant, topic, partition, messageID, ) if err != nil { return 0, fmt.Errorf("delete message: %w", err) @@ -260,10 +290,10 @@ func (s *AdminStore) DeleteMessage(ctx context.Context, topic string, messageID } // PurgeTopic deletes all messages for a topic. -func (s *AdminStore) PurgeTopic(ctx context.Context, topic string) (int64, error) { +func (s *AdminStore) PurgeTopic(ctx context.Context, tenant string, topic string) (int64, error) { result, err := s.db.ExecContext(ctx, - fmt.Sprintf("DELETE FROM %s WHERE topic = ?", mysql.MessagesTableName), - topic, + fmt.Sprintf("DELETE FROM %s WHERE tenant = ? AND topic = ?", mysql.MessagesTableName), + tenant, topic, ) if err != nil { return 0, fmt.Errorf("purge topic: %w", err) @@ -273,7 +303,7 @@ func (s *AdminStore) PurgeTopic(ctx context.Context, topic string) (int64, error // RequeueDLQ moves a message from the DLQ topic back to its original topic. // This is done transactionally: read from DLQ, insert into original topic, delete from DLQ. -func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID string, dlqSuffix string) error { +func (s *AdminStore) RequeueDLQ(ctx context.Context, tenant string, topic string, partition string, messageID string, dlqSuffix string) error { dlqTopic := topic + dlqSuffix tx, err := s.db.BeginTx(ctx, nil) @@ -285,13 +315,12 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str // Read the DLQ message var payload []byte var metadataJSON []byte - var partitionKey string var createdAt, publishedAt int64 err = tx.QueryRowContext(ctx, - fmt.Sprintf("SELECT payload, metadata, partition_key, created_at, published_at FROM %s WHERE topic = ? AND id = ?", mysql.MessagesTableName), - dlqTopic, messageID, - ).Scan(&payload, &metadataJSON, &partitionKey, &createdAt, &publishedAt) + fmt.Sprintf("SELECT payload, metadata, created_at, published_at FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ?", mysql.MessagesTableName), + tenant, dlqTopic, partition, messageID, + ).Scan(&payload, &metadataJSON, &createdAt, &publishedAt) if err == sql.ErrNoRows { return fmt.Errorf("message %q not found in DLQ topic %q", messageID, dlqTopic) } @@ -302,8 +331,8 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str // Insert into original topic with reset fields nowMs := time.Now().UnixMilli() _, err = tx.ExecContext(ctx, - fmt.Sprintf("INSERT INTO %s (topic, partition_key, id, payload, metadata, created_at, published_at, failed_at, failure_count, last_error, original_topic) VALUES (?, ?, ?, ?, ?, ?, ?, 0, 0, '', '')", mysql.MessagesTableName), - topic, partitionKey, messageID, payload, metadataJSON, createdAt, nowMs, + fmt.Sprintf("INSERT INTO %s (tenant, topic, partition_key, id, payload, metadata, created_at, published_at, failed_at, failure_count, last_error, original_topic) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, '', '')", mysql.MessagesTableName), + tenant, topic, partition, messageID, payload, metadataJSON, createdAt, nowMs, ) if err != nil { return fmt.Errorf("insert requeued message: %w", err) @@ -311,8 +340,8 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str // Delete from DLQ _, err = tx.ExecContext(ctx, - fmt.Sprintf("DELETE FROM %s WHERE topic = ? AND id = ?", mysql.MessagesTableName), - dlqTopic, messageID, + fmt.Sprintf("DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ?", mysql.MessagesTableName), + tenant, dlqTopic, partition, messageID, ) if err != nil { return fmt.Errorf("delete dlq message: %w", err) @@ -321,19 +350,32 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str return tx.Commit() } -// ListOffsets returns consumer group offsets, optionally filtered by group. -func (s *AdminStore) ListOffsets(ctx context.Context, consumerGroup string) ([]OffsetInfo, error) { +// ListOffsets returns consumer group offsets within scope, optionally filtered by group. +func (s *AdminStore) ListOffsets(ctx context.Context, scope TenantScope, consumerGroup string) ([]OffsetInfo, error) { + if err := scope.validate(); err != nil { + return nil, err + } var rows *sql.Rows var err error - if consumerGroup != "" { + if !scope.AllTenants && consumerGroup != "" { rows, err = s.db.QueryContext(ctx, - fmt.Sprintf("SELECT consumer_group, topic, partition_key, offset_acked, updated_at FROM %s WHERE consumer_group = ? ORDER BY consumer_group, topic, partition_key", mysql.OffsetsTableName), + fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, offset_acked, updated_at FROM %s WHERE tenant = ? AND consumer_group = ? ORDER BY tenant, consumer_group, topic, partition_key", mysql.OffsetsTableName), + scope.Tenant, consumerGroup, + ) + } else if !scope.AllTenants { + rows, err = s.db.QueryContext(ctx, + fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, offset_acked, updated_at FROM %s WHERE tenant = ? ORDER BY tenant, consumer_group, topic, partition_key", mysql.OffsetsTableName), + scope.Tenant, + ) + } else if consumerGroup != "" { + rows, err = s.db.QueryContext(ctx, + fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, offset_acked, updated_at FROM %s WHERE consumer_group = ? ORDER BY tenant, consumer_group, topic, partition_key", mysql.OffsetsTableName), consumerGroup, ) } else { rows, err = s.db.QueryContext(ctx, - fmt.Sprintf("SELECT consumer_group, topic, partition_key, offset_acked, updated_at FROM %s ORDER BY consumer_group, topic, partition_key", mysql.OffsetsTableName), + fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, offset_acked, updated_at FROM %s ORDER BY tenant, consumer_group, topic, partition_key", mysql.OffsetsTableName), ) } if err != nil { @@ -344,7 +386,7 @@ func (s *AdminStore) ListOffsets(ctx context.Context, consumerGroup string) ([]O var offsets []OffsetInfo for rows.Next() { var o OffsetInfo - if err := rows.Scan(&o.ConsumerGroup, &o.Topic, &o.PartitionKey, &o.OffsetAcked, &o.UpdatedAt); err != nil { + if err := rows.Scan(&o.Tenant, &o.ConsumerGroup, &o.Topic, &o.PartitionKey, &o.OffsetAcked, &o.UpdatedAt); err != nil { return nil, fmt.Errorf("scan offset row: %w", err) } offsets = append(offsets, o) @@ -353,11 +395,11 @@ func (s *AdminStore) ListOffsets(ctx context.Context, consumerGroup string) ([]O } // ResetOffset updates the acked offset for a consumer group/topic/partition. -func (s *AdminStore) ResetOffset(ctx context.Context, consumerGroup, topic, partition string, offset int64) (int64, error) { +func (s *AdminStore) ResetOffset(ctx context.Context, tenant, consumerGroup, topic, partition string, offset int64) (int64, error) { nowMs := time.Now().UnixMilli() result, err := s.db.ExecContext(ctx, - fmt.Sprintf("UPDATE %s SET offset_acked = ?, updated_at = ? WHERE consumer_group = ? AND topic = ? AND partition_key = ?", mysql.OffsetsTableName), - offset, nowMs, consumerGroup, topic, partition, + fmt.Sprintf("UPDATE %s SET offset_acked = ?, updated_at = ? WHERE tenant = ? AND topic = ? AND partition_key = ? AND consumer_group = ?", mysql.OffsetsTableName), + offset, nowMs, tenant, topic, partition, consumerGroup, ) if err != nil { return 0, fmt.Errorf("reset offset: %w", err) @@ -365,11 +407,18 @@ func (s *AdminStore) ResetOffset(ctx context.Context, consumerGroup, topic, part return result.RowsAffected() } -// ListLeases returns all partition leases. -func (s *AdminStore) ListLeases(ctx context.Context) ([]LeaseInfo, error) { - rows, err := s.db.QueryContext(ctx, - fmt.Sprintf("SELECT consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s ORDER BY consumer_group, topic, partition_key", mysql.PartitionLeasesTableName), - ) +// ListLeases returns partition leases within scope. +func (s *AdminStore) ListLeases(ctx context.Context, scope TenantScope) ([]LeaseInfo, error) { + if err := scope.validate(); err != nil { + return nil, err + } + query := fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s ORDER BY tenant, consumer_group, topic, partition_key", mysql.PartitionLeasesTableName) + var args []any + if !scope.AllTenants { + query = fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s WHERE tenant = ? ORDER BY tenant, consumer_group, topic, partition_key", mysql.PartitionLeasesTableName) + args = append(args, scope.Tenant) + } + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("list leases: %w", err) } @@ -378,7 +427,7 @@ func (s *AdminStore) ListLeases(ctx context.Context) ([]LeaseInfo, error) { var leases []LeaseInfo for rows.Next() { var l LeaseInfo - if err := rows.Scan(&l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.LeasedBy, &l.LeasedAt, &l.LeaseRenewedAt); err != nil { + if err := rows.Scan(&l.Tenant, &l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.LeasedBy, &l.LeasedAt, &l.LeaseRenewedAt); err != nil { return nil, fmt.Errorf("scan lease row: %w", err) } leases = append(leases, l) @@ -388,6 +437,8 @@ func (s *AdminStore) ListLeases(ctx context.Context) ([]LeaseInfo, error) { // LagInfo contains consumer lag information for a single partition. type LagInfo struct { + // Tenant identifies the queue tenant. + Tenant string // ConsumerGroup is the consumer group name ConsumerGroup string // Topic is the topic being consumed @@ -404,23 +455,23 @@ type LagInfo struct { // ConsumerLag returns per-partition lag for each consumer group on a topic. // Lag = max message offset in partition - consumer group's acked offset. -func (s *AdminStore) ConsumerLag(ctx context.Context, topic string) ([]LagInfo, error) { +func (s *AdminStore) ConsumerLag(ctx context.Context, tenant string, topic string) ([]LagInfo, error) { query := fmt.Sprintf(` - SELECT o.consumer_group, o.topic, o.partition_key, o.offset_acked, + SELECT o.tenant, o.consumer_group, o.topic, o.partition_key, o.offset_acked, COALESCE(m.latest_offset, 0) AS latest_offset FROM %s o LEFT JOIN ( - SELECT topic, partition_key, MAX(`+"`offset`"+`) AS latest_offset + SELECT tenant, topic, partition_key, MAX(`+"`offset`"+`) AS latest_offset FROM %s - WHERE topic = ? - GROUP BY topic, partition_key - ) m ON o.topic = m.topic AND o.partition_key = m.partition_key - WHERE o.topic = ? + WHERE tenant = ? AND topic = ? + GROUP BY tenant, topic, partition_key + ) m ON o.tenant = m.tenant AND o.topic = m.topic AND o.partition_key = m.partition_key + WHERE o.tenant = ? AND o.topic = ? ORDER BY o.consumer_group, o.partition_key`, mysql.OffsetsTableName, mysql.MessagesTableName, ) - rows, err := s.db.QueryContext(ctx, query, topic, topic) + rows, err := s.db.QueryContext(ctx, query, tenant, topic, tenant, topic) if err != nil { return nil, fmt.Errorf("consumer lag: %w", err) } @@ -429,7 +480,7 @@ func (s *AdminStore) ConsumerLag(ctx context.Context, topic string) ([]LagInfo, var results []LagInfo for rows.Next() { var l LagInfo - if err := rows.Scan(&l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.AckedOffset, &l.LatestOffset); err != nil { + if err := rows.Scan(&l.Tenant, &l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.AckedOffset, &l.LatestOffset); err != nil { return nil, fmt.Errorf("scan lag row: %w", err) } l.Lag = l.LatestOffset - l.AckedOffset @@ -444,12 +495,18 @@ func (s *AdminStore) ConsumerLag(ctx context.Context, topic string) ([]LagInfo, // StaleLeases returns leases whose lease_renewed_at is older than the threshold. // thresholdMs is the staleness threshold in milliseconds — leases not renewed // within this duration from now are considered stale. -func (s *AdminStore) StaleLeases(ctx context.Context, thresholdMs int64) ([]LeaseInfo, error) { +func (s *AdminStore) StaleLeases(ctx context.Context, scope TenantScope, thresholdMs int64) ([]LeaseInfo, error) { + if err := scope.validate(); err != nil { + return nil, err + } cutoff := time.Now().UnixMilli() - thresholdMs - rows, err := s.db.QueryContext(ctx, - fmt.Sprintf("SELECT consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s WHERE lease_renewed_at < ? ORDER BY lease_renewed_at", mysql.PartitionLeasesTableName), - cutoff, - ) + query := fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s WHERE lease_renewed_at < ? ORDER BY lease_renewed_at", mysql.PartitionLeasesTableName) + args := []any{cutoff} + if !scope.AllTenants { + query = fmt.Sprintf("SELECT tenant, consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at FROM %s WHERE tenant = ? AND lease_renewed_at < ? ORDER BY lease_renewed_at", mysql.PartitionLeasesTableName) + args = []any{scope.Tenant, cutoff} + } + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("stale leases: %w", err) } @@ -458,7 +515,7 @@ func (s *AdminStore) StaleLeases(ctx context.Context, thresholdMs int64) ([]Leas var leases []LeaseInfo for rows.Next() { var l LeaseInfo - if err := rows.Scan(&l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.LeasedBy, &l.LeasedAt, &l.LeaseRenewedAt); err != nil { + if err := rows.Scan(&l.Tenant, &l.ConsumerGroup, &l.Topic, &l.PartitionKey, &l.LeasedBy, &l.LeasedAt, &l.LeaseRenewedAt); err != nil { return nil, fmt.Errorf("scan stale lease row: %w", err) } leases = append(leases, l) @@ -467,10 +524,10 @@ func (s *AdminStore) StaleLeases(ctx context.Context, thresholdMs int64) ([]Leas } // ReleaseLease force-releases a partition lease. -func (s *AdminStore) ReleaseLease(ctx context.Context, consumerGroup, topic, partition string) (int64, error) { +func (s *AdminStore) ReleaseLease(ctx context.Context, tenant, consumerGroup, topic, partition string) (int64, error) { result, err := s.db.ExecContext(ctx, - fmt.Sprintf("DELETE FROM %s WHERE consumer_group = ? AND topic = ? AND partition_key = ?", mysql.PartitionLeasesTableName), - consumerGroup, topic, partition, + fmt.Sprintf("DELETE FROM %s WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ?", mysql.PartitionLeasesTableName), + tenant, consumerGroup, topic, partition, ) if err != nil { return 0, fmt.Errorf("release lease: %w", err) diff --git a/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go b/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go index 46e07c2b5..c51755a42 100644 --- a/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go +++ b/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go @@ -30,15 +30,17 @@ func TestListTopics(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"topic", "count"}). - AddRow("orders", 10). - AddRow("payments", 5) - mock.ExpectQuery("SELECT topic, COUNT\\(\\*\\) FROM queue_messages GROUP BY topic ORDER BY topic"). + rows := sqlmock.NewRows([]string{"tenant", "topic", "count"}). + AddRow("acme", "orders", 10). + AddRow("acme", "payments", 5) + mock.ExpectQuery("SELECT tenant, topic, COUNT\\(\\*\\) FROM queue_messages WHERE tenant = \\? GROUP BY tenant, topic ORDER BY tenant, topic"). + WithArgs("acme"). WillReturnRows(rows) - topics, err := store.ListTopics(context.Background()) + topics, err := store.ListTopics(context.Background(), TenantScope{Tenant: "acme"}) require.NoError(t, err) assert.Len(t, topics, 2) + assert.Equal(t, "acme", topics[0].Tenant) assert.Equal(t, "orders", topics[0].Topic) assert.Equal(t, int64(10), topics[0].MessageCount) assert.Equal(t, "payments", topics[1].Topic) @@ -53,11 +55,11 @@ func TestListTopicsEmpty(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"topic", "count"}) - mock.ExpectQuery("SELECT topic, COUNT\\(\\*\\) FROM queue_messages GROUP BY topic ORDER BY topic"). + rows := sqlmock.NewRows([]string{"tenant", "topic", "count"}) + mock.ExpectQuery("SELECT tenant, topic, COUNT\\(\\*\\) FROM queue_messages GROUP BY tenant, topic ORDER BY tenant, topic"). WillReturnRows(rows) - topics, err := store.ListTopics(context.Background()) + topics, err := store.ListTopics(context.Background(), TenantScope{AllTenants: true}) require.NoError(t, err) assert.Empty(t, topics) assert.NoError(t, mock.ExpectationsWereMet()) @@ -71,27 +73,28 @@ func TestGetTopicStats(t *testing.T) { store := NewAdminStore(db) // Total messages - mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM queue_messages WHERE topic = \\?"). - WithArgs("orders"). + mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM queue_messages WHERE tenant = \\? AND topic = \\?"). + WithArgs("acme", "orders"). WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(100)) // DLQ count - mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM queue_messages WHERE topic = \\?"). - WithArgs("orders_dlq"). + mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM queue_messages WHERE tenant = \\? AND topic = \\?"). + WithArgs("acme", "orders_dlq"). WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(3)) // Distinct partitions - mock.ExpectQuery("SELECT COUNT\\(DISTINCT partition_key\\) FROM queue_messages WHERE topic = \\?"). - WithArgs("orders"). + mock.ExpectQuery("SELECT COUNT\\(DISTINCT partition_key\\) FROM queue_messages WHERE tenant = \\? AND topic = \\?"). + WithArgs("acme", "orders"). WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(4)) // Consumer groups - mock.ExpectQuery("SELECT COUNT\\(DISTINCT consumer_group\\) FROM queue_offsets WHERE topic = \\?"). - WithArgs("orders"). + mock.ExpectQuery("SELECT COUNT\\(DISTINCT consumer_group\\) FROM queue_offsets WHERE tenant = \\? AND topic = \\?"). + WithArgs("acme", "orders"). WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(2)) - stats, err := store.GetTopicStats(context.Background(), "orders", "_dlq") + stats, err := store.GetTopicStats(context.Background(), "acme", "orders", "_dlq") require.NoError(t, err) + assert.Equal(t, "acme", stats.Tenant) assert.Equal(t, "orders", stats.Topic) assert.Equal(t, int64(100), stats.TotalMessages) assert.Equal(t, int64(3), stats.DLQCount) @@ -107,17 +110,18 @@ func TestListMessages(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"offset", "id", "topic", "partition_key", "created_at", "published_at"}). - AddRow(1, "msg-1", "orders", "repo-1", 1000, 1000). - AddRow(2, "msg-2", "orders", "repo-1", 2000, 2000) - mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE topic = \\? ORDER BY `offset` LIMIT \\?"). - WithArgs("orders", 50). + rows := sqlmock.NewRows([]string{"tenant", "offset", "id", "topic", "partition_key", "created_at", "published_at"}). + AddRow("acme", 1, "msg-1", "orders", "repo-1", 1000, 1000). + AddRow("acme", 2, "msg-2", "orders", "repo-1", 2000, 2000) + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? ORDER BY `offset` LIMIT \\?"). + WithArgs("acme", "orders", 50). WillReturnRows(rows) - messages, err := store.ListMessages(context.Background(), "orders", "", 50) + messages, err := store.ListMessages(context.Background(), "acme", "orders", "", 50) require.NoError(t, err) assert.Len(t, messages, 2) assert.Equal(t, "msg-1", messages[0].ID) + assert.Equal(t, "acme", messages[0].Tenant) assert.Equal(t, int64(1), messages[0].Offset) assert.Equal(t, "msg-2", messages[1].ID) assert.NoError(t, mock.ExpectationsWereMet()) @@ -130,13 +134,13 @@ func TestListMessagesWithPartition(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"offset", "id", "topic", "partition_key", "created_at", "published_at"}). - AddRow(1, "msg-1", "orders", "repo-1", 1000, 1000) - mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE topic = \\? AND partition_key = \\? ORDER BY `offset` LIMIT \\?"). - WithArgs("orders", "repo-1", 10). + rows := sqlmock.NewRows([]string{"tenant", "offset", "id", "topic", "partition_key", "created_at", "published_at"}). + AddRow("acme", 1, "msg-1", "orders", "repo-1", 1000, 1000) + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? ORDER BY `offset` LIMIT \\?"). + WithArgs("acme", "orders", "repo-1", 10). WillReturnRows(rows) - messages, err := store.ListMessages(context.Background(), "orders", "repo-1", 10) + messages, err := store.ListMessages(context.Background(), "acme", "orders", "repo-1", 10) require.NoError(t, err) assert.Len(t, messages, 1) assert.Equal(t, "repo-1", messages[0].PartitionKey) @@ -150,16 +154,17 @@ func TestInspectMessage(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"offset", "id", "topic", "partition_key", "created_at", "published_at", "payload", "metadata", "failed_at", "failure_count", "last_error", "original_topic"}). - AddRow(1, "msg-1", "orders", "repo-1", 1000, 1000, []byte("hello"), []byte(`{"key":"val"}`), 0, 0, "", "") - mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE topic = \\? AND id = \\?"). - WithArgs("orders", "msg-1"). + rows := sqlmock.NewRows([]string{"tenant", "offset", "id", "topic", "partition_key", "created_at", "published_at", "payload", "metadata", "failed_at", "failure_count", "last_error", "original_topic"}). + AddRow("acme", 1, "msg-1", "orders", "repo-1", 1000, 1000, []byte("hello"), []byte(`{"key":"val"}`), 0, 0, "", "") + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? AND id = \\?"). + WithArgs("acme", "orders", "repo-1", "msg-1"). WillReturnRows(rows) - detail, found, err := store.InspectMessage(context.Background(), "orders", "msg-1") + detail, found, err := store.InspectMessage(context.Background(), "acme", "orders", "repo-1", "msg-1") require.NoError(t, err) assert.True(t, found) assert.Equal(t, "msg-1", detail.ID) + assert.Equal(t, "acme", detail.Tenant) assert.Equal(t, []byte("hello"), detail.Payload) assert.Equal(t, "val", detail.Metadata["key"]) assert.Equal(t, int64(0), detail.FailedAt) @@ -173,12 +178,12 @@ func TestInspectMessageNotFound(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"offset", "id", "topic", "partition_key", "created_at", "published_at", "payload", "metadata", "failed_at", "failure_count", "last_error", "original_topic"}) - mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE topic = \\? AND id = \\?"). - WithArgs("orders", "missing"). + rows := sqlmock.NewRows([]string{"tenant", "offset", "id", "topic", "partition_key", "created_at", "published_at", "payload", "metadata", "failed_at", "failure_count", "last_error", "original_topic"}) + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? AND id = \\?"). + WithArgs("acme", "orders", "repo-1", "missing"). WillReturnRows(rows) - _, found, err := store.InspectMessage(context.Background(), "orders", "missing") + _, found, err := store.InspectMessage(context.Background(), "acme", "orders", "repo-1", "missing") require.NoError(t, err) assert.False(t, found) assert.NoError(t, mock.ExpectationsWereMet()) @@ -191,11 +196,11 @@ func TestDeleteMessage(t *testing.T) { store := NewAdminStore(db) - mock.ExpectExec("DELETE FROM queue_messages WHERE topic = \\? AND id = \\?"). - WithArgs("orders", "msg-1"). + mock.ExpectExec("DELETE FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? AND id = \\?"). + WithArgs("acme", "orders", "repo-1", "msg-1"). WillReturnResult(sqlmock.NewResult(0, 1)) - affected, err := store.DeleteMessage(context.Background(), "orders", "msg-1") + affected, err := store.DeleteMessage(context.Background(), "acme", "orders", "repo-1", "msg-1") require.NoError(t, err) assert.Equal(t, int64(1), affected) assert.NoError(t, mock.ExpectationsWereMet()) @@ -208,11 +213,11 @@ func TestPurgeTopic(t *testing.T) { store := NewAdminStore(db) - mock.ExpectExec("DELETE FROM queue_messages WHERE topic = \\?"). - WithArgs("orders"). + mock.ExpectExec("DELETE FROM queue_messages WHERE tenant = \\? AND topic = \\?"). + WithArgs("acme", "orders"). WillReturnResult(sqlmock.NewResult(0, 42)) - affected, err := store.PurgeTopic(context.Background(), "orders") + affected, err := store.PurgeTopic(context.Background(), "acme", "orders") require.NoError(t, err) assert.Equal(t, int64(42), affected) assert.NoError(t, mock.ExpectationsWereMet()) @@ -226,19 +231,19 @@ func TestRequeueDLQ(t *testing.T) { store := NewAdminStore(db) mock.ExpectBegin() - mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE topic = \\? AND id = \\?"). - WithArgs("orders_dlq", "msg-1"). - WillReturnRows(sqlmock.NewRows([]string{"payload", "metadata", "partition_key", "created_at", "published_at"}). - AddRow([]byte("data"), []byte(`{}`), "repo-1", 1000, 1000)) + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? AND id = \\?"). + WithArgs("acme", "orders_dlq", "repo-1", "msg-1"). + WillReturnRows(sqlmock.NewRows([]string{"payload", "metadata", "created_at", "published_at"}). + AddRow([]byte("data"), []byte(`{}`), 1000, 1000)) mock.ExpectExec("INSERT INTO queue_messages"). - WithArgs("orders", "repo-1", "msg-1", []byte("data"), []byte(`{}`), int64(1000), sqlmock.AnyArg()). + WithArgs("acme", "orders", "repo-1", "msg-1", []byte("data"), []byte(`{}`), int64(1000), sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectExec("DELETE FROM queue_messages WHERE topic = \\? AND id = \\?"). - WithArgs("orders_dlq", "msg-1"). + mock.ExpectExec("DELETE FROM queue_messages WHERE tenant = \\? AND topic = \\? AND partition_key = \\? AND id = \\?"). + WithArgs("acme", "orders_dlq", "repo-1", "msg-1"). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() - err = store.RequeueDLQ(context.Background(), "orders", "msg-1", "_dlq") + err = store.RequeueDLQ(context.Background(), "acme", "orders", "repo-1", "msg-1", "_dlq") require.NoError(t, err) assert.NoError(t, mock.ExpectationsWereMet()) } @@ -250,14 +255,15 @@ func TestListOffsets(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "offset_acked", "updated_at"}). - AddRow("group-1", "orders", "repo-1", 100, 5000) + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "offset_acked", "updated_at"}). + AddRow("acme", "group-1", "orders", "repo-1", 100, 5000) mock.ExpectQuery("SELECT .+ FROM queue_offsets ORDER BY"). WillReturnRows(rows) - offsets, err := store.ListOffsets(context.Background(), "") + offsets, err := store.ListOffsets(context.Background(), TenantScope{AllTenants: true}, "") require.NoError(t, err) assert.Len(t, offsets, 1) + assert.Equal(t, "acme", offsets[0].Tenant) assert.Equal(t, "group-1", offsets[0].ConsumerGroup) assert.Equal(t, int64(100), offsets[0].OffsetAcked) assert.NoError(t, mock.ExpectationsWereMet()) @@ -270,13 +276,13 @@ func TestListOffsetsFiltered(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "offset_acked", "updated_at"}). - AddRow("group-1", "orders", "repo-1", 100, 5000) - mock.ExpectQuery("SELECT .+ FROM queue_offsets WHERE consumer_group = \\?"). - WithArgs("group-1"). + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "offset_acked", "updated_at"}). + AddRow("acme", "group-1", "orders", "repo-1", 100, 5000) + mock.ExpectQuery("SELECT .+ FROM queue_offsets WHERE tenant = \\? AND consumer_group = \\?"). + WithArgs("acme", "group-1"). WillReturnRows(rows) - offsets, err := store.ListOffsets(context.Background(), "group-1") + offsets, err := store.ListOffsets(context.Background(), TenantScope{Tenant: "acme"}, "group-1") require.NoError(t, err) assert.Len(t, offsets, 1) assert.NoError(t, mock.ExpectationsWereMet()) @@ -290,10 +296,10 @@ func TestResetOffset(t *testing.T) { store := NewAdminStore(db) mock.ExpectExec("UPDATE queue_offsets SET offset_acked = \\?, updated_at = \\?"). - WithArgs(int64(0), sqlmock.AnyArg(), "group-1", "orders", "repo-1"). + WithArgs(int64(0), sqlmock.AnyArg(), "acme", "orders", "repo-1", "group-1"). WillReturnResult(sqlmock.NewResult(0, 1)) - affected, err := store.ResetOffset(context.Background(), "group-1", "orders", "repo-1", 0) + affected, err := store.ResetOffset(context.Background(), "acme", "group-1", "orders", "repo-1", 0) require.NoError(t, err) assert.Equal(t, int64(1), affected) assert.NoError(t, mock.ExpectationsWereMet()) @@ -306,14 +312,16 @@ func TestListLeases(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}). - AddRow("group-1", "orders", "repo-1", "worker-1", 1000, 2000) - mock.ExpectQuery("SELECT .+ FROM queue_partition_leases ORDER BY"). + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}). + AddRow("acme", "group-1", "orders", "repo-1", "worker-1", 1000, 2000) + mock.ExpectQuery("SELECT .+ FROM queue_partition_leases WHERE tenant = \\? ORDER BY"). + WithArgs("acme"). WillReturnRows(rows) - leases, err := store.ListLeases(context.Background()) + leases, err := store.ListLeases(context.Background(), TenantScope{Tenant: "acme"}) require.NoError(t, err) assert.Len(t, leases, 1) + assert.Equal(t, "acme", leases[0].Tenant) assert.Equal(t, "worker-1", leases[0].LeasedBy) assert.Equal(t, int64(1000), leases[0].LeasedAt) assert.NoError(t, mock.ExpectationsWereMet()) @@ -326,11 +334,11 @@ func TestReleaseLease(t *testing.T) { store := NewAdminStore(db) - mock.ExpectExec("DELETE FROM queue_partition_leases WHERE consumer_group = \\? AND topic = \\? AND partition_key = \\?"). - WithArgs("group-1", "orders", "repo-1"). + mock.ExpectExec("DELETE FROM queue_partition_leases WHERE tenant = \\? AND consumer_group = \\? AND topic = \\? AND partition_key = \\?"). + WithArgs("acme", "group-1", "orders", "repo-1"). WillReturnResult(sqlmock.NewResult(0, 1)) - affected, err := store.ReleaseLease(context.Background(), "group-1", "orders", "repo-1") + affected, err := store.ReleaseLease(context.Background(), "acme", "group-1", "orders", "repo-1") require.NoError(t, err) assert.Equal(t, int64(1), affected) assert.NoError(t, mock.ExpectationsWereMet()) @@ -349,16 +357,17 @@ func TestConsumerLag(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "offset_acked", "latest_offset"}). - AddRow("group-1", "orders", "repo-1", 50, 100). - AddRow("group-1", "orders", "repo-2", 75, 75) + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "offset_acked", "latest_offset"}). + AddRow("acme", "group-1", "orders", "repo-1", 50, 100). + AddRow("acme", "group-1", "orders", "repo-2", 75, 75) mock.ExpectQuery("SELECT .+ FROM queue_offsets .+ LEFT JOIN"). - WithArgs("orders", "orders"). + WithArgs("acme", "orders", "acme", "orders"). WillReturnRows(rows) - lags, err := store.ConsumerLag(context.Background(), "orders") + lags, err := store.ConsumerLag(context.Background(), "acme", "orders") require.NoError(t, err) assert.Len(t, lags, 2) + assert.Equal(t, "acme", lags[0].Tenant) assert.Equal(t, int64(50), lags[0].Lag) assert.Equal(t, int64(100), lags[0].LatestOffset) assert.Equal(t, int64(50), lags[0].AckedOffset) @@ -374,13 +383,13 @@ func TestConsumerLagNoMessages(t *testing.T) { store := NewAdminStore(db) // Consumer has offset but no messages remain (all acked and deleted) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "offset_acked", "latest_offset"}). - AddRow("group-1", "orders", "repo-1", 100, 0) + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "offset_acked", "latest_offset"}). + AddRow("acme", "group-1", "orders", "repo-1", 100, 0) mock.ExpectQuery("SELECT .+ FROM queue_offsets .+ LEFT JOIN"). - WithArgs("orders", "orders"). + WithArgs("acme", "orders", "acme", "orders"). WillReturnRows(rows) - lags, err := store.ConsumerLag(context.Background(), "orders") + lags, err := store.ConsumerLag(context.Background(), "acme", "orders") require.NoError(t, err) assert.Len(t, lags, 1) assert.Equal(t, int64(0), lags[0].Lag) // clamped to 0, not negative @@ -394,15 +403,16 @@ func TestStaleLeases(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}). - AddRow("group-1", "orders", "repo-1", "worker-1", 1000, 2000) - mock.ExpectQuery("SELECT .+ FROM queue_partition_leases WHERE lease_renewed_at < \\?"). - WithArgs(sqlmock.AnyArg()). + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}). + AddRow("acme", "group-1", "orders", "repo-1", "worker-1", 1000, 2000) + mock.ExpectQuery("SELECT .+ FROM queue_partition_leases WHERE tenant = \\? AND lease_renewed_at < \\?"). + WithArgs("acme", sqlmock.AnyArg()). WillReturnRows(rows) - leases, err := store.StaleLeases(context.Background(), 60000) + leases, err := store.StaleLeases(context.Background(), TenantScope{Tenant: "acme"}, 60000) require.NoError(t, err) assert.Len(t, leases, 1) + assert.Equal(t, "acme", leases[0].Tenant) assert.Equal(t, "worker-1", leases[0].LeasedBy) assert.NoError(t, mock.ExpectationsWereMet()) } @@ -414,13 +424,63 @@ func TestStaleLeasesEmpty(t *testing.T) { store := NewAdminStore(db) - rows := sqlmock.NewRows([]string{"consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}) + rows := sqlmock.NewRows([]string{"tenant", "consumer_group", "topic", "partition_key", "leased_by", "leased_at", "lease_renewed_at"}) mock.ExpectQuery("SELECT .+ FROM queue_partition_leases WHERE lease_renewed_at < \\?"). WithArgs(sqlmock.AnyArg()). WillReturnRows(rows) - leases, err := store.StaleLeases(context.Background(), 60000) + leases, err := store.StaleLeases(context.Background(), TenantScope{AllTenants: true}, 60000) require.NoError(t, err) assert.Empty(t, leases) assert.NoError(t, mock.ExpectationsWereMet()) } + +func TestTenantScopeRequired(t *testing.T) { + db, _, err := sqlmock.New() + require.NoError(t, err) + defer db.Close() + + store := NewAdminStore(db) + tests := []struct { + name string + scope TenantScope + call func(TenantScope) error + }{ + { + name: "list topics without scope", + call: func(scope TenantScope) error { + _, err := store.ListTopics(context.Background(), scope) + return err + }, + }, + { + name: "list offsets with conflicting scope", + scope: TenantScope{Tenant: "acme", AllTenants: true}, + call: func(scope TenantScope) error { + _, err := store.ListOffsets(context.Background(), scope, "") + return err + }, + }, + { + name: "list leases without scope", + call: func(scope TenantScope) error { + _, err := store.ListLeases(context.Background(), scope) + return err + }, + }, + { + name: "stale leases with conflicting scope", + scope: TenantScope{Tenant: "acme", AllTenants: true}, + call: func(scope TenantScope) error { + _, err := store.StaleLeases(context.Background(), scope, 60000) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Error(t, tt.call(tt.scope)) + }) + } +} diff --git a/platform/extension/messagequeue/mysql/delivery_state_store.go b/platform/extension/messagequeue/mysql/delivery_state_store.go index 2794ce91e..8bf17facd 100644 --- a/platform/extension/messagequeue/mysql/delivery_state_store.go +++ b/platform/extension/messagequeue/mysql/delivery_state_store.go @@ -45,10 +45,10 @@ func newDeliveryStateStore(db *sql.DB, logger *zap.SugaredLogger, scope tally.Sc // Returns the resulting retry_count after the operation. // // The INSERT and subsequent SELECT are not in a transaction. This is safe because -// partition leasing guarantees a single writer per (consumer_group, topic, partition_key) +// partition leasing guarantees a single writer per (tenant, consumer_group, topic, partition_key) // — only the lease holder calls MarkDelivered for a given partition, so no concurrent // mutation can occur between the two statements. -func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (_ int, retErr error) { +func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (_ int, retErr error) { op := metrics.Begin(s.scope, "mark_delivered", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -62,17 +62,17 @@ func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup // postponed reset must come after it. A postponed redelivery is a deliberate // wait, not a failure — it is exempt from the increment and consumes the flag. _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count, postponed) - VALUES (?, ?, ?, ?, FALSE, ?, 0, FALSE) + INSERT INTO %s (tenant, consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count, postponed) + VALUES (?, ?, ?, ?, ?, FALSE, ?, 0, FALSE) ON DUPLICATE KEY UPDATE invisible_until = IF(acked = FALSE, VALUES(invisible_until), invisible_until), retry_count = IF(acked = FALSE AND postponed = FALSE, retry_count + 1, retry_count), postponed = IF(acked = FALSE, FALSE, postponed) `, DeliveryStateTableName), - consumerGroup, topic, partitionKey, offset, invisibleUntil) + tenant, consumerGroup, topic, partitionKey, offset, invisibleUntil) if err != nil { - return 0, fmt.Errorf("mark delivered topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return 0, fmt.Errorf("mark delivered tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } // Read retry_count after INSERT/UPDATE to get the current value. @@ -82,10 +82,10 @@ func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup var retryCount int err = s.db.QueryRowContext(ctx, fmt.Sprintf(` SELECT retry_count FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? - `, DeliveryStateTableName), consumerGroup, topic, partitionKey, offset).Scan(&retryCount) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? + `, DeliveryStateTableName), tenant, consumerGroup, topic, partitionKey, offset).Scan(&retryCount) if err != nil { - return 0, fmt.Errorf("get retry count after mark delivered topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return 0, fmt.Errorf("get retry count after mark delivered tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } return retryCount, nil @@ -93,7 +93,7 @@ func (s *sqldeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup // ExtendVisibility extends the visibility timeout for an in-flight message // without incrementing retry_count. Used by ExtendVisibilityTimeout. -func (s *sqldeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retErr error) { +func (s *sqldeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retErr error) { op := metrics.Begin(s.scope, "extend_visibility", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -105,17 +105,18 @@ func (s *sqldeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGr result, err := s.db.ExecContext(ctx, fmt.Sprintf(` UPDATE %s SET invisible_until = ? - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? AND acked = FALSE + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? AND acked = FALSE `, DeliveryStateTableName), - invisibleUntil, consumerGroup, topic, partitionKey, offset) + invisibleUntil, tenant, consumerGroup, topic, partitionKey, offset) if err != nil { - return fmt.Errorf("extend visibility topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return fmt.Errorf("extend visibility tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } rowsAffected, raErr := result.RowsAffected() if raErr == nil && rowsAffected == 0 { s.logger.Warnw("extend visibility matched no rows, lease may have expired or message already acked", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, "offset", offset, @@ -126,21 +127,21 @@ func (s *sqldeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGr } // MarkAcked sets acked = TRUE to indicate this group has processed the message. -func (s *sqldeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (retErr error) { +func (s *sqldeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) (retErr error) { op := metrics.Begin(s.scope, "mark_acked", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) defer func() { op.Complete(retErr) }() _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count) - VALUES (?, ?, ?, ?, TRUE, 0, 0) + INSERT INTO %s (tenant, consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count) + VALUES (?, ?, ?, ?, ?, TRUE, 0, 0) ON DUPLICATE KEY UPDATE acked = TRUE `, DeliveryStateTableName), - consumerGroup, topic, partitionKey, offset) + tenant, consumerGroup, topic, partitionKey, offset) if err != nil { - return fmt.Errorf("mark acked topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return fmt.Errorf("mark acked tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } return nil @@ -148,27 +149,27 @@ func (s *sqldeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, to // MarkNacked makes the message eligible for redelivery after delayMs. // retry_count is NOT incremented here — it is incremented by MarkDelivered on redelivery. -func (s *sqldeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, delayMs int64) (retErr error) { +func (s *sqldeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, delayMs int64) (retErr error) { op := metrics.Begin(s.scope, "mark_nacked", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) defer func() { op.Complete(retErr) }() if delayMs < 0 || delayMs > maxRetryBackoffMs { - return fmt.Errorf("mark nacked topic=%s partition=%s offset=%d: retry delay %d is outside [0, %d]", topic, partitionKey, offset, delayMs, maxRetryBackoffMs) + return fmt.Errorf("mark nacked tenant=%s topic=%s partition=%s offset=%d: retry delay %d is outside [0, %d]", tenant, topic, partitionKey, offset, delayMs, maxRetryBackoffMs) } invisibleUntil := time.Now().UnixMilli() + delayMs _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count) - VALUES (?, ?, ?, ?, FALSE, ?, 0) + INSERT INTO %s (tenant, consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count) + VALUES (?, ?, ?, ?, ?, FALSE, ?, 0) ON DUPLICATE KEY UPDATE invisible_until = IF(acked = FALSE, VALUES(invisible_until), invisible_until) `, DeliveryStateTableName), - consumerGroup, topic, partitionKey, offset, invisibleUntil) + tenant, consumerGroup, topic, partitionKey, offset, invisibleUntil) if err != nil { - return fmt.Errorf("mark nacked topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return fmt.Errorf("mark nacked tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } return nil @@ -180,7 +181,7 @@ func (s *sqldeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, t // MarkDelivered from the retry_count increment. The reset restarts failure // accounting — a completed delivery that chose to wait has demonstrated the // message is processable. -func (s *sqldeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, delayMs int64) (retErr error) { +func (s *sqldeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, delayMs int64) (retErr error) { op := metrics.Begin(s.scope, "mark_postponed", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -190,17 +191,17 @@ func (s *sqldeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup invisibleUntil := now + delayMs _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count, postponed) - VALUES (?, ?, ?, ?, FALSE, ?, 0, TRUE) + INSERT INTO %s (tenant, consumer_group, topic, partition_key, message_offset, acked, invisible_until, retry_count, postponed) + VALUES (?, ?, ?, ?, ?, FALSE, ?, 0, TRUE) ON DUPLICATE KEY UPDATE invisible_until = IF(acked = FALSE, VALUES(invisible_until), invisible_until), retry_count = IF(acked = FALSE, 0, retry_count), postponed = IF(acked = FALSE, TRUE, postponed) `, DeliveryStateTableName), - consumerGroup, topic, partitionKey, offset, invisibleUntil) + tenant, consumerGroup, topic, partitionKey, offset, invisibleUntil) if err != nil { - return fmt.Errorf("mark postponed topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return fmt.Errorf("mark postponed tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } return nil @@ -208,7 +209,7 @@ func (s *sqldeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup // GetDeliveryState returns the full delivery state for a message offset. // Returns (state, found, error). found=false means no row (never delivered). -func (s *sqldeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (_ DeliveryState, _ bool, retErr error) { +func (s *sqldeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) (_ DeliveryState, _ bool, retErr error) { op := metrics.Begin(s.scope, "get_delivery_state", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -217,14 +218,14 @@ func (s *sqldeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGr var state DeliveryState err := s.db.QueryRowContext(ctx, fmt.Sprintf(` SELECT acked, invisible_until, retry_count, postponed FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? - `, DeliveryStateTableName), consumerGroup, topic, partitionKey, offset).Scan(&state.Acked, &state.InvisibleUntil, &state.RetryCount, &state.Postponed) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset = ? + `, DeliveryStateTableName), tenant, consumerGroup, topic, partitionKey, offset).Scan(&state.Acked, &state.InvisibleUntil, &state.RetryCount, &state.Postponed) if err == sql.ErrNoRows { return DeliveryState{}, false, nil } if err != nil { - return DeliveryState{}, false, fmt.Errorf("get delivery state topic=%s partition=%s offset=%d: %w", topic, partitionKey, offset, err) + return DeliveryState{}, false, fmt.Errorf("get delivery state tenant=%s topic=%s partition=%s offset=%d: %w", tenant, topic, partitionKey, offset, err) } return state, true, nil @@ -234,7 +235,7 @@ func (s *sqldeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGr // delivery state rows that are behind it. // offsets are the actual message offsets above the current watermark (from messageStore). // Returns the new watermark (highest contiguous acked offset from currentWatermark). -func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGroup, topic, partitionKey string, currentWatermark int64, offsets []int64) (_ int64, retErr error) { +func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, currentWatermark int64, offsets []int64) (_ int64, retErr error) { op := metrics.Begin(s.scope, "advance_watermark", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -246,8 +247,8 @@ func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGr // Batch-fetch delivery state for the provided offsets. placeholders := make([]byte, 0, len(offsets)*2-1) - args := make([]interface{}, 0, 3+len(offsets)) - args = append(args, consumerGroup, topic, partitionKey) + args := make([]interface{}, 0, 4+len(offsets)) + args = append(args, tenant, consumerGroup, topic, partitionKey) for i, offset := range offsets { if i > 0 { placeholders = append(placeholders, ',') @@ -258,11 +259,11 @@ func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGr rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` SELECT message_offset, acked FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset IN (%s) `, DeliveryStateTableName, string(placeholders)), args...) if err != nil { - return currentWatermark, fmt.Errorf("query delivery state for watermark topic=%s partition=%s: %w", topic, partitionKey, err) + return currentWatermark, fmt.Errorf("query delivery state for watermark tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } defer rows.Close() @@ -272,12 +273,12 @@ func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGr var offset int64 var acked bool if err := rows.Scan(&offset, &acked); err != nil { - return currentWatermark, fmt.Errorf("scan delivery state topic=%s partition=%s: %w", topic, partitionKey, err) + return currentWatermark, fmt.Errorf("scan delivery state tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } ackedMap[offset] = acked } if err := rows.Err(); err != nil { - return currentWatermark, fmt.Errorf("delivery state iteration topic=%s partition=%s: %w", topic, partitionKey, err) + return currentWatermark, fmt.Errorf("delivery state iteration tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } // Walk message offsets in order. Advance while contiguous acked. @@ -299,12 +300,13 @@ func (s *sqldeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGr if newWatermark > currentWatermark { _, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset <= ? - `, DeliveryStateTableName), consumerGroup, topic, partitionKey, newWatermark) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND message_offset <= ? + `, DeliveryStateTableName), tenant, consumerGroup, topic, partitionKey, newWatermark) if err != nil { metrics.NamedCounter(s.scope, "advance_watermark", "cleanup_errors", 1, metrics.NewTag("topic", topic)) s.logger.Warnw("failed to clean up delivery state behind watermark, will retry on next advance", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, "watermark", newWatermark, diff --git a/platform/extension/messagequeue/mysql/delivery_state_store_test.go b/platform/extension/messagequeue/mysql/delivery_state_store_test.go index 919294821..417948b55 100644 --- a/platform/extension/messagequeue/mysql/delivery_state_store_test.go +++ b/platform/extension/messagequeue/mysql/delivery_state_store_test.go @@ -75,25 +75,25 @@ func TestDeliveryStateStore_MarkDelivered(t *testing.T) { if tt.execErr { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). WillReturnError(assert.AnError) } else { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) if tt.queryErr { mock.ExpectQuery("SELECT retry_count FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnError(assert.AnError) } else { mock.ExpectQuery("SELECT retry_count FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnRows(sqlmock.NewRows([]string{"retry_count"}).AddRow(tt.wantRetryCount)) } } - retryCount, err := store.MarkDelivered(context.Background(), "group-1", "orders", "part-1", 5, 30000) + retryCount, err := store.MarkDelivered(context.Background(), "group-1", testTenant, "orders", "part-1", 5, 30000) if tt.execErr || tt.queryErr { require.Error(t, err) @@ -128,15 +128,15 @@ func TestDeliveryStateStore_ExtendVisibility(t *testing.T) { if tt.wantErr { mock.ExpectExec("UPDATE queue_delivery_state"). - WithArgs(sqlmock.AnyArg(), "group-1", "orders", "part-1", int64(5)). + WithArgs(sqlmock.AnyArg(), testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnError(assert.AnError) } else { mock.ExpectExec("UPDATE queue_delivery_state"). - WithArgs(sqlmock.AnyArg(), "group-1", "orders", "part-1", int64(5)). + WithArgs(sqlmock.AnyArg(), testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnResult(sqlmock.NewResult(0, 1)) } - err := store.ExtendVisibility(context.Background(), "group-1", "orders", "part-1", 5, 60000) + err := store.ExtendVisibility(context.Background(), "group-1", testTenant, "orders", "part-1", 5, 60000) if tt.wantErr { require.Error(t, err) @@ -170,15 +170,15 @@ func TestDeliveryStateStore_MarkAcked(t *testing.T) { if tt.wantErr { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnError(assert.AnError) } else { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnResult(sqlmock.NewResult(1, 1)) } - err := store.MarkAcked(context.Background(), "group-1", "orders", "part-1", 5) + err := store.MarkAcked(context.Background(), "group-1", testTenant, "orders", "part-1", 5) if tt.wantErr { require.Error(t, err) @@ -214,15 +214,15 @@ func TestDeliveryStateStore_MarkNacked(t *testing.T) { if tt.wantErr { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), invisibleUntil). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), invisibleUntil). WillReturnError(assert.AnError) } else { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), invisibleUntil). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), invisibleUntil). WillReturnResult(sqlmock.NewResult(1, 1)) } - err := store.MarkNacked(context.Background(), "group-1", "orders", "part-1", 5, retryDelayMs) + err := store.MarkNacked(context.Background(), "group-1", testTenant, "orders", "part-1", 5, retryDelayMs) if tt.wantErr { require.Error(t, err) @@ -248,7 +248,7 @@ func TestDeliveryStateStore_MarkNackedRejectsInvalidDelay(t *testing.T) { store, db, mock := newTestDeliveryStateStoreWithMock(t) defer db.Close() - err := store.MarkNacked(context.Background(), "group-1", "orders", "part-1", 5, tt.delayMs) + err := store.MarkNacked(context.Background(), "group-1", testTenant, "orders", "part-1", 5, tt.delayMs) require.ErrorContains(t, err, "is outside") assert.NoError(t, mock.ExpectationsWereMet()) }) @@ -277,15 +277,15 @@ func TestDeliveryStateStore_MarkPostponed(t *testing.T) { if tt.wantErr { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). WillReturnError(assert.AnError) } else { mock.ExpectExec("INSERT INTO queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5), sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) } - err := store.MarkPostponed(context.Background(), "group-1", "orders", "part-1", 5, 5000) + err := store.MarkPostponed(context.Background(), "group-1", testTenant, "orders", "part-1", 5, 5000) if tt.wantErr { require.Error(t, err) @@ -355,20 +355,20 @@ func TestDeliveryStateStore_GetDeliveryState(t *testing.T) { if tt.wantErr { mock.ExpectQuery("SELECT acked, invisible_until, retry_count, postponed FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnError(assert.AnError) } else if tt.noRows { mock.ExpectQuery("SELECT acked, invisible_until, retry_count, postponed FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnRows(sqlmock.NewRows([]string{"acked", "invisible_until", "retry_count", "postponed"})) } else { mock.ExpectQuery("SELECT acked, invisible_until, retry_count, postponed FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", int64(5)). + WithArgs(testTenant, "group-1", "orders", "part-1", int64(5)). WillReturnRows(sqlmock.NewRows([]string{"acked", "invisible_until", "retry_count", "postponed"}). AddRow(tt.acked, tt.invisibleUntil, tt.retryCount, tt.postponed)) } - state, found, err := store.GetDeliveryState(context.Background(), "group-1", "orders", "part-1", 5) + state, found, err := store.GetDeliveryState(context.Background(), "group-1", testTenant, "orders", "part-1", 5) if tt.wantErr { require.Error(t, err) @@ -495,8 +495,8 @@ func TestDeliveryStateStore_AdvanceWatermark(t *testing.T) { // Delivery state query is only issued if there are offsets if len(tt.offsets) > 0 { - dsArgs := make([]driver.Value, 0, 3+len(tt.offsets)) - dsArgs = append(dsArgs, "group-1", "orders", "part-1") + dsArgs := make([]driver.Value, 0, 4+len(tt.offsets)) + dsArgs = append(dsArgs, testTenant, "group-1", "orders", "part-1") for _, offset := range tt.offsets { dsArgs = append(dsArgs, offset) } @@ -518,11 +518,11 @@ func TestDeliveryStateStore_AdvanceWatermark(t *testing.T) { if tt.expectCleanup { mock.ExpectExec("DELETE FROM queue_delivery_state"). - WithArgs("group-1", "orders", "part-1", tt.expectWatermark). + WithArgs(testTenant, "group-1", "orders", "part-1", tt.expectWatermark). WillReturnResult(sqlmock.NewResult(0, tt.expectWatermark-tt.currentWatermark)) } - watermark, err := store.AdvanceWatermark(context.Background(), "group-1", "orders", "part-1", tt.currentWatermark, tt.offsets) + watermark, err := store.AdvanceWatermark(context.Background(), "group-1", testTenant, "orders", "part-1", tt.currentWatermark, tt.offsets) if tt.dsQueryErr { require.Error(t, err) diff --git a/platform/extension/messagequeue/mysql/identifier.go b/platform/extension/messagequeue/mysql/identifier.go new file mode 100644 index 000000000..732fa6c49 --- /dev/null +++ b/platform/extension/messagequeue/mysql/identifier.go @@ -0,0 +1,68 @@ +// 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 mysql + +import ( + "fmt" + "strings" + "unicode/utf8" +) + +const maxIdentifierLength = 255 + +func validateASCIIIdentifier(name, value string) error { + if len(value) > maxIdentifierLength { + return fmt.Errorf("%s exceeds %d bytes", name, maxIdentifierLength) + } + for i := range len(value) { + if value[i] > 0x7f { + return fmt.Errorf("%s must contain only ASCII characters", name) + } + if value[i] == 0 { + return fmt.Errorf("%s must not contain NUL", name) + } + } + return nil +} + +func validateTextIdentifier(name, value string) error { + if !utf8.ValidString(value) { + return fmt.Errorf("%s is not valid UTF-8", name) + } + if utf8.RuneCountInString(value) > maxIdentifierLength { + return fmt.Errorf("%s exceeds %d characters", name, maxIdentifierLength) + } + return nil +} + +func normalizeTenants(tenants []string) ([]string, error) { + normalized := make([]string, 0, len(tenants)) + seen := make(map[string]struct{}, len(tenants)) + for _, tenant := range tenants { + tenant = strings.TrimSpace(tenant) + if tenant == "" { + continue + } + if err := validateASCIIIdentifier("tenant", tenant); err != nil { + return nil, err + } + if _, exists := seen[tenant]; exists { + continue + } + seen[tenant] = struct{}{} + normalized = append(normalized, tenant) + } + return normalized, nil +} diff --git a/platform/extension/messagequeue/mysql/message_store.go b/platform/extension/messagequeue/mysql/message_store.go index 78f9c9bf5..981bccd69 100644 --- a/platform/extension/messagequeue/mysql/message_store.go +++ b/platform/extension/messagequeue/mysql/message_store.go @@ -47,14 +47,14 @@ func newMessageStore(db *sql.DB, logger *zap.SugaredLogger, scope tally.Scope) m // Insert inserts messages into the messages table. // -// Publishes are idempotent on the (topic, partition_key, id) unique key: a +// Publishes are idempotent on the (tenant, topic, partition_key, id) unique key: a // repeated publish for the same key is silently treated as success and does // not overwrite the original payload. This matches the queue_messages schema's // documented intent ("Supports: INSERT ... ON DUPLICATE KEY to enforce // idempotent publishes") and lets callers safely retry publishes (e.g. a // second Cancel RPC for the same request) without surfacing 1062 duplicate-key // errors. -func (s *sqlmessageStore) Insert(ctx context.Context, topic string, messages []entityqueue.Message) (retErr error) { +func (s *sqlmessageStore) Insert(ctx context.Context, tenant string, topic string, messages []entityqueue.Message) (retErr error) { op := metrics.Begin(s.scope, "insert", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -63,13 +63,14 @@ func (s *sqlmessageStore) Insert(ctx context.Context, topic string, messages []e } s.logger.Debugw("inserting messages", + logTenant, tenant, logTopic, topic, "count", len(messages), ) tx, err := s.db.BeginTx(ctx, nil) if err != nil { - return fmt.Errorf("begin transaction topic=%s: %w", topic, err) + return fmt.Errorf("begin transaction tenant=%s topic=%s: %w", tenant, topic, err) } defer tx.Rollback() @@ -80,26 +81,43 @@ func (s *sqlmessageStore) Insert(ctx context.Context, topic string, messages []e // is NULL rather than an empty sentinel: it is a JSON column, which rejects // ''. stmt, err := tx.PrepareContext(ctx, fmt.Sprintf(` - INSERT INTO %s (topic, id, payload, metadata, partition_key, created_at, published_at, failed_at, failure_count, last_error, original_topic, failure_detail) - VALUES (?, ?, ?, ?, ?, ?, ?, 0, 0, '', '', NULL) + INSERT INTO %s (tenant, topic, id, payload, metadata, partition_key, created_at, published_at, failed_at, failure_count, last_error, original_topic, failure_detail) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, '', '', NULL) ON DUPLICATE KEY UPDATE topic = topic `, MessagesTableName)) if err != nil { - return fmt.Errorf("prepare statement topic=%s: %w", topic, err) + return fmt.Errorf("prepare statement tenant=%s topic=%s: %w", tenant, topic, err) } defer stmt.Close() now := time.Now().UnixMilli() for _, msg := range messages { + rowTenant := tenant + if rowTenant == "" { + rowTenant = msg.Tenant + } + if rowTenant == "" { + return fmt.Errorf("insert message topic=%s message=%s: tenant is required", topic, msg.ID) + } + if msg.Tenant != "" && msg.Tenant != rowTenant { + return fmt.Errorf("insert message tenant=%s topic=%s message=%s: message tenant %q does not match", rowTenant, topic, msg.ID, msg.Tenant) + } + messageWithTenant := msg.Copy() + messageWithTenant.Tenant = rowTenant + if err := entityqueue.ValidateTenantMetadata(messageWithTenant); err != nil { + return fmt.Errorf("insert message tenant=%s topic=%s message=%s: %w", rowTenant, topic, msg.ID, err) + } + var metadataJSON []byte if len(msg.Metadata) > 0 { metadataJSON, err = json.Marshal(msg.Metadata) if err != nil { - return fmt.Errorf("marshal metadata topic=%s message=%s: %w", topic, msg.ID, err) + return fmt.Errorf("marshal metadata tenant=%s topic=%s message=%s: %w", rowTenant, topic, msg.ID, err) } } _, err = stmt.ExecContext(ctx, + rowTenant, topic, msg.ID, msg.Payload, @@ -109,15 +127,16 @@ func (s *sqlmessageStore) Insert(ctx context.Context, topic string, messages []e msg.PublishedAt, ) if err != nil { - return fmt.Errorf("insert message topic=%s message=%s partition=%s: %w", topic, msg.ID, msg.PartitionKey, err) + return fmt.Errorf("insert message tenant=%s topic=%s message=%s partition=%s: %w", rowTenant, topic, msg.ID, msg.PartitionKey, err) } } if err := tx.Commit(); err != nil { - return fmt.Errorf("commit transaction topic=%s: %w", topic, err) + return fmt.Errorf("commit transaction tenant=%s topic=%s: %w", tenant, topic, err) } s.logger.Debugw("inserted messages", + logTenant, tenant, logTopic, topic, "count", len(messages), ) @@ -125,17 +144,17 @@ func (s *sqlmessageStore) Insert(ctx context.Context, topic string, messages []e return nil } -// Delete deletes a message by topic, partition key, and ID -func (s *sqlmessageStore) Delete(ctx context.Context, topic string, partitionKey string, messageID string) (retErr error) { +// Delete deletes a message by tenant, topic, partition key, and ID +func (s *sqlmessageStore) Delete(ctx context.Context, tenant string, topic string, partitionKey string, messageID string) (retErr error) { op := metrics.Begin(s.scope, "delete", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - DELETE FROM %s WHERE topic = ? AND partition_key = ? AND id = ? - `, MessagesTableName), topic, partitionKey, messageID) + DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ? + `, MessagesTableName), tenant, topic, partitionKey, messageID) if err != nil { - return fmt.Errorf("delete message topic=%s partition=%s message=%s: %w", topic, partitionKey, messageID, err) + return fmt.Errorf("delete message tenant=%s topic=%s partition=%s message=%s: %w", tenant, topic, partitionKey, messageID, err) } return nil @@ -143,19 +162,19 @@ func (s *sqlmessageStore) Delete(ctx context.Context, topic string, partitionKey // FetchByOffset fetches messages with offset > currentOffset for a specific partition. // Messages are fetched from the immutable log; no per-message mutation occurs. -func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, partitionKey string, currentOffset int64, limit int) (_ []messageRow, retErr error) { +func (s *sqlmessageStore) FetchByOffset(ctx context.Context, tenant string, topic string, partitionKey string, currentOffset int64, limit int) (_ []messageRow, retErr error) { op := metrics.Begin(s.scope, "fetch", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` - SELECT offset, id, payload, metadata, partition_key, published_at, failed_at, failure_count, last_error, original_topic, failure_detail + SELECT tenant, offset, id, payload, metadata, partition_key, published_at, failed_at, failure_count, last_error, original_topic, failure_detail FROM %s - WHERE topic = ? AND partition_key = ? AND offset > ? + WHERE tenant = ? AND topic = ? AND partition_key = ? AND offset > ? ORDER BY offset LIMIT ? - `, MessagesTableName), topic, partitionKey, currentOffset, limit) + `, MessagesTableName), tenant, topic, partitionKey, currentOffset, limit) if err != nil { - return nil, fmt.Errorf("query messages topic=%s partition=%s: %w", topic, partitionKey, err) + return nil, fmt.Errorf("query messages tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } defer rows.Close() @@ -163,6 +182,7 @@ func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, parti for rows.Next() { var ( + rowTenant string offset int64 id string payload []byte @@ -176,14 +196,14 @@ func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, parti failureDetail []byte ) - if err := rows.Scan(&offset, &id, &payload, &metadataJSON, &partKey, &publishedAtMilli, &failedAt, &failureCount, &lastError, &originalTopic, &failureDetail); err != nil { - return nil, fmt.Errorf("scan row topic=%s partition=%s: %w", topic, partitionKey, err) + if err := rows.Scan(&rowTenant, &offset, &id, &payload, &metadataJSON, &partKey, &publishedAtMilli, &failedAt, &failureCount, &lastError, &originalTopic, &failureDetail); err != nil { + return nil, fmt.Errorf("scan row tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } var metadata map[string]string if len(metadataJSON) > 0 { if err := json.Unmarshal(metadataJSON, &metadata); err != nil { - return nil, fmt.Errorf("unmarshal metadata topic=%s partition=%s message=%s: %w", topic, partitionKey, id, err) + return nil, fmt.Errorf("unmarshal metadata tenant=%s topic=%s partition=%s message=%s: %w", tenant, topic, partitionKey, id, err) } } if metadata == nil { @@ -191,6 +211,7 @@ func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, parti } results = append(results, messageRow{ + Tenant: rowTenant, Offset: offset, ID: id, Payload: payload, @@ -206,10 +227,11 @@ func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, parti } if err := rows.Err(); err != nil { - return nil, fmt.Errorf("row iteration topic=%s partition=%s: %w", topic, partitionKey, err) + return nil, fmt.Errorf("row iteration tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } s.logger.Debugw("fetched messages", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, "count", len(results), @@ -226,7 +248,7 @@ func (s *sqlmessageStore) FetchByOffset(ctx context.Context, topic string, parti // row stays readable without decoding anything, and its subjects and detail // into failure_detail. A failure with no structure leaves failure_detail NULL, // which is what an unattributed dead letter looks like. -func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, topic string, partitionKey string, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) (retErr error) { +func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, tenant string, topic string, partitionKey string, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) (retErr error) { op := metrics.Begin(s.scope, "move_to_dlq", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -235,7 +257,7 @@ func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, topic string, partition failureDetail, err := failure.Encode(f) if err != nil { - return fmt.Errorf("encode failure detail topic=%s message=%s: %w", topic, messageID, err) + return fmt.Errorf("encode failure detail tenant=%s topic=%s message=%s: %w", tenant, topic, messageID, err) } // Bind NULL explicitly when there is no structure. A nil []byte would leave // the column's value up to the driver, and an empty string is not valid @@ -247,7 +269,7 @@ func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, topic string, partition tx, err := s.db.BeginTx(ctx, nil) if err != nil { - return fmt.Errorf("begin transaction topic=%s message=%s: %w", topic, messageID, err) + return fmt.Errorf("begin transaction tenant=%s topic=%s message=%s: %w", tenant, topic, messageID, err) } defer tx.Rollback() @@ -263,43 +285,44 @@ func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, topic string, partition err = tx.QueryRowContext(ctx, fmt.Sprintf(` SELECT payload, metadata, partition_key, created_at, published_at FROM %s - WHERE topic = ? AND partition_key = ? AND id = ? - `, MessagesTableName), topic, partitionKey, messageID).Scan(&payload, &metadataJSON, &fetchPartKey, &createdAtMilli, &publishedAtMilli) + WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ? + `, MessagesTableName), tenant, topic, partitionKey, messageID).Scan(&payload, &metadataJSON, &fetchPartKey, &createdAtMilli, &publishedAtMilli) if err != nil { if err == sql.ErrNoRows { // Message already deleted or doesn't exist s.logger.Debugw("message not found for DLQ move", + logTenant, tenant, logTopic, topic, logMessageID, messageID, ) return nil } - return fmt.Errorf("fetch message for DLQ topic=%s partition=%s message=%s: %w", topic, partitionKey, messageID, err) + return fmt.Errorf("fetch message for DLQ tenant=%s topic=%s partition=%s message=%s: %w", tenant, topic, partitionKey, messageID, err) } // Insert into queue_messages table with DLQ topic name and DLQ-specific fields. now := time.Now().UnixMilli() _, err = tx.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (topic, id, payload, metadata, partition_key, created_at, published_at, failed_at, failure_count, last_error, original_topic, failure_detail) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - `, MessagesTableName), dlqTopic, messageID, payload, metadataJSON, fetchPartKey, createdAtMilli, publishedAtMilli, now, failureCount, f.Message, topic, failureDetailArg) + INSERT INTO %s (tenant, topic, id, payload, metadata, partition_key, created_at, published_at, failed_at, failure_count, last_error, original_topic, failure_detail) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, MessagesTableName), tenant, dlqTopic, messageID, payload, metadataJSON, fetchPartKey, createdAtMilli, publishedAtMilli, now, failureCount, f.Message, topic, failureDetailArg) if err != nil { - return fmt.Errorf("insert into DLQ topic=%s dlq=%s partition=%s message=%s: %w", topic, dlqTopic, partitionKey, messageID, err) + return fmt.Errorf("insert into DLQ tenant=%s topic=%s dlq=%s partition=%s message=%s: %w", tenant, topic, dlqTopic, partitionKey, messageID, err) } // Delete from original topic _, err = tx.ExecContext(ctx, fmt.Sprintf(` - DELETE FROM %s WHERE topic = ? AND partition_key = ? AND id = ? - `, MessagesTableName), topic, partitionKey, messageID) + DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND id = ? + `, MessagesTableName), tenant, topic, partitionKey, messageID) if err != nil { - return fmt.Errorf("delete from main table topic=%s partition=%s message=%s: %w", topic, partitionKey, messageID, err) + return fmt.Errorf("delete from main table tenant=%s topic=%s partition=%s message=%s: %w", tenant, topic, partitionKey, messageID, err) } if err := tx.Commit(); err != nil { - return fmt.Errorf("commit DLQ transaction topic=%s message=%s: %w", topic, messageID, err) + return fmt.Errorf("commit DLQ transaction tenant=%s topic=%s message=%s: %w", tenant, topic, messageID, err) } return nil @@ -309,7 +332,7 @@ func (s *sqlmessageStore) MoveToDLQ(ctx context.Context, topic string, partition // The caller provides minAckedOffset (from offsetStore), keeping messageStore // free of cross-table queries. // Returns the number of rows deleted. -func (s *sqlmessageStore) GarbageCollect(ctx context.Context, topic string, partitionKey string, minAckedOffset int64) (_ int64, retErr error) { +func (s *sqlmessageStore) GarbageCollect(ctx context.Context, tenant string, topic string, partitionKey string, minAckedOffset int64) (_ int64, retErr error) { op := metrics.Begin(s.scope, "gc", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -319,11 +342,11 @@ func (s *sqlmessageStore) GarbageCollect(ctx context.Context, topic string, part // Delete messages up to the minimum acked offset result, err := s.db.ExecContext(ctx, fmt.Sprintf(` - DELETE FROM %s WHERE topic = ? AND partition_key = ? AND offset <= ? - `, MessagesTableName), topic, partitionKey, minAckedOffset) + DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND offset <= ? + `, MessagesTableName), tenant, topic, partitionKey, minAckedOffset) if err != nil { - return 0, fmt.Errorf("garbage collect messages topic=%s partition=%s: %w", topic, partitionKey, err) + return 0, fmt.Errorf("garbage collect messages tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } // RowsAffected error is swallowed because the DELETE query itself succeeded. @@ -332,6 +355,7 @@ func (s *sqlmessageStore) GarbageCollect(ctx context.Context, topic string, part deleted, err := result.RowsAffected() if err != nil { s.logger.Warnw("garbage collect succeeded but row count unavailable (driver diagnostic failure), no impact on correctness", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, logError, err, @@ -339,6 +363,7 @@ func (s *sqlmessageStore) GarbageCollect(ctx context.Context, topic string, part } if deleted > 0 { s.logger.Debugw("garbage collected messages", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, "deleted", deleted, @@ -351,18 +376,18 @@ func (s *sqlmessageStore) GarbageCollect(ctx context.Context, topic string, part } // GetOffsetsAbove returns message offsets above afterOffset for a partition, ordered ascending. -func (s *sqlmessageStore) GetOffsetsAbove(ctx context.Context, topic string, partitionKey string, afterOffset int64, limit int) (_ []int64, retErr error) { +func (s *sqlmessageStore) GetOffsetsAbove(ctx context.Context, tenant string, topic string, partitionKey string, afterOffset int64, limit int) (_ []int64, retErr error) { op := metrics.Begin(s.scope, "get_offsets_above", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` SELECT offset FROM %s - WHERE topic = ? AND partition_key = ? AND offset > ? + WHERE tenant = ? AND topic = ? AND partition_key = ? AND offset > ? ORDER BY offset ASC LIMIT ? - `, MessagesTableName), topic, partitionKey, afterOffset, limit) + `, MessagesTableName), tenant, topic, partitionKey, afterOffset, limit) if err != nil { - return nil, fmt.Errorf("query offsets topic=%s partition=%s: %w", topic, partitionKey, err) + return nil, fmt.Errorf("query offsets tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } defer rows.Close() @@ -370,12 +395,12 @@ func (s *sqlmessageStore) GetOffsetsAbove(ctx context.Context, topic string, par for rows.Next() { var offset int64 if err := rows.Scan(&offset); err != nil { - return nil, fmt.Errorf("scan offset topic=%s partition=%s: %w", topic, partitionKey, err) + return nil, fmt.Errorf("scan offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } offsets = append(offsets, offset) } if err := rows.Err(); err != nil { - return nil, fmt.Errorf("offset iteration topic=%s partition=%s: %w", topic, partitionKey, err) + return nil, fmt.Errorf("offset iteration tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } return offsets, nil diff --git a/platform/extension/messagequeue/mysql/message_store_test.go b/platform/extension/messagequeue/mysql/message_store_test.go index acd83d730..e726a5393 100644 --- a/platform/extension/messagequeue/mysql/message_store_test.go +++ b/platform/extension/messagequeue/mysql/message_store_test.go @@ -57,8 +57,8 @@ func TestMessageStore_Insert(t *testing.T) { { name: "successful insert with multiple messages", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, - {ID: "msg2", Payload: []byte("payload2"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, + {Tenant: testTenant, ID: "msg2", Payload: []byte("payload2"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, }, setup: func(mock sqlmock.Sqlmock, messages []entityqueue.Message) { mock.ExpectBegin() @@ -77,13 +77,28 @@ func TestMessageStore_Insert(t *testing.T) { setup: func(mock sqlmock.Sqlmock, messages []entityqueue.Message) {}, wantErr: false, }, + { + name: "conflicting queue metadata is rejected", + messages: []entityqueue.Message{{ + Tenant: testTenant, + ID: "msg-conflict", + PartitionKey: "part1", + Metadata: map[string]string{entityqueue.MetadataKeyQueueName: "other-tenant"}, + }}, + setup: func(mock sqlmock.Sqlmock, messages []entityqueue.Message) { + mock.ExpectBegin() + mock.ExpectPrepare("INSERT INTO queue_messages") + mock.ExpectRollback() + }, + wantErr: true, + }, { // Regression: re-publishing the same (topic, partition_key, id) tuple // must succeed silently. sqlmock returns 0 affected rows to simulate // MySQL's ON DUPLICATE KEY UPDATE swallowing the unique-key collision. name: "duplicate publish is idempotent", messages: []entityqueue.Message{ - {ID: "msg-dup", Payload: []byte("payload"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, + {Tenant: testTenant, ID: "msg-dup", Payload: []byte("payload"), PartitionKey: "part1", PublishedAt: time.Now().UnixMilli()}, }, setup: func(mock sqlmock.Sqlmock, messages []entityqueue.Message) { mock.ExpectBegin() @@ -104,7 +119,7 @@ func TestMessageStore_Insert(t *testing.T) { tt.setup(mock, tt.messages) ctx := context.Background() - err := store.Insert(ctx, "test_topic", tt.messages) + err := store.Insert(ctx, testTenant, "test_topic", tt.messages) if tt.wantErr { require.Error(t, err) @@ -126,10 +141,10 @@ func TestMessageStore_Delete(t *testing.T) { messageID := "msg1" mock.ExpectExec("DELETE FROM queue_messages"). - WithArgs(topic, partitionKey, messageID). + WithArgs(testTenant, topic, partitionKey, messageID). WillReturnResult(sqlmock.NewResult(0, 1)) - err := store.Delete(ctx, topic, partitionKey, messageID) + err := store.Delete(ctx, testTenant, topic, partitionKey, messageID) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } @@ -145,14 +160,14 @@ func TestMessageStore_FetchByOffset(t *testing.T) { limit := 10 // Mock query results (no transaction, simple SELECT) - rows := sqlmock.NewRows([]string{"offset", "id", "payload", "metadata", "partition_key", "published_at", "failed_at", "failure_count", "last_error", "original_topic", "failure_detail"}). - AddRow(int64(1), "msg1", []byte("payload1"), []byte("{}"), "part1", time.Now().UnixMilli(), int64(0), 0, "", "", nil) + rows := sqlmock.NewRows([]string{"tenant", "offset", "id", "payload", "metadata", "partition_key", "published_at", "failed_at", "failure_count", "last_error", "original_topic", "failure_detail"}). + AddRow(testTenant, int64(1), "msg1", []byte("payload1"), []byte("{}"), "part1", time.Now().UnixMilli(), int64(0), 0, "", "", nil) mock.ExpectQuery("SELECT (.+) FROM queue_messages"). - WithArgs(topic, partitionKey, currentOffset, limit). + WithArgs(testTenant, topic, partitionKey, currentOffset, limit). WillReturnRows(rows) - results, err := store.FetchByOffset(ctx, topic, partitionKey, currentOffset, limit) + results, err := store.FetchByOffset(ctx, testTenant, topic, partitionKey, currentOffset, limit) require.NoError(t, err) require.Len(t, results, 1) require.Equal(t, "msg1", results[0].ID) @@ -180,25 +195,25 @@ func TestMessageStore_MoveToDLQ(t *testing.T) { AddRow([]byte("payload1"), []byte(`{"key":"value"}`), "part1", time.Now().UnixMilli(), time.Now().UnixMilli()) mock.ExpectQuery("SELECT (.+) FROM queue_messages"). - WithArgs(topic, partitionKey, messageID). + WithArgs(testTenant, topic, partitionKey, messageID). WillReturnRows(rows) // Expect insert into queue_messages with DLQ topic. The failure's message // goes to last_error; failure_detail is NULL because this failure names no // subjects — see TestMessageStore_MoveToDLQ_WritesFailureDetail. mock.ExpectExec("INSERT INTO queue_messages"). - WithArgs(dlqTopic, messageID, sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), failureCount, lastError, topic, nil). + WithArgs(testTenant, dlqTopic, messageID, sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), failureCount, lastError, topic, nil). WillReturnResult(sqlmock.NewResult(1, 1)) // Expect delete from main table (now includes partition_key in WHERE) mock.ExpectExec("DELETE FROM queue_messages"). - WithArgs(topic, partitionKey, messageID). + WithArgs(testTenant, topic, partitionKey, messageID). WillReturnResult(sqlmock.NewResult(0, 1)) // Expect commit mock.ExpectCommit() - err := store.MoveToDLQ(ctx, topic, partitionKey, messageID, failureCount, failure.New(lastError), dlqTopicSuffix) + err := store.MoveToDLQ(ctx, testTenant, topic, partitionKey, messageID, failureCount, failure.New(lastError), dlqTopicSuffix) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } @@ -218,18 +233,18 @@ func TestMessageStore_MoveToDLQ_WritesFailureDetail(t *testing.T) { mock.ExpectBegin() mock.ExpectQuery("SELECT (.+) FROM queue_messages"). - WithArgs("test_topic", "part1", "msg1"). + WithArgs(testTenant, "test_topic", "part1", "msg1"). WillReturnRows(sqlmock.NewRows([]string{"payload", "metadata", "partition_key", "created_at", "published_at"}). AddRow([]byte("payload1"), nil, "part1", int64(1), int64(2))) mock.ExpectExec("INSERT INTO queue_messages"). - WithArgs("test_topic_dlq", "msg1", sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), 3, "speculator failed", "test_topic", encoded). + WithArgs(testTenant, "test_topic_dlq", "msg1", sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), 3, "speculator failed", "test_topic", encoded). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectExec("DELETE FROM queue_messages"). - WithArgs("test_topic", "part1", "msg1"). + WithArgs(testTenant, "test_topic", "part1", "msg1"). WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() - require.NoError(t, store.MoveToDLQ(context.Background(), "test_topic", "part1", "msg1", 3, f, "_dlq")) + require.NoError(t, store.MoveToDLQ(context.Background(), testTenant, "test_topic", "part1", "msg1", 3, f, "_dlq")) require.NoError(t, mock.ExpectationsWereMet()) } @@ -274,7 +289,7 @@ func TestMessageStore_GetOffsetsAbove(t *testing.T) { if tt.wantErr { mock.ExpectQuery("SELECT offset FROM queue_messages"). - WithArgs("test_topic", "part-1", tt.afterOffset, tt.limit). + WithArgs(testTenant, "test_topic", "part-1", tt.afterOffset, tt.limit). WillReturnError(fmt.Errorf("db error")) } else { rows := sqlmock.NewRows([]string{"offset"}) @@ -282,11 +297,11 @@ func TestMessageStore_GetOffsetsAbove(t *testing.T) { rows.AddRow(offset) } mock.ExpectQuery("SELECT offset FROM queue_messages"). - WithArgs("test_topic", "part-1", tt.afterOffset, tt.limit). + WithArgs(testTenant, "test_topic", "part-1", tt.afterOffset, tt.limit). WillReturnRows(rows) } - offsets, err := store.GetOffsetsAbove(context.Background(), "test_topic", "part-1", tt.afterOffset, tt.limit) + offsets, err := store.GetOffsetsAbove(context.Background(), testTenant, "test_topic", "part-1", tt.afterOffset, tt.limit) if tt.wantErr { require.Error(t, err) @@ -333,16 +348,16 @@ func TestMessageStore_GarbageCollect(t *testing.T) { if tt.minAckedOffset > 0 { if tt.deleteErr { mock.ExpectExec("DELETE FROM queue_messages"). - WithArgs("test_topic", "part-1", tt.minAckedOffset). + WithArgs(testTenant, "test_topic", "part-1", tt.minAckedOffset). WillReturnError(fmt.Errorf("db error")) } else { mock.ExpectExec("DELETE FROM queue_messages"). - WithArgs("test_topic", "part-1", tt.minAckedOffset). + WithArgs(testTenant, "test_topic", "part-1", tt.minAckedOffset). WillReturnResult(sqlmock.NewResult(0, tt.wantDeleted)) } } - deleted, err := store.GarbageCollect(context.Background(), "test_topic", "part-1", tt.minAckedOffset) + deleted, err := store.GarbageCollect(context.Background(), testTenant, "test_topic", "part-1", tt.minAckedOffset) if tt.wantErr { require.Error(t, err) diff --git a/platform/extension/messagequeue/mysql/mock_stores.go b/platform/extension/messagequeue/mysql/mock_stores.go index 614c93d11..a9be425e5 100644 --- a/platform/extension/messagequeue/mysql/mock_stores.go +++ b/platform/extension/messagequeue/mysql/mock_stores.go @@ -43,90 +43,90 @@ func (m *MockmessageStore) EXPECT() *MockmessageStoreMockRecorder { } // Delete mocks base method. -func (m *MockmessageStore) Delete(ctx context.Context, topic, partitionKey, messageID string) error { +func (m *MockmessageStore) Delete(ctx context.Context, tenant, topic, partitionKey, messageID string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Delete", ctx, topic, partitionKey, messageID) + ret := m.ctrl.Call(m, "Delete", ctx, tenant, topic, partitionKey, messageID) ret0, _ := ret[0].(error) return ret0 } // Delete indicates an expected call of Delete. -func (mr *MockmessageStoreMockRecorder) Delete(ctx, topic, partitionKey, messageID any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) Delete(ctx, tenant, topic, partitionKey, messageID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockmessageStore)(nil).Delete), ctx, topic, partitionKey, messageID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockmessageStore)(nil).Delete), ctx, tenant, topic, partitionKey, messageID) } // FetchByOffset mocks base method. -func (m *MockmessageStore) FetchByOffset(ctx context.Context, topic, partitionKey string, currentOffset int64, limit int) ([]messageRow, error) { +func (m *MockmessageStore) FetchByOffset(ctx context.Context, tenant, topic, partitionKey string, currentOffset int64, limit int) ([]messageRow, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "FetchByOffset", ctx, topic, partitionKey, currentOffset, limit) + ret := m.ctrl.Call(m, "FetchByOffset", ctx, tenant, topic, partitionKey, currentOffset, limit) ret0, _ := ret[0].([]messageRow) ret1, _ := ret[1].(error) return ret0, ret1 } // FetchByOffset indicates an expected call of FetchByOffset. -func (mr *MockmessageStoreMockRecorder) FetchByOffset(ctx, topic, partitionKey, currentOffset, limit any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) FetchByOffset(ctx, tenant, topic, partitionKey, currentOffset, limit any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FetchByOffset", reflect.TypeOf((*MockmessageStore)(nil).FetchByOffset), ctx, topic, partitionKey, currentOffset, limit) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FetchByOffset", reflect.TypeOf((*MockmessageStore)(nil).FetchByOffset), ctx, tenant, topic, partitionKey, currentOffset, limit) } // GarbageCollect mocks base method. -func (m *MockmessageStore) GarbageCollect(ctx context.Context, topic, partitionKey string, minAckedOffset int64) (int64, error) { +func (m *MockmessageStore) GarbageCollect(ctx context.Context, tenant, topic, partitionKey string, minAckedOffset int64) (int64, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GarbageCollect", ctx, topic, partitionKey, minAckedOffset) + ret := m.ctrl.Call(m, "GarbageCollect", ctx, tenant, topic, partitionKey, minAckedOffset) ret0, _ := ret[0].(int64) ret1, _ := ret[1].(error) return ret0, ret1 } // GarbageCollect indicates an expected call of GarbageCollect. -func (mr *MockmessageStoreMockRecorder) GarbageCollect(ctx, topic, partitionKey, minAckedOffset any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) GarbageCollect(ctx, tenant, topic, partitionKey, minAckedOffset any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GarbageCollect", reflect.TypeOf((*MockmessageStore)(nil).GarbageCollect), ctx, topic, partitionKey, minAckedOffset) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GarbageCollect", reflect.TypeOf((*MockmessageStore)(nil).GarbageCollect), ctx, tenant, topic, partitionKey, minAckedOffset) } // GetOffsetsAbove mocks base method. -func (m *MockmessageStore) GetOffsetsAbove(ctx context.Context, topic, partitionKey string, afterOffset int64, limit int) ([]int64, error) { +func (m *MockmessageStore) GetOffsetsAbove(ctx context.Context, tenant, topic, partitionKey string, afterOffset int64, limit int) ([]int64, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetOffsetsAbove", ctx, topic, partitionKey, afterOffset, limit) + ret := m.ctrl.Call(m, "GetOffsetsAbove", ctx, tenant, topic, partitionKey, afterOffset, limit) ret0, _ := ret[0].([]int64) ret1, _ := ret[1].(error) return ret0, ret1 } // GetOffsetsAbove indicates an expected call of GetOffsetsAbove. -func (mr *MockmessageStoreMockRecorder) GetOffsetsAbove(ctx, topic, partitionKey, afterOffset, limit any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) GetOffsetsAbove(ctx, tenant, topic, partitionKey, afterOffset, limit any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOffsetsAbove", reflect.TypeOf((*MockmessageStore)(nil).GetOffsetsAbove), ctx, topic, partitionKey, afterOffset, limit) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOffsetsAbove", reflect.TypeOf((*MockmessageStore)(nil).GetOffsetsAbove), ctx, tenant, topic, partitionKey, afterOffset, limit) } // Insert mocks base method. -func (m *MockmessageStore) Insert(ctx context.Context, topic string, messages []messagequeue.Message) error { +func (m *MockmessageStore) Insert(ctx context.Context, tenant, topic string, messages []messagequeue.Message) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Insert", ctx, topic, messages) + ret := m.ctrl.Call(m, "Insert", ctx, tenant, topic, messages) ret0, _ := ret[0].(error) return ret0 } // Insert indicates an expected call of Insert. -func (mr *MockmessageStoreMockRecorder) Insert(ctx, topic, messages any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) Insert(ctx, tenant, topic, messages any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Insert", reflect.TypeOf((*MockmessageStore)(nil).Insert), ctx, topic, messages) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Insert", reflect.TypeOf((*MockmessageStore)(nil).Insert), ctx, tenant, topic, messages) } // MoveToDLQ mocks base method. -func (m *MockmessageStore) MoveToDLQ(ctx context.Context, topic, partitionKey, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) error { +func (m *MockmessageStore) MoveToDLQ(ctx context.Context, tenant, topic, partitionKey, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MoveToDLQ", ctx, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix) + ret := m.ctrl.Call(m, "MoveToDLQ", ctx, tenant, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix) ret0, _ := ret[0].(error) return ret0 } // MoveToDLQ indicates an expected call of MoveToDLQ. -func (mr *MockmessageStoreMockRecorder) MoveToDLQ(ctx, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix any) *gomock.Call { +func (mr *MockmessageStoreMockRecorder) MoveToDLQ(ctx, tenant, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MoveToDLQ", reflect.TypeOf((*MockmessageStore)(nil).MoveToDLQ), ctx, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MoveToDLQ", reflect.TypeOf((*MockmessageStore)(nil).MoveToDLQ), ctx, tenant, topic, partitionKey, messageID, failureCount, f, dlqTopicSuffix) } // MockoffsetStore is a mock of offsetStore interface. @@ -154,38 +154,38 @@ func (m *MockoffsetStore) EXPECT() *MockoffsetStoreMockRecorder { } // DeleteOffset mocks base method. -func (m *MockoffsetStore) DeleteOffset(ctx context.Context, topic, partitionKey, consumerGroup string) error { +func (m *MockoffsetStore) DeleteOffset(ctx context.Context, tenant, topic, partitionKey, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "DeleteOffset", ctx, topic, partitionKey, consumerGroup) + ret := m.ctrl.Call(m, "DeleteOffset", ctx, tenant, topic, partitionKey, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // DeleteOffset indicates an expected call of DeleteOffset. -func (mr *MockoffsetStoreMockRecorder) DeleteOffset(ctx, topic, partitionKey, consumerGroup any) *gomock.Call { +func (mr *MockoffsetStoreMockRecorder) DeleteOffset(ctx, tenant, topic, partitionKey, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOffset", reflect.TypeOf((*MockoffsetStore)(nil).DeleteOffset), ctx, topic, partitionKey, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOffset", reflect.TypeOf((*MockoffsetStore)(nil).DeleteOffset), ctx, tenant, topic, partitionKey, consumerGroup) } // GetAckedOffset mocks base method. -func (m *MockoffsetStore) GetAckedOffset(ctx context.Context, topic, partitionKey, consumerGroup string) (int64, error) { +func (m *MockoffsetStore) GetAckedOffset(ctx context.Context, tenant, topic, partitionKey, consumerGroup string) (int64, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetAckedOffset", ctx, topic, partitionKey, consumerGroup) + ret := m.ctrl.Call(m, "GetAckedOffset", ctx, tenant, topic, partitionKey, consumerGroup) ret0, _ := ret[0].(int64) ret1, _ := ret[1].(error) return ret0, ret1 } // GetAckedOffset indicates an expected call of GetAckedOffset. -func (mr *MockoffsetStoreMockRecorder) GetAckedOffset(ctx, topic, partitionKey, consumerGroup any) *gomock.Call { +func (mr *MockoffsetStoreMockRecorder) GetAckedOffset(ctx, tenant, topic, partitionKey, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).GetAckedOffset), ctx, topic, partitionKey, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).GetAckedOffset), ctx, tenant, topic, partitionKey, consumerGroup) } // GetMinAckedOffset mocks base method. -func (m *MockoffsetStore) GetMinAckedOffset(ctx context.Context, topic, partitionKey string) (int64, bool, error) { +func (m *MockoffsetStore) GetMinAckedOffset(ctx context.Context, tenant, topic, partitionKey string) (int64, bool, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetMinAckedOffset", ctx, topic, partitionKey) + ret := m.ctrl.Call(m, "GetMinAckedOffset", ctx, tenant, topic, partitionKey) ret0, _ := ret[0].(int64) ret1, _ := ret[1].(bool) ret2, _ := ret[2].(error) @@ -193,37 +193,37 @@ func (m *MockoffsetStore) GetMinAckedOffset(ctx context.Context, topic, partitio } // GetMinAckedOffset indicates an expected call of GetMinAckedOffset. -func (mr *MockoffsetStoreMockRecorder) GetMinAckedOffset(ctx, topic, partitionKey any) *gomock.Call { +func (mr *MockoffsetStoreMockRecorder) GetMinAckedOffset(ctx, tenant, topic, partitionKey any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMinAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).GetMinAckedOffset), ctx, topic, partitionKey) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMinAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).GetMinAckedOffset), ctx, tenant, topic, partitionKey) } // Initialize mocks base method. -func (m *MockoffsetStore) Initialize(ctx context.Context, topic, partitionKey, consumerGroup string) error { +func (m *MockoffsetStore) Initialize(ctx context.Context, tenant, topic, partitionKey, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Initialize", ctx, topic, partitionKey, consumerGroup) + ret := m.ctrl.Call(m, "Initialize", ctx, tenant, topic, partitionKey, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // Initialize indicates an expected call of Initialize. -func (mr *MockoffsetStoreMockRecorder) Initialize(ctx, topic, partitionKey, consumerGroup any) *gomock.Call { +func (mr *MockoffsetStoreMockRecorder) Initialize(ctx, tenant, topic, partitionKey, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Initialize", reflect.TypeOf((*MockoffsetStore)(nil).Initialize), ctx, topic, partitionKey, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Initialize", reflect.TypeOf((*MockoffsetStore)(nil).Initialize), ctx, tenant, topic, partitionKey, consumerGroup) } // UpdateAckedOffset mocks base method. -func (m *MockoffsetStore) UpdateAckedOffset(ctx context.Context, topic, partitionKey string, offset int64, consumerGroup string) error { +func (m *MockoffsetStore) UpdateAckedOffset(ctx context.Context, tenant, topic, partitionKey string, offset int64, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateAckedOffset", ctx, topic, partitionKey, offset, consumerGroup) + ret := m.ctrl.Call(m, "UpdateAckedOffset", ctx, tenant, topic, partitionKey, offset, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // UpdateAckedOffset indicates an expected call of UpdateAckedOffset. -func (mr *MockoffsetStoreMockRecorder) UpdateAckedOffset(ctx, topic, partitionKey, offset, consumerGroup any) *gomock.Call { +func (mr *MockoffsetStoreMockRecorder) UpdateAckedOffset(ctx, tenant, topic, partitionKey, offset, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).UpdateAckedOffset), ctx, topic, partitionKey, offset, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateAckedOffset", reflect.TypeOf((*MockoffsetStore)(nil).UpdateAckedOffset), ctx, tenant, topic, partitionKey, offset, consumerGroup) } // MockpartitionLeaseStore is a mock of partitionLeaseStore interface. @@ -251,9 +251,9 @@ func (m *MockpartitionLeaseStore) EXPECT() *MockpartitionLeaseStoreMockRecorder } // DiscoverAndAcquirePartitions mocks base method. -func (m *MockpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Context, topic, subscriberName, consumerGroup string, leaseDurationMs int64, maxPartitions int) (int, []string, error) { +func (m *MockpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Context, tenant, topic, subscriberName, consumerGroup string, leaseDurationMs int64, maxPartitions int) (int, []string, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "DiscoverAndAcquirePartitions", ctx, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) + ret := m.ctrl.Call(m, "DiscoverAndAcquirePartitions", ctx, tenant, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) ret0, _ := ret[0].(int) ret1, _ := ret[1].([]string) ret2, _ := ret[2].(error) @@ -261,96 +261,96 @@ func (m *MockpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Conte } // DiscoverAndAcquirePartitions indicates an expected call of DiscoverAndAcquirePartitions. -func (mr *MockpartitionLeaseStoreMockRecorder) DiscoverAndAcquirePartitions(ctx, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) DiscoverAndAcquirePartitions(ctx, tenant, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverAndAcquirePartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).DiscoverAndAcquirePartitions), ctx, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverAndAcquirePartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).DiscoverAndAcquirePartitions), ctx, tenant, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) } // GetAllLeases mocks base method. -func (m *MockpartitionLeaseStore) GetAllLeases(ctx context.Context, topic, consumerGroup string) ([]leaseInfo, error) { +func (m *MockpartitionLeaseStore) GetAllLeases(ctx context.Context, tenant, topic, consumerGroup string) ([]leaseInfo, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetAllLeases", ctx, topic, consumerGroup) + ret := m.ctrl.Call(m, "GetAllLeases", ctx, tenant, topic, consumerGroup) ret0, _ := ret[0].([]leaseInfo) ret1, _ := ret[1].(error) return ret0, ret1 } // GetAllLeases indicates an expected call of GetAllLeases. -func (mr *MockpartitionLeaseStoreMockRecorder) GetAllLeases(ctx, topic, consumerGroup any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) GetAllLeases(ctx, tenant, topic, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllLeases", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetAllLeases), ctx, topic, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllLeases", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetAllLeases), ctx, tenant, topic, consumerGroup) } // GetLeasedPartitions mocks base method. -func (m *MockpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic, subscriberName, consumerGroup string) ([]string, error) { +func (m *MockpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, tenant, topic, subscriberName, consumerGroup string) ([]string, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetLeasedPartitions", ctx, topic, subscriberName, consumerGroup) + ret := m.ctrl.Call(m, "GetLeasedPartitions", ctx, tenant, topic, subscriberName, consumerGroup) ret0, _ := ret[0].([]string) ret1, _ := ret[1].(error) return ret0, ret1 } // GetLeasedPartitions indicates an expected call of GetLeasedPartitions. -func (mr *MockpartitionLeaseStoreMockRecorder) GetLeasedPartitions(ctx, topic, subscriberName, consumerGroup any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) GetLeasedPartitions(ctx, tenant, topic, subscriberName, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLeasedPartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetLeasedPartitions), ctx, topic, subscriberName, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLeasedPartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetLeasedPartitions), ctx, tenant, topic, subscriberName, consumerGroup) } // PurgeStale mocks base method. -func (m *MockpartitionLeaseStore) PurgeStale(ctx context.Context, topic, consumerGroup string, olderThanMs int64) error { +func (m *MockpartitionLeaseStore) PurgeStale(ctx context.Context, tenant, topic, consumerGroup string, olderThanMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "PurgeStale", ctx, topic, consumerGroup, olderThanMs) + ret := m.ctrl.Call(m, "PurgeStale", ctx, tenant, topic, consumerGroup, olderThanMs) ret0, _ := ret[0].(error) return ret0 } // PurgeStale indicates an expected call of PurgeStale. -func (mr *MockpartitionLeaseStoreMockRecorder) PurgeStale(ctx, topic, consumerGroup, olderThanMs any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) PurgeStale(ctx, tenant, topic, consumerGroup, olderThanMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PurgeStale", reflect.TypeOf((*MockpartitionLeaseStore)(nil).PurgeStale), ctx, topic, consumerGroup, olderThanMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PurgeStale", reflect.TypeOf((*MockpartitionLeaseStore)(nil).PurgeStale), ctx, tenant, topic, consumerGroup, olderThanMs) } // ReleaseLease mocks base method. -func (m *MockpartitionLeaseStore) ReleaseLease(ctx context.Context, topic, partitionKey, subscriberName, consumerGroup string) error { +func (m *MockpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant, topic, partitionKey, subscriberName, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "ReleaseLease", ctx, topic, partitionKey, subscriberName, consumerGroup) + ret := m.ctrl.Call(m, "ReleaseLease", ctx, tenant, topic, partitionKey, subscriberName, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // ReleaseLease indicates an expected call of ReleaseLease. -func (mr *MockpartitionLeaseStoreMockRecorder) ReleaseLease(ctx, topic, partitionKey, subscriberName, consumerGroup any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) ReleaseLease(ctx, tenant, topic, partitionKey, subscriberName, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReleaseLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).ReleaseLease), ctx, topic, partitionKey, subscriberName, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReleaseLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).ReleaseLease), ctx, tenant, topic, partitionKey, subscriberName, consumerGroup) } // RenewLease mocks base method. -func (m *MockpartitionLeaseStore) RenewLease(ctx context.Context, topic, partitionKey, subscriberName, consumerGroup string, leaseDurationMs int64) error { +func (m *MockpartitionLeaseStore) RenewLease(ctx context.Context, tenant, topic, partitionKey, subscriberName, consumerGroup string, leaseDurationMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "RenewLease", ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) + ret := m.ctrl.Call(m, "RenewLease", ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) ret0, _ := ret[0].(error) return ret0 } // RenewLease indicates an expected call of RenewLease. -func (mr *MockpartitionLeaseStoreMockRecorder) RenewLease(ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) RenewLease(ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RenewLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).RenewLease), ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RenewLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).RenewLease), ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) } // TryAcquireLease mocks base method. -func (m *MockpartitionLeaseStore) TryAcquireLease(ctx context.Context, topic, partitionKey, subscriberName, consumerGroup string, leaseDurationMs int64) (bool, error) { +func (m *MockpartitionLeaseStore) TryAcquireLease(ctx context.Context, tenant, topic, partitionKey, subscriberName, consumerGroup string, leaseDurationMs int64) (bool, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "TryAcquireLease", ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) + ret := m.ctrl.Call(m, "TryAcquireLease", ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) ret0, _ := ret[0].(bool) ret1, _ := ret[1].(error) return ret0, ret1 } // TryAcquireLease indicates an expected call of TryAcquireLease. -func (mr *MockpartitionLeaseStoreMockRecorder) TryAcquireLease(ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs any) *gomock.Call { +func (mr *MockpartitionLeaseStoreMockRecorder) TryAcquireLease(ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TryAcquireLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).TryAcquireLease), ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TryAcquireLease", reflect.TypeOf((*MockpartitionLeaseStore)(nil).TryAcquireLease), ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) } // MocksubscriberHeartbeatStore is a mock of subscriberHeartbeatStore interface. @@ -378,60 +378,60 @@ func (m *MocksubscriberHeartbeatStore) EXPECT() *MocksubscriberHeartbeatStoreMoc } // ActiveSubscribers mocks base method. -func (m *MocksubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, topic, consumerGroup string, staleDurationMs int64) ([]string, error) { +func (m *MocksubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, tenant, topic, consumerGroup string, staleDurationMs int64) ([]string, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "ActiveSubscribers", ctx, topic, consumerGroup, staleDurationMs) + ret := m.ctrl.Call(m, "ActiveSubscribers", ctx, tenant, topic, consumerGroup, staleDurationMs) ret0, _ := ret[0].([]string) ret1, _ := ret[1].(error) return ret0, ret1 } // ActiveSubscribers indicates an expected call of ActiveSubscribers. -func (mr *MocksubscriberHeartbeatStoreMockRecorder) ActiveSubscribers(ctx, topic, consumerGroup, staleDurationMs any) *gomock.Call { +func (mr *MocksubscriberHeartbeatStoreMockRecorder) ActiveSubscribers(ctx, tenant, topic, consumerGroup, staleDurationMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ActiveSubscribers", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).ActiveSubscribers), ctx, topic, consumerGroup, staleDurationMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ActiveSubscribers", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).ActiveSubscribers), ctx, tenant, topic, consumerGroup, staleDurationMs) } // Deregister mocks base method. -func (m *MocksubscriberHeartbeatStore) Deregister(ctx context.Context, topic, subscriberName, consumerGroup string) error { +func (m *MocksubscriberHeartbeatStore) Deregister(ctx context.Context, tenant, topic, subscriberName, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Deregister", ctx, topic, subscriberName, consumerGroup) + ret := m.ctrl.Call(m, "Deregister", ctx, tenant, topic, subscriberName, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // Deregister indicates an expected call of Deregister. -func (mr *MocksubscriberHeartbeatStoreMockRecorder) Deregister(ctx, topic, subscriberName, consumerGroup any) *gomock.Call { +func (mr *MocksubscriberHeartbeatStoreMockRecorder) Deregister(ctx, tenant, topic, subscriberName, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Deregister", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).Deregister), ctx, topic, subscriberName, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Deregister", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).Deregister), ctx, tenant, topic, subscriberName, consumerGroup) } // Heartbeat mocks base method. -func (m *MocksubscriberHeartbeatStore) Heartbeat(ctx context.Context, topic, subscriberName, consumerGroup string) error { +func (m *MocksubscriberHeartbeatStore) Heartbeat(ctx context.Context, tenant, topic, subscriberName, consumerGroup string) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Heartbeat", ctx, topic, subscriberName, consumerGroup) + ret := m.ctrl.Call(m, "Heartbeat", ctx, tenant, topic, subscriberName, consumerGroup) ret0, _ := ret[0].(error) return ret0 } // Heartbeat indicates an expected call of Heartbeat. -func (mr *MocksubscriberHeartbeatStoreMockRecorder) Heartbeat(ctx, topic, subscriberName, consumerGroup any) *gomock.Call { +func (mr *MocksubscriberHeartbeatStoreMockRecorder) Heartbeat(ctx, tenant, topic, subscriberName, consumerGroup any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Heartbeat", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).Heartbeat), ctx, topic, subscriberName, consumerGroup) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Heartbeat", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).Heartbeat), ctx, tenant, topic, subscriberName, consumerGroup) } // PurgeStale mocks base method. -func (m *MocksubscriberHeartbeatStore) PurgeStale(ctx context.Context, topic, consumerGroup string, olderThanMs int64) error { +func (m *MocksubscriberHeartbeatStore) PurgeStale(ctx context.Context, tenant, topic, consumerGroup string, olderThanMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "PurgeStale", ctx, topic, consumerGroup, olderThanMs) + ret := m.ctrl.Call(m, "PurgeStale", ctx, tenant, topic, consumerGroup, olderThanMs) ret0, _ := ret[0].(error) return ret0 } // PurgeStale indicates an expected call of PurgeStale. -func (mr *MocksubscriberHeartbeatStoreMockRecorder) PurgeStale(ctx, topic, consumerGroup, olderThanMs any) *gomock.Call { +func (mr *MocksubscriberHeartbeatStoreMockRecorder) PurgeStale(ctx, tenant, topic, consumerGroup, olderThanMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PurgeStale", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).PurgeStale), ctx, topic, consumerGroup, olderThanMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PurgeStale", reflect.TypeOf((*MocksubscriberHeartbeatStore)(nil).PurgeStale), ctx, tenant, topic, consumerGroup, olderThanMs) } // MockdeliveryStateStore is a mock of deliveryStateStore interface. @@ -459,38 +459,38 @@ func (m *MockdeliveryStateStore) EXPECT() *MockdeliveryStateStoreMockRecorder { } // AdvanceWatermark mocks base method. -func (m *MockdeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGroup, topic, partitionKey string, currentWatermark int64, offsets []int64) (int64, error) { +func (m *MockdeliveryStateStore) AdvanceWatermark(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, currentWatermark int64, offsets []int64) (int64, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "AdvanceWatermark", ctx, consumerGroup, topic, partitionKey, currentWatermark, offsets) + ret := m.ctrl.Call(m, "AdvanceWatermark", ctx, consumerGroup, tenant, topic, partitionKey, currentWatermark, offsets) ret0, _ := ret[0].(int64) ret1, _ := ret[1].(error) return ret0, ret1 } // AdvanceWatermark indicates an expected call of AdvanceWatermark. -func (mr *MockdeliveryStateStoreMockRecorder) AdvanceWatermark(ctx, consumerGroup, topic, partitionKey, currentWatermark, offsets any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) AdvanceWatermark(ctx, consumerGroup, tenant, topic, partitionKey, currentWatermark, offsets any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AdvanceWatermark", reflect.TypeOf((*MockdeliveryStateStore)(nil).AdvanceWatermark), ctx, consumerGroup, topic, partitionKey, currentWatermark, offsets) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AdvanceWatermark", reflect.TypeOf((*MockdeliveryStateStore)(nil).AdvanceWatermark), ctx, consumerGroup, tenant, topic, partitionKey, currentWatermark, offsets) } // ExtendVisibility mocks base method. -func (m *MockdeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGroup, topic, partitionKey string, offset, visibilityTimeoutMs int64) error { +func (m *MockdeliveryStateStore) ExtendVisibility(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset, visibilityTimeoutMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "ExtendVisibility", ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs) + ret := m.ctrl.Call(m, "ExtendVisibility", ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs) ret0, _ := ret[0].(error) return ret0 } // ExtendVisibility indicates an expected call of ExtendVisibility. -func (mr *MockdeliveryStateStoreMockRecorder) ExtendVisibility(ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) ExtendVisibility(ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExtendVisibility", reflect.TypeOf((*MockdeliveryStateStore)(nil).ExtendVisibility), ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExtendVisibility", reflect.TypeOf((*MockdeliveryStateStore)(nil).ExtendVisibility), ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs) } // GetDeliveryState mocks base method. -func (m *MockdeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (DeliveryState, bool, error) { +func (m *MockdeliveryStateStore) GetDeliveryState(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) (DeliveryState, bool, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetDeliveryState", ctx, consumerGroup, topic, partitionKey, offset) + ret := m.ctrl.Call(m, "GetDeliveryState", ctx, consumerGroup, tenant, topic, partitionKey, offset) ret0, _ := ret[0].(DeliveryState) ret1, _ := ret[1].(bool) ret2, _ := ret[2].(error) @@ -498,64 +498,64 @@ func (m *MockdeliveryStateStore) GetDeliveryState(ctx context.Context, consumerG } // GetDeliveryState indicates an expected call of GetDeliveryState. -func (mr *MockdeliveryStateStoreMockRecorder) GetDeliveryState(ctx, consumerGroup, topic, partitionKey, offset any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) GetDeliveryState(ctx, consumerGroup, tenant, topic, partitionKey, offset any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeliveryState", reflect.TypeOf((*MockdeliveryStateStore)(nil).GetDeliveryState), ctx, consumerGroup, topic, partitionKey, offset) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeliveryState", reflect.TypeOf((*MockdeliveryStateStore)(nil).GetDeliveryState), ctx, consumerGroup, tenant, topic, partitionKey, offset) } // MarkAcked mocks base method. -func (m *MockdeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) error { +func (m *MockdeliveryStateStore) MarkAcked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkAcked", ctx, consumerGroup, topic, partitionKey, offset) + ret := m.ctrl.Call(m, "MarkAcked", ctx, consumerGroup, tenant, topic, partitionKey, offset) ret0, _ := ret[0].(error) return ret0 } // MarkAcked indicates an expected call of MarkAcked. -func (mr *MockdeliveryStateStoreMockRecorder) MarkAcked(ctx, consumerGroup, topic, partitionKey, offset any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) MarkAcked(ctx, consumerGroup, tenant, topic, partitionKey, offset any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkAcked", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkAcked), ctx, consumerGroup, topic, partitionKey, offset) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkAcked", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkAcked), ctx, consumerGroup, tenant, topic, partitionKey, offset) } // MarkDelivered mocks base method. -func (m *MockdeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup, topic, partitionKey string, offset, visibilityTimeoutMs int64) (int, error) { +func (m *MockdeliveryStateStore) MarkDelivered(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset, visibilityTimeoutMs int64) (int, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkDelivered", ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs) + ret := m.ctrl.Call(m, "MarkDelivered", ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs) ret0, _ := ret[0].(int) ret1, _ := ret[1].(error) return ret0, ret1 } // MarkDelivered indicates an expected call of MarkDelivered. -func (mr *MockdeliveryStateStoreMockRecorder) MarkDelivered(ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) MarkDelivered(ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkDelivered", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkDelivered), ctx, consumerGroup, topic, partitionKey, offset, visibilityTimeoutMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkDelivered", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkDelivered), ctx, consumerGroup, tenant, topic, partitionKey, offset, visibilityTimeoutMs) } // MarkNacked mocks base method. -func (m *MockdeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, topic, partitionKey string, offset, delayMs int64) error { +func (m *MockdeliveryStateStore) MarkNacked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset, delayMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkNacked", ctx, consumerGroup, topic, partitionKey, offset, delayMs) + ret := m.ctrl.Call(m, "MarkNacked", ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs) ret0, _ := ret[0].(error) return ret0 } // MarkNacked indicates an expected call of MarkNacked. -func (mr *MockdeliveryStateStoreMockRecorder) MarkNacked(ctx, consumerGroup, topic, partitionKey, offset, delayMs any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) MarkNacked(ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkNacked", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkNacked), ctx, consumerGroup, topic, partitionKey, offset, delayMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkNacked", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkNacked), ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs) } // MarkPostponed mocks base method. -func (m *MockdeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup, topic, partitionKey string, offset, delayMs int64) error { +func (m *MockdeliveryStateStore) MarkPostponed(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset, delayMs int64) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkPostponed", ctx, consumerGroup, topic, partitionKey, offset, delayMs) + ret := m.ctrl.Call(m, "MarkPostponed", ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs) ret0, _ := ret[0].(error) return ret0 } // MarkPostponed indicates an expected call of MarkPostponed. -func (mr *MockdeliveryStateStoreMockRecorder) MarkPostponed(ctx, consumerGroup, topic, partitionKey, offset, delayMs any) *gomock.Call { +func (mr *MockdeliveryStateStoreMockRecorder) MarkPostponed(ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPostponed", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkPostponed), ctx, consumerGroup, topic, partitionKey, offset, delayMs) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPostponed", reflect.TypeOf((*MockdeliveryStateStore)(nil).MarkPostponed), ctx, consumerGroup, tenant, topic, partitionKey, offset, delayMs) } diff --git a/platform/extension/messagequeue/mysql/offset_store.go b/platform/extension/messagequeue/mysql/offset_store.go index a6ed71b34..62757c67e 100644 --- a/platform/extension/messagequeue/mysql/offset_store.go +++ b/platform/extension/messagequeue/mysql/offset_store.go @@ -38,8 +38,8 @@ func newOffsetStore(db *sql.DB, scope tally.Scope) offsetStore { } } -// Initialize creates an offset entry for a topic+partition if it doesn't exist -func (s *sqloffsetStore) Initialize(ctx context.Context, topic string, partitionKey string, consumerGroup string) (retErr error) { +// Initialize creates an offset entry for a tenant+topic+partition if it doesn't exist +func (s *sqloffsetStore) Initialize(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "initialize", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -49,19 +49,19 @@ func (s *sqloffsetStore) Initialize(ctx context.Context, topic string, partition // Try to insert, ignore if already exists _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT IGNORE INTO %s (consumer_group, topic, partition_key, offset_acked, updated_at) - VALUES (?, ?, ?, 0, ?) - `, OffsetsTableName), consumerGroup, topic, partitionKey, now) + INSERT IGNORE INTO %s (tenant, topic, partition_key, consumer_group, offset_acked, updated_at) + VALUES (?, ?, ?, ?, 0, ?) + `, OffsetsTableName), tenant, topic, partitionKey, consumerGroup, now) if err != nil { - return fmt.Errorf("initialize offset topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("initialize offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } return nil } -// GetAckedOffset returns the current acked offset for a topic+partition -func (s *sqloffsetStore) GetAckedOffset(ctx context.Context, topic string, partitionKey string, consumerGroup string) (_ int64, retErr error) { +// GetAckedOffset returns the current acked offset for a tenant+topic+partition +func (s *sqloffsetStore) GetAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) (_ int64, retErr error) { op := metrics.Begin(s.scope, "get_acked_offset", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -69,8 +69,8 @@ func (s *sqloffsetStore) GetAckedOffset(ctx context.Context, topic string, parti var offset int64 err := s.db.QueryRowContext(ctx, fmt.Sprintf(` - SELECT offset_acked FROM %s WHERE consumer_group = ? AND topic = ? AND partition_key = ? - `, OffsetsTableName), consumerGroup, topic, partitionKey).Scan(&offset) + SELECT offset_acked FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND consumer_group = ? + `, OffsetsTableName), tenant, topic, partitionKey, consumerGroup).Scan(&offset) if err == sql.ErrNoRows { // Partition not yet initialized, return 0 @@ -78,14 +78,14 @@ func (s *sqloffsetStore) GetAckedOffset(ctx context.Context, topic string, parti } if err != nil { - return 0, fmt.Errorf("get acked offset topic=%s partition=%s: %w", topic, partitionKey, err) + return 0, fmt.Errorf("get acked offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } return offset, nil } -// UpdateAckedOffset updates the offset_acked for a topic+partition (only if new offset is greater) -func (s *sqloffsetStore) UpdateAckedOffset(ctx context.Context, topic string, partitionKey string, offset int64, consumerGroup string) (retErr error) { +// UpdateAckedOffset updates the offset_acked for a tenant+topic+partition (only if new offset is greater) +func (s *sqloffsetStore) UpdateAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string, offset int64, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "update_acked_offset", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) @@ -96,30 +96,30 @@ func (s *sqloffsetStore) UpdateAckedOffset(ctx context.Context, topic string, pa _, err := s.db.ExecContext(ctx, fmt.Sprintf(` UPDATE %s SET offset_acked = ?, updated_at = ? - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND offset_acked < ? - `, OffsetsTableName), offset, now, consumerGroup, topic, partitionKey, offset) + WHERE tenant = ? AND topic = ? AND partition_key = ? AND consumer_group = ? AND offset_acked < ? + `, OffsetsTableName), offset, now, tenant, topic, partitionKey, consumerGroup, offset) if err != nil { - return fmt.Errorf("update acked offset topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("update acked offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } return nil } // GetMinAckedOffset returns the minimum offset_acked across all consumer groups -// for a topic+partition. Returns (0, false, nil) if no offset rows exist. -func (s *sqloffsetStore) GetMinAckedOffset(ctx context.Context, topic string, partitionKey string) (_ int64, _ bool, retErr error) { +// for a tenant+topic+partition. Returns (0, false, nil) if no offset rows exist. +func (s *sqloffsetStore) GetMinAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string) (_ int64, _ bool, retErr error) { op := metrics.Begin(s.scope, "get_min_acked_offset", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() var minOffset int64 err := s.db.QueryRowContext(ctx, fmt.Sprintf(` - SELECT COALESCE(MIN(offset_acked), 0) FROM %s WHERE topic = ? AND partition_key = ? - `, OffsetsTableName), topic, partitionKey).Scan(&minOffset) + SELECT COALESCE(MIN(offset_acked), 0) FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? + `, OffsetsTableName), tenant, topic, partitionKey).Scan(&minOffset) if err != nil { - return 0, false, fmt.Errorf("query min acked offset topic=%s partition=%s: %w", topic, partitionKey, err) + return 0, false, fmt.Errorf("query min acked offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } if minOffset == 0 { @@ -131,18 +131,18 @@ func (s *sqloffsetStore) GetMinAckedOffset(ctx context.Context, topic string, pa // DeleteOffset removes one consumer group's offset row for a partition. // Idempotent — see the offsetStore interface doc. -func (s *sqloffsetStore) DeleteOffset(ctx context.Context, topic string, partitionKey string, consumerGroup string) (retErr error) { +func (s *sqloffsetStore) DeleteOffset(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "delete_offset", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic), metrics.NewTag("consumer_group", consumerGroup)) defer func() { op.Complete(retErr) }() _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - DELETE FROM %s WHERE consumer_group = ? AND topic = ? AND partition_key = ? - `, OffsetsTableName), consumerGroup, topic, partitionKey) + DELETE FROM %s WHERE tenant = ? AND topic = ? AND partition_key = ? AND consumer_group = ? + `, OffsetsTableName), tenant, topic, partitionKey, consumerGroup) if err != nil { - return fmt.Errorf("delete offset topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("delete offset tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } return nil diff --git a/platform/extension/messagequeue/mysql/offset_store_test.go b/platform/extension/messagequeue/mysql/offset_store_test.go index 112067c8f..e0c78044c 100644 --- a/platform/extension/messagequeue/mysql/offset_store_test.go +++ b/platform/extension/messagequeue/mysql/offset_store_test.go @@ -27,6 +27,7 @@ import ( const ( testConsumerGroup = "test-consumer" testSubscriberName = "test-subscriber" + testTenant = "test-tenant" ) func setupoffsetStoreTest(t *testing.T) (*sql.DB, sqlmock.Sqlmock, offsetStore) { @@ -49,10 +50,10 @@ func TestOffsetStore_Initialize(t *testing.T) { partitionKey := "part1" mock.ExpectExec("INSERT IGNORE INTO queue_offsets"). - WithArgs(testConsumerGroup, topic, partitionKey, sqlmock.AnyArg()). + WithArgs(testTenant, topic, partitionKey, testConsumerGroup, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) - err := store.Initialize(ctx, topic, partitionKey, testConsumerGroup) + err := store.Initialize(ctx, testTenant, topic, partitionKey, testConsumerGroup) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } @@ -69,7 +70,7 @@ func TestOffsetStore_GetAckedOffset(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { rows := sqlmock.NewRows([]string{"offset_acked"}).AddRow(int64(100)) mock.ExpectQuery("SELECT offset_acked FROM queue_offsets"). - WithArgs(testConsumerGroup, "test_topic", "part1"). + WithArgs(testTenant, "test_topic", "part1", testConsumerGroup). WillReturnRows(rows) }, expectedOffset: 100, @@ -79,7 +80,7 @@ func TestOffsetStore_GetAckedOffset(t *testing.T) { name: "offset not found returns zero", setup: func(mock sqlmock.Sqlmock) { mock.ExpectQuery("SELECT offset_acked FROM queue_offsets"). - WithArgs(testConsumerGroup, "test_topic", "part1"). + WithArgs(testTenant, "test_topic", "part1", testConsumerGroup). WillReturnError(sql.ErrNoRows) }, expectedOffset: 0, @@ -98,7 +99,7 @@ func TestOffsetStore_GetAckedOffset(t *testing.T) { tt.setup(mock) - offset, err := store.GetAckedOffset(ctx, topic, partitionKey, testConsumerGroup) + offset, err := store.GetAckedOffset(ctx, testTenant, topic, partitionKey, testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { @@ -120,10 +121,10 @@ func TestOffsetStore_UpdateAckedOffset(t *testing.T) { offset := int64(150) mock.ExpectExec("UPDATE queue_offsets"). - WithArgs(offset, sqlmock.AnyArg(), testConsumerGroup, topic, partitionKey, offset). + WithArgs(offset, sqlmock.AnyArg(), testTenant, topic, partitionKey, testConsumerGroup, offset). WillReturnResult(sqlmock.NewResult(0, 1)) - err := store.UpdateAckedOffset(ctx, topic, partitionKey, offset, testConsumerGroup) + err := store.UpdateAckedOffset(ctx, testTenant, topic, partitionKey, offset, testConsumerGroup) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } @@ -167,15 +168,15 @@ func TestOffsetStore_GetMinAckedOffset(t *testing.T) { if tt.queryErr { mock.ExpectQuery("SELECT COALESCE\\(MIN\\(offset_acked\\), 0\\) FROM queue_offsets"). - WithArgs("test_topic", "part-1"). + WithArgs(testTenant, "test_topic", "part-1"). WillReturnError(fmt.Errorf("db error")) } else { mock.ExpectQuery("SELECT COALESCE\\(MIN\\(offset_acked\\), 0\\) FROM queue_offsets"). - WithArgs("test_topic", "part-1"). + WithArgs(testTenant, "test_topic", "part-1"). WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(tt.minOffset)) } - offset, found, err := store.GetMinAckedOffset(context.Background(), "test_topic", "part-1") + offset, found, err := store.GetMinAckedOffset(context.Background(), testTenant, "test_topic", "part-1") if tt.wantErr { require.Error(t, err) @@ -199,7 +200,7 @@ func TestOffsetStore_DeleteOffset(t *testing.T) { name: "deletes the consumer group's offset row", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_offsets"). - WithArgs(testConsumerGroup, "test_topic", "part-1"). + WithArgs(testTenant, "test_topic", "part-1", testConsumerGroup). WillReturnResult(sqlmock.NewResult(0, 1)) }, }, @@ -207,7 +208,7 @@ func TestOffsetStore_DeleteOffset(t *testing.T) { name: "idempotent - row already gone", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_offsets"). - WithArgs(testConsumerGroup, "test_topic", "part-1"). + WithArgs(testTenant, "test_topic", "part-1", testConsumerGroup). WillReturnResult(sqlmock.NewResult(0, 0)) }, }, @@ -215,7 +216,7 @@ func TestOffsetStore_DeleteOffset(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_offsets"). - WithArgs(testConsumerGroup, "test_topic", "part-1"). + WithArgs(testTenant, "test_topic", "part-1", testConsumerGroup). WillReturnError(fmt.Errorf("db error")) }, wantErr: true, @@ -229,7 +230,7 @@ func TestOffsetStore_DeleteOffset(t *testing.T) { tt.setup(mock) - err := store.DeleteOffset(context.Background(), "test_topic", "part-1", testConsumerGroup) + err := store.DeleteOffset(context.Background(), testTenant, "test_topic", "part-1", testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { diff --git a/platform/extension/messagequeue/mysql/partition_lease_store.go b/platform/extension/messagequeue/mysql/partition_lease_store.go index d13bcb963..8d022c1b8 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store.go @@ -43,7 +43,7 @@ func newPartitionLeaseStore(db *sql.DB, logger *zap.SugaredLogger, scope tally.S } // TryAcquireLease attempts to acquire or renew a lease for a partition -func (s *sqlpartitionLeaseStore) TryAcquireLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (_ bool, retErr error) { +func (s *sqlpartitionLeaseStore) TryAcquireLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (_ bool, retErr error) { op := metrics.Begin(s.scope, "try_acquire_lease", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -52,35 +52,36 @@ func (s *sqlpartitionLeaseStore) TryAcquireLease(ctx context.Context, topic stri // Try to insert or update stale lease _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO %s (tenant, consumer_group, topic, partition_key, leased_by, leased_at, lease_renewed_at) + VALUES (?, ?, ?, ?, ?, ?, ?) ON DUPLICATE KEY UPDATE leased_by = IF(lease_renewed_at < ?, VALUES(leased_by), leased_by), leased_at = IF(lease_renewed_at < ?, VALUES(leased_at), leased_at), lease_renewed_at = IF(lease_renewed_at < ?, VALUES(lease_renewed_at), lease_renewed_at) `, PartitionLeasesTableName), - consumerGroup, topic, partitionKey, subscriberName, now, now, + tenant, consumerGroup, topic, partitionKey, subscriberName, now, now, staleThreshold, staleThreshold, staleThreshold) if err != nil { - return false, fmt.Errorf("acquire lease topic=%s partition=%s: %w", topic, partitionKey, err) + return false, fmt.Errorf("acquire lease tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } // Check if we own the lease var owner string err = s.db.QueryRowContext(ctx, fmt.Sprintf(` SELECT leased_by FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? - `, PartitionLeasesTableName), consumerGroup, topic, partitionKey).Scan(&owner) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? + `, PartitionLeasesTableName), tenant, consumerGroup, topic, partitionKey).Scan(&owner) if err != nil { - return false, fmt.Errorf("check lease ownership topic=%s partition=%s: %w", topic, partitionKey, err) + return false, fmt.Errorf("check lease ownership tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } acquired := owner == subscriberName if acquired { metrics.NamedCounter(s.scope, "try_acquire_lease", "acquired", 1, metrics.NewTag("topic", topic)) s.logger.Debugw("acquired lease", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, ) @@ -92,7 +93,7 @@ func (s *sqlpartitionLeaseStore) TryAcquireLease(ctx context.Context, topic stri } // RenewLease renews the lease for a partition owned by this worker -func (s *sqlpartitionLeaseStore) RenewLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (retErr error) { +func (s *sqlpartitionLeaseStore) RenewLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (retErr error) { op := metrics.Begin(s.scope, "renew_lease", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -101,16 +102,16 @@ func (s *sqlpartitionLeaseStore) RenewLease(ctx context.Context, topic string, p result, err := s.db.ExecContext(ctx, fmt.Sprintf(` UPDATE %s SET lease_renewed_at = ? - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND leased_by = ? - `, PartitionLeasesTableName), now, consumerGroup, topic, partitionKey, subscriberName) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND leased_by = ? + `, PartitionLeasesTableName), now, tenant, consumerGroup, topic, partitionKey, subscriberName) if err != nil { - return fmt.Errorf("renew lease topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("renew lease tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } rows, err := result.RowsAffected() if err != nil { - return fmt.Errorf("check renewal result topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("check renewal result tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } if rows == 0 { @@ -118,6 +119,7 @@ func (s *sqlpartitionLeaseStore) RenewLease(ctx context.Context, topic string, p } s.logger.Debugw("renewed lease", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, ) @@ -126,17 +128,17 @@ func (s *sqlpartitionLeaseStore) RenewLease(ctx context.Context, topic string, p } // ReleaseLease releases the lease for a partition owned by this worker -func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "release_lease", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() result, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s - WHERE consumer_group = ? AND topic = ? AND partition_key = ? AND leased_by = ? - `, PartitionLeasesTableName), consumerGroup, topic, partitionKey, subscriberName) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND partition_key = ? AND leased_by = ? + `, PartitionLeasesTableName), tenant, consumerGroup, topic, partitionKey, subscriberName) if err != nil { - return fmt.Errorf("release lease topic=%s partition=%s: %w", topic, partitionKey, err) + return fmt.Errorf("release lease tenant=%s topic=%s partition=%s: %w", tenant, topic, partitionKey, err) } // RowsAffected error is swallowed because the DELETE query itself succeeded. @@ -145,6 +147,7 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, topic string, rows, err := result.RowsAffected() if err != nil { s.logger.Warnw("failed to get rows affected after release lease", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, logError, err, @@ -152,6 +155,7 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, topic string, } if rows > 0 { s.logger.Debugw("released lease", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, ) @@ -161,17 +165,17 @@ func (s *sqlpartitionLeaseStore) ReleaseLease(ctx context.Context, topic string, } // GetLeasedPartitions returns all partitions currently leased by this worker -func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string) (_ []string, retErr error) { +func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) (_ []string, retErr error) { op := metrics.Begin(s.scope, "get_leased_partitions", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` SELECT partition_key FROM %s - WHERE consumer_group = ? AND topic = ? AND leased_by = ? - `, PartitionLeasesTableName), consumerGroup, topic, subscriberName) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND leased_by = ? + `, PartitionLeasesTableName), tenant, consumerGroup, topic, subscriberName) if err != nil { - return nil, fmt.Errorf("get leased partitions topic=%s: %w", topic, err) + return nil, fmt.Errorf("get leased partitions tenant=%s topic=%s: %w", tenant, topic, err) } defer rows.Close() @@ -179,16 +183,17 @@ func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic for rows.Next() { var partition string if err := rows.Scan(&partition); err != nil { - return nil, fmt.Errorf("scan partition topic=%s: %w", topic, err) + return nil, fmt.Errorf("scan partition tenant=%s topic=%s: %w", tenant, topic, err) } partitions = append(partitions, partition) } if err := rows.Err(); err != nil { - return nil, fmt.Errorf("row iteration topic=%s: %w", topic, err) + return nil, fmt.Errorf("row iteration tenant=%s topic=%s: %w", tenant, topic, err) } s.logger.Debugw("retrieved leased partitions", + logTenant, tenant, logTopic, topic, "count", len(partitions), ) @@ -197,18 +202,18 @@ func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic } // GetAllLeases returns the lease row for every partition currently leased -// under (topic, consumerGroup) by any subscriber. -func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, topic string, consumerGroup string) (_ []leaseInfo, retErr error) { +// under (tenant, topic, consumerGroup) by any subscriber. +func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, tenant string, topic string, consumerGroup string) (_ []leaseInfo, retErr error) { op := metrics.Begin(s.scope, "get_all_leases", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` SELECT partition_key, leased_by, lease_renewed_at FROM %s - WHERE consumer_group = ? AND topic = ? - `, PartitionLeasesTableName), consumerGroup, topic) + WHERE tenant = ? AND consumer_group = ? AND topic = ? + `, PartitionLeasesTableName), tenant, consumerGroup, topic) if err != nil { - return nil, fmt.Errorf("get all leases topic=%s: %w", topic, err) + return nil, fmt.Errorf("get all leases tenant=%s topic=%s: %w", tenant, topic, err) } defer rows.Close() @@ -216,13 +221,13 @@ func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, topic string, for rows.Next() { var lease leaseInfo if err := rows.Scan(&lease.PartitionKey, &lease.LeasedBy, &lease.LeaseRenewedAt); err != nil { - return nil, fmt.Errorf("scan lease topic=%s: %w", topic, err) + return nil, fmt.Errorf("scan lease tenant=%s topic=%s: %w", tenant, topic, err) } leases = append(leases, lease) } if err := rows.Err(); err != nil { - return nil, fmt.Errorf("row iteration topic=%s: %w", topic, err) + return nil, fmt.Errorf("row iteration tenant=%s topic=%s: %w", tenant, topic, err) } return leases, nil @@ -230,7 +235,7 @@ func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, topic string, // PurgeStale deletes lease rows not renewed within olderThanMs. See the // partitionLeaseStore interface doc. -func (s *sqlpartitionLeaseStore) PurgeStale(ctx context.Context, topic string, consumerGroup string, olderThanMs int64) (retErr error) { +func (s *sqlpartitionLeaseStore) PurgeStale(ctx context.Context, tenant string, topic string, consumerGroup string, olderThanMs int64) (retErr error) { op := metrics.Begin(s.scope, "purge_stale", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -238,11 +243,11 @@ func (s *sqlpartitionLeaseStore) PurgeStale(ctx context.Context, topic string, c result, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s - WHERE consumer_group = ? AND topic = ? AND lease_renewed_at < ? - `, PartitionLeasesTableName), consumerGroup, topic, threshold) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND lease_renewed_at < ? + `, PartitionLeasesTableName), tenant, consumerGroup, topic, threshold) if err != nil { - return fmt.Errorf("failed to purge stale leases: %w", err) + return fmt.Errorf("failed to purge stale leases tenant=%s topic=%s: %w", tenant, topic, err) } // RowsAffected error is swallowed because the DELETE itself succeeded; @@ -250,6 +255,7 @@ func (s *sqlpartitionLeaseStore) PurgeStale(ctx context.Context, topic string, c if deleted, err := result.RowsAffected(); err == nil && deleted > 0 { metrics.NamedCounter(s.scope, "purge_stale", "rows_deleted", deleted, metrics.NewTag("topic", topic)) s.logger.Debugw("purged stale leases", + logTenant, tenant, logTopic, topic, "deleted", deleted, ) @@ -271,7 +277,7 @@ func (s *sqlpartitionLeaseStore) PurgeStale(ctx context.Context, topic string, c // write on a contended lease row. The classification is advisory (a lease // can expire or renew between the read and the attempt); TryAcquireLease // remains the atomic arbiter. -func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string, leaseDurationMs int64, maxPartitions int) (_ int, _ []string, retErr error) { +func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string, leaseDurationMs int64, maxPartitions int) (_ int, _ []string, retErr error) { op := metrics.Begin(s.scope, "discover_and_acquire", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -281,10 +287,10 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex // making them permanently unprocessable. The maxPartitions cap only limits how // many leases this subscriber acquires, not how many partitions are visible. rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` - SELECT DISTINCT partition_key FROM %s WHERE topic = ? ORDER BY partition_key - `, MessagesTableName), topic) + SELECT DISTINCT partition_key FROM %s WHERE tenant = ? AND topic = ? ORDER BY partition_key + `, MessagesTableName), tenant, topic) if err != nil { - return 0, nil, fmt.Errorf("discover partitions topic=%s: %w", topic, err) + return 0, nil, fmt.Errorf("discover partitions tenant=%s topic=%s: %w", tenant, topic, err) } defer rows.Close() @@ -292,16 +298,17 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex for rows.Next() { var partitionKey string if err := rows.Scan(&partitionKey); err != nil { - return 0, nil, fmt.Errorf("scan partition key topic=%s: %w", topic, err) + return 0, nil, fmt.Errorf("scan partition key tenant=%s topic=%s: %w", tenant, topic, err) } partitions = append(partitions, partitionKey) } if err := rows.Err(); err != nil { - return 0, nil, fmt.Errorf("row iteration topic=%s: %w", topic, err) + return 0, nil, fmt.Errorf("row iteration tenant=%s topic=%s: %w", tenant, topic, err) } s.logger.Debugw("discovered partitions", + logTenant, tenant, logTopic, topic, "count", len(partitions), ) @@ -309,9 +316,9 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex // One read of every lease row classifies the discovered partitions: // self-owned (count toward the cap, no re-probe), validly held by // another subscriber (skip), or unleased/stale (acquisition candidates). - allLeases, err := s.GetAllLeases(ctx, topic, consumerGroup) + allLeases, err := s.GetAllLeases(ctx, tenant, topic, consumerGroup) if err != nil { - return 0, nil, fmt.Errorf("get all leases for acquisition topic=%s: %w", topic, err) + return 0, nil, fmt.Errorf("get all leases for acquisition tenant=%s topic=%s: %w", tenant, topic, err) } staleThreshold := currentTimeMillis() - leaseDurationMs ownedCount := 0 @@ -347,6 +354,7 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex // Enforce maxPartitions cap using local count if maxPartitions > 0 && ownedCount >= maxPartitions { s.logger.Debugw("reached max partitions cap, stopping acquisition", + logTenant, tenant, logTopic, topic, "max_partitions", maxPartitions, "owned_count", ownedCount, @@ -354,12 +362,13 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex break } - acquired, err := s.TryAcquireLease(ctx, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) + acquired, err := s.TryAcquireLease(ctx, tenant, topic, partitionKey, subscriberName, consumerGroup, leaseDurationMs) if err != nil { // Per-partition error is swallowed because one partition's DB failure // should not prevent acquiring leases for other partitions. The failed // partition is retried on the next discovery cycle. s.logger.Errorw("failed to acquire lease for partition", + logTenant, tenant, logTopic, topic, logPartitionKey, partitionKey, logError, err, @@ -376,6 +385,7 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex metrics.NamedCounter(s.scope, "discover_and_acquire", "partitions_acquired", int64(acquiredCount), metrics.NewTag("topic", topic)) metrics.NamedCounter(s.scope, "discover_and_acquire", "lease_aware_skipped", int64(skippedCount), metrics.NewTag("topic", topic)) s.logger.Debugw("completed partition discovery and acquisition", + logTenant, tenant, logTopic, topic, "discovered_count", len(partitions), "acquired_count", acquiredCount, diff --git a/platform/extension/messagequeue/mysql/partition_lease_store_test.go b/platform/extension/messagequeue/mysql/partition_lease_store_test.go index 1a6a83e9f..17f278049 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store_test.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store_test.go @@ -51,11 +51,11 @@ func TestPartitionLeaseStore_TryAcquireLease(t *testing.T) { name: "successfully acquire lease", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("INSERT INTO queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", "part1", testSubscriberName, sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName, sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) rows := sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName) mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", "part1"). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1"). WillReturnRows(rows) }, acquired: true, @@ -68,7 +68,7 @@ func TestPartitionLeaseStore_TryAcquireLease(t *testing.T) { WillReturnResult(sqlmock.NewResult(1, 1)) rows := sqlmock.NewRows([]string{"leased_by"}).AddRow("other-worker") mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", "part1"). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1"). WillReturnRows(rows) }, acquired: false, @@ -87,7 +87,7 @@ func TestPartitionLeaseStore_TryAcquireLease(t *testing.T) { tt.setup(mock) - acquired, err := store.TryAcquireLease(ctx, topic, partitionKey, testSubscriberName, testConsumerGroup, testLeaseDurationMs) + acquired, err := store.TryAcquireLease(ctx, testTenant, topic, partitionKey, testSubscriberName, testConsumerGroup, testLeaseDurationMs) if tt.wantErr { require.Error(t, err) } else { @@ -109,7 +109,7 @@ func TestPartitionLeaseStore_RenewLease(t *testing.T) { name: "successfully renew lease", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE queue_partition_leases"). - WithArgs(sqlmock.AnyArg(), testConsumerGroup, "test_topic", "part1", testSubscriberName). + WithArgs(sqlmock.AnyArg(), testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) }, wantErr: false, @@ -118,7 +118,7 @@ func TestPartitionLeaseStore_RenewLease(t *testing.T) { name: "lease not owned", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE queue_partition_leases"). - WithArgs(sqlmock.AnyArg(), testConsumerGroup, "test_topic", "part1", testSubscriberName). + WithArgs(sqlmock.AnyArg(), testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 0)) }, wantErr: true, @@ -136,7 +136,7 @@ func TestPartitionLeaseStore_RenewLease(t *testing.T) { tt.setup(mock) - err := store.RenewLease(ctx, topic, partitionKey, testSubscriberName, testConsumerGroup, testLeaseDurationMs) + err := store.RenewLease(ctx, testTenant, topic, partitionKey, testSubscriberName, testConsumerGroup, testLeaseDurationMs) if tt.wantErr { require.Error(t, err) } else { @@ -157,7 +157,7 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { name: "successfully release lease", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", "part1", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) }, wantErr: false, @@ -166,7 +166,7 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { name: "idempotent - already released", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", "part1", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", "part1", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 0)) }, wantErr: false, @@ -184,7 +184,7 @@ func TestPartitionLeaseStore_ReleaseLease(t *testing.T) { tt.setup(mock) - err := store.ReleaseLease(ctx, topic, partitionKey, testSubscriberName, testConsumerGroup) + err := store.ReleaseLease(ctx, testTenant, topic, partitionKey, testSubscriberName, testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { @@ -208,10 +208,10 @@ func TestPartitionLeaseStore_GetLeasedPartitions(t *testing.T) { AddRow("part3") mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, topic, testSubscriberName). + WithArgs(testTenant, testConsumerGroup, topic, testSubscriberName). WillReturnRows(rows) - partitions, err := store.GetLeasedPartitions(ctx, topic, testSubscriberName, testConsumerGroup) + partitions, err := store.GetLeasedPartitions(ctx, testTenant, topic, testSubscriberName, testConsumerGroup) require.NoError(t, err) require.Len(t, partitions, 3) require.Equal(t, []string{"part1", "part2", "part3"}, partitions) @@ -231,7 +231,7 @@ func TestPartitionLeaseStore_GetAllLeases(t *testing.T) { AddRow("part1", testSubscriberName, int64(1000)). AddRow("part2", "other-worker", int64(2000)) mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(rows) }, want: []leaseInfo{ @@ -243,7 +243,7 @@ func TestPartitionLeaseStore_GetAllLeases(t *testing.T) { name: "no leases returns empty", setup: func(mock sqlmock.Sqlmock) { mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows([]string{"partition_key", "leased_by", "lease_renewed_at"})) }, want: nil, @@ -257,7 +257,7 @@ func TestPartitionLeaseStore_GetAllLeases(t *testing.T) { tt.setup(mock) - leases, err := store.GetAllLeases(context.Background(), "test_topic", testConsumerGroup) + leases, err := store.GetAllLeases(context.Background(), testTenant, "test_topic", testConsumerGroup) require.NoError(t, err) require.Equal(t, tt.want, leases) require.NoError(t, mock.ExpectationsWereMet()) @@ -277,7 +277,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { rows.AddRow(pk) } mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). + WithArgs(testTenant, "test_topic"). WillReturnRows(rows) } @@ -302,7 +302,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1", "part2") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns). AddRow("part2", "other-worker", freshMs)) // Only unleased part1 is attempted; part2's fresh lease is @@ -317,7 +317,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns). AddRow("part1", "other-worker", staleMs)) expectAcquire(mock, testSubscriberName) @@ -330,7 +330,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1", "part2") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns). AddRow("part1", testSubscriberName, freshMs)) // Only part2 is attempted; renewal of part1 is the lease @@ -345,7 +345,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1", "part2", "part3") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns)) // part1 and part2 acquired; part3 never attempted at the cap. expectAcquire(mock, testSubscriberName) @@ -359,7 +359,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1", "part2", "part3") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns). AddRow("existing1", testSubscriberName, freshMs). AddRow("existing2", testSubscriberName, freshMs)) @@ -374,7 +374,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1", "part2") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns). AddRow("existing1", testSubscriberName, freshMs). AddRow("existing2", testSubscriberName, freshMs)) @@ -388,7 +388,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { expectDiscover(mock, "part1") mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic"). + WithArgs(testTenant, testConsumerGroup, "test_topic"). WillReturnRows(sqlmock.NewRows(leaseColumns)) // Attempted while unleased, but another subscriber won the // atomic acquire between the read and the write. @@ -405,7 +405,7 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { tt.setup(mock) - acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(context.Background(), "test_topic", testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) + acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(context.Background(), testTenant, "test_topic", testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) require.NoError(t, err) require.Equal(t, tt.wantAcquired, acquired) require.NotNil(t, discoveredPartitions) @@ -424,7 +424,7 @@ func TestPartitionLeaseStore_PurgeStale(t *testing.T) { name: "deletes rows older than threshold", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 2)) }, }, @@ -432,7 +432,7 @@ func TestPartitionLeaseStore_PurgeStale(t *testing.T) { name: "no stale rows is a no-op", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 0)) }, }, @@ -440,7 +440,7 @@ func TestPartitionLeaseStore_PurgeStale(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnError(fmt.Errorf("db error")) }, wantErr: true, @@ -454,7 +454,7 @@ func TestPartitionLeaseStore_PurgeStale(t *testing.T) { tt.setup(mock) - err := store.PurgeStale(context.Background(), "test_topic", testConsumerGroup, 300_000) + err := store.PurgeStale(context.Background(), testTenant, "test_topic", testConsumerGroup, 300_000) if tt.wantErr { require.Error(t, err) } else { diff --git a/platform/extension/messagequeue/mysql/publisher.go b/platform/extension/messagequeue/mysql/publisher.go index fc9ecfe34..a587b23a5 100644 --- a/platform/extension/messagequeue/mysql/publisher.go +++ b/platform/extension/messagequeue/mysql/publisher.go @@ -57,7 +57,36 @@ func (p *publisher) Publish(ctx context.Context, topic string, message entityque return ErrPublisherClosed } - if err := p.messageStore.Insert(ctx, topic, []entityqueue.Message{message}); err != nil { + if message.Tenant == "" { + return fmt.Errorf("publish: message tenant is required") + } + if err := entityqueue.ValidateTenantMetadata(message); err != nil { + return fmt.Errorf("publish: %w", err) + } + for _, identifier := range []struct { + name string + value string + }{ + {name: "tenant", value: message.Tenant}, + {name: "topic", value: topic}, + } { + if err := validateASCIIIdentifier(identifier.name, identifier.value); err != nil { + return fmt.Errorf("publish: %w", err) + } + } + for _, identifier := range []struct { + name string + value string + }{ + {name: "message ID", value: message.ID}, + {name: "partition key", value: message.PartitionKey}, + } { + if err := validateTextIdentifier(identifier.name, identifier.value); err != nil { + return fmt.Errorf("publish: %w", err) + } + } + + if err := p.messageStore.Insert(ctx, message.Tenant, topic, []entityqueue.Message{message}); err != nil { return fmt.Errorf("publish message store insert error: %w", err) } diff --git a/platform/extension/messagequeue/mysql/publisher_test.go b/platform/extension/messagequeue/mysql/publisher_test.go index 115c72313..b59ca0bff 100644 --- a/platform/extension/messagequeue/mysql/publisher_test.go +++ b/platform/extension/messagequeue/mysql/publisher_test.go @@ -18,6 +18,7 @@ import ( "context" "errors" "fmt" + "strings" "testing" "github.com/stretchr/testify/require" @@ -43,6 +44,8 @@ func setupPublisherTest(t *testing.T, mockStore *MockmessageStore) extqueue.Publ } func TestPublisher_Publish(t *testing.T) { + overlong := strings.Repeat("x", maxIdentifierLength+1) + noStoreCall := func(*MockmessageStore) {} tests := []struct { name string topic string @@ -54,24 +57,24 @@ func TestPublisher_Publish(t *testing.T) { name: "publish single message", topic: "test_topic", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "test_topic", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "test_topic", gomock.Any()).Return(nil).Times(1) }, }, { name: "publish multiple messages", topic: "multi_topic", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, - {ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, - {ID: "msg3", Payload: []byte("p3"), PartitionKey: "part2", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg3", Payload: []byte("p3"), PartitionKey: "part2", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "multi_topic", gomock.Any()).Return(nil).Times(3) + m.EXPECT().Insert(gomock.Any(), testTenant, "multi_topic", gomock.Any()).Return(nil).Times(3) }, }, { @@ -88,6 +91,7 @@ func TestPublisher_Publish(t *testing.T) { topic: "metadata_topic", messages: []entityqueue.Message{ { + Tenant: testTenant, ID: "msg_meta", Payload: []byte("payload"), PartitionKey: "part1", @@ -97,20 +101,82 @@ func TestPublisher_Publish(t *testing.T) { }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "metadata_topic", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "metadata_topic", gomock.Any()).Return(nil).Times(1) }, }, { name: "publish with valid topic name - hyphens", topic: "topic-with-dash", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("p"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("p"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "topic-with-dash", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "topic-with-dash", gomock.Any()).Return(nil).Times(1) }, }, + { + name: "rejects overlong tenant", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: overlong, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong topic", + topic: overlong, + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong message ID", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: overlong, PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong partition key", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: overlong}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "accepts UTF-8 message ID and partition key", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg-é", PartitionKey: strings.Repeat("é", maxIdentifierLength)}}, + setupMock: func(m *MockmessageStore) { + m.EXPECT().Insert(gomock.Any(), testTenant, "test_topic", gomock.Any()).Return(nil) + }, + }, + { + name: "rejects non-ASCII tenant", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: "tenant-é", ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects non-ASCII topic", + topic: "topic-é", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects conflicting queue metadata", + topic: "test_topic", + messages: []entityqueue.Message{{ + Tenant: testTenant, + ID: "msg1", + PartitionKey: "part1", + Metadata: map[string]string{entityqueue.MetadataKeyQueueName: "other-tenant"}, + }}, + wantErr: true, + setupMock: noStoreCall, + }, } for _, tt := range tests { @@ -181,7 +247,7 @@ func TestPublisher_PublishMetrics(t *testing.T) { defer ctrl.Finish() mockStore := NewMockmessageStore(ctrl) - mockStore.EXPECT().Insert(gomock.Any(), "metrics_test", gomock.Any()).Return(nil).Times(2) + mockStore.EXPECT().Insert(gomock.Any(), testTenant, "metrics_test", gomock.Any()).Return(nil).Times(2) pub := setupPublisherTest(t, mockStore) @@ -190,8 +256,8 @@ func TestPublisher_PublishMetrics(t *testing.T) { // Publish some messages messages := []entityqueue.Message{ - {ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, - {ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, } for _, msg := range messages { @@ -211,7 +277,7 @@ func TestPublisher_ConcurrentPublish(t *testing.T) { const messagesPerGoroutine = 5 mockStore := NewMockmessageStore(ctrl) - mockStore.EXPECT().Insert(gomock.Any(), "concurrent_topic", gomock.Any()).Return(nil).Times(numGoroutines * messagesPerGoroutine) + mockStore.EXPECT().Insert(gomock.Any(), testTenant, "concurrent_topic", gomock.Any()).Return(nil).Times(numGoroutines * messagesPerGoroutine) pub := setupPublisherTest(t, mockStore) @@ -224,6 +290,7 @@ func TestPublisher_ConcurrentPublish(t *testing.T) { go func(id int) { for j := 0; j < messagesPerGoroutine; j++ { msg := entityqueue.Message{ + Tenant: testTenant, ID: fmt.Sprintf("msg_%d_%d", id, j), Payload: []byte(fmt.Sprintf("payload_%d_%d", id, j)), PartitionKey: fmt.Sprintf("part_%d", id), @@ -246,7 +313,7 @@ func TestPublisher_PublishContextCancellation(t *testing.T) { defer ctrl.Finish() mockStore := NewMockmessageStore(ctrl) - mockStore.EXPECT().Insert(gomock.Any(), "test_topic", gomock.Any()).Return(context.Canceled).Times(1) + mockStore.EXPECT().Insert(gomock.Any(), testTenant, "test_topic", gomock.Any()).Return(context.Canceled).Times(1) pub := setupPublisherTest(t, mockStore) @@ -255,6 +322,7 @@ func TestPublisher_PublishContextCancellation(t *testing.T) { cancel() msg := entityqueue.NewMessage("msg1", []byte("payload"), "part1", nil) + msg.Tenant = testTenant // Should fail with context cancelled error err := pub.Publish(ctx, "test_topic", msg) diff --git a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql index 27625cdc8..e561aef2c 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql @@ -8,14 +8,17 @@ -- acked = FALSE, invisible_until <= now → ready for (re-)delivery CREATE TABLE IF NOT EXISTS queue_delivery_state ( + -- tenant is the shard isolation identity + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, + -- Consumer group this delivery state belongs to - consumer_group VARCHAR(255) NOT NULL, + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic of the message - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition key of the message - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Offset of the message in the immutable log message_offset BIGINT UNSIGNED NOT NULL, @@ -24,18 +27,13 @@ CREATE TABLE IF NOT EXISTS queue_delivery_state ( acked BOOLEAN NOT NULL DEFAULT FALSE, -- Visibility timeout (epoch milliseconds) - -- Only meaningful when acked = FALSE. - -- Future timestamp = in-flight or nack delay, 0/past = ready for delivery. invisible_until BIGINT UNSIGNED NOT NULL DEFAULT 0, -- Number of times this message has been redelivered to this consumer group retry_count INT UNSIGNED NOT NULL DEFAULT 0, -- Whether the last delivery was postponed (deliberate wait, not a failure). - -- While set and invisible, the message is a barrier: its partition is not - -- consumed past it. The next delivery is exempt from the retry_count - -- increment and clears the flag. postponed BOOLEAN NOT NULL DEFAULT FALSE, - PRIMARY KEY (consumer_group, topic, partition_key, message_offset) + PRIMARY KEY (tenant, consumer_group, topic, partition_key, message_offset) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/schema/queue_messages.sql b/platform/extension/messagequeue/mysql/schema/queue_messages.sql index a3a655365..87a178fec 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_messages.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_messages.sql @@ -1,21 +1,23 @@ -- MESSAGES TABLE (Immutable Log) --- Single table for all topics. Partition key determines distribution across workers. +-- Single table for all topics. tenant is the Vitess vindex; partition_key orders work within a tenant. -- Messages are append-only; per-consumer-group delivery tracking is in queue_delivery_state. --- Example: topic="merge_queue", partition_key="uber/cadence" +-- Example: tenant="monorepo/main", topic="merge_queue", partition_key="uber/cadence" CREATE TABLE IF NOT EXISTS queue_messages ( - -- Auto-incrementing global offset for ordering - offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY, + -- tenant is the shard isolation identity (SubmitQueue maps queueName here at wiring) + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, - -- Topic identifies the queue type - topic VARCHAR(255) NOT NULL, + -- Topic identifies the pipeline stage + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, - -- Partition key for distributing work across workers - -- Example: repo ID, user ID, tenant ID - partition_key VARCHAR(255) NOT NULL, + -- Partition key for distributing work across workers within a tenant + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, + + -- Auto-incrementing offset for ordering within (tenant, topic, partition_key) + offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, -- Message identification - id VARCHAR(255) NOT NULL, + id VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Message data payload BLOB NOT NULL, @@ -27,26 +29,16 @@ CREATE TABLE IF NOT EXISTS queue_messages ( -- DLQ-specific fields (0/"" for normal messages, populated for DLQ messages) failed_at BIGINT UNSIGNED NOT NULL, - -- failure_count stores how many times the message failed on the ORIGINAL topic before moving to DLQ failure_count INT UNSIGNED NOT NULL, last_error TEXT NOT NULL, - original_topic VARCHAR(255) NOT NULL, - -- failure_detail holds the structured half of the failure: which entities it - -- was about, plus free-form context. last_error keeps the human-readable - -- message, so this column never has to be decoded to read one, and a plain - -- SELECT last_error stays useful. - -- - -- NULL rather than an empty-string sentinel like its neighbours: a JSON - -- column rejects '' as invalid, and NULL is the honest reading of a failure - -- that recorded no structure — including every row written before this - -- column existed, and the retry-limit backstop, which has none to record. + original_topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, failure_detail JSON, - -- Supports: SELECT ... WHERE topic=? AND partition_key=? AND offset > ? ORDER BY offset - -- Used by subscribers to poll for messages within their assigned partition - INDEX idx_topic_partition_offset (topic, partition_key, offset), + PRIMARY KEY (tenant, topic, partition_key, offset), -- Supports: INSERT ... ON DUPLICATE KEY to enforce idempotent publishes - -- Also enables efficient lookups for message updates/deletes by ID - UNIQUE KEY idx_topic_partition_id (topic, partition_key, id) + UNIQUE KEY idx_tenant_topic_partition_id (tenant, topic, partition_key, id), + + -- InnoDB requires AUTO_INCREMENT column to be leftmost on some index + KEY idx_offset (offset) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql index 0a1717693..45e8439ff 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql @@ -1,20 +1,19 @@ -- CONSUMER OFFSETS TABLE --- Tracks consumption progress per consumer group + topic + partition. +-- Tracks consumption progress per consumer group + tenant + topic + partition. -- Each partition has independent offset tracking for crash recovery. --- --- The primary key (consumer_group, topic, partition_key) serves as the main --- lookup index for all queries in offsetStore. No additional indexes are needed --- because all queries filter by the full primary key or a left prefix of it. CREATE TABLE IF NOT EXISTS queue_offsets ( - -- Consumer group consuming the topic - consumer_group VARCHAR(255) NOT NULL, + -- tenant is the shard isolation identity + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic being consumed - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition being consumed - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, + + -- Consumer group consuming the topic + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Last offset that was successfully acked for this partition offset_acked BIGINT UNSIGNED NOT NULL, @@ -22,13 +21,5 @@ CREATE TABLE IF NOT EXISTS queue_offsets ( -- Last update timestamp (epoch milliseconds) updated_at BIGINT UNSIGNED NOT NULL, - -- Primary key ensures each consumer group has one offset per topic/partition. - -- Supports: INSERT ... ON DUPLICATE KEY UPDATE for idempotent offset updates. - -- Also enables efficient lookups: SELECT ... WHERE consumer_group=? AND topic=? AND partition_key=? - -- Left-prefix covers: SELECT ... WHERE consumer_group=? (all offsets for a group) - PRIMARY KEY (consumer_group, topic, partition_key), - - -- Supports: SELECT ... WHERE topic=? - -- Used for querying all consumer groups consuming a specific topic - INDEX idx_topic (topic) + PRIMARY KEY (tenant, topic, partition_key, consumer_group) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql index e1015764e..47f1eebfe 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql @@ -1,38 +1,28 @@ -- PARTITION LEASES TABLE -- Tracks which worker has leased which partition for exclusive processing. -- Workers must renew leases to maintain ownership; stale leases can be stolen. --- --- The primary key (consumer_group, topic, partition_key) serves as the main --- lookup index. Queries by leased_by always include consumer_group and topic, --- so the primary key's left-prefix is sufficient. CREATE TABLE IF NOT EXISTS queue_partition_leases ( + -- tenant is the shard isolation identity + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, + -- Consumer group (e.g., "orchestrator") - consumer_group VARCHAR(255) NOT NULL, + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic being consumed - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition that is leased - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Worker that owns the lease (e.g., "worker-1") - leased_by VARCHAR(255) NOT NULL, + leased_by VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- When lease was acquired (epoch milliseconds) leased_at BIGINT UNSIGNED NOT NULL, -- Last lease renewal timestamp (epoch milliseconds) - -- Used to detect stale leases lease_renewed_at BIGINT UNSIGNED NOT NULL, - -- Primary key ensures each partition can only be leased by one worker per consumer group. - -- Supports: INSERT ... ON DUPLICATE KEY UPDATE for lease acquisition and renewal. - -- Also enables efficient lookups: SELECT ... WHERE consumer_group=? AND topic=? AND partition_key=? - -- Left-prefix covers: SELECT ... WHERE consumer_group=? AND topic=? AND leased_by=? - PRIMARY KEY (consumer_group, topic, partition_key), - - -- Supports: SELECT ... WHERE lease_renewed_at0 means deregistered at that time. deregistered_at BIGINT UNSIGNED NOT NULL, - PRIMARY KEY (consumer_group, topic, subscriber_name) + PRIMARY KEY (tenant, consumer_group, topic, subscriber_name) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/sql.go b/platform/extension/messagequeue/mysql/sql.go index de4e3b8db..62fa18b95 100644 --- a/platform/extension/messagequeue/mysql/sql.go +++ b/platform/extension/messagequeue/mysql/sql.go @@ -57,10 +57,19 @@ type Params struct { // OnSignal receives typed subscriber lifecycle signals (HookSignal). // Nil in production; used by integration tests for event-driven waits. OnSignal chan HookSignal + + // Tenants is the configured shard isolation list. Empty permits publishing + // but causes Subscribe to return ErrInvalidConfig. + Tenants []string } // NewQueue creates a new SQL-based queue func NewQueue(params Params) (extqueue.Queue, error) { + tenants, err := normalizeTenants(params.Tenants) + if err != nil { + return nil, fmt.Errorf("invalid queue tenants: %w", err) + } + // Test connection if err := params.DB.Ping(); err != nil { return nil, fmt.Errorf("failed to ping database: %w", err) @@ -102,6 +111,7 @@ func NewQueue(params Params) (extqueue.Queue, error) { leaseStore, heartbeatStore, deliveryStateStore, + tenants, ) subscriber.OnSignal = params.OnSignal diff --git a/platform/extension/messagequeue/mysql/sql_test.go b/platform/extension/messagequeue/mysql/sql_test.go index 6eadafbd6..216f3315a 100644 --- a/platform/extension/messagequeue/mysql/sql_test.go +++ b/platform/extension/messagequeue/mysql/sql_test.go @@ -108,6 +108,53 @@ func TestNewQueue(t *testing.T) { require.Error(t, err) assert.Nil(t, q) }) + t.Run("normalizes and copies tenants", func(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.MonitorPingsOption(true)) + require.NoError(t, err) + defer db.Close() + + mock.ExpectPing() + tenants := []string{" tenant-a ", "", "tenant-b", "tenant-a"} + q, err := NewQueue(Params{ + DB: db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NewTestScope("test", nil), + Tenants: tenants, + }) + require.NoError(t, err) + tenants[0] = "mutated" + + impl := q.(*queueImpl) + sub := impl.subscriber.(*subscriber) + assert.Equal(t, []string{"tenant-a", "tenant-b"}, sub.tenants) + assert.NoError(t, q.Close()) + require.NoError(t, mock.ExpectationsWereMet()) + }) + t.Run("rejects invalid tenants before opening the queue", func(t *testing.T) { + for _, tt := range []struct { + name string + tenant string + }{ + {name: "non-ASCII", tenant: "tenant-é"}, + {name: "NUL", tenant: "tenant-\x00a"}, + } { + t.Run(tt.name, func(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.MonitorPingsOption(true)) + require.NoError(t, err) + mock.ExpectClose() + + q, err := NewQueue(Params{ + DB: db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NewTestScope("test", nil), + Tenants: []string{tt.tenant}, + }) + require.Error(t, err) + assert.Nil(t, q) + require.NoError(t, db.Close()) + }) + } + }) } func TestQueue_Publisher(t *testing.T) { diff --git a/platform/extension/messagequeue/mysql/stores.go b/platform/extension/messagequeue/mysql/stores.go index 199eeb314..2689f25a2 100644 --- a/platform/extension/messagequeue/mysql/stores.go +++ b/platform/extension/messagequeue/mysql/stores.go @@ -34,6 +34,8 @@ const ( // messageRow represents a row from the messages table (internal use only) type messageRow struct { + // Tenant is the shard isolation identity + Tenant string // Offset is the auto-incrementing sequence number for message ordering within a partition Offset int64 // ID is the unique message identifier @@ -63,57 +65,42 @@ type messageRow struct { // messageStore handles message table operations (internal use only) type messageStore interface { // Insert inserts messages into the topic table. - Insert(ctx context.Context, topic string, messages []entityqueue.Message) error + Insert(ctx context.Context, tenant string, topic string, messages []entityqueue.Message) error - // Delete deletes a message by topic, partition key, and ID - Delete(ctx context.Context, topic string, partitionKey string, messageID string) error + // Delete deletes a message by tenant, topic, partition key, and ID + Delete(ctx context.Context, tenant string, topic string, partitionKey string, messageID string) error // FetchByOffset fetches messages with offset > currentOffset for a specific partition. - // Per-consumer-group visibility is handled by the deliveryStateStore. - FetchByOffset(ctx context.Context, topic string, partitionKey string, currentOffset int64, limit int) ([]messageRow, error) + FetchByOffset(ctx context.Context, tenant string, topic string, partitionKey string, currentOffset int64, limit int) ([]messageRow, error) // MoveToDLQ moves a message to the dead letter queue - // dlqTopicSuffix is appended to the original topic to form the DLQ topic name - // f is split across the row: its message into last_error, its structured - // half into failure_detail. - MoveToDLQ(ctx context.Context, topic string, partitionKey string, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) error + MoveToDLQ(ctx context.Context, tenant string, topic string, partitionKey string, messageID string, failureCount int, f failure.Failure, dlqTopicSuffix string) error // GarbageCollect deletes messages with offset <= minAckedOffset. - // The caller (subscriber) is responsible for computing minAckedOffset from the - // offsetStore, keeping messageStore free of cross-table queries. - // Returns the number of rows deleted. - GarbageCollect(ctx context.Context, topic string, partitionKey string, minAckedOffset int64) (int64, error) + GarbageCollect(ctx context.Context, tenant string, topic string, partitionKey string, minAckedOffset int64) (int64, error) // GetOffsetsAbove returns message offsets above afterOffset for a partition, - // ordered ascending, up to limit rows. Used by the subscriber to drive - // watermark advancement without requiring a cross-table JOIN in the delivery - // state store. Watermark advancement is incremental and idempotent, so - // limiting the result set is safe — it converges over multiple calls. - GetOffsetsAbove(ctx context.Context, topic string, partitionKey string, afterOffset int64, limit int) ([]int64, error) + // ordered ascending, up to limit rows. + GetOffsetsAbove(ctx context.Context, tenant string, topic string, partitionKey string, afterOffset int64, limit int) ([]int64, error) } // offsetStore handles offset table operations for per-partition offset tracking (internal use only) type offsetStore interface { - // Initialize creates an offset entry for a topic+partition if it doesn't exist - Initialize(ctx context.Context, topic string, partitionKey string, consumerGroup string) error + // Initialize creates an offset entry for a tenant+topic+partition if it doesn't exist + Initialize(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) error - // GetAckedOffset returns the current acked offset for a topic+partition - GetAckedOffset(ctx context.Context, topic string, partitionKey string, consumerGroup string) (int64, error) + // GetAckedOffset returns the current acked offset for a tenant+topic+partition + GetAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) (int64, error) - // UpdateAckedOffset updates the offset_acked for a topic+partition (only if new offset is greater) - UpdateAckedOffset(ctx context.Context, topic string, partitionKey string, offset int64, consumerGroup string) error + // UpdateAckedOffset updates the offset_acked for a tenant+topic+partition (only if new offset is greater) + UpdateAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string, offset int64, consumerGroup string) error // GetMinAckedOffset returns the minimum offset_acked across all consumer groups - // for a topic+partition. Returns (0, false, nil) if no offset rows exist. - // Used by the subscriber to compute the GC threshold without messageStore - // needing to query the offsets table. - GetMinAckedOffset(ctx context.Context, topic string, partitionKey string) (offset int64, found bool, err error) + // for a tenant+topic+partition. Returns (0, false, nil) if no offset rows exist. + GetMinAckedOffset(ctx context.Context, tenant string, topic string, partitionKey string) (offset int64, found bool, err error) // DeleteOffset removes one consumer group's offset row for a partition. - // Callers use this when retiring a fully-drained partition; Initialize - // recreates the row if the partition ever receives messages again. - // Idempotent: no-op if the row is already gone. - DeleteOffset(ctx context.Context, topic string, partitionKey string, consumerGroup string) error + DeleteOffset(ctx context.Context, tenant string, topic string, partitionKey string, consumerGroup string) error } // leaseInfo describes one partition's current lease row (internal use only) @@ -130,63 +117,41 @@ type leaseInfo struct { // partitionLeaseStore handles partition lease operations (internal use only) type partitionLeaseStore interface { // TryAcquireLease attempts to acquire or renew a lease for a partition - // Returns true if lease is acquired/owned by this worker - // leaseDurationMs is how long the lease is valid (in milliseconds) - TryAcquireLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (bool, error) + TryAcquireLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) (bool, error) // RenewLease renews the lease for a partition owned by this worker - // leaseDurationMs is how long the lease is valid (in milliseconds) - RenewLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) error + RenewLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string, leaseDurationMs int64) error // ReleaseLease releases the lease for a partition owned by this worker - ReleaseLease(ctx context.Context, topic string, partitionKey string, subscriberName string, consumerGroup string) error + ReleaseLease(ctx context.Context, tenant string, topic string, partitionKey string, subscriberName string, consumerGroup string) error // GetLeasedPartitions returns all partitions currently leased by this worker - GetLeasedPartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string) ([]string, error) + GetLeasedPartitions(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) ([]string, error) // GetAllLeases returns the lease row for every partition currently leased - // under (topic, consumerGroup) by any subscriber. One PK-prefix read that - // lets acquisition skip partitions validly held by other subscribers - // instead of write-probing every lease row each discovery tick. - GetAllLeases(ctx context.Context, topic string, consumerGroup string) ([]leaseInfo, error) - - // PurgeStale deletes lease rows not renewed within olderThanMs. Backstop - // for holders that crashed while owning a drained partition: acquisition - // only probes discovered partitions, so a stale lease on a partition - // with no messages is otherwise never refreshed or removed. Deleting a - // stale row is equivalent to lease expiry — a concurrent renewal makes - // the row fresh and the age predicate skips it. - PurgeStale(ctx context.Context, topic string, consumerGroup string, olderThanMs int64) error + // under (tenant, topic, consumerGroup) by any subscriber. + GetAllLeases(ctx context.Context, tenant string, topic string, consumerGroup string) ([]leaseInfo, error) + + // PurgeStale deletes lease rows not renewed within olderThanMs. + PurgeStale(ctx context.Context, tenant string, topic string, consumerGroup string, olderThanMs int64) error // DiscoverAndAcquirePartitions discovers partitions from messages table and tries to acquire leases. - // Returns the number of new leases acquired and the full list of discovered partitions. - // leaseDurationMs is how long the lease is valid (in milliseconds) - // maxPartitions limits how many total partitions this subscriber can own (0 = unlimited) - DiscoverAndAcquirePartitions(ctx context.Context, topic string, subscriberName string, consumerGroup string, leaseDurationMs int64, maxPartitions int) (acquiredCount int, discoveredPartitions []string, err error) + DiscoverAndAcquirePartitions(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string, leaseDurationMs int64, maxPartitions int) (acquiredCount int, discoveredPartitions []string, err error) } // subscriberHeartbeatStore handles subscriber heartbeat operations for fair partition leasing (internal use only) type subscriberHeartbeatStore interface { // Heartbeat registers or renews a subscriber's heartbeat - Heartbeat(ctx context.Context, topic string, subscriberName string, consumerGroup string) error + Heartbeat(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) error // ActiveSubscribers returns the names of subscribers with a recent heartbeat. - // staleDurationMs defines the staleness threshold: subscribers without a heartbeat - // within this duration are considered dead. - ActiveSubscribers(ctx context.Context, topic string, consumerGroup string, staleDurationMs int64) ([]string, error) - - // Deregister removes a subscriber's heartbeat row. Hard delete: the row - // is not needed once the subscriber is gone, and subscriber names are - // unique per process (hostname-pid), so rows would otherwise accumulate - // forever across deploys. Re-subscribing re-inserts via Heartbeat. - Deregister(ctx context.Context, topic string, subscriberName string, consumerGroup string) error - - // PurgeStale deletes heartbeat rows whose last heartbeat is older than - // olderThanMs. Backstop for subscribers that never deregistered - // (crashes, SIGKILL): without it the table grows monotonically since - // every process registers under a fresh name. Deleting a live-but-stalled - // subscriber's row is harmless — its next heartbeat re-inserts it. - PurgeStale(ctx context.Context, topic string, consumerGroup string, olderThanMs int64) error + ActiveSubscribers(ctx context.Context, tenant string, topic string, consumerGroup string, staleDurationMs int64) ([]string, error) + + // Deregister removes a subscriber's heartbeat row. + Deregister(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) error + + // PurgeStale deletes heartbeat rows whose last heartbeat is older than olderThanMs. + PurgeStale(ctx context.Context, tenant string, topic string, consumerGroup string, olderThanMs int64) error } // DeliveryState represents the full per-message delivery tracking state. @@ -206,33 +171,24 @@ type DeliveryState struct { // deliveryStateStore handles per-consumer-group delivery tracking (internal use only) type deliveryStateStore interface { // MarkDelivered inserts a row marking message as in-flight for this consumer group. - // Increments retry_count on redelivery (ON DUPLICATE KEY UPDATE), except when the - // row is marked postponed — that delivery is exempt and clears the postponed flag. - // Returns the resulting retry_count after the operation. - MarkDelivered(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retryCount int, err error) + MarkDelivered(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) (retryCount int, err error) // ExtendVisibility extends the visibility timeout for an in-flight message - // without incrementing retry_count. - ExtendVisibility(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) error + ExtendVisibility(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, visibilityTimeoutMs int64) error // MarkAcked sets acked = TRUE to indicate this group has processed the message. - MarkAcked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) error + MarkAcked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) error // MarkNacked makes the message eligible for redelivery after delayMs. - MarkNacked(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, delayMs int64) error + MarkNacked(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, delayMs int64) error - // MarkPostponed sets invisible_until = now + delay, resets retry_count, and - // sets the postponed flag. The message becomes a partition barrier until it - // redelivers, and the redelivery does not count as a failure. - MarkPostponed(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64, delayMs int64) error + // MarkPostponed sets invisible_until = now + delay, resets retry_count, and sets the postponed flag. + MarkPostponed(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64, delayMs int64) error // GetDeliveryState returns the full delivery state for a message offset. - // Returns (state, found, error). found=false means no row (never delivered). - GetDeliveryState(ctx context.Context, consumerGroup, topic, partitionKey string, offset int64) (DeliveryState, bool, error) + GetDeliveryState(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, offset int64) (DeliveryState, bool, error) // AdvanceWatermark computes the new contiguous acked watermark and cleans up // delivery state rows behind it. - // offsets are the actual message offsets above the current watermark (from messageStore). - // Returns the new watermark (highest contiguous acked offset from currentWatermark). - AdvanceWatermark(ctx context.Context, consumerGroup, topic, partitionKey string, currentWatermark int64, offsets []int64) (int64, error) + AdvanceWatermark(ctx context.Context, consumerGroup, tenant, topic, partitionKey string, currentWatermark int64, offsets []int64) (int64, error) } diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index f96ce8af4..c9d4aa8f0 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -34,10 +34,6 @@ import ( ) const ( - // workerStopTimeout is the maximum time to wait for a partition worker to - // exit after its context is cancelled. - workerStopTimeout = 30 * time.Second - // leaseReleaseTimeout is the timeout for releasing partition leases during // shutdown. Uses a fresh context since the subscription context is cancelled. leaseReleaseTimeout = 30 * time.Second @@ -112,6 +108,8 @@ type subscriber struct { leaseStore partitionLeaseStore heartbeatStore subscriberHeartbeatStore deliveryStateStore deliveryStateStore + tenants []string + shutdownTimeout time.Duration mu sync.RWMutex closed bool @@ -130,10 +128,6 @@ type subscription struct { deliveryCh chan extqueue.Delivery cancelFunc context.CancelFunc - // wg tracks the single managePartitions supervisor goroutine. - // Close() waits on this to know the entire subscription is shut down. - wg sync.WaitGroup - // done is closed once managePartitions has exited, which implies // deliveryCh is already closed. Subscribe consults it to tell a live // subscription from one whose ctx was cancelled independently of Close, @@ -151,27 +145,82 @@ type subscription struct { // Only accessed by the managePartitions goroutine for reads/reconciliation, // but mutations are protected by workersMu since stopPartitionWorker may // be called during shutdown. - workers map[string]*partitionWorker + workers map[entityqueue.PartitionIdentity]*partitionWorker workersMu sync.Mutex // lastDiscoveredPartitions is cached from the most recent - // DiscoverAndAcquirePartitions call. Used by fairShareCap during - // rebalance to avoid a redundant discovery query. - lastDiscoveredPartitions []string - - // drainedSince tracks, per owned partition absent from discovery, when - // this subscriber first observed it drained (no stored messages left). - // Drives idle-lease release: partitions drained beyond the grace period - // are released so fully-consumed short-lived partition keys don't hold - // leases, offset rows, and polling workers forever. Only accessed by the - // single managePartitions goroutine — no locking needed. - drainedSince map[string]time.Time + // DiscoverAndAcquirePartitions calls. + // Used by fairShareCap during rebalance to avoid a redundant discovery query. + lastDiscoveredPartitions []entityqueue.PartitionIdentity + + // drainedSince tracks, per owned (tenant, partition) absent from discovery, + // when this subscriber first observed it drained (no stored messages left). + // Keys include both tenant and partition. Drives idle-lease release: partitions + // drained beyond the grace period are released so fully-consumed short-lived + // partition keys don't hold leases, offset rows, and polling workers forever. + // Only accessed by the single managePartitions goroutine — no locking needed. + drainedSince map[entityqueue.PartitionIdentity]time.Time +} + +func partitionKeysForTenant(partitions []entityqueue.PartitionIdentity, tenant string) []string { + out := make([]string, 0) + for _, partition := range partitions { + if partition.Tenant == tenant { + out = append(out, partition.PartitionKey) + } + } + return out +} + +func sortPartitionIdentities(partitions []entityqueue.PartitionIdentity) { + sort.Slice(partitions, func(i, j int) bool { + if partitions[i].Tenant != partitions[j].Tenant { + return partitions[i].Tenant < partitions[j].Tenant + } + return partitions[i].PartitionKey < partitions[j].PartitionKey + }) +} + +type tenantOperationResult[T any] struct { + tenant string + value T + err error +} + +func runTenantOperations[T any]( + ctx context.Context, + tenants []string, + timeout time.Duration, + operation func(context.Context, string) (T, error), +) []tenantOperationResult[T] { + results := make(chan tenantOperationResult[T], len(tenants)) + for _, tenant := range tenants { + go func() { + result := tenantOperationResult[T]{tenant: tenant} + defer func() { results <- result }() + tenantCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + result.value, result.err = operation(tenantCtx, tenant) + }() + } + + byTenant := make(map[string]tenantOperationResult[T], len(tenants)) + for range tenants { + result := <-results + byTenant[result.tenant] = result + } + ordered := make([]tenantOperationResult[T], 0, len(tenants)) + for _, tenant := range tenants { + ordered = append(ordered, byTenant[tenant]) + } + return ordered } // partitionWorker handles polling and delivering messages for a single partition. // Each worker runs in its own goroutine, polling the DB on a ticker and sending // deliveries to the shared deliveryCh. type partitionWorker struct { + tenant string partitionKey string sub *subscription subscriber *subscriber @@ -197,6 +246,7 @@ type sqlDelivery struct { // Backend-specific fields for ack/nack subscriber *subscriber + tenant string topic string partitionKey string offset int64 @@ -229,6 +279,7 @@ func newSQLDelivery( attempt int, metadata map[string]string, subscriber *subscriber, + tenant string, topic string, partitionKey string, offset int64, @@ -246,6 +297,7 @@ func newSQLDelivery( receivedAt: time.Now().UnixMilli(), metadata: metadata, subscriber: subscriber, + tenant: tenant, topic: topic, partitionKey: partitionKey, offset: offset, @@ -296,7 +348,7 @@ func (d *sqlDelivery) Ack(ctx context.Context) error { // Mark as acked in delivery state (per consumer group). // Watermark advancement is deferred to the poll loop to reduce per-ack // latency from 4-5 DB round trips to 1. - if err := d.subscriber.deliveryStateStore.MarkAcked(ctx, d.consumerGroup, d.topic, d.partitionKey, d.offset); err != nil { + if err := d.subscriber.deliveryStateStore.MarkAcked(ctx, d.consumerGroup, d.tenant, d.topic, d.partitionKey, d.offset); err != nil { return err } @@ -338,7 +390,7 @@ func (d *sqlDelivery) Nack(ctx context.Context, f failure.Failure) error { } retryDelayMs := retryBackoffMs(d.retry, d.attempt) - if err := d.subscriber.deliveryStateStore.MarkNacked(ctx, d.consumerGroup, d.topic, d.partitionKey, d.offset, retryDelayMs); err != nil { + if err := d.subscriber.deliveryStateStore.MarkNacked(ctx, d.consumerGroup, d.tenant, d.topic, d.partitionKey, d.offset, retryDelayMs); err != nil { return err } @@ -364,7 +416,7 @@ func (d *sqlDelivery) Postpone(ctx context.Context, delayMs int64) error { // Mark as postponed in delivery state (per consumer group): invisible for // the delay, retry_count reset, partition barrier until redelivery. - if err := d.subscriber.deliveryStateStore.MarkPostponed(ctx, d.consumerGroup, d.topic, d.partitionKey, d.offset, delayMs); err != nil { + if err := d.subscriber.deliveryStateStore.MarkPostponed(ctx, d.consumerGroup, d.tenant, d.topic, d.partitionKey, d.offset, delayMs); err != nil { return err } @@ -400,7 +452,7 @@ func (d *sqlDelivery) deadLetter(ctx context.Context, f failure.Failure) error { if d.dlqConfig.Enabled { // Move to DLQ if err := d.subscriber.messageStore.MoveToDLQ( - ctx, d.topic, d.partitionKey, d.messageID, d.attempt, f, d.dlqConfig.TopicSuffix, + ctx, d.tenant, d.topic, d.partitionKey, d.messageID, d.attempt, f, d.dlqConfig.TopicSuffix, ); err != nil { return fmt.Errorf("failed to move message to DLQ: %w", err) } @@ -408,7 +460,7 @@ func (d *sqlDelivery) deadLetter(ctx context.Context, f failure.Failure) error { // Mark as acked in delivery state. Watermark advancement is deferred // to the poll loop, same as Ack. - if err := d.subscriber.deliveryStateStore.MarkAcked(ctx, d.consumerGroup, d.topic, d.partitionKey, d.offset); err != nil { + if err := d.subscriber.deliveryStateStore.MarkAcked(ctx, d.consumerGroup, d.tenant, d.topic, d.partitionKey, d.offset); err != nil { return fmt.Errorf("mark acked after DLQ move: %w", err) } @@ -431,14 +483,14 @@ func (d *sqlDelivery) ExtendVisibilityTimeout(ctx context.Context, durationMilli } // Extend visibility without incrementing retry_count - if err := d.subscriber.deliveryStateStore.ExtendVisibility(ctx, d.consumerGroup, d.topic, d.partitionKey, d.offset, durationMillis); err != nil { + if err := d.subscriber.deliveryStateStore.ExtendVisibility(ctx, d.consumerGroup, d.tenant, d.topic, d.partitionKey, d.offset, durationMillis); err != nil { return err } return nil } -func NewSubscriber(logger *zap.SugaredLogger, scope tally.Scope, messageStore messageStore, offsetStore offsetStore, leaseStore partitionLeaseStore, heartbeatStore subscriberHeartbeatStore, deliveryStateStore deliveryStateStore) *subscriber { +func NewSubscriber(logger *zap.SugaredLogger, scope tally.Scope, messageStore messageStore, offsetStore offsetStore, leaseStore partitionLeaseStore, heartbeatStore subscriberHeartbeatStore, deliveryStateStore deliveryStateStore, tenants []string) *subscriber { return &subscriber{ logger: logger.Named("subscriber"), scope: scope.SubScope("subscriber"), @@ -447,6 +499,8 @@ func NewSubscriber(logger *zap.SugaredLogger, scope tally.Scope, messageStore me leaseStore: leaseStore, heartbeatStore: heartbeatStore, deliveryStateStore: deliveryStateStore, + tenants: tenants, + shutdownTimeout: subscriptionShutdownTimeout, subscriptions: make(map[string]*subscription), } } @@ -463,24 +517,24 @@ func (s *subscriber) emitSignal(sig HookSignal) { // advanceWatermark advances offset_acked to the highest contiguous acked offset. // All operations are idempotent — safe to call from multiple paths (Reject, retry-limit, // poll loop) and safe to retry on failure. -func (s *subscriber) advanceWatermark(ctx context.Context, consumerGroup, topic, partitionKey string) error { - currentOffset, err := s.offsetStore.GetAckedOffset(ctx, topic, partitionKey, consumerGroup) +func (s *subscriber) advanceWatermark(ctx context.Context, tenant, consumerGroup, topic, partitionKey string) error { + currentOffset, err := s.offsetStore.GetAckedOffset(ctx, tenant, topic, partitionKey, consumerGroup) if err != nil { return fmt.Errorf("get acked offset for watermark advance: %w", err) } - offsets, err := s.messageStore.GetOffsetsAbove(ctx, topic, partitionKey, currentOffset, watermarkAdvancementLimit) + offsets, err := s.messageStore.GetOffsetsAbove(ctx, tenant, topic, partitionKey, currentOffset, watermarkAdvancementLimit) if err != nil { return fmt.Errorf("get message offsets for watermark advance: %w", err) } - newWatermark, err := s.deliveryStateStore.AdvanceWatermark(ctx, consumerGroup, topic, partitionKey, currentOffset, offsets) + newWatermark, err := s.deliveryStateStore.AdvanceWatermark(ctx, consumerGroup, tenant, topic, partitionKey, currentOffset, offsets) if err != nil { return fmt.Errorf("advance watermark: %w", err) } if newWatermark > currentOffset { - if err := s.offsetStore.UpdateAckedOffset(ctx, topic, partitionKey, newWatermark, consumerGroup); err != nil { + if err := s.offsetStore.UpdateAckedOffset(ctx, tenant, topic, partitionKey, newWatermark, consumerGroup); err != nil { return fmt.Errorf("update acked offset after watermark advance: %w", err) } } @@ -495,23 +549,43 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu op := metrics.Begin(s.scope, "subscribe", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() + for _, identifier := range []struct { + name string + value string + }{ + {name: "topic", value: topic}, + {name: "consumer group", value: config.ConsumerGroup}, + {name: "subscriber name", value: config.SubscriberName}, + } { + if err := validateASCIIIdentifier(identifier.name, identifier.value); err != nil { + return nil, fmt.Errorf("subscribe topic %q: %w: %v", topic, ErrInvalidConfig, err) + } + } + if err := validateRetryConfig(config.Retry); err != nil { + return nil, fmt.Errorf("subscribe topic %q: %w: %v", topic, ErrInvalidConfig, err) + } + + s.subMu.Lock() + defer s.subMu.Unlock() + s.mu.RLock() closed := s.closed + tenants := s.tenants s.mu.RUnlock() - if closed { return nil, ErrSubscriberClosed } - if err := validateRetryConfig(config.Retry); err != nil { - return nil, fmt.Errorf("subscribe topic %q: %w: %v", topic, ErrInvalidConfig, err) + if len(tenants) == 0 { + return nil, fmt.Errorf("subscribe topic %q: %w: no tenants configured", topic, ErrInvalidConfig) + } + for _, tenant := range tenants { + if err := validateASCIIIdentifier("tenant", tenant); err != nil { + return nil, fmt.Errorf("subscribe topic %q: %w: %v", topic, ErrInvalidConfig, err) + } } - // Create subscription key (topic + consumer group must be unique) subKey := topic + ":" + config.ConsumerGroup - s.subMu.Lock() - defer s.subMu.Unlock() - // A subscription whose supervisor already exited (its ctx was cancelled // without Close, which would have cleared the map) has a closed // deliveryCh. Evicting it here and falling through to build a fresh one @@ -544,7 +618,7 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu deliveryCh: make(chan extqueue.Delivery, config.BatchSize*2), cancelFunc: cancel, done: make(chan struct{}), - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } s.subscriptions[subKey] = sub @@ -552,7 +626,6 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu // Start the supervisor goroutine. It will discover partitions, acquire // leases, and spawn per-partition worker goroutines. The supervisor runs // until the subscription context is cancelled (via Close or explicit cancel). - sub.wg.Add(1) go s.managePartitions(subCtx, sub) s.logger.Debugw("subscription created", "topic", topic, "consumer_group", config.ConsumerGroup, "subscriber_name", config.SubscriberName) @@ -568,7 +641,7 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu // // Goroutine hierarchy: // -// managePartitions (this goroutine) <- supervisor, tracked by sub.wg +// managePartitions (this goroutine) <- supervisor, tracked by sub.done // +-- partitionWorker("part-1") <- tracked by sub.workerWg // +-- partitionWorker("part-2") // +-- partitionWorker("part-N") @@ -579,9 +652,8 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu // 3. workerWg.Wait(): blocks until all workers have fully exited -- this ensures // no worker can send on deliveryCh after step 4 // 4. close(deliveryCh): safe because step 3 guarantees no senders remain -// 5. managePartitions returns -> done and wg.Done() fire -> Close() unblocks +// 5. managePartitions returns -> done closes -> Close() unblocks func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { - defer sub.wg.Done() // Deferred so every exit path marks the subscription stale for Subscribe. // Must not take s.subMu: Close holds it across cancelFunc()+wg.Wait() for // every subscription, so locking here would deadlock against the very @@ -614,8 +686,13 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { // fair shares until the first leaseTicker fires. // Initial heartbeat failure is non-fatal — the next leaseTicker fires within // LeaseRenewalIntervalMs and retries. - if err := s.sendHeartbeat(ctx, sub); err != nil { - s.logger.Errorw("initial heartbeat failed", append(logFields, "error", err)...) + tenantLeaseTimeout := time.Duration(cfg.LeaseRenewalIntervalMs) * time.Millisecond + for _, result := range runTenantOperations(ctx, s.tenants, tenantLeaseTimeout, func(tenantCtx context.Context, tenant string) (struct{}, error) { + return struct{}{}, s.sendHeartbeat(tenantCtx, sub, tenant) + }) { + if result.err != nil { + s.logger.Errorw("initial heartbeat failed", append(logFields, "tenant", result.tenant, "error", result.err)...) + } } for { @@ -642,60 +719,63 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { return case <-leaseTicker.C: - // Fetch leased partitions once for this tick — shared by rebalance - // and renewLeases to avoid redundant queries. - leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - s.logger.Errorw("get leased partitions failed", append(logFields, "error", err)...) - // Skip rebalance+renew on this tick; retry next tick. - if err := s.sendHeartbeat(ctx, sub); err != nil { - s.logger.Errorw("heartbeat failed during lease error recovery", append(logFields, "error", err)...) + runTenantOperations(ctx, s.tenants, tenantLeaseTimeout, func(tenantCtx context.Context, tenant string) (struct{}, error) { + tenantFields := append(logFields, "tenant", tenant) + // Fetch leased partitions once for this tenant tick — shared by + // rebalance and renewLeases to avoid redundant queries. + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + s.logger.Errorw("get leased partitions failed", append(tenantFields, "error", err)...) + // Skip rebalance+renew for this tenant; retry next tick. + if err := s.sendHeartbeat(tenantCtx, sub, tenant); err != nil { + s.logger.Errorw("heartbeat failed during lease error recovery", append(tenantFields, "error", err)...) + } + return struct{}{}, nil } - s.emitSignal(SignalPartitionUpdate) - continue - } - // Rebalance, renew, and heartbeat are independent operations. - // Each can fail without affecting the others — the next tick retries. - // Renewal covers only the partitions kept after shedding; renewing - // a just-released lease would spuriously fail with ErrLeaseExpired. - released, err := s.rebalance(ctx, sub, leasedPartitions) - if err != nil { - s.logger.Errorw("rebalance failed", append(logFields, "error", err)...) - } - kept := leasedPartitions - if len(released) > 0 { - releasedSet := make(map[string]struct{}, len(released)) - for _, pk := range released { - releasedSet[pk] = struct{}{} + // Rebalance, renew, and heartbeat are independent operations. + // Each can fail without affecting the others — the next tick retries. + // Renewal covers only the partitions kept after shedding; renewing + // a just-released lease would spuriously fail with ErrLeaseExpired. + released, err := s.rebalance(tenantCtx, sub, tenant, leasedPartitions) + if err != nil { + s.logger.Errorw("rebalance failed", append(tenantFields, "error", err)...) } - kept = make([]string, 0, len(leasedPartitions)) - for _, pk := range leasedPartitions { - if _, ok := releasedSet[pk]; !ok { - kept = append(kept, pk) + kept := leasedPartitions + if len(released) > 0 { + releasedSet := make(map[string]struct{}, len(released)) + for _, pk := range released { + releasedSet[pk] = struct{}{} + } + kept = make([]string, 0, len(leasedPartitions)) + for _, pk := range leasedPartitions { + if _, ok := releasedSet[pk]; !ok { + kept = append(kept, pk) + } } } - } - if err := s.renewLeases(ctx, sub, kept); err != nil { - s.logger.Errorw("lease renewal failed", append(logFields, "error", err)...) - } - if err := s.sendHeartbeat(ctx, sub); err != nil { - s.logger.Errorw("periodic heartbeat failed", append(logFields, "error", err)...) - } - // Purge heartbeat rows abandoned by subscribers that never - // deregistered (crashes) — without this the table grows - // monotonically, since every process registers under a fresh - // hostname-pid name. - if err := s.heartbeatStore.PurgeStale(ctx, sub.topic, cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs); err != nil { - s.logger.Errorw("stale heartbeat purge failed", append(logFields, "error", err)...) - } - // Purge lease rows abandoned by holders that crashed while - // owning a drained partition — acquisition only probes - // discovered partitions, so nothing else ever refreshes or - // removes a stale lease on a partition with no messages. - if err := s.leaseStore.PurgeStale(ctx, sub.topic, cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs); err != nil { - s.logger.Errorw("stale lease purge failed", append(logFields, "error", err)...) - } + if err := s.renewLeases(tenantCtx, sub, tenant, kept); err != nil { + s.logger.Errorw("lease renewal failed", append(tenantFields, "error", err)...) + } + if err := s.sendHeartbeat(tenantCtx, sub, tenant); err != nil { + s.logger.Errorw("periodic heartbeat failed", append(tenantFields, "error", err)...) + } + // Purge heartbeat rows abandoned by subscribers that never + // deregistered (crashes) — without this the table grows + // monotonically, since every process registers under a fresh + // hostname-pid name. + if err := s.heartbeatStore.PurgeStale(tenantCtx, tenant, sub.topic, cfg.ConsumerGroup, heartbeatPurgeAfterLeaseDurations*cfg.LeaseDurationMs); err != nil { + s.logger.Errorw("stale heartbeat purge failed", append(tenantFields, "error", err)...) + } + // Purge lease rows abandoned by holders that crashed while + // owning a drained partition — acquisition only probes + // discovered partitions, so nothing else ever refreshes or + // removes a stale lease on a partition with no messages. + if err := s.leaseStore.PurgeStale(tenantCtx, tenant, sub.topic, cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs); err != nil { + s.logger.Errorw("stale lease purge failed", append(tenantFields, "error", err)...) + } + return struct{}{}, nil + }) s.emitSignal(SignalPartitionUpdate) case <-discoveryTicker.C: @@ -722,119 +802,164 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { // Uses fair share to limit how many partitions this subscriber acquires; // uncapped skips the fair-share cap entirely (the orphan sweep). func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subscription, uncapped bool) error { - cfg := sub.config - - // Get current leased partitions for fair share computation. - leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - return fmt.Errorf("get leased partitions: %w", err) + if len(s.tenants) == 0 { + return nil } - // Use cached discovered partitions from last tick for fair share cap. - // On the first tick, lastDiscoveredPartitions is nil → fairShareCap sees - // only owned partitions, so a joiner's first-tick cap floors at 1 and - // ramps once discovery is cached. + cfg := sub.config + sub.workersMu.Lock() - cachedDiscovered := sub.lastDiscoveredPartitions + cachedDiscovered := append([]entityqueue.PartitionIdentity(nil), sub.lastDiscoveredPartitions...) sub.workersMu.Unlock() - // maxPartitions == 0 means unlimited (the orphan sweep, or an - // uncontended single subscriber via fairShareCap). - maxPartitions := 0 - if !uncapped { - maxPartitions, err = s.fairShareCap(ctx, sub, leasedPartitions, cachedDiscovered) + discoveredByTenant := make(map[string][]string, len(s.tenants)) + leasedByTenant := make(map[string][]string, len(s.tenants)) + var discoveryErrs []error + + type tenantDiscovery struct { + discovered []string + leased []string + } + discoveryTimeout := max( + time.Duration(cfg.PartitionDiscoveryIntervalMs)*time.Millisecond, + time.Duration(cfg.LeaseRenewalIntervalMs)*time.Millisecond, + ) + results := runTenantOperations(ctx, s.tenants, discoveryTimeout, func(tenantCtx context.Context, tenant string) (tenantDiscovery, error) { + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return tenantDiscovery{}, fmt.Errorf("get leased partitions tenant=%s: %w", tenant, err) + } + + cachedForTenant := partitionKeysForTenant(cachedDiscovered, tenant) + + maxPartitions := 0 + if !uncapped { + maxPartitions, err = s.fairShareCap(tenantCtx, sub, tenant, leasedPartitions, cachedForTenant) + if err != nil { + return tenantDiscovery{}, fmt.Errorf("compute fair share cap tenant=%s: %w", tenant, err) + } + } + + _, discoveredPartitions, err := s.leaseStore.DiscoverAndAcquirePartitions(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, maxPartitions) if err != nil { - return fmt.Errorf("compute fair share cap: %w", err) + return tenantDiscovery{}, fmt.Errorf("discover and acquire partitions tenant=%s: %w", tenant, err) } + + leasedPartitions, err = s.leaseStore.GetLeasedPartitions(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return tenantDiscovery{}, fmt.Errorf("get leased partitions after acquire tenant=%s: %w", tenant, err) + } + return tenantDiscovery{discovered: discoveredPartitions, leased: leasedPartitions}, nil + }) + + for _, result := range results { + if result.err != nil { + discoveryErrs = append(discoveryErrs, result.err) + continue + } + discoveredByTenant[result.tenant] = result.value.discovered + leasedByTenant[result.tenant] = result.value.leased } - // Discover and try to acquire leases for new partitions. - // Returns discovered partitions to cache for the next tick. - _, discoveredPartitions, err := s.leaseStore.DiscoverAndAcquirePartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, maxPartitions) - if err != nil { - return fmt.Errorf("discover and acquire partitions: %w", err) + allDiscovered := make([]entityqueue.PartitionIdentity, 0) + allLeased := make([]entityqueue.PartitionIdentity, 0) + nextDrainedSince := make(map[entityqueue.PartitionIdentity]time.Time, len(sub.drainedSince)) + grace := time.Duration(idleLeaseReleaseAfterLeaseDurations*cfg.LeaseDurationMs) * time.Millisecond + now := time.Now() + var expired []entityqueue.PartitionIdentity + + for _, tenant := range s.tenants { + discoveredPartitions, succeeded := discoveredByTenant[tenant] + if !succeeded { + // Unconfirmed leases must not keep workers polling: a peer can acquire after expiry. + for _, pk := range partitionKeysForTenant(cachedDiscovered, tenant) { + allDiscovered = append(allDiscovered, entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: pk}) + } + continue + } + + leasedPartitions := leasedByTenant[tenant] + tenantDiscovered := make([]entityqueue.PartitionIdentity, 0, len(discoveredPartitions)) + for _, pk := range discoveredPartitions { + partition := entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: pk} + tenantDiscovered = append(tenantDiscovered, partition) + allDiscovered = append(allDiscovered, partition) + } + tenantLeased := make([]entityqueue.PartitionIdentity, 0, len(leasedPartitions)) + for _, pk := range leasedPartitions { + partition := entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: pk} + tenantLeased = append(tenantLeased, partition) + allLeased = append(allLeased, partition) + } + + previouslyDrained := make(map[entityqueue.PartitionIdentity]time.Time) + for partition, since := range sub.drainedSince { + if partition.Tenant == tenant { + previouslyDrained[partition] = since + } + } + tracked, tenantExpired := updateDrainedTracking(previouslyDrained, tenantLeased, tenantDiscovered, grace, now) + for key, since := range tracked { + nextDrainedSince[key] = since + } + expired = append(expired, tenantExpired...) } - // Cache discovered partitions for fairShareCap reuse by rebalance and next tick. sub.workersMu.Lock() - sub.lastDiscoveredPartitions = discoveredPartitions + sub.lastDiscoveredPartitions = allDiscovered sub.workersMu.Unlock() - // Refresh leased partitions after acquisition (new leases may have been acquired) - leasedPartitions, err = s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - return fmt.Errorf("get leased partitions after acquire: %w", err) - } - - // Idle-lease release: an owned partition absent from discovery has no - // stored messages left — everything was consumed and garbage-collected. - // Held past the grace period, such a lease buys nothing (a worker - // polling an empty partition forever) and on topics with short-lived - // partition keys it leaks a lease row, an offsets row, and a goroutine - // per key ever used. Release drops the partition entirely: reconcile - // stops its worker, and if a message arrives later the partition - // reappears in discovery and is reacquired like any new partition. - grace := time.Duration(idleLeaseReleaseAfterLeaseDurations*cfg.LeaseDurationMs) * time.Millisecond - var expired []string - sub.drainedSince, expired = updateDrainedTracking(sub.drainedSince, leasedPartitions, discoveredPartitions, grace, time.Now()) + sub.drainedSince = nextDrainedSince + sortPartitionIdentities(expired) if len(expired) > 0 { - released := make(map[string]struct{}, len(expired)) - for _, pk := range expired { - // Delete this consumer group's offsets row first, while the - // lease still guarantees exclusive ownership — nobody else can - // be initializing the partition concurrently. Initialize - // recreates the row if the partition ever comes back. - if err := s.offsetStore.DeleteOffset(ctx, sub.topic, pk, cfg.ConsumerGroup); err != nil { - // Retried next tick — the lease is still held, so the - // partition stays tracked as drained. - s.logger.Errorw("delete offsets for drained partition failed", + released := make(map[entityqueue.PartitionIdentity]struct{}, len(expired)) + for _, partition := range expired { + tenant := partition.Tenant + pk := partition.PartitionKey + if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + s.logger.Errorw("release lease for drained partition failed", + "tenant", tenant, "topic", sub.topic, "partition_key", pk, "error", err, ) continue } - if err := s.leaseStore.ReleaseLease(ctx, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - // Offsets row already deleted — harmless (the partition is - // empty; Initialize recreates it on resurrection). Release - // is retried next tick. - s.logger.Errorw("release lease for drained partition failed", + released[partition] = struct{}{} + s.stopPartitionWorker(sub, partition) + + if err := s.offsetStore.DeleteOffset(ctx, tenant, sub.topic, pk, cfg.ConsumerGroup); err != nil { + s.logger.Errorw("delete offsets for drained partition failed", + "tenant", tenant, "topic", sub.topic, "partition_key", pk, "error", err, ) continue } - released[pk] = struct{}{} - delete(sub.drainedSince, pk) - - // Stop the worker immediately rather than waiting for the - // reconcile at the end of this tick: if a message arrived in the - // window just before the release, another subscriber can acquire - // the partition right away, and the old worker must not poll - // alongside it. Mirrors the shed path in rebalance. - s.stopPartitionWorker(sub, pk) + delete(sub.drainedSince, partition) metrics.NamedCounter(s.scope, "idle_lease", "released", 1, metrics.NewTag("topic", sub.topic)) s.logger.Infow("released idle partition lease", + "tenant", tenant, "topic", sub.topic, "consumer_group", cfg.ConsumerGroup, "partition_key", pk, ) } if len(released) > 0 { - kept := make([]string, 0, len(leasedPartitions)) - for _, pk := range leasedPartitions { - if _, ok := released[pk]; !ok { - kept = append(kept, pk) + kept := make([]entityqueue.PartitionIdentity, 0, len(allLeased)) + for _, partition := range allLeased { + if _, ok := released[partition]; !ok { + kept = append(kept, partition) } } - leasedPartitions = kept + allLeased = kept } } - s.reconcilePartitionWorkers(ctx, sub, leasedPartitions) - return nil + s.reconcilePartitionWorkers(ctx, sub, allLeased) + return errors.Join(discoveryErrs...) } // updateDrainedTracking recomputes, for every owned partition absent from @@ -846,28 +971,34 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc // first-seen times carry over so the clock accumulates across ticks. Returns // the updated tracking map and the partitions drained for at least grace // (release candidates), sorted for determinism. -func updateDrainedTracking(prev map[string]time.Time, owned []string, discovered []string, grace time.Duration, now time.Time) (map[string]time.Time, []string) { - discoveredSet := make(map[string]struct{}, len(discovered)) - for _, pk := range discovered { - discoveredSet[pk] = struct{}{} - } - - next := make(map[string]time.Time) - var expired []string - for _, pk := range owned { - if _, live := discoveredSet[pk]; live { +func updateDrainedTracking( + prev map[entityqueue.PartitionIdentity]time.Time, + owned []entityqueue.PartitionIdentity, + discovered []entityqueue.PartitionIdentity, + grace time.Duration, + now time.Time, +) (map[entityqueue.PartitionIdentity]time.Time, []entityqueue.PartitionIdentity) { + discoveredSet := make(map[entityqueue.PartitionIdentity]struct{}, len(discovered)) + for _, partition := range discovered { + discoveredSet[partition] = struct{}{} + } + + next := make(map[entityqueue.PartitionIdentity]time.Time) + var expired []entityqueue.PartitionIdentity + for _, partition := range owned { + if _, live := discoveredSet[partition]; live { continue } - since, tracked := prev[pk] + since, tracked := prev[partition] if !tracked { since = now } - next[pk] = since + next[partition] = since if now.Sub(since) >= grace { - expired = append(expired, pk) + expired = append(expired, partition) } } - sort.Strings(expired) + sortPartitionIdentities(expired) return next, expired } @@ -878,52 +1009,53 @@ func updateDrainedTracking(prev map[string]time.Time, owned []string, discovered // Thread safety: only called from the single managePartitions goroutine, so the // snapshot of workers read under the lock does not change between unlock and the // subsequent start/stop calls. The lock is held briefly to read state, then -// released before blocking operations (stop may wait up to workerStopTimeout). -func (s *subscriber) reconcilePartitionWorkers(ctx context.Context, sub *subscription, currentLeases []string) { - leaseSet := make(map[string]struct{}, len(currentLeases)) - for _, pk := range currentLeases { - leaseSet[pk] = struct{}{} +// released before worker cancellation. +func (s *subscriber) reconcilePartitionWorkers(ctx context.Context, sub *subscription, currentLeases []entityqueue.PartitionIdentity) { + leaseSet := make(map[entityqueue.PartitionIdentity]struct{}, len(currentLeases)) + for _, partition := range currentLeases { + leaseSet[partition] = struct{}{} } sub.workersMu.Lock() // Find workers to stop (no longer leased) - var toStop []string - for pk := range sub.workers { - if _, ok := leaseSet[pk]; !ok { - toStop = append(toStop, pk) + var toStop []entityqueue.PartitionIdentity + for partition := range sub.workers { + if _, ok := leaseSet[partition]; !ok { + toStop = append(toStop, partition) } } // Find partitions to start (newly leased) - var toStart []string - for _, pk := range currentLeases { - if _, ok := sub.workers[pk]; !ok { - toStart = append(toStart, pk) + var toStart []entityqueue.PartitionIdentity + for _, partition := range currentLeases { + if _, ok := sub.workers[partition]; !ok { + toStart = append(toStart, partition) } } sub.workersMu.Unlock() // Stop workers for partitions no longer leased - for _, pk := range toStop { - s.stopPartitionWorker(sub, pk) + for _, partition := range toStop { + s.stopPartitionWorker(sub, partition) } // Start workers for newly leased partitions - for _, pk := range toStart { - s.startPartitionWorker(ctx, sub, pk) + for _, partition := range toStart { + s.startPartitionWorker(ctx, sub, partition) } } // startPartitionWorker creates and starts a worker goroutine for a partition. // The worker is tracked in sub.workers (for reconciliation) and sub.workerWg // (for shutdown synchronization). -func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription, partitionKey string) { +func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription, partition entityqueue.PartitionIdentity) { workerCtx, cancel := context.WithCancel(ctx) w := &partitionWorker{ - partitionKey: partitionKey, + tenant: partition.Tenant, + partitionKey: partition.PartitionKey, sub: sub, subscriber: s, cancelFunc: cancel, @@ -931,15 +1063,16 @@ func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription } sub.workersMu.Lock() - sub.workers[partitionKey] = w + sub.workers[partition] = w sub.workersMu.Unlock() sub.workerWg.Add(1) go w.run(workerCtx) s.logger.Debugw("started partition worker", + "tenant", partition.Tenant, "topic", sub.topic, - "partition_key", partitionKey, + "partition_key", partition.PartitionKey, ) } @@ -949,13 +1082,9 @@ func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription // worker's context is cancelled, so its DB calls will fail and it will exit // imminently. workerWg still tracks the old goroutine, so Close() blocks until // it fully exits -- preventing sends on a closed deliveryCh. -// -// The select with workerStopTimeout is purely for observability: if the worker -// takes longer than expected to exit, a warning is logged but no action is needed -// since workerWg handles the hard guarantee. -func (s *subscriber) stopPartitionWorker(sub *subscription, partitionKey string) { +func (s *subscriber) stopPartitionWorker(sub *subscription, partition entityqueue.PartitionIdentity) { sub.workersMu.Lock() - w, ok := sub.workers[partitionKey] + w, ok := sub.workers[partition] if !ok { sub.workersMu.Unlock() return @@ -968,34 +1097,21 @@ func (s *subscriber) stopPartitionWorker(sub *subscription, partitionKey string) // The old worker's context is cancelled so it will exit imminently. // workerWg still tracks it for shutdown -- Close() won't return until it exits. sub.workersMu.Lock() - delete(sub.workers, partitionKey) + delete(sub.workers, partition) sub.workersMu.Unlock() - - select { - case <-w.done: - s.logger.Debugw("stopped partition worker", - "topic", sub.topic, - "partition_key", partitionKey, - ) - case <-time.After(workerStopTimeout): - s.logger.Warnw("partition worker stop timeout, worker will drain in background", - "topic", sub.topic, - "partition_key", partitionKey, - ) - } } // stopAllWorkers stops all partition workers for a subscription. func (s *subscriber) stopAllWorkers(sub *subscription) { sub.workersMu.Lock() - keys := make([]string, 0, len(sub.workers)) - for pk := range sub.workers { - keys = append(keys, pk) + partitions := make([]entityqueue.PartitionIdentity, 0, len(sub.workers)) + for partition := range sub.workers { + partitions = append(partitions, partition) } sub.workersMu.Unlock() - for _, pk := range keys { - s.stopPartitionWorker(sub, pk) + for _, partition := range partitions { + s.stopPartitionWorker(sub, partition) } } @@ -1029,6 +1145,7 @@ func (w *partitionWorker) run(ctx context.Context) { // the resulting error is part of normal teardown. if errors.Is(err, context.Canceled) && errors.Is(ctx.Err(), context.Canceled) { w.subscriber.logger.Infow("poll canceled while stopping partition worker", + "tenant", w.tenant, "topic", w.sub.topic, "partition_key", w.partitionKey, "consumer_group", w.sub.config.ConsumerGroup, @@ -1037,6 +1154,7 @@ func (w *partitionWorker) run(ctx context.Context) { return } w.subscriber.logger.Errorw("poll failed", + "tenant", w.tenant, "topic", w.sub.topic, "partition_key", w.partitionKey, "consumer_group", w.sub.config.ConsumerGroup, @@ -1062,6 +1180,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { s := w.subscriber sub := w.sub cfg := sub.config + tenant := w.tenant partitionKey := w.partitionKey op := metrics.Begin(s.scope, "poll", metrics.StorageLatencyBuckets, @@ -1071,20 +1190,20 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // Initialize offset for this partition once per worker lifetime if !w.offsetInitialized { - if err := s.offsetStore.Initialize(ctx, sub.topic, partitionKey, cfg.ConsumerGroup); err != nil { + if err := s.offsetStore.Initialize(ctx, tenant, sub.topic, partitionKey, cfg.ConsumerGroup); err != nil { return fmt.Errorf("initialize offset: %w", err) } w.offsetInitialized = true } // Get current offset for this partition - currentOffset, err := s.offsetStore.GetAckedOffset(ctx, sub.topic, partitionKey, cfg.ConsumerGroup) + currentOffset, err := s.offsetStore.GetAckedOffset(ctx, tenant, sub.topic, partitionKey, cfg.ConsumerGroup) if err != nil { return fmt.Errorf("get acked offset: %w", err) } // Fetch messages from the immutable log. - rows, err := s.messageStore.FetchByOffset(ctx, sub.topic, partitionKey, currentOffset, cfg.BatchSize) + rows, err := s.messageStore.FetchByOffset(ctx, tenant, sub.topic, partitionKey, currentOffset, cfg.BatchSize) if err != nil { return fmt.Errorf("fetch messages: %w", err) } @@ -1093,7 +1212,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { for _, row := range rows { // Check per-consumer-group deliverability via delivery state. // Single query replaces separate IsDeliverable + GetRetryCount calls. - state, found, err := s.deliveryStateStore.GetDeliveryState(ctx, cfg.ConsumerGroup, sub.topic, partitionKey, row.Offset) + state, found, err := s.deliveryStateStore.GetDeliveryState(ctx, cfg.ConsumerGroup, tenant, sub.topic, partitionKey, row.Offset) if err != nil { return fmt.Errorf("get delivery state offset=%d: %w", row.Offset, err) } @@ -1115,7 +1234,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // Mark as delivered (in-flight) in delivery state. // Returns the resulting retry_count, avoiding a separate GetRetryCount call. - retryCount, err := s.deliveryStateStore.MarkDelivered(ctx, cfg.ConsumerGroup, sub.topic, partitionKey, row.Offset, cfg.VisibilityTimeoutMs) + retryCount, err := s.deliveryStateStore.MarkDelivered(ctx, cfg.ConsumerGroup, tenant, sub.topic, partitionKey, row.Offset, cfg.VisibilityTimeoutMs) if err != nil { return fmt.Errorf("mark delivered offset=%d: %w", row.Offset, err) } @@ -1140,14 +1259,14 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // visibility timeout expired unacked. if cfg.DLQ.Enabled { retryLimitFailure := failure.New("exceeded retry limit") - if err := s.messageStore.MoveToDLQ(ctx, sub.topic, partitionKey, row.ID, retryCount, retryLimitFailure, cfg.DLQ.TopicSuffix); err != nil { + if err := s.messageStore.MoveToDLQ(ctx, tenant, sub.topic, partitionKey, row.ID, retryCount, retryLimitFailure, cfg.DLQ.TopicSuffix); err != nil { return fmt.Errorf("move to DLQ message=%s: %w", row.ID, err) } } // Mark as acked so watermark can advance past it. // Watermark advancement is deferred to the poll loop. - if err := s.deliveryStateStore.MarkAcked(ctx, cfg.ConsumerGroup, sub.topic, partitionKey, row.Offset); err != nil { + if err := s.deliveryStateStore.MarkAcked(ctx, cfg.ConsumerGroup, tenant, sub.topic, partitionKey, row.Offset); err != nil { return fmt.Errorf("mark acked after retry limit message=%s: %w", row.ID, err) } continue @@ -1156,6 +1275,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // Create message (value type) msg := entityqueue.NewMessage(row.ID, row.Payload, row.PartitionKey, row.Metadata) msg.PublishedAt = row.PublishedAt + msg.Tenant = row.Tenant // Calculate message age for metrics messageAge := time.Duration(time.Now().UnixMilli()-row.PublishedAt) * time.Millisecond @@ -1215,6 +1335,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { retryCount+1, // RetryCount is 0-based, Attempt is 1-based deliveryMetadata, s, + tenant, sub.topic, partitionKey, row.Offset, @@ -1238,7 +1359,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { // Advance watermark periodically (on every poll tick). // This is deferred from Ack() to reduce per-ack latency to 1 DB call. // advanceWatermark is idempotent and incremental — safe to call every tick. - if err := s.advanceWatermark(ctx, cfg.ConsumerGroup, sub.topic, partitionKey); err != nil { + if err := s.advanceWatermark(ctx, tenant, cfg.ConsumerGroup, sub.topic, partitionKey); err != nil { s.logger.Warnw("watermark advancement failed", "topic", sub.topic, "partition_key", partitionKey, @@ -1275,7 +1396,7 @@ func (w *partitionWorker) pollAndDeliver(ctx context.Context) (retErr error) { func (w *partitionWorker) garbageCollect(ctx context.Context) error { s := w.subscriber - minOffset, found, err := s.offsetStore.GetMinAckedOffset(ctx, w.sub.topic, w.partitionKey) + minOffset, found, err := s.offsetStore.GetMinAckedOffset(ctx, w.tenant, w.sub.topic, w.partitionKey) if err != nil { return fmt.Errorf("get min acked offset: %w", err) } @@ -1283,7 +1404,7 @@ func (w *partitionWorker) garbageCollect(ctx context.Context) error { return nil } - if _, err := s.messageStore.GarbageCollect(ctx, w.sub.topic, w.partitionKey, minOffset); err != nil { + if _, err := s.messageStore.GarbageCollect(ctx, w.tenant, w.sub.topic, w.partitionKey, minOffset); err != nil { return fmt.Errorf("delete messages: %w", err) } @@ -1291,12 +1412,12 @@ func (w *partitionWorker) garbageCollect(ctx context.Context) error { } // renewLeases renews leases for all partitions owned by this worker. -func (s *subscriber) renewLeases(ctx context.Context, sub *subscription, leasedPartitions []string) error { +func (s *subscriber) renewLeases(ctx context.Context, sub *subscription, tenant string, leasedPartitions []string) error { cfg := sub.config for _, partitionKey := range leasedPartitions { - if err := s.leaseStore.RenewLease(ctx, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs); err != nil { - return fmt.Errorf("renew lease partition=%s: %w", partitionKey, err) + if err := s.leaseStore.RenewLease(ctx, tenant, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs); err != nil { + return fmt.Errorf("renew lease tenant=%s partition=%s: %w", tenant, partitionKey, err) } } return nil @@ -1305,24 +1426,35 @@ func (s *subscriber) renewLeases(ctx context.Context, sub *subscription, leasedP // releaseAllLeases releases all leases for a topic. func (s *subscriber) releaseAllLeases(ctx context.Context, sub *subscription) error { cfg := sub.config - leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - return fmt.Errorf("get leased partitions for release: %w", err) - } + var releaseErrs []error + timeout := time.Duration(cfg.LeaseRenewalIntervalMs) * time.Millisecond + results := runTenantOperations(ctx, s.tenants, timeout, func(tenantCtx context.Context, tenant string) (struct{}, error) { + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return struct{}{}, fmt.Errorf("get leased partitions for release tenant=%s: %w", tenant, err) + } - for _, partitionKey := range leasedPartitions { - if err := s.leaseStore.ReleaseLease(ctx, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("release lease partition=%s: %w", partitionKey, err) + var tenantErrs []error + for _, partitionKey := range leasedPartitions { + if err := s.leaseStore.ReleaseLease(tenantCtx, tenant, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + tenantErrs = append(tenantErrs, fmt.Errorf("release lease tenant=%s partition=%s: %w", tenant, partitionKey, err)) + } + } + return struct{}{}, errors.Join(tenantErrs...) + }) + for _, result := range results { + if result.err != nil { + releaseErrs = append(releaseErrs, result.err) } } - return nil + return errors.Join(releaseErrs...) } // sendHeartbeat sends a heartbeat for this subscriber. -func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription) error { +func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription, tenant string) error { cfg := sub.config - if err := s.heartbeatStore.Heartbeat(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("heartbeat: %w", err) + if err := s.heartbeatStore.Heartbeat(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + return fmt.Errorf("heartbeat tenant=%s: %w", tenant, err) } return nil } @@ -1330,10 +1462,21 @@ func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription) error // deregisterHeartbeat removes this subscriber's heartbeat entry during shutdown. func (s *subscriber) deregisterHeartbeat(ctx context.Context, sub *subscription) error { cfg := sub.config - if err := s.heartbeatStore.Deregister(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("deregister heartbeat: %w", err) + var deregistrationErrs []error + timeout := time.Duration(cfg.LeaseRenewalIntervalMs) * time.Millisecond + results := runTenantOperations(ctx, s.tenants, timeout, func(tenantCtx context.Context, tenant string) (struct{}, error) { + err := s.heartbeatStore.Deregister(tenantCtx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return struct{}{}, fmt.Errorf("deregister heartbeat tenant=%s: %w", tenant, err) + } + return struct{}{}, nil + }) + for _, result := range results { + if result.err != nil { + deregistrationErrs = append(deregistrationErrs, result.err) + } } - return nil + return errors.Join(deregistrationErrs...) } // rebalance checks if this subscriber holds more partitions than its fair share @@ -1342,15 +1485,14 @@ func (s *subscriber) deregisterHeartbeat(ctx context.Context, sub *subscription) // renewing a just-released lease would spuriously fail with ErrLeaseExpired. // The owned slice is never mutated (the caller shares it with lease renewal). // On error, partitions released before the failure are still returned. -func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []string) (released []string, retErr error) { +func (s *subscriber) rebalance(ctx context.Context, sub *subscription, tenant string, owned []string) (released []string, retErr error) { cfg := sub.config - // Use cached discovered partitions from the most recent discovery tick. sub.workersMu.Lock() - discoveredPartitions := sub.lastDiscoveredPartitions + discoveredPartitions := partitionKeysForTenant(sub.lastDiscoveredPartitions, tenant) sub.workersMu.Unlock() - maxPart, err := s.fairShareCap(ctx, sub, owned, discoveredPartitions) + maxPart, err := s.fairShareCap(ctx, sub, tenant, owned, discoveredPartitions) if err != nil { return nil, fmt.Errorf("compute fair share cap: %w", err) } @@ -1358,23 +1500,20 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []s return nil, nil } - // Sort a copy deterministically so the same partitions are released - // across runs without reordering the caller's slice. sortedOwned := make([]string, len(owned)) copy(sortedOwned, owned) sort.Strings(sortedOwned) - // Release excess partitions for _, pk := range sortedOwned[maxPart:] { - if err := s.leaseStore.ReleaseLease(ctx, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { return released, fmt.Errorf("release partition %s during rebalance: %w", pk, err) } released = append(released, pk) - // Stop the worker immediately to prevent duplicate processing. - s.stopPartitionWorker(sub, pk) + s.stopPartitionWorker(sub, entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: pk}) s.logger.Infow("released partition for rebalance", + "tenant", tenant, "topic", sub.topic, "consumer_group", cfg.ConsumerGroup, "partition_key", pk, @@ -1400,10 +1539,10 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []s // cap implies another under its cap (rebalance sheds, the peer acquires), // and an unleased partition implies a subscriber with spare cap to claim it // — neither a starved subscriber nor a leftover partition is a stable state. -func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, owned []string, discoveredPartitions []string) (int, error) { +func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, tenant string, owned []string, discoveredPartitions []string) (int, error) { cfg := sub.config - active, err := s.heartbeatStore.ActiveSubscribers(ctx, sub.topic, cfg.ConsumerGroup, cfg.LeaseDurationMs) + active, err := s.heartbeatStore.ActiveSubscribers(ctx, tenant, sub.topic, cfg.ConsumerGroup, cfg.LeaseDurationMs) if err != nil { return 0, err } @@ -1457,52 +1596,54 @@ func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, owned return maxPart, nil } -// Close gracefully shuts down the subscriber and all its subscriptions. -// -// For each subscription: -// 1. Cancels the subscription context, triggering managePartitions shutdown -// 2. Wraps sub.wg.Wait() in a goroutine with subscriptionShutdownTimeout so -// Close() does not block indefinitely if a subscription hangs -// 3. managePartitions internally handles stopping workers and closing deliveryCh -// (see managePartitions shutdown sequence) +// Close cancels every subscription and waits up to subscriptionShutdownTimeout +// for each supervisor. Timed-out supervisors finish asynchronously. func (s *subscriber) Close() (retErr error) { op := metrics.Begin(s.scope, "close", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() s.mu.Lock() - defer s.mu.Unlock() - if s.closed { + s.mu.Unlock() return nil } + s.closed = true + s.mu.Unlock() s.logger.Infow("closing subscriber") s.subMu.Lock() - defer s.subMu.Unlock() - - // Cancel all subscriptions + closing := make([]*subscription, 0, len(s.subscriptions)) for _, sub := range s.subscriptions { s.logger.Debugw("closing subscription", "topic", sub.topic, "consumer_group", sub.config.ConsumerGroup, ) sub.cancelFunc() + closing = append(closing, sub) + } + s.subscriptions = make(map[string]*subscription) + s.subMu.Unlock() - // Wait for the managePartitions goroutine to finish. We wrap the - // blocking Wait in a goroutine so we can enforce a timeout -- if - // managePartitions is stuck, we log a warning and move on rather - // than blocking Close() indefinitely. - done := make(chan struct{}) - go func() { - sub.wg.Wait() - close(done) - }() - + for _, sub := range closing { + timer := time.NewTimer(s.shutdownTimeout) select { - case <-done: - // Graceful shutdown completed - case <-time.After(subscriptionShutdownTimeout): + case <-sub.done: + timer.Stop() + case <-timer.C: + retErr = errors.Join(retErr, fmt.Errorf( + "subscription shutdown timed out: topic=%q consumer_group=%q", + sub.topic, + sub.config.ConsumerGroup, + )) + metrics.NamedCounter( + s.scope, + "close", + "subscription_timeout", + 1, + metrics.NewTag("topic", sub.topic), + metrics.NewTag("consumer_group", sub.config.ConsumerGroup), + ) s.logger.Warnw("subscription shutdown timeout", "topic", sub.topic, "consumer_group", sub.config.ConsumerGroup, @@ -1510,12 +1651,8 @@ func (s *subscriber) Close() (retErr error) { } } - s.subscriptions = make(map[string]*subscription) - - s.closed = true - s.logger.Infow("subscriber closed") - return nil + return retErr } func retryBackoffMs(retry extqueue.RetryConfig, attempt int) int64 { diff --git a/platform/extension/messagequeue/mysql/subscriber_heartbeat_store.go b/platform/extension/messagequeue/mysql/subscriber_heartbeat_store.go index f12a2f7c7..06645be4e 100644 --- a/platform/extension/messagequeue/mysql/subscriber_heartbeat_store.go +++ b/platform/extension/messagequeue/mysql/subscriber_heartbeat_store.go @@ -44,27 +44,27 @@ func newSubscriberHeartbeatStore(db *sql.DB, logger *zap.SugaredLogger, scope ta } // Heartbeat registers or renews a subscriber's heartbeat. -func (s *sqlSubscriberHeartbeatStore) Heartbeat(ctx context.Context, topic string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlSubscriberHeartbeatStore) Heartbeat(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "heartbeat", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() now := s.nowFunc().UnixMilli() _, err := s.db.ExecContext(ctx, fmt.Sprintf(` - INSERT INTO %s (consumer_group, topic, subscriber_name, heartbeat_at, deregistered_at) - VALUES (?, ?, ?, ?, 0) + INSERT INTO %s (tenant, consumer_group, topic, subscriber_name, heartbeat_at, deregistered_at) + VALUES (?, ?, ?, ?, ?, 0) ON DUPLICATE KEY UPDATE heartbeat_at = VALUES(heartbeat_at), deregistered_at = 0 - `, SubscriberHeartbeatsTableName), consumerGroup, topic, subscriberName, now) + `, SubscriberHeartbeatsTableName), tenant, consumerGroup, topic, subscriberName, now) if err != nil { - return fmt.Errorf("failed to send heartbeat: %w", err) + return fmt.Errorf("failed to send heartbeat tenant=%s topic=%s: %w", tenant, topic, err) } return nil } // ActiveSubscribers returns the names of subscribers with a heartbeat newer than the stale threshold. -func (s *sqlSubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, topic string, consumerGroup string, staleDurationMs int64) (_ []string, retErr error) { +func (s *sqlSubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, tenant string, topic string, consumerGroup string, staleDurationMs int64) (_ []string, retErr error) { op := metrics.Begin(s.scope, "active_subscribers", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() @@ -72,10 +72,10 @@ func (s *sqlSubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, top rows, err := s.db.QueryContext(ctx, fmt.Sprintf(` SELECT subscriber_name FROM %s - WHERE consumer_group = ? AND topic = ? AND heartbeat_at >= ? AND deregistered_at = 0 - `, SubscriberHeartbeatsTableName), consumerGroup, topic, staleThreshold) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND heartbeat_at >= ? AND deregistered_at = 0 + `, SubscriberHeartbeatsTableName), tenant, consumerGroup, topic, staleThreshold) if err != nil { - return nil, fmt.Errorf("failed to query active subscribers: %w", err) + return nil, fmt.Errorf("failed to query active subscribers tenant=%s topic=%s: %w", tenant, topic, err) } defer rows.Close() @@ -83,16 +83,17 @@ func (s *sqlSubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, top for rows.Next() { var name string if err := rows.Scan(&name); err != nil { - return nil, fmt.Errorf("failed to scan subscriber name: %w", err) + return nil, fmt.Errorf("failed to scan subscriber name tenant=%s topic=%s: %w", tenant, topic, err) } names = append(names, name) } if err := rows.Err(); err != nil { - return nil, fmt.Errorf("row iteration error: %w", err) + return nil, fmt.Errorf("row iteration error tenant=%s topic=%s: %w", tenant, topic, err) } s.logger.Debugw("found active subscribers", + logTenant, tenant, logTopic, topic, "count", len(names), "subscribers", names, @@ -103,20 +104,21 @@ func (s *sqlSubscriberHeartbeatStore) ActiveSubscribers(ctx context.Context, top // Deregister removes a subscriber's heartbeat row (hard delete — see the // subscriberHeartbeatStore interface doc). Idempotent: no-op if already gone. -func (s *sqlSubscriberHeartbeatStore) Deregister(ctx context.Context, topic string, subscriberName string, consumerGroup string) (retErr error) { +func (s *sqlSubscriberHeartbeatStore) Deregister(ctx context.Context, tenant string, topic string, subscriberName string, consumerGroup string) (retErr error) { op := metrics.Begin(s.scope, "deregister", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() _, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s - WHERE consumer_group = ? AND topic = ? AND subscriber_name = ? - `, SubscriberHeartbeatsTableName), consumerGroup, topic, subscriberName) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND subscriber_name = ? + `, SubscriberHeartbeatsTableName), tenant, consumerGroup, topic, subscriberName) if err != nil { - return fmt.Errorf("failed to deregister subscriber: %w", err) + return fmt.Errorf("failed to deregister subscriber tenant=%s topic=%s: %w", tenant, topic, err) } s.logger.Debugw("deregistered subscriber", + logTenant, tenant, logTopic, topic, "subscriber_name", subscriberName, ) @@ -126,7 +128,7 @@ func (s *sqlSubscriberHeartbeatStore) Deregister(ctx context.Context, topic stri // PurgeStale deletes heartbeat rows older than olderThanMs for the topic and // consumer group. See the subscriberHeartbeatStore interface doc. -func (s *sqlSubscriberHeartbeatStore) PurgeStale(ctx context.Context, topic string, consumerGroup string, olderThanMs int64) (retErr error) { +func (s *sqlSubscriberHeartbeatStore) PurgeStale(ctx context.Context, tenant string, topic string, consumerGroup string, olderThanMs int64) (retErr error) { op := metrics.Begin(s.scope, "purge_stale", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() @@ -134,11 +136,11 @@ func (s *sqlSubscriberHeartbeatStore) PurgeStale(ctx context.Context, topic stri result, err := s.db.ExecContext(ctx, fmt.Sprintf(` DELETE FROM %s - WHERE consumer_group = ? AND topic = ? AND heartbeat_at < ? - `, SubscriberHeartbeatsTableName), consumerGroup, topic, threshold) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND heartbeat_at < ? + `, SubscriberHeartbeatsTableName), tenant, consumerGroup, topic, threshold) if err != nil { - return fmt.Errorf("failed to purge stale heartbeats: %w", err) + return fmt.Errorf("failed to purge stale heartbeats tenant=%s topic=%s: %w", tenant, topic, err) } // RowsAffected error is swallowed because the DELETE itself succeeded; @@ -146,6 +148,7 @@ func (s *sqlSubscriberHeartbeatStore) PurgeStale(ctx context.Context, topic stri if deleted, err := result.RowsAffected(); err == nil && deleted > 0 { metrics.NamedCounter(s.scope, "purge_stale", "rows_deleted", deleted, metrics.NewTag("topic", topic)) s.logger.Debugw("purged stale heartbeats", + logTenant, tenant, logTopic, topic, "deleted", deleted, ) diff --git a/platform/extension/messagequeue/mysql/subscriber_heartbeat_store_test.go b/platform/extension/messagequeue/mysql/subscriber_heartbeat_store_test.go index 8aa8fa88b..37b85ad78 100644 --- a/platform/extension/messagequeue/mysql/subscriber_heartbeat_store_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_heartbeat_store_test.go @@ -48,7 +48,7 @@ func TestSubscriberHeartbeatStore_Heartbeat(t *testing.T) { name: "successfully send heartbeat", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("INSERT INTO queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) }, wantErr: false, @@ -57,7 +57,7 @@ func TestSubscriberHeartbeatStore_Heartbeat(t *testing.T) { name: "update existing heartbeat", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("INSERT INTO queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 2)) // ON DUPLICATE KEY UPDATE returns 2 for update }, wantErr: false, @@ -66,7 +66,7 @@ func TestSubscriberHeartbeatStore_Heartbeat(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("INSERT INTO queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). WillReturnError(fmt.Errorf("db error")) }, wantErr: true, @@ -81,7 +81,7 @@ func TestSubscriberHeartbeatStore_Heartbeat(t *testing.T) { ctx := context.Background() tt.setup(mock) - err := store.Heartbeat(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err := store.Heartbeat(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { @@ -105,7 +105,7 @@ func TestSubscriberHeartbeatStore_ActiveSubscribers(t *testing.T) { rows := sqlmock.NewRows([]string{"subscriber_name"}). AddRow("sub-1").AddRow("sub-2").AddRow("sub-3") mock.ExpectQuery("SELECT subscriber_name"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnRows(rows) }, wantNames: []string{"sub-1", "sub-2", "sub-3"}, @@ -116,7 +116,7 @@ func TestSubscriberHeartbeatStore_ActiveSubscribers(t *testing.T) { setup: func(mock sqlmock.Sqlmock) { rows := sqlmock.NewRows([]string{"subscriber_name"}) mock.ExpectQuery("SELECT subscriber_name"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnRows(rows) }, wantNames: nil, @@ -126,7 +126,7 @@ func TestSubscriberHeartbeatStore_ActiveSubscribers(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectQuery("SELECT subscriber_name"). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnError(fmt.Errorf("db error")) }, wantNames: nil, @@ -142,7 +142,7 @@ func TestSubscriberHeartbeatStore_ActiveSubscribers(t *testing.T) { ctx := context.Background() tt.setup(mock) - names, err := store.ActiveSubscribers(ctx, "test_topic", testConsumerGroup, testLeaseDurationMs) + names, err := store.ActiveSubscribers(ctx, testTenant, "test_topic", testConsumerGroup, testLeaseDurationMs) if tt.wantErr { require.Error(t, err) } else { @@ -163,10 +163,10 @@ func TestSubscriberHeartbeatStore_ActiveSubscribers_ExcludesDeregistered(t *test // Verify the query filters by deregistered_at = 0 rows := sqlmock.NewRows([]string{"subscriber_name"}).AddRow("sub-1").AddRow("sub-2") mock.ExpectQuery(`SELECT subscriber_name FROM queue_subscriber_heartbeats.*deregistered_at = 0`). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnRows(rows) - names, err := store.ActiveSubscribers(ctx, "test_topic", testConsumerGroup, testLeaseDurationMs) + names, err := store.ActiveSubscribers(ctx, testTenant, "test_topic", testConsumerGroup, testLeaseDurationMs) require.NoError(t, err) require.Equal(t, []string{"sub-1", "sub-2"}, names) require.NoError(t, mock.ExpectationsWereMet()) @@ -181,10 +181,10 @@ func TestSubscriberHeartbeatStore_Deregister_HardDelete(t *testing.T) { // Verify deregister deletes the row outright — subscriber names are // unique per process, so soft-deleted rows would accumulate forever. mock.ExpectExec(`DELETE FROM queue_subscriber_heartbeats`). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) - err := store.Deregister(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err := store.Deregister(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } @@ -199,7 +199,7 @@ func TestSubscriberHeartbeatStore_PurgeStale(t *testing.T) { name: "deletes rows older than threshold", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec(`DELETE FROM queue_subscriber_heartbeats`). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 3)) }, }, @@ -207,7 +207,7 @@ func TestSubscriberHeartbeatStore_PurgeStale(t *testing.T) { name: "no stale rows is a no-op", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec(`DELETE FROM queue_subscriber_heartbeats`). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(0, 0)) }, }, @@ -215,7 +215,7 @@ func TestSubscriberHeartbeatStore_PurgeStale(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec(`DELETE FROM queue_subscriber_heartbeats`). - WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", sqlmock.AnyArg()). WillReturnError(fmt.Errorf("db error")) }, wantErr: true, @@ -229,7 +229,7 @@ func TestSubscriberHeartbeatStore_PurgeStale(t *testing.T) { tt.setup(mock) - err := store.PurgeStale(context.Background(), "test_topic", testConsumerGroup, 300_000) + err := store.PurgeStale(context.Background(), testTenant, "test_topic", testConsumerGroup, 300_000) if tt.wantErr { require.Error(t, err) } else { @@ -248,26 +248,26 @@ func TestSubscriberHeartbeatStore_ReRegistration(t *testing.T) { // Step 1: Initial heartbeat registers the subscriber mock.ExpectExec("INSERT INTO queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) // Step 2: Deregister deletes the subscriber's row mock.ExpectExec("DELETE FROM queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) // Step 3: Heartbeat again re-registers with a fresh insert mock.ExpectExec("INSERT INTO queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName, sqlmock.AnyArg()). WillReturnResult(sqlmock.NewResult(1, 1)) - err := store.Heartbeat(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err := store.Heartbeat(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) require.NoError(t, err) - err = store.Deregister(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err = store.Deregister(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) require.NoError(t, err) - err = store.Heartbeat(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err = store.Heartbeat(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) @@ -283,7 +283,7 @@ func TestSubscriberHeartbeatStore_Deregister(t *testing.T) { name: "successfully deregister", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 1)) }, wantErr: false, @@ -292,7 +292,7 @@ func TestSubscriberHeartbeatStore_Deregister(t *testing.T) { name: "idempotent - already deregistered", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName). WillReturnResult(sqlmock.NewResult(0, 0)) }, wantErr: false, @@ -301,7 +301,7 @@ func TestSubscriberHeartbeatStore_Deregister(t *testing.T) { name: "database error", setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("DELETE FROM queue_subscriber_heartbeats"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). + WithArgs(testTenant, testConsumerGroup, "test_topic", testSubscriberName). WillReturnError(fmt.Errorf("db error")) }, wantErr: true, @@ -316,7 +316,7 @@ func TestSubscriberHeartbeatStore_Deregister(t *testing.T) { ctx := context.Background() tt.setup(mock) - err := store.Deregister(ctx, "test_topic", testSubscriberName, testConsumerGroup) + err := store.Deregister(ctx, testTenant, "test_topic", testSubscriberName, testConsumerGroup) if tt.wantErr { require.Error(t, err) } else { diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index 086147981..834d81762 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -19,6 +19,7 @@ import ( "errors" "fmt" "math" + "strings" "testing" "time" @@ -39,9 +40,10 @@ import ( // so a test only names the parts it cares about. func newDeliveryForTest(sub *subscriber, attempt int, dlq extqueue.DLQConfig, retry extqueue.RetryConfig) *sqlDelivery { msg := entityqueue.NewMessage("msg-1", []byte("payload"), "part-1", nil) + msg.Tenant = testTenant return newSQLDelivery( msg, "1", attempt, nil, - sub, "test_topic", "part-1", 100, "msg-1", "test-group", + sub, testTenant, "test_topic", "part-1", 100, "msg-1", "test-group", dlq, retry, failure.Failure{}, false, ) } @@ -50,25 +52,87 @@ func testSubscriptionConfig() extqueue.SubscriptionConfig { return extqueue.DefaultSubscriptionConfig("test-subscriber", "test-consumer") } +func TestRunTenantOperationsConcurrentlyWithDeadlines(t *testing.T) { + slowStarted := make(chan struct{}) + releaseSlow := make(chan struct{}) + fastCompleted := make(chan struct{}) + operationCompleted := make(chan []tenantOperationResult[string]) + + go func() { + operationCompleted <- runTenantOperations( + context.Background(), + []string{"slow", "fast"}, + time.Hour, + func(ctx context.Context, tenant string) (string, error) { + _, hasDeadline := ctx.Deadline() + if !hasDeadline { + return "", errors.New("tenant operation has no deadline") + } + if tenant == "slow" { + close(slowStarted) + <-releaseSlow + } else { + close(fastCompleted) + } + return tenant, nil + }, + ) + }() + + <-slowStarted + <-fastCompleted + close(releaseSlow) + results := <-operationCompleted + assert.Equal(t, []tenantOperationResult[string]{ + {tenant: "slow", value: "slow"}, + {tenant: "fast", value: "fast"}, + }, results) +} + +func TestRunTenantOperationsPropagatesCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan string, 2) + operationCompleted := make(chan []tenantOperationResult[struct{}]) + + go func() { + operationCompleted <- runTenantOperations( + ctx, + []string{"tenant-a", "tenant-b"}, + time.Hour, + func(ctx context.Context, tenant string) (struct{}, error) { + started <- tenant + <-ctx.Done() + return struct{}{}, ctx.Err() + }, + ) + }() + + assert.ElementsMatch(t, []string{"tenant-a", "tenant-b"}, []string{<-started, <-started}) + cancel() + for _, result := range <-operationCompleted { + assert.ErrorIs(t, result.err, context.Canceled) + } +} + // newTestHeartbeatStore creates a mock heartbeat store that allows all calls func newTestHeartbeatStore(ctrl *gomock.Controller) *MocksubscriberHeartbeatStore { mockHB := NewMocksubscriberHeartbeatStore(ctrl) - mockHB.EXPECT().Heartbeat(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockHB.EXPECT().ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{"self"}, nil).AnyTimes() - mockHB.EXPECT().Deregister(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockHB.EXPECT().PurgeStale(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockHB.EXPECT().Heartbeat(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockHB.EXPECT().ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{"self"}, nil).AnyTimes() + mockHB.EXPECT().Deregister(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockHB.EXPECT().PurgeStale(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() return mockHB } // newTestDeliveryStateStore creates a mock delivery state store that allows all calls func newTestDeliveryStateStore(ctrl *gomock.Controller) *MockdeliveryStateStore { mockDS := NewMockdeliveryStateStore(ctrl) - mockDS.EXPECT().MarkDelivered(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(0, nil).AnyTimes() - mockDS.EXPECT().MarkAcked(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockDS.EXPECT().MarkNacked(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockDS.EXPECT().GetDeliveryState(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(DeliveryState{}, false, nil).AnyTimes() - mockDS.EXPECT().AdvanceWatermark(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockDS.EXPECT().ExtendVisibility(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockDS.EXPECT().MarkDelivered(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(0, nil).AnyTimes() + mockDS.EXPECT().MarkAcked(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockDS.EXPECT().MarkNacked(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockDS.EXPECT().GetDeliveryState(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(DeliveryState{}, false, nil).AnyTimes() + mockDS.EXPECT().AdvanceWatermark(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockDS.EXPECT().ExtendVisibility(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() return mockDS } @@ -78,9 +142,9 @@ func setupSubscriberTest(t *testing.T, mockMessageStore *MockmessageStore, mockO mockHeartbeatStore := newTestHeartbeatStore(ctrl) mockDeliveryStateStore := newTestDeliveryStateStore(ctrl) // Allow watermark advancement calls from poll loop - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - return NewSubscriber(zaptest.NewLogger(t).Sugar().Named("subscriber"), tally.NoopScope.SubScope("subscriber"), mockMessageStore, mockOffsetStore, mockLeaseStore, mockHeartbeatStore, mockDeliveryStateStore) + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + return NewSubscriber(zaptest.NewLogger(t).Sugar().Named("subscriber"), tally.NoopScope.SubScope("subscriber"), mockMessageStore, mockOffsetStore, mockLeaseStore, mockHeartbeatStore, mockDeliveryStateStore, []string{testTenant}) } func TestSubscriber_Subscribe(t *testing.T) { @@ -119,7 +183,7 @@ func TestSubscriber_Subscribe(t *testing.T) { // Reached via releaseAllLeases on the shutdown path, and by the // discovery ticker if it fires before teardown. - mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) // Close waits for managePartitions to exit; a bare cancel would only @@ -178,6 +242,51 @@ func TestSubscriber_SubscribeRejectsInvalidRetryConfig(t *testing.T) { } } +func TestSubscriber_SubscribeRejectsInvalidIdentifiers(t *testing.T) { + overlong := strings.Repeat("x", maxIdentifierLength+1) + validConfig := testSubscriptionConfig() + tests := []struct { + name string + topic string + config extqueue.SubscriptionConfig + tenants []string + }{ + {name: "no tenants", topic: "test_topic", config: validConfig}, + {name: "overlong topic", topic: overlong, config: validConfig, tenants: []string{testTenant}}, + {name: "overlong consumer group", topic: "test_topic", config: func() extqueue.SubscriptionConfig { + cfg := validConfig + cfg.ConsumerGroup = overlong + return cfg + }(), tenants: []string{testTenant}}, + {name: "overlong subscriber name", topic: "test_topic", config: func() extqueue.SubscriptionConfig { + cfg := validConfig + cfg.SubscriberName = overlong + return cfg + }(), tenants: []string{testTenant}}, + {name: "overlong tenant", topic: "test_topic", config: validConfig, tenants: []string{overlong}}, + {name: "non-ASCII tenant", topic: "test_topic", config: validConfig, tenants: []string{"tenant-é"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sub := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + nil, + nil, + nil, + nil, + nil, + tt.tenants, + ) + + ch, err := sub.Subscribe(context.Background(), tt.topic, tt.config) + require.Nil(t, ch) + require.ErrorIs(t, err, ErrInvalidConfig) + }) + } +} + func TestSubscriber_SubscribeContextCancellation(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -185,7 +294,7 @@ func TestSubscriber_SubscribeContextCancellation(t *testing.T) { mockMessageStore := NewMockmessageStore(ctrl) mockOffsetStore := NewMockoffsetStore(ctrl) mockLeaseStore := NewMockpartitionLeaseStore(ctrl) - mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) defer func() { @@ -211,7 +320,7 @@ func TestSubscriber_SubscribeReplacesStaleSubscription(t *testing.T) { mockMessageStore := NewMockmessageStore(ctrl) mockOffsetStore := NewMockoffsetStore(ctrl) mockLeaseStore := NewMockpartitionLeaseStore(ctrl) - mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) defer func() { @@ -279,6 +388,7 @@ func TestSQLDelivery_Ack(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -290,7 +400,7 @@ func TestSQLDelivery_Ack(t *testing.T) { if !tt.alreadyAcked { // Ack only calls MarkAcked — watermark is deferred to poll loop mockDeliveryState.EXPECT().MarkAcked( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), ).Return(tt.markAckedErr) } @@ -346,6 +456,7 @@ func TestSQLDelivery_Postpone(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -356,7 +467,7 @@ func TestSQLDelivery_Postpone(t *testing.T) { if !tt.alreadyAcked { mockDeliveryState.EXPECT().MarkPostponed( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), int64(5000), + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), int64(5000), ).Return(tt.markPostponedErr) } @@ -424,6 +535,7 @@ func TestSQLDelivery_Reject(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{ @@ -439,19 +551,19 @@ func TestSQLDelivery_Reject(t *testing.T) { if tt.expectMoveDLQ { mockMsgStore.EXPECT().MoveToDLQ( - gomock.Any(), "test_topic", "part-1", "msg-1", 1, failure.New("bad payload"), "_dlq", + gomock.Any(), testTenant, "test_topic", "part-1", "msg-1", 1, failure.New("bad payload"), "_dlq", ).Return(tt.moveToDLQErr) if tt.moveToDLQErr == nil { mockDeliveryState.EXPECT().MarkAcked( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), ).Return(nil) } } if tt.expectAck { mockDeliveryState.EXPECT().MarkAcked( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), ).Return(nil) } @@ -557,6 +669,7 @@ func TestSQLDelivery_NackDeadLettersWhenBudgetSpent(t *testing.T) { NewMockpartitionLeaseStore(ctrl), newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{Enabled: true, TopicSuffix: "_dlq"} @@ -566,14 +679,14 @@ func TestSQLDelivery_NackDeadLettersWhenBudgetSpent(t *testing.T) { if tt.wantDLQ { mockMsgStore.EXPECT().MoveToDLQ( - gomock.Any(), "test_topic", "part-1", "msg-1", tt.attempt, f, "_dlq", + gomock.Any(), testTenant, "test_topic", "part-1", "msg-1", tt.attempt, f, "_dlq", ).Return(nil) mockDeliveryState.EXPECT().MarkAcked( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), ).Return(nil) } else { mockDeliveryState.EXPECT().MarkNacked( - gomock.Any(), "test-group", "test_topic", "part-1", int64(100), tt.wantRetryDelayMs, + gomock.Any(), "test-group", testTenant, "test_topic", "part-1", int64(100), tt.wantRetryDelayMs, ).Return(nil) } @@ -598,6 +711,7 @@ func TestSQLDelivery_FailureAbsentOnNormalDelivery(t *testing.T) { NewMockpartitionLeaseStore(ctrl), newTestHeartbeatStore(ctrl), NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -648,7 +762,7 @@ func TestSubscriber_Close(t *testing.T) { mockLeaseStore := NewMockpartitionLeaseStore(ctrl) // Expect lease operations during cleanup - mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) ctx := context.Background() @@ -681,6 +795,71 @@ func TestSubscriber_Close(t *testing.T) { } } +func TestSubscriber_CloseReturnsErrorOnSupervisorTimeout(t *testing.T) { + ctrl := gomock.NewController(t) + subscriberImpl := setupSubscriberTest( + t, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), + ) + sub := subscriberImpl.(*subscriber) + sub.shutdownTimeout = 0 + sub.subscriptions["test-topic:test-consumer"] = &subscription{ + topic: "test-topic", + config: testSubscriptionConfig(), + cancelFunc: func() {}, + done: make(chan struct{}), + } + + require.Error(t, sub.Close()) + assert.True(t, sub.closed) + + deliveries, err := sub.Subscribe(context.Background(), "another-topic", testSubscriptionConfig()) + require.ErrorIs(t, err, ErrSubscriberClosed) + assert.Nil(t, deliveries) +} + +func TestSubscriber_SubscribeDuringCloseDoesNotWaitForSupervisor(t *testing.T) { + ctrl := gomock.NewController(t) + subscriberImpl := setupSubscriberTest( + t, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), + ) + sub := subscriberImpl.(*subscriber) + sub.shutdownTimeout = time.Hour + hung := make(chan struct{}) + sub.subscriptions["test-topic:test-consumer"] = &subscription{ + topic: "test-topic", + config: testSubscriptionConfig(), + cancelFunc: func() {}, + done: hung, + } + + closeErr := make(chan error, 1) + go func() { + closeErr <- sub.Close() + }() + + for { + sub.mu.RLock() + closed := sub.closed + sub.mu.RUnlock() + if closed { + break + } + } + + deliveries, err := sub.Subscribe(context.Background(), "another-topic", testSubscriptionConfig()) + require.ErrorIs(t, err, ErrSubscriberClosed) + assert.Nil(t, deliveries) + + close(hung) + require.NoError(t, <-closeErr) +} + // TestSubscriber_ReconcilePartitionWorkers tests that workers are started/stopped // based on lease changes. func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { @@ -722,15 +901,16 @@ func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), newTestDeliveryStateStore(ctrl), + []string{testTenant}, ) // Allow offset initialization, fetch, and watermark calls from workers - mockOffsetStore.EXPECT().Initialize(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), false, nil).AnyTimes() + mockOffsetStore.EXPECT().Initialize(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), false, nil).AnyTimes() ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -739,23 +919,23 @@ func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { topic: "test_topic", config: testSubscriptionConfig(), deliveryCh: make(chan extqueue.Delivery, 100), - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } // Start initial workers - s.reconcilePartitionWorkers(ctx, sub, tt.initialLeases) + s.reconcilePartitionWorkers(ctx, sub, tenantPartitionKeys(testTenant, tt.initialLeases)) sub.workersMu.Lock() assert.Equal(t, len(tt.initialLeases), len(sub.workers)) sub.workersMu.Unlock() // Reconcile with updated leases - s.reconcilePartitionWorkers(ctx, sub, tt.updatedLeases) + s.reconcilePartitionWorkers(ctx, sub, tenantPartitionKeys(testTenant, tt.updatedLeases)) sub.workersMu.Lock() assert.Equal(t, len(tt.updatedLeases), len(sub.workers)) for _, pk := range tt.updatedLeases { - assert.Contains(t, sub.workers, pk) + assert.Contains(t, sub.workers, entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: pk}) } sub.workersMu.Unlock() @@ -766,6 +946,315 @@ func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { } } +func TestSubscriber_ReconcilePartitionWorkersKeepsTenantIdentity(t *testing.T) { + ctrl := gomock.NewController(t) + mockMessageStore := NewMockmessageStore(ctrl) + mockOffsetStore := NewMockoffsetStore(ctrl) + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + mockMessageStore, + mockOffsetStore, + NewMockpartitionLeaseStore(ctrl), + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{"tenant-a", "tenant-b"}, + ) + + mockOffsetStore.EXPECT().Initialize(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), false, nil).AnyTimes() + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + + ctx, cancel := context.WithCancel(context.Background()) + sub := &subscription{ + topic: "test-topic", + config: testSubscriptionConfig(), + deliveryCh: make(chan extqueue.Delivery, 2), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), + } + partitions := []entityqueue.PartitionIdentity{ + {Tenant: "tenant-a", PartitionKey: "shared"}, + {Tenant: "tenant-b", PartitionKey: "shared"}, + } + + s.reconcilePartitionWorkers(ctx, sub, partitions) + require.Len(t, sub.workers, 2) + for _, partition := range partitions { + worker, ok := sub.workers[partition] + require.True(t, ok) + assert.Equal(t, partition.Tenant, worker.tenant) + assert.Equal(t, partition.PartitionKey, worker.partitionKey) + } + + cancel() + s.stopAllWorkers(sub) +} + +func TestSubscriber_DiscoverAndReconcileWorkersIsolatesTenantFailures(t *testing.T) { + const ( + tenantBefore = "tenant-before" + tenantFailed = "tenant-failed" + tenantAfter = "tenant-after" + tenantFailedLast = "tenant-failed-last" + ) + + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + discoveryErr := errors.New("tenant store unavailable") + lastDiscoveryErr := errors.New("last tenant store unavailable") + + cfg := testSubscriptionConfig() + cfg.PollIntervalMs = int64(time.Hour / time.Millisecond) + + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantBefore, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, nil) + mockLeaseStore.EXPECT(). + DiscoverAndAcquirePartitions(gomock.Any(), tenantBefore, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, 0). + Return(1, []string{"before-new"}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantBefore, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{"before-new"}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantFailed, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, discoveryErr) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantAfter, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, nil) + mockLeaseStore.EXPECT(). + DiscoverAndAcquirePartitions(gomock.Any(), tenantAfter, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, 0). + Return(1, []string{"after-new"}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantAfter, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{"after-new"}, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenantFailedLast, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, lastDiscoveryErr) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{tenantBefore, tenantFailed, tenantAfter, tenantFailedLast}, + ) + + failedWorkerDone := make(chan struct{}) + close(failedWorkerDone) + failedWorkerKey := entityqueue.PartitionIdentity{Tenant: tenantFailed, PartitionKey: "failed-existing"} + failedDrainSince := time.Now().Add(-time.Hour) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 3), + workers: map[entityqueue.PartitionIdentity]*partitionWorker{ + failedWorkerKey: { + cancelFunc: func() {}, + done: failedWorkerDone, + }, + }, + lastDiscoveredPartitions: []entityqueue.PartitionIdentity{{Tenant: tenantFailed, PartitionKey: "failed-discovered"}}, + drainedSince: map[entityqueue.PartitionIdentity]time.Time{failedWorkerKey: failedDrainSince}, + } + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(func() { + cancel() + s.stopAllWorkers(sub) + sub.workerWg.Wait() + }) + + err := s.discoverAndReconcileWorkers(ctx, sub, true) + require.ErrorIs(t, err, discoveryErr) + require.ErrorIs(t, err, lastDiscoveryErr) + + sub.workersMu.Lock() + workerKeys := make([]entityqueue.PartitionIdentity, 0, len(sub.workers)) + for key := range sub.workers { + workerKeys = append(workerKeys, key) + } + discovered := append([]entityqueue.PartitionIdentity(nil), sub.lastDiscoveredPartitions...) + sub.workersMu.Unlock() + + assert.ElementsMatch(t, []entityqueue.PartitionIdentity{ + {Tenant: tenantBefore, PartitionKey: "before-new"}, + {Tenant: tenantAfter, PartitionKey: "after-new"}, + }, workerKeys) + assert.ElementsMatch(t, []entityqueue.PartitionIdentity{ + {Tenant: tenantBefore, PartitionKey: "before-new"}, + {Tenant: tenantFailed, PartitionKey: "failed-discovered"}, + {Tenant: tenantAfter, PartitionKey: "after-new"}, + }, discovered) + assert.NotContains(t, sub.drainedSince, failedWorkerKey) +} + +func TestSubscriber_DiscoverFailureStopsUnconfirmedWorkers(t *testing.T) { + const tenant = "tenant-failed" + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + cfg := testSubscriptionConfig() + cfg.PollIntervalMs = int64(time.Hour / time.Millisecond) + discoveryErr := errors.New("tenant store unavailable") + + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), tenant, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, discoveryErr) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{tenant}, + ) + + failedWorkerDone := make(chan struct{}) + close(failedWorkerDone) + failedWorkerKey := entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: "p1"} + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 1), + workers: map[entityqueue.PartitionIdentity]*partitionWorker{ + failedWorkerKey: { + cancelFunc: func() {}, + done: failedWorkerDone, + }, + }, + } + + err := s.discoverAndReconcileWorkers(context.Background(), sub, true) + require.ErrorIs(t, err, discoveryErr) + assert.NotContains(t, sub.workers, failedWorkerKey) +} + +func TestSubscriber_DrainedPartitionKeepsOffsetWhenLeaseReleaseFails(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + cfg := testSubscriptionConfig() + partition := entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "drained"} + + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), testTenant, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{partition.PartitionKey}, nil) + mockLeaseStore.EXPECT(). + DiscoverAndAcquirePartitions(gomock.Any(), testTenant, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, 0). + Return(0, nil, nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), testTenant, "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{partition.PartitionKey}, nil) + mockLeaseStore.EXPECT(). + ReleaseLease(gomock.Any(), testTenant, "test-topic", partition.PartitionKey, cfg.SubscriberName, cfg.ConsumerGroup). + Return(errors.New("release failed")) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + NewMockmessageStore(ctrl), + NewMockoffsetStore(ctrl), + mockLeaseStore, + newTestHeartbeatStore(ctrl), + newTestDeliveryStateStore(ctrl), + []string{testTenant}, + ) + workerDone := make(chan struct{}) + close(workerDone) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery), + workers: map[entityqueue.PartitionIdentity]*partitionWorker{ + partition: { + cancelFunc: func() {}, + done: workerDone, + }, + }, + drainedSince: map[entityqueue.PartitionIdentity]time.Time{partition: time.Now().Add(-time.Hour)}, + } + + require.NoError(t, s.discoverAndReconcileWorkers(context.Background(), sub, true)) + assert.Contains(t, sub.drainedSince, partition) +} + +func TestSubscriber_ReleaseAllLeasesContinuesAfterErrors(t *testing.T) { + ctrl := gomock.NewController(t) + mockLeaseStore := NewMockpartitionLeaseStore(ctrl) + releaseErr := errors.New("release failed") + discoveryErr := errors.New("lease lookup failed") + cfg := testSubscriptionConfig() + + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), "tenant-1", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{"p1", "p2"}, nil) + mockLeaseStore.EXPECT(). + ReleaseLease(gomock.Any(), "tenant-1", "test-topic", "p1", cfg.SubscriberName, cfg.ConsumerGroup). + Return(releaseErr) + mockLeaseStore.EXPECT(). + ReleaseLease(gomock.Any(), "tenant-1", "test-topic", "p2", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), "tenant-2", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil, discoveryErr) + mockLeaseStore.EXPECT(). + GetLeasedPartitions(gomock.Any(), "tenant-3", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return([]string{"p3"}, nil) + mockLeaseStore.EXPECT(). + ReleaseLease(gomock.Any(), "tenant-3", "test-topic", "p3", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + mockLeaseStore, NewMocksubscriberHeartbeatStore(ctrl), + NewMockdeliveryStateStore(ctrl), + []string{"tenant-1", "tenant-2", "tenant-3"}, + ) + sub := &subscription{topic: "test-topic", config: cfg} + + err := s.releaseAllLeases(context.Background(), sub) + require.ErrorIs(t, err, releaseErr) + require.ErrorIs(t, err, discoveryErr) +} + +func TestSubscriber_DeregisterHeartbeatContinuesAfterErrors(t *testing.T) { + ctrl := gomock.NewController(t) + mockHeartbeatStore := NewMocksubscriberHeartbeatStore(ctrl) + firstErr := errors.New("first deregistration failed") + lastErr := errors.New("last deregistration failed") + cfg := testSubscriptionConfig() + + mockHeartbeatStore.EXPECT(). + Deregister(gomock.Any(), "tenant-1", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(firstErr) + mockHeartbeatStore.EXPECT(). + Deregister(gomock.Any(), "tenant-2", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(nil) + mockHeartbeatStore.EXPECT(). + Deregister(gomock.Any(), "tenant-3", "test-topic", cfg.SubscriberName, cfg.ConsumerGroup). + Return(lastErr) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), mockHeartbeatStore, + NewMockdeliveryStateStore(ctrl), + []string{"tenant-1", "tenant-2", "tenant-3"}, + ) + sub := &subscription{topic: "test-topic", config: cfg} + + err := s.deregisterHeartbeat(context.Background(), sub) + require.ErrorIs(t, err, firstErr) + require.ErrorIs(t, err, lastErr) +} + // TestSubscriber_PartitionWorkerPollAndDeliver verifies a partition worker delivers messages. func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { ctrl := gomock.NewController(t) @@ -784,6 +1273,7 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -792,14 +1282,14 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { topic: "test_topic", config: cfg, deliveryCh: deliveryCh, - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } ctx := context.Background() - mockOffsetStore.EXPECT().Initialize(gomock.Any(), "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) + mockOffsetStore.EXPECT().Initialize(gomock.Any(), testTenant, "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) // GetAckedOffset is called twice: once by pollAndDeliver, once by advanceWatermark - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), "test_topic", "part-1", cfg.ConsumerGroup).Return(int64(0), nil).Times(2) + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), testTenant, "test_topic", "part-1", cfg.ConsumerGroup).Return(int64(0), nil).Times(2) row := messageRow{ ID: "msg-1", @@ -808,19 +1298,20 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { Payload: []byte("payload"), PublishedAt: time.Now().UnixMilli(), } - mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), "test_topic", "part-1", int64(0), cfg.BatchSize). + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), testTenant, "test_topic", "part-1", int64(0), cfg.BatchSize). Return([]messageRow{row}, nil) // Delivery state checks — GetDeliveryState returns not-found (new message) - mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", int64(1)).Return(DeliveryState{}, false, nil) + mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", int64(1)).Return(DeliveryState{}, false, nil) // MarkDelivered returns retry count 0 (first delivery) - mockDeliveryState.EXPECT().MarkDelivered(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", int64(1), cfg.VisibilityTimeoutMs).Return(0, nil) + mockDeliveryState.EXPECT().MarkDelivered(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", int64(1), cfg.VisibilityTimeoutMs).Return(0, nil) // advanceWatermark called at end of pollAndDeliver - mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), "test_topic", "part-1", int64(0), watermarkAdvancementLimit).Return([]int64{1}, nil) - mockDeliveryState.EXPECT().AdvanceWatermark(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", int64(0), []int64{1}).Return(int64(0), nil) + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), testTenant, "test_topic", "part-1", int64(0), watermarkAdvancementLimit).Return([]int64{1}, nil) + mockDeliveryState.EXPECT().AdvanceWatermark(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", int64(0), []int64{1}).Return(int64(0), nil) w := &partitionWorker{ + tenant: testTenant, partitionKey: "part-1", sub: sub, subscriber: s, @@ -889,9 +1380,10 @@ func TestSubscriber_PollAndDeliver_GCOnBusyTicks(t *testing.T) { topic: "test_topic", config: cfg, deliveryCh: deliveryCh, - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } w := &partitionWorker{ + tenant: testTenant, partitionKey: "part-1", sub: sub, subscriber: s, @@ -906,13 +1398,13 @@ func TestSubscriber_PollAndDeliver_GCOnBusyTicks(t *testing.T) { PublishedAt: time.Now().UnixMilli(), } // Every poll delivers one message, so the partition never idles. - mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), "test_topic", "part-1", int64(0), cfg.BatchSize). + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), testTenant, "test_topic", "part-1", int64(0), cfg.BatchSize). Return([]messageRow{row}, nil).Times(3) - mockOffsetStore.EXPECT().Initialize(gomock.Any(), "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) + mockOffsetStore.EXPECT().Initialize(gomock.Any(), testTenant, "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) // The counter reaches gcTickInterval on the second busy tick. - mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), "test_topic", "part-1").Return(int64(1), true, nil) - mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), "test_topic", "part-1", int64(1)).Return(int64(1), nil) + mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), testTenant, "test_topic", "part-1").Return(int64(1), true, nil) + mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), testTenant, "test_topic", "part-1", int64(1)).Return(int64(1), nil) ctx := context.Background() for i := 0; i < 3; i++ { @@ -974,6 +1466,7 @@ func TestSubscriber_PollAndDeliver_PostponedBarrier(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -982,37 +1475,38 @@ func TestSubscriber_PollAndDeliver_PostponedBarrier(t *testing.T) { topic: "test_topic", config: cfg, deliveryCh: deliveryCh, - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } ctx := context.Background() - mockOffsetStore.EXPECT().Initialize(gomock.Any(), "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), "test_topic", "part-1", cfg.ConsumerGroup).Return(int64(0), nil).Times(2) + mockOffsetStore.EXPECT().Initialize(gomock.Any(), testTenant, "test_topic", "part-1", cfg.ConsumerGroup).Return(nil) + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), testTenant, "test_topic", "part-1", cfg.ConsumerGroup).Return(int64(0), nil).Times(2) rows := []messageRow{ {ID: "msg-1", Offset: 1, PartitionKey: "part-1", Payload: []byte("p1"), PublishedAt: time.Now().UnixMilli()}, {ID: "msg-2", Offset: 2, PartitionKey: "part-1", Payload: []byte("p2"), PublishedAt: time.Now().UnixMilli()}, {ID: "msg-3", Offset: 3, PartitionKey: "part-1", Payload: []byte("p3"), PublishedAt: time.Now().UnixMilli()}, } - mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), "test_topic", "part-1", int64(0), cfg.BatchSize). + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), testTenant, "test_topic", "part-1", int64(0), cfg.BatchSize). Return(rows, nil) - mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", int64(1)). + mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", int64(1)). Return(tt.firstRowState, true, nil) if tt.expectDeliveries > 0 { for _, offset := range []int64{2, 3} { - mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", offset). + mockDeliveryState.EXPECT().GetDeliveryState(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", offset). Return(DeliveryState{}, false, nil) - mockDeliveryState.EXPECT().MarkDelivered(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", offset, cfg.VisibilityTimeoutMs). + mockDeliveryState.EXPECT().MarkDelivered(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", offset, cfg.VisibilityTimeoutMs). Return(0, nil) } } - mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), "test_topic", "part-1", int64(0), watermarkAdvancementLimit).Return(nil, nil) - mockDeliveryState.EXPECT().AdvanceWatermark(gomock.Any(), cfg.ConsumerGroup, "test_topic", "part-1", int64(0), gomock.Nil()).Return(int64(0), nil) + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), testTenant, "test_topic", "part-1", int64(0), watermarkAdvancementLimit).Return(nil, nil) + mockDeliveryState.EXPECT().AdvanceWatermark(gomock.Any(), cfg.ConsumerGroup, testTenant, "test_topic", "part-1", int64(0), gomock.Nil()).Return(int64(0), nil) w := &partitionWorker{ + tenant: testTenant, partitionKey: "part-1", sub: sub, subscriber: s, @@ -1052,15 +1546,16 @@ func TestSubscriber_StopAllWorkers(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), newTestDeliveryStateStore(ctrl), + []string{testTenant}, ) // Allow worker polling and watermark advancement - mockOffsetStore.EXPECT().Initialize(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() - mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), false, nil).AnyTimes() + mockOffsetStore.EXPECT().Initialize(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockMessageStore.EXPECT().FetchByOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockMessageStore.EXPECT().GetOffsetsAbove(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mockMessageStore.EXPECT().GarbageCollect(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mockOffsetStore.EXPECT().GetMinAckedOffset(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(int64(0), false, nil).AnyTimes() ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -1069,13 +1564,13 @@ func TestSubscriber_StopAllWorkers(t *testing.T) { topic: "test_topic", config: testSubscriptionConfig(), deliveryCh: make(chan extqueue.Delivery, 100), - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } // Start 3 workers - s.startPartitionWorker(ctx, sub, "part-1") - s.startPartitionWorker(ctx, sub, "part-2") - s.startPartitionWorker(ctx, sub, "part-3") + s.startPartitionWorker(ctx, sub, entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "part-1"}) + s.startPartitionWorker(ctx, sub, entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "part-2"}) + s.startPartitionWorker(ctx, sub, entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: "part-3"}) sub.workersMu.Lock() assert.Equal(t, 3, len(sub.workers)) @@ -1137,9 +1632,9 @@ func TestPartitionWorker_RunPollErrorLogging(t *testing.T) { mockLeaseStore := NewMockpartitionLeaseStore(ctrl) pollStarted := make(chan struct{}, 1) - mockOffsetStore.EXPECT().Initialize(gomock.Any(), "test_topic", "part-1", "test-consumer").Return(nil) - mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), "test_topic", "part-1", "test-consumer").DoAndReturn( - func(ctx context.Context, _, _, _ string) (int64, error) { + mockOffsetStore.EXPECT().Initialize(gomock.Any(), testTenant, "test_topic", "part-1", "test-consumer").Return(nil) + mockOffsetStore.EXPECT().GetAckedOffset(gomock.Any(), testTenant, "test_topic", "part-1", "test-consumer").DoAndReturn( + func(ctx context.Context, _, _, _, _ string) (int64, error) { select { case pollStarted <- struct{}{}: default: @@ -1157,6 +1652,7 @@ func TestPartitionWorker_RunPollErrorLogging(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), newTestDeliveryStateStore(ctrl), + []string{testTenant}, ) s.OnSignal = make(chan HookSignal, 1) cfg := testSubscriptionConfig() @@ -1167,6 +1663,7 @@ func TestPartitionWorker_RunPollErrorLogging(t *testing.T) { deliveryCh: make(chan extqueue.Delivery, 1), } worker := &partitionWorker{ + tenant: testTenant, partitionKey: "part-1", sub: sub, subscriber: s, @@ -1280,7 +1777,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { ctrl := gomock.NewController(t) mockHB := NewMocksubscriberHeartbeatStore(ctrl) mockHB.EXPECT(). - ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(tt.active, nil). AnyTimes() @@ -1289,13 +1786,14 @@ func TestSubscriber_FairShareCap(t *testing.T) { NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sub := &subscription{ topic: "test-topic", config: extqueue.DefaultSubscriptionConfig(tt.self, "test-cg"), } - got, err := s.fairShareCap(context.Background(), sub, tt.owned, tt.discovered) + got, err := s.fairShareCap(context.Background(), sub, testTenant, tt.owned, tt.discovered) require.NoError(t, err) assert.Equal(t, tt.want, got) }) @@ -1315,7 +1813,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { ctrl := gomock.NewController(t) mockHB := NewMocksubscriberHeartbeatStore(ctrl) mockHB.EXPECT(). - ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(active, nil). AnyTimes() s := NewSubscriber( @@ -1323,6 +1821,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sum := 0 @@ -1331,7 +1830,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { topic: "test-topic", config: extqueue.DefaultSubscriptionConfig(self, "test-cg"), } - cap, err := s.fairShareCap(context.Background(), sub, nil, partitionKeysN(p)) + cap, err := s.fairShareCap(context.Background(), sub, testTenant, nil, partitionKeysN(p)) require.NoError(t, err) sum += cap } @@ -1350,37 +1849,46 @@ func partitionKeysN(n int) []string { return keys } +func tenantPartitionKeys(tenant string, partitions []string) []entityqueue.PartitionIdentity { + keys := make([]entityqueue.PartitionIdentity, len(partitions)) + for i, partition := range partitions { + keys[i] = entityqueue.PartitionIdentity{Tenant: tenant, PartitionKey: partition} + } + return keys +} + func TestSubscriber_RebalanceReleasesExcess(t *testing.T) { ctrl := gomock.NewController(t) // Two active subscribers, four partitions: self is rank 0 -> cap 2. mockHB := NewMocksubscriberHeartbeatStore(ctrl) mockHB.EXPECT(). - ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return([]string{"s1", "s2"}, nil) // The lexicographically largest partitions beyond the cap are released. mockLease := NewMockpartitionLeaseStore(ctrl) mockLease.EXPECT(). - ReleaseLease(gomock.Any(), "test-topic", "pk-c", "s1", "test-cg"). + ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-c", "s1", "test-cg"). Return(nil) mockLease.EXPECT(). - ReleaseLease(gomock.Any(), "test-topic", "pk-d", "s1", "test-cg"). + ReleaseLease(gomock.Any(), testTenant, "test-topic", "pk-d", "s1", "test-cg"). Return(nil) s := NewSubscriber( zaptest.NewLogger(t).Sugar(), tally.NoopScope, NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), mockLease, mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sub := &subscription{ topic: "test-topic", config: extqueue.DefaultSubscriptionConfig("s1", "test-cg"), - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), } owned := []string{"pk-d", "pk-a", "pk-c", "pk-b"} - released, err := s.rebalance(context.Background(), sub, owned) + released, err := s.rebalance(context.Background(), sub, testTenant, owned) require.NoError(t, err) assert.Equal(t, []string{"pk-c", "pk-d"}, released) // The caller's slice is shared with lease renewal and must not be @@ -1394,7 +1902,7 @@ func TestSubscriber_RebalanceUnderCapReleasesNothing(t *testing.T) { mockHB := NewMocksubscriberHeartbeatStore(ctrl) mockHB.EXPECT(). - ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return([]string{"s1", "s2"}, nil) // No ReleaseLease expectations: owning exactly the cap sheds nothing. @@ -1402,17 +1910,18 @@ func TestSubscriber_RebalanceUnderCapReleasesNothing(t *testing.T) { zaptest.NewLogger(t).Sugar(), tally.NoopScope, NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sub := &subscription{ topic: "test-topic", config: extqueue.DefaultSubscriptionConfig("s1", "test-cg"), - workers: make(map[string]*partitionWorker), + workers: make(map[entityqueue.PartitionIdentity]*partitionWorker), // Four known partitions across two subscribers -> rank-0 cap is 2: // owning exactly the cap must shed nothing. - lastDiscoveredPartitions: []string{"pk-a", "pk-b", "pk-c", "pk-d"}, + lastDiscoveredPartitions: tenantPartitionKeys(testTenant, []string{"pk-a", "pk-b", "pk-c", "pk-d"}), } - released, err := s.rebalance(context.Background(), sub, []string{"pk-a", "pk-b"}) + released, err := s.rebalance(context.Background(), sub, testTenant, []string{"pk-a", "pk-b"}) require.NoError(t, err) assert.Empty(t, released) } @@ -1421,52 +1930,55 @@ func TestUpdateDrainedTracking(t *testing.T) { now := time.UnixMilli(1_000_000) earlier := now.Add(-time.Minute) grace := 30 * time.Second + partition := func(key string) entityqueue.PartitionIdentity { + return entityqueue.PartitionIdentity{Tenant: testTenant, PartitionKey: key} + } tests := []struct { name string - prev map[string]time.Time - owned []string - discovered []string - wantTracked map[string]time.Time - wantExpired []string + prev map[entityqueue.PartitionIdentity]time.Time + owned []entityqueue.PartitionIdentity + discovered []entityqueue.PartitionIdentity + wantTracked map[entityqueue.PartitionIdentity]time.Time + wantExpired []entityqueue.PartitionIdentity }{ { name: "owned and discovered is not tracked", - owned: []string{"p1"}, - discovered: []string{"p1"}, - wantTracked: map[string]time.Time{}, + owned: []entityqueue.PartitionIdentity{partition("p1")}, + discovered: []entityqueue.PartitionIdentity{partition("p1")}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{}, }, { name: "freshly drained starts tracking now", - owned: []string{"p1"}, + owned: []entityqueue.PartitionIdentity{partition("p1")}, discovered: nil, - wantTracked: map[string]time.Time{"p1": now}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{partition("p1"): now}, }, { name: "already tracked keeps original since time", - prev: map[string]time.Time{"p1": now.Add(-10 * time.Second)}, - owned: []string{"p1"}, - wantTracked: map[string]time.Time{"p1": now.Add(-10 * time.Second)}, + prev: map[entityqueue.PartitionIdentity]time.Time{partition("p1"): now.Add(-10 * time.Second)}, + owned: []entityqueue.PartitionIdentity{partition("p1")}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{partition("p1"): now.Add(-10 * time.Second)}, }, { name: "drained past grace expires sorted", - prev: map[string]time.Time{"p-b": earlier, "p-a": earlier, "p-young": now.Add(-time.Second)}, - owned: []string{"p-b", "p-a", "p-young"}, - wantTracked: map[string]time.Time{"p-a": earlier, "p-b": earlier, "p-young": now.Add(-time.Second)}, - wantExpired: []string{"p-a", "p-b"}, + prev: map[entityqueue.PartitionIdentity]time.Time{partition("p-b"): earlier, partition("p-a"): earlier, partition("p-young"): now.Add(-time.Second)}, + owned: []entityqueue.PartitionIdentity{partition("p-b"), partition("p-a"), partition("p-young")}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{partition("p-a"): earlier, partition("p-b"): earlier, partition("p-young"): now.Add(-time.Second)}, + wantExpired: []entityqueue.PartitionIdentity{partition("p-a"), partition("p-b")}, }, { name: "rediscovered partition resets the clock", - prev: map[string]time.Time{"p1": earlier}, - owned: []string{"p1"}, - discovered: []string{"p1"}, - wantTracked: map[string]time.Time{}, + prev: map[entityqueue.PartitionIdentity]time.Time{partition("p1"): earlier}, + owned: []entityqueue.PartitionIdentity{partition("p1")}, + discovered: []entityqueue.PartitionIdentity{partition("p1")}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{}, }, { name: "no longer owned is dropped from tracking", - prev: map[string]time.Time{"gone": earlier}, - owned: []string{"p1"}, - wantTracked: map[string]time.Time{"p1": now}, + prev: map[entityqueue.PartitionIdentity]time.Time{partition("gone"): earlier}, + owned: []entityqueue.PartitionIdentity{partition("p1")}, + wantTracked: map[entityqueue.PartitionIdentity]time.Time{partition("p1"): now}, }, } diff --git a/platform/hook/publisher.go b/platform/hook/publisher.go index 07dc3b4bd..0048fada7 100644 --- a/platform/hook/publisher.go +++ b/platform/hook/publisher.go @@ -23,9 +23,9 @@ import ( "github.com/uber/submitqueue/platform/publish" ) -// Publish sends one hook event to the domain's hook topic, partitioned by -// partitionKey. The topic key is not a parameter: a domain runs a single hook -// topic, and the caller's registry is what binds that key to a wire topic. +// Publish sends one hook event to the domain's hook topic for tenant, +// partitioned by partitionKey. The topic key is not a parameter: a domain runs +// a single hook topic, and the caller's registry binds that key to a wire topic. // // The event id is the message id, so a redelivery republishing the same event // dedups into the original message instead of enqueuing a second one. Callers @@ -37,6 +37,7 @@ import ( func Publish( ctx context.Context, registry consumer.TopicRegistry, + tenant string, event *basehook.HookEvent, partitionKey string, ) error { @@ -49,9 +50,12 @@ func Publish( return fmt.Errorf("failed to serialize hook event %s: %w", event.GetId(), err) } - if err := publish.Message( - ctx, registry, basehook.TopicKeyHook, publish.IntentID(event.GetId()), body, partitionKey, - ); err != nil { + if err := publish.Message(ctx, registry, basehook.TopicKeyHook, publish.MessageParams{ + Tenant: tenant, + ID: publish.IntentID(event.GetId()), + Payload: body, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish hook event %s: %w", event.GetId(), err) } return nil diff --git a/platform/hook/publisher_test.go b/platform/hook/publisher_test.go index c9db7b0ef..ca07233b7 100644 --- a/platform/hook/publisher_test.go +++ b/platform/hook/publisher_test.go @@ -31,6 +31,7 @@ import ( const ( testEventID = "stovepipe/validation.repository.recorded/request/7/0" testPartitionKey = "request/7" + testTenant = "monorepo/main" ) func testEvent() *basehook.HookEvent { @@ -69,12 +70,14 @@ func TestPublish(t *testing.T) { ctrl := gomock.NewController(t) registry, published := registryWithHookTopic(t, ctrl, nil) - require.NoError(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.NoError(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) // The event id is the message id, so a redelivery republishing the same // event dedups into the original message. assert.Equal(t, testEventID, published.ID) assert.Equal(t, testPartitionKey, published.PartitionKey) + assert.Equal(t, testTenant, published.Tenant) + assert.Equal(t, testTenant, published.Metadata[entityqueue.MetadataKeyQueueName]) decoded := &basehook.HookEvent{} require.NoError(t, basehook.Unmarshal(published.Payload, decoded)) @@ -105,7 +108,7 @@ func TestPublish_RejectsMalformedEvent(t *testing.T) { require.NoError(t, err) // No Publish expectation: a malformed event must not reach the queue. - require.Error(t, Publish(context.Background(), registry, tt.event, testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, tt.event, testPartitionKey)) }) } } @@ -114,12 +117,12 @@ func TestPublish_PropagatesPublishFailure(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := registryWithHookTopic(t, ctrl, errors.New("boom")) - require.Error(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) } func TestPublish_FailsWhenHookTopicIsUnregistered(t *testing.T) { registry, err := consumer.NewTopicRegistry(nil) require.NoError(t, err) - require.Error(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) } diff --git a/platform/publish/publish.go b/platform/publish/publish.go index 53cbc7ed4..6fd749225 100644 --- a/platform/publish/publish.go +++ b/platform/publish/publish.go @@ -35,29 +35,37 @@ import ( "github.com/uber/submitqueue/platform/consumer" ) -// Message publishes payload to the topic registered for key. Allowlisted -// delivery context is propagated as message metadata. +// MessageParams describes a message to publish. +type MessageParams struct { + // Tenant selects the queue shard. + Tenant string + // ID identifies the message for deduplication. + ID string + // Payload is the serialized message body. + Payload []byte + // PartitionKey selects the ordered partition. + PartitionKey string + // Metadata contains side-band delivery attributes. + Metadata map[string]string +} + +// Message publishes params.Payload to the topic registered for key. +// Params.Tenant selects the shard and is propagated as queue-name metadata. // -// msgID selects the dedup behavior, so the caller must choose it deliberately. -// The queue deduplicates on (topic, partition key, message ID) against every +// Params.ID selects the dedup behavior, so the caller must choose it deliberately. +// The queue deduplicates on (tenant, topic, partition key, message ID) against every // row it has not garbage-collected yet, consumed ones included — a window with // no upper bound on a busy partition. A publish that collides is reported as a // success and writes nothing, and nothing retries it. // -// Build msgID with IntentID: name the entity the message is about and the cause +// Build params.ID with IntentID: name the entity the message is about and the cause // this particular message exists for. A retry of the same cause then dedups, // which is what makes redelivery safe, while a new cause about the same entity // can never be swallowed by an older row. -func Message(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, msgID string, payload []byte, partitionKey string) error { - return MessageWithMetadata(ctx, registry, key, msgID, payload, partitionKey, nil) -} - -// MessageWithMetadata is Message with side-band message metadata (headers/attributes) -// attached to the delivery. Use it to carry diagnostic context that is not part of -// the payload — the backend persists and redelivers metadata alongside the message. -// Allowlisted delivery context, currently only the queue name, is propagated unless -// the caller supplies that metadata key explicitly. -func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, msgID string, payload []byte, partitionKey string, metadata map[string]string) error { +func Message(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, params MessageParams) error { + if params.Tenant == "" { + return fmt.Errorf("tenant is required") + } q, ok := registry.Queue(key) if !ok { return fmt.Errorf("no queue registered for topic key %s", key) @@ -67,24 +75,18 @@ func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, k return fmt.Errorf("no topic name registered for topic key %s", key) } - msg := entityqueue.NewMessage(msgID, payload, partitionKey, metadataFromContext(ctx, metadata)) - return q.Publisher().Publish(ctx, topicName, msg) -} - -func metadataFromContext(ctx context.Context, metadata map[string]string) map[string]string { - metadata = maps.Clone(metadata) - if _, exists := metadata[entityqueue.MetadataKeyQueueName]; exists { - return metadata - } - queueName, ok := entityqueue.QueueName(ctx) - if !ok || queueName == "" { - return metadata + if queueName, exists := params.Metadata[entityqueue.MetadataKeyQueueName]; exists && queueName != params.Tenant { + return fmt.Errorf("queue-name metadata %q does not match tenant %q", queueName, params.Tenant) } + metadata := maps.Clone(params.Metadata) if metadata == nil { metadata = make(map[string]string) } - metadata[entityqueue.MetadataKeyQueueName] = queueName - return metadata + metadata[entityqueue.MetadataKeyQueueName] = params.Tenant + + msg := entityqueue.NewMessage(params.ID, params.Payload, params.PartitionKey, metadata) + msg.Tenant = params.Tenant + return q.Publisher().Publish(ctx, topicName, msg) } // IntentID names the occasion to publish rather than the entity published diff --git a/platform/publish/publish_test.go b/platform/publish/publish_test.go index 23ec08ae7..0e6583131 100644 --- a/platform/publish/publish_test.go +++ b/platform/publish/publish_test.go @@ -55,15 +55,21 @@ func TestMessage(t *testing.T) { return nil }) - err := Message(context.Background(), registry, testKey, "msg-1", []byte("payload"), "partition-1") + err := Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "tenant-1", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) require.NoError(t, err) + assert.Equal(t, "tenant-1", published.Tenant) assert.Equal(t, "msg-1", published.ID) assert.Equal(t, []byte("payload"), published.Payload) assert.Equal(t, "partition-1", published.PartitionKey) - assert.Empty(t, published.Metadata) + assert.Equal(t, "tenant-1", published.Metadata[entityqueue.MetadataKeyQueueName]) } -func TestMessage_PropagatesQueueNameFromContext(t *testing.T) { +func TestMessage_PropagatesTenantAsQueueName(t *testing.T) { ctrl := gomock.NewController(t) registry, publisher := newTestRegistry(t, ctrl) @@ -75,12 +81,17 @@ func TestMessage_PropagatesQueueNameFromContext(t *testing.T) { return nil }) - ctx := entityqueue.WithQueueName(context.Background(), "monorepo/main") - require.NoError(t, Message(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1")) + require.NoError(t, Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "monorepo/main", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + })) assert.Equal(t, "monorepo/main", published.Metadata[entityqueue.MetadataKeyQueueName]) + assert.Equal(t, "monorepo/main", published.Tenant) } -func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { +func TestMessage_MergesMetadataWithoutMutatingInput(t *testing.T) { ctrl := gomock.NewController(t) registry, publisher := newTestRegistry(t, ctrl) @@ -93,38 +104,58 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { }) metadata := map[string]string{"failure_reason": "build failed"} - ctx := entityqueue.WithQueueName(context.Background(), "monorepo/main") - require.NoError(t, MessageWithMetadata(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1", metadata)) + require.NoError(t, Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "monorepo/main", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + Metadata: metadata, + })) assert.Equal(t, map[string]string{ "failure_reason": "build failed", entityqueue.MetadataKeyQueueName: "monorepo/main", }, published.Metadata) + assert.Equal(t, "monorepo/main", published.Tenant) assert.Equal(t, map[string]string{"failure_reason": "build failed"}, metadata) } -func TestMessageWithMetadata_ExplicitQueueNameWins(t *testing.T) { +func TestMessage_RejectsQueueNameDifferentFromTenant(t *testing.T) { ctrl := gomock.NewController(t) - registry, publisher := newTestRegistry(t, ctrl) - - var published entityqueue.Message - publisher.EXPECT(). - Publish(gomock.Any(), "test-topic", gomock.Any()). - DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { - published = msg - return nil - }) + registry, _ := newTestRegistry(t, ctrl) - ctx := entityqueue.WithQueueName(context.Background(), "inbound") metadata := map[string]string{entityqueue.MetadataKeyQueueName: "outbound"} - require.NoError(t, MessageWithMetadata(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1", metadata)) - assert.Equal(t, "outbound", published.Metadata[entityqueue.MetadataKeyQueueName]) + err := Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "inbound", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + Metadata: metadata, + }) + require.Error(t, err) } func TestMessage_UnregisteredKey(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := newTestRegistry(t, ctrl) - err := Message(context.Background(), registry, "unregistered-key", "msg-1", []byte("payload"), "partition-1") + err := Message(context.Background(), registry, "unregistered-key", MessageParams{ + Tenant: "tenant-1", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) + require.Error(t, err) +} + +func TestMessage_RequiresTenant(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newTestRegistry(t, ctrl) + + err := Message(context.Background(), registry, testKey, MessageParams{ + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) require.Error(t, err) } diff --git a/runway/controller/dlq/BUILD.bazel b/runway/controller/dlq/BUILD.bazel index 98668bc37..bda796008 100644 --- a/runway/controller/dlq/BUILD.bazel +++ b/runway/controller/dlq/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//api/runway/messagequeue:go_default_library", "//api/runway/messagequeue/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/runway/controller/dlq/dlq.go b/runway/controller/dlq/dlq.go index d53f25c74..f06838e46 100644 --- a/runway/controller/dlq/dlq.go +++ b/runway/controller/dlq/dlq.go @@ -43,6 +43,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" runwaypb "github.com/uber/submitqueue/api/runway/messagequeue/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -122,6 +123,11 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er ) return nil } + if err := entityqueue.ValidatePayloadQueue(msg, request.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + c.logger.Errorw("dlq reconcile: tenant/payload queue mismatch, dropping", "err", err) + return nil + } reason := meta["dlq.last_error"] if reason == "" { @@ -164,8 +170,12 @@ func (c *Controller) publish(ctx context.Context, result *runwaymq.MergeResult, return fmt.Errorf("failed to serialize merge result: %w", err) } - if err := publish.Message(ctx, c.registry, c.signalTopicKey, - publish.IntentID(result.GetId(), "dlq"), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, c.signalTopicKey, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId(), "dlq"), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/dlq/dlq_test.go b/runway/controller/dlq/dlq_test.go index 360852834..3a7299086 100644 --- a/runway/controller/dlq/dlq_test.go +++ b/runway/controller/dlq/dlq_test.go @@ -46,6 +46,7 @@ type publishedMsg struct { func newDelivery(t *testing.T, ctrl *gomock.Controller, payload []byte, meta map[string]string) *consumermock.MockDelivery { t.Helper() msg := entityqueue.NewMessage(testID, payload, testPartitionKey, nil) + msg.Tenant = testQueue d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Metadata().Return(meta).AnyTimes() @@ -112,6 +113,7 @@ func TestProcess_DecodableRepublishesFailure(t *testing.T) { require.Len(t, *published, 1) got := (*published)[0] assert.Equal(t, "merge-signal", got.topic) + assert.Equal(t, testQueue, got.msg.Tenant) result := &runwaymq.MergeResult{} require.NoError(t, runwaymq.Unmarshal(got.msg.Payload, result)) @@ -130,3 +132,21 @@ func TestProcess_UndecodableAcksAndPublishesNothing(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) assert.Empty(t, *published) } + +func TestProcess_TenantPayloadQueueMismatchAcksAndPublishesNothing(t *testing.T) { + ctrl := gomock.NewController(t) + registry, published := newRegistry(t, ctrl) + controller := newController(t, registry) + + req := &runwaymq.MergeRequest{Id: testID, QueueName: testQueue} + payload, err := runwaymq.Marshal(req) + require.NoError(t, err) + msg := entityqueue.NewMessage(testID, payload, testPartitionKey, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + d.EXPECT().Metadata().Return(map[string]string{"dlq.original_topic": "runway-merge"}).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), d)) + assert.Empty(t, *published) +} diff --git a/runway/controller/merge/BUILD.bazel b/runway/controller/merge/BUILD.bazel index 17434cd1a..805db4a1f 100644 --- a/runway/controller/merge/BUILD.bazel +++ b/runway/controller/merge/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//api/runway/messagequeue:go_default_library", "//api/runway/messagequeue/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/runway/controller/merge/merge.go b/runway/controller/merge/merge.go index 8a98cb3d4..0d9c918e1 100644 --- a/runway/controller/merge/merge.go +++ b/runway/controller/merge/merge.go @@ -30,6 +30,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" runwaypb "github.com/uber/submitqueue/api/runway/messagequeue/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -86,6 +87,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize merge request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, request.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } c.logger.Infow("received merge request", "id", request.Id, @@ -153,7 +157,12 @@ func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, result return fmt.Errorf("failed to serialize merge result: %w", err) } - if err := publish.Message(ctx, c.registry, key, publish.IntentID(result.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/merge/merge_test.go b/runway/controller/merge/merge_test.go index c7900863c..2bd691f1d 100644 --- a/runway/controller/merge/merge_test.go +++ b/runway/controller/merge/merge_test.go @@ -43,6 +43,7 @@ const ( func newDelivery(t *testing.T, ctrl *gomock.Controller, payload []byte) *consumermock.MockDelivery { t.Helper() msg := entityqueue.NewMessage(testID, payload, testPartitionKey, nil) + msg.Tenant = testQueue d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -118,11 +119,13 @@ func TestProcess_Success(t *testing.T) { var gotTopic string var gotPayload []byte + var gotTenant string pub := queuemock.NewMockPublisher(ctrl) pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, topic string, msg entityqueue.Message) error { gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ) @@ -145,6 +148,7 @@ func TestProcess_Success(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) assert.Equal(t, "merge-signal", gotTopic) + assert.Equal(t, testQueue, gotTenant) result := &runwaymq.MergeResult{} require.NoError(t, runwaymq.Unmarshal(gotPayload, result)) assert.Equal(t, testID, result.Id) @@ -154,6 +158,20 @@ func TestProcess_Success(t *testing.T) { assert.Equal(t, "abc123", result.Steps[0].Outputs[0].Id) } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + factory := mergermock.NewMockFactory(ctrl) + registry, _ := newRegistry(t, ctrl, nil) + controller := newController(t, factory, registry) + request := &runwaymq.MergeRequest{Id: testID, QueueName: testQueue} + msg := entityqueue.NewMessage(testID, requestPayload(t, request), testPartitionKey, nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + func TestProcess_MergeConflict(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/runway/controller/mergeconflictcheck/BUILD.bazel b/runway/controller/mergeconflictcheck/BUILD.bazel index dc9ac9274..fd5ae575f 100644 --- a/runway/controller/mergeconflictcheck/BUILD.bazel +++ b/runway/controller/mergeconflictcheck/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//api/runway/messagequeue:go_default_library", "//api/runway/messagequeue/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/runway/controller/mergeconflictcheck/mergeconflictcheck.go b/runway/controller/mergeconflictcheck/mergeconflictcheck.go index 202f879b7..ca9f678ea 100644 --- a/runway/controller/mergeconflictcheck/mergeconflictcheck.go +++ b/runway/controller/mergeconflictcheck/mergeconflictcheck.go @@ -30,6 +30,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" runwaypb "github.com/uber/submitqueue/api/runway/messagequeue/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -86,6 +87,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize merge request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, request.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } c.logger.Infow("received merge-conflict-check request", "id", request.Id, @@ -151,7 +155,12 @@ func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, result return fmt.Errorf("failed to serialize merge result: %w", err) } - if err := publish.Message(ctx, c.registry, key, publish.IntentID(result.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go b/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go index 13215239b..8e4bd201a 100644 --- a/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go +++ b/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go @@ -43,6 +43,7 @@ const ( func newDelivery(t *testing.T, ctrl *gomock.Controller, payload []byte) *consumermock.MockDelivery { t.Helper() msg := entityqueue.NewMessage(testID, payload, testPartitionKey, nil) + msg.Tenant = testQueue d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -118,11 +119,13 @@ func TestProcess_Success(t *testing.T) { var gotTopic string var gotPayload []byte + var gotTenant string pub := queuemock.NewMockPublisher(ctrl) pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, topic string, msg entityqueue.Message) error { gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ) @@ -145,12 +148,27 @@ func TestProcess_Success(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) assert.Equal(t, "merge-conflict-check-signal", gotTopic) + assert.Equal(t, testQueue, gotTenant) result := &runwaymq.MergeResult{} require.NoError(t, runwaymq.Unmarshal(gotPayload, result)) assert.Equal(t, testID, result.Id) assert.Equal(t, runwaypb.Outcome_SUCCEEDED, result.Outcome) } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + factory := mergermock.NewMockFactory(ctrl) + registry, _ := newRegistry(t, ctrl, nil) + controller := newController(t, factory, registry) + request := &runwaymq.MergeRequest{Id: testID, QueueName: testQueue} + msg := entityqueue.NewMessage(testID, requestPayload(t, request), testPartitionKey, nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + func TestProcess_MergeConflict(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/service/messagequeue/BUILD.bazel b/service/messagequeue/BUILD.bazel new file mode 100644 index 000000000..253af8ddd --- /dev/null +++ b/service/messagequeue/BUILD.bazel @@ -0,0 +1,18 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["tenant.go"], + importpath = "github.com/uber/submitqueue/service/messagequeue", + visibility = ["//visibility:public"], +) + +go_test( + name = "go_default_test", + srcs = ["tenant_test.go"], + embed = [":go_default_library"], + deps = [ + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/service/messagequeue/tenant.go b/service/messagequeue/tenant.go new file mode 100644 index 000000000..ec78f937d --- /dev/null +++ b/service/messagequeue/tenant.go @@ -0,0 +1,104 @@ +// 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 messagequeue holds message-queue configuration shared by service wiring. +package messagequeue + +import ( + "fmt" + "slices" + "strings" +) + +const maxTenantLength = 255 + +// ParseRequiredTenants parses a comma-separated tenant list. +// Tenants are unique, non-empty ASCII identifiers of at most 255 bytes. +func ParseRequiredTenants(value string) ([]string, error) { + seen := make(map[string]struct{}) + var tenants []string + for _, part := range strings.Split(value, ",") { + tenant := strings.TrimSpace(part) + if tenant == "" { + continue + } + if len(tenant) > maxTenantLength { + return nil, fmt.Errorf("tenant %q exceeds %d bytes", tenant, maxTenantLength) + } + for i := range len(tenant) { + if tenant[i] > 0x7f { + return nil, fmt.Errorf("tenant %q must contain only ASCII characters", tenant) + } + if tenant[i] == 0 { + return nil, fmt.Errorf("tenant %q must not contain NUL", tenant) + } + } + if _, exists := seen[tenant]; exists { + continue + } + seen[tenant] = struct{}{} + tenants = append(tenants, tenant) + } + if len(tenants) == 0 { + return nil, fmt.Errorf("tenant list must contain at least one tenant") + } + return tenants, nil +} + +// ValidateTenantSetsEqual requires both named configurations to contain the +// same tenants. Input order and duplicate entries do not affect equality. +func ValidateTenantSetsEqual(firstName string, first []string, secondName string, second []string) error { + onlyFirst, onlySecond := tenantSetDifferences(first, second) + if len(onlyFirst) == 0 && len(onlySecond) == 0 { + return nil + } + return fmt.Errorf("%s and %s tenant sets differ: only in %s=%v, only in %s=%v", + firstName, secondName, firstName, onlyFirst, secondName, onlySecond) +} + +// ValidateTenantSubset requires every tenant in subset to be present in +// superset. Input order and duplicate entries do not affect membership. +func ValidateTenantSubset(supersetName string, superset []string, subsetName string, subset []string) error { + _, onlySubset := tenantSetDifferences(superset, subset) + if len(onlySubset) == 0 { + return nil + } + return fmt.Errorf("%s contains tenants not present in %s: %v", subsetName, supersetName, onlySubset) +} + +func tenantSetDifferences(first, second []string) (onlyFirst, onlySecond []string) { + firstSet := tenantSet(first) + secondSet := tenantSet(second) + for tenant := range firstSet { + if _, exists := secondSet[tenant]; !exists { + onlyFirst = append(onlyFirst, tenant) + } + } + for tenant := range secondSet { + if _, exists := firstSet[tenant]; !exists { + onlySecond = append(onlySecond, tenant) + } + } + slices.Sort(onlyFirst) + slices.Sort(onlySecond) + return onlyFirst, onlySecond +} + +func tenantSet(tenants []string) map[string]struct{} { + set := make(map[string]struct{}, len(tenants)) + for _, tenant := range tenants { + set[tenant] = struct{}{} + } + return set +} diff --git a/service/messagequeue/tenant_test.go b/service/messagequeue/tenant_test.go new file mode 100644 index 000000000..7fc74c943 --- /dev/null +++ b/service/messagequeue/tenant_test.go @@ -0,0 +1,110 @@ +// 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 messagequeue + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseRequiredTenants(t *testing.T) { + tests := []struct { + name string + value string + want []string + wantErr bool + }{ + { + name: "comma-separated tenants", + value: " monorepo/main, monorepo/release ", + want: []string{"monorepo/main", "monorepo/release"}, + }, + { + name: "duplicates removed in first-seen order", + value: "monorepo/main,monorepo/release,monorepo/main", + want: []string{"monorepo/main", "monorepo/release"}, + }, + {name: "empty", wantErr: true}, + {name: "whitespace and commas", value: " , , ", wantErr: true}, + {name: "tenant exceeds byte limit", value: strings.Repeat("x", maxTenantLength+1), wantErr: true}, + {name: "tenant contains non-ASCII characters", value: "monorepo/café", wantErr: true}, + {name: "tenant contains NUL", value: "monorepo/\x00main", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseRequiredTenants(tt.value) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestValidateTenantSetsEqual(t *testing.T) { + tests := []struct { + name string + first []string + second []string + wantErr bool + }{ + {name: "same tenants", first: []string{"a", "b"}, second: []string{"b", "a"}}, + {name: "duplicates ignored", first: []string{"a", "a"}, second: []string{"a"}}, + {name: "first has extra tenant", first: []string{"a", "b"}, second: []string{"a"}, wantErr: true}, + {name: "second has extra tenant", first: []string{"a"}, second: []string{"a", "b"}, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateTenantSetsEqual("first", tt.first, "second", tt.second) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + }) + } +} + +func TestValidateTenantSubset(t *testing.T) { + tests := []struct { + name string + superset []string + subset []string + wantErr bool + }{ + {name: "subset", superset: []string{"a", "b"}, subset: []string{"b"}}, + {name: "equal", superset: []string{"a"}, subset: []string{"a"}}, + {name: "duplicates ignored", superset: []string{"a"}, subset: []string{"a", "a"}}, + {name: "unknown tenant", superset: []string{"a"}, subset: []string{"b"}, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateTenantSubset("superset", tt.superset, "subset", tt.subset) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + }) + } +} diff --git a/service/stovepipe/server/BUILD.bazel b/service/stovepipe/server/BUILD.bazel index 509957e26..321f5ea23 100644 --- a/service/stovepipe/server/BUILD.bazel +++ b/service/stovepipe/server/BUILD.bazel @@ -21,6 +21,7 @@ go_library( "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//platform/hook:go_default_library", + "//service/messagequeue:go_default_library", "//service/stovepipe/server/mapper:go_default_library", "//stovepipe/controller:go_default_library", "//stovepipe/controller/build:go_default_library", diff --git a/service/stovepipe/server/main.go b/service/stovepipe/server/main.go index 1e412c4be..bacab144d 100644 --- a/service/stovepipe/server/main.go +++ b/service/stovepipe/server/main.go @@ -43,6 +43,7 @@ import ( extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" platformhook "github.com/uber/submitqueue/platform/hook" + servicemq "github.com/uber/submitqueue/service/messagequeue" "github.com/uber/submitqueue/service/stovepipe/server/mapper" "github.com/uber/submitqueue/stovepipe/controller" "github.com/uber/submitqueue/stovepipe/controller/build" @@ -254,11 +255,17 @@ func run() error { } defer queueDB.Close() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) + if err != nil { + return fmt.Errorf("failed to configure queue subscribers: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -340,6 +347,7 @@ func run() error { storageFty, materializer, registry, + tenants, ) srv := &StovepipeServer{ pingController: pingController, diff --git a/stovepipe/controller/BUILD.bazel b/stovepipe/controller/BUILD.bazel index d56dcd402..f87da20ff 100644 --- a/stovepipe/controller/BUILD.bazel +++ b/stovepipe/controller/BUILD.bazel @@ -10,7 +10,6 @@ go_library( visibility = ["//visibility:public"], deps = [ "//api/stovepipe/protopb:go_default_library", - "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/errs:go_default_library", "//platform/extension/counter:go_default_library", diff --git a/stovepipe/controller/build/BUILD.bazel b/stovepipe/controller/build/BUILD.bazel index ecb3e1cbb..ec1689ee5 100644 --- a/stovepipe/controller/build/BUILD.bazel +++ b/stovepipe/controller/build/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/stovepipe/controller/build", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/errs:go_default_library", "//platform/metrics:go_default_library", diff --git a/stovepipe/controller/build/build.go b/stovepipe/controller/build/build.go index 0cdebdc17..49aec5401 100644 --- a/stovepipe/controller/build/build.go +++ b/stovepipe/controller/build/build.go @@ -24,6 +24,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/metrics" @@ -88,6 +89,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: a malformed message will never succeed regardless of retries. return fmt.Errorf("failed to deserialize build request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, br.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: br.GetQueueName()}) if err != nil { metrics.NamedCounter(c.metricsScope, _opName, "storage_resolve_errors", 1, metrics.TagsFromContext(ctx)...) @@ -174,7 +178,12 @@ func (c *Controller) publishBuildSignal(ctx context.Context, buildID, queue stri return fmt.Errorf("failed to serialize build signal: %w", err) } - return publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuildSignal, publish.IntentID(buildID), payload, buildID) + return publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuildSignal, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(buildID), + Payload: payload, + PartitionKey: buildID, + }) } // Name returns the controller name for logging and metrics. diff --git a/stovepipe/controller/build/build_test.go b/stovepipe/controller/build/build_test.go index e5c9b7bc9..f7db59c70 100644 --- a/stovepipe/controller/build/build_test.go +++ b/stovepipe/controller/build/build_test.go @@ -119,6 +119,19 @@ func TestProcessTagsDeserializeErrorsFromContext(t *testing.T) { assert.EqualValues(t, 1, counter.Value()) } +func TestProcessRejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c, _ := newController(t, ctrl) + msg := entityqueue.NewMessage(testID, buildPayload(t, testID), testQueue, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + err := c.Process(queueContext(), d) + require.Error(t, err) + assert.False(t, errs.IsRetryable(err)) +} + func TestPublishBuildSignalCarriesQueueMetadata(t *testing.T) { ctrl := gomock.NewController(t) c, m := newController(t, ctrl) @@ -132,13 +145,16 @@ func TestPublishBuildSignalCarriesQueueMetadata(t *testing.T) { require.NoError(t, c.publishBuildSignal(queueContext(), testBuildID, testQueue)) assert.Equal(t, testBuildID, got.PartitionKey) + assert.Equal(t, testQueue, got.Tenant) assert.Equal(t, testQueue, got.Metadata[entityqueue.MetadataKeyQueueName]) } func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte) consumer.Delivery { t.Helper() d := consumermock.NewMockDelivery(ctrl) - d.EXPECT().Message().Return(entityqueue.NewMessage(testID, payload, testID, nil)).AnyTimes() + msg := entityqueue.NewMessage(testID, payload, testID, nil) + msg.Tenant = testQueue + d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() return d } diff --git a/stovepipe/controller/buildsignal/BUILD.bazel b/stovepipe/controller/buildsignal/BUILD.bazel index 5cb922589..743dc3621 100644 --- a/stovepipe/controller/buildsignal/BUILD.bazel +++ b/stovepipe/controller/buildsignal/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/stovepipe/controller/buildsignal", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/stovepipe/controller/buildsignal/buildsignal.go b/stovepipe/controller/buildsignal/buildsignal.go index f8aca8bfd..c2c029ef1 100644 --- a/stovepipe/controller/buildsignal/buildsignal.go +++ b/stovepipe/controller/buildsignal/buildsignal.go @@ -26,6 +26,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -108,6 +109,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: a malformed message will never succeed regardless of retries. return fmt.Errorf("failed to deserialize build signal: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, sig.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: sig.GetQueueName()}) if err != nil { metrics.NamedCounter(c.metricsScope, _opName, "storage_resolve_errors", 1, metrics.TagsFromContext(ctx)...) @@ -358,7 +362,12 @@ func (c *Controller) publishRecord(ctx context.Context, requestID, queue string) if err != nil { return fmt.Errorf("failed to serialize record: %w", err) } - return publish.Message(ctx, c.registry, stovepipemq.TopicKeyRecord, publish.IntentID(requestID), payload, requestID) + return publish.Message(ctx, c.registry, stovepipemq.TopicKeyRecord, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(requestID), + Payload: payload, + PartitionKey: requestID, + }) } // Name returns the controller name for logging and metrics. diff --git a/stovepipe/controller/buildsignal/buildsignal_test.go b/stovepipe/controller/buildsignal/buildsignal_test.go index d20ee7289..f65ba9d58 100644 --- a/stovepipe/controller/buildsignal/buildsignal_test.go +++ b/stovepipe/controller/buildsignal/buildsignal_test.go @@ -110,10 +110,25 @@ func TestProcessTagsMetricsWithQueue(t *testing.T) { assert.EqualValues(t, 1, counter.Value()) } +func TestProcessRejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c, _ := newController(t, ctrl) + msg := entityqueue.NewMessage(testBuildID, buildSignalPayload(t, testID), testQueue, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + err := c.Process(queueContext(), d) + require.Error(t, err) + assert.False(t, errs.IsRetryable(err)) +} + func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte) *consumermock.MockDelivery { t.Helper() d := consumermock.NewMockDelivery(ctrl) - d.EXPECT().Message().Return(entityqueue.NewMessage(testBuildID, payload, testBuildID, nil)).AnyTimes() + msg := entityqueue.NewMessage(testBuildID, payload, testBuildID, nil) + msg.Tenant = testQueue + d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() return d } @@ -456,6 +471,7 @@ func TestPublishRecordCarriesRequestID(t *testing.T) { assert.Equal(t, testID, payload.Id) assert.Equal(t, testID, got.ID) assert.Equal(t, testID, got.PartitionKey) + assert.Equal(t, "monorepo/main", got.Tenant) assert.Equal(t, "monorepo/main", got.Metadata[entityqueue.MetadataKeyQueueName]) } diff --git a/stovepipe/controller/dlq/BUILD.bazel b/stovepipe/controller/dlq/BUILD.bazel index 929be0543..cff04e401 100644 --- a/stovepipe/controller/dlq/BUILD.bazel +++ b/stovepipe/controller/dlq/BUILD.bazel @@ -11,6 +11,7 @@ go_library( importpath = "github.com/uber/submitqueue/stovepipe/controller/dlq", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//stovepipe/core/messagequeue:go_default_library", diff --git a/stovepipe/controller/dlq/build.go b/stovepipe/controller/dlq/build.go index 1bee58da3..02f28d35e 100644 --- a/stovepipe/controller/dlq/build.go +++ b/stovepipe/controller/dlq/build.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" stovepipemq "github.com/uber/submitqueue/stovepipe/core/messagequeue" @@ -60,11 +61,16 @@ func NewDLQBuildController( // Process drives the request named by a dead-lettered BuildRequest to failed. func (c *buildController) Process(ctx context.Context, delivery consumer.Delivery) error { + msg := delivery.Message() buildRequest := &stovepipemq.BuildRequest{} - if err := stovepipemq.Unmarshal(delivery.Message().Payload, buildRequest); err != nil { + if err := stovepipemq.Unmarshal(msg.Payload, buildRequest); err != nil { metrics.NamedCounter(c.metricsScope, _buildOpName, "deserialize_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("failed to decode dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, buildRequest.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, _buildOpName, "queue_identity_errors", 1, metrics.TagsFromContext(ctx)...) + return nil + } if buildRequest.Id == "" { metrics.NamedCounter(c.metricsScope, _buildOpName, "empty_id_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("build dlq payload decoded to empty request id") diff --git a/stovepipe/controller/dlq/build_test.go b/stovepipe/controller/dlq/build_test.go index fa9c68341..cbb1113b2 100644 --- a/stovepipe/controller/dlq/build_test.go +++ b/stovepipe/controller/dlq/build_test.go @@ -56,6 +56,7 @@ func TestBuildProcess(t *testing.T) { tests := []struct { name string payload []byte + tenant string setup func(m dlqMocks) wantErr bool wantMetric string @@ -79,6 +80,12 @@ func TestBuildProcess(t *testing.T) { wantErr: true, wantMetric: "test.build_dlq_controller.build_dlq.empty_id_errors+queue=monorepo/main", }, + { + name: "payload queue mismatch is acked", + tenant: "monorepo/release", + wantErr: false, + wantMetric: "test.build_dlq_controller.build_dlq.queue_identity_errors+queue=monorepo/main", + }, } for _, tt := range tests { @@ -93,7 +100,7 @@ func TestBuildProcess(t *testing.T) { if payload == nil { payload = buildPayload(t, testID) } - err := controller.Process(queueContext(), delivery(t, ctrl, payload)) + err := controller.Process(queueContext(), delivery(t, ctrl, payload, tt.tenant)) if tt.wantMetric != "" { counter, ok := mocks.metricsScope.Snapshot().Counters()[tt.wantMetric] require.True(t, ok) diff --git a/stovepipe/controller/dlq/buildsignal.go b/stovepipe/controller/dlq/buildsignal.go index c42629bad..189541c8d 100644 --- a/stovepipe/controller/dlq/buildsignal.go +++ b/stovepipe/controller/dlq/buildsignal.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" stovepipemq "github.com/uber/submitqueue/stovepipe/core/messagequeue" @@ -96,6 +97,10 @@ func (c *buildSignalController) Process(ctx context.Context, delivery consumer.D // without saying so. return fmt.Errorf("failed to decode dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, sig.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, _buildSignalOpName, "queue_identity_errors", 1, metrics.TagsFromContext(ctx)...) + return nil + } if sig.Id == "" { metrics.NamedCounter(c.metricsScope, _buildSignalOpName, "empty_id_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("dlq payload decoded to empty build id") diff --git a/stovepipe/controller/dlq/buildsignal_test.go b/stovepipe/controller/dlq/buildsignal_test.go index 8d3447072..278e5b764 100644 --- a/stovepipe/controller/dlq/buildsignal_test.go +++ b/stovepipe/controller/dlq/buildsignal_test.go @@ -84,6 +84,7 @@ func TestBuildSignalProcess(t *testing.T) { tests := []struct { name string payload []byte + tenant string setup func(m buildSignalDLQMocks) wantErr bool wantMetric string @@ -156,6 +157,13 @@ func TestBuildSignalProcess(t *testing.T) { wantErr: true, wantMetric: "test.buildsignal_dlq_controller.buildsignal_dlq.empty_id_errors+queue=monorepo/main", }, + { + name: "payload queue mismatch is acked", + tenant: "monorepo/release", + setup: func(buildSignalDLQMocks) {}, + wantErr: false, + wantMetric: "test.buildsignal_dlq_controller.buildsignal_dlq.queue_identity_errors+queue=monorepo/main", + }, } for _, tt := range tests { @@ -169,7 +177,7 @@ func TestBuildSignalProcess(t *testing.T) { payload = buildSignalPayload(t, testBuildID) } - err := c.Process(queueContext(), delivery(t, ctrl, payload)) + err := c.Process(queueContext(), delivery(t, ctrl, payload, tt.tenant)) if tt.wantMetric != "" { counter, ok := m.metricsScope.Snapshot().Counters()[tt.wantMetric] require.True(t, ok) diff --git a/stovepipe/controller/dlq/dlq_test.go b/stovepipe/controller/dlq/dlq_test.go index b296e2104..a3bf9e07f 100644 --- a/stovepipe/controller/dlq/dlq_test.go +++ b/stovepipe/controller/dlq/dlq_test.go @@ -73,10 +73,16 @@ func newController(t *testing.T, ctrl *gomock.Controller) (consumer.Controller, return c, m } -func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte) consumer.Delivery { +func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte, tenants ...string) consumer.Delivery { t.Helper() + tenant := testQueue + if len(tenants) > 0 && tenants[0] != "" { + tenant = tenants[0] + } + msg := entityqueue.NewMessage(testID, payload, testQueue, nil) + msg.Tenant = tenant d := consumermock.NewMockDelivery(ctrl) - d.EXPECT().Message().Return(entityqueue.NewMessage(testID, payload, testQueue, nil)).AnyTimes() + d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(4).AnyTimes() d.EXPECT().Metadata().Return(map[string]string{ "dlq.original_topic": "process", @@ -106,6 +112,7 @@ func TestProcess(t *testing.T) { tests := []struct { name string payload []byte + tenant string setup func(m dlqMocks) wantErr bool wantMetric string @@ -210,6 +217,13 @@ func TestProcess(t *testing.T) { wantErr: true, wantMetric: "test.process_dlq_controller.process_dlq.empty_id_errors+queue=monorepo/main", }, + { + name: "payload queue mismatch is acked", + tenant: "monorepo/release", + setup: func(m dlqMocks) {}, + wantErr: false, + wantMetric: "test.process_dlq_controller.process_dlq.queue_identity_errors+queue=monorepo/main", + }, } for _, tt := range tests { @@ -223,7 +237,7 @@ func TestProcess(t *testing.T) { payload = processPayload(t, testID) } - err := c.Process(queueContext(), delivery(t, ctrl, payload)) + err := c.Process(queueContext(), delivery(t, ctrl, payload, tt.tenant)) if tt.wantMetric != "" { counter, ok := m.metricsScope.Snapshot().Counters()[tt.wantMetric] require.True(t, ok) diff --git a/stovepipe/controller/dlq/request.go b/stovepipe/controller/dlq/request.go index 89f4a63cd..1a2079284 100644 --- a/stovepipe/controller/dlq/request.go +++ b/stovepipe/controller/dlq/request.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" stovepipemq "github.com/uber/submitqueue/stovepipe/core/messagequeue" @@ -84,6 +85,10 @@ func (c *requestController) Process(ctx context.Context, delivery consumer.Deliv // non-terminal. return fmt.Errorf("failed to decode dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, pr.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, _opName, "queue_identity_errors", 1, metrics.TagsFromContext(ctx)...) + return nil + } if pr.Id == "" { metrics.NamedCounter(c.metricsScope, _opName, "empty_id_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("dlq payload decoded to empty request id") diff --git a/stovepipe/controller/ingest.go b/stovepipe/controller/ingest.go index 1c9022f0a..2ffd6cb81 100644 --- a/stovepipe/controller/ingest.go +++ b/stovepipe/controller/ingest.go @@ -20,7 +20,6 @@ import ( "fmt" "github.com/uber-go/tally" - entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/extension/counter" @@ -56,13 +55,14 @@ func IsInvalidRequest(err error) bool { // the request ID onto the process stage. Ingestion is idempotent: a re-reported head resolves to // the already-minted request and republishes it only while it remains accepted. type IngestController struct { - logger *zap.SugaredLogger - metricsScope tally.Scope - counters counter.Factory - sourceControl sourcecontrol.Factory - stores storage.Factory - materializer requestlog.Materializer - registry consumer.TopicRegistry + logger *zap.SugaredLogger + metricsScope tally.Scope + counters counter.Factory + sourceControl sourcecontrol.Factory + stores storage.Factory + materializer requestlog.Materializer + registry consumer.TopicRegistry + configuredQueues map[string]struct{} } // NewIngestController creates a new instance of the stovepipe ingest controller. It publishes @@ -75,15 +75,21 @@ func NewIngestController( stores storage.Factory, materializer requestlog.Materializer, registry consumer.TopicRegistry, + configuredQueues []string, ) *IngestController { + queueNames := make(map[string]struct{}, len(configuredQueues)) + for _, queue := range configuredQueues { + queueNames[queue] = struct{}{} + } return &IngestController{ - logger: logger, - metricsScope: scope.SubScope("ingest_controller"), - counters: counters, - sourceControl: sourceControl, - stores: stores, - materializer: materializer, - registry: registry, + logger: logger, + metricsScope: scope.SubScope("ingest_controller"), + counters: counters, + sourceControl: sourceControl, + stores: stores, + materializer: materializer, + registry: registry, + configuredQueues: queueNames, } } @@ -106,6 +112,9 @@ func (c *IngestController) Ingest(ctx context.Context, req entity.IngestRequest) return entity.IngestResult{}, fmt.Errorf("requires the request to have a queue name specified: %w", ErrInvalidRequest) } queue := req.Queue + if _, ok := c.configuredQueues[queue]; !ok { + return entity.IngestResult{}, fmt.Errorf("queue %q is not configured: %w", queue, ErrInvalidRequest) + } store, err := c.stores.For(storage.Config{QueueName: queue}) if err != nil { @@ -312,8 +321,12 @@ func (c *IngestController) publishProcess(ctx context.Context, id, queue string) return fmt.Errorf("failed to serialize process request: %w", err) } - ctx = entityqueue.WithQueueName(ctx, queue) - if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyProcess, publish.IntentID(id), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyProcess, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(id), + Payload: payload, + PartitionKey: queue, + }); err != nil { return fmt.Errorf("failed to publish process request: %w", err) } return nil diff --git a/stovepipe/controller/ingest_test.go b/stovepipe/controller/ingest_test.go index f500110c6..74492d19d 100644 --- a/stovepipe/controller/ingest_test.go +++ b/stovepipe/controller/ingest_test.go @@ -97,7 +97,7 @@ func newIngestController(t *testing.T, ctrl *gomock.Controller) (*IngestControll }) require.NoError(t, err) - c := NewIngestController(zap.NewNop().Sugar(), tally.NewTestScope("test", nil), staticCounterFactory{counter: m.counter}, m.factory, staticStorageFactory{store: store}, m.materializer, registry) + c := NewIngestController(zap.NewNop().Sugar(), tally.NewTestScope("test", nil), staticCounterFactory{counter: m.counter}, m.factory, staticStorageFactory{store: store}, m.materializer, registry, []string{testQueue}) return c, m } @@ -155,6 +155,7 @@ func TestPublishProcessCarriesQueueMetadata(t *testing.T) { require.NoError(t, c.publishProcess(context.Background(), "request/monorepo/main/7", testQueue)) assert.Equal(t, testQueue, got.PartitionKey) + assert.Equal(t, testQueue, got.Tenant) assert.Equal(t, testQueue, got.Metadata[entityqueue.MetadataKeyQueueName]) } @@ -233,6 +234,13 @@ func TestIngestController_Ingest(t *testing.T) { wantErr: true, wantInvalid: true, }, + { + name: "unconfigured queue is invalid", + queue: "monorepo/unconfigured", + setup: func(m ingestMocks) {}, + wantErr: true, + wantInvalid: true, + }, { name: "unknown queue head is invalid", queue: testQueue, diff --git a/stovepipe/controller/process/BUILD.bazel b/stovepipe/controller/process/BUILD.bazel index ab1a36b30..73508f0a1 100644 --- a/stovepipe/controller/process/BUILD.bazel +++ b/stovepipe/controller/process/BUILD.bazel @@ -7,6 +7,7 @@ go_library( visibility = ["//visibility:public"], deps = [ "//api/base/hook:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/errs:go_default_library", "//platform/hook:go_default_library", diff --git a/stovepipe/controller/process/process.go b/stovepipe/controller/process/process.go index 75e264784..624ab7f1b 100644 --- a/stovepipe/controller/process/process.go +++ b/stovepipe/controller/process/process.go @@ -25,6 +25,7 @@ import ( "github.com/uber-go/tally" basehook "github.com/uber/submitqueue/api/base/hook" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" platformhook "github.com/uber/submitqueue/platform/hook" @@ -94,6 +95,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: a malformed message will never succeed regardless of retries. return fmt.Errorf("failed to deserialize process request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, pr.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: pr.GetQueueName()}) if err != nil { metrics.NamedCounter(c.metricsScope, _opName, "storage_resolve_errors", 1, metrics.TagsFromContext(ctx)...) @@ -483,7 +487,12 @@ func (c *Controller) publishBuild(ctx context.Context, id, queue string) error { return fmt.Errorf("failed to serialize build request: %w", err) } - if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuild, publish.IntentID(id), payload, id); err != nil { + if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuild, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(id), + Payload: payload, + PartitionKey: id, + }); err != nil { return fmt.Errorf("failed to publish build request: %w", err) } return nil @@ -498,7 +507,7 @@ func (c *Controller) publishBuild(ctx context.Context, id, queue string) error { // Partitioning by request id matches the process topic's own, carrying // per-request ordering across the seam. func (c *Controller) publishHookEvent(ctx context.Context, request entity.Request, event *basehook.HookEvent) error { - if err := platformhook.Publish(ctx, c.registry, event, request.ID); err != nil { + if err := platformhook.Publish(ctx, c.registry, request.Queue, event, request.ID); err != nil { metrics.NamedCounter(c.metricsScope, _opName, "hook_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("failed to announce %s for request %s: %w", event.GetType(), request.ID, err) } diff --git a/stovepipe/controller/process/process_test.go b/stovepipe/controller/process/process_test.go index 6328a754b..b98da7381 100644 --- a/stovepipe/controller/process/process_test.go +++ b/stovepipe/controller/process/process_test.go @@ -111,7 +111,9 @@ func newControllerWithScope(t *testing.T, ctrl *gomock.Controller, scope tally.S func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte) *consumermock.MockDelivery { t.Helper() d := consumermock.NewMockDelivery(ctrl) - d.EXPECT().Message().Return(entityqueue.NewMessage(testID, payload, testQueue, nil)).AnyTimes() + msg := entityqueue.NewMessage(testID, payload, testQueue, nil) + msg.Tenant = testQueue + d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() return d } @@ -123,6 +125,20 @@ func processPayload(t *testing.T, id string) []byte { return b } +func TestProcessRejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c, _ := newController(t, ctrl) + msg := entityqueue.NewMessage(testID, processPayload(t, testID), testQueue, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + err := c.Process(queueContext(testQueue), d) + + require.Error(t, err) + assert.False(t, errs.IsRetryable(err)) +} + func acceptedRequest(id string) entity.Request { return entity.Request{ ID: id, @@ -183,6 +199,7 @@ func expectStartValidationAnnounce(t *testing.T, m processMocks, id string) { Publish(gomock.Any(), "stovepipe-hook", gomock.AssignableToTypeOf(entityqueue.Message{})). DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { assert.Equal(t, id, msg.PartitionKey) + assert.Equal(t, testQueue, msg.Tenant) event := &basehook.HookEvent{} require.NoError(t, basehook.Unmarshal(msg.Payload, event)) assert.Equal(t, string(hookevent.TypeValidationRepositoryStarted), event.GetType()) @@ -202,6 +219,7 @@ func expectBuildPublish(t *testing.T, m processMocks, id string) { DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { assert.Equal(t, id, msg.ID) assert.Equal(t, id, msg.PartitionKey) + assert.Equal(t, testQueue, msg.Tenant) assert.Equal(t, testQueue, msg.Metadata[entityqueue.MetadataKeyQueueName]) buildReq := &stovepipemq.BuildRequest{} require.NoError(t, stovepipemq.Unmarshal(msg.Payload, buildReq)) diff --git a/stovepipe/controller/record/BUILD.bazel b/stovepipe/controller/record/BUILD.bazel index 378ebfe32..432f52668 100644 --- a/stovepipe/controller/record/BUILD.bazel +++ b/stovepipe/controller/record/BUILD.bazel @@ -7,6 +7,7 @@ go_library( visibility = ["//visibility:public"], deps = [ "//api/base/hook:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/hook:go_default_library", "//platform/metrics:go_default_library", diff --git a/stovepipe/controller/record/record.go b/stovepipe/controller/record/record.go index 93366aec7..12b207b4a 100644 --- a/stovepipe/controller/record/record.go +++ b/stovepipe/controller/record/record.go @@ -35,6 +35,7 @@ import ( "github.com/uber-go/tally" basehook "github.com/uber/submitqueue/api/base/hook" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" platformhook "github.com/uber/submitqueue/platform/hook" "github.com/uber/submitqueue/platform/metrics" @@ -110,6 +111,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: a malformed message will never succeed regardless of retries. return fmt.Errorf("failed to deserialize record: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, rec.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: rec.GetQueueName()}) if err != nil { metrics.NamedCounter(c.metricsScope, _opName, "storage_resolve_errors", 1, metrics.TagsFromContext(ctx)...) @@ -489,7 +493,7 @@ func (c *Controller) promote(ctx context.Context, request entity.Request) error // Partitioning by request id matches the record topic's own, carrying // per-request ordering across the seam. func (c *Controller) publishHookEvent(ctx context.Context, request entity.Request, event *basehook.HookEvent) error { - if err := platformhook.Publish(ctx, c.registry, event, request.ID); err != nil { + if err := platformhook.Publish(ctx, c.registry, request.Queue, event, request.ID); err != nil { metrics.NamedCounter(c.metricsScope, _opName, "hook_errors", 1, metrics.TagsFromContext(ctx)...) return fmt.Errorf("failed to announce %s for request %s: %w", event.GetType(), request.ID, err) } diff --git a/stovepipe/controller/record/record_test.go b/stovepipe/controller/record/record_test.go index e4eee5b87..ca1658b7b 100644 --- a/stovepipe/controller/record/record_test.go +++ b/stovepipe/controller/record/record_test.go @@ -78,6 +78,7 @@ type recordMocks struct { // plumbing carrying it. Setting err makes the publish fail. type hookRecorder struct { events []*basehook.HookEvent + tenant string err error } @@ -160,6 +161,7 @@ func newControllerForTopic(t *testing.T, ctrl *gomock.Controller, topicKey consu return err } m.hooks.events = append(m.hooks.events, event) + m.hooks.tenant = msg.Tenant return nil }).AnyTimes() @@ -195,7 +197,9 @@ func TestControllerIdentity(t *testing.T) { func delivery(t *testing.T, ctrl *gomock.Controller, payload []byte) *consumermock.MockDelivery { t.Helper() d := consumermock.NewMockDelivery(ctrl) - d.EXPECT().Message().Return(entityqueue.NewMessage(testID, payload, testID, nil)).AnyTimes() + msg := entityqueue.NewMessage(testID, payload, testID, nil) + msg.Tenant = testQueue + d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() return d } @@ -287,6 +291,7 @@ func TestProcess_AdvancesBookmarkOnSuccess(t *testing.T) { m.sourceControl.EXPECT().Promote(gomock.Any(), testURI).Return(nil) require.NoError(t, c.Process(queueContext(), delivery(t, ctrl, recordPayload(t, testID)))) + assert.Equal(t, testQueue, m.hooks.tenant) assert.Equal(t, tt.wantURI, written.LastGreenURI) assert.Equal(t, testID, written.LastGreenRequestID) @@ -976,6 +981,17 @@ func TestProcess_MalformedPayloadFails(t *testing.T) { require.Error(t, c.Process(queueContext(), delivery(t, ctrl, []byte("not-protojson")))) } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c, _ := newController(t, ctrl) + msg := entityqueue.NewMessage(testID, recordPayload(t, testID), testID, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, c.Process(queueContext(), d)) +} + func TestProcess_QueueMismatchFails(t *testing.T) { ctrl := gomock.NewController(t) c, m := newController(t, ctrl) @@ -987,6 +1003,10 @@ func TestProcess_QueueMismatchFails(t *testing.T) { payload, err := stovepipemq.Marshal(&stovepipemq.Record{Id: testID, QueueName: "monorepo/other"}) require.NoError(t, err) + msg := entityqueue.NewMessage(testID, payload, testID, nil) + msg.Tenant = "monorepo/other" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() - require.Error(t, c.Process(queueContext(), delivery(t, ctrl, payload))) + require.Error(t, c.Process(queueContext(), d)) } diff --git a/submitqueue/core/request/log.go b/submitqueue/core/request/log.go index 598d9a1c7..94d606719 100644 --- a/submitqueue/core/request/log.go +++ b/submitqueue/core/request/log.go @@ -52,8 +52,12 @@ func PublishLog(ctx context.Context, registry consumer.TopicRegistry, logEntry e cause = append(cause, occurrence) } - if err := publish.Message(ctx, registry, topickey.TopicKeyLog, - publish.IntentID(logEntry.RequestID, cause...), payload, partitionKey); err != nil { + if err := publish.Message(ctx, registry, topickey.TopicKeyLog, publish.MessageParams{ + Tenant: logEntry.Queue, + ID: publish.IntentID(logEntry.RequestID, cause...), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/core/request/log_test.go b/submitqueue/core/request/log_test.go index ea3314795..044aac039 100644 --- a/submitqueue/core/request/log_test.go +++ b/submitqueue/core/request/log_test.go @@ -128,10 +128,12 @@ func TestPublishLog_MessageIDScopedByStatus(t *testing.T) { ctrl := gomock.NewController(t) var ids []string + var tenants []string mockPub := queuemock.NewMockPublisher(ctrl) mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, _ string, msg entityqueue.Message) error { ids = append(ids, msg.ID) + tenants = append(tenants, msg.Tenant) return nil }, ).AnyTimes() @@ -161,6 +163,7 @@ func TestPublishLog_MessageIDScopedByStatus(t *testing.T) { "req/1/cancelled", "req/1/started", }, ids) + require.Equal(t, []string{"req", "req", "req", "req"}, tenants) } // TestPublishLog_MessageIDScopedByOccurrence locks in the recurring half of the diff --git a/submitqueue/core/request/terminate_test.go b/submitqueue/core/request/terminate_test.go index ac98c0098..fac4cfc09 100644 --- a/submitqueue/core/request/terminate_test.go +++ b/submitqueue/core/request/terminate_test.go @@ -124,7 +124,7 @@ func TestTerminateRequest(t *testing.T) { "already in target terminal state republishes log": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - already := entity.Request{ID: requestID, State: entity.RequestStateError, Version: 5} + already := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateError, Version: 5} rs.EXPECT().Get(gomock.Any(), requestID).Return(already, nil) }, wantResult: TerminationResult{ @@ -140,7 +140,7 @@ func TestTerminateRequest(t *testing.T) { "already in target terminal state returns republish error": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - already := entity.Request{ID: requestID, State: entity.RequestStateError, Version: 5} + already := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateError, Version: 5} rs.EXPECT().Get(gomock.Any(), requestID).Return(already, nil) }, publishErr: fmt.Errorf("connection refused"), @@ -150,7 +150,7 @@ func TestTerminateRequest(t *testing.T) { "diverged terminal state is left untouched": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - landed := entity.Request{ID: requestID, State: entity.RequestStateLanded, Version: 7} + landed := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateLanded, Version: 7} rs.EXPECT().Get(gomock.Any(), requestID).Return(landed, nil) }, wantResult: TerminationResult{ diff --git a/submitqueue/gateway/controller/cancel.go b/submitqueue/gateway/controller/cancel.go index 644ae3448..cbaafd667 100644 --- a/submitqueue/gateway/controller/cancel.go +++ b/submitqueue/gateway/controller/cancel.go @@ -139,8 +139,12 @@ func (c *cancelController) publishToQueue(ctx context.Context, cancelRequest ent // The request ID is the message ID with no cause: a request is cancelled at // most once, so a second Cancel for one already being cancelled is meant to // dedup rather than enqueue redundant work. - if err := publish.Message(ctx, c.registry, topickey.TopicKeyCancel, - publish.IntentID(cancelRequest.ID), payload, cancelRequest.ID); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyCancel, publish.MessageParams{ + Tenant: cancelRequest.Queue, + ID: publish.IntentID(cancelRequest.ID), + Payload: payload, + PartitionKey: cancelRequest.ID, + }); err != nil { return fmt.Errorf("failed to publish cancel request message: %w", err) } diff --git a/submitqueue/gateway/controller/cancel_test.go b/submitqueue/gateway/controller/cancel_test.go index 45cefe752..554079906 100644 --- a/submitqueue/gateway/controller/cancel_test.go +++ b/submitqueue/gateway/controller/cancel_test.go @@ -153,6 +153,7 @@ func TestCancel_PublishesToQueue(t *testing.T) { assert.Equal(t, "cancel", publishedTopic) assert.Equal(t, "my-queue/7", publishedMessage.ID) + assert.Equal(t, "my-queue", publishedMessage.Tenant) assert.Equal(t, "my-queue/7", publishedMessage.PartitionKey) deserialized, err := entity.CancelRequestFromBytes(publishedMessage.Payload) diff --git a/submitqueue/gateway/controller/land.go b/submitqueue/gateway/controller/land.go index 0a1a3708d..d03b537a8 100644 --- a/submitqueue/gateway/controller/land.go +++ b/submitqueue/gateway/controller/land.go @@ -212,8 +212,12 @@ func (c *landController) publishToQueue(ctx context.Context, landRequest entity. // retry of this same publish dedups instead of enqueuing it twice // - Payload: serialized LandRequest entity // - Partition key: landRequest.Queue (ensures ordering per queue) - if err := publish.Message(ctx, c.registry, topickey.TopicKeyStart, - publish.IntentID(landRequest.ID), payload, landRequest.Queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyStart, publish.MessageParams{ + Tenant: landRequest.Queue, + ID: publish.IntentID(landRequest.ID), + Payload: payload, + PartitionKey: landRequest.Queue, + }); err != nil { return fmt.Errorf("failed to publish land request message: %w", err) } diff --git a/submitqueue/gateway/controller/land_test.go b/submitqueue/gateway/controller/land_test.go index c957ef23b..52b25f8ee 100644 --- a/submitqueue/gateway/controller/land_test.go +++ b/submitqueue/gateway/controller/land_test.go @@ -459,6 +459,7 @@ func TestLand_PublishesToQueue(t *testing.T) { // Verify message was published to the topic registered under TopicKeyStart assert.Equal(t, "start", publishedTopic) assert.Equal(t, "test-queue/123", publishedMessage.ID) + assert.Equal(t, "test-queue", publishedMessage.Tenant) assert.Equal(t, "test-queue", publishedMessage.PartitionKey) // Verify payload can be deserialized diff --git a/submitqueue/gateway/controller/log/BUILD.bazel b/submitqueue/gateway/controller/log/BUILD.bazel index 829cfb333..3c22824ae 100644 --- a/submitqueue/gateway/controller/log/BUILD.bazel +++ b/submitqueue/gateway/controller/log/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/gateway/controller/log", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//submitqueue/core/request:go_default_library", diff --git a/submitqueue/gateway/controller/log/log.go b/submitqueue/gateway/controller/log/log.go index 15a63b198..ef02de7ff 100644 --- a/submitqueue/gateway/controller/log/log.go +++ b/submitqueue/gateway/controller/log/log.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" requestcore "github.com/uber/submitqueue/submitqueue/core/request" @@ -76,6 +77,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: malformed messages will never succeed regardless of retry count return fmt.Errorf("failed to deserialize request log: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, logEntry.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } c.logger.Debugw("received request log entry", "request_id", logEntry.RequestID, diff --git a/submitqueue/gateway/controller/log/log_test.go b/submitqueue/gateway/controller/log/log_test.go index d943e70ee..01cea18f1 100644 --- a/submitqueue/gateway/controller/log/log_test.go +++ b/submitqueue/gateway/controller/log/log_test.go @@ -103,6 +103,9 @@ func TestController_Process(t *testing.T) { } controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, tt.setupStore(ctrl), topickey.TopicKeyLog, "gateway-log") msg := entityqueue.NewMessage("test-queue/1", payload, "test-queue", nil) + if tt.logEntry != nil { + msg.Tenant = tt.logEntry.Queue + } delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -117,6 +120,20 @@ func TestController_Process(t *testing.T) { } } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + logEntry := newRequestLog("test-queue/1", entity.RequestStatusStarted, 1, "", nil) + payload, err := logEntry.ToBytes() + require.NoError(t, err) + controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, newUnusedMaterializer(ctrl), topickey.TopicKeyLog, "gateway-log") + msg := entityqueue.NewMessage(logEntry.RequestID, payload, logEntry.Queue, nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + func newLogControllerStore(ctrl *gomock.Controller, insertErr, getErr, updateErr, queueErr error) *requestcore.Materializer { store := storagemock.NewMockStorage(ctrl) logStore := storagemock.NewMockRequestLogStore(ctrl) diff --git a/submitqueue/orchestrator/controller/batch/BUILD.bazel b/submitqueue/orchestrator/controller/batch/BUILD.bazel index 0270e7e89..aac1f9e63 100644 --- a/submitqueue/orchestrator/controller/batch/BUILD.bazel +++ b/submitqueue/orchestrator/controller/batch/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/batch", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/extension/counter:go_default_library", "//platform/metrics:go_default_library", diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 0b60b32d6..dec1c73fc 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/extension/counter" "github.com/uber/submitqueue/platform/metrics" @@ -87,6 +88,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize request ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, rid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: rid.Queue}) if err != nil { @@ -211,8 +215,12 @@ func (c *Controller) publishToDependencyAnalysis(ctx context.Context, batch enti return fmt.Errorf("failed to serialize batch ID: %w", err) } - if err := publish.Message(ctx, c.registry, topickey.TopicKeyDependencyAnalysis, - publish.IntentID(batch.ID), payload, batch.Queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyDependencyAnalysis, publish.MessageParams{ + Tenant: batch.Queue, + ID: publish.IntentID(batch.ID), + Payload: payload, + PartitionKey: batch.Queue, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index 68bfc7820..4227deb8b 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -40,8 +40,8 @@ import ( ) // requestIDPayload serializes a RequestID to JSON bytes for test message payloads. -func requestIDPayload(t *testing.T, id string) []byte { - payload, err := entity.RequestID{ID: id}.ToBytes() +func requestIDPayload(t *testing.T, id, queue string) []byte { + payload, err := entity.RequestID{ID: id, Queue: queue}.ToBytes() require.NoError(t, err) return payload } @@ -127,13 +127,16 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M func newDelivery(t *testing.T, ctrl *gomock.Controller, request entity.Request, payloadQueue string) *consumermock.MockDelivery { t.Helper() - payload := requestIDPayload(t, request.ID) + payload := requestIDPayload(t, request.ID, request.Queue) + tenant := request.Queue if payloadQueue != "" { bytes, err := entity.RequestID{ID: request.ID, Queue: payloadQueue}.ToBytes() require.NoError(t, err) payload = bytes + tenant = payloadQueue } msg := entityqueue.NewMessage(request.ID, payload, request.Queue, nil) + msg.Tenant = tenant delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -159,6 +162,19 @@ func TestController_Process_Success(t *testing.T) { require.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, testRequest(), ""))) } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + controller := newTestController(t, ctrl, newSequentialCounter(ctrl), storagemock.NewMockStorage(ctrl), nil) + request := testRequest() + payload := requestIDPayload(t, request.ID, request.Queue) + msg := entityqueue.NewMessage(request.ID, payload, request.Queue, nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + // A payload whose queue disagrees with the request's authoritative queue is // rejected without touching the counter or the batch store. func TestController_Process_QueueMismatchRejected(t *testing.T) { diff --git a/submitqueue/orchestrator/controller/build/BUILD.bazel b/submitqueue/orchestrator/controller/build/BUILD.bazel index ca53a5de8..1803101cf 100644 --- a/submitqueue/orchestrator/controller/build/BUILD.bazel +++ b/submitqueue/orchestrator/controller/build/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/build", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/build/build.go b/submitqueue/orchestrator/controller/build/build.go index afa80cc99..207646e3c 100644 --- a/submitqueue/orchestrator/controller/build/build.go +++ b/submitqueue/orchestrator/controller/build/build.go @@ -34,6 +34,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -105,6 +106,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize batch ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: bid.Queue}) if err != nil { @@ -381,7 +385,12 @@ func (c *Controller) publishBuildSignal(ctx context.Context, buildID, queue stri return fmt.Errorf("failed to serialize build ID: %w", err) } - if err := publish.Message(ctx, c.registry, topickey.TopicKeyBuildSignal, publish.IntentID(buildID), payload, buildID); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyBuildSignal, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(buildID), + Payload: payload, + PartitionKey: buildID, + }); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to buildsignal: %w", err) } diff --git a/submitqueue/orchestrator/controller/build/build_test.go b/submitqueue/orchestrator/controller/build/build_test.go index dcb9a69cb..74f23690b 100644 --- a/submitqueue/orchestrator/controller/build/build_test.go +++ b/submitqueue/orchestrator/controller/build/build_test.go @@ -58,7 +58,7 @@ func (f staticBuildRunnerFactory) For(buildrunner.Config) (buildrunner.BuildRunn func batchIDPayload(t *testing.T, id string) []byte { t.Helper() - payload, err := entity.BatchID{ID: id}.ToBytes() + payload, err := entity.BatchID{ID: id, Queue: "test-queue"}.ToBytes() require.NoError(t, err) return payload } @@ -154,12 +154,24 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, batch entity.Batch func processAttempt(t *testing.T, ctrl *gomock.Controller, c *Controller, attempt int) error { t.Helper() msg := entityqueue.NewMessage("msg-1", batchIDPayload(t, headID), "test-queue", nil) + msg.Tenant = "test-queue" d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(attempt).AnyTimes() return c.Process(context.Background(), d) } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c, _ := newTestController(t, ctrl, headBatch(entity.BatchStateCreated)) + msg := entityqueue.NewMessage("msg-1", batchIDPayload(t, headID), "test-queue", nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, c.Process(context.Background(), d)) +} + // process delivers the head's batch ID as a first delivery. func process(t *testing.T, ctrl *gomock.Controller, c *Controller) error { t.Helper() @@ -175,6 +187,7 @@ func expectSignal(t *testing.T, deps *testDeps, buildID string) { got, err := entity.BuildIDFromBytes(msg.Payload) require.NoError(t, err) assert.Equal(t, buildID, got.ID) + assert.Equal(t, "test-queue", msg.Tenant) assert.Equal(t, buildID, msg.PartitionKey, "polls partition per build so one slow build cannot block a head's others") return nil diff --git a/submitqueue/orchestrator/controller/buildsignal/BUILD.bazel b/submitqueue/orchestrator/controller/buildsignal/BUILD.bazel index 1b5d0e82e..9022b3dda 100644 --- a/submitqueue/orchestrator/controller/buildsignal/BUILD.bazel +++ b/submitqueue/orchestrator/controller/buildsignal/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/buildsignal", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/buildsignal/buildsignal.go b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go index 9c28b8693..18ce0e052 100644 --- a/submitqueue/orchestrator/controller/buildsignal/buildsignal.go +++ b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go @@ -46,6 +46,7 @@ import ( "strconv" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -143,6 +144,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: malformed messages will never succeed. return fmt.Errorf("failed to deserialize build ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, buildID.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: buildID.Queue}) if err != nil { @@ -440,7 +444,12 @@ func (c *Controller) publishBatchID(ctx context.Context, key consumer.TopicKey, if err != nil { return fmt.Errorf("failed to serialize batch ID: %w", err) } - return publish.Message(ctx, c.registry, key, msgID, payload, queue) + return publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: msgID, + Payload: payload, + PartitionKey: queue, + }) } // Name returns the controller name for logging and metrics. diff --git a/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go b/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go index 9e02cda19..c9184c91f 100644 --- a/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go +++ b/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go @@ -166,9 +166,10 @@ func (h *testHarness) wanted() { // is returned so tests can expect a Hold for the next poll. func delivery(t *testing.T, ctrl *gomock.Controller) *consumermock.MockDelivery { t.Helper() - payload, err := entity.BuildID{ID: testBuildID}.ToBytes() + payload, err := entity.BuildID{ID: testBuildID, Queue: "test-queue"}.ToBytes() require.NoError(t, err) msg := entityqueue.NewMessage(testBuildID, payload, testBuildID, nil) + msg.Tenant = "test-queue" d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -197,6 +198,19 @@ func TestController_Identity(t *testing.T) { var _ consumer.Controller = h.controller } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl, entity.BatchStateSpeculating) + payload, err := entity.BuildID{ID: testBuildID, Queue: "test-queue"}.ToBytes() + require.NoError(t, err) + msg := entityqueue.NewMessage(testBuildID, payload, testBuildID, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, h.controller.Process(context.Background(), d)) +} + // The poll loop records status on the build and wakes the run. It must never // write the path set — that is the speculate run's state, and it is read here // only as the kill list. @@ -216,7 +230,7 @@ func TestProcess_RecordsStatusAndNeverWritesThePathSet(t *testing.T) { h := newTestHarness(t, ctrl, entity.BatchStateSpeculating) h.builds.EXPECT().Get(gomock.Any(), testBuildID).Return(testBuild(entity.BuildStatusRunning), nil) - h.br.EXPECT().Status(gomock.Any(), entity.BuildID{ID: testBuildID}).Return(tt.status, nil, nil) + h.br.EXPECT().Status(gomock.Any(), entity.BuildID{ID: testBuildID, Queue: "test-queue"}).Return(tt.status, nil, nil) h.builds.EXPECT().Update(gomock.Any(), testBuild(tt.status)).Return(nil) h.speculatePub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil) @@ -253,6 +267,7 @@ func TestProcess_SpeculateWakeUpIsNamedForTheObservedStatus(t *testing.T) { h.speculatePub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).DoAndReturn( func(_ context.Context, _ string, msg entityqueue.Message) error { id = msg.ID + assert.Equal(t, "test-queue", msg.Tenant) return nil }, ) @@ -384,7 +399,7 @@ func TestProcess_StopsUnwantedBuilds(t *testing.T) { h.builds.EXPECT().Get(gomock.Any(), testBuildID).Return(testBuild(entity.BuildStatusRunning), nil) h.br.EXPECT().Status(gomock.Any(), gomock.Any()).Return(entity.BuildStatusRunning, nil, nil) - h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID}).Return(nil) + h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID, Queue: "test-queue"}).Return(nil) h.speculatePub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil) d := delivery(t, ctrl) @@ -497,7 +512,7 @@ func TestProcess_CancelFailureDoesNotFailThePoll(t *testing.T) { h.builds.EXPECT().Get(gomock.Any(), testBuildID).Return(testBuild(entity.BuildStatusRunning), nil) h.br.EXPECT().Status(gomock.Any(), gomock.Any()).Return(entity.BuildStatusRunning, nil, nil) - h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID}). + h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID, Queue: "test-queue"}). Return(errors.New("runner unavailable")) h.speculatePub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil) @@ -539,7 +554,7 @@ func TestProcess_HaltedBatchStillRuns(t *testing.T) { h.builds.EXPECT().Get(gomock.Any(), testBuildID).Return(testBuild(entity.BuildStatusRunning), nil) h.br.EXPECT().Status(gomock.Any(), gomock.Any()).Return(entity.BuildStatusRunning, nil, nil) - h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID}).Return(nil) + h.br.EXPECT().Cancel(gomock.Any(), entity.BuildID{ID: testBuildID, Queue: "test-queue"}).Return(nil) h.speculatePub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil) d := delivery(t, ctrl) d.EXPECT().Hold(PollDelayRunningMs) diff --git a/submitqueue/orchestrator/controller/cancel/BUILD.bazel b/submitqueue/orchestrator/controller/cancel/BUILD.bazel index df905bd86..8dfe860ae 100644 --- a/submitqueue/orchestrator/controller/cancel/BUILD.bazel +++ b/submitqueue/orchestrator/controller/cancel/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/cancel", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 243e3af32..aa571c20d 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -55,6 +55,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -110,6 +111,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize cancel request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, cancelReq.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: cancelReq.Queue}) if err != nil { @@ -350,7 +354,12 @@ func (c *Controller) publishBatchID(ctx context.Context, key consumer.TopicKey, if err != nil { return fmt.Errorf("failed to serialize batch ID: %w", err) } - return publish.Message(ctx, c.registry, key, publish.UniqueID(batchID), payload, queue) + return publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: publish.UniqueID(batchID), + Payload: payload, + PartitionKey: queue, + }) } // Name returns the controller name for logging and metrics. diff --git a/submitqueue/orchestrator/controller/cancel/cancel_test.go b/submitqueue/orchestrator/controller/cancel/cancel_test.go index 01274bafa..f1eac4dab 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel_test.go +++ b/submitqueue/orchestrator/controller/cancel/cancel_test.go @@ -61,7 +61,7 @@ func requestWithState(request entity.Request, state entity.RequestState) entity. // cancelPayload serializes a CancelRequest to JSON bytes for test message payloads. func cancelPayload(t *testing.T, id, reason string) []byte { - payload, err := entity.CancelRequest{ID: id, Reason: reason}.ToBytes() + payload, err := entity.CancelRequest{ID: id, Queue: "q", Reason: reason}.ToBytes() require.NoError(t, err) return payload } @@ -91,6 +91,7 @@ func newController(t *testing.T, store storage.Storage, registry consumer.TopicR func newDelivery(t *testing.T, ctrl *gomock.Controller, payload []byte, partitionKey string) consumer.Delivery { msg := entityqueue.NewMessage("cancel-msg", payload, partitionKey, nil) + msg.Tenant = "q" d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -136,6 +137,17 @@ func TestNewController(t *testing.T) { var _ consumer.Controller = controller } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + controller := newController(t, storagemock.NewMockStorage(ctrl), consumer.TopicRegistry{}) + msg := entityqueue.NewMessage("cancel-msg", cancelPayload(t, "q/1", ""), "q", nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), d)) +} + func TestProcess_AlreadyTerminal_NoOp(t *testing.T) { ctrl := gomock.NewController(t) registry, pub := newRegistry(t, ctrl) diff --git a/submitqueue/orchestrator/controller/conclude/BUILD.bazel b/submitqueue/orchestrator/controller/conclude/BUILD.bazel index 91992843d..493cdfc06 100644 --- a/submitqueue/orchestrator/controller/conclude/BUILD.bazel +++ b/submitqueue/orchestrator/controller/conclude/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/conclude", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//submitqueue/core/request:go_default_library", diff --git a/submitqueue/orchestrator/controller/conclude/conclude.go b/submitqueue/orchestrator/controller/conclude/conclude.go index 8f7c28d38..171ee93b5 100644 --- a/submitqueue/orchestrator/controller/conclude/conclude.go +++ b/submitqueue/orchestrator/controller/conclude/conclude.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" corerequest "github.com/uber/submitqueue/submitqueue/core/request" @@ -74,6 +75,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, "process", "deserialize_errors", 1) return fmt.Errorf("failed to deserialize batch ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: bid.Queue}) if err != nil { diff --git a/submitqueue/orchestrator/controller/conclude/conclude_test.go b/submitqueue/orchestrator/controller/conclude/conclude_test.go index f39910da0..03fad71b3 100644 --- a/submitqueue/orchestrator/controller/conclude/conclude_test.go +++ b/submitqueue/orchestrator/controller/conclude/conclude_test.go @@ -48,7 +48,7 @@ func requestWithState(request entity.Request, state entity.RequestState) entity. // batchIDPayload serializes a BatchID to JSON bytes for test message payloads. func batchIDPayload(t *testing.T, id string) []byte { - payload, err := entity.BatchID{ID: id}.ToBytes() + payload, err := entity.BatchID{ID: id, Queue: "test-queue"}.ToBytes() require.NoError(t, err) return payload } @@ -96,6 +96,17 @@ func TestNewController(t *testing.T) { assert.Equal(t, "conclude", controller.Name()) } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + controller, _ := newTestController(t, ctrl, nil, false) + msg := entityqueue.NewMessage("test-queue/batch/1", batchIDPayload(t, "test-queue/batch/1"), "test-queue", nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + func TestController_Process(t *testing.T) { tests := []struct { name string @@ -126,12 +137,12 @@ func TestController_Process(t *testing.T) { mockRequestStore := storagemock.NewMockRequestStore(ctrl) request1 := entity.Request{ - ID: "test-queue/1", Version: 2, State: entity.RequestStateProcessing, + ID: "test-queue/1", Queue: "test-queue", Version: 2, State: entity.RequestStateProcessing, } mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(request1, nil) mockRequestStore.EXPECT().Update(gomock.Any(), requestWithState(request1, entity.RequestStateLanded), int32(2), int32(3)).Return(nil) request2 := entity.Request{ - ID: "test-queue/2", Version: 3, State: entity.RequestStateProcessing, + ID: "test-queue/2", Queue: "test-queue", Version: 3, State: entity.RequestStateProcessing, } mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/2").Return(request2, nil) mockRequestStore.EXPECT().Update(gomock.Any(), requestWithState(request2, entity.RequestStateLanded), int32(3), int32(4)).Return(nil) @@ -164,7 +175,7 @@ func TestController_Process(t *testing.T) { mockRequestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "test-queue/5", Version: 1, State: entity.RequestStateProcessing, + ID: "test-queue/5", Queue: "test-queue", Version: 1, State: entity.RequestStateProcessing, } mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/5").Return(request, nil) mockRequestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -197,7 +208,7 @@ func TestController_Process(t *testing.T) { mockRequestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "test-queue/10", Version: 4, State: entity.RequestStateProcessing, + ID: "test-queue/10", Queue: "test-queue", Version: 4, State: entity.RequestStateProcessing, } mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/10").Return(request, nil) mockRequestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateCancelled), int32(4), int32(5)).Return(nil) @@ -232,7 +243,7 @@ func TestController_Process(t *testing.T) { // must NOT be called — gomock will fail the test if it is. mockRequestStore := storagemock.NewMockRequestStore(ctrl) mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/20").Return(entity.Request{ - ID: "test-queue/20", Version: 7, State: entity.RequestStateLanded, + ID: "test-queue/20", Queue: "test-queue", Version: 7, State: entity.RequestStateLanded, }, nil) mockStorage := storagemock.NewMockStorage(ctrl) @@ -266,7 +277,7 @@ func TestController_Process(t *testing.T) { // and must not attempt UpdateState. mockRequestStore := storagemock.NewMockRequestStore(ctrl) mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/30").Return(entity.Request{ - ID: "test-queue/30", Version: 5, State: entity.RequestStateCancelled, + ID: "test-queue/30", Queue: "test-queue", Version: 5, State: entity.RequestStateCancelled, }, nil) mockStorage := storagemock.NewMockStorage(ctrl) @@ -385,7 +396,7 @@ func TestController_Process(t *testing.T) { mockRequestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "test-queue/1", Version: 2, State: entity.RequestStateProcessing, + ID: "test-queue/1", Queue: "test-queue", Version: 2, State: entity.RequestStateProcessing, } mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(request, nil) mockRequestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateLanded), int32(2), int32(3)).Return(storage.ErrVersionMismatch) @@ -434,6 +445,7 @@ func TestController_Process(t *testing.T) { controller, _ := newTestController(t, ctrl, mockStorage, tt.expectLogPublish) msg := entityqueue.NewMessage(tt.batch.ID, batchIDPayload(t, tt.batch.ID), tt.batch.Queue, nil) + msg.Tenant = tt.batch.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -491,6 +503,7 @@ func TestController_Process_FailedBatchCarriesReasonToRequestLog(t *testing.T) { ) msg := entityqueue.NewMessage(batch.ID, batchIDPayload(t, batch.ID), batch.Queue, map[string]string{topickey.MetadataKeyFailureReason: reason}) + msg.Tenant = batch.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -513,6 +526,7 @@ func TestController_Process_StorageFailure(t *testing.T) { controller, _ := newTestController(t, ctrl, mockStorage, false) msg := entityqueue.NewMessage("test-queue/batch/1", batchIDPayload(t, "test-queue/batch/1"), "test-queue", nil) + msg.Tenant = "test-queue" delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() diff --git a/submitqueue/orchestrator/controller/dependencyanalysis/BUILD.bazel b/submitqueue/orchestrator/controller/dependencyanalysis/BUILD.bazel index fe5b08e60..79c891d48 100644 --- a/submitqueue/orchestrator/controller/dependencyanalysis/BUILD.bazel +++ b/submitqueue/orchestrator/controller/dependencyanalysis/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/dependencyanalysis", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go index cb36da1d8..6d00894a6 100644 --- a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go +++ b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go @@ -54,6 +54,7 @@ import ( "slices" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -114,6 +115,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize batch ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: bid.Queue}) if err != nil { @@ -440,7 +444,12 @@ func (c *Controller) publishToSpeculate(ctx context.Context, batch entity.Batch) return fmt.Errorf("failed to serialize batch ID: %w", err) } - if err := publish.Message(ctx, c.registry, topickey.TopicKeySpeculate, publish.IntentID(batch.ID), payload, batch.Queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeySpeculate, publish.MessageParams{ + Tenant: batch.Queue, + ID: publish.IntentID(batch.ID), + Payload: payload, + PartitionKey: batch.Queue, + }); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish batch ID to speculate topic: %w", err) } diff --git a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis_test.go b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis_test.go index 1a04e7032..d06363c9c 100644 --- a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis_test.go +++ b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis_test.go @@ -137,6 +137,7 @@ func newDelivery(t *testing.T, ctrl *gomock.Controller, id, queue string) *consu t.Helper() msg := entityqueue.NewMessage(id, batchIDPayload(t, id, queue), queue, nil) + msg.Tenant = queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -176,6 +177,17 @@ func TestNewController(t *testing.T) { var _ consumer.Controller = controller } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + controller := newTestController(t, ctrl, storagemock.NewMockStorage(ctrl), nil, nil) + msg := entityqueue.NewMessage("test-queue/batch/1", batchIDPayload(t, "test-queue/batch/1", "test-queue"), "test-queue", nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + // The whole point of the stage: resolve what the batch must serialize behind, // record it on both sides of the graph, and promote it to Created. func TestController_Process_AnalyzesAndTransitionsToCreated(t *testing.T) { diff --git a/submitqueue/orchestrator/controller/dlq/BUILD.bazel b/submitqueue/orchestrator/controller/dlq/BUILD.bazel index ffcbd2eb1..002291134 100644 --- a/submitqueue/orchestrator/controller/dlq/BUILD.bazel +++ b/submitqueue/orchestrator/controller/dlq/BUILD.bazel @@ -16,6 +16,7 @@ go_library( visibility = ["//visibility:public"], deps = [ "//api/runway/messagequeue:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/dlq/batch.go b/submitqueue/orchestrator/controller/dlq/batch.go index 5111926bf..7dc7f9976 100644 --- a/submitqueue/orchestrator/controller/dlq/batch.go +++ b/submitqueue/orchestrator/controller/dlq/batch.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/submitqueue/entity" @@ -86,6 +87,10 @@ func (c *batchController) Process(ctx context.Context, delivery consumer.Deliver metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to decode batch id from dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } if bid.ID == "" { metrics.NamedCounter(c.metricsScope, opName, "empty_id_errors", 1) return fmt.Errorf("dlq payload decoded to empty batch id") diff --git a/submitqueue/orchestrator/controller/dlq/batch_test.go b/submitqueue/orchestrator/controller/dlq/batch_test.go index a125271df..27a43f2a1 100644 --- a/submitqueue/orchestrator/controller/dlq/batch_test.go +++ b/submitqueue/orchestrator/controller/dlq/batch_test.go @@ -53,7 +53,7 @@ func TestDLQBatchController_Process_FailsAndFansOut(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 1, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -69,13 +69,24 @@ func TestDLQBatchController_Process_FailsAndFansOut(t *testing.T) { c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") - payload, err := entity.BatchID{ID: "q/batch/9"}.ToBytes() + payload, err := entity.BatchID{ID: "q/batch/9", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) require.NoError(t, c.Process(context.Background(), delivery)) } +func TestDLQBatchController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") + + payload, err := entity.BatchID{ID: "q/batch/9", Queue: "q"}.ToBytes() + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQBatchController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) @@ -95,7 +106,7 @@ func TestDLQBatchController_Process_EmptyIDFails(t *testing.T) { store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") - payload, err := entity.BatchID{ID: ""}.ToBytes() + payload, err := entity.BatchID{ID: "", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) diff --git a/submitqueue/orchestrator/controller/dlq/buildsignal.go b/submitqueue/orchestrator/controller/dlq/buildsignal.go index f6b61c44c..21089b869 100644 --- a/submitqueue/orchestrator/controller/dlq/buildsignal.go +++ b/submitqueue/orchestrator/controller/dlq/buildsignal.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/submitqueue/entity" @@ -78,6 +79,10 @@ func (c *buildSignalController) Process(ctx context.Context, delivery consumer.D metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to decode build id from dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, buildID.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } if buildID.ID == "" { metrics.NamedCounter(c.metricsScope, opName, "empty_id_errors", 1) return fmt.Errorf("dlq payload decoded to empty build id") diff --git a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go index acbd3009b..f3bf64768 100644 --- a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go @@ -60,7 +60,7 @@ func TestDLQBuildSignalController_Process_FansOutToBatch(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 1, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -77,13 +77,24 @@ func TestDLQBuildSignalController_Process_FansOutToBatch(t *testing.T) { c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") - payload, err := entity.BuildID{ID: "build-1"}.ToBytes() + payload, err := entity.BuildID{ID: "build-1", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) require.NoError(t, c.Process(context.Background(), delivery)) } +func TestDLQBuildSignalController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") + + payload, err := entity.BuildID{ID: "build-1", Queue: "q"}.ToBytes() + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQBuildSignalController_Process_BuildNotFoundIsNoOp(t *testing.T) { ctrl := gomock.NewController(t) @@ -96,7 +107,7 @@ func TestDLQBuildSignalController_Process_BuildNotFoundIsNoOp(t *testing.T) { c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") - payload, err := entity.BuildID{ID: "build-1"}.ToBytes() + payload, err := entity.BuildID{ID: "build-1", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) @@ -117,7 +128,7 @@ func TestDLQBuildSignalController_Process_BuildMissingBatchIsNoOp(t *testing.T) c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") - payload, err := entity.BuildID{ID: "build-1"}.ToBytes() + payload, err := entity.BuildID{ID: "build-1", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) diff --git a/submitqueue/orchestrator/controller/dlq/dlq_test.go b/submitqueue/orchestrator/controller/dlq/dlq_test.go index 8d7f63b98..390537220 100644 --- a/submitqueue/orchestrator/controller/dlq/dlq_test.go +++ b/submitqueue/orchestrator/controller/dlq/dlq_test.go @@ -82,7 +82,7 @@ func TestFailRequest_TerminalStates(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ - ID: "q/1", Version: 5, State: tt.state, + ID: "q/1", Queue: "q", Version: 5, State: tt.state, }, nil) store := storagemock.NewMockStorage(ctrl) @@ -114,7 +114,7 @@ func TestFailRequest_CancellingTransitionsToError(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 7, State: entity.RequestStateCancelling, + ID: "q/1", Queue: "q", Version: 7, State: entity.RequestStateCancelling, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(7), int32(8)).Return(nil) @@ -139,7 +139,7 @@ func TestFailRequest_TransitionsToError(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 3, State: entity.RequestStateValidated, + ID: "q/1", Queue: "q", Version: 3, State: entity.RequestStateValidated, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(3), int32(4)).Return(nil) @@ -167,7 +167,7 @@ func TestFailRequest_LogPublishErrorPropagates(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 3, State: entity.RequestStateValidated, + ID: "q/1", Queue: "q", Version: 3, State: entity.RequestStateValidated, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(3), int32(4)).Return(nil) @@ -228,12 +228,12 @@ func TestFailBatch_TransitionsAndFansOut(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request1 := entity.Request{ - ID: "q/1", Version: 2, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 2, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request1, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request1, entity.RequestStateError), int32(2), int32(3)).Return(nil) request2 := entity.Request{ - ID: "q/2", Version: 1, State: entity.RequestStateProcessing, + ID: "q/2", Queue: "q", Version: 1, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/2").Return(request2, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request2, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -263,7 +263,7 @@ func TestFailBatch_FailedFansOutForRepair(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 2, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 2, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(2), int32(3)).Return(nil) @@ -319,7 +319,7 @@ func TestFailBatch_CancellingTransitionsToFailed(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 3, State: entity.RequestStateCancelling, + ID: "q/1", Queue: "q", Version: 3, State: entity.RequestStateCancelling, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(3), int32(4)).Return(nil) diff --git a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal.go b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal.go index c51561288..14850b8e7 100644 --- a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal.go +++ b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal.go @@ -20,6 +20,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/submitqueue/extension/storage" @@ -74,6 +75,10 @@ func (c *mergeConflictSignalController) Process(ctx context.Context, delivery co metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to decode merge conflict check result from dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, result.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } store, err := c.stores.For(storage.Config{QueueName: result.GetQueueName()}) if err != nil { diff --git a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go index bf23a1d6a..210b9d6bd 100644 --- a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go @@ -46,7 +46,7 @@ func TestDLQMergeConflictSignalController_Process_ReconcilesRequest(t *testing.T requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 1, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -61,13 +61,24 @@ func TestDLQMergeConflictSignalController_Process_ReconcilesRequest(t *testing.T c := NewDLQMergeConflictSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, TopicKey(runwaymq.TopicKeyMergeConflictCheckSignal), "orchestrator-mergeconflictsignal-dlq") - payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/1", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) + payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/1", QueueName: "q", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) require.NoError(t, c.Process(context.Background(), delivery)) } +func TestDLQMergeConflictSignalController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := NewDLQMergeConflictSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeConflictCheckSignal), "orchestrator-mergeconflictsignal-dlq") + + payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/1", QueueName: "q", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQMergeConflictSignalController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/submitqueue/orchestrator/controller/dlq/mergesignal.go b/submitqueue/orchestrator/controller/dlq/mergesignal.go index 8eb3d9e4f..3f10b2752 100644 --- a/submitqueue/orchestrator/controller/dlq/mergesignal.go +++ b/submitqueue/orchestrator/controller/dlq/mergesignal.go @@ -20,6 +20,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/submitqueue/extension/storage" @@ -73,6 +74,10 @@ func (c *mergeSignalController) Process(ctx context.Context, delivery consumer.D metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to decode merge result from dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, result.GetQueueName()); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } store, err := c.stores.For(storage.Config{QueueName: result.GetQueueName()}) if err != nil { diff --git a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go index 69a122ae4..6c6a4e6da 100644 --- a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go @@ -56,7 +56,7 @@ func TestDLQMergeSignalController_Process_ReconcilesBatch(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 1, State: entity.RequestStateProcessing, + ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateProcessing, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -72,13 +72,24 @@ func TestDLQMergeSignalController_Process_ReconcilesBatch(t *testing.T) { c := NewDLQMergeSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, TopicKey(runwaymq.TopicKeyMergeSignal), "orchestrator-mergesignal-dlq") - payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/batch/1", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) + payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/batch/1", QueueName: "q", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) require.NoError(t, c.Process(context.Background(), delivery)) } +func TestDLQMergeSignalController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := NewDLQMergeSignalController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeSignal), "orchestrator-mergesignal-dlq") + + payload, err := runwaymq.Marshal(&runwaymq.MergeResult{Id: "q/batch/1", QueueName: "q", Outcome: runwaypb.Outcome_FAILED, Reason: "boom"}) + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQMergeSignalController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/submitqueue/orchestrator/controller/dlq/publisher_test.go b/submitqueue/orchestrator/controller/dlq/publisher_test.go index ea4192ff6..e767cbd78 100644 --- a/submitqueue/orchestrator/controller/dlq/publisher_test.go +++ b/submitqueue/orchestrator/controller/dlq/publisher_test.go @@ -37,6 +37,7 @@ func newTestLogRegistry( func(_ context.Context, _ string, message entityqueue.Message) error { logEntry, err := entity.RequestLogFromBytes(message.Payload) require.NoError(t, err) + require.Equal(t, logEntry.Queue, message.Tenant) return publishFn(logEntry) }, ).Times(publishCount) diff --git a/submitqueue/orchestrator/controller/dlq/request.go b/submitqueue/orchestrator/controller/dlq/request.go index cae5ab41f..9a078dcf9 100644 --- a/submitqueue/orchestrator/controller/dlq/request.go +++ b/submitqueue/orchestrator/controller/dlq/request.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" corebatch "github.com/uber/submitqueue/submitqueue/core/batch" @@ -119,6 +120,10 @@ func (c *requestController) Process(ctx context.Context, delivery consumer.Deliv // acked and dropped after the error is logged. return fmt.Errorf("failed to decode dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, rid.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } if rid.ID == "" { metrics.NamedCounter(c.metricsScope, opName, "empty_id_errors", 1) return fmt.Errorf("dlq payload decoded to empty request id") diff --git a/submitqueue/orchestrator/controller/dlq/request_test.go b/submitqueue/orchestrator/controller/dlq/request_test.go index adc34d064..ed058ceb5 100644 --- a/submitqueue/orchestrator/controller/dlq/request_test.go +++ b/submitqueue/orchestrator/controller/dlq/request_test.go @@ -49,7 +49,7 @@ func TestDLQRequestController_Process_LandRequestPayload(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/1", Version: 1, State: entity.RequestStateStarted, + ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateStarted, } requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -72,12 +72,23 @@ func TestDLQRequestController_Process_LandRequestPayload(t *testing.T) { require.NoError(t, c.Process(context.Background(), delivery)) } +func TestDLQRequestController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") + + payload, err := entity.RequestID{ID: "q/1", Queue: "q"}.ToBytes() + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQRequestController_Process_CancelRequestPayload(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/7", Version: 2, State: entity.RequestStateBatched, + ID: "q/7", Queue: "q", Version: 2, State: entity.RequestStateBatched, } requestStore.EXPECT().Get(gomock.Any(), "q/7").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(2), int32(3)).Return(nil) @@ -93,7 +104,7 @@ func TestDLQRequestController_Process_CancelRequestPayload(t *testing.T) { c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, DecodeCancelRequestID, TopicKey(topickey.TopicKeyCancel), "orchestrator-cancel-dlq") - payload, err := entity.CancelRequest{ID: "q/7", Reason: "user"}.ToBytes() + payload, err := entity.CancelRequest{ID: "q/7", Queue: "q", Reason: "user"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) @@ -105,7 +116,7 @@ func TestDLQRequestController_Process_RequestIDPayload(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) request := entity.Request{ - ID: "q/3", Version: 1, State: entity.RequestStateValidated, + ID: "q/3", Queue: "q", Version: 1, State: entity.RequestStateValidated, } requestStore.EXPECT().Get(gomock.Any(), "q/3").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -122,7 +133,7 @@ func TestDLQRequestController_Process_RequestIDPayload(t *testing.T) { c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, registry, DecodeRequestID, TopicKey(topickey.TopicKeyBatch), "orchestrator-batch-dlq") - payload, err := entity.RequestID{ID: "q/3"}.ToBytes() + payload, err := entity.RequestID{ID: "q/3", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) @@ -134,7 +145,7 @@ func TestDLQRequestController_Process_DifferentTerminalOutcomeSkips(t *testing.T requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ - ID: "q/1", Version: 5, State: entity.RequestStateLanded, + ID: "q/1", Queue: "q", Version: 5, State: entity.RequestStateLanded, }, nil) store := storagemock.NewMockStorage(ctrl) @@ -144,7 +155,7 @@ func TestDLQRequestController_Process_DifferentTerminalOutcomeSkips(t *testing.T c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") - payload, err := entity.RequestID{ID: "q/1"}.ToBytes() + payload, err := entity.RequestID{ID: "q/1", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) @@ -176,7 +187,7 @@ func TestDLQRequestController_Process_EmptyIDFails(t *testing.T) { c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), staticStorageFactory{store: store}, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") - payload, err := entity.RequestID{ID: ""}.ToBytes() + payload, err := entity.RequestID{ID: "", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) @@ -190,10 +201,17 @@ func newMockDelivery(ctrl *gomock.Controller, payload []byte) *consumermock.Mock return newMockDeliveryWithFailure(ctrl, payload, failure.Failure{}, false) } -// newMockDeliveryWithFailure builds a delivery that also reports a recorded -// failure, as a redelivery from a DLQ topic does. +func newMockDeliveryWithTenant(ctrl *gomock.Controller, payload []byte, tenant string) *consumermock.MockDelivery { + return newMockDeliveryWithTenantAndFailure(ctrl, payload, tenant, failure.Failure{}, false) +} + func newMockDeliveryWithFailure(ctrl *gomock.Controller, payload []byte, f failure.Failure, failed bool) *consumermock.MockDelivery { + return newMockDeliveryWithTenantAndFailure(ctrl, payload, "q", f, failed) +} + +func newMockDeliveryWithTenantAndFailure(ctrl *gomock.Controller, payload []byte, tenant string, f failure.Failure, failed bool) *consumermock.MockDelivery { msg := queue.NewMessage("dlq-msg-1", payload, "", nil) + msg.Tenant = tenant d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -252,7 +270,7 @@ func TestDLQRequestController_Process_SkipsRequestOwnedByLiveBatch(t *testing.T) func TestDLQRequestController_Process_FailsWhenEveryBatchIsTerminal(t *testing.T) { ctrl := gomock.NewController(t) - request := entity.Request{ID: "q/1", Version: 1, State: entity.RequestStateBatched} + request := entity.Request{ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateBatched} requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -296,7 +314,7 @@ func TestDLQRequestController_Process_FailsWhenCreatingBatchNeverClaimed(t *test t.Run(string(state), func(t *testing.T) { ctrl := gomock.NewController(t) - request := entity.Request{ID: "q/1", Version: 1, State: state} + request := entity.Request{ID: "q/1", Queue: "q", Version: 1, State: state} requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil).Times(2) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -338,7 +356,7 @@ func TestDLQRequestController_Process_SkipsWhenCreatingBatchAlreadyClaimed(t *te requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1"). - Return(entity.Request{ID: "q/1", Version: 2, State: entity.RequestStateBatched}, nil) + Return(entity.Request{ID: "q/1", Queue: "q", Version: 2, State: entity.RequestStateBatched}, nil) // Update must NOT be called — the batch owns the outcome. associations := storagemock.NewMockRequestBatchStore(ctrl) diff --git a/submitqueue/orchestrator/controller/dlq/speculate.go b/submitqueue/orchestrator/controller/dlq/speculate.go index 7e6fe5876..80bb2ece9 100644 --- a/submitqueue/orchestrator/controller/dlq/speculate.go +++ b/submitqueue/orchestrator/controller/dlq/speculate.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -90,6 +91,10 @@ func (c *speculateController) Process(ctx context.Context, delivery consumer.Del metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to decode batch id from dlq payload: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "queue_identity_errors", 1) + return nil + } if bid.ID == "" { metrics.NamedCounter(c.metricsScope, opName, "empty_id_errors", 1) return fmt.Errorf("dlq payload decoded to empty batch id") @@ -203,7 +208,12 @@ func (c *speculateController) retrigger(ctx context.Context, store storage.Stora // A distinct message ID every time: the queue deduplicates on // (topic, partition, ID) against rows it has not collected yet, so reusing // the batch ID would make this wake-up a silent no-op. - if err := publish.Message(ctx, c.registry, topickey.TopicKeySpeculate, publish.UniqueID(next), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeySpeculate, publish.MessageParams{ + Tenant: queue, + ID: publish.UniqueID(next), + Payload: payload, + PartitionKey: queue, + }); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to re-trigger speculation for queue %s: %w", queue, err) } diff --git a/submitqueue/orchestrator/controller/dlq/speculate_test.go b/submitqueue/orchestrator/controller/dlq/speculate_test.go index 3cb627fbe..a23e2252b 100644 --- a/submitqueue/orchestrator/controller/dlq/speculate_test.go +++ b/submitqueue/orchestrator/controller/dlq/speculate_test.go @@ -127,7 +127,7 @@ func TestDLQSpeculateController_Process_Attribution(t *testing.T) { batchStore.EXPECT().Get(gomock.Any(), tt.wantFailedBatch).Return(blamed, nil) batchStore.EXPECT().Update(gomock.Any(), batchWithState(blamed, entity.BatchStateFailed), int32(2), int32(3)).Return(nil) - request := entity.Request{ID: "q/1", Version: 1, State: entity.RequestStateProcessing} + request := entity.Request{ID: "q/1", Queue: "q", Version: 1, State: entity.RequestStateProcessing} requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateError), int32(1), int32(2)).Return(nil) @@ -270,13 +270,24 @@ func TestDLQSpeculateController_Process_MalformedPayloadFails(t *testing.T) { require.Error(t, c.Process(context.Background(), delivery)) } +func TestDLQSpeculateController_Process_TenantPayloadQueueMismatchAcks(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + c := newSpeculateController(consumer.TopicRegistry{}, store, t) + + payload, err := entity.BatchID{ID: "q/batch/named", Queue: "q"}.ToBytes() + require.NoError(t, err) + + require.NoError(t, c.Process(context.Background(), newMockDeliveryWithTenant(ctrl, payload, "other-queue"))) +} + func TestDLQSpeculateController_Process_EmptyIDFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) c := newSpeculateController(consumer.TopicRegistry{}, store, t) - payload, err := entity.BatchID{ID: ""}.ToBytes() + payload, err := entity.BatchID{ID: "", Queue: "q"}.ToBytes() require.NoError(t, err) delivery := newMockDelivery(ctrl, payload) diff --git a/submitqueue/orchestrator/controller/merge/BUILD.bazel b/submitqueue/orchestrator/controller/merge/BUILD.bazel index a0eb8359a..ff3b98537 100644 --- a/submitqueue/orchestrator/controller/merge/BUILD.bazel +++ b/submitqueue/orchestrator/controller/merge/BUILD.bazel @@ -10,6 +10,7 @@ go_library( "//api/base/mergestrategy/protopb:go_default_library", "//api/runway/messagequeue:go_default_library", "//platform/base/mergestrategy:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/merge/merge.go b/submitqueue/orchestrator/controller/merge/merge.go index a9657457b..4f6cfdf87 100644 --- a/submitqueue/orchestrator/controller/merge/merge.go +++ b/submitqueue/orchestrator/controller/merge/merge.go @@ -32,6 +32,7 @@ import ( strategypb "github.com/uber/submitqueue/api/base/mergestrategy/protopb" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" "github.com/uber/submitqueue/platform/base/mergestrategy" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -101,6 +102,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize batch ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: bid.Queue}) if err != nil { @@ -229,7 +233,12 @@ func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, req *ru return fmt.Errorf("failed to serialize merge request: %w", err) } - if err := publish.Message(ctx, c.registry, key, publish.IntentID(req.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: req.GetQueueName(), + ID: publish.IntentID(req.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/merge/merge_test.go b/submitqueue/orchestrator/controller/merge/merge_test.go index 90c63cd6c..0ad460ef6 100644 --- a/submitqueue/orchestrator/controller/merge/merge_test.go +++ b/submitqueue/orchestrator/controller/merge/merge_test.go @@ -47,14 +47,15 @@ type staticStorageFactory struct{ store storage.Storage } // For returns the fixed store aggregate for any queue. func (f staticStorageFactory) For(storage.Config) (storage.Storage, error) { return f.store, nil } -func batchIDPayload(t *testing.T, id string) []byte { - payload, err := entity.BatchID{ID: id}.ToBytes() +func batchIDPayload(t *testing.T, id, queue string) []byte { + payload, err := entity.BatchID{ID: id, Queue: queue}.ToBytes() require.NoError(t, err) return payload } func newDelivery(t *testing.T, ctrl *gomock.Controller, batchID, partitionKey string) *consumermock.MockDelivery { - msg := entityqueue.NewMessage(batchID, batchIDPayload(t, batchID), partitionKey, nil) + msg := entityqueue.NewMessage(batchID, batchIDPayload(t, batchID, partitionKey), partitionKey, nil) + msg.Tenant = partitionKey delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -128,6 +129,17 @@ func TestNewController(t *testing.T) { var _ consumer.Controller = c } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c := newController(t, storagemock.NewMockStorage(ctrl), consumer.TopicRegistry{}) + msg := entityqueue.NewMessage("test-queue/batch/1", batchIDPayload(t, "test-queue/batch/1", "test-queue"), "test-queue", nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, c.Process(context.Background(), delivery)) +} + func TestProcess_PublishesFullPayloadToRunway(t *testing.T) { ctrl := gomock.NewController(t) @@ -164,6 +176,7 @@ func TestProcess_PublishesFullPayloadToRunway(t *testing.T) { var gotTopic string var gotPayload []byte + var gotTenant string pub := queuemock.NewMockPublisher(ctrl) pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, topic string, msg entityqueue.Message) error { @@ -172,6 +185,7 @@ func TestProcess_PublishesFullPayloadToRunway(t *testing.T) { } gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ).AnyTimes() @@ -188,6 +202,7 @@ func TestProcess_PublishesFullPayloadToRunway(t *testing.T) { // Full payload published to runway, keyed by the batch id (the correlation id). assert.Equal(t, "runway-merge", gotTopic) + assert.Equal(t, batch.Queue, gotTenant) got := &runwaymq.MergeRequest{} require.NoError(t, runwaymq.Unmarshal(gotPayload, got)) assert.Equal(t, batch.ID, got.Id) diff --git a/submitqueue/orchestrator/controller/mergeconflictsignal/BUILD.bazel b/submitqueue/orchestrator/controller/mergeconflictsignal/BUILD.bazel index 94f0c2827..33d929693 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/BUILD.bazel +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//api/runway/messagequeue:go_default_library", "//api/runway/messagequeue/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go index dbfb69b5c..9fb29456b 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go @@ -26,6 +26,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" runwaypb "github.com/uber/submitqueue/api/runway/messagequeue/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -88,6 +89,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize merge conflict check result: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, result.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: result.GetQueueName()}) if err != nil { @@ -207,7 +211,12 @@ func (c *Controller) publishRequestID(ctx context.Context, key consumer.TopicKey return fmt.Errorf("failed to serialize request ID: %w", err) } - if err := publish.Message(ctx, c.registry, key, publish.IntentID(requestID), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(requestID), + Payload: payload, + PartitionKey: queue, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go index c71dd9fb7..1403555cd 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go @@ -48,12 +48,16 @@ func requestWithState(request entity.Request, state entity.RequestState) entity. } func resultPayload(t *testing.T, res runwaymq.MergeResult) []byte { + if res.QueueName == "" { + res.QueueName = testQueue + } payload, err := runwaymq.Marshal(&res) require.NoError(t, err) return payload } func newDelivery(ctrl *gomock.Controller, msg entityqueue.Message) *consumermock.MockDelivery { + msg.Tenant = testQueue d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -65,6 +69,20 @@ const ( testQueue = "test-queue" ) +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, staticStorageFactory{store: store}, consumer.TopicRegistry{}, + runwaymq.TopicKeyMergeConflictCheckSignal, "orchestrator-mergeconflictsignal") + res := runwaymq.MergeResult{Id: testRequestID, Outcome: runwaypb.Outcome_SUCCEEDED} + msg := entityqueue.NewMessage(testRequestID, resultPayload(t, res), testQueue, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), d)) +} + func TestProcess_MergeablePublishesToBatch(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel b/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel index 9163af3db..69a777d48 100644 --- a/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel +++ b/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//api/runway/messagequeue:go_default_library", "//api/runway/messagequeue/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go index 7611dd837..796e855ec 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go @@ -28,6 +28,7 @@ import ( "github.com/uber-go/tally" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" runwaypb "github.com/uber/submitqueue/api/runway/messagequeue/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -90,6 +91,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize merge result: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, result.GetQueueName()); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: result.GetQueueName()}) if err != nil { @@ -207,7 +211,13 @@ func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, msgID, return fmt.Errorf("failed to serialize batch ID: %w", err) } - if err := publish.MessageWithMetadata(ctx, c.registry, key, msgID, payload, queue, metadata); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: msgID, + Payload: payload, + PartitionKey: queue, + Metadata: metadata, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go index 73b78c3e3..b1235901c 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go @@ -61,12 +61,16 @@ const ( ) func resultPayload(t *testing.T, res runwaymq.MergeResult) []byte { + if res.QueueName == "" { + res.QueueName = testQueue + } payload, err := runwaymq.Marshal(&res) require.NoError(t, err) return payload } func newDelivery(ctrl *gomock.Controller, msg entityqueue.Message) *consumermock.MockDelivery { + msg.Tenant = testQueue d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -78,7 +82,8 @@ func newDelivery(ctrl *gomock.Controller, msg entityqueue.Message) *consumermock func recordingRegistry(t *testing.T, ctrl *gomock.Controller, got *[]string) consumer.TopicRegistry { pub := queuemock.NewMockPublisher(ctrl) pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, topic string, _ entityqueue.Message) error { + func(_ context.Context, topic string, msg entityqueue.Message) error { + assert.Equal(t, testQueue, msg.Tenant) *got = append(*got, topic) return nil }, @@ -117,6 +122,18 @@ func TestNewController(t *testing.T) { var _ consumer.Controller = c } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + c := newController(t, storagemock.NewMockStorage(ctrl), recordingRegistry(t, ctrl, new([]string))) + res := runwaymq.MergeResult{Id: testBatchID, Outcome: runwaypb.Outcome_SUCCEEDED} + msg := entityqueue.NewMessage(testBatchID, resultPayload(t, res), testQueue, nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, c.Process(context.Background(), d)) +} + func TestProcess_MergedAdvancesBatch(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/submitqueue/orchestrator/controller/speculate/BUILD.bazel b/submitqueue/orchestrator/controller/speculate/BUILD.bazel index c0fd13583..c223ee7e2 100644 --- a/submitqueue/orchestrator/controller/speculate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/speculate/BUILD.bazel @@ -17,6 +17,7 @@ go_library( visibility = ["//visibility:public"], deps = [ "//platform/base/failure:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/errs:go_default_library", "//platform/metrics:go_default_library", diff --git a/submitqueue/orchestrator/controller/speculate/run_test.go b/submitqueue/orchestrator/controller/speculate/run_test.go index 2c96b74e8..6c0ee85e6 100644 --- a/submitqueue/orchestrator/controller/speculate/run_test.go +++ b/submitqueue/orchestrator/controller/speculate/run_test.go @@ -271,6 +271,7 @@ func TestRun_DispatchStampsQueueAndPartitionsByHead(t *testing.T) { require.NoError(t, err) assert.Equal(t, head, got.ID) assert.Equal(t, "q", got.Queue, "the payload must name the real queue, not the partition key") + assert.Equal(t, "q", h.messages[0].Tenant, "the tenant must name the real queue, not the partition key") assert.Equal(t, head, h.messages[0].PartitionKey, "heads dispatch in parallel, so the batch is the partition key") } diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index ec7c90425..f5540313a 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -20,6 +20,7 @@ import ( "github.com/uber-go/tally" "github.com/uber/submitqueue/platform/base/failure" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/metrics" @@ -99,6 +100,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) return fmt.Errorf("failed to deserialize batch ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, bid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: bid.Queue}) if err != nil { @@ -235,7 +239,13 @@ func (c *Controller) publishBatchIDWithMetadata(ctx context.Context, key consume if err != nil { return fmt.Errorf("failed to serialize batch ID: %w", err) } - return publish.MessageWithMetadata(ctx, c.registry, key, msgID, payload, partitionKey, metadata) + return publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: msgID, + Payload: payload, + PartitionKey: partitionKey, + Metadata: metadata, + }) } // attributed records what a failure was about and counts it by subject type. diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 1fe4ef271..d924c4fb9 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate_test.go +++ b/submitqueue/orchestrator/controller/speculate/speculate_test.go @@ -84,7 +84,7 @@ func (h *procHarness) listsInFlight(batches ...entity.Batch) { func batchIDPayload(t *testing.T, id string) []byte { t.Helper() - payload, err := entity.BatchID{ID: id}.ToBytes() + payload, err := entity.BatchID{ID: id, Queue: "test-queue"}.ToBytes() require.NoError(t, err) return payload } @@ -173,6 +173,7 @@ func newProcHarness(t *testing.T, ctrl *gomock.Controller, publishErr error) *pr func (h *procHarness) process(t *testing.T, ctrl *gomock.Controller, batchID string) error { t.Helper() msg := entityqueue.NewMessage(batchID, batchIDPayload(t, batchID), "test-queue", nil) + msg.Tenant = "test-queue" d := consumermock.NewMockDelivery(ctrl) d.EXPECT().Message().Return(msg).AnyTimes() d.EXPECT().Attempt().Return(1).AnyTimes() @@ -190,6 +191,17 @@ func TestNewController(t *testing.T) { var _ consumer.Controller = h.controller } +func TestProcess_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + h := newProcHarness(t, ctrl, nil) + msg := entityqueue.NewMessage("test-queue/batch/1", batchIDPayload(t, "test-queue/batch/1"), "test-queue", nil) + msg.Tenant = "other-queue" + d := consumermock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, h.controller.Process(context.Background(), d)) +} + // A Created batch is admitted so the Speculator can act on it, and must not // reach an outcome on the same message — nothing has been built yet. func TestProcess_AdmitsCreatedBatch(t *testing.T) { diff --git a/submitqueue/orchestrator/controller/start/BUILD.bazel b/submitqueue/orchestrator/controller/start/BUILD.bazel index b86bcd768..dbd2ced1c 100644 --- a/submitqueue/orchestrator/controller/start/BUILD.bazel +++ b/submitqueue/orchestrator/controller/start/BUILD.bazel @@ -6,6 +6,7 @@ go_library( importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/start", visibility = ["//visibility:public"], deps = [ + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/start/start.go b/submitqueue/orchestrator/controller/start/start.go index 5dbfbd04f..d6b54deee 100644 --- a/submitqueue/orchestrator/controller/start/start.go +++ b/submitqueue/orchestrator/controller/start/start.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -81,6 +82,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Non-retryable: malformed messages will never succeed regardless of retry count return fmt.Errorf("failed to deserialize land request: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, landRequest.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: landRequest.Queue}) if err != nil { @@ -149,7 +153,12 @@ func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, request return fmt.Errorf("failed to serialize request ID: %w", err) } - if err := publish.Message(ctx, c.registry, key, publish.IntentID(requestID), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(requestID), + Payload: payload, + PartitionKey: queue, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/start/start_test.go b/submitqueue/orchestrator/controller/start/start_test.go index 12b88450d..cb0e8a4ac 100644 --- a/submitqueue/orchestrator/controller/start/start_test.go +++ b/submitqueue/orchestrator/controller/start/start_test.go @@ -90,6 +90,7 @@ func makeDelivery(t *testing.T, ctrl *gomock.Controller, lr entity.LandRequest) require.NoError(t, err) msg := entityqueue.NewMessage(lr.ID, payload, lr.Queue, nil) + msg.Tenant = lr.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -120,6 +121,23 @@ func TestController_Process_Success(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + controller := newTestController(t, ctrl, newMockStorage(ctrl), nil) + request := entity.LandRequest{ID: "test-queue/123", Queue: "test-queue"} + payload, err := request.ToBytes() + require.NoError(t, err) + msg := entityqueue.NewMessage(request.ID, payload, request.Queue, nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + err = controller.Process(context.Background(), delivery) + + require.Error(t, err) + assert.False(t, errs.IsRetryable(err)) +} + func TestController_Process_InvalidJSON(t *testing.T) { ctrl := gomock.NewController(t) controller := newTestController(t, ctrl, newMockStorage(ctrl), nil) diff --git a/submitqueue/orchestrator/controller/validate/BUILD.bazel b/submitqueue/orchestrator/controller/validate/BUILD.bazel index af59d4255..11ba47349 100644 --- a/submitqueue/orchestrator/controller/validate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/validate/BUILD.bazel @@ -10,6 +10,7 @@ go_library( "//api/base/mergestrategy/protopb:go_default_library", "//api/runway/messagequeue:go_default_library", "//platform/base/mergestrategy:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", "//platform/publish:go_default_library", diff --git a/submitqueue/orchestrator/controller/validate/validate.go b/submitqueue/orchestrator/controller/validate/validate.go index 9eacab3cf..dbb8e7e17 100644 --- a/submitqueue/orchestrator/controller/validate/validate.go +++ b/submitqueue/orchestrator/controller/validate/validate.go @@ -25,6 +25,7 @@ import ( strategypb "github.com/uber/submitqueue/api/base/mergestrategy/protopb" runwaymq "github.com/uber/submitqueue/api/runway/messagequeue" "github.com/uber/submitqueue/platform/base/mergestrategy" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" coremetrics "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/platform/publish" @@ -97,6 +98,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er coremetrics.NamedCounter(c.metricsScope, "process", "deserialize_errors", 1) return fmt.Errorf("failed to deserialize request ID: %w", err) } + if err := entityqueue.ValidatePayloadQueue(msg, rid.Queue); err != nil { + return fmt.Errorf("invalid message identity: %w", err) + } store, err := c.stores.For(storage.Config{QueueName: rid.Queue}) if err != nil { @@ -321,7 +325,12 @@ func (c *Controller) publishMergeCheck(ctx context.Context, req *runwaymq.MergeR return fmt.Errorf("failed to serialize merge conflict check request: %w", err) } - if err := publish.Message(ctx, c.registry, c.runwayTopicKey, publish.IntentID(req.GetId()), payload, req.GetQueueName()); err != nil { + if err := publish.Message(ctx, c.registry, c.runwayTopicKey, publish.MessageParams{ + Tenant: req.GetQueueName(), + ID: publish.IntentID(req.GetId()), + Payload: payload, + PartitionKey: req.GetQueueName(), + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/validate/validate_test.go b/submitqueue/orchestrator/controller/validate/validate_test.go index 4778d408c..bb2bced08 100644 --- a/submitqueue/orchestrator/controller/validate/validate_test.go +++ b/submitqueue/orchestrator/controller/validate/validate_test.go @@ -55,7 +55,7 @@ func requestWithState(request entity.Request, state entity.RequestState) entity. // requestIDPayload serializes a RequestID to JSON bytes for test message payloads. func requestIDPayload(t *testing.T, id string) []byte { - payload, err := entity.RequestID{ID: id}.ToBytes() + payload, err := entity.RequestID{ID: id, Queue: "test-queue"}.ToBytes() require.NoError(t, err) return payload } @@ -173,6 +173,7 @@ func TestController_Process_Success(t *testing.T) { controller := newTestController(t, ctrl, store, newMockChangeStore(ctrl), nil) msg := entityqueue.NewMessage("test-queue/123", requestIDPayload(t, request.ID), "test-queue", nil) + msg.Tenant = "test-queue" delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -180,6 +181,18 @@ func TestController_Process_Success(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) } +func TestController_Process_RejectsTenantPayloadQueueMismatch(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + controller := newTestController(t, ctrl, store, storagemock.NewMockChangeStore(ctrl), nil) + msg := entityqueue.NewMessage("test-queue/123", requestIDPayload(t, "test-queue/123"), "test-queue", nil) + msg.Tenant = "other-queue" + delivery := consumermock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + + require.Error(t, controller.Process(context.Background(), delivery)) +} + // TestController_Process_PublishesCheckToRunway verifies the full merge-conflict // check request is published to runway's merge-conflict-check queue (keyed by // the request id, the client-owned correlation id) on the happy path. @@ -201,6 +214,7 @@ func TestController_Process_PublishesCheckToRunway(t *testing.T) { var gotTopic string var gotPayload []byte + var gotTenant string mockPub := queuemock.NewMockPublisher(ctrl) mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, topic string, msg entityqueue.Message) error { @@ -209,6 +223,7 @@ func TestController_Process_PublishesCheckToRunway(t *testing.T) { } gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ).AnyTimes() @@ -227,6 +242,7 @@ func TestController_Process_PublishesCheckToRunway(t *testing.T) { controller := NewController(logger, tally.NoopScope, staticStorageFactory{store: store}, registry, cpFactory, nil, runwaymq.TopicKeyMergeConflictCheck, topickey.TopicKeyValidate, "orchestrator-validate") msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -235,6 +251,7 @@ func TestController_Process_PublishesCheckToRunway(t *testing.T) { // Full payload published to runway, keyed by the request id (the correlation id). assert.Equal(t, "merge-conflict-check", gotTopic) + assert.Equal(t, request.Queue, gotTenant) got := &runwaymq.MergeRequest{} require.NoError(t, runwaymq.Unmarshal(gotPayload, got)) assert.Equal(t, request.ID, got.Id) @@ -286,6 +303,7 @@ func TestController_Process_ClaimsChangeRecordsWithDetails(t *testing.T) { controller := newTestController(t, ctrl, store, cs, nil) msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -304,6 +322,7 @@ func TestController_Process_StorageFailure(t *testing.T) { controller := newTestController(t, ctrl, store, newMockChangeStore(ctrl), nil) msg := entityqueue.NewMessage("test-queue/123", requestIDPayload(t, "test-queue/123"), "test-queue", nil) + msg.Tenant = "test-queue" delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -329,6 +348,7 @@ func TestController_Process_PublishFailure(t *testing.T) { controller := newTestController(t, ctrl, store, newMockChangeStore(ctrl), fmt.Errorf("publish failed")) msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -507,6 +527,7 @@ func TestController_Process_DuplicateDetection(t *testing.T) { controller := newTestController(t, ctrl, store, cs, nil) msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -547,6 +568,7 @@ func TestController_Process_ChangeStoreQueryFailure(t *testing.T) { controller := newTestController(t, ctrl, store, cs, nil) msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -586,6 +608,7 @@ func TestController_Process_TerminalShortCircuit(t *testing.T) { controller := newTestController(t, ctrl, store, cs, fmt.Errorf("should not publish")) msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -637,6 +660,7 @@ func TestController_Process_CustomValidatorPasses(t *testing.T) { controller := NewController(logger, tally.NoopScope, staticStorageFactory{store: store}, registry, cpFactory, mockValidatorFactory, runwaymq.TopicKeyMergeConflictCheck, topickey.TopicKeyValidate, "orchestrator-validate") msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -696,6 +720,7 @@ func TestController_Process_CustomValidatorFails(t *testing.T) { controller := NewController(logger, tally.NoopScope, staticStorageFactory{store: store}, registry, cpFactory, mockValidatorFactory, runwaymq.TopicKeyMergeConflictCheck, topickey.TopicKeyValidate, "orchestrator-validate") msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() @@ -758,6 +783,7 @@ func TestController_Process_CustomValidatorFailure_TerminationPublishFails(t *te controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, staticStorageFactory{store: store}, registry, cpFactory, mockValidatorFactory, runwaymq.TopicKeyMergeConflictCheck, topickey.TopicKeyValidate, "orchestrator-validate") msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + msg.Tenant = request.Queue delivery := consumermock.NewMockDelivery(ctrl) delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() diff --git a/test/e2e/submitqueue/BUILD.bazel b/test/e2e/submitqueue/BUILD.bazel index 680b0f6d3..fb03d6cad 100644 --- a/test/e2e/submitqueue/BUILD.bazel +++ b/test/e2e/submitqueue/BUILD.bazel @@ -39,6 +39,7 @@ go_test( "//api/runway/messagequeue:go_default_library", "//api/submitqueue/gateway/protopb:go_default_library", "//api/submitqueue/orchestrator/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/extension/consumergate:go_default_library", "//platform/extension/consumergate/file:go_default_library", diff --git a/test/e2e/submitqueue/harness_test.go b/test/e2e/submitqueue/harness_test.go index 4ee147f49..68427ef0d 100644 --- a/test/e2e/submitqueue/harness_test.go +++ b/test/e2e/submitqueue/harness_test.go @@ -35,6 +35,7 @@ import ( 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" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/extension/consumergate" queuemysql "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" @@ -286,6 +287,7 @@ func (s *E2EIntegrationSuite) redeliverBatchMessage(req request) { DB: s.queueDB, Logger: zap.NewNop(), MetricsScope: tally.NoopScope, + Tenants: []string{req.queue}, }) require.NoError(t, err, "failed to open the queue for a manual publish") defer func() { require.NoError(t, queue.Close()) }() @@ -298,8 +300,12 @@ func (s *E2EIntegrationSuite) redeliverBatchMessage(req request) { payload, err := entity.RequestID{ID: req.sqid, Queue: req.queue}.ToBytes() require.NoError(t, err) - require.NoError(t, publish.Message(s.ctx, registry, topickey.TopicKeyBatch, - publish.UniqueID(req.sqid), payload, req.queue), "failed to redeliver the batch message") + require.NoError(t, publish.Message(entityqueue.WithQueueName(s.ctx, req.queue), registry, topickey.TopicKeyBatch, publish.MessageParams{ + Tenant: req.queue, + ID: publish.UniqueID(req.sqid), + Payload: payload, + PartitionKey: req.queue, + }), "failed to redeliver the batch message") s.log.Logf("Redelivered the batch message for %s", req.sqid) } @@ -376,9 +382,15 @@ func (s *E2EIntegrationSuite) strandInCreated(queue, batchID string) { // 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) { +func (s *E2EIntegrationSuite) closeGate(tenant, consumerGroup, partitionKey, reason string) { t := s.T() - key := consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey} + key := consumergate.Key{ + ConsumerGroup: consumerGroup, + Partition: entityqueue.PartitionIdentity{ + Tenant: tenant, + PartitionKey: partitionKey, + }, + } require.NoError(t, s.gate.Close(s.ctx, key, consumergate.Metadata{ Reason: reason, CreatedBy: "e2e-suite", @@ -390,9 +402,15 @@ func (s *E2EIntegrationSuite) closeGate(consumerGroup, partitionKey, reason stri // 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) { +func (s *E2EIntegrationSuite) openGate(tenant, consumerGroup, partitionKey string) { t := s.T() - key := consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: partitionKey} + key := consumergate.Key{ + ConsumerGroup: consumerGroup, + Partition: entityqueue.PartitionIdentity{ + Tenant: tenant, + 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) } diff --git a/test/e2e/submitqueue/suite_test.go b/test/e2e/submitqueue/suite_test.go index 2ecc0973d..e35dfdae9 100644 --- a/test/e2e/submitqueue/suite_test.go +++ b/test/e2e/submitqueue/suite_test.go @@ -325,10 +325,10 @@ func (s *E2EIntegrationSuite) TestDependentBatch_BypassesMergingDependency() { const gateGroup = "runway-merge" gateTopic := runwaymq.TopicKeyMerge.String() - s.closeGate(gateGroup, queue, "e2e: hold the lead merge so the dependent finishes building first") + s.closeGate(queue, gateGroup, queue, "e2e: hold the lead merge so the dependent finishes building first") // 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) + defer s.openGate(queue, gateGroup, queue) lead := s.land(queue, "github://github.example.com/uber/e2e-chain/pull/1/abcdef0123456789abcdef0123456789abcdef01") s.log.Logf("Landed lead request %s; awaiting its merge to park", lead.sqid) @@ -360,7 +360,7 @@ func (s *E2EIntegrationSuite) TestDependentBatch_BypassesMergingDependency() { // Start: the lead merges, and its fan-out is now the only thing that can // move the dependent. - s.openGate(gateGroup, queue) + s.openGate(queue, gateGroup, queue) s.awaitUnparked(gateGroup, gateTopic, leadBatch) s.awaitStatus(lead, entity.RequestStatusLanded) @@ -380,8 +380,8 @@ func (s *E2EIntegrationSuite) TestDependentBatch_BypassedHeadLandsFirst() { lead := s.land(queue, "github://github.example.com/uber/e2e-bypass/pull/1/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") heldBatch := s.awaitBatchID(lead) - s.closeGate(gateGroup, heldBatch, "e2e: hold the leader's build so the follower can fully cover it") - defer s.openGate(gateGroup, heldBatch) + s.closeGate(queue, gateGroup, heldBatch, "e2e: hold the leader's build so the follower can fully cover it") + defer s.openGate(queue, gateGroup, heldBatch) s.awaitBatchState(queue, heldBatch, entity.BatchStateSpeculating) follower := s.land(queue, "github://github.example.com/uber/e2e-bypass/pull/2/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb") @@ -395,7 +395,7 @@ func (s *E2EIntegrationSuite) TestDependentBatch_BypassedHeadLandsFirst() { assert.Equal(t, entity.RequestStatusSpeculating, s.mustStatus(lead), "the follower must land while its dependency is still held") - s.openGate(gateGroup, heldBatch) + s.openGate(queue, gateGroup, heldBatch) s.awaitStatus(lead, entity.RequestStatusLanded) s.assertStatusesInOrder(follower, @@ -426,8 +426,8 @@ func (s *E2EIntegrationSuite) TestDependentBatch_NoBypassWhenCoverageIsIncomplet follower := s.land(queue, "github://github.example.com/uber/e2e-nobypass/pull/2/dddddddddddddddddddddddddddddddddddddddd") heldBatch := s.awaitBatchID(follower) - s.closeGate(gateGroup, heldBatch, "e2e: hold the follower's builds so only the seeded path exists") - defer s.openGate(gateGroup, heldBatch) + s.closeGate(queue, gateGroup, heldBatch, "e2e: hold the follower's builds so only the seeded path exists") + defer s.openGate(queue, gateGroup, heldBatch) // Strand the follower in Created, then seed exactly one passed path: the // guess that the lead succeeds. Its builds are parked, so nothing can add @@ -461,7 +461,7 @@ func (s *E2EIntegrationSuite) TestDependentBatch_NoBypassWhenCoverageIsIncomplet // Let the queue finish. The follower's own speculative builds were parked // on the held partition; releasing it lets any surviving work drain. How // the queue ultimately converges is exercised by the other tests. - s.openGate(gateGroup, heldBatch) + s.openGate(queue, gateGroup, heldBatch) s.awaitStatus(trigger, entity.RequestStatusLanded) } @@ -557,8 +557,8 @@ func (s *E2EIntegrationSuite) TestLand_DependentBatch_BypassesAnUnresolvedDepend const gateGroup = "orchestrator" leaderBatch := queue + "/batch/1" - s.closeGate(gateGroup, leaderBatch, "e2e: hold the leader's build so its dependent speculates first") - defer s.openGate(gateGroup, leaderBatch) + s.closeGate(queue, gateGroup, leaderBatch, "e2e: hold the leader's build so its dependent speculates first") + defer s.openGate(queue, gateGroup, leaderBatch) leader := s.land(queue, "github://github.example.com/uber/e2e-respeculate/pull/1/1111111111111111111111111111111111111111?sq-fake=build-fail") follower := s.land(queue, "github://github.example.com/uber/e2e-respeculate/pull/2/2222222222222222222222222222222222222222") @@ -572,7 +572,7 @@ func (s *E2EIntegrationSuite) TestLand_DependentBatch_BypassesAnUnresolvedDepend // Release the leader only after the follower has landed. Its later failure is // one of the outcomes the follower already validated. - s.openGate(gateGroup, leaderBatch) + s.openGate(queue, gateGroup, leaderBatch) assert.Equal(s.T(), entity.RequestStatusError, s.awaitTerminal(leader), "the leader's build carries a failure marker, so it must not land") @@ -632,10 +632,10 @@ func (s *E2EIntegrationSuite) TestCancel_CaughtPreBatch_NeverLands() { const gateGroup = "runway-mergeconflictcheck" gateTopic := runwaymq.TopicKeyMergeConflictCheck.String() - s.closeGate(gateGroup, queue, "e2e: hold merge-conflict check to catch cancel pre-batch") + s.closeGate(queue, 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) + defer s.openGate(queue, gateGroup, queue) req := 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", req.sqid) @@ -661,7 +661,7 @@ func (s *E2EIntegrationSuite) TestCancel_CaughtPreBatch_NeverLands() { "operating store should show request %s terminal cancelled while its check is parked", req.sqid) // Start the controller again and prove the parked delivery cleared the gate. - s.openGate(gateGroup, queue) + s.openGate(queue, gateGroup, queue) s.awaitUnparked(gateGroup, gateTopic, req.sqid) // Sentinel on the same queue: its landing proves the stale signal ahead of @@ -697,10 +697,10 @@ func (s *E2EIntegrationSuite) TestBatchRedelivery_DoesNotEnrolTheRequestTwice() // the build topic partitions by batch ID. const heldBatch = queue + "/batch/1" - s.closeGate(gateGroup, heldBatch, "e2e: hold the build so the request stays in flight for the redelivery") + s.closeGate(queue, gateGroup, heldBatch, "e2e: hold the build so the request stays in flight for the redelivery") // 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, heldBatch) + defer s.openGate(queue, gateGroup, heldBatch) req := s.land(queue, "github://github.example.com/uber/e2e-redelivery/pull/1/abcdef0123456789abcdef0123456789abcdef01") require.Equal(t, heldBatch, s.awaitBatchID(req), "the first batch of a fresh queue must be batch/1") @@ -715,7 +715,7 @@ func (s *E2EIntegrationSuite) TestBatchRedelivery_DoesNotEnrolTheRequestTwice() assert.Equal(t, []string{heldBatch}, s.batchIDsFor(req), "the redelivery must resume the existing batch, not mint another") - s.openGate(gateGroup, heldBatch) + s.openGate(queue, gateGroup, heldBatch) s.awaitStatus(req, entity.RequestStatusLanded) s.awaitStatus(settle, entity.RequestStatusLanded) } @@ -739,8 +739,8 @@ func (s *E2EIntegrationSuite) TestStrandedBatch_IsAdmittedByALaterRun() { const gateGroup = "orchestrator" const heldBatch = queue + "/batch/1" - s.closeGate(gateGroup, heldBatch, "e2e: hold the build so the batch can be stranded while still in flight") - defer s.openGate(gateGroup, heldBatch) + s.closeGate(queue, gateGroup, heldBatch, "e2e: hold the build so the batch can be stranded while still in flight") + defer s.openGate(queue, gateGroup, heldBatch) req := s.land(queue, "github://github.example.com/uber/e2e-strand/pull/1/abcdef0123456789abcdef0123456789abcdef01") require.Equal(t, heldBatch, s.awaitBatchID(req), "the first batch of a fresh queue must be batch/1") @@ -757,7 +757,7 @@ func (s *E2EIntegrationSuite) TestStrandedBatch_IsAdmittedByALaterRun() { s.awaitBatchState(queue, heldBatch, entity.BatchStateSpeculating) - s.openGate(gateGroup, heldBatch) + s.openGate(queue, gateGroup, heldBatch) s.awaitStatus(req, entity.RequestStatusLanded) s.awaitStatus(trigger, entity.RequestStatusLanded) } diff --git a/test/integration/extension/messagequeue/mysql/BUILD.bazel b/test/integration/extension/messagequeue/mysql/BUILD.bazel index e6803cf72..942aeff28 100644 --- a/test/integration/extension/messagequeue/mysql/BUILD.bazel +++ b/test/integration/extension/messagequeue/mysql/BUILD.bazel @@ -2,7 +2,10 @@ load("@rules_go//go:def.bzl", "go_test") go_test( name = "go_default_test", - srcs = ["queue_test.go"], + srcs = [ + "queue_test.go", + "tenant_isolation_test.go", + ], data = [ "docker-compose.yml", "//platform/extension/messagequeue/mysql/schema", diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index a87503019..2e756beaf 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -38,6 +38,8 @@ import ( "github.com/uber/submitqueue/test/testutil" ) +const testTenant = "test-tenant" + type SQLQueueIntegrationSuite struct { suite.Suite ctx context.Context @@ -50,6 +52,25 @@ func TestSQLQueueIntegration(t *testing.T) { suite.Run(t, new(SQLQueueIntegrationSuite)) } +func (s *SQLQueueIntegrationSuite) testQueueParams(t *testing.T, extra func(*queueMySQL.Params)) queueMySQL.Params { + p := queueMySQL.Params{ + DB: s.db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + Tenants: []string{testTenant}, + } + if extra != nil { + extra(&p) + } + return p +} + +func testMessage(id string, payload []byte, partitionKey string, metadata map[string]string) entityqueue.Message { + msg := entityqueue.NewMessage(id, payload, partitionKey, metadata) + msg.Tenant = testTenant + return msg +} + func (s *SQLQueueIntegrationSuite) SetupSuite() { t := s.T() s.ctx = context.Background() @@ -99,6 +120,64 @@ func (s *SQLQueueIntegrationSuite) TearDownSuite() { // Cleanup handled automatically by testutil.ComposeStack } +func (s *SQLQueueIntegrationSuite) TestIndexedIdentifiersFitInnoDBKeyLimit() { + var asciiIdentifiers int + err := s.db.QueryRowContext(s.ctx, ` + SELECT COUNT(DISTINCT s.TABLE_NAME, s.COLUMN_NAME) + FROM information_schema.STATISTICS AS s + JOIN information_schema.COLUMNS AS c + ON c.TABLE_SCHEMA = s.TABLE_SCHEMA + AND c.TABLE_NAME = s.TABLE_NAME + AND c.COLUMN_NAME = s.COLUMN_NAME + WHERE s.TABLE_SCHEMA = DATABASE() + AND c.DATA_TYPE = 'varchar' + AND c.CHARACTER_MAXIMUM_LENGTH = 255 + AND c.CHARACTER_SET_NAME = 'ascii' + AND c.COLLATION_NAME = 'ascii_bin' + `).Scan(&asciiIdentifiers) + require.NoError(s.T(), err) + assert.Equal(s.T(), 15, asciiIdentifiers) + + var utf8Identifiers int + err = s.db.QueryRowContext(s.ctx, ` + SELECT COUNT(DISTINCT s.TABLE_NAME, s.COLUMN_NAME) + FROM information_schema.STATISTICS AS s + JOIN information_schema.COLUMNS AS c + ON c.TABLE_SCHEMA = s.TABLE_SCHEMA + AND c.TABLE_NAME = s.TABLE_NAME + AND c.COLUMN_NAME = s.COLUMN_NAME + WHERE s.TABLE_SCHEMA = DATABASE() + AND c.DATA_TYPE = 'varchar' + AND c.CHARACTER_MAXIMUM_LENGTH = 255 + AND c.CHARACTER_SET_NAME = 'utf8mb4' + AND c.COLLATION_NAME = 'utf8mb4_bin' + `).Scan(&utf8Identifiers) + require.NoError(s.T(), err) + assert.Equal(s.T(), 5, utf8Identifiers) +} + +func (s *SQLQueueIntegrationSuite) TestOffsetsPrimaryKeyOrder() { + rows, err := s.db.QueryContext(s.ctx, ` + SELECT COLUMN_NAME + FROM information_schema.KEY_COLUMN_USAGE + WHERE TABLE_SCHEMA = DATABASE() + AND TABLE_NAME = 'queue_offsets' + AND CONSTRAINT_NAME = 'PRIMARY' + ORDER BY ORDINAL_POSITION + `) + require.NoError(s.T(), err) + defer rows.Close() + + var columns []string + for rows.Next() { + var column string + require.NoError(s.T(), rows.Scan(&column)) + columns = append(columns, column) + } + require.NoError(s.T(), rows.Err()) + assert.Equal(s.T(), []string{"tenant", "topic", "partition_key", "consumer_group"}, columns) +} + // testSubConfig returns a SubscriptionConfig with short lease/visibility // timeouts for fast integration tests. The defaults (30s lease, 60s visibility) // would make crash recovery tests wait 90s of real wall-clock time since the @@ -322,7 +401,7 @@ func waitForLag( t.Helper() for { - lags, err := admin.ConsumerLag(ctx, topic) + lags, err := admin.ConsumerLag(ctx, testTenant, topic) require.NoError(t, err) var actual int64 = -1 @@ -345,11 +424,7 @@ func (s *SQLQueueIntegrationSuite) TestPublishAndSubscribe() { t := s.T() // Create queue - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -365,13 +440,13 @@ func (s *SQLQueueIntegrationSuite) TestPublishAndSubscribe() { require.NoError(t, err) // Publish messages with various metadata scenarios - msg1 := entityqueue.NewMessage("msg-1", []byte("hello"), "partition-1", map[string]string{ + msg1 := testMessage("msg-1", []byte("hello"), "partition-1", map[string]string{ "key1": "value1", "key2": "value2", "trace_id": "abc123", }) - msg2 := entityqueue.NewMessage("msg-2", []byte("world"), "partition-1", nil) + msg2 := testMessage("msg-2", []byte("world"), "partition-1", nil) err = publisher.Publish(s.ctx, topic, msg1) require.NoError(t, err) @@ -416,11 +491,7 @@ func (s *SQLQueueIntegrationSuite) TestPublishAndSubscribe() { func (s *SQLQueueIntegrationSuite) TestSubscriberPerPartitionIsolation() { t := s.T() - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -437,8 +508,8 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPerPartitionIsolation() { require.NoError(t, err) // Publish 1 message to partition-a and 1 to partition-b - msgA := entityqueue.NewMessage("iso-msg-a", []byte("data-a"), "partition-a", nil) - msgB := entityqueue.NewMessage("iso-msg-b", []byte("data-b"), "partition-b", nil) + msgA := testMessage("iso-msg-a", []byte("data-a"), "partition-a", nil) + msgB := testMessage("iso-msg-b", []byte("data-b"), "partition-b", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msgA)) require.NoError(t, publisher.Publish(s.ctx, topic, msgB)) t.Logf("Published 1 message to partition-a and 1 to partition-b") @@ -472,11 +543,7 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPerPartitionIsolation() { func (s *SQLQueueIntegrationSuite) TestSubscriberPartitionOrderPreserved() { t := s.T() - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -492,7 +559,7 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPartitionOrderPreserved() { for i := 0; i < numMessages; i++ { msgID := fmt.Sprintf("order-msg-%03d", i) publishedIDs[i] = msgID - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("payload-%d", i)), partitionKey, nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("payload-%d", i)), partitionKey, nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages to partition %s", numMessages, partitionKey) @@ -525,11 +592,7 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPartitionOrderPreserved() { func (s *SQLQueueIntegrationSuite) TestMultiplePartitions() { t := s.T() - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -549,8 +612,8 @@ func (s *SQLQueueIntegrationSuite) TestMultiplePartitions() { expectedCount := len(partitions) * 2 // 2 messages per partition for _, partition := range partitions { - msg1 := entityqueue.NewMessage(partition+"-msg-1", []byte("data-1"), partition, nil) - msg2 := entityqueue.NewMessage(partition+"-msg-2", []byte("data-2"), partition, nil) + msg1 := testMessage(partition+"-msg-1", []byte("data-1"), partition, nil) + msg2 := testMessage(partition+"-msg-2", []byte("data-2"), partition, nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg1)) require.NoError(t, publisher.Publish(s.ctx, topic, msg2)) @@ -572,12 +635,9 @@ func (s *SQLQueueIntegrationSuite) TestVisibilityTimeoutAndRetry() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -594,7 +654,7 @@ func (s *SQLQueueIntegrationSuite) TestVisibilityTimeoutAndRetry() { require.NoError(t, err) // Publish a message - msg := entityqueue.NewMessage("retry-msg", []byte("test"), "retry-partition", nil) + msg := testMessage("retry-msg", []byte("test"), "retry-partition", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) t.Logf("Published message, expecting visibility timeout retry") @@ -629,7 +689,7 @@ func (s *SQLQueueIntegrationSuite) TestVisibilityTimeoutAndRetry() { t.Logf("Test 2: Visibility timeout retry") // Publish another message - msg2 := entityqueue.NewMessage("retry-msg-2", []byte("test2"), "retry-partition", nil) + msg2 := testMessage("retry-msg-2", []byte("test2"), "retry-partition", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg2)) // Receive first time @@ -654,12 +714,9 @@ func (s *SQLQueueIntegrationSuite) TestNackBackoff() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -673,7 +730,7 @@ func (s *SQLQueueIntegrationSuite) TestNackBackoff() { deliveryChan, err := q.Subscriber().Subscribe(s.ctx, "nack_backoff_topic", subConfig) require.NoError(t, err) require.NoError(t, q.Publisher().Publish(s.ctx, "nack_backoff_topic", - entityqueue.NewMessage("retry-msg", []byte("test"), "retry-partition", nil))) + testMessage("retry-msg", []byte("test"), "retry-partition", nil))) firstDelivery := receive(t, deliveryChan) assert.Equal(t, 1, firstDelivery.Attempt()) @@ -689,12 +746,9 @@ func (s *SQLQueueIntegrationSuite) TestIdempotentPublish() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -710,7 +764,7 @@ func (s *SQLQueueIntegrationSuite) TestIdempotentPublish() { require.NoError(t, err) // Publish same message twice - msg := entityqueue.NewMessage("same-id", []byte("duplicate"), "same-partition", nil) + msg := testMessage("same-id", []byte("duplicate"), "same-partition", nil) err1 := publisher.Publish(s.ctx, topic, msg) require.NoError(t, err1) @@ -746,12 +800,9 @@ func (s *SQLQueueIntegrationSuite) TestDedupOutlivesConsumption() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -766,7 +817,7 @@ func (s *SQLQueueIntegrationSuite) TestDedupOutlivesConsumption() { // The first publish for an entity: delivered, acked, and now awaiting // collection rather than gone. require.NoError(t, publisher.Publish(s.ctx, topic, - entityqueue.NewMessage("batch-1", []byte("announced"), "queue-1", nil))) + testMessage("batch-1", []byte("announced"), "queue-1", nil))) first := receive(t, deliveryChan) require.Equal(t, "batch-1", first.Message().ID) require.NoError(t, first.Ack(s.ctx)) @@ -774,12 +825,12 @@ func (s *SQLQueueIntegrationSuite) TestDedupOutlivesConsumption() { // A later, unrelated event about the same entity, published under the same // ID. It reports success and is never delivered. require.NoError(t, publisher.Publish(s.ctx, topic, - entityqueue.NewMessage("batch-1", []byte("woken"), "queue-1", nil))) + testMessage("batch-1", []byte("woken"), "queue-1", nil))) assertNoDelivery(t, deliveryChan, signalCh, queueMySQL.SignalDeliveryCheck, 3) // Naming the cause is what gets it through. require.NoError(t, publisher.Publish(s.ctx, topic, - entityqueue.NewMessage("batch-1/merged", []byte("woken"), "queue-1", nil))) + testMessage("batch-1/merged", []byte("woken"), "queue-1", nil))) second := receive(t, deliveryChan) assert.Equal(t, "batch-1/merged", second.Message().ID) assert.Equal(t, []byte("woken"), second.Message().Payload) @@ -789,11 +840,7 @@ func (s *SQLQueueIntegrationSuite) TestDedupOutlivesConsumption() { func (s *SQLQueueIntegrationSuite) TestConcurrentPublishers() { t := s.T() - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -817,7 +864,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentPublishers() { for i := 0; i < numPublishers; i++ { go func(publisherID int) { for j := 0; j < messagesPerPublisher; j++ { - msg := entityqueue.NewMessage( + msg := testMessage( t.Name()+"-"+string(rune(publisherID))+"-"+string(rune(j)), []byte("concurrent"), "concurrent-partition", @@ -846,11 +893,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentPublishers() { func (s *SQLQueueIntegrationSuite) TestCrashRecovery() { t := s.T() - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) publisher := q1.Publisher() @@ -866,7 +909,7 @@ func (s *SQLQueueIntegrationSuite) TestCrashRecovery() { require.NoError(t, err) // Publish message - msg := entityqueue.NewMessage("crash-msg", []byte("test-crash"), "crash-partition", nil) + msg := testMessage("crash-msg", []byte("test-crash"), "crash-partition", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) // Worker 1 receives but doesn't ack (simulating crash) @@ -880,11 +923,7 @@ func (s *SQLQueueIntegrationSuite) TestCrashRecovery() { // Start worker 2 with same consumer group — it will poll and find the // message after lease + visibility timeout expire in the DB - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q2.Close() @@ -912,19 +951,11 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroups() { topic := "multi_group_topic" // Create two different consumer groups - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q1.Close() - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q2.Close() @@ -949,7 +980,7 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroups() { for i := 0; i < numMessages; i++ { msgID := fmt.Sprintf("msg-%d", i) messageIDs[i] = msgID - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), "partition-1", nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), "partition-1", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages to topic", numMessages) @@ -993,19 +1024,11 @@ func (s *SQLQueueIntegrationSuite) TestMultipleWorkersInConsumerGroup() { consumerGroup := "shared-group" // Create two workers in same consumer group - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q1.Close() - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q2.Close() @@ -1032,7 +1055,7 @@ func (s *SQLQueueIntegrationSuite) TestMultipleWorkersInConsumerGroup() { messageIDs[i] = msgID // Use different partition keys to allow distribution partitionKey := fmt.Sprintf("partition-%d", i%3) - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), partitionKey, nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), partitionKey, nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages to topic across multiple partitions", numMessages) @@ -1067,11 +1090,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentSubscribers() { totalMessages := numSubscribers * messagesPerSubscriber // Create publisher - pubQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + pubQueue, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQueue.Close() @@ -1082,11 +1101,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentSubscribers() { var deliveryChans []<-chan extqueue.Delivery for i := 0; i < numSubscribers; i++ { - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) queues = append(queues, q) @@ -1111,7 +1126,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentSubscribers() { for i := 0; i < totalMessages; i++ { msgID := fmt.Sprintf("concurrent-msg-%d", i) partitionKey := fmt.Sprintf("partition-%d", i%5) - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), partitionKey, nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), partitionKey, nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages", totalMessages) @@ -1140,12 +1155,9 @@ func (s *SQLQueueIntegrationSuite) TestDeadLetterQueue() { topic := "dlq_topic" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -1162,7 +1174,7 @@ func (s *SQLQueueIntegrationSuite) TestDeadLetterQueue() { require.NoError(t, err) // Publish a message that will fail - msg := entityqueue.NewMessage("poison-msg", []byte("poison"), "partition-1", nil) + msg := testMessage("poison-msg", []byte("poison"), "partition-1", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) t.Logf("Published poison message, will nack repeatedly") @@ -1243,11 +1255,7 @@ func (s *SQLQueueIntegrationSuite) TestMessageOrderingWithinPartition() { topic := "ordering_topic" partitionKey := "ordered-partition" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -1266,7 +1274,7 @@ func (s *SQLQueueIntegrationSuite) TestMessageOrderingWithinPartition() { for i := 0; i < numMessages; i++ { msgID := fmt.Sprintf("msg-%03d", i) messageIDs[i] = msgID - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("order-%d", i)), partitionKey, nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("order-%d", i)), partitionKey, nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages to same partition: %s", numMessages, partitionKey) @@ -1295,11 +1303,7 @@ func (s *SQLQueueIntegrationSuite) TestLateSubscriber() { topic := "late_subscriber_topic" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -1311,7 +1315,7 @@ func (s *SQLQueueIntegrationSuite) TestLateSubscriber() { for i := 0; i < numMessages; i++ { msgID := fmt.Sprintf("early-msg-%d", i) messageIDs[i] = msgID - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), "partition-1", nil) + msg := testMessage(msgID, []byte(fmt.Sprintf("data-%d", i)), "partition-1", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages BEFORE subscribing", numMessages) @@ -1348,12 +1352,9 @@ func (s *SQLQueueIntegrationSuite) TestEmptyTopicSubscribe() { topic := "empty_topic" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -1374,7 +1375,7 @@ func (s *SQLQueueIntegrationSuite) TestEmptyTopicSubscribe() { // Now publish a message publisher := q.Publisher() - msg := entityqueue.NewMessage("late-msg", []byte("data"), "partition-1", nil) + msg := testMessage("late-msg", []byte("data"), "partition-1", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) t.Logf("Published message to previously-empty topic") @@ -1391,11 +1392,7 @@ func (s *SQLQueueIntegrationSuite) TestGracefulShutdownDuringProcessing() { topic := "shutdown_topic" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) publisher := q.Publisher() @@ -1409,7 +1406,7 @@ func (s *SQLQueueIntegrationSuite) TestGracefulShutdownDuringProcessing() { // Publish messages numMessages := 5 for i := 0; i < numMessages; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("msg-%d", i), []byte("data"), "partition-1", nil) + msg := testMessage(fmt.Sprintf("msg-%d", i), []byte("data"), "partition-1", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } t.Logf("Published %d messages", numMessages) @@ -1438,11 +1435,7 @@ func (s *SQLQueueIntegrationSuite) TestGracefulShutdownDuringProcessing() { // Start new subscriber to verify all messages are redelivered. // Messages become visible after visibility timeout expires in DB. t.Logf("Starting new subscriber to verify message recovery...") - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q2.Close() @@ -1479,20 +1472,18 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ListTopicsAfterPublish() { t := s.T() topic := "admin_list_topics_test" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() // Publish messages publisher := q.Publisher() - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-1", []byte("a"), "p1", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-2", []byte("b"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-1", []byte("a"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-2", []byte("b"), "p1", nil))) // Verify via AdminStore admin := queueAdmin.NewAdminStore(s.db) - topics, err := admin.ListTopics(s.ctx) + topics, err := admin.ListTopics(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}) require.NoError(t, err) found := false @@ -1509,19 +1500,17 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_TopicStatsAfterPublish() { t := s.T() topic := "admin_stats_test" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() publisher := q.Publisher() - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("s1", []byte("x"), "p1", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("s2", []byte("y"), "p2", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("s3", []byte("z"), "p2", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("s1", []byte("x"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("s2", []byte("y"), "p2", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("s3", []byte("z"), "p2", nil))) admin := queueAdmin.NewAdminStore(s.db) - stats, err := admin.GetTopicStats(s.ctx, topic, "_dlq") + stats, err := admin.GetTopicStats(s.ctx, testTenant, topic, "_dlq") require.NoError(t, err) assert.Equal(t, int64(3), stats.TotalMessages) @@ -1535,18 +1524,17 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_InspectMessage() { t := s.T() topic := "admin_inspect_test" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() metadata := map[string]string{"env": "test", "trace": "abc"} publisher := q.Publisher() - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("inspect-1", []byte("payload-data"), "p1", metadata))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("inspect-1", []byte("payload-data"), "p1", metadata))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("inspect-1", []byte("other-partition"), "p2", nil))) admin := queueAdmin.NewAdminStore(s.db) - detail, found, err := admin.InspectMessage(s.ctx, topic, "inspect-1") + detail, found, err := admin.InspectMessage(s.ctx, testTenant, topic, "p1", "inspect-1") require.NoError(t, err) assert.True(t, found) assert.Equal(t, "inspect-1", detail.ID) @@ -1556,6 +1544,11 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_InspectMessage() { assert.Equal(t, "test", detail.Metadata["env"]) assert.Equal(t, "abc", detail.Metadata["trace"]) + otherDetail, found, err := admin.InspectMessage(s.ctx, testTenant, topic, "p2", "inspect-1") + require.NoError(t, err) + assert.True(t, found) + assert.Equal(t, []byte("other-partition"), otherDetail.Payload) + t.Logf("Inspect message verified: id=%s payload=%s metadata=%v", detail.ID, string(detail.Payload), detail.Metadata) } @@ -1563,36 +1556,39 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_DeleteAndPurge() { t := s.T() topic := "admin_delete_test" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() publisher := q.Publisher() - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("del-1", []byte("a"), "p1", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("del-2", []byte("b"), "p1", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("del-3", []byte("c"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("del-1", []byte("a"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("del-1", []byte("other"), "p2", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("del-2", []byte("b"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("del-3", []byte("c"), "p1", nil))) admin := queueAdmin.NewAdminStore(s.db) // Delete single message - affected, err := admin.DeleteMessage(s.ctx, topic, "del-1") + affected, err := admin.DeleteMessage(s.ctx, testTenant, topic, "p1", "del-1") require.NoError(t, err) assert.Equal(t, int64(1), affected) // Verify it's gone - _, found, err := admin.InspectMessage(s.ctx, topic, "del-1") + _, found, err := admin.InspectMessage(s.ctx, testTenant, topic, "p1", "del-1") require.NoError(t, err) assert.False(t, found) + _, found, err = admin.InspectMessage(s.ctx, testTenant, topic, "p2", "del-1") + require.NoError(t, err) + assert.True(t, found) + // Purge remaining - affected, err = admin.PurgeTopic(s.ctx, topic) + affected, err = admin.PurgeTopic(s.ctx, testTenant, topic) require.NoError(t, err) - assert.Equal(t, int64(2), affected) + assert.Equal(t, int64(3), affected) // Verify topic is empty - msgs, err := admin.ListMessages(s.ctx, topic, "", 50) + msgs, err := admin.ListMessages(s.ctx, testTenant, topic, "", 50) require.NoError(t, err) assert.Empty(t, msgs) } @@ -1603,9 +1599,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ConsumerLagAfterPartialAck() { topic := "admin_lag_test" consumerGroup := "lag-consumer" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -1614,7 +1608,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ConsumerLagAfterPartialAck() { // Publish 5 messages to same partition for i := 0; i < 5; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("lag-%d", i), []byte("data"), "lag-partition", nil) + msg := testMessage(fmt.Sprintf("lag-%d", i), []byte("data"), "lag-partition", nil) require.NoError(t, publisher.Publish(s.ctx, topic, msg)) } @@ -1631,7 +1625,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ConsumerLagAfterPartialAck() { // Check consumer lag — should show lag > 0 admin := queueAdmin.NewAdminStore(s.db) - lags, err := admin.ConsumerLag(s.ctx, topic) + lags, err := admin.ConsumerLag(s.ctx, testTenant, topic) require.NoError(t, err) require.NotEmpty(t, lags) @@ -1654,12 +1648,9 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { consumerGroup := "lease-consumer" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -1667,7 +1658,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { subscriber := q.Subscriber() // Publish and subscribe to create leases and offsets - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("lo-1", []byte("a"), "p1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("lo-1", []byte("a"), "p1", nil))) subConfig := extqueue.DefaultSubscriptionConfig("admin-worker-1", consumerGroup) subConfig.PartitionDiscoveryIntervalMs = 100 @@ -1688,7 +1679,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { for !offsetAdvanced { _, ok := <-signalCh require.True(t, ok, "signal channel closed before offset advanced") - offsets, err := admin.ListOffsets(s.ctx, consumerGroup) + offsets, err := admin.ListOffsets(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}, consumerGroup) require.NoError(t, err) for _, o := range offsets { if o.Topic == topic && o.OffsetAcked > 0 { @@ -1698,7 +1689,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { } // Verify leases are visible - leases, err := admin.ListLeases(s.ctx) + leases, err := admin.ListLeases(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}) require.NoError(t, err) var leaseFound bool @@ -1714,7 +1705,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { assert.True(t, leaseFound, "should find lease for consumer group %q", consumerGroup) // Verify offsets are visible - offsets, err := admin.ListOffsets(s.ctx, consumerGroup) + offsets, err := admin.ListOffsets(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}, consumerGroup) require.NoError(t, err) var offsetFound bool @@ -1734,9 +1725,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { topic := "admin_reset_test" consumerGroup := "reset-consumer" - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -1744,7 +1733,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { subscriber := q.Subscriber() // Publish, subscribe, ack — creates offsets and leases - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("r1", []byte("a"), "rp1", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("r1", []byte("a"), "rp1", nil))) subConfig := extqueue.DefaultSubscriptionConfig("reset-worker", consumerGroup) subConfig.PartitionDiscoveryIntervalMs = 100 @@ -1758,12 +1747,12 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { admin := queueAdmin.NewAdminStore(s.db) // Reset offset to 0 - affected, err := admin.ResetOffset(s.ctx, consumerGroup, topic, "rp1", 0) + affected, err := admin.ResetOffset(s.ctx, testTenant, consumerGroup, topic, "rp1", 0) require.NoError(t, err) assert.Equal(t, int64(1), affected) // Verify offset was reset - offsets, err := admin.ListOffsets(s.ctx, consumerGroup) + offsets, err := admin.ListOffsets(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}, consumerGroup) require.NoError(t, err) for _, o := range offsets { if o.Topic == topic && o.PartitionKey == "rp1" { @@ -1772,12 +1761,12 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { } // Release the lease - affected, err = admin.ReleaseLease(s.ctx, consumerGroup, topic, "rp1") + affected, err = admin.ReleaseLease(s.ctx, testTenant, consumerGroup, topic, "rp1") require.NoError(t, err) assert.Equal(t, int64(1), affected) // Verify lease is gone - leases, err := admin.ListLeases(s.ctx) + leases, err := admin.ListLeases(s.ctx, queueAdmin.TenantScope{Tenant: testTenant}) require.NoError(t, err) for _, l := range leases { if l.ConsumerGroup == consumerGroup && l.Topic == topic && l.PartitionKey == "rp1" { @@ -1795,8 +1784,8 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { // consumer group. func getPartitionLeases(db *sql.DB, topic, consumerGroup string) (map[string][]string, error) { rows, err := db.Query( - "SELECT leased_by, partition_key FROM queue_partition_leases WHERE topic = ? AND consumer_group = ? ORDER BY leased_by, partition_key", - topic, consumerGroup, + "SELECT leased_by, partition_key FROM queue_partition_leases WHERE tenant = ? AND topic = ? AND consumer_group = ? ORDER BY leased_by, partition_key", + testTenant, topic, consumerGroup, ) if err != nil { return nil, err @@ -1824,22 +1813,19 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_EvenDistribution() { signalCh := make(chan queueMySQL.HookSignal, 100) // Publish one message per partition so they are discoverable. - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("rb-even-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-even-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } // S1: subscribe, should acquire all 4 partitions (only subscriber). - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q1.Close() @@ -1852,10 +1838,9 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_EvenDistribution() { }, "S1 should acquire all 4 partitions") // S2: subscribe. After rebalancing, each should own 2. - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q2.Close() @@ -1880,29 +1865,25 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_SubscriberLeaves() { signalCh := make(chan queueMySQL.HookSignal, 100) // Publish messages. - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("rb-leave-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-leave-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } // S1 + S2 start, wait for 2+2 split. - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q1.Close() - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) // no defer close — we close explicitly below @@ -1930,8 +1911,8 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_SubscriberLeaves() { var s2Rows int require.NoError(t, s.db.QueryRowContext(s.ctx, ` SELECT COUNT(*) FROM queue_subscriber_heartbeats - WHERE consumer_group = ? AND topic = ? AND subscriber_name = ? - `, consumerGroup, topic, "s2").Scan(&s2Rows)) + WHERE tenant = ? AND consumer_group = ? AND topic = ? AND subscriber_name = ? + `, testTenant, consumerGroup, topic, "s2").Scan(&s2Rows)) assert.Equal(t, 0, s2Rows, "closed subscriber's heartbeat row must be deleted") t.Logf("Subscriber leave verified: S1 owns all 4 partitions after S2 departed") @@ -1946,28 +1927,24 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OddPartitions() { signalCh := make(chan queueMySQL.HookSignal, 100) - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("rb-odd-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-odd-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q1.Close() - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q2.Close() @@ -2003,14 +1980,12 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_NoOrphans() { signalCh := make(chan queueMySQL.HookSignal, 100) - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("rb-orphan-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-orphan-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } @@ -2018,10 +1993,9 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_NoOrphans() { queues := make([]extqueue.Queue, 3) subNames := []string{"s1", "s2", "s3"} for i, name := range subNames { - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) queues[i] = q _, err = q.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig(name, consumerGroup)) @@ -2063,14 +2037,12 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_MoreSubscribersThanPartitions() signalCh := make(chan queueMySQL.HookSignal, 100) - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("rb-excess-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-excess-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } @@ -2078,10 +2050,9 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_MoreSubscribersThanPartitions() subNames := []string{"s1", "s2", "s3", "s4"} var queues []extqueue.Queue for _, name := range subNames { - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) queues = append(queues, q) _, err = q.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig(name, consumerGroup)) @@ -2123,26 +2094,23 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_NoStarvation_UnevenSplit() { signalCh := make(chan queueMySQL.HookSignal, 100) - pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + pubQ, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer pubQ.Close() const partitionCount = 12 for i := 0; i < partitionCount; i++ { pk := fmt.Sprintf("pk-%02d", i) - msg := entityqueue.NewMessage(fmt.Sprintf("rb-starve-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("rb-starve-%d", i), []byte("x"), pk, nil) require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) } subNames := []string{"s1", "s2", "s3", "s4", "s5"} var queues []extqueue.Queue for _, name := range subNames { - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) queues = append(queues, q) // Nothing is acked in this test; a high retry budget keeps the @@ -2194,14 +2162,12 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OrphanSweep() { consumerGroup := "rebalance-sweep-cg" partitions := []string{"pk-a", "pk-b", "pk-c"} - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() for i, pk := range partitions { - msg := entityqueue.NewMessage(fmt.Sprintf("sweep-%d", i), []byte("x"), pk, nil) + msg := testMessage(fmt.Sprintf("sweep-%d", i), []byte("x"), pk, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) } @@ -2212,10 +2178,10 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OrphanSweep() { futureMs := time.Now().Add(10 * time.Minute).UnixMilli() for i := 0; i < 2; i++ { _, err := s.db.ExecContext(s.ctx, ` - INSERT INTO queue_subscriber_heartbeats (consumer_group, topic, subscriber_name, heartbeat_at, deregistered_at) - VALUES (?, ?, ?, ?, 0) + INSERT INTO queue_subscriber_heartbeats (tenant, consumer_group, topic, subscriber_name, heartbeat_at, deregistered_at) + VALUES (?, ?, ?, ?, ?, 0) ON DUPLICATE KEY UPDATE heartbeat_at = VALUES(heartbeat_at), deregistered_at = 0 - `, consumerGroup, topic, fmt.Sprintf("phantom-%d", i), futureMs) + `, testTenant, consumerGroup, topic, fmt.Sprintf("phantom-%d", i), futureMs) require.NoError(t, err) } @@ -2252,10 +2218,9 @@ func (s *SQLQueueIntegrationSuite) TestIdleLeaseRelease() { partition := "pk-idle" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2267,7 +2232,7 @@ func (s *SQLQueueIntegrationSuite) TestIdleLeaseRelease() { deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, cfg) require.NoError(t, err) - msg := entityqueue.NewMessage("idle-1", []byte("x"), partition, nil) + msg := testMessage("idle-1", []byte("x"), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) delivery := receive(t, deliveryChan) @@ -2279,8 +2244,8 @@ func (s *SQLQueueIntegrationSuite) TestIdleLeaseRelease() { rowCount := func(table string) int { var n int require.NoError(t, s.db.QueryRowContext(s.ctx, - "SELECT COUNT(*) FROM "+table+" WHERE consumer_group = ? AND topic = ?", - consumerGroup, topic).Scan(&n)) + "SELECT COUNT(*) FROM "+table+" WHERE tenant = ? AND consumer_group = ? AND topic = ?", + testTenant, consumerGroup, topic).Scan(&n)) return n } waitForCondition(t, signalCh, func() bool { @@ -2289,7 +2254,7 @@ func (s *SQLQueueIntegrationSuite) TestIdleLeaseRelease() { // Resurrection: a new message re-creates the partition through normal // discovery and is delivered like any other. - msg2 := entityqueue.NewMessage("idle-2", []byte("y"), partition, nil) + msg2 := testMessage("idle-2", []byte("y"), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg2)) delivery2 := receive(t, deliveryChan) @@ -2311,10 +2276,9 @@ func (s *SQLQueueIntegrationSuite) TestGCReclaimsAckedRowsUnderContinuousTraffic consumerGroup := "gc-busy-cg" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2323,7 +2287,7 @@ func (s *SQLQueueIntegrationSuite) TestGCReclaimsAckedRowsUnderContinuousTraffic const initialBatch = 200 for i := 0; i < initialBatch; i++ { require.NoError(t, q.Publisher().Publish(s.ctx, topic, - entityqueue.NewMessage(fmt.Sprintf("gc-%d", i), []byte("x"), partition, nil))) + testMessage(fmt.Sprintf("gc-%d", i), []byte("x"), partition, nil))) } // Fast poll so the 100-tick GC cadence elapses quickly; at the 100ms @@ -2336,8 +2300,8 @@ func (s *SQLQueueIntegrationSuite) TestGCReclaimsAckedRowsUnderContinuousTraffic countMessages := func() int { var n int require.NoError(t, s.db.QueryRowContext(s.ctx, - "SELECT COUNT(*) FROM queue_messages WHERE topic = ? AND partition_key = ?", - topic, partition).Scan(&n)) + "SELECT COUNT(*) FROM queue_messages WHERE tenant = ? AND topic = ? AND partition_key = ?", + testTenant, topic, partition).Scan(&n)) return n } @@ -2354,7 +2318,7 @@ func (s *SQLQueueIntegrationSuite) TestGCReclaimsAckedRowsUnderContinuousTraffic const continuousTrafficIterations = 150 for i := 0; i < continuousTrafficIterations; i++ { require.NoError(t, q.Publisher().Publish(s.ctx, topic, - entityqueue.NewMessage(fmt.Sprintf("gc-busy-%d", i), []byte("y"), partition, nil))) + testMessage(fmt.Sprintf("gc-busy-%d", i), []byte("y"), partition, nil))) delivery := receive(t, deliveryChan) require.NoError(t, delivery.Ack(s.ctx)) // Drain signals so the worker's blocking send cannot stall the traffic loop. @@ -2372,9 +2336,7 @@ func (s *SQLQueueIntegrationSuite) TestGCReclaimsAckedRowsUnderContinuousTraffic func (s *SQLQueueIntegrationSuite) TestInFlightMessageDoesNotBlockOtherMessages() { t := s.T() - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q.Close() @@ -2392,7 +2354,7 @@ func (s *SQLQueueIntegrationSuite) TestInFlightMessageDoesNotBlockOtherMessages( // Publish the first message alone and receive it, leaving it in flight // (un-finalized, invisible) at the lowest offset of the partition. - msg1 := entityqueue.NewMessage("msg-1", []byte("payload-1"), partition, nil) + msg1 := testMessage("msg-1", []byte("payload-1"), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg1)) d1 := receive(t, deliveryCh) assert.Equal(t, "msg-1", d1.Message().ID) @@ -2401,7 +2363,7 @@ func (s *SQLQueueIntegrationSuite) TestInFlightMessageDoesNotBlockOtherMessages( // Later offsets must still be deliverable despite the invisible msg-1 — // the opposite of a postponed message, which is a barrier. for i := 2; i <= 3; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), partition, nil) + msg := testMessage(fmt.Sprintf("msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) } @@ -2427,12 +2389,9 @@ func (s *SQLQueueIntegrationSuite) TestPostponeBlocksPartitionUntilDue() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2452,7 +2411,7 @@ func (s *SQLQueueIntegrationSuite) TestPostponeBlocksPartitionUntilDue() { // barrier must be in place before the later messages exist — deliveries // already fetched into the in-memory buffer are past the barrier by // design (it acts at the fetch layer). - msg1 := entityqueue.NewMessage("msg-1", []byte("payload-1"), partition, nil) + msg1 := testMessage("msg-1", []byte("payload-1"), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg1)) d1 := receive(t, deliveryCh) @@ -2463,7 +2422,7 @@ func (s *SQLQueueIntegrationSuite) TestPostponeBlocksPartitionUntilDue() { // Publish two more messages behind the postponed one for i := 2; i <= 3; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), partition, nil) + msg := testMessage(fmt.Sprintf("msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) } @@ -2497,12 +2456,9 @@ func (s *SQLQueueIntegrationSuite) TestPostponeResetsRetryBudget() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2515,7 +2471,7 @@ func (s *SQLQueueIntegrationSuite) TestPostponeResetsRetryBudget() { deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) - msg := entityqueue.NewMessage("wait-then-poison", []byte("payload"), "partition-1", nil) + msg := testMessage("wait-then-poison", []byte("payload"), "partition-1", nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) // First delivery: postpone briefly — a deliberate wait, not a failure @@ -2557,12 +2513,9 @@ func (s *SQLQueueIntegrationSuite) TestBatchSizeOneStrictSerialization() { t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2579,7 +2532,7 @@ func (s *SQLQueueIntegrationSuite) TestBatchSizeOneStrictSerialization() { // Publish 5 messages for i := 1; i <= 5; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("serial-%d", i), []byte(strconv.Itoa(i)), partition, nil) + msg := testMessage(fmt.Sprintf("serial-%d", i), []byte(strconv.Itoa(i)), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) } @@ -2605,12 +2558,9 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroupsIndependentState() t := s.T() signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2632,7 +2582,7 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroupsIndependentState() // Publish 2 messages for i := 1; i <= 2; i++ { - msg := entityqueue.NewMessage(fmt.Sprintf("shared-%d", i), []byte(strconv.Itoa(i)), partition, nil) + msg := testMessage(fmt.Sprintf("shared-%d", i), []byte(strconv.Itoa(i)), partition, nil) require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) } @@ -2680,19 +2630,15 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRejectDoesNotLoseMessages() { topic := "crash_reject_topic" - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) publisher := q1.Publisher() // Publish 3 messages to the same partition - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-A", []byte("A"), "same-part", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-B", []byte("B"), "same-part", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-C", []byte("C"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-A", []byte("A"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-B", []byte("B"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-C", []byte("C"), "same-part", nil))) // Subscribe with short timeouts for fast test subConfig := testSubConfig("worker-1", "crash-reject-cg") @@ -2726,12 +2672,9 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRejectDoesNotLoseMessages() { // Start worker-2 with same consumer group — it polls and finds msg-C // after lease + visibility expire in the DB signalCh := make(chan queueMySQL.HookSignal, 100) - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q2.Close() @@ -2768,7 +2711,7 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRejectDoesNotLoseMessages() { // Wait for the poll loop so advanceWatermark has run after all acks. waitForSignal(t, signalCh, queueMySQL.SignalDeliveryCheck) admin := queueAdmin.NewAdminStore(s.db) - lags, err := admin.ConsumerLag(s.ctx, topic) + lags, err := admin.ConsumerLag(s.ctx, testTenant, topic) require.NoError(t, err) for _, lag := range lags { if lag.ConsumerGroup == "crash-reject-cg" { @@ -2788,19 +2731,15 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRetryLimitDoesNotLoseMessages() topic := "crash_retry_limit_topic" - q1, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q1, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) publisher := q1.Publisher() // Publish 3 messages to the same partition - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-A", []byte("A"), "same-part", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-B", []byte("B"), "same-part", nil))) - require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("msg-C", []byte("C"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-A", []byte("A"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-B", []byte("B"), "same-part", nil))) + require.NoError(t, publisher.Publish(s.ctx, topic, testMessage("msg-C", []byte("C"), "same-part", nil))) // MaxAttempts=2: msg-B needs nack → redeliver → retry_count=2 → auto-DLQ. // Use standard visibility (2s) instead of 30s — event-driven waits make @@ -2851,11 +2790,7 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRetryLimitDoesNotLoseMessages() // Start worker-2 with same consumer group — it polls and finds messages // after lease + visibility expire in the DB - q2, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - }) + q2, err := queueMySQL.NewQueue(s.testQueueParams(t, nil)) require.NoError(t, err) defer q2.Close() @@ -2892,12 +2827,9 @@ func (s *SQLQueueIntegrationSuite) TestWatermarkAdvancesContiguously() { topic := "watermark_contiguous_topic" signalCh := make(chan queueMySQL.HookSignal, 100) - q, err := queueMySQL.NewQueue(queueMySQL.Params{ - DB: s.db, - Logger: zaptest.NewLogger(t), - MetricsScope: tally.NoopScope, - OnSignal: signalCh, - }) + q, err := queueMySQL.NewQueue(s.testQueueParams(t, func(p *queueMySQL.Params) { + p.OnSignal = signalCh + })) require.NoError(t, err) defer q.Close() @@ -2905,7 +2837,7 @@ func (s *SQLQueueIntegrationSuite) TestWatermarkAdvancesContiguously() { // Publish 5 messages to the same partition for i := 1; i <= 5; i++ { - msg := entityqueue.NewMessage( + msg := testMessage( fmt.Sprintf("wm-msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), "wm-part", diff --git a/test/integration/extension/messagequeue/mysql/tenant_isolation_test.go b/test/integration/extension/messagequeue/mysql/tenant_isolation_test.go new file mode 100644 index 000000000..be0167575 --- /dev/null +++ b/test/integration/extension/messagequeue/mysql/tenant_isolation_test.go @@ -0,0 +1,168 @@ +// 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 mysql + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber-go/tally" + "github.com/uber/submitqueue/platform/base/failure" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" + queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" + "go.uber.org/zap/zaptest" +) + +func (s *SQLQueueIntegrationSuite) tenantTopicRowCount(t *testing.T, table, tenant, topic string) int { + t.Helper() + var count int + err := s.db.QueryRowContext( + s.ctx, + "SELECT COUNT(*) FROM "+table+" WHERE tenant = ? AND topic = ?", + tenant, + topic, + ).Scan(&count) + require.NoError(t, err) + return count +} + +func (s *SQLQueueIntegrationSuite) TestTenantIsolationWithEqualMessageIdentities() { + t := s.T() + + const ( + tenantA = "tenant-a" + tenantB = "tenant-b" + topic = "tenant_isolation_topic" + partitionKey = "共享-partition" + messageID = "共享-message" + consumerGroup = "tenant-isolation-consumer" + ) + signalCh := make(chan queueMySQL.HookSignal, 100) + + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: s.db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + Tenants: []string{tenantA, tenantB}, + OnSignal: signalCh, + }) + require.NoError(t, err) + defer q.Close() + + cfg := testSubConfig("tenant-isolation-worker", consumerGroup) + cfg.VisibilityTimeoutMs = cfg.LeaseDurationMs * 10 + deliveries, err := q.Subscriber().Subscribe(s.ctx, topic, cfg) + require.NoError(t, err) + + for _, tenant := range []string{tenantA, tenantB} { + msg := entityqueue.NewMessage(messageID, []byte(tenant), partitionKey, nil) + msg.Tenant = tenant + require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) + } + + received := make(map[string]string, 2) + receivedDeliveries := make(map[string]extqueue.Delivery, 2) + receiveN(t, deliveries, 2, func(delivery extqueue.Delivery, _ int) { + msg := delivery.Message() + received[msg.Tenant] = string(msg.Payload) + receivedDeliveries[msg.Tenant] = delivery + }) + assert.Equal(t, map[string]string{tenantA: tenantA, tenantB: tenantB}, received) + + for _, table := range []string{ + "queue_messages", + "queue_delivery_state", + "queue_offsets", + "queue_partition_leases", + "queue_subscriber_heartbeats", + } { + for _, tenant := range []string{tenantA, tenantB} { + assert.Equal(t, 1, s.tenantTopicRowCount(t, table, tenant, topic), "%s rows for %s", table, tenant) + } + } + + require.NoError(t, receivedDeliveries[tenantB].Postpone(s.ctx, cfg.LeaseDurationMs*10)) + require.NoError(t, receivedDeliveries[tenantA].Ack(s.ctx)) + waitForCondition(t, signalCh, func() bool { + return s.tenantTopicRowCount(t, "queue_messages", tenantA, topic) == 0 + }, "acked tenant message was garbage-collected") + + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_messages", tenantB, topic)) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_delivery_state", tenantB, topic)) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_offsets", tenantB, topic)) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_partition_leases", tenantB, topic)) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_subscriber_heartbeats", tenantB, topic)) + + var tenantBOffset int64 + err = s.db.QueryRowContext( + s.ctx, + "SELECT offset_acked FROM queue_offsets WHERE tenant = ? AND topic = ? AND partition_key = ? AND consumer_group = ?", + tenantB, + topic, + partitionKey, + consumerGroup, + ).Scan(&tenantBOffset) + require.NoError(t, err) + assert.Zero(t, tenantBOffset) +} + +func (s *SQLQueueIntegrationSuite) TestTenantIsolationWhenMovingToDLQ() { + t := s.T() + + const ( + tenantA = "tenant-dlq-a" + tenantB = "tenant-dlq-b" + topic = "tenant_isolation_dlq_topic" + partitionKey = "shared-partition" + messageID = "shared-message" + ) + + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: s.db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + Tenants: []string{tenantA, tenantB}, + }) + require.NoError(t, err) + defer q.Close() + + cfg := testSubConfig("tenant-dlq-worker", "tenant-dlq-consumer") + cfg.VisibilityTimeoutMs = cfg.LeaseDurationMs * 10 + deliveries, err := q.Subscriber().Subscribe(s.ctx, topic, cfg) + require.NoError(t, err) + + for _, tenant := range []string{tenantA, tenantB} { + msg := entityqueue.NewMessage(messageID, []byte(tenant), partitionKey, nil) + msg.Tenant = tenant + require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) + } + + receivedDeliveries := make(map[string]extqueue.Delivery, 2) + receiveN(t, deliveries, 2, func(delivery extqueue.Delivery, _ int) { + receivedDeliveries[delivery.Message().Tenant] = delivery + }) + + reason := failure.New("tenant-scoped failure", failure.Subject{Type: "message", ID: messageID}) + require.NoError(t, receivedDeliveries[tenantA].Reject(s.ctx, reason)) + + assert.Zero(t, s.tenantTopicRowCount(t, "queue_messages", tenantA, topic)) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_messages", tenantA, topic+"_dlq")) + assert.Equal(t, 1, s.tenantTopicRowCount(t, "queue_messages", tenantB, topic)) + assert.Zero(t, s.tenantTopicRowCount(t, "queue_messages", tenantB, topic+"_dlq")) + + require.NoError(t, receivedDeliveries[tenantB].Ack(s.ctx)) +} diff --git a/tool/linter/queueshard/main.go b/tool/linter/queueshard/main.go index 3b32f95f2..d900bc11d 100644 --- a/tool/linter/queueshard/main.go +++ b/tool/linter/queueshard/main.go @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Command queueshard checks that every domain table is shardable by queue: its -// primary key must lead with the queue column, and no secondary index may span -// queues. +// Command queueshard checks that every table is shardable by its owning key: +// its primary key must lead with the shard column, and no secondary index may +// span shards. // // A table is shardable by queue when one queue's rows are unreachable through // another queue's binding. That holds exactly when the queue is the leading @@ -32,23 +32,20 @@ import ( "strings" ) -// queueColumns are the column names that identify the owning queue. Most tables -// call it "queue"; a table whose rows *are* queues (stovepipe's queue table) -// names it "name" because the row's identity is the queue itself. -var queueColumns = map[string]bool{ - "queue": true, - "name": true, +// schemaShardColumns identifies the allowed shard columns for each schema root. +var schemaShardColumns = map[string]map[string]bool{ + "submitqueue/extension/storage/mysql/schema": {"queue": true, "name": true}, + "stovepipe/extension/storage/mysql/schema": {"queue": true, "name": true}, + "platform/extension/counter/mysql/schema": {"queue": true, "name": true}, + "platform/extension/messagequeue/mysql/schema": {"tenant": true}, } // schemaRoots are the directories scanned for table definitions. -// -// platform/extension/messagequeue is deliberately absent: it is a message-queue -// backend keyed by (consumer_group, topic, partition_key), not a domain table -// set, and sharding it is tracked separately. var schemaRoots = []string{ "submitqueue/extension/storage/mysql/schema", "stovepipe/extension/storage/mysql/schema", "platform/extension/counter/mysql/schema", + "platform/extension/messagequeue/mysql/schema", } var ( @@ -76,6 +73,7 @@ func main() { var violations []violation var checked int for _, schemaRoot := range schemaRoots { + shardColumns := schemaShardColumns[schemaRoot] files, err := filepath.Glob(filepath.Join(root, schemaRoot, "*.sql")) if err != nil { fmt.Fprintf(os.Stderr, "error globbing %s: %v\n", schemaRoot, err) @@ -95,28 +93,27 @@ func main() { if relErr != nil { rel = file } - found, tableViolations := check(rel, string(content)) + found, tableViolations := check(rel, string(content), shardColumns) checked += found violations = append(violations, tableViolations...) } } if len(violations) > 0 { - fmt.Fprintf(os.Stderr, "%d table(s) are not shardable by queue:\n\n", len(violations)) + fmt.Fprintf(os.Stderr, "%d table(s) are not shardable:\n\n", len(violations)) for _, v := range violations { fmt.Fprintf(os.Stderr, " %s: table %q %s\n", v.file, v.table, v.problem) } - fmt.Fprintf(os.Stderr, "\nEvery table's primary key must lead with the queue column, and no\n") - fmt.Fprintf(os.Stderr, "secondary index may span queues, so that one queue's rows are\n") - fmt.Fprintf(os.Stderr, "unreachable through another queue's binding.\n") + fmt.Fprintf(os.Stderr, "\nEvery table's primary key must lead with its schema's shard column, and no\n") + fmt.Fprintf(os.Stderr, "secondary index may span shards.\n") os.Exit(1) } - fmt.Printf("All %d tables are shardable by queue.\n", checked) + fmt.Printf("All %d tables are shardable.\n", checked) } // check returns the number of tables found in content and any violations. -func check(file, content string) (int, []violation) { +func check(file, content string, shardColumns map[string]bool) (int, []violation) { var violations []violation matches := createTableRe.FindAllStringSubmatch(content, -1) for _, match := range matches { @@ -132,23 +129,29 @@ func check(file, content string) (int, []violation) { violations = append(violations, violation{file, table, "has an empty PRIMARY KEY"}) continue } - if !queueColumns[columns[0]] { + if !shardColumns[columns[0]] { violations = append(violations, violation{ file, table, - fmt.Sprintf("leads its PRIMARY KEY with %q, not the queue column", columns[0]), + fmt.Sprintf("leads its PRIMARY KEY with %q, not a shard column", columns[0]), }) } for _, idx := range indexRe.FindAllStringSubmatch(body, -1) { idxColumns := splitColumns(idx[2]) - if len(idxColumns) == 0 || !queueColumns[idxColumns[0]] { + if table == "queue_messages" && idx[1] == "idx_offset" && + len(idxColumns) == 1 && idxColumns[0] == "offset" { + // InnoDB requires AUTO_INCREMENT to be the leading column of some + // index; queue_messages keeps a dedicated offset-only index for that. + continue + } + if len(idxColumns) == 0 || !shardColumns[idxColumns[0]] { lead := "(empty)" if len(idxColumns) > 0 { lead = idxColumns[0] } violations = append(violations, violation{ file, table, - fmt.Sprintf("has index %q leading with %q, which spans queues", idx[1], lead), + fmt.Sprintf("has index %q leading with %q, which spans shards", idx[1], lead), }) } } diff --git a/tool/linter/queueshard/main_test.go b/tool/linter/queueshard/main_test.go index 2b5887b4f..551d93d2c 100644 --- a/tool/linter/queueshard/main_test.go +++ b/tool/linter/queueshard/main_test.go @@ -47,7 +47,7 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, wantViolations: 1, - wantProblem: `leads its PRIMARY KEY with "request_id", not the queue column`, + wantProblem: `leads its PRIMARY KEY with "request_id", not a shard column`, }, { name: "queue present but not leading is rejected", @@ -87,7 +87,7 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, wantViolations: 1, - wantProblem: `has index "idx_status" leading with "status", which spans queues`, + wantProblem: `has index "idx_status" leading with "status", which spans shards`, }, { name: "a queue-leading secondary index passes", @@ -100,6 +100,40 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, }, + { + name: "queue_messages offset index passes", + schema: "CREATE TABLE IF NOT EXISTS queue_messages (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + }, + { + name: "offset index on another table is rejected", + schema: "CREATE TABLE IF NOT EXISTS queue_offsets (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + wantViolations: 1, + wantProblem: `has index "idx_offset" leading with "offset", which spans shards`, + }, + { + name: "differently named queue_messages offset index is rejected", + schema: "CREATE TABLE IF NOT EXISTS queue_messages (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_unscoped_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + wantViolations: 1, + wantProblem: `has index "idx_unscoped_offset" leading with "offset", which spans shards`, + }, { name: "a table with no primary key is rejected", schema: "CREATE TABLE IF NOT EXISTS loose (\n" + @@ -118,7 +152,7 @@ func TestCheck(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - tables, violations := check("test.sql", tt.schema) + tables, violations := check("test.sql", tt.schema, map[string]bool{"queue": true, "name": true, "tenant": true}) assert.Equal(t, tt.wantTables, tables) require.Len(t, violations, tt.wantViolations) if tt.wantProblem != "" { From 58adda28471b579f2303cbf1019df11b8cde5502 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Mon, 7 Sep 2026 23:05:58 -0700 Subject: [PATCH 2/2] docs(messagequeue): comment ascii vs utf8mb4 column encodings in schema ## Summary ### Why? The CREATE TABLE statements encode two different VARCHAR(255) contracts. A reader opening the SQL should see the byte vs character limits, binary collations, and why tenant stays ASCII. ### What? - Comment ascii/ascii_bin vs utf8mb4/utf8mb4_bin on every queue table, with the full encoding note on queue_messages. Co-authored-by: Cursor --- .../messagequeue/mysql/schema/queue_delivery_state.sql | 2 ++ .../messagequeue/mysql/schema/queue_messages.sql | 9 ++++++++- .../messagequeue/mysql/schema/queue_offsets.sql | 2 ++ .../messagequeue/mysql/schema/queue_partition_leases.sql | 2 ++ .../mysql/schema/queue_subscriber_heartbeats.sql | 2 ++ 5 files changed, 16 insertions(+), 1 deletion(-) diff --git a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql index e561aef2c..47cf66b05 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql @@ -1,6 +1,8 @@ -- DELIVERY STATE TABLE -- Per-consumer-group delivery tracking for messages in the immutable log. -- Tracks visibility, ack state, and retry count independently per consumer group. +-- tenant/topic/consumer_group: VARCHAR(255) ascii/ascii_bin (255 bytes, byte-wise compare). +-- partition_key: VARCHAR(255) utf8mb4/utf8mb4_bin (255 Unicode chars). See queue_messages.sql. -- -- State encoding: -- acked = TRUE → processed, never redeliver diff --git a/platform/extension/messagequeue/mysql/schema/queue_messages.sql b/platform/extension/messagequeue/mysql/schema/queue_messages.sql index 87a178fec..a2693032b 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_messages.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_messages.sql @@ -1,7 +1,14 @@ -- MESSAGES TABLE (Immutable Log) --- Single table for all topics. tenant is the Vitess vindex; partition_key orders work within a tenant. +-- Single table for all topics. tenant is the shard key; partition_key orders work within a tenant. -- Messages are append-only; per-consumer-group delivery tracking is in queue_delivery_state. -- Example: tenant="monorepo/main", topic="merge_queue", partition_key="uber/cadence" +-- +-- Identifier encodings (InnoDB max key 3072 bytes; utf8mb4 counted at 4 bytes/char): +-- ascii COLLATE ascii_bin: bytes 0x00-0x7F, VARCHAR(255)=255 bytes, byte-wise compare (case-sensitive). +-- tenant, topic, original_topic (and consumer_group / subscriber_name / leased_by on other tables). +-- utf8mb4 COLLATE utf8mb4_bin: Unicode, VARCHAR(255)=255 chars (<=1020 bytes), code-point compare. +-- partition_key, id (caller-chosen; may be non-ASCII). +-- NOT NULL rejects SQL NULL; the Go backend also rejects empty operational IDs and non-ASCII tenants. CREATE TABLE IF NOT EXISTS queue_messages ( -- tenant is the shard isolation identity (SubmitQueue maps queueName here at wiring) diff --git a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql index 45e8439ff..5342dbd7f 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql @@ -1,6 +1,8 @@ -- CONSUMER OFFSETS TABLE -- Tracks consumption progress per consumer group + tenant + topic + partition. -- Each partition has independent offset tracking for crash recovery. +-- tenant/topic/consumer_group: VARCHAR(255) ascii/ascii_bin (255 bytes, byte-wise compare). +-- partition_key: VARCHAR(255) utf8mb4/utf8mb4_bin (255 Unicode chars). See queue_messages.sql. CREATE TABLE IF NOT EXISTS queue_offsets ( -- tenant is the shard isolation identity diff --git a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql index 47f1eebfe..f401a3670 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql @@ -1,6 +1,8 @@ -- PARTITION LEASES TABLE -- Tracks which worker has leased which partition for exclusive processing. -- Workers must renew leases to maintain ownership; stale leases can be stolen. +-- tenant/topic/consumer_group/leased_by: VARCHAR(255) ascii/ascii_bin (255 bytes, byte-wise compare). +-- partition_key: VARCHAR(255) utf8mb4/utf8mb4_bin (255 Unicode chars). See queue_messages.sql. CREATE TABLE IF NOT EXISTS queue_partition_leases ( -- tenant is the shard isolation identity diff --git a/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql b/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql index d4e7d22f5..6ccba4241 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql @@ -1,6 +1,8 @@ -- SUBSCRIBER HEARTBEATS TABLE -- Tracks active subscribers for fair partition leasing per tenant. -- Each subscriber registers itself with periodic heartbeat renewal. +-- tenant/topic/consumer_group/subscriber_name: VARCHAR(255) ascii/ascii_bin (255 bytes, byte-wise compare). +-- See queue_messages.sql for the ascii vs utf8mb4 identifier encodings. CREATE TABLE IF NOT EXISTS queue_subscriber_heartbeats ( -- tenant is the shard isolation identity