From c3c5b164de4913205706a6b6374b997d42298e81 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 4 Sep 2026 07:14:25 -0400 Subject: [PATCH 1/4] feat(messagequeue): shard MySQL queue tables by tenant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vitess needs a stable vindex that is not the Kafka-style partition key. Prefix every queue table and store query with tenant, discover partitions from the configured tenant list, and map SubmitQueue queueName onto tenant at wiring so one busy queue cannot scatter across shards. ## Summary ### Why? The MySQL message queue keyed every table on (topic, partition_key). That cannot be a Vitess vindex: SubmitQueue already uses partition_key for ordering (for example a build ID on build-signal), so hashing it would scatter one queue across shards and still mix tenants on a given shard. Isolation has to be an explicit tenant column that SubmitQueue maps from queueName at the wiring boundary. ### What? - Add Message.Tenant and require it on publish; platform/publish stamps it from queue_name metadata or context. - Rewrite the five queue schemas so PKs lead with tenant; drop cross-tenant indexes; keep KEY idx_offset for InnoDB AUTO_INCREMENT. - Drive subscriber discovery, leases, heartbeats, and fair-share from Params.Tenants instead of SELECT DISTINCT partition_key. - Wire tenants from SubmitQueue/Runway queue config and Stovepipe MQ_TENANTS. - Require --tenant on topic-scoped admin commands; teach queueshard the MQ schema root. ## Test Plan ✅ `./tool/bazel test //platform/extension/messagequeue/... //tool/linter/queueshard/...` Co-authored-by: Cursor # Conflicts: # platform/extension/messagequeue/mysql/schema/queue_messages.sql # service/runway/server/config.go # test/integration/extension/messagequeue/mysql/queue_test.go # Please enter the commit message for your changes. Lines starting # with '#' will be kept; you may remove them yourself if you want to. # An empty message aborts the commit. # # interactive rebase in progress; onto c9ee7a84 # Last command done (1 command done): # pick 8940d26f # feat(messagequeue): shard MySQL queue tables by tenant # Next command to do (1 remaining command): # pick c79649f8 # fix(messagequeue): harden tenant shard isolation # You are currently rebasing branch 'preetam/mq-tenant-shard' on 'c9ee7a84'. # # Changes to be committed: # modified: doc/rfc/index.md # new file: doc/rfc/messagequeue-tenant-sharding.md # modified: platform/base/messagequeue/message.go # modified: platform/extension/messagequeue/mysql/BUILD.bazel # modified: platform/extension/messagequeue/mysql/constants.go # modified: platform/extension/messagequeue/mysql/ctl/commands.go # modified: platform/extension/messagequeue/mysql/ctl/lib/admin.go # modified: platform/extension/messagequeue/mysql/ctl/lib/admin_test.go # modified: platform/extension/messagequeue/mysql/delivery_state_store.go # modified: platform/extension/messagequeue/mysql/delivery_state_store_test.go # modified: platform/extension/messagequeue/mysql/message_store.go # modified: platform/extension/messagequeue/mysql/message_store_test.go # modified: platform/extension/messagequeue/mysql/mock_stores.go # modified: platform/extension/messagequeue/mysql/offset_store.go # modified: platform/extension/messagequeue/mysql/offset_store_test.go # modified: platform/extension/messagequeue/mysql/partition_lease_store.go # modified: platform/extension/messagequeue/mysql/partition_lease_store_test.go # modified: platform/extension/messagequeue/mysql/publisher.go # modified: platform/extension/messagequeue/mysql/publisher_test.go # modified: platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql # modified: platform/extension/messagequeue/mysql/schema/queue_messages.sql # modified: platform/extension/messagequeue/mysql/schema/queue_offsets.sql # modified: platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql # modified: platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql # modified: platform/extension/messagequeue/mysql/sql.go # modified: platform/extension/messagequeue/mysql/stores.go # modified: platform/extension/messagequeue/mysql/subscriber.go # modified: platform/extension/messagequeue/mysql/subscriber_heartbeat_store.go # modified: platform/extension/messagequeue/mysql/subscriber_heartbeat_store_test.go # modified: platform/extension/messagequeue/mysql/subscriber_test.go # new file: platform/extension/messagequeue/mysql/tenants.go # modified: platform/publish/publish.go # modified: platform/publish/publish_test.go # modified: service/runway/server/config.go # modified: service/runway/server/main.go # modified: service/stovepipe/docker-compose.yml # modified: service/stovepipe/server/main.go # modified: service/submitqueue/gateway/server/main.go # modified: service/submitqueue/orchestrator/server/main.go # modified: test/e2e/runway/harness_test.go # modified: test/e2e/runway/suite_test.go # modified: test/e2e/submitqueue/BUILD.bazel # modified: test/e2e/submitqueue/harness_test.go # modified: test/integration/extension/messagequeue/mysql/queue_test.go # modified: test/integration/submitqueue/core/consumer/consumer_test.go # modified: test/integration/submitqueue/gateway/BUILD.bazel # modified: test/integration/submitqueue/gateway/suite_test.go # modified: tool/linter/queueshard/main.go # modified: tool/linter/queueshard/main_test.go # --- doc/rfc/index.md | 1 + doc/rfc/messagequeue-tenant-sharding.md | 91 +++ platform/base/messagequeue/message.go | 5 + .../extension/messagequeue/mysql/BUILD.bazel | 1 + .../extension/messagequeue/mysql/constants.go | 1 + .../messagequeue/mysql/ctl/commands.go | 97 ++- .../messagequeue/mysql/ctl/lib/admin.go | 161 +++-- .../messagequeue/mysql/ctl/lib/admin_test.go | 169 +++--- .../mysql/delivery_state_store.go | 86 +-- .../mysql/delivery_state_store_test.go | 52 +- .../messagequeue/mysql/message_store.go | 114 ++-- .../messagequeue/mysql/message_store_test.go | 48 +- .../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 | 6 +- .../messagequeue/mysql/publisher_test.go | 31 +- .../mysql/schema/queue_delivery_state.sql | 10 +- .../mysql/schema/queue_messages.sql | 36 +- .../mysql/schema/queue_offsets.sql | 19 +- .../mysql/schema/queue_partition_leases.sql | 18 +- .../schema/queue_subscriber_heartbeats.sql | 16 +- platform/extension/messagequeue/mysql/sql.go | 5 + .../extension/messagequeue/mysql/stores.go | 134 ++--- .../messagequeue/mysql/subscriber.go | 392 ++++++------ .../mysql/subscriber_heartbeat_store.go | 41 +- .../mysql/subscriber_heartbeat_store_test.go | 52 +- .../messagequeue/mysql/subscriber_test.go | 181 +++--- .../extension/messagequeue/mysql/tenants.go | 31 + platform/publish/publish.go | 5 + platform/publish/publish_test.go | 3 + service/runway/server/config.go | 8 + service/runway/server/main.go | 6 + service/stovepipe/docker-compose.yml | 1 + service/stovepipe/server/main.go | 7 + service/submitqueue/gateway/server/main.go | 33 +- .../submitqueue/orchestrator/server/main.go | 27 +- test/e2e/runway/harness_test.go | 1 + test/e2e/runway/suite_test.go | 9 + test/e2e/submitqueue/BUILD.bazel | 1 + test/e2e/submitqueue/harness_test.go | 4 +- .../messagequeue/mysql/queue_test.go | 567 +++++++----------- .../core/consumer/consumer_test.go | 17 +- .../submitqueue/gateway/BUILD.bazel | 1 + .../submitqueue/gateway/suite_test.go | 8 +- tool/linter/queueshard/main.go | 50 +- tool/linter/queueshard/main_test.go | 6 +- 49 files changed, 1598 insertions(+), 1417 deletions(-) create mode 100644 doc/rfc/messagequeue-tenant-sharding.md create mode 100644 platform/extension/messagequeue/mysql/tenants.go diff --git a/doc/rfc/index.md b/doc/rfc/index.md index d3319e661..5fa4ec3f0 100644 --- a/doc/rfc/index.md +++ b/doc/rfc/index.md @@ -5,6 +5,7 @@ Design documents and technical proposals, grouped by scope. Shared/cross-cutting ## Shared - [SQL-Based Distributed Queue](sql-queue-rfc.md) - MySQL-based distributed message queue with partition leasing and at-least-once delivery (used by SubmitQueue, Stovepipe, and other repo-local services) +- [Message Queue Tenant Sharding](messagequeue-tenant-sharding.md) - Per-tenant Vitess vindex on the platform MySQL message queue; SubmitQueue maps `queueName` to `tenant` at wiring - [Message Queue Contract](messagequeue-contract.md) - How queue payloads are defined (Protobuf, serialized as protobuf JSON), located by audience (external in `api/{domain}/messagequeue/`, internal in `{domain}/core/messagequeue/`), bound to topics (the `topics` proto option), and enforced by Bazel visibility - [Consumer Gate](consumer-gate.md) - Stopping and starting individual queue controllers at runtime via a consumer-side check: blocked deliveries are recorded as parked and postponed back to the queue (re-checked on redelivery), gate state as a separate extension with a file-based first implementation shared by tests and operators - [Consumer Hold](consumer-hold.md) - Fourth delivery outcome letting a controller postpone its delivery: the message becomes a partition barrier that pauses consumption for a chosen delay, redelivers in order, and does not count as a failure toward dead-lettering diff --git a/doc/rfc/messagequeue-tenant-sharding.md b/doc/rfc/messagequeue-tenant-sharding.md new file mode 100644 index 000000000..5455c6ccb --- /dev/null +++ b/doc/rfc/messagequeue-tenant-sharding.md @@ -0,0 +1,91 @@ +# RFC: Per-Tenant Vitess Sharding for the MySQL Message Queue + +## Metadata + +| Field | Value | +|-------|-------| +| **Author** | Preetam Dwivedi | +| **Status** | In Review | +| **Created** | 2026-09-03 | + +## Summary + +The platform MySQL message queue gains a domain-agnostic `tenant` column on every table. Vitess routes on `tenant`; every primary key and hot-path query leads with it. SubmitQueue, Stovepipe, and Runway map their `queueName` onto `tenant` at the wiring boundary. `partition_key` remains the ordering unit within a tenant and is not the vindex. + +## Background + +Domain storage (`request`, `batch`, `counter`, …) already shards by a leading `queue` column. The message queue backend was deliberately excluded from `tool/linter/queueshard` because it keyed rows by `(consumer_group, topic, partition_key)` with a global `AUTO_INCREMENT offset` — fine for a single MySQL instance, not for Vitess. + +SubmitQueue's business queue name already flows through publish metadata (`queue_name`) and consumer context (`WithQueueName`). Several pipeline stages use a different `partition_key` (build ID, request ID) so those deliveries serialize independently while still belonging to one business queue. Sharding on `partition_key` would split one queue across shards; the MQ needs a dedicated isolation column. + +## Naming + +| Layer | Column / field | Meaning | +|-------|----------------|---------| +| Platform MQ schema | `tenant` | Vitess vindex; opaque to the backend | +| Platform `Message` | `Tenant` | Persisted shard identity | +| SubmitQueue domain | `queue` / `queueName` | Same string as `tenant` at wiring | +| Platform MQ schema | `partition_key` | Ordering unit within `(tenant, topic)` | + +The MQ schema does not use `queue` — that word is overloaded (SubmitQueue domain, `Queue` interface, `queue_*` table prefix). + +## Schema + +Every table's primary key leads with `tenant`. Secondary indexes that do not lead with `tenant` are removed. + +### `queue_messages` + +- PK: `(tenant, topic, partition_key, offset)` +- Unique: `(tenant, topic, partition_key, id)` +- `offset` is per-partition, not global; fetch is `WHERE tenant=? AND topic=? AND partition_key=? AND offset>? ORDER BY offset` + +### `queue_delivery_state` + +- PK: `(tenant, consumer_group, topic, partition_key, message_offset)` + +### `queue_offsets` + +- PK: `(tenant, consumer_group, topic, partition_key)` +- Drop `idx_topic` + +### `queue_partition_leases` + +- PK: `(tenant, consumer_group, topic, partition_key)` +- Drop `idx_lease_renewed`; purge is scoped to `(tenant, consumer_group, topic)` + +### `queue_subscriber_heartbeats` + +- PK: `(tenant, consumer_group, topic, subscriber_name)` + +DLQ moves rewrite `topic` to `original + suffix` and keep `tenant` + `partition_key` on the same shard. + +## Subscriber discovery + +Today partition discovery runs `SELECT DISTINCT partition_key FROM queue_messages WHERE topic=?`, which scatter-gathers across all Vitess shards. + +The subscriber takes a configured tenant list (SubmitQueue: queue names from YAML). Discovery becomes: + +```sql +SELECT DISTINCT partition_key FROM queue_messages +WHERE tenant = ? AND topic = ? +ORDER BY partition_key +``` + +Fair-share, orphan sweep, and idle-lease release run per `(tenant, topic)`, not across all tenants on a topic. + +## Publish + +`platform/publish` stamps `Message.Tenant` from context metadata (`queue_name`). Empty tenant on publish is rejected. `PartitionKey` is unchanged. + +## Wiring + +One `extqueue.Queue` and one VTGate DSN per service. `NewQueue` / subscriber `Params` carry `Tenants []string`. Service `main.go` fills that from configured queue names. + +## Out of scope + +Live migration of existing Stovepipe prod queue databases (expand/contract, backfill, dual-write). This RFC describes a breaking greenfield schema; prod cutover is a separate exercise. + +## Related + +- [SQL-Based Distributed Queue](sql-queue-rfc.md) +- [Modular Queue Wiring](submitqueue/modular-queue-wiring.md) diff --git a/platform/base/messagequeue/message.go b/platform/base/messagequeue/message.go index 77ef1f00e..7c62872c2 100644 --- a/platform/base/messagequeue/message.go +++ b/platform/base/messagequeue/message.go @@ -48,6 +48,10 @@ type Message struct { // 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/extension/messagequeue/mysql/BUILD.bazel b/platform/extension/messagequeue/mysql/BUILD.bazel index 12b2f7a95..e22f38c10 100644 --- a/platform/extension/messagequeue/mysql/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/BUILD.bazel @@ -15,6 +15,7 @@ go_library( "stores.go", "subscriber.go", "subscriber_heartbeat_store.go", + "tenants.go", ], importpath = "github.com/uber/submitqueue/platform/extension/messagequeue/mysql", visibility = ["//visibility:public"], 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/commands.go b/platform/extension/messagequeue/mysql/ctl/commands.go index afb351265..f188a323e 100644 --- a/platform/extension/messagequeue/mysql/ctl/commands.go +++ b/platform/extension/messagequeue/mysql/ctl/commands.go @@ -123,10 +123,10 @@ func newListTopicsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { 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 @@ -135,12 +135,12 @@ func newListTopicsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { } 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 +149,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 +160,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,20 +197,22 @@ 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, 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, messageID) if err != nil { return err } @@ -219,6 +224,7 @@ func newInspectMessageCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command } headers := []string{"FIELD", "VALUE"} rows := [][]string{ + {"Tenant", detail.Tenant}, {"Offset", strconv.FormatInt(detail.Offset, 10)}, {"ID", detail.ID}, {"Topic", detail.Topic}, @@ -244,15 +250,17 @@ 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(&messageID, "message-id", "", "Message ID (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") cmd.MarkFlagRequired("message-id") return cmd } func newDeleteMessageCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Command { - var topic, messageID string + var tenant, topic, messageID string cmd := &cobra.Command{ Use: "delete-message", Short: "Delete a specific message", @@ -260,7 +268,7 @@ func newDeleteMessageCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Com if err := confirmAction(*noInteractive, fmt.Sprintf("Delete message %q from topic %q?", messageID, topic)); err != nil { return err } - affected, err := (*store).DeleteMessage(cmd.Context(), topic, messageID) + affected, err := (*store).DeleteMessage(cmd.Context(), tenant, topic, messageID) if err != nil { return err } @@ -272,15 +280,17 @@ func newDeleteMessageCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Com return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Tenant name (required)") cmd.Flags().StringVar(&topic, "topic", "", "Topic name (required)") cmd.Flags().StringVar(&messageID, "message-id", "", "Message ID (required)") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") 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 +298,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 +306,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 +342,40 @@ 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, 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, 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(&messageID, "message-id", "", "Message ID (required)") cmd.Flags().StringVar(&dlqSuffix, "dlq-suffix", "_dlq", "DLQ topic suffix") + cmd.MarkFlagRequired("tenant") cmd.MarkFlagRequired("topic") 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 +384,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 +392,32 @@ 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 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(), tenant, 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 +429,13 @@ func newListOffsetsCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Filter by tenant") 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 +444,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 +456,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 +469,23 @@ func newResetOffsetCmd(store **lib.AdminStore, noInteractive *bool) *cobra.Comma } func newListLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { - return &cobra.Command{ + var tenant string + 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(), tenant) 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 +498,17 @@ func newListLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Filter by tenant") + 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 +531,21 @@ 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 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(), tenant, thresholdMs) if err != nil { return err } @@ -527,10 +556,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 +573,13 @@ func newStaleLeasesCmd(store **lib.AdminStore, jsonOut *bool) *cobra.Command { return nil }, } + cmd.Flags().StringVar(&tenant, "tenant", "", "Filter by tenant") 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 +587,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 +599,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/lib/admin.go b/platform/extension/messagequeue/mysql/ctl/lib/admin.go index 8ccdc77a2..b9eeee7d2 100644 --- a/platform/extension/messagequeue/mysql/ctl/lib/admin.go +++ b/platform/extension/messagequeue/mysql/ctl/lib/admin.go @@ -37,6 +37,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 +72,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 +88,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 +106,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 +116,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 @@ -123,7 +133,7 @@ type TopicStats struct { // 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", + "SELECT tenant, topic, COUNT(*) FROM %s GROUP BY tenant, topic ORDER BY tenant, topic", mysql.MessagesTableName, ) rows, err := s.db.QueryContext(ctx, query) @@ -135,7 +145,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 +154,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 +169,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 +178,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 +187,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 +198,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 +221,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 +230,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, 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 id = ?", mysql.MessagesTableName), + tenant, topic, 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 } @@ -248,10 +258,10 @@ func (s *AdminStore) InspectMessage(ctx context.Context, topic string, messageID } // DeleteMessage deletes a specific message by topic and ID. -func (s *AdminStore) DeleteMessage(ctx context.Context, topic string, messageID string) (int64, error) { +func (s *AdminStore) DeleteMessage(ctx context.Context, tenant string, topic 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 id = ?", mysql.MessagesTableName), + tenant, topic, messageID, ) if err != nil { return 0, fmt.Errorf("delete message: %w", err) @@ -260,10 +270,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 +283,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, messageID string, dlqSuffix string) error { dlqTopic := topic + dlqSuffix tx, err := s.db.BeginTx(ctx, nil) @@ -289,8 +299,8 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str 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, + fmt.Sprintf("SELECT payload, metadata, partition_key, created_at, published_at FROM %s WHERE tenant = ? AND topic = ? AND id = ?", mysql.MessagesTableName), + tenant, dlqTopic, messageID, ).Scan(&payload, &metadataJSON, &partitionKey, &createdAt, &publishedAt) if err == sql.ErrNoRows { return fmt.Errorf("message %q not found in DLQ topic %q", messageID, dlqTopic) @@ -302,8 +312,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, partitionKey, messageID, payload, metadataJSON, createdAt, nowMs, ) if err != nil { return fmt.Errorf("insert requeued message: %w", err) @@ -311,8 +321,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 id = ?", mysql.MessagesTableName), + tenant, dlqTopic, messageID, ) if err != nil { return fmt.Errorf("delete dlq message: %w", err) @@ -322,18 +332,28 @@ func (s *AdminStore) RequeueDLQ(ctx context.Context, topic string, messageID str } // ListOffsets returns consumer group offsets, optionally filtered by group. -func (s *AdminStore) ListOffsets(ctx context.Context, consumerGroup string) ([]OffsetInfo, error) { +func (s *AdminStore) ListOffsets(ctx context.Context, tenant string, consumerGroup string) ([]OffsetInfo, error) { var rows *sql.Rows var err error - if consumerGroup != "" { + if tenant != "" && 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), + tenant, consumerGroup, + ) + } else if tenant != "" { + 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), + 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 +364,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 +373,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 consumer_group = ? AND topic = ? AND partition_key = ?", mysql.OffsetsTableName), + offset, nowMs, tenant, consumerGroup, topic, partition, ) if err != nil { return 0, fmt.Errorf("reset offset: %w", err) @@ -366,10 +386,14 @@ func (s *AdminStore) ResetOffset(ctx context.Context, consumerGroup, topic, part } // 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), - ) +func (s *AdminStore) ListLeases(ctx context.Context, tenant string) ([]LeaseInfo, error) { + 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 tenant != "" { + 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, tenant) + } + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("list leases: %w", err) } @@ -378,7 +402,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 +412,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 +430,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 +455,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 +470,15 @@ 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, tenant string, thresholdMs int64) ([]LeaseInfo, error) { 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 tenant != "" { + 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{tenant, cutoff} + } + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return nil, fmt.Errorf("stale leases: %w", err) } @@ -458,7 +487,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 +496,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..951dc75ff 100644 --- a/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go +++ b/platform/extension/messagequeue/mysql/ctl/lib/admin_test.go @@ -30,15 +30,16 @@ 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 GROUP BY tenant, topic ORDER BY tenant, topic"). WillReturnRows(rows) topics, err := store.ListTopics(context.Background()) 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,8 +54,8 @@ 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()) @@ -71,27 +72,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 +109,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 +133,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 +153,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 id = \\?"). + WithArgs("acme", "orders", "msg-1"). WillReturnRows(rows) - detail, found, err := store.InspectMessage(context.Background(), "orders", "msg-1") + detail, found, err := store.InspectMessage(context.Background(), "acme", "orders", "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 +177,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 id = \\?"). + WithArgs("acme", "orders", "missing"). WillReturnRows(rows) - _, found, err := store.InspectMessage(context.Background(), "orders", "missing") + _, found, err := store.InspectMessage(context.Background(), "acme", "orders", "missing") require.NoError(t, err) assert.False(t, found) assert.NoError(t, mock.ExpectationsWereMet()) @@ -191,11 +195,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 id = \\?"). + WithArgs("acme", "orders", "msg-1"). WillReturnResult(sqlmock.NewResult(0, 1)) - affected, err := store.DeleteMessage(context.Background(), "orders", "msg-1") + affected, err := store.DeleteMessage(context.Background(), "acme", "orders", "msg-1") require.NoError(t, err) assert.Equal(t, int64(1), affected) assert.NoError(t, mock.ExpectationsWereMet()) @@ -208,11 +212,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 +230,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"). + mock.ExpectQuery("SELECT .+ FROM queue_messages WHERE tenant = \\? AND topic = \\? AND id = \\?"). + WithArgs("acme", "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.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 id = \\?"). + WithArgs("acme", "orders_dlq", "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", "msg-1", "_dlq") require.NoError(t, err) assert.NoError(t, mock.ExpectationsWereMet()) } @@ -250,14 +254,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(), "", "") 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 +275,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(), "acme", "group-1") require.NoError(t, err) assert.Len(t, offsets, 1) assert.NoError(t, mock.ExpectationsWereMet()) @@ -290,10 +295,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", "group-1", "orders", "repo-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 +311,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(), "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 +333,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 +356,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 +382,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 +402,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(), "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,12 +423,12 @@ 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(), "", 60000) require.NoError(t, err) assert.Empty(t, leases) assert.NoError(t, mock.ExpectationsWereMet()) 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/message_store.go b/platform/extension/messagequeue/mysql/message_store.go index 78f9c9bf5..ecb85f0fa 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,38 @@ 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) + } + 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 +122,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 +139,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 +157,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 +177,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 +191,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 +206,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 +222,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 +243,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 +252,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 +264,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 +280,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 +327,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 +337,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 +350,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 +358,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 +371,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 +390,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..b8bb8e8de 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() @@ -83,7 +83,7 @@ func TestMessageStore_Insert(t *testing.T) { // 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 +104,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 +126,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 +145,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 +180,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 +218,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 +274,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 +282,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 +333,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..bb4a56579 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, consumer_group, topic, partition_key, offset_acked, updated_at) + VALUES (?, ?, ?, ?, 0, ?) + `, OffsetsTableName), tenant, consumerGroup, topic, partitionKey, 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 consumer_group = ? AND topic = ? AND partition_key = ? + `, OffsetsTableName), tenant, consumerGroup, topic, partitionKey).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 consumer_group = ? AND topic = ? AND partition_key = ? AND offset_acked < ? + `, OffsetsTableName), offset, now, tenant, consumerGroup, topic, partitionKey, 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 consumer_group = ? AND topic = ? AND partition_key = ? + `, OffsetsTableName), tenant, consumerGroup, topic, partitionKey) 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..f9ba7d8d5 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, testConsumerGroup, topic, partitionKey, 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, testConsumerGroup, "test_topic", "part1"). 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, testConsumerGroup, "test_topic", "part1"). 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, testConsumerGroup, topic, partitionKey, 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, testConsumerGroup, "test_topic", "part-1"). 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, testConsumerGroup, "test_topic", "part-1"). 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, testConsumerGroup, "test_topic", "part-1"). 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..af8cd2433 100644 --- a/platform/extension/messagequeue/mysql/publisher.go +++ b/platform/extension/messagequeue/mysql/publisher.go @@ -57,7 +57,11 @@ 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 := 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..26f109d55 100644 --- a/platform/extension/messagequeue/mysql/publisher_test.go +++ b/platform/extension/messagequeue/mysql/publisher_test.go @@ -54,24 +54,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 +88,7 @@ func TestPublisher_Publish(t *testing.T) { topic: "metadata_topic", messages: []entityqueue.Message{ { + Tenant: testTenant, ID: "msg_meta", Payload: []byte("payload"), PartitionKey: "part1", @@ -97,18 +98,18 @@ 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) }, }, } @@ -181,7 +182,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 +191,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 +212,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 +225,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 +248,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 +257,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..b084c206f 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql @@ -8,6 +8,9 @@ -- 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) NOT NULL, + -- Consumer group this delivery state belongs to consumer_group VARCHAR(255) 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..82efbb12e 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_messages.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_messages.sql @@ -1,19 +1,21 @@ -- 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) NOT NULL, - -- Topic identifies the queue type + -- Topic identifies the pipeline stage topic VARCHAR(255) NOT NULL, - -- Partition key for distributing work across workers - -- Example: repo ID, user ID, tenant ID + -- Partition key for distributing work across workers within a tenant partition_key VARCHAR(255) 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, @@ -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. 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..74f0c080f 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql @@ -1,12 +1,11 @@ -- 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 ( + -- tenant is the shard isolation identity + tenant VARCHAR(255) NOT NULL, + -- Consumer group consuming the topic consumer_group VARCHAR(255) 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, consumer_group, topic, partition_key) ) 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..d8dffdd34 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql @@ -1,12 +1,11 @@ -- 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) NOT NULL, + -- Consumer group (e.g., "orchestrator") consumer_group VARCHAR(255) NOT NULL, @@ -23,16 +22,7 @@ CREATE TABLE IF NOT EXISTS queue_partition_leases ( 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..ce817bc10 100644 --- a/platform/extension/messagequeue/mysql/sql.go +++ b/platform/extension/messagequeue/mysql/sql.go @@ -57,6 +57,10 @@ 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 this subscriber + // discovers and leases partitions for. Empty skips partition discovery. + Tenants []string } // NewQueue creates a new SQL-based queue @@ -102,6 +106,7 @@ func NewQueue(params Params) (extqueue.Queue, error) { leaseStore, heartbeatStore, deliveryStateStore, + params.Tenants, ) subscriber.OnSignal = params.OnSignal 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..994758d0e 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -21,6 +21,7 @@ import ( "math" "sort" "strconv" + "strings" "sync" "time" @@ -112,6 +113,7 @@ type subscriber struct { leaseStore partitionLeaseStore heartbeatStore subscriberHeartbeatStore deliveryStateStore deliveryStateStore + tenants []string mu sync.RWMutex closed bool @@ -155,23 +157,46 @@ type subscription struct { workersMu sync.Mutex // lastDiscoveredPartitions is cached from the most recent - // DiscoverAndAcquirePartitions call. Used by fairShareCap during - // rebalance to avoid a redundant discovery query. + // DiscoverAndAcquirePartitions calls, keyed by workerKey(tenant, partition). + // 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 tracks, per owned (tenant, partition) absent from discovery, + // when this subscriber first observed it drained (no stored messages left). + // Keys are workerKey(tenant, 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[string]time.Time } +// workerKey identifies a worker as tenant + partition; the same partition_key +// under two tenants must not share a worker, drain timer, or discovery cache entry. +func workerKey(tenant, partitionKey string) string { + return tenant + "\x00" + partitionKey +} + +func splitWorkerKey(key string) (tenant, partitionKey string) { + tenant, partitionKey, _ = strings.Cut(key, "\x00") + return +} + +func partitionKeysForTenant(keys []string, tenant string) []string { + out := make([]string, 0) + for _, key := range keys { + kTenant, pk := splitWorkerKey(key) + if kTenant == tenant { + out = append(out, pk) + } + } + return out +} + // 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 +222,7 @@ type sqlDelivery struct { // Backend-specific fields for ack/nack subscriber *subscriber + tenant string topic string partitionKey string offset int64 @@ -229,6 +255,7 @@ func newSQLDelivery( attempt int, metadata map[string]string, subscriber *subscriber, + tenant string, topic string, partitionKey string, offset int64, @@ -246,6 +273,7 @@ func newSQLDelivery( receivedAt: time.Now().UnixMilli(), metadata: metadata, subscriber: subscriber, + tenant: tenant, topic: topic, partitionKey: partitionKey, offset: offset, @@ -296,7 +324,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 +366,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 +392,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 +428,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 +436,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 +459,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 +475,7 @@ func NewSubscriber(logger *zap.SugaredLogger, scope tally.Scope, messageStore me leaseStore: leaseStore, heartbeatStore: heartbeatStore, deliveryStateStore: deliveryStateStore, + tenants: tenants, subscriptions: make(map[string]*subscription), } } @@ -463,24 +492,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) } } @@ -614,8 +643,10 @@ 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)...) + for _, tenant := range s.tenants { + if err := s.sendHeartbeat(ctx, sub, tenant); err != nil { + s.logger.Errorw("initial heartbeat failed", append(logFields, "tenant", tenant, "error", err)...) + } } for { @@ -642,59 +673,61 @@ 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)...) + for _, tenant := range s.tenants { + 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(ctx, 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(ctx, sub, tenant); err != nil { + s.logger.Errorw("heartbeat failed during lease error recovery", append(tenantFields, "error", err)...) + } + continue } - 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(ctx, 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(ctx, sub, tenant, kept); err != nil { + s.logger.Errorw("lease renewal failed", append(tenantFields, "error", err)...) + } + if err := s.sendHeartbeat(ctx, 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(ctx, 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(ctx, tenant, sub.topic, cfg.ConsumerGroup, leasePurgeAfterLeaseDurations*cfg.LeaseDurationMs); err != nil { + s.logger.Errorw("stale lease purge failed", append(tenantFields, "error", err)...) + } } s.emitSignal(SignalPartitionUpdate) @@ -722,118 +755,107 @@ 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 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) + allDiscovered := make([]string, 0) + allLeased := make([]string, 0) + + for _, tenant := range s.tenants { + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) if err != nil { - return fmt.Errorf("compute fair share cap: %w", err) + return fmt.Errorf("get leased partitions tenant=%s: %w", tenant, err) } - } - // 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) + cachedForTenant := partitionKeysForTenant(cachedDiscovered, tenant) + + maxPartitions := 0 + if !uncapped { + maxPartitions, err = s.fairShareCap(ctx, sub, tenant, leasedPartitions, cachedForTenant) + if err != nil { + return fmt.Errorf("compute fair share cap tenant=%s: %w", tenant, err) + } + } + + _, discoveredPartitions, err := s.leaseStore.DiscoverAndAcquirePartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, maxPartitions) + if err != nil { + return fmt.Errorf("discover and acquire partitions tenant=%s: %w", tenant, err) + } + + leasedPartitions, err = s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return fmt.Errorf("get leased partitions after acquire tenant=%s: %w", tenant, err) + } + + for _, pk := range discoveredPartitions { + allDiscovered = append(allDiscovered, workerKey(tenant, pk)) + } + for _, pk := range leasedPartitions { + allLeased = append(allLeased, workerKey(tenant, pk)) + } } - // 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, expired = updateDrainedTracking(sub.drainedSince, allLeased, allDiscovered, grace, time.Now()) 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. + for _, key := range expired { + tenant, pk := splitWorkerKey(key) + 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 } - 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. + 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 } - released[pk] = struct{}{} - delete(sub.drainedSince, pk) + released[key] = struct{}{} + delete(sub.drainedSince, key) - // 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) + s.stopPartitionWorker(sub, key) 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([]string, 0, len(allLeased)) + for _, key := range allLeased { + if _, ok := released[key]; !ok { + kept = append(kept, key) } } - leasedPartitions = kept + allLeased = kept } } - s.reconcilePartitionWorkers(ctx, sub, leasedPartitions) + s.reconcilePartitionWorkers(ctx, sub, allLeased) return nil } @@ -911,18 +933,21 @@ func (s *subscriber) reconcilePartitionWorkers(ctx context.Context, sub *subscri } // Start workers for newly leased partitions - for _, pk := range toStart { - s.startPartitionWorker(ctx, sub, pk) + for _, key := range toStart { + tenant, pk := splitWorkerKey(key) + s.startPartitionWorker(ctx, sub, tenant, pk) } } // 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, tenant, partitionKey string) { workerCtx, cancel := context.WithCancel(ctx) + key := workerKey(tenant, partitionKey) w := &partitionWorker{ + tenant: tenant, partitionKey: partitionKey, sub: sub, subscriber: s, @@ -931,13 +956,14 @@ func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription } sub.workersMu.Lock() - sub.workers[partitionKey] = w + sub.workers[key] = w sub.workersMu.Unlock() sub.workerWg.Add(1) go w.run(workerCtx) s.logger.Debugw("started partition worker", + "tenant", tenant, "topic", sub.topic, "partition_key", partitionKey, ) @@ -953,9 +979,9 @@ func (s *subscriber) startPartitionWorker(ctx context.Context, sub *subscription // 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, key string) { sub.workersMu.Lock() - w, ok := sub.workers[partitionKey] + w, ok := sub.workers[key] if !ok { sub.workersMu.Unlock() return @@ -968,17 +994,20 @@ 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, key) sub.workersMu.Unlock() + tenant, partitionKey := splitWorkerKey(key) select { case <-w.done: s.logger.Debugw("stopped partition worker", + "tenant", tenant, "topic", sub.topic, "partition_key", partitionKey, ) case <-time.After(workerStopTimeout): s.logger.Warnw("partition worker stop timeout, worker will drain in background", + "tenant", tenant, "topic", sub.topic, "partition_key", partitionKey, ) @@ -1029,6 +1058,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 +1067,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 +1093,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 +1103,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 +1125,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 +1147,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 +1172,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 +1188,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 +1248,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 +1272,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 +1309,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 +1317,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 +1325,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 +1339,26 @@ 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) - } + for _, tenant := range s.tenants { + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + return 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) + for _, partitionKey := range leasedPartitions { + if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + return fmt.Errorf("release lease tenant=%s partition=%s: %w", tenant, partitionKey, err) + } } } return nil } // 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,8 +1366,10 @@ 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) + for _, tenant := range s.tenants { + if err := s.heartbeatStore.Deregister(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + return fmt.Errorf("deregister heartbeat tenant=%s: %w", tenant, err) + } } return nil } @@ -1342,15 +1380,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 +1395,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, workerKey(tenant, pk)) s.logger.Infow("released partition for rebalance", + "tenant", tenant, "topic", sub.topic, "consumer_group", cfg.ConsumerGroup, "partition_key", pk, @@ -1400,10 +1434,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 } 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..bc49e99a6 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -39,9 +39,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, ) } @@ -53,22 +54,22 @@ func testSubscriptionConfig() extqueue.SubscriptionConfig { // 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 +79,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 +120,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 @@ -185,7 +186,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 +212,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 +280,7 @@ func TestSQLDelivery_Ack(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -290,7 +292,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 +348,7 @@ func TestSQLDelivery_Postpone(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -356,7 +359,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 +427,7 @@ func TestSQLDelivery_Reject(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{ @@ -439,19 +443,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 +561,7 @@ func TestSQLDelivery_NackDeadLettersWhenBudgetSpent(t *testing.T) { NewMockpartitionLeaseStore(ctrl), newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{Enabled: true, TopicSuffix: "_dlq"} @@ -566,14 +571,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 +603,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 +654,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() @@ -722,15 +728,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() @@ -743,19 +750,19 @@ func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { } // 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, workerKey(testTenant, pk)) } sub.workersMu.Unlock() @@ -784,6 +791,7 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -797,9 +805,9 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { 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 +816,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, @@ -892,6 +901,7 @@ func TestSubscriber_PollAndDeliver_GCOnBusyTicks(t *testing.T) { workers: make(map[string]*partitionWorker), } w := &partitionWorker{ + tenant: testTenant, partitionKey: "part-1", sub: sub, subscriber: s, @@ -906,13 +916,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 +984,7 @@ func TestSubscriber_PollAndDeliver_PostponedBarrier(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -987,32 +998,33 @@ func TestSubscriber_PollAndDeliver_PostponedBarrier(t *testing.T) { 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 +1064,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() @@ -1073,9 +1086,9 @@ func TestSubscriber_StopAllWorkers(t *testing.T) { } // 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, testTenant, "part-1") + s.startPartitionWorker(ctx, sub, testTenant, "part-2") + s.startPartitionWorker(ctx, sub, testTenant, "part-3") sub.workersMu.Lock() assert.Equal(t, 3, len(sub.workers)) @@ -1137,9 +1150,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 +1170,7 @@ func TestPartitionWorker_RunPollErrorLogging(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), newTestDeliveryStateStore(ctrl), + []string{testTenant}, ) s.OnSignal = make(chan HookSignal, 1) cfg := testSubscriptionConfig() @@ -1167,6 +1181,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 +1295,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 +1304,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 +1331,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 +1339,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sum := 0 @@ -1331,7 +1348,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,28 +1367,37 @@ func partitionKeysN(n int) []string { return keys } +func tenantPartitionKeys(tenant string, partitions []string) []string { + keys := make([]string, len(partitions)) + for i, partition := range partitions { + keys[i] = workerKey(tenant, 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", @@ -1380,7 +1406,7 @@ func TestSubscriber_RebalanceReleasesExcess(t *testing.T) { } 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 +1420,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,6 +1428,7 @@ 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", @@ -1409,10 +1436,10 @@ func TestSubscriber_RebalanceUnderCapReleasesNothing(t *testing.T) { workers: make(map[string]*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) } diff --git a/platform/extension/messagequeue/mysql/tenants.go b/platform/extension/messagequeue/mysql/tenants.go new file mode 100644 index 000000000..f3f191d89 --- /dev/null +++ b/platform/extension/messagequeue/mysql/tenants.go @@ -0,0 +1,31 @@ +// Copyright (c) 2025 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 ( + "os" + "strings" +) + +// ParseTenantsFromEnv reads MQ_TENANTS env var (comma-separated). +func ParseTenantsFromEnv() []string { + var out []string + for _, part := range strings.Split(os.Getenv("MQ_TENANTS"), ",") { + if trimmed := strings.TrimSpace(part); trimmed != "" { + out = append(out, trimmed) + } + } + return out +} diff --git a/platform/publish/publish.go b/platform/publish/publish.go index 53cbc7ed4..1e0497bc4 100644 --- a/platform/publish/publish.go +++ b/platform/publish/publish.go @@ -68,6 +68,11 @@ func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, k } msg := entityqueue.NewMessage(msgID, payload, partitionKey, metadataFromContext(ctx, metadata)) + if tenant := msg.Metadata[entityqueue.MetadataKeyQueueName]; tenant != "" { + msg.Tenant = tenant + } else if queueName, ok := entityqueue.QueueName(ctx); ok { + msg.Tenant = queueName + } return q.Publisher().Publish(ctx, topicName, msg) } diff --git a/platform/publish/publish_test.go b/platform/publish/publish_test.go index 23ec08ae7..adfce441d 100644 --- a/platform/publish/publish_test.go +++ b/platform/publish/publish_test.go @@ -78,6 +78,7 @@ func TestMessage_PropagatesQueueNameFromContext(t *testing.T) { ctx := entityqueue.WithQueueName(context.Background(), "monorepo/main") require.NoError(t, Message(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1")) assert.Equal(t, "monorepo/main", published.Metadata[entityqueue.MetadataKeyQueueName]) + assert.Equal(t, "monorepo/main", published.Tenant) } func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { @@ -99,6 +100,7 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { "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) } @@ -118,6 +120,7 @@ func TestMessageWithMetadata_ExplicitQueueNameWins(t *testing.T) { 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]) + assert.Equal(t, "outbound", published.Tenant) } func TestMessage_UnregisteredKey(t *testing.T) { diff --git a/service/runway/server/config.go b/service/runway/server/config.go index 94b64a09e..522faeaef 100644 --- a/service/runway/server/config.go +++ b/service/runway/server/config.go @@ -44,6 +44,14 @@ type mergeConfig struct { Queues []namedQueueMergeConfig `yaml:"queues"` } +func tenantNamesFromMergeConfig(cfg mergeConfig) []string { + names := make([]string, 0, len(cfg.Queues)) + for _, q := range cfg.Queues { + names = append(names, q.Name) + } + return names +} + // namedQueueMergeConfig is one queue's entry. type namedQueueMergeConfig struct { // Name is the queue this entry configures. diff --git a/service/runway/server/main.go b/service/runway/server/main.go index dc9985001..947769ab1 100644 --- a/service/runway/server/main.go +++ b/service/runway/server/main.go @@ -135,11 +135,17 @@ func run() error { } defer queueDB.Close() + mergeCfg, err := loadMergeConfigFromEnv(logger) + if err != nil { + return fmt.Errorf("failed to load merge config: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenantNamesFromMergeConfig(mergeCfg), }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) diff --git a/service/stovepipe/docker-compose.yml b/service/stovepipe/docker-compose.yml index ab4a82e10..c99d0c18b 100644 --- a/service/stovepipe/docker-compose.yml +++ b/service/stovepipe/docker-compose.yml @@ -70,6 +70,7 @@ services: # Level for the queue's own logs; info by default so its per-message # chatter does not bury the rest of the service at debug. - QUEUE_LOG_LEVEL=${QUEUE_LOG_LEVEL:-} + - MQ_TENANTS=${MQ_TENANTS:-monorepo/main} - HOSTNAME=stovepipe-dev depends_on: mysql-app: diff --git a/service/stovepipe/server/main.go b/service/stovepipe/server/main.go index 1e412c4be..ae6e0cf63 100644 --- a/service/stovepipe/server/main.go +++ b/service/stovepipe/server/main.go @@ -259,6 +259,7 @@ func run() error { Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: parseMQTenantsFromEnv(), }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -409,6 +410,12 @@ func run() error { return err } +// parseMQTenantsFromEnv reads MQ_TENANTS. Stovepipe accepts any queue at ingest +// but discovery requires configured tenants. +func parseMQTenantsFromEnv() []string { + return queueMySQL.ParseTenantsFromEnv() +} + // registerPrimaryControllers creates the primary-pipeline queue controllers and // registers them with c, returning how many were registered. func registerPrimaryControllers( diff --git a/service/submitqueue/gateway/server/main.go b/service/submitqueue/gateway/server/main.go index ba0d91d59..cc6c839db 100644 --- a/service/submitqueue/gateway/server/main.go +++ b/service/submitqueue/gateway/server/main.go @@ -241,12 +241,33 @@ func run() error { } defer queueDB.Close() - // Initialize queue + // Load queue configurations from YAML. Path is required so the gateway + // can reject requests for unknown queues at the edge. Queue names are + // also the MQ tenant list and must be known before NewQueue so the + // subscriber can discover partitions. + queueConfigPath := os.Getenv("QUEUE_CONFIG_PATH") + if queueConfigPath == "" { + return fmt.Errorf("QUEUE_CONFIG_PATH environment variable is required") + } + queueConfigs, err := yamlqueueconfig.NewStore(queueConfigPath) + if err != nil { + return fmt.Errorf("failed to load queue configs: %w", err) + } + configuredQueues, err := queueConfigs.List(ctx) + if err != nil { + return fmt.Errorf("failed to list queue configs: %w", err) + } + tenants := make([]string, 0, len(configuredQueues)) + for _, q := range configuredQueues { + tenants = append(tenants, q.Name) + } + 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) @@ -314,16 +335,6 @@ func run() error { if err != nil { return fmt.Errorf("failed to create storage: %w", err) } - // Load queue configurations from YAML. Path is required so the gateway - // can reject requests for unknown queues at the edge. - queueConfigPath := os.Getenv("QUEUE_CONFIG_PATH") - if queueConfigPath == "" { - return fmt.Errorf("QUEUE_CONFIG_PATH environment variable is required") - } - queueConfigs, err := yamlqueueconfig.NewStore(queueConfigPath) - if err != nil { - return fmt.Errorf("failed to load queue configs: %w", err) - } // Create controllers and wrap them for gRPC. Every store is queue-scoped and // resolves through the factory adapter; land/cancel/log share one materializer. diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index fe781b28a..25ecb1f0e 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -156,12 +156,21 @@ func run() error { } defer queueDB.Close() - // Initialize queue + // Build per-queue extension profiles (host-private). Each queue resolves + // to its own set of extension implementations (conflict analyzer, …), + // falling back to a baseline profile for queues without an explicit entry. + storageFty := storageFactory{backend: store} + profilesCfg, err := loadProfilesConfigFromEnv(logger) + if err != nil { + return fmt.Errorf("failed to load extension profiles: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenantNamesFromProfiles(profilesCfg), }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -176,14 +185,6 @@ func run() error { subscriberName = fmt.Sprintf("orchestrator-%d", time.Now().Unix()) } - // Build per-queue extension profiles (host-private). Each queue resolves - // to its own set of extension implementations (conflict analyzer, …), - // falling back to a baseline profile for queues without an explicit entry. - storageFty := storageFactory{backend: store} - profilesCfg, err := loadProfilesConfigFromEnv(logger) - if err != nil { - return fmt.Errorf("failed to load extension profiles: %w", err) - } profiles, err := newProfiles(ctx, logger, scope, changeset.New(storageFty), storageFty, profilesCfg) if err != nil { return fmt.Errorf("failed to build profiles: %w", err) @@ -341,6 +342,14 @@ func loadProfilesConfigFromEnv(logger *zap.Logger) (profilesConfig, error) { return cfg, nil } +func tenantNamesFromProfiles(cfg profilesConfig) []string { + names := make([]string, 0, len(cfg.Queues)) + for _, q := range cfg.Queues { + names = append(names, q.Name) + } + return names +} + // defaultProfilesConfig is the example topology used when no configuration file // is supplied: fake edge integrations everywhere, with the conflict analyzer // varied per queue to exercise its different behaviors. diff --git a/test/e2e/runway/harness_test.go b/test/e2e/runway/harness_test.go index 1fdf43e32..8e7dfdd99 100644 --- a/test/e2e/runway/harness_test.go +++ b/test/e2e/runway/harness_test.go @@ -168,6 +168,7 @@ func (s *RunwayE2ESuite) publishRaw(topic, id, partitionKey string, payload []by t := s.T() msg := entityqueue.NewMessage(id, payload, partitionKey, nil) + msg.Tenant = partitionKey require.NoError(t, s.queue.Publisher().Publish(s.ctx, topic, msg), "failed to publish %s to %s", id, topic) s.log.Logf("published %s to %s (partition %s)", id, topic, partitionKey) diff --git a/test/e2e/runway/suite_test.go b/test/e2e/runway/suite_test.go index 868bd9922..4fedcb77a 100644 --- a/test/e2e/runway/suite_test.go +++ b/test/e2e/runway/suite_test.go @@ -60,6 +60,14 @@ import ( "go.uber.org/zap/zaptest" ) +var runwayTestTenants = []string{ + "e2e-runway/merge", + "e2e-runway/check", + "e2e-runway/failed", + "e2e-runway/dlq", + "e2e-runway/undecodable", +} + type RunwayE2ESuite struct { suite.Suite ctx context.Context @@ -114,6 +122,7 @@ func (s *RunwayE2ESuite) SetupSuite() { DB: s.queueDB, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, + Tenants: runwayTestTenants, }) require.NoError(t, err, "failed to create queue client") t.Cleanup(func() { s.queue.Close() }) 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..08a49bbb3 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,7 +300,7 @@ 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, + require.NoError(t, publish.Message(entityqueue.WithQueueName(s.ctx, req.queue), registry, topickey.TopicKeyBatch, publish.UniqueID(req.sqid), payload, req.queue), "failed to redeliver the batch message") s.log.Logf("Redelivered the batch message for %s", req.sqid) } diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index a87503019..2006711e7 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() @@ -322,7 +343,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 +366,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 +382,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 +433,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 +450,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 +485,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 +501,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 +534,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 +554,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 +577,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 +596,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 +631,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 +656,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 +672,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 +688,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 +706,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 +742,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 +759,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 +767,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 +782,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 +806,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 +835,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 +851,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 +865,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 +893,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 +922,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 +966,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 +997,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 +1032,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 +1043,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 +1068,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 +1097,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 +1116,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 +1197,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 +1216,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 +1245,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 +1257,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 +1294,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 +1317,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 +1334,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 +1348,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 +1377,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,16 +1414,14 @@ 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) @@ -1509,19 +1442,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 +1466,16 @@ 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))) admin := queueAdmin.NewAdminStore(s.db) - detail, found, err := admin.InspectMessage(s.ctx, topic, "inspect-1") + detail, found, err := admin.InspectMessage(s.ctx, testTenant, topic, "inspect-1") require.NoError(t, err) assert.True(t, found) assert.Equal(t, "inspect-1", detail.ID) @@ -1563,36 +1492,34 @@ 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-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, "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, "del-1") require.NoError(t, err) assert.False(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) // 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 +1530,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 +1539,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 +1556,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 +1579,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 +1589,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 +1610,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, testTenant, consumerGroup) require.NoError(t, err) for _, o := range offsets { if o.Topic == topic && o.OffsetAcked > 0 { @@ -1698,7 +1620,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { } // Verify leases are visible - leases, err := admin.ListLeases(s.ctx) + leases, err := admin.ListLeases(s.ctx, testTenant) require.NoError(t, err) var leaseFound bool @@ -1714,7 +1636,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, testTenant, consumerGroup) require.NoError(t, err) var offsetFound bool @@ -1734,9 +1656,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 +1664,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 +1678,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, testTenant, consumerGroup) require.NoError(t, err) for _, o := range offsets { if o.Topic == topic && o.PartitionKey == "rp1" { @@ -1772,12 +1692,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, testTenant) require.NoError(t, err) for _, l := range leases { if l.ConsumerGroup == consumerGroup && l.Topic == topic && l.PartitionKey == "rp1" { @@ -1795,8 +1715,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 +1744,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 +1769,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 +1796,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 +1842,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 +1858,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 +1911,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 +1924,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 +1968,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 +1981,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 +2025,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 +2093,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 +2109,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 +2149,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 +2163,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 +2175,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 +2185,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 +2207,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 +2218,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 +2231,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 +2249,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 +2267,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 +2285,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 +2294,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 +2320,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 +2342,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 +2353,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 +2387,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 +2402,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 +2444,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 +2463,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 +2489,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 +2513,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 +2561,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 +2603,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 +2642,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 +2662,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 +2721,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 +2758,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 +2768,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/submitqueue/core/consumer/consumer_test.go b/test/integration/submitqueue/core/consumer/consumer_test.go index fd0ed9846..010a74bdb 100644 --- a/test/integration/submitqueue/core/consumer/consumer_test.go +++ b/test/integration/submitqueue/core/consumer/consumer_test.go @@ -54,6 +54,8 @@ const testTimeout = 10 * time.Second // stopTimeoutMs is the timeout in milliseconds for consumer.Stop(). const stopTimeoutMs = 10000 +const testTenant = "test-queue" + type ConsumerIntegrationSuite struct { suite.Suite ctx context.Context @@ -110,11 +112,18 @@ func (s *ConsumerIntegrationSuite) newQueue(t *testing.T) extqueue.Queue { DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, + Tenants: []string{testTenant}, }) require.NoError(t, err) return q } +func newTestMessage(id string, payload []byte, partitionKey string, metadata map[string]string) entityqueue.Message { + msg := entityqueue.NewMessage(id, payload, partitionKey, metadata) + msg.Tenant = testTenant + return msg +} + // newConsumer creates a consumer with a TopicRegistry wired to the given queue and topic. func (s *ConsumerIntegrationSuite) newConsumer(t *testing.T, q extqueue.Queue, topicKey consumer.TopicKey, topicName string, consumerGroup string) consumer.Consumer { t.Helper() @@ -197,7 +206,7 @@ func (s *ConsumerIntegrationSuite) TestConsumerPerPartitionIsolation() { require.NoError(t, c.Start(s.ctx)) // Publish to partition-a, wait for it to start blocking - msgA := entityqueue.NewMessage("iso-a", []byte("data-a"), "partition-a", nil) + msgA := newTestMessage("iso-a", []byte("data-a"), "partition-a", nil) require.NoError(t, publisher.Publish(s.ctx, topicName, msgA)) select { @@ -208,7 +217,7 @@ func (s *ConsumerIntegrationSuite) TestConsumerPerPartitionIsolation() { } // Now publish to partition-b — should be processed even though partition-a is blocked - msgB := entityqueue.NewMessage("iso-b", []byte("data-b"), "partition-b", nil) + msgB := newTestMessage("iso-b", []byte("data-b"), "partition-b", nil) require.NoError(t, publisher.Publish(s.ctx, topicName, msgB)) select { @@ -243,7 +252,7 @@ func (s *ConsumerIntegrationSuite) TestConsumerPartitionOrdering() { for i := range numMessages { msgID := fmt.Sprintf("order-%03d", i) publishedIDs[i] = msgID - msg := entityqueue.NewMessage(msgID, []byte(fmt.Sprintf("payload-%d", i)), "single-partition", nil) + msg := newTestMessage(msgID, []byte(fmt.Sprintf("payload-%d", i)), "single-partition", nil) require.NoError(t, publisher.Publish(s.ctx, topicName, msg)) } s.log.Logf("Published %d messages to single-partition", numMessages) @@ -316,7 +325,7 @@ func (s *ConsumerIntegrationSuite) TestConsumerMultiPartitionThroughput() { numPartitions := 3 for i := range numPartitions { partition := fmt.Sprintf("tp-partition-%d", i) - msg := entityqueue.NewMessage(fmt.Sprintf("tp-msg-%d", i), []byte("data"), partition, nil) + msg := newTestMessage(fmt.Sprintf("tp-msg-%d", i), []byte("data"), partition, nil) require.NoError(t, publisher.Publish(s.ctx, topicName, msg)) } s.log.Logf("Published 1 message to each of %d partitions", numPartitions) diff --git a/test/integration/submitqueue/gateway/BUILD.bazel b/test/integration/submitqueue/gateway/BUILD.bazel index ec8bc96bf..fe437299e 100644 --- a/test/integration/submitqueue/gateway/BUILD.bazel +++ b/test/integration/submitqueue/gateway/BUILD.bazel @@ -18,6 +18,7 @@ go_test( "//api/base/change/protopb:go_default_library", "//api/base/mergestrategy/protopb:go_default_library", "//api/submitqueue/gateway/protopb:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//submitqueue/core/request:go_default_library", diff --git a/test/integration/submitqueue/gateway/suite_test.go b/test/integration/submitqueue/gateway/suite_test.go index 6e3da7612..e8b29edb1 100644 --- a/test/integration/submitqueue/gateway/suite_test.go +++ b/test/integration/submitqueue/gateway/suite_test.go @@ -38,6 +38,7 @@ import ( changepb "github.com/uber/submitqueue/api/base/change/protopb" mergestrategypb "github.com/uber/submitqueue/api/base/mergestrategy/protopb" pb "github.com/uber/submitqueue/api/submitqueue/gateway/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" corerequest "github.com/uber/submitqueue/submitqueue/core/request" @@ -264,6 +265,8 @@ func (s *GatewayIntegrationSuite) TestReadAPIErrorCodes() { // entry to storage, observable through the request-summary RPC. func (s *GatewayIntegrationSuite) TestRequestLogConsumer() { t := s.T() + const sqid = "log-consumer-test/1" + const logQueue = "log-consumer-test" // Build a publisher against the shared queue database. NewQueue only wires up // stores; nothing consumes until a subscriber is started, so this publish-only @@ -272,6 +275,7 @@ func (s *GatewayIntegrationSuite) TestRequestLogConsumer() { DB: s.queueDB, Logger: zap.NewNop(), MetricsScope: tally.NoopScope, + Tenants: []string{logQueue}, }) require.NoError(t, err, "failed to create queue publisher") defer queue.Close() @@ -281,8 +285,6 @@ func (s *GatewayIntegrationSuite) TestRequestLogConsumer() { }) require.NoError(t, err, "failed to create topic registry") - const sqid = "log-consumer-test/1" - const logQueue = "log-consumer-test" store, err := mysqlstorage.NewStorage(s.db, tally.NoopScope) require.NoError(t, err) logQueueStore, err := store.For(logQueue) @@ -293,7 +295,7 @@ func (s *GatewayIntegrationSuite) TestRequestLogConsumer() { } require.NoError(t, logQueueStore.GetRequestSummaryStore().Create(s.ctx, summary)) logEntry := entity.NewRequestStatusLog(logQueue, sqid, entity.RequestStatusStarted, 1, "", nil) - require.NoError(t, corerequest.PublishLog(s.ctx, registry, logEntry, sqid, ""), + require.NoError(t, corerequest.PublishLog(entityqueue.WithQueueName(s.ctx, logQueue), registry, logEntry, sqid, ""), "failed to publish request log to log topic") s.log.Logf("Published 'started' log for sqid=%s; waiting for gateway consumer to persist it", sqid) diff --git a/tool/linter/queueshard/main.go b/tool/linter/queueshard/main.go index 3b32f95f2..2d41dbe2c 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,28 @@ 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 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..d2ff4542d 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", @@ -118,7 +118,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}) assert.Equal(t, tt.wantTables, tables) require.Len(t, violations, tt.wantViolations) if tt.wantProblem != "" { From 718d045afe2ba2029844414ebd3b2c58145dcb7d Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Mon, 7 Sep 2026 12:17:03 -0700 Subject: [PATCH 2/4] fix(messagequeue): harden tenant shard isolation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? The initial tenant-sharding implementation could exceed InnoDB's composite-key limit, infer incomplete tenant lists from extension overrides, and let one tenant's storage failure block unrelated tenants. Stovepipe could also accept work for a tenant its subscribers would never discover. ### What? - Use readable ASCII VARCHAR identifiers for tenants, topics, and consumer names while preserving utf8mb4 partition keys and message IDs. - Validate backend identifier contracts before database access and require an authoritative MQ_TENANTS list for consumer services. - Reject unconfigured Stovepipe queues before side effects. - Isolate discovery and shutdown failures per tenant and aggregate errors after all tenants are attempted. - Narrow the queueshard AUTO_INCREMENT exception and add schema, cross-tenant ACK/GC, and DLQ isolation coverage. ## Test Plan ✅ `./tool/bazel test //platform/extension/messagequeue/mysql/... //tool/linter/queueshard/... --test_output=errors` ✅ `./tool/bazel test //test/integration/extension/messagequeue/mysql:go_default_test --test_output=errors` ✅ `./tool/bazel build //service/runway/server/... //service/stovepipe/server/... //service/submitqueue/orchestrator/server/...` ✅ `make fmt && make gazelle` Co-authored-by: Cursor # Conflicts: # service/runway/server/config.go # service/runway/server/main.go # Please enter the commit message for your changes. Lines starting # with '#' will be kept; you may remove them yourself if you want to. # An empty message aborts the commit. # # interactive rebase in progress; onto c9ee7a84 # Last commands done (2 commands done): # pick 8940d26f # feat(messagequeue): shard MySQL queue tables by tenant # pick c79649f8 # fix(messagequeue): harden tenant shard isolation # No commands remaining. # You are currently rebasing branch 'preetam/mq-tenant-shard' on 'c9ee7a84'. # # Changes to be committed: # modified: doc/rfc/messagequeue-tenant-sharding.md # modified: platform/extension/messagequeue/mysql/BUILD.bazel # new file: platform/extension/messagequeue/mysql/identifier.go # modified: platform/extension/messagequeue/mysql/publisher.go # modified: platform/extension/messagequeue/mysql/publisher_test.go # modified: platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql # modified: platform/extension/messagequeue/mysql/schema/queue_messages.sql # modified: platform/extension/messagequeue/mysql/schema/queue_offsets.sql # modified: platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql # modified: platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql # modified: platform/extension/messagequeue/mysql/subscriber.go # modified: platform/extension/messagequeue/mysql/subscriber_test.go # modified: platform/extension/messagequeue/mysql/tenants.go # new file: platform/extension/messagequeue/mysql/tenants_test.go # modified: service/runway/server/config.go # modified: service/runway/server/docker-compose.yml # modified: service/runway/server/main.go # modified: service/stovepipe/server/main.go # modified: service/submitqueue/docker-compose.yml # modified: service/submitqueue/orchestrator/server/docker-compose.yml # modified: service/submitqueue/orchestrator/server/main.go # modified: stovepipe/controller/ingest.go # modified: stovepipe/controller/ingest_test.go # modified: test/integration/extension/messagequeue/mysql/BUILD.bazel # modified: test/integration/extension/messagequeue/mysql/queue_test.go # new file: test/integration/extension/messagequeue/mysql/tenant_isolation_test.go # modified: tool/linter/queueshard/main.go # modified: tool/linter/queueshard/main_test.go # --- doc/rfc/messagequeue-tenant-sharding.md | 11 +- .../extension/messagequeue/mysql/BUILD.bazel | 2 + .../messagequeue/mysql/identifier.go | 44 ++++ .../extension/messagequeue/mysql/publisher.go | 22 ++ .../messagequeue/mysql/publisher_test.go | 53 +++++ .../mysql/schema/queue_delivery_state.sql | 8 +- .../mysql/schema/queue_messages.sql | 10 +- .../mysql/schema/queue_offsets.sql | 8 +- .../mysql/schema/queue_partition_leases.sql | 10 +- .../schema/queue_subscriber_heartbeats.sql | 8 +- .../messagequeue/mysql/subscriber.go | 116 +++++++-- .../messagequeue/mysql/subscriber_test.go | 225 ++++++++++++++++++ .../extension/messagequeue/mysql/tenants.go | 15 ++ .../messagequeue/mysql/tenants_test.go | 72 ++++++ service/runway/server/config.go | 8 - service/runway/server/docker-compose.yml | 1 + service/runway/server/main.go | 6 +- service/stovepipe/server/main.go | 14 +- service/submitqueue/docker-compose.yml | 2 + .../orchestrator/server/docker-compose.yml | 1 + .../submitqueue/orchestrator/server/main.go | 14 +- stovepipe/controller/ingest.go | 38 +-- stovepipe/controller/ingest_test.go | 9 +- .../extension/messagequeue/mysql/BUILD.bazel | 5 +- .../messagequeue/mysql/queue_test.go | 36 +++ .../mysql/tenant_isolation_test.go | 168 +++++++++++++ tool/linter/queueshard/main.go | 3 +- tool/linter/queueshard/main_test.go | 36 ++- 28 files changed, 855 insertions(+), 90 deletions(-) create mode 100644 platform/extension/messagequeue/mysql/identifier.go create mode 100644 platform/extension/messagequeue/mysql/tenants_test.go create mode 100644 test/integration/extension/messagequeue/mysql/tenant_isolation_test.go diff --git a/doc/rfc/messagequeue-tenant-sharding.md b/doc/rfc/messagequeue-tenant-sharding.md index 5455c6ccb..ee09c407f 100644 --- a/doc/rfc/messagequeue-tenant-sharding.md +++ b/doc/rfc/messagequeue-tenant-sharding.md @@ -31,13 +31,14 @@ The MQ schema does not use `queue` — that word is overloaded (SubmitQueue doma ## Schema -Every table's primary key leads with `tenant`. Secondary indexes that do not lead with `tenant` are removed. +Every table's primary key leads with `tenant`. Tenant, topic, consumer-group, and subscriber identifiers use `VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin`; partition keys and message IDs use explicit `VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin`. This keeps operational identifiers readable, preserves unrestricted UTF-8 ordering keys and IDs, and keeps the largest composite key within InnoDB's 3072-byte limit. The backend validates each contract before database access. Secondary indexes that do not lead with `tenant` are removed except `queue_messages.idx_offset`, which InnoDB requires because the `AUTO_INCREMENT offset` column must be leftmost in an index. ### `queue_messages` - PK: `(tenant, topic, partition_key, offset)` - Unique: `(tenant, topic, partition_key, id)` -- `offset` is per-partition, not global; fetch is `WHERE tenant=? AND topic=? AND partition_key=? AND offset>? ORDER BY offset` +- Required InnoDB index: `idx_offset (offset)` for the `AUTO_INCREMENT` column +- `offset` is allocated from a shard-wide monotonic sequence and used as an ordering cursor within each partition; fetch is `WHERE tenant=? AND topic=? AND partition_key=? AND offset>? ORDER BY offset` ### `queue_delivery_state` @@ -63,7 +64,7 @@ DLQ moves rewrite `topic` to `original + suffix` and keep `tenant` + `partition_ Today partition discovery runs `SELECT DISTINCT partition_key FROM queue_messages WHERE topic=?`, which scatter-gathers across all Vitess shards. -The subscriber takes a configured tenant list (SubmitQueue: queue names from YAML). Discovery becomes: +The subscriber takes an explicit configured tenant list from `MQ_TENANTS`. Consumer processes reject an empty list at startup; Stovepipe also rejects ingest requests for names outside the list. Discovery becomes: ```sql SELECT DISTINCT partition_key FROM queue_messages @@ -71,7 +72,7 @@ WHERE tenant = ? AND topic = ? ORDER BY partition_key ``` -Fair-share, orphan sweep, and idle-lease release run per `(tenant, topic)`, not across all tenants on a topic. +Fair-share, orphan sweep, and idle-lease release run per `(tenant, topic)`, not across all tenants on a topic. Discovery and shutdown attempt every configured tenant and aggregate errors so one unavailable shard does not block unrelated tenants. ## Publish @@ -79,7 +80,7 @@ Fair-share, orphan sweep, and idle-lease release run per `(tenant, topic)`, not ## Wiring -One `extqueue.Queue` and one VTGate DSN per service. `NewQueue` / subscriber `Params` carry `Tenants []string`. Service `main.go` fills that from configured queue names. +One `extqueue.Queue` and one VTGate DSN per service. `NewQueue` / subscriber `Params` carry `Tenants []string`. Consumer service wiring parses the authoritative comma-separated `MQ_TENANTS` list once and passes it to the backend and any ingress validation. ## Out of scope diff --git a/platform/extension/messagequeue/mysql/BUILD.bazel b/platform/extension/messagequeue/mysql/BUILD.bazel index e22f38c10..b62131192 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", @@ -42,6 +43,7 @@ go_test( "sql_test.go", "subscriber_heartbeat_store_test.go", "subscriber_test.go", + "tenants_test.go", ], embed = [":go_default_library"], deps = [ diff --git a/platform/extension/messagequeue/mysql/identifier.go b/platform/extension/messagequeue/mysql/identifier.go new file mode 100644 index 000000000..a0a10da5c --- /dev/null +++ b/platform/extension/messagequeue/mysql/identifier.go @@ -0,0 +1,44 @@ +// 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" + "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) + } + } + 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 +} diff --git a/platform/extension/messagequeue/mysql/publisher.go b/platform/extension/messagequeue/mysql/publisher.go index af8cd2433..d7b4251f6 100644 --- a/platform/extension/messagequeue/mysql/publisher.go +++ b/platform/extension/messagequeue/mysql/publisher.go @@ -60,6 +60,28 @@ func (p *publisher) Publish(ctx context.Context, topic string, message entityque if message.Tenant == "" { return fmt.Errorf("publish: message tenant is required") } + 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 26f109d55..d7e166efa 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 @@ -112,6 +115,56 @@ func TestPublisher_Publish(t *testing.T) { 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, + }, } for _, tt := range tests { diff --git a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql index b084c206f..e561aef2c 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql @@ -9,16 +9,16 @@ CREATE TABLE IF NOT EXISTS queue_delivery_state ( -- tenant is the shard isolation identity - tenant VARCHAR(255) NOT NULL, + 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, diff --git a/platform/extension/messagequeue/mysql/schema/queue_messages.sql b/platform/extension/messagequeue/mysql/schema/queue_messages.sql index 82efbb12e..87a178fec 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_messages.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_messages.sql @@ -5,19 +5,19 @@ CREATE TABLE IF NOT EXISTS queue_messages ( -- tenant is the shard isolation identity (SubmitQueue maps queueName here at wiring) - tenant VARCHAR(255) NOT NULL, + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic identifies the pipeline stage - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition key for distributing work across workers within a tenant - partition_key VARCHAR(255) NOT NULL, + 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, @@ -31,7 +31,7 @@ CREATE TABLE IF NOT EXISTS queue_messages ( failed_at BIGINT UNSIGNED NOT NULL, failure_count INT UNSIGNED NOT NULL, last_error TEXT NOT NULL, - original_topic VARCHAR(255) NOT NULL, + original_topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, failure_detail JSON, PRIMARY KEY (tenant, topic, partition_key, offset), diff --git a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql index 74f0c080f..4fa0d0c5d 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql @@ -4,16 +4,16 @@ CREATE TABLE IF NOT EXISTS queue_offsets ( -- tenant is the shard isolation identity - tenant VARCHAR(255) NOT NULL, + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Consumer group consuming the topic - 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 being consumed - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Last offset that was successfully acked for this partition offset_acked BIGINT UNSIGNED NOT NULL, diff --git a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql index d8dffdd34..47f1eebfe 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql @@ -4,19 +4,19 @@ CREATE TABLE IF NOT EXISTS queue_partition_leases ( -- tenant is the shard isolation identity - tenant VARCHAR(255) NOT NULL, + 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, diff --git a/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql b/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql index 49b712dc0..d4e7d22f5 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_subscriber_heartbeats.sql @@ -4,16 +4,16 @@ CREATE TABLE IF NOT EXISTS queue_subscriber_heartbeats ( -- tenant is the shard isolation identity - tenant VARCHAR(255) NOT NULL, + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- consumer_group identifies the consumer group this subscriber belongs to - consumer_group VARCHAR(255) NOT NULL, + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- topic is the topic this subscriber is consuming from - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- subscriber_name uniquely identifies this subscriber within the consumer group - subscriber_name VARCHAR(255) NOT NULL, + subscriber_name VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- heartbeat_at is the Unix timestamp in milliseconds of the last heartbeat heartbeat_at BIGINT UNSIGNED NOT NULL, diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index 994758d0e..9ed6b98d2 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -531,6 +531,26 @@ func (s *subscriber) Subscribe(ctx context.Context, topic string, config extqueu if closed { return nil, ErrSubscriberClosed } + if len(s.tenants) == 0 { + return nil, fmt.Errorf("subscribe topic %q: %w: no tenants configured", topic, ErrInvalidConfig) + } + 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) + } + } + for _, tenant := range s.tenants { + if err := validateASCIIIdentifier("tenant", tenant); 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) } @@ -762,16 +782,22 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc cfg := sub.config sub.workersMu.Lock() - cachedDiscovered := sub.lastDiscoveredPartitions + cachedDiscovered := append([]string(nil), sub.lastDiscoveredPartitions...) + existingWorkers := make([]string, 0, len(sub.workers)) + for key := range sub.workers { + existingWorkers = append(existingWorkers, key) + } sub.workersMu.Unlock() - allDiscovered := make([]string, 0) - allLeased := make([]string, 0) + discoveredByTenant := make(map[string][]string, len(s.tenants)) + leasedByTenant := make(map[string][]string, len(s.tenants)) + var discoveryErrs []error for _, tenant := range s.tenants { leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) if err != nil { - return fmt.Errorf("get leased partitions tenant=%s: %w", tenant, err) + discoveryErrs = append(discoveryErrs, fmt.Errorf("get leased partitions tenant=%s: %w", tenant, err)) + continue } cachedForTenant := partitionKeysForTenant(cachedDiscovered, tenant) @@ -780,35 +806,86 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc if !uncapped { maxPartitions, err = s.fairShareCap(ctx, sub, tenant, leasedPartitions, cachedForTenant) if err != nil { - return fmt.Errorf("compute fair share cap tenant=%s: %w", tenant, err) + discoveryErrs = append(discoveryErrs, fmt.Errorf("compute fair share cap tenant=%s: %w", tenant, err)) + continue } } _, discoveredPartitions, err := s.leaseStore.DiscoverAndAcquirePartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, maxPartitions) if err != nil { - return fmt.Errorf("discover and acquire partitions tenant=%s: %w", tenant, err) + discoveryErrs = append(discoveryErrs, fmt.Errorf("discover and acquire partitions tenant=%s: %w", tenant, err)) + continue } leasedPartitions, err = s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) if err != nil { - return fmt.Errorf("get leased partitions after acquire tenant=%s: %w", tenant, err) + discoveryErrs = append(discoveryErrs, fmt.Errorf("get leased partitions after acquire tenant=%s: %w", tenant, err)) + continue } + discoveredByTenant[tenant] = discoveredPartitions + leasedByTenant[tenant] = leasedPartitions + } + + allDiscovered := make([]string, 0) + allLeased := make([]string, 0) + nextDrainedSince := make(map[string]time.Time, len(sub.drainedSince)) + grace := time.Duration(idleLeaseReleaseAfterLeaseDurations*cfg.LeaseDurationMs) * time.Millisecond + now := time.Now() + var expired []string + + for _, tenant := range s.tenants { + discoveredPartitions, succeeded := discoveredByTenant[tenant] + if !succeeded { + for _, pk := range partitionKeysForTenant(cachedDiscovered, tenant) { + allDiscovered = append(allDiscovered, workerKey(tenant, pk)) + } + for _, pk := range partitionKeysForTenant(existingWorkers, tenant) { + allLeased = append(allLeased, workerKey(tenant, pk)) + } + for key, since := range sub.drainedSince { + keyTenant, _ := splitWorkerKey(key) + if keyTenant == tenant { + nextDrainedSince[key] = since + } + } + continue + } + + leasedPartitions := leasedByTenant[tenant] + tenantDiscovered := make([]string, 0, len(discoveredPartitions)) for _, pk := range discoveredPartitions { - allDiscovered = append(allDiscovered, workerKey(tenant, pk)) + key := workerKey(tenant, pk) + tenantDiscovered = append(tenantDiscovered, key) + allDiscovered = append(allDiscovered, key) } + tenantLeased := make([]string, 0, len(leasedPartitions)) for _, pk := range leasedPartitions { - allLeased = append(allLeased, workerKey(tenant, pk)) + key := workerKey(tenant, pk) + tenantLeased = append(tenantLeased, key) + allLeased = append(allLeased, key) } + + previouslyDrained := make(map[string]time.Time) + for key, since := range sub.drainedSince { + keyTenant, _ := splitWorkerKey(key) + if keyTenant == tenant { + previouslyDrained[key] = since + } + } + tracked, tenantExpired := updateDrainedTracking(previouslyDrained, tenantLeased, tenantDiscovered, grace, now) + for key, since := range tracked { + nextDrainedSince[key] = since + } + expired = append(expired, tenantExpired...) } sub.workersMu.Lock() sub.lastDiscoveredPartitions = allDiscovered sub.workersMu.Unlock() - grace := time.Duration(idleLeaseReleaseAfterLeaseDurations*cfg.LeaseDurationMs) * time.Millisecond - var expired []string - sub.drainedSince, expired = updateDrainedTracking(sub.drainedSince, allLeased, allDiscovered, grace, time.Now()) + sub.drainedSince = nextDrainedSince + sort.Strings(expired) if len(expired) > 0 { released := make(map[string]struct{}, len(expired)) for _, key := range expired { @@ -856,7 +933,7 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc } s.reconcilePartitionWorkers(ctx, sub, allLeased) - return nil + return errors.Join(discoveryErrs...) } // updateDrainedTracking recomputes, for every owned partition absent from @@ -1339,19 +1416,21 @@ func (s *subscriber) renewLeases(ctx context.Context, sub *subscription, tenant // releaseAllLeases releases all leases for a topic. func (s *subscriber) releaseAllLeases(ctx context.Context, sub *subscription) error { cfg := sub.config + var releaseErrs []error for _, tenant := range s.tenants { leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) if err != nil { - return fmt.Errorf("get leased partitions for release tenant=%s: %w", tenant, err) + releaseErrs = append(releaseErrs, fmt.Errorf("get leased partitions for release tenant=%s: %w", tenant, err)) + continue } for _, partitionKey := range leasedPartitions { if err := s.leaseStore.ReleaseLease(ctx, tenant, sub.topic, partitionKey, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("release lease tenant=%s partition=%s: %w", tenant, partitionKey, err) + releaseErrs = append(releaseErrs, fmt.Errorf("release lease tenant=%s partition=%s: %w", tenant, partitionKey, err)) } } } - return nil + return errors.Join(releaseErrs...) } // sendHeartbeat sends a heartbeat for this subscriber. @@ -1366,12 +1445,13 @@ func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription, tenan // deregisterHeartbeat removes this subscriber's heartbeat entry during shutdown. func (s *subscriber) deregisterHeartbeat(ctx context.Context, sub *subscription) error { cfg := sub.config + var deregistrationErrs []error for _, tenant := range s.tenants { if err := s.heartbeatStore.Deregister(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("deregister heartbeat tenant=%s: %w", tenant, err) + deregistrationErrs = append(deregistrationErrs, fmt.Errorf("deregister heartbeat tenant=%s: %w", tenant, err)) } } - return nil + return errors.Join(deregistrationErrs...) } // rebalance checks if this subscriber holds more partitions than its fair share diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index bc49e99a6..3b2745e08 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" @@ -179,6 +180,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() @@ -773,6 +819,185 @@ func TestSubscriber_ReconcilePartitionWorkers(t *testing.T) { } } +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) + + gomock.InOrder( + 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 := workerKey(tenantFailed, "failed-existing") + failedDrainSince := time.Now().Add(-time.Hour) + sub := &subscription{ + topic: "test-topic", + config: cfg, + deliveryCh: make(chan extqueue.Delivery, 3), + workers: map[string]*partitionWorker{ + failedWorkerKey: { + cancelFunc: func() {}, + done: failedWorkerDone, + }, + }, + lastDiscoveredPartitions: []string{workerKey(tenantFailed, "failed-discovered")}, + drainedSince: map[string]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([]string, 0, len(sub.workers)) + for key := range sub.workers { + workerKeys = append(workerKeys, key) + } + discovered := append([]string(nil), sub.lastDiscoveredPartitions...) + sub.workersMu.Unlock() + + assert.ElementsMatch(t, []string{ + workerKey(tenantBefore, "before-new"), + failedWorkerKey, + workerKey(tenantAfter, "after-new"), + }, workerKeys) + assert.ElementsMatch(t, []string{ + workerKey(tenantBefore, "before-new"), + workerKey(tenantFailed, "failed-discovered"), + workerKey(tenantAfter, "after-new"), + }, discovered) + assert.Equal(t, failedDrainSince, sub.drainedSince[failedWorkerKey]) +} + +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() + + gomock.InOrder( + 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() + + gomock.InOrder( + 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) diff --git a/platform/extension/messagequeue/mysql/tenants.go b/platform/extension/messagequeue/mysql/tenants.go index f3f191d89..0eb1a5745 100644 --- a/platform/extension/messagequeue/mysql/tenants.go +++ b/platform/extension/messagequeue/mysql/tenants.go @@ -15,6 +15,7 @@ package mysql import ( + "fmt" "os" "strings" ) @@ -29,3 +30,17 @@ func ParseTenantsFromEnv() []string { } return out } + +// ParseRequiredTenantsFromEnv reads MQ_TENANTS and rejects an empty tenant list. +func ParseRequiredTenantsFromEnv() ([]string, error) { + tenants := ParseTenantsFromEnv() + if len(tenants) == 0 { + return nil, fmt.Errorf("MQ_TENANTS must contain at least one tenant") + } + for _, tenant := range tenants { + if err := validateASCIIIdentifier("tenant", tenant); err != nil { + return nil, fmt.Errorf("invalid MQ_TENANTS: %w", err) + } + } + return tenants, nil +} diff --git a/platform/extension/messagequeue/mysql/tenants_test.go b/platform/extension/messagequeue/mysql/tenants_test.go new file mode 100644 index 000000000..7a9c9ed5d --- /dev/null +++ b/platform/extension/messagequeue/mysql/tenants_test.go @@ -0,0 +1,72 @@ +// 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 ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseRequiredTenantsFromEnv(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: "empty", + wantErr: true, + }, + { + name: "whitespace and commas", + value: " , , ", + wantErr: true, + }, + { + name: "tenant exceeds byte limit", + value: strings.Repeat("x", maxIdentifierLength+1), + wantErr: true, + }, + { + name: "tenant contains non-ASCII characters", + value: "monorepo/café", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("MQ_TENANTS", tt.value) + + got, err := ParseRequiredTenantsFromEnv() + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/service/runway/server/config.go b/service/runway/server/config.go index 522faeaef..94b64a09e 100644 --- a/service/runway/server/config.go +++ b/service/runway/server/config.go @@ -44,14 +44,6 @@ type mergeConfig struct { Queues []namedQueueMergeConfig `yaml:"queues"` } -func tenantNamesFromMergeConfig(cfg mergeConfig) []string { - names := make([]string, 0, len(cfg.Queues)) - for _, q := range cfg.Queues { - names = append(names, q.Name) - } - return names -} - // namedQueueMergeConfig is one queue's entry. type namedQueueMergeConfig struct { // Name is the queue this entry configures. diff --git a/service/runway/server/docker-compose.yml b/service/runway/server/docker-compose.yml index 1706886dc..72a4e66a4 100644 --- a/service/runway/server/docker-compose.yml +++ b/service/runway/server/docker-compose.yml @@ -52,6 +52,7 @@ services: # Level for the queue's own logs; info by default so its per-message # chatter does not bury the rest of the service at debug. - QUEUE_LOG_LEVEL=${QUEUE_LOG_LEVEL:-} + - MQ_TENANTS=${MQ_TENANTS:-test-queue,e2e-runway/merge,e2e-runway/check,e2e-runway/failed,e2e-runway/dlq,e2e-runway/undecodable} - HOSTNAME=runway-dev depends_on: mysql-queue: diff --git a/service/runway/server/main.go b/service/runway/server/main.go index 947769ab1..9eb2fe0f5 100644 --- a/service/runway/server/main.go +++ b/service/runway/server/main.go @@ -135,9 +135,9 @@ func run() error { } defer queueDB.Close() - mergeCfg, err := loadMergeConfigFromEnv(logger) + tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() if err != nil { - return fmt.Errorf("failed to load merge config: %w", err) + return fmt.Errorf("failed to configure queue subscribers: %w", err) } mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ @@ -145,7 +145,7 @@ func run() error { Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), - Tenants: tenantNamesFromMergeConfig(mergeCfg), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) diff --git a/service/stovepipe/server/main.go b/service/stovepipe/server/main.go index ae6e0cf63..2de12fc0e 100644 --- a/service/stovepipe/server/main.go +++ b/service/stovepipe/server/main.go @@ -254,12 +254,17 @@ func run() error { } defer queueDB.Close() + tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() + 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: parseMQTenantsFromEnv(), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -341,6 +346,7 @@ func run() error { storageFty, materializer, registry, + tenants, ) srv := &StovepipeServer{ pingController: pingController, @@ -410,12 +416,6 @@ func run() error { return err } -// parseMQTenantsFromEnv reads MQ_TENANTS. Stovepipe accepts any queue at ingest -// but discovery requires configured tenants. -func parseMQTenantsFromEnv() []string { - return queueMySQL.ParseTenantsFromEnv() -} - // registerPrimaryControllers creates the primary-pipeline queue controllers and // registers them with c, returning how many were registered. func registerPrimaryControllers( diff --git a/service/submitqueue/docker-compose.yml b/service/submitqueue/docker-compose.yml index c7a11d77d..04917a705 100644 --- a/service/submitqueue/docker-compose.yml +++ b/service/submitqueue/docker-compose.yml @@ -107,6 +107,7 @@ services: # Level for the queue's own logs; info by default so its per-message # chatter does not bury the rest of the service at debug. - QUEUE_LOG_LEVEL=${QUEUE_LOG_LEVEL:-} + - MQ_TENANTS=${MQ_TENANTS:-test-queue,e2e-test-queue,e2e-cancel-queue,e2e-chain-queue,e2e-redelivery-queue,e2e-strand-queue,e2e-conflict-error-queue,e2e-git-queue,demo-queue,e2e-respeculate-queue} - HOSTNAME=orchestrator-dev # Consumer-gate state shared with the host (see header comment) - CONSUMER_GATE_DIR=/var/submitqueue/consumergate @@ -138,6 +139,7 @@ services: # Level for the queue's own logs; info by default so its per-message # chatter does not bury the rest of the service at debug. - QUEUE_LOG_LEVEL=${QUEUE_LOG_LEVEL:-} + - MQ_TENANTS=${MQ_TENANTS:-test-queue,e2e-test-queue,e2e-cancel-queue,e2e-chain-queue,e2e-redelivery-queue,e2e-strand-queue,e2e-conflict-error-queue,e2e-git-queue,demo-queue,e2e-respeculate-queue} - HOSTNAME=runway-dev # Consumer-gate state shared with the host (see header comment) - CONSUMER_GATE_DIR=/var/submitqueue/consumergate diff --git a/service/submitqueue/orchestrator/server/docker-compose.yml b/service/submitqueue/orchestrator/server/docker-compose.yml index 890478a11..8eb57ace5 100644 --- a/service/submitqueue/orchestrator/server/docker-compose.yml +++ b/service/submitqueue/orchestrator/server/docker-compose.yml @@ -65,6 +65,7 @@ services: # Level for the queue's own logs; info by default so its per-message # chatter does not bury the rest of the service at debug. - QUEUE_LOG_LEVEL=${QUEUE_LOG_LEVEL:-} + - MQ_TENANTS=${MQ_TENANTS:-test-queue} - HOSTNAME=orchestrator-dev depends_on: mysql-app: diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 25ecb1f0e..65c4bf86e 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -164,13 +164,17 @@ func run() error { if err != nil { return fmt.Errorf("failed to load extension profiles: %w", err) } + tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() + 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: tenantNamesFromProfiles(profilesCfg), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -342,14 +346,6 @@ func loadProfilesConfigFromEnv(logger *zap.Logger) (profilesConfig, error) { return cfg, nil } -func tenantNamesFromProfiles(cfg profilesConfig) []string { - names := make([]string, 0, len(cfg.Queues)) - for _, q := range cfg.Queues { - names = append(names, q.Name) - } - return names -} - // defaultProfilesConfig is the example topology used when no configuration file // is supplied: fake edge integrations everywhere, with the conflict analyzer // varied per queue to exercise its different behaviors. diff --git a/stovepipe/controller/ingest.go b/stovepipe/controller/ingest.go index 1c9022f0a..5011693ce 100644 --- a/stovepipe/controller/ingest.go +++ b/stovepipe/controller/ingest.go @@ -56,13 +56,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 +76,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 +113,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 { diff --git a/stovepipe/controller/ingest_test.go b/stovepipe/controller/ingest_test.go index f500110c6..05f4d8a43 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 } @@ -233,6 +233,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/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 2006711e7..0f54ef176 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -120,6 +120,42 @@ 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) +} + // 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 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..a259c0a98 --- /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 consumer_group = ? AND topic = ? AND partition_key = ?", + tenantB, + consumerGroup, + topic, + partitionKey, + ).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 2d41dbe2c..d900bc11d 100644 --- a/tool/linter/queueshard/main.go +++ b/tool/linter/queueshard/main.go @@ -138,7 +138,8 @@ func check(file, content string, shardColumns map[string]bool) (int, []violation for _, idx := range indexRe.FindAllStringSubmatch(body, -1) { idxColumns := splitColumns(idx[2]) - if len(idxColumns) == 1 && idxColumns[0] == "offset" { + 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 diff --git a/tool/linter/queueshard/main_test.go b/tool/linter/queueshard/main_test.go index d2ff4542d..551d93d2c 100644 --- a/tool/linter/queueshard/main_test.go +++ b/tool/linter/queueshard/main_test.go @@ -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, map[string]bool{"queue": true, "name": true}) + 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 6a6bf07de082ba190e04be04eb91029633dbbb65 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Mon, 7 Sep 2026 13:23:27 -0700 Subject: [PATCH 3/4] refactor(messagequeue): make publish tenant explicit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Tenant selects the physical MQ shard and is therefore a first-class publish input, not metadata to infer implicitly. Parsing MQ_TENANTS is service wiring behavior rather than a responsibility of the MySQL backend. ### What? - Require tenant explicitly in platform publish and hook APIs, propagate it to queue_name metadata, and reject conflicting metadata. - Pass each producer's authoritative queue through SubmitQueue, Runway, Stovepipe, hooks, and tests independently from partition keys. - Move required-tenant parsing into a shared service/messagequeue wiring package. ## Test Plan ✅ Focused unit tests across publish, hook, service wiring, Runway, Stovepipe, and SubmitQueue controllers ✅ `./tool/bazel build //...` ✅ `make fmt && make gazelle` Co-authored-by: Cursor --- doc/rfc/messagequeue-tenant-sharding.md | 2 +- .../extension/messagequeue/mysql/BUILD.bazel | 2 - .../extension/messagequeue/mysql/tenants.go | 46 ---------------- platform/hook/publisher.go | 9 ++-- platform/hook/publisher_test.go | 11 ++-- platform/publish/publish.go | 43 ++++++--------- platform/publish/publish_test.go | 41 +++++++------- runway/controller/dlq/dlq.go | 2 +- runway/controller/dlq/dlq_test.go | 1 + runway/controller/merge/merge.go | 2 +- runway/controller/merge/merge_test.go | 3 ++ .../mergeconflictcheck/mergeconflictcheck.go | 2 +- .../mergeconflictcheck_test.go | 3 ++ service/messagequeue/BUILD.bazel | 18 +++++++ service/messagequeue/tenant.go | 53 +++++++++++++++++++ .../messagequeue/tenant_test.go | 35 ++++-------- service/runway/server/BUILD.bazel | 1 + service/runway/server/main.go | 3 +- service/stovepipe/server/BUILD.bazel | 1 + service/stovepipe/server/main.go | 3 +- .../orchestrator/server/BUILD.bazel | 1 + .../submitqueue/orchestrator/server/main.go | 3 +- stovepipe/controller/BUILD.bazel | 1 - stovepipe/controller/build/build.go | 2 +- stovepipe/controller/build/build_test.go | 1 + .../controller/buildsignal/buildsignal.go | 2 +- .../buildsignal/buildsignal_test.go | 1 + stovepipe/controller/ingest.go | 4 +- stovepipe/controller/ingest_test.go | 1 + stovepipe/controller/process/process.go | 4 +- stovepipe/controller/process/process_test.go | 2 + stovepipe/controller/record/record.go | 2 +- stovepipe/controller/record/record_test.go | 3 ++ submitqueue/core/request/log.go | 2 +- submitqueue/core/request/log_test.go | 3 ++ submitqueue/core/request/terminate_test.go | 6 +-- submitqueue/gateway/controller/cancel.go | 2 +- submitqueue/gateway/controller/cancel_test.go | 1 + submitqueue/gateway/controller/land.go | 2 +- submitqueue/gateway/controller/land_test.go | 1 + .../orchestrator/controller/batch/batch.go | 2 +- .../orchestrator/controller/build/build.go | 2 +- .../controller/build/build_test.go | 1 + .../controller/buildsignal/buildsignal.go | 2 +- .../buildsignal/buildsignal_test.go | 1 + .../orchestrator/controller/cancel/cancel.go | 2 +- .../controller/conclude/conclude_test.go | 14 ++--- .../dependencyanalysis/dependencyanalysis.go | 2 +- .../orchestrator/controller/dlq/batch_test.go | 2 +- .../controller/dlq/buildsignal_test.go | 2 +- .../orchestrator/controller/dlq/dlq_test.go | 16 +++--- .../dlq/mergeconflictsignal_test.go | 2 +- .../controller/dlq/mergesignal_test.go | 2 +- .../controller/dlq/publisher_test.go | 1 + .../controller/dlq/request_test.go | 14 ++--- .../orchestrator/controller/dlq/speculate.go | 2 +- .../controller/dlq/speculate_test.go | 2 +- .../orchestrator/controller/merge/merge.go | 2 +- .../controller/merge/merge_test.go | 3 ++ .../mergeconflictsignal.go | 2 +- .../controller/mergesignal/mergesignal.go | 2 +- .../mergesignal/mergesignal_test.go | 3 +- .../controller/speculate/run_test.go | 1 + .../controller/speculate/speculate.go | 2 +- .../orchestrator/controller/start/start.go | 2 +- .../controller/validate/validate.go | 2 +- .../controller/validate/validate_test.go | 3 ++ test/e2e/submitqueue/harness_test.go | 2 +- 68 files changed, 227 insertions(+), 191 deletions(-) delete mode 100644 platform/extension/messagequeue/mysql/tenants.go create mode 100644 service/messagequeue/BUILD.bazel create mode 100644 service/messagequeue/tenant.go rename platform/extension/messagequeue/mysql/tenants_test.go => service/messagequeue/tenant_test.go (64%) diff --git a/doc/rfc/messagequeue-tenant-sharding.md b/doc/rfc/messagequeue-tenant-sharding.md index ee09c407f..8a9d5d445 100644 --- a/doc/rfc/messagequeue-tenant-sharding.md +++ b/doc/rfc/messagequeue-tenant-sharding.md @@ -76,7 +76,7 @@ Fair-share, orphan sweep, and idle-lease release run per `(tenant, topic)`, not ## Publish -`platform/publish` stamps `Message.Tenant` from context metadata (`queue_name`). Empty tenant on publish is rejected. `PartitionKey` is unchanged. +Every `platform/publish` call supplies tenant explicitly. The package stamps `Message.Tenant` and mirrors it into `queue_name` delivery metadata, rejecting empty tenants and conflicting caller metadata. `PartitionKey` is unchanged. ## Wiring diff --git a/platform/extension/messagequeue/mysql/BUILD.bazel b/platform/extension/messagequeue/mysql/BUILD.bazel index b62131192..fbc1f83b9 100644 --- a/platform/extension/messagequeue/mysql/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/BUILD.bazel @@ -16,7 +16,6 @@ go_library( "stores.go", "subscriber.go", "subscriber_heartbeat_store.go", - "tenants.go", ], importpath = "github.com/uber/submitqueue/platform/extension/messagequeue/mysql", visibility = ["//visibility:public"], @@ -43,7 +42,6 @@ go_test( "sql_test.go", "subscriber_heartbeat_store_test.go", "subscriber_test.go", - "tenants_test.go", ], embed = [":go_default_library"], deps = [ diff --git a/platform/extension/messagequeue/mysql/tenants.go b/platform/extension/messagequeue/mysql/tenants.go deleted file mode 100644 index 0eb1a5745..000000000 --- a/platform/extension/messagequeue/mysql/tenants.go +++ /dev/null @@ -1,46 +0,0 @@ -// Copyright (c) 2025 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" - "os" - "strings" -) - -// ParseTenantsFromEnv reads MQ_TENANTS env var (comma-separated). -func ParseTenantsFromEnv() []string { - var out []string - for _, part := range strings.Split(os.Getenv("MQ_TENANTS"), ",") { - if trimmed := strings.TrimSpace(part); trimmed != "" { - out = append(out, trimmed) - } - } - return out -} - -// ParseRequiredTenantsFromEnv reads MQ_TENANTS and rejects an empty tenant list. -func ParseRequiredTenantsFromEnv() ([]string, error) { - tenants := ParseTenantsFromEnv() - if len(tenants) == 0 { - return nil, fmt.Errorf("MQ_TENANTS must contain at least one tenant") - } - for _, tenant := range tenants { - if err := validateASCIIIdentifier("tenant", tenant); err != nil { - return nil, fmt.Errorf("invalid MQ_TENANTS: %w", err) - } - } - return tenants, nil -} diff --git a/platform/hook/publisher.go b/platform/hook/publisher.go index 07dc3b4bd..71a7a0624 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 { @@ -50,7 +51,7 @@ func Publish( } if err := publish.Message( - ctx, registry, basehook.TopicKeyHook, publish.IntentID(event.GetId()), body, partitionKey, + ctx, registry, basehook.TopicKeyHook, tenant, publish.IntentID(event.GetId()), body, partitionKey, ); err != nil { return fmt.Errorf("failed to publish hook event %s: %w", event.GetId(), err) } 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 1e0497bc4..87a0d7ac7 100644 --- a/platform/publish/publish.go +++ b/platform/publish/publish.go @@ -35,11 +35,11 @@ import ( "github.com/uber/submitqueue/platform/consumer" ) -// Message publishes payload to the topic registered for key. Allowlisted -// delivery context is propagated as message metadata. +// Message publishes payload to the topic registered for key. Tenant selects +// the shard and is propagated as the delivery's 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 +// 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. @@ -48,16 +48,18 @@ import ( // 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) +func Message(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, tenant, msgID string, payload []byte, partitionKey string) error { + return MessageWithMetadata(ctx, registry, key, tenant, 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 { +// Tenant is also propagated as the delivery's queue-name context. +func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, tenant, msgID string, payload []byte, partitionKey string, metadata map[string]string) error { + if 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,29 +69,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)) - if tenant := msg.Metadata[entityqueue.MetadataKeyQueueName]; tenant != "" { - msg.Tenant = tenant - } else if queueName, ok := entityqueue.QueueName(ctx); ok { - msg.Tenant = queueName + if queueName, exists := metadata[entityqueue.MetadataKeyQueueName]; exists && queueName != tenant { + return fmt.Errorf("queue-name metadata %q does not match tenant %q", queueName, tenant) } - 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 metadata == nil { metadata = make(map[string]string) } - metadata[entityqueue.MetadataKeyQueueName] = queueName - return metadata + metadata[entityqueue.MetadataKeyQueueName] = tenant + + msg := entityqueue.NewMessage(msgID, payload, partitionKey, metadata) + msg.Tenant = 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 adfce441d..61feec51f 100644 --- a/platform/publish/publish_test.go +++ b/platform/publish/publish_test.go @@ -55,15 +55,16 @@ 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, "tenant-1", "msg-1", []byte("payload"), "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,8 +76,7 @@ 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, "monorepo/main", "msg-1", []byte("payload"), "partition-1")) assert.Equal(t, "monorepo/main", published.Metadata[entityqueue.MetadataKeyQueueName]) assert.Equal(t, "monorepo/main", published.Tenant) } @@ -94,8 +94,7 @@ 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, MessageWithMetadata(context.Background(), registry, testKey, "monorepo/main", "msg-1", []byte("payload"), "partition-1", metadata)) assert.Equal(t, map[string]string{ "failure_reason": "build failed", entityqueue.MetadataKeyQueueName: "monorepo/main", @@ -104,30 +103,28 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { assert.Equal(t, map[string]string{"failure_reason": "build failed"}, metadata) } -func TestMessageWithMetadata_ExplicitQueueNameWins(t *testing.T) { +func TestMessageWithMetadata_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]) - assert.Equal(t, "outbound", published.Tenant) + err := MessageWithMetadata(context.Background(), registry, testKey, "inbound", "msg-1", []byte("payload"), "partition-1", 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", "tenant-1", "msg-1", []byte("payload"), "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, "", "msg-1", []byte("payload"), "partition-1") require.Error(t, err) } diff --git a/runway/controller/dlq/dlq.go b/runway/controller/dlq/dlq.go index d53f25c74..f88ce1339 100644 --- a/runway/controller/dlq/dlq.go +++ b/runway/controller/dlq/dlq.go @@ -165,7 +165,7 @@ func (c *Controller) publish(ctx context.Context, result *runwaymq.MergeResult, } if err := publish.Message(ctx, c.registry, c.signalTopicKey, - publish.IntentID(result.GetId(), "dlq"), payload, partitionKey); err != nil { + result.GetQueueName(), publish.IntentID(result.GetId(), "dlq"), payload, 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..b5cb7df7c 100644 --- a/runway/controller/dlq/dlq_test.go +++ b/runway/controller/dlq/dlq_test.go @@ -112,6 +112,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)) diff --git a/runway/controller/merge/merge.go b/runway/controller/merge/merge.go index 8a98cb3d4..4a8e63633 100644 --- a/runway/controller/merge/merge.go +++ b/runway/controller/merge/merge.go @@ -153,7 +153,7 @@ 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, result.GetQueueName(), publish.IntentID(result.GetId()), payload, 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..4aa915c6c 100644 --- a/runway/controller/merge/merge_test.go +++ b/runway/controller/merge/merge_test.go @@ -118,11 +118,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 +147,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) diff --git a/runway/controller/mergeconflictcheck/mergeconflictcheck.go b/runway/controller/mergeconflictcheck/mergeconflictcheck.go index 202f879b7..c58f2c5b2 100644 --- a/runway/controller/mergeconflictcheck/mergeconflictcheck.go +++ b/runway/controller/mergeconflictcheck/mergeconflictcheck.go @@ -151,7 +151,7 @@ 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, result.GetQueueName(), publish.IntentID(result.GetId()), payload, 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..b1269a90a 100644 --- a/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go +++ b/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go @@ -118,11 +118,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 +147,7 @@ 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) 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..bd45dbcb3 --- /dev/null +++ b/service/messagequeue/tenant.go @@ -0,0 +1,53 @@ +// 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" + "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 _, 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 +} diff --git a/platform/extension/messagequeue/mysql/tenants_test.go b/service/messagequeue/tenant_test.go similarity index 64% rename from platform/extension/messagequeue/mysql/tenants_test.go rename to service/messagequeue/tenant_test.go index 7a9c9ed5d..950a8b80b 100644 --- a/platform/extension/messagequeue/mysql/tenants_test.go +++ b/service/messagequeue/tenant_test.go @@ -4,7 +4,7 @@ // 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 +// 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, @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mysql +package messagequeue import ( "strings" @@ -22,7 +22,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestParseRequiredTenantsFromEnv(t *testing.T) { +func TestParseRequiredTenants(t *testing.T) { tests := []struct { name string value string @@ -35,32 +35,19 @@ func TestParseRequiredTenantsFromEnv(t *testing.T) { 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", maxIdentifierLength+1), - wantErr: true, - }, - { - name: "tenant contains non-ASCII characters", - value: "monorepo/café", - wantErr: true, + 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}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - t.Setenv("MQ_TENANTS", tt.value) - - got, err := ParseRequiredTenantsFromEnv() - + got, err := ParseRequiredTenants(tt.value) if tt.wantErr { require.Error(t, err) return diff --git a/service/runway/server/BUILD.bazel b/service/runway/server/BUILD.bazel index 113cd6e23..49d483c0c 100644 --- a/service/runway/server/BUILD.bazel +++ b/service/runway/server/BUILD.bazel @@ -37,6 +37,7 @@ go_library( "//runway/extension/merger/fake:go_default_library", "//runway/extension/merger/git:go_default_library", "//runway/extension/merger/noop:go_default_library", + "//service/messagequeue:go_default_library", "@com_github_go_sql_driver_mysql//:go_default_library", "@com_github_uber_go_tally//:go_default_library", "@in_gopkg_yaml_v3//:go_default_library", diff --git a/service/runway/server/main.go b/service/runway/server/main.go index 9eb2fe0f5..c50613a67 100644 --- a/service/runway/server/main.go +++ b/service/runway/server/main.go @@ -51,6 +51,7 @@ import ( "github.com/uber/submitqueue/runway/extension/merger/fake" gitmerger "github.com/uber/submitqueue/runway/extension/merger/git" "github.com/uber/submitqueue/runway/extension/merger/noop" + servicemq "github.com/uber/submitqueue/service/messagequeue" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/reflection" @@ -135,7 +136,7 @@ func run() error { } defer queueDB.Close() - tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) if err != nil { return fmt.Errorf("failed to configure queue subscribers: %w", 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 2de12fc0e..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,7 +255,7 @@ func run() error { } defer queueDB.Close() - tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) if err != nil { return fmt.Errorf("failed to configure queue subscribers: %w", err) } diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 33d9b694b..3ced8b0c7 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -33,6 +33,7 @@ go_library( "//platform/githubactions:go_default_library", "//platform/http:go_default_library", "//platform/pipeline:go_default_library", + "//service/messagequeue:go_default_library", "//submitqueue/core/changeset:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/buildrunner:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 65c4bf86e..f54749563 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -42,6 +42,7 @@ import ( hooknoop "github.com/uber/submitqueue/platform/extension/hook/noop" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" "github.com/uber/submitqueue/platform/pipeline" + servicemq "github.com/uber/submitqueue/service/messagequeue" "github.com/uber/submitqueue/submitqueue/core/changeset" "github.com/uber/submitqueue/submitqueue/extension/storage" mysqlstorage "github.com/uber/submitqueue/submitqueue/extension/storage/mysql" @@ -164,7 +165,7 @@ func run() error { if err != nil { return fmt.Errorf("failed to load extension profiles: %w", err) } - tenants, err := queueMySQL.ParseRequiredTenantsFromEnv() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) if err != nil { return fmt.Errorf("failed to configure queue subscribers: %w", err) } 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.go b/stovepipe/controller/build/build.go index 0cdebdc17..71d27e2f1 100644 --- a/stovepipe/controller/build/build.go +++ b/stovepipe/controller/build/build.go @@ -174,7 +174,7 @@ 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, queue, publish.IntentID(buildID), payload, 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..b458997fb 100644 --- a/stovepipe/controller/build/build_test.go +++ b/stovepipe/controller/build/build_test.go @@ -132,6 +132,7 @@ 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]) } diff --git a/stovepipe/controller/buildsignal/buildsignal.go b/stovepipe/controller/buildsignal/buildsignal.go index f8aca8bfd..92b08b212 100644 --- a/stovepipe/controller/buildsignal/buildsignal.go +++ b/stovepipe/controller/buildsignal/buildsignal.go @@ -358,7 +358,7 @@ 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, queue, publish.IntentID(requestID), payload, 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..e0bec0fcb 100644 --- a/stovepipe/controller/buildsignal/buildsignal_test.go +++ b/stovepipe/controller/buildsignal/buildsignal_test.go @@ -456,6 +456,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/ingest.go b/stovepipe/controller/ingest.go index 5011693ce..13d7ee96d 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" @@ -322,8 +321,7 @@ 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, queue, publish.IntentID(id), payload, 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 05f4d8a43..74492d19d 100644 --- a/stovepipe/controller/ingest_test.go +++ b/stovepipe/controller/ingest_test.go @@ -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]) } diff --git a/stovepipe/controller/process/process.go b/stovepipe/controller/process/process.go index 75e264784..d2624a062 100644 --- a/stovepipe/controller/process/process.go +++ b/stovepipe/controller/process/process.go @@ -483,7 +483,7 @@ 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, queue, publish.IntentID(id), payload, id); err != nil { return fmt.Errorf("failed to publish build request: %w", err) } return nil @@ -498,7 +498,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..99843cb9e 100644 --- a/stovepipe/controller/process/process_test.go +++ b/stovepipe/controller/process/process_test.go @@ -183,6 +183,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 +203,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/record.go b/stovepipe/controller/record/record.go index 93366aec7..b4486520a 100644 --- a/stovepipe/controller/record/record.go +++ b/stovepipe/controller/record/record.go @@ -489,7 +489,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..bcd8ef2f8 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() @@ -287,6 +289,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) diff --git a/submitqueue/core/request/log.go b/submitqueue/core/request/log.go index 598d9a1c7..4e04ffded 100644 --- a/submitqueue/core/request/log.go +++ b/submitqueue/core/request/log.go @@ -53,7 +53,7 @@ func PublishLog(ctx context.Context, registry consumer.TopicRegistry, logEntry e } if err := publish.Message(ctx, registry, topickey.TopicKeyLog, - publish.IntentID(logEntry.RequestID, cause...), payload, partitionKey); err != nil { + logEntry.Queue, publish.IntentID(logEntry.RequestID, cause...), payload, 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..5dad0c909 100644 --- a/submitqueue/gateway/controller/cancel.go +++ b/submitqueue/gateway/controller/cancel.go @@ -140,7 +140,7 @@ func (c *cancelController) publishToQueue(ctx context.Context, cancelRequest ent // 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 { + cancelRequest.Queue, publish.IntentID(cancelRequest.ID), payload, 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..d8b1ba287 100644 --- a/submitqueue/gateway/controller/land.go +++ b/submitqueue/gateway/controller/land.go @@ -213,7 +213,7 @@ func (c *landController) publishToQueue(ctx context.Context, landRequest entity. // - 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 { + landRequest.Queue, publish.IntentID(landRequest.ID), payload, 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/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 0b60b32d6..deb49e4aa 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -212,7 +212,7 @@ func (c *Controller) publishToDependencyAnalysis(ctx context.Context, batch enti } if err := publish.Message(ctx, c.registry, topickey.TopicKeyDependencyAnalysis, - publish.IntentID(batch.ID), payload, batch.Queue); err != nil { + batch.Queue, publish.IntentID(batch.ID), payload, batch.Queue); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/build/build.go b/submitqueue/orchestrator/controller/build/build.go index afa80cc99..ca422ab69 100644 --- a/submitqueue/orchestrator/controller/build/build.go +++ b/submitqueue/orchestrator/controller/build/build.go @@ -381,7 +381,7 @@ 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, queue, publish.IntentID(buildID), payload, 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..85ee83048 100644 --- a/submitqueue/orchestrator/controller/build/build_test.go +++ b/submitqueue/orchestrator/controller/build/build_test.go @@ -175,6 +175,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/buildsignal.go b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go index 9c28b8693..7ca07709d 100644 --- a/submitqueue/orchestrator/controller/buildsignal/buildsignal.go +++ b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go @@ -440,7 +440,7 @@ 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, queue, msgID, payload, 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..e15b35a46 100644 --- a/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go +++ b/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go @@ -253,6 +253,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 }, ) diff --git a/submitqueue/orchestrator/controller/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 243e3af32..119079885 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -350,7 +350,7 @@ 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, queue, publish.UniqueID(batchID), payload, queue) } // Name returns the controller name for logging and metrics. diff --git a/submitqueue/orchestrator/controller/conclude/conclude_test.go b/submitqueue/orchestrator/controller/conclude/conclude_test.go index f39910da0..b884a0d54 100644 --- a/submitqueue/orchestrator/controller/conclude/conclude_test.go +++ b/submitqueue/orchestrator/controller/conclude/conclude_test.go @@ -126,12 +126,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 +164,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 +197,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 +232,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 +266,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 +385,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) diff --git a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go index cb36da1d8..53d9a08fb 100644 --- a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go +++ b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go @@ -440,7 +440,7 @@ 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, batch.Queue, publish.IntentID(batch.ID), payload, 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/dlq/batch_test.go b/submitqueue/orchestrator/controller/dlq/batch_test.go index a125271df..78bcb116c 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) diff --git a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go index acbd3009b..d5da09dfe 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) 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_test.go b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go index bf23a1d6a..4ecf28409 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) diff --git a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go index 69a122ae4..927f111df 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) 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_test.go b/submitqueue/orchestrator/controller/dlq/request_test.go index adc34d064..70525e871 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) @@ -77,7 +77,7 @@ func TestDLQRequestController_Process_CancelRequestPayload(t *testing.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) @@ -105,7 +105,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) @@ -134,7 +134,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) @@ -252,7 +252,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 +296,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 +338,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..d8b62ed10 100644 --- a/submitqueue/orchestrator/controller/dlq/speculate.go +++ b/submitqueue/orchestrator/controller/dlq/speculate.go @@ -203,7 +203,7 @@ 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, queue, publish.UniqueID(next), payload, 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..a82473fc0 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) diff --git a/submitqueue/orchestrator/controller/merge/merge.go b/submitqueue/orchestrator/controller/merge/merge.go index a9657457b..3005a8901 100644 --- a/submitqueue/orchestrator/controller/merge/merge.go +++ b/submitqueue/orchestrator/controller/merge/merge.go @@ -229,7 +229,7 @@ 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, req.GetQueueName(), publish.IntentID(req.GetId()), payload, 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..87b477bef 100644 --- a/submitqueue/orchestrator/controller/merge/merge_test.go +++ b/submitqueue/orchestrator/controller/merge/merge_test.go @@ -164,6 +164,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 +173,7 @@ func TestProcess_PublishesFullPayloadToRunway(t *testing.T) { } gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ).AnyTimes() @@ -188,6 +190,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/mergeconflictsignal.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go index dbfb69b5c..645dcf8e7 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go @@ -207,7 +207,7 @@ 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, queue, publish.IntentID(requestID), payload, queue); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go index 7611dd837..4fef98c9f 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go @@ -207,7 +207,7 @@ 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.MessageWithMetadata(ctx, c.registry, key, queue, msgID, payload, queue, 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..53905d4c0 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go @@ -78,7 +78,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 }, 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..eec340ec1 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -235,7 +235,7 @@ 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.MessageWithMetadata(ctx, c.registry, key, queue, msgID, payload, partitionKey, metadata) } // attributed records what a failure was about and counts it by subject type. diff --git a/submitqueue/orchestrator/controller/start/start.go b/submitqueue/orchestrator/controller/start/start.go index 5dbfbd04f..b80e4fae4 100644 --- a/submitqueue/orchestrator/controller/start/start.go +++ b/submitqueue/orchestrator/controller/start/start.go @@ -149,7 +149,7 @@ 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, queue, publish.IntentID(requestID), payload, queue); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/validate/validate.go b/submitqueue/orchestrator/controller/validate/validate.go index 9eacab3cf..62420ef7a 100644 --- a/submitqueue/orchestrator/controller/validate/validate.go +++ b/submitqueue/orchestrator/controller/validate/validate.go @@ -321,7 +321,7 @@ 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, req.GetQueueName(), publish.IntentID(req.GetId()), payload, 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..740737ba5 100644 --- a/submitqueue/orchestrator/controller/validate/validate_test.go +++ b/submitqueue/orchestrator/controller/validate/validate_test.go @@ -201,6 +201,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 +210,7 @@ func TestController_Process_PublishesCheckToRunway(t *testing.T) { } gotTopic = topic gotPayload = msg.Payload + gotTenant = msg.Tenant return nil }, ).AnyTimes() @@ -235,6 +237,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) diff --git a/test/e2e/submitqueue/harness_test.go b/test/e2e/submitqueue/harness_test.go index 08a49bbb3..4a5ba652e 100644 --- a/test/e2e/submitqueue/harness_test.go +++ b/test/e2e/submitqueue/harness_test.go @@ -301,7 +301,7 @@ func (s *E2EIntegrationSuite) redeliverBatchMessage(req request) { require.NoError(t, err) require.NoError(t, publish.Message(entityqueue.WithQueueName(s.ctx, req.queue), registry, topickey.TopicKeyBatch, - publish.UniqueID(req.sqid), payload, req.queue), "failed to redeliver the batch message") + req.queue, publish.UniqueID(req.sqid), payload, req.queue), "failed to redeliver the batch message") s.log.Logf("Redelivered the batch message for %s", req.sqid) } From cfce51e81b3d9661fa1426f27caf4fcf5a8de076 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Mon, 7 Sep 2026 13:32:12 -0700 Subject: [PATCH 4/4] refactor(messagequeue): consolidate publish arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Adding tenant to the positional publish API made future message attributes likely to trigger another repository-wide signature migration and left metadata as a separate method variant. ### What? Replace Message and MessageWithMetadata with one Message function taking documented MessageParams. Named fields keep tenant, message identity, payload, partitioning, and optional metadata explicit at every call site while allowing additive evolution. ## Test Plan ✅ Platform publish and hook tests ✅ 28 focused controller and wiring test suites ✅ `./tool/bazel build //...` ✅ `make fmt && make gazelle` Co-authored-by: Cursor --- platform/hook/publisher.go | 9 ++-- platform/publish/publish.go | 46 ++++++++++-------- platform/publish/publish_test.go | 47 +++++++++++++++---- runway/controller/dlq/dlq.go | 8 +++- runway/controller/merge/merge.go | 7 ++- .../mergeconflictcheck/mergeconflictcheck.go | 7 ++- stovepipe/controller/build/build.go | 7 ++- .../controller/buildsignal/buildsignal.go | 7 ++- stovepipe/controller/ingest.go | 7 ++- stovepipe/controller/process/process.go | 7 ++- submitqueue/core/request/log.go | 8 +++- submitqueue/gateway/controller/cancel.go | 8 +++- submitqueue/gateway/controller/land.go | 8 +++- .../orchestrator/controller/batch/batch.go | 8 +++- .../orchestrator/controller/build/build.go | 7 ++- .../controller/buildsignal/buildsignal.go | 7 ++- .../orchestrator/controller/cancel/cancel.go | 7 ++- .../dependencyanalysis/dependencyanalysis.go | 7 ++- .../orchestrator/controller/dlq/speculate.go | 7 ++- .../orchestrator/controller/merge/merge.go | 7 ++- .../mergeconflictsignal.go | 7 ++- .../controller/mergesignal/mergesignal.go | 8 +++- .../controller/speculate/speculate.go | 8 +++- .../orchestrator/controller/start/start.go | 7 ++- .../controller/validate/validate.go | 7 ++- test/e2e/submitqueue/harness_test.go | 8 +++- 26 files changed, 211 insertions(+), 60 deletions(-) diff --git a/platform/hook/publisher.go b/platform/hook/publisher.go index 71a7a0624..0048fada7 100644 --- a/platform/hook/publisher.go +++ b/platform/hook/publisher.go @@ -50,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, tenant, 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/publish/publish.go b/platform/publish/publish.go index 87a0d7ac7..6fd749225 100644 --- a/platform/publish/publish.go +++ b/platform/publish/publish.go @@ -35,29 +35,35 @@ import ( "github.com/uber/submitqueue/platform/consumer" ) -// Message publishes payload to the topic registered for key. Tenant selects -// the shard and is propagated as the delivery's queue-name 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. +// 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, tenant, msgID string, payload []byte, partitionKey string) error { - return MessageWithMetadata(ctx, registry, key, tenant, 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. -// Tenant is also propagated as the delivery's queue-name context. -func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, tenant, msgID string, payload []byte, partitionKey string, metadata map[string]string) error { - if tenant == "" { +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) @@ -69,17 +75,17 @@ func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, k return fmt.Errorf("no topic name registered for topic key %s", key) } - if queueName, exists := metadata[entityqueue.MetadataKeyQueueName]; exists && queueName != tenant { - return fmt.Errorf("queue-name metadata %q does not match tenant %q", queueName, tenant) + 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(metadata) + metadata := maps.Clone(params.Metadata) if metadata == nil { metadata = make(map[string]string) } - metadata[entityqueue.MetadataKeyQueueName] = tenant + metadata[entityqueue.MetadataKeyQueueName] = params.Tenant - msg := entityqueue.NewMessage(msgID, payload, partitionKey, metadata) - msg.Tenant = tenant + msg := entityqueue.NewMessage(params.ID, params.Payload, params.PartitionKey, metadata) + msg.Tenant = params.Tenant return q.Publisher().Publish(ctx, topicName, msg) } diff --git a/platform/publish/publish_test.go b/platform/publish/publish_test.go index 61feec51f..0e6583131 100644 --- a/platform/publish/publish_test.go +++ b/platform/publish/publish_test.go @@ -55,7 +55,12 @@ func TestMessage(t *testing.T) { return nil }) - err := Message(context.Background(), registry, testKey, "tenant-1", "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) @@ -76,12 +81,17 @@ func TestMessage_PropagatesTenantAsQueueName(t *testing.T) { return nil }) - require.NoError(t, Message(context.Background(), registry, testKey, "monorepo/main", "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) @@ -94,7 +104,13 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { }) metadata := map[string]string{"failure_reason": "build failed"} - require.NoError(t, MessageWithMetadata(context.Background(), registry, testKey, "monorepo/main", "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", @@ -103,12 +119,18 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { assert.Equal(t, map[string]string{"failure_reason": "build failed"}, metadata) } -func TestMessageWithMetadata_RejectsQueueNameDifferentFromTenant(t *testing.T) { +func TestMessage_RejectsQueueNameDifferentFromTenant(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := newTestRegistry(t, ctrl) metadata := map[string]string{entityqueue.MetadataKeyQueueName: "outbound"} - err := MessageWithMetadata(context.Background(), registry, testKey, "inbound", "msg-1", []byte("payload"), "partition-1", metadata) + err := Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "inbound", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + Metadata: metadata, + }) require.Error(t, err) } @@ -116,7 +138,12 @@ func TestMessage_UnregisteredKey(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := newTestRegistry(t, ctrl) - err := Message(context.Background(), registry, "unregistered-key", "tenant-1", "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) } @@ -124,7 +151,11 @@ func TestMessage_RequiresTenant(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := newTestRegistry(t, ctrl) - err := Message(context.Background(), registry, testKey, "", "msg-1", []byte("payload"), "partition-1") + 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/dlq.go b/runway/controller/dlq/dlq.go index f88ce1339..a922409d5 100644 --- a/runway/controller/dlq/dlq.go +++ b/runway/controller/dlq/dlq.go @@ -164,8 +164,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, - result.GetQueueName(), 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/merge/merge.go b/runway/controller/merge/merge.go index 4a8e63633..cb44930ec 100644 --- a/runway/controller/merge/merge.go +++ b/runway/controller/merge/merge.go @@ -153,7 +153,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, result.GetQueueName(), 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.go b/runway/controller/mergeconflictcheck/mergeconflictcheck.go index c58f2c5b2..56affff06 100644 --- a/runway/controller/mergeconflictcheck/mergeconflictcheck.go +++ b/runway/controller/mergeconflictcheck/mergeconflictcheck.go @@ -151,7 +151,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, result.GetQueueName(), 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/stovepipe/controller/build/build.go b/stovepipe/controller/build/build.go index 71d27e2f1..bed49654e 100644 --- a/stovepipe/controller/build/build.go +++ b/stovepipe/controller/build/build.go @@ -174,7 +174,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, queue, 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/buildsignal/buildsignal.go b/stovepipe/controller/buildsignal/buildsignal.go index 92b08b212..05069634c 100644 --- a/stovepipe/controller/buildsignal/buildsignal.go +++ b/stovepipe/controller/buildsignal/buildsignal.go @@ -358,7 +358,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, queue, 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/ingest.go b/stovepipe/controller/ingest.go index 13d7ee96d..2ffd6cb81 100644 --- a/stovepipe/controller/ingest.go +++ b/stovepipe/controller/ingest.go @@ -321,7 +321,12 @@ func (c *IngestController) publishProcess(ctx context.Context, id, queue string) return fmt.Errorf("failed to serialize process request: %w", err) } - if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyProcess, queue, 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/process/process.go b/stovepipe/controller/process/process.go index d2624a062..267764bb6 100644 --- a/stovepipe/controller/process/process.go +++ b/stovepipe/controller/process/process.go @@ -483,7 +483,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, queue, 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 diff --git a/submitqueue/core/request/log.go b/submitqueue/core/request/log.go index 4e04ffded..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, - logEntry.Queue, 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/gateway/controller/cancel.go b/submitqueue/gateway/controller/cancel.go index 5dad0c909..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, - cancelRequest.Queue, 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/land.go b/submitqueue/gateway/controller/land.go index d8b1ba287..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, - landRequest.Queue, 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/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index deb49e4aa..ca4cec263 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -211,8 +211,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, - batch.Queue, 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/build/build.go b/submitqueue/orchestrator/controller/build/build.go index ca422ab69..7897e04a5 100644 --- a/submitqueue/orchestrator/controller/build/build.go +++ b/submitqueue/orchestrator/controller/build/build.go @@ -381,7 +381,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, queue, 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/buildsignal/buildsignal.go b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go index 7ca07709d..60f6ece30 100644 --- a/submitqueue/orchestrator/controller/buildsignal/buildsignal.go +++ b/submitqueue/orchestrator/controller/buildsignal/buildsignal.go @@ -440,7 +440,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, queue, 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/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 119079885..28ba5d3d9 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -350,7 +350,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, queue, 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/dependencyanalysis/dependencyanalysis.go b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go index 53d9a08fb..60496c69e 100644 --- a/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go +++ b/submitqueue/orchestrator/controller/dependencyanalysis/dependencyanalysis.go @@ -440,7 +440,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, batch.Queue, 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/dlq/speculate.go b/submitqueue/orchestrator/controller/dlq/speculate.go index d8b62ed10..63853c82c 100644 --- a/submitqueue/orchestrator/controller/dlq/speculate.go +++ b/submitqueue/orchestrator/controller/dlq/speculate.go @@ -203,7 +203,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, queue, 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/merge/merge.go b/submitqueue/orchestrator/controller/merge/merge.go index 3005a8901..b15ac6574 100644 --- a/submitqueue/orchestrator/controller/merge/merge.go +++ b/submitqueue/orchestrator/controller/merge/merge.go @@ -229,7 +229,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, req.GetQueueName(), 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/mergeconflictsignal/mergeconflictsignal.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go index 645dcf8e7..a00a6ca7b 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go @@ -207,7 +207,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, queue, 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/mergesignal/mergesignal.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go index 4fef98c9f..8d15c1023 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go @@ -207,7 +207,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, queue, 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/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index eec340ec1..d97ea65ae 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -235,7 +235,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, queue, 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/start/start.go b/submitqueue/orchestrator/controller/start/start.go index b80e4fae4..61df2f290 100644 --- a/submitqueue/orchestrator/controller/start/start.go +++ b/submitqueue/orchestrator/controller/start/start.go @@ -149,7 +149,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, queue, 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/validate/validate.go b/submitqueue/orchestrator/controller/validate/validate.go index 62420ef7a..a6f1208a6 100644 --- a/submitqueue/orchestrator/controller/validate/validate.go +++ b/submitqueue/orchestrator/controller/validate/validate.go @@ -321,7 +321,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, req.GetQueueName(), 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/test/e2e/submitqueue/harness_test.go b/test/e2e/submitqueue/harness_test.go index 4a5ba652e..cd92821cf 100644 --- a/test/e2e/submitqueue/harness_test.go +++ b/test/e2e/submitqueue/harness_test.go @@ -300,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(entityqueue.WithQueueName(s.ctx, req.queue), registry, topickey.TopicKeyBatch, - req.queue, 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) }