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..8a9d5d445 --- /dev/null +++ b/doc/rfc/messagequeue-tenant-sharding.md @@ -0,0 +1,92 @@ +# 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`. 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)` +- 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` + +- 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 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 +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. Discovery and shutdown attempt every configured tenant and aggregate errors so one unavailable shard does not block unrelated tenants. + +## Publish + +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 + +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 + +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..fbc1f83b9 100644 --- a/platform/extension/messagequeue/mysql/BUILD.bazel +++ b/platform/extension/messagequeue/mysql/BUILD.bazel @@ -6,6 +6,7 @@ go_library( "constants.go", "delivery_state_store.go", "errors.go", + "identifier.go", "message_store.go", "mock_stores.go", "offset_store.go", diff --git a/platform/extension/messagequeue/mysql/constants.go b/platform/extension/messagequeue/mysql/constants.go index 1d5c7fc9a..ca9da3adc 100644 --- a/platform/extension/messagequeue/mysql/constants.go +++ b/platform/extension/messagequeue/mysql/constants.go @@ -18,6 +18,7 @@ package mysql const ( // Common log field names (used extensively across all stores) + logTenant = "tenant" logTopic = "topic" logPartitionKey = "partition_key" logMessageID = "message_id" diff --git a/platform/extension/messagequeue/mysql/ctl/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/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/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..d7b4251f6 100644 --- a/platform/extension/messagequeue/mysql/publisher.go +++ b/platform/extension/messagequeue/mysql/publisher.go @@ -57,7 +57,33 @@ 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") + } + for _, identifier := range []struct { + name string + value string + }{ + {name: "tenant", value: message.Tenant}, + {name: "topic", value: topic}, + } { + if err := validateASCIIIdentifier(identifier.name, identifier.value); err != nil { + return fmt.Errorf("publish: %w", err) + } + } + for _, identifier := range []struct { + name string + value string + }{ + {name: "message ID", value: message.ID}, + {name: "partition key", value: message.PartitionKey}, + } { + if err := validateTextIdentifier(identifier.name, identifier.value); err != nil { + return fmt.Errorf("publish: %w", err) + } + } + + if err := p.messageStore.Insert(ctx, message.Tenant, topic, []entityqueue.Message{message}); err != nil { return fmt.Errorf("publish message store insert error: %w", err) } diff --git a/platform/extension/messagequeue/mysql/publisher_test.go b/platform/extension/messagequeue/mysql/publisher_test.go index 115c72313..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 @@ -54,24 +57,24 @@ func TestPublisher_Publish(t *testing.T) { name: "publish single message", topic: "test_topic", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("payload1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "test_topic", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "test_topic", gomock.Any()).Return(nil).Times(1) }, }, { name: "publish multiple messages", topic: "multi_topic", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, - {ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, - {ID: "msg3", Payload: []byte("p3"), PartitionKey: "part2", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("p1"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg2", Payload: []byte("p2"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg3", Payload: []byte("p3"), PartitionKey: "part2", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "multi_topic", gomock.Any()).Return(nil).Times(3) + m.EXPECT().Insert(gomock.Any(), testTenant, "multi_topic", gomock.Any()).Return(nil).Times(3) }, }, { @@ -88,6 +91,7 @@ func TestPublisher_Publish(t *testing.T) { topic: "metadata_topic", messages: []entityqueue.Message{ { + Tenant: testTenant, ID: "msg_meta", Payload: []byte("payload"), PartitionKey: "part1", @@ -97,20 +101,70 @@ func TestPublisher_Publish(t *testing.T) { }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "metadata_topic", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "metadata_topic", gomock.Any()).Return(nil).Times(1) }, }, { name: "publish with valid topic name - hyphens", topic: "topic-with-dash", messages: []entityqueue.Message{ - {ID: "msg1", Payload: []byte("p"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, + {Tenant: testTenant, ID: "msg1", Payload: []byte("p"), PartitionKey: "part1", PublishedAt: fixedTimestamp}, }, wantErr: false, setupMock: func(m *MockmessageStore) { - m.EXPECT().Insert(gomock.Any(), "topic-with-dash", gomock.Any()).Return(nil).Times(1) + m.EXPECT().Insert(gomock.Any(), testTenant, "topic-with-dash", gomock.Any()).Return(nil).Times(1) }, }, + { + name: "rejects overlong tenant", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: overlong, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong topic", + topic: overlong, + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong message ID", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: overlong, PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects overlong partition key", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: overlong}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "accepts UTF-8 message ID and partition key", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg-é", PartitionKey: strings.Repeat("é", maxIdentifierLength)}}, + setupMock: func(m *MockmessageStore) { + m.EXPECT().Insert(gomock.Any(), testTenant, "test_topic", gomock.Any()).Return(nil) + }, + }, + { + name: "rejects non-ASCII tenant", + topic: "test_topic", + messages: []entityqueue.Message{{Tenant: "tenant-é", ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, + { + name: "rejects non-ASCII topic", + topic: "topic-é", + messages: []entityqueue.Message{{Tenant: testTenant, ID: "msg1", PartitionKey: "part1"}}, + wantErr: true, + setupMock: noStoreCall, + }, } for _, tt := range tests { @@ -181,7 +235,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 +244,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 +265,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 +278,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 +301,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 +310,7 @@ func TestPublisher_PublishContextCancellation(t *testing.T) { cancel() msg := entityqueue.NewMessage("msg1", []byte("payload"), "part1", nil) + msg.Tenant = testTenant // Should fail with context cancelled error err := pub.Publish(ctx, "test_topic", msg) diff --git a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql index 27625cdc8..e561aef2c 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_delivery_state.sql @@ -8,14 +8,17 @@ -- acked = FALSE, invisible_until <= now → ready for (re-)delivery CREATE TABLE IF NOT EXISTS queue_delivery_state ( + -- tenant is the shard isolation identity + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, + -- Consumer group this delivery state belongs to - consumer_group VARCHAR(255) NOT NULL, + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic of the message - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition key of the message - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Offset of the message in the immutable log message_offset BIGINT UNSIGNED NOT NULL, @@ -24,18 +27,13 @@ CREATE TABLE IF NOT EXISTS queue_delivery_state ( acked BOOLEAN NOT NULL DEFAULT FALSE, -- Visibility timeout (epoch milliseconds) - -- Only meaningful when acked = FALSE. - -- Future timestamp = in-flight or nack delay, 0/past = ready for delivery. invisible_until BIGINT UNSIGNED NOT NULL DEFAULT 0, -- Number of times this message has been redelivered to this consumer group retry_count INT UNSIGNED NOT NULL DEFAULT 0, -- Whether the last delivery was postponed (deliberate wait, not a failure). - -- While set and invisible, the message is a barrier: its partition is not - -- consumed past it. The next delivery is exempt from the retry_count - -- increment and clears the flag. postponed BOOLEAN NOT NULL DEFAULT FALSE, - PRIMARY KEY (consumer_group, topic, partition_key, message_offset) + PRIMARY KEY (tenant, consumer_group, topic, partition_key, message_offset) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/schema/queue_messages.sql b/platform/extension/messagequeue/mysql/schema/queue_messages.sql index a3a655365..87a178fec 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_messages.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_messages.sql @@ -1,21 +1,23 @@ -- MESSAGES TABLE (Immutable Log) --- Single table for all topics. Partition key determines distribution across workers. +-- Single table for all topics. tenant is the Vitess vindex; partition_key orders work within a tenant. -- Messages are append-only; per-consumer-group delivery tracking is in queue_delivery_state. --- Example: topic="merge_queue", partition_key="uber/cadence" +-- Example: tenant="monorepo/main", topic="merge_queue", partition_key="uber/cadence" CREATE TABLE IF NOT EXISTS queue_messages ( - -- Auto-incrementing global offset for ordering - offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY, + -- tenant is the shard isolation identity (SubmitQueue maps queueName here at wiring) + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, - -- Topic identifies the queue type - topic VARCHAR(255) NOT NULL, + -- Topic identifies the pipeline stage + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, - -- Partition key for distributing work across workers - -- Example: repo ID, user ID, tenant ID - partition_key VARCHAR(255) NOT NULL, + -- Partition key for distributing work across workers within a tenant + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, + + -- Auto-incrementing offset for ordering within (tenant, topic, partition_key) + offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, -- Message identification - id VARCHAR(255) NOT NULL, + id VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Message data payload BLOB NOT NULL, @@ -27,26 +29,16 @@ CREATE TABLE IF NOT EXISTS queue_messages ( -- DLQ-specific fields (0/"" for normal messages, populated for DLQ messages) failed_at BIGINT UNSIGNED NOT NULL, - -- failure_count stores how many times the message failed on the ORIGINAL topic before moving to DLQ failure_count INT UNSIGNED NOT NULL, last_error TEXT NOT NULL, - original_topic VARCHAR(255) NOT NULL, - -- failure_detail holds the structured half of the failure: which entities it - -- was about, plus free-form context. last_error keeps the human-readable - -- message, so this column never has to be decoded to read one, and a plain - -- SELECT last_error stays useful. - -- - -- NULL rather than an empty-string sentinel like its neighbours: a JSON - -- column rejects '' as invalid, and NULL is the honest reading of a failure - -- that recorded no structure — including every row written before this - -- column existed, and the retry-limit backstop, which has none to record. + original_topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, failure_detail JSON, - -- Supports: SELECT ... WHERE topic=? AND partition_key=? AND offset > ? ORDER BY offset - -- Used by subscribers to poll for messages within their assigned partition - INDEX idx_topic_partition_offset (topic, partition_key, offset), + PRIMARY KEY (tenant, topic, partition_key, offset), -- Supports: INSERT ... ON DUPLICATE KEY to enforce idempotent publishes - -- Also enables efficient lookups for message updates/deletes by ID - UNIQUE KEY idx_topic_partition_id (topic, partition_key, id) + UNIQUE KEY idx_tenant_topic_partition_id (tenant, topic, partition_key, id), + + -- InnoDB requires AUTO_INCREMENT column to be leftmost on some index + KEY idx_offset (offset) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql index 0a1717693..4fa0d0c5d 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_offsets.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_offsets.sql @@ -1,20 +1,19 @@ -- CONSUMER OFFSETS TABLE --- Tracks consumption progress per consumer group + topic + partition. +-- Tracks consumption progress per consumer group + tenant + topic + partition. -- Each partition has independent offset tracking for crash recovery. --- --- The primary key (consumer_group, topic, partition_key) serves as the main --- lookup index for all queries in offsetStore. No additional indexes are needed --- because all queries filter by the full primary key or a left prefix of it. CREATE TABLE IF NOT EXISTS queue_offsets ( + -- tenant is the shard isolation identity + 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, @@ -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..47f1eebfe 100644 --- a/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql +++ b/platform/extension/messagequeue/mysql/schema/queue_partition_leases.sql @@ -1,38 +1,28 @@ -- PARTITION LEASES TABLE -- Tracks which worker has leased which partition for exclusive processing. -- Workers must renew leases to maintain ownership; stale leases can be stolen. --- --- The primary key (consumer_group, topic, partition_key) serves as the main --- lookup index. Queries by leased_by always include consumer_group and topic, --- so the primary key's left-prefix is sufficient. CREATE TABLE IF NOT EXISTS queue_partition_leases ( + -- tenant is the shard isolation identity + tenant VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, + -- Consumer group (e.g., "orchestrator") - consumer_group VARCHAR(255) NOT NULL, + consumer_group VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Topic being consumed - topic VARCHAR(255) NOT NULL, + topic VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- Partition that is leased - partition_key VARCHAR(255) NOT NULL, + partition_key VARCHAR(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL, -- Worker that owns the lease (e.g., "worker-1") - leased_by VARCHAR(255) NOT NULL, + leased_by VARCHAR(255) CHARACTER SET ascii COLLATE ascii_bin NOT NULL, -- When lease was acquired (epoch milliseconds) leased_at BIGINT UNSIGNED NOT NULL, -- Last lease renewal timestamp (epoch milliseconds) - -- Used to detect stale leases lease_renewed_at BIGINT UNSIGNED NOT NULL, - -- Primary key ensures each partition can only be leased by one worker per consumer group. - -- Supports: INSERT ... ON DUPLICATE KEY UPDATE for lease acquisition and renewal. - -- Also enables efficient lookups: SELECT ... WHERE consumer_group=? AND topic=? AND partition_key=? - -- Left-prefix covers: SELECT ... WHERE consumer_group=? AND topic=? AND leased_by=? - PRIMARY KEY (consumer_group, topic, partition_key), - - -- Supports: SELECT ... WHERE lease_renewed_at0 means deregistered at that time. deregistered_at BIGINT UNSIGNED NOT NULL, - PRIMARY KEY (consumer_group, topic, subscriber_name) + PRIMARY KEY (tenant, consumer_group, topic, subscriber_name) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/platform/extension/messagequeue/mysql/sql.go b/platform/extension/messagequeue/mysql/sql.go index de4e3b8db..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..9ed6b98d2 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) } } @@ -502,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) } @@ -614,8 +663,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 +693,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,119 +775,165 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { // Uses fair share to limit how many partitions this subscriber acquires; // uncapped skips the fair-share cap entirely (the orphan sweep). func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subscription, uncapped bool) error { - cfg := sub.config - - // Get current leased partitions for fair share computation. - leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - return fmt.Errorf("get leased partitions: %w", err) + if len(s.tenants) == 0 { + return nil } - // Use cached discovered partitions from last tick for fair share cap. - // On the first tick, lastDiscoveredPartitions is nil → fairShareCap sees - // only owned partitions, so a joiner's first-tick cap floors at 1 and - // ramps once discovery is cached. + cfg := sub.config + sub.workersMu.Lock() - cachedDiscovered := sub.lastDiscoveredPartitions + cachedDiscovered := append([]string(nil), sub.lastDiscoveredPartitions...) + existingWorkers := make([]string, 0, len(sub.workers)) + for key := range sub.workers { + existingWorkers = append(existingWorkers, key) + } sub.workersMu.Unlock() - // maxPartitions == 0 means unlimited (the orphan sweep, or an - // uncontended single subscriber via fairShareCap). - maxPartitions := 0 - if !uncapped { - maxPartitions, err = s.fairShareCap(ctx, sub, leasedPartitions, cachedDiscovered) + discoveredByTenant := make(map[string][]string, len(s.tenants)) + leasedByTenant := make(map[string][]string, len(s.tenants)) + var discoveryErrs []error + + 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) + discoveryErrs = append(discoveryErrs, fmt.Errorf("get leased partitions tenant=%s: %w", tenant, err)) + continue } - } - // 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) - // Cache discovered partitions for fairShareCap reuse by rebalance and next tick. - sub.workersMu.Lock() - sub.lastDiscoveredPartitions = discoveredPartitions - sub.workersMu.Unlock() + maxPartitions := 0 + if !uncapped { + maxPartitions, err = s.fairShareCap(ctx, sub, tenant, leasedPartitions, cachedForTenant) + if err != nil { + discoveryErrs = append(discoveryErrs, fmt.Errorf("compute fair share cap tenant=%s: %w", tenant, err)) + continue + } + } - // 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) + _, discoveredPartitions, err := s.leaseStore.DiscoverAndAcquirePartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup, cfg.LeaseDurationMs, maxPartitions) + if err != nil { + 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 { + discoveryErrs = append(discoveryErrs, fmt.Errorf("get leased partitions after acquire tenant=%s: %w", tenant, err)) + continue + } + + discoveredByTenant[tenant] = discoveredPartitions + leasedByTenant[tenant] = leasedPartitions } - // 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. + 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 - sub.drainedSince, expired = updateDrainedTracking(sub.drainedSince, leasedPartitions, discoveredPartitions, grace, time.Now()) + + 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 { + key := workerKey(tenant, pk) + tenantDiscovered = append(tenantDiscovered, key) + allDiscovered = append(allDiscovered, key) + } + tenantLeased := make([]string, 0, len(leasedPartitions)) + for _, pk := range leasedPartitions { + 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() + + sub.drainedSince = nextDrainedSince + sort.Strings(expired) if len(expired) > 0 { released := make(map[string]struct{}, len(expired)) - for _, pk := range expired { - // Delete this consumer group's offsets row first, while the - // lease still guarantees exclusive ownership — nobody else can - // be initializing the partition concurrently. Initialize - // recreates the row if the partition ever comes back. - if err := s.offsetStore.DeleteOffset(ctx, sub.topic, pk, cfg.ConsumerGroup); err != nil { - // Retried next tick — the lease is still held, so the - // partition stays tracked as drained. + 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) - return nil + s.reconcilePartitionWorkers(ctx, sub, allLeased) + return errors.Join(discoveryErrs...) } // updateDrainedTracking recomputes, for every owned partition absent from @@ -911,18 +1010,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 +1033,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 +1056,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 +1071,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 +1135,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 +1144,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 +1170,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 +1180,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 +1202,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 +1224,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 +1249,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 +1265,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 +1325,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 +1349,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 +1386,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 +1394,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 +1402,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 +1416,28 @@ func (s *subscriber) renewLeases(ctx context.Context, sub *subscription, leasedP // releaseAllLeases releases all leases for a topic. func (s *subscriber) releaseAllLeases(ctx context.Context, sub *subscription) error { cfg := sub.config - leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) - if err != nil { - return fmt.Errorf("get leased partitions for release: %w", err) - } + var releaseErrs []error + for _, tenant := range s.tenants { + leasedPartitions, err := s.leaseStore.GetLeasedPartitions(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup) + if err != nil { + 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, 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 { + 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. -func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription) error { +func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription, tenant string) error { cfg := sub.config - if err := s.heartbeatStore.Heartbeat(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("heartbeat: %w", err) + if err := s.heartbeatStore.Heartbeat(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + return fmt.Errorf("heartbeat tenant=%s: %w", tenant, err) } return nil } @@ -1330,10 +1445,13 @@ func (s *subscriber) sendHeartbeat(ctx context.Context, sub *subscription) error // deregisterHeartbeat removes this subscriber's heartbeat entry during shutdown. func (s *subscriber) deregisterHeartbeat(ctx context.Context, sub *subscription) error { cfg := sub.config - if err := s.heartbeatStore.Deregister(ctx, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("deregister heartbeat: %w", err) + var deregistrationErrs []error + for _, tenant := range s.tenants { + if err := s.heartbeatStore.Deregister(ctx, tenant, sub.topic, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { + 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 @@ -1342,15 +1460,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 +1475,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 +1514,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..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" @@ -39,9 +40,10 @@ import ( // so a test only names the parts it cares about. func newDeliveryForTest(sub *subscriber, attempt int, dlq extqueue.DLQConfig, retry extqueue.RetryConfig) *sqlDelivery { msg := entityqueue.NewMessage("msg-1", []byte("payload"), "part-1", nil) + msg.Tenant = testTenant return newSQLDelivery( msg, "1", attempt, nil, - sub, "test_topic", "part-1", 100, "msg-1", "test-group", + sub, testTenant, "test_topic", "part-1", 100, "msg-1", "test-group", dlq, retry, failure.Failure{}, false, ) } @@ -53,22 +55,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 +80,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 +121,7 @@ func TestSubscriber_Subscribe(t *testing.T) { // Reached via releaseAllLeases on the shutdown path, and by the // discovery ticker if it fires before teardown. - mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() + mockLeaseStore.EXPECT().GetLeasedPartitions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return([]string{}, nil).AnyTimes() sub := setupSubscriberTest(t, mockMessageStore, mockOffsetStore, mockLeaseStore) // Close waits for managePartitions to exit; a bare cancel would only @@ -178,6 +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() @@ -185,7 +232,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 +258,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 +326,7 @@ func TestSQLDelivery_Ack(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -290,7 +338,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 +394,7 @@ func TestSQLDelivery_Postpone(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) d := newDeliveryForTest(sub, 1, extqueue.DLQConfig{}, extqueue.RetryConfig{}) @@ -356,7 +405,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 +473,7 @@ func TestSQLDelivery_Reject(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{ @@ -439,19 +489,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 +607,7 @@ func TestSQLDelivery_NackDeadLettersWhenBudgetSpent(t *testing.T) { NewMockpartitionLeaseStore(ctrl), newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) dlqConfig := extqueue.DLQConfig{Enabled: true, TopicSuffix: "_dlq"} @@ -566,14 +617,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 +649,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 +700,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 +774,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 +796,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() @@ -766,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) @@ -784,6 +1016,7 @@ func TestSubscriber_PartitionWorkerPollAndDeliver(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -797,9 +1030,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 +1041,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 +1126,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 +1141,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 +1209,7 @@ func TestSubscriber_PollAndDeliver_PostponedBarrier(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), mockDeliveryState, + []string{testTenant}, ) cfg := testSubscriptionConfig() @@ -987,32 +1223,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 +1289,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 +1311,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 +1375,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 +1395,7 @@ func TestPartitionWorker_RunPollErrorLogging(t *testing.T) { mockLeaseStore, newTestHeartbeatStore(ctrl), newTestDeliveryStateStore(ctrl), + []string{testTenant}, ) s.OnSignal = make(chan HookSignal, 1) cfg := testSubscriptionConfig() @@ -1167,6 +1406,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 +1520,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 +1529,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 +1556,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 +1564,7 @@ func TestSubscriber_FairShareCap(t *testing.T) { NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + []string{testTenant}, ) sum := 0 @@ -1331,7 +1573,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 +1592,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 +1631,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 +1645,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 +1653,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 +1661,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/hook/publisher.go b/platform/hook/publisher.go index 07dc3b4bd..0048fada7 100644 --- a/platform/hook/publisher.go +++ b/platform/hook/publisher.go @@ -23,9 +23,9 @@ import ( "github.com/uber/submitqueue/platform/publish" ) -// Publish sends one hook event to the domain's hook topic, partitioned by -// partitionKey. The topic key is not a parameter: a domain runs a single hook -// topic, and the caller's registry is what binds that key to a wire topic. +// Publish sends one hook event to the domain's hook topic for tenant, +// partitioned by partitionKey. The topic key is not a parameter: a domain runs +// a single hook topic, and the caller's registry binds that key to a wire topic. // // The event id is the message id, so a redelivery republishing the same event // dedups into the original message instead of enqueuing a second one. Callers @@ -37,6 +37,7 @@ import ( func Publish( ctx context.Context, registry consumer.TopicRegistry, + tenant string, event *basehook.HookEvent, partitionKey string, ) error { @@ -49,9 +50,12 @@ func Publish( return fmt.Errorf("failed to serialize hook event %s: %w", event.GetId(), err) } - if err := publish.Message( - ctx, registry, basehook.TopicKeyHook, publish.IntentID(event.GetId()), body, partitionKey, - ); err != nil { + if err := publish.Message(ctx, registry, basehook.TopicKeyHook, publish.MessageParams{ + Tenant: tenant, + ID: publish.IntentID(event.GetId()), + Payload: body, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish hook event %s: %w", event.GetId(), err) } return nil diff --git a/platform/hook/publisher_test.go b/platform/hook/publisher_test.go index c9db7b0ef..ca07233b7 100644 --- a/platform/hook/publisher_test.go +++ b/platform/hook/publisher_test.go @@ -31,6 +31,7 @@ import ( const ( testEventID = "stovepipe/validation.repository.recorded/request/7/0" testPartitionKey = "request/7" + testTenant = "monorepo/main" ) func testEvent() *basehook.HookEvent { @@ -69,12 +70,14 @@ func TestPublish(t *testing.T) { ctrl := gomock.NewController(t) registry, published := registryWithHookTopic(t, ctrl, nil) - require.NoError(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.NoError(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) // The event id is the message id, so a redelivery republishing the same // event dedups into the original message. assert.Equal(t, testEventID, published.ID) assert.Equal(t, testPartitionKey, published.PartitionKey) + assert.Equal(t, testTenant, published.Tenant) + assert.Equal(t, testTenant, published.Metadata[entityqueue.MetadataKeyQueueName]) decoded := &basehook.HookEvent{} require.NoError(t, basehook.Unmarshal(published.Payload, decoded)) @@ -105,7 +108,7 @@ func TestPublish_RejectsMalformedEvent(t *testing.T) { require.NoError(t, err) // No Publish expectation: a malformed event must not reach the queue. - require.Error(t, Publish(context.Background(), registry, tt.event, testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, tt.event, testPartitionKey)) }) } } @@ -114,12 +117,12 @@ func TestPublish_PropagatesPublishFailure(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := registryWithHookTopic(t, ctrl, errors.New("boom")) - require.Error(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) } func TestPublish_FailsWhenHookTopicIsUnregistered(t *testing.T) { registry, err := consumer.NewTopicRegistry(nil) require.NoError(t, err) - require.Error(t, Publish(context.Background(), registry, testEvent(), testPartitionKey)) + require.Error(t, Publish(context.Background(), registry, testTenant, testEvent(), testPartitionKey)) } diff --git a/platform/publish/publish.go b/platform/publish/publish.go index 53cbc7ed4..6fd749225 100644 --- a/platform/publish/publish.go +++ b/platform/publish/publish.go @@ -35,29 +35,37 @@ import ( "github.com/uber/submitqueue/platform/consumer" ) -// Message publishes payload to the topic registered for key. Allowlisted -// delivery context is propagated as message metadata. +// MessageParams describes a message to publish. +type MessageParams struct { + // Tenant selects the queue shard. + Tenant string + // ID identifies the message for deduplication. + ID string + // Payload is the serialized message body. + Payload []byte + // PartitionKey selects the ordered partition. + PartitionKey string + // Metadata contains side-band delivery attributes. + Metadata map[string]string +} + +// Message publishes params.Payload to the topic registered for key. +// Params.Tenant selects the shard and is propagated as queue-name metadata. // -// msgID selects the dedup behavior, so the caller must choose it deliberately. -// The queue deduplicates on (topic, partition key, message ID) against every +// Params.ID selects the dedup behavior, so the caller must choose it deliberately. +// The queue deduplicates on (tenant, topic, partition key, message ID) against every // row it has not garbage-collected yet, consumed ones included — a window with // no upper bound on a busy partition. A publish that collides is reported as a // success and writes nothing, and nothing retries it. // -// Build msgID with IntentID: name the entity the message is about and the cause +// Build params.ID with IntentID: name the entity the message is about and the cause // this particular message exists for. A retry of the same cause then dedups, // which is what makes redelivery safe, while a new cause about the same entity // can never be swallowed by an older row. -func Message(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, msgID string, payload []byte, partitionKey string) error { - return MessageWithMetadata(ctx, registry, key, msgID, payload, partitionKey, nil) -} - -// MessageWithMetadata is Message with side-band message metadata (headers/attributes) -// attached to the delivery. Use it to carry diagnostic context that is not part of -// the payload — the backend persists and redelivers metadata alongside the message. -// Allowlisted delivery context, currently only the queue name, is propagated unless -// the caller supplies that metadata key explicitly. -func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, msgID string, payload []byte, partitionKey string, metadata map[string]string) error { +func Message(ctx context.Context, registry consumer.TopicRegistry, key consumer.TopicKey, params MessageParams) error { + if params.Tenant == "" { + return fmt.Errorf("tenant is required") + } q, ok := registry.Queue(key) if !ok { return fmt.Errorf("no queue registered for topic key %s", key) @@ -67,24 +75,18 @@ func MessageWithMetadata(ctx context.Context, registry consumer.TopicRegistry, k return fmt.Errorf("no topic name registered for topic key %s", key) } - msg := entityqueue.NewMessage(msgID, payload, partitionKey, metadataFromContext(ctx, metadata)) - return q.Publisher().Publish(ctx, topicName, msg) -} - -func metadataFromContext(ctx context.Context, metadata map[string]string) map[string]string { - metadata = maps.Clone(metadata) - if _, exists := metadata[entityqueue.MetadataKeyQueueName]; exists { - return metadata - } - queueName, ok := entityqueue.QueueName(ctx) - if !ok || queueName == "" { - return metadata + if queueName, exists := params.Metadata[entityqueue.MetadataKeyQueueName]; exists && queueName != params.Tenant { + return fmt.Errorf("queue-name metadata %q does not match tenant %q", queueName, params.Tenant) } + metadata := maps.Clone(params.Metadata) if metadata == nil { metadata = make(map[string]string) } - metadata[entityqueue.MetadataKeyQueueName] = queueName - return metadata + metadata[entityqueue.MetadataKeyQueueName] = params.Tenant + + msg := entityqueue.NewMessage(params.ID, params.Payload, params.PartitionKey, metadata) + msg.Tenant = params.Tenant + return q.Publisher().Publish(ctx, topicName, msg) } // IntentID names the occasion to publish rather than the entity published diff --git a/platform/publish/publish_test.go b/platform/publish/publish_test.go index 23ec08ae7..0e6583131 100644 --- a/platform/publish/publish_test.go +++ b/platform/publish/publish_test.go @@ -55,15 +55,21 @@ func TestMessage(t *testing.T) { return nil }) - err := Message(context.Background(), registry, testKey, "msg-1", []byte("payload"), "partition-1") + err := Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "tenant-1", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) require.NoError(t, err) + assert.Equal(t, "tenant-1", published.Tenant) assert.Equal(t, "msg-1", published.ID) assert.Equal(t, []byte("payload"), published.Payload) assert.Equal(t, "partition-1", published.PartitionKey) - assert.Empty(t, published.Metadata) + assert.Equal(t, "tenant-1", published.Metadata[entityqueue.MetadataKeyQueueName]) } -func TestMessage_PropagatesQueueNameFromContext(t *testing.T) { +func TestMessage_PropagatesTenantAsQueueName(t *testing.T) { ctrl := gomock.NewController(t) registry, publisher := newTestRegistry(t, ctrl) @@ -75,12 +81,17 @@ func TestMessage_PropagatesQueueNameFromContext(t *testing.T) { return nil }) - ctx := entityqueue.WithQueueName(context.Background(), "monorepo/main") - require.NoError(t, Message(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1")) + require.NoError(t, Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "monorepo/main", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + })) assert.Equal(t, "monorepo/main", published.Metadata[entityqueue.MetadataKeyQueueName]) + assert.Equal(t, "monorepo/main", published.Tenant) } -func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { +func TestMessage_MergesMetadataWithoutMutatingInput(t *testing.T) { ctrl := gomock.NewController(t) registry, publisher := newTestRegistry(t, ctrl) @@ -93,38 +104,58 @@ func TestMessageWithMetadata_MergesContextWithoutMutatingInput(t *testing.T) { }) metadata := map[string]string{"failure_reason": "build failed"} - ctx := entityqueue.WithQueueName(context.Background(), "monorepo/main") - require.NoError(t, MessageWithMetadata(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1", metadata)) + require.NoError(t, Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "monorepo/main", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + Metadata: metadata, + })) assert.Equal(t, map[string]string{ "failure_reason": "build failed", entityqueue.MetadataKeyQueueName: "monorepo/main", }, published.Metadata) + assert.Equal(t, "monorepo/main", published.Tenant) assert.Equal(t, map[string]string{"failure_reason": "build failed"}, metadata) } -func TestMessageWithMetadata_ExplicitQueueNameWins(t *testing.T) { +func TestMessage_RejectsQueueNameDifferentFromTenant(t *testing.T) { ctrl := gomock.NewController(t) - registry, publisher := newTestRegistry(t, ctrl) - - var published entityqueue.Message - publisher.EXPECT(). - Publish(gomock.Any(), "test-topic", gomock.Any()). - DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { - published = msg - return nil - }) + registry, _ := newTestRegistry(t, ctrl) - ctx := entityqueue.WithQueueName(context.Background(), "inbound") metadata := map[string]string{entityqueue.MetadataKeyQueueName: "outbound"} - require.NoError(t, MessageWithMetadata(ctx, registry, testKey, "msg-1", []byte("payload"), "partition-1", metadata)) - assert.Equal(t, "outbound", published.Metadata[entityqueue.MetadataKeyQueueName]) + err := Message(context.Background(), registry, testKey, MessageParams{ + Tenant: "inbound", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + Metadata: metadata, + }) + require.Error(t, err) } func TestMessage_UnregisteredKey(t *testing.T) { ctrl := gomock.NewController(t) registry, _ := newTestRegistry(t, ctrl) - err := Message(context.Background(), registry, "unregistered-key", "msg-1", []byte("payload"), "partition-1") + err := Message(context.Background(), registry, "unregistered-key", MessageParams{ + Tenant: "tenant-1", + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) + require.Error(t, err) +} + +func TestMessage_RequiresTenant(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newTestRegistry(t, ctrl) + + err := Message(context.Background(), registry, testKey, MessageParams{ + ID: "msg-1", + Payload: []byte("payload"), + PartitionKey: "partition-1", + }) require.Error(t, err) } diff --git a/runway/controller/dlq/dlq.go b/runway/controller/dlq/dlq.go index d53f25c74..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, - publish.IntentID(result.GetId(), "dlq"), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, c.signalTopicKey, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId(), "dlq"), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/dlq/dlq_test.go b/runway/controller/dlq/dlq_test.go index 360852834..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..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, publish.IntentID(result.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/merge/merge_test.go b/runway/controller/merge/merge_test.go index c7900863c..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..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, publish.IntentID(result.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: result.GetQueueName(), + ID: publish.IntentID(result.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go b/runway/controller/mergeconflictcheck/mergeconflictcheck_test.go index 13215239b..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/service/messagequeue/tenant_test.go b/service/messagequeue/tenant_test.go new file mode 100644 index 000000000..950a8b80b --- /dev/null +++ b/service/messagequeue/tenant_test.go @@ -0,0 +1,59 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package messagequeue + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseRequiredTenants(t *testing.T) { + tests := []struct { + name string + value string + want []string + wantErr bool + }{ + { + name: "comma-separated tenants", + value: " monorepo/main, monorepo/release ", + want: []string{"monorepo/main", "monorepo/release"}, + }, + { + name: "duplicates removed in first-seen order", + value: "monorepo/main,monorepo/release,monorepo/main", + want: []string{"monorepo/main", "monorepo/release"}, + }, + {name: "empty", wantErr: true}, + {name: "whitespace and commas", value: " , , ", wantErr: true}, + {name: "tenant exceeds byte limit", value: strings.Repeat("x", maxTenantLength+1), wantErr: true}, + {name: "tenant contains non-ASCII characters", value: "monorepo/café", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseRequiredTenants(tt.value) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} 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/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 dc9985001..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,11 +136,17 @@ func run() error { } defer queueDB.Close() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) + if err != nil { + return fmt.Errorf("failed to configure queue subscribers: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) 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/BUILD.bazel b/service/stovepipe/server/BUILD.bazel index 509957e26..321f5ea23 100644 --- a/service/stovepipe/server/BUILD.bazel +++ b/service/stovepipe/server/BUILD.bazel @@ -21,6 +21,7 @@ go_library( "//platform/extension/messagequeue:go_default_library", "//platform/extension/messagequeue/mysql:go_default_library", "//platform/hook:go_default_library", + "//service/messagequeue:go_default_library", "//service/stovepipe/server/mapper:go_default_library", "//stovepipe/controller:go_default_library", "//stovepipe/controller/build:go_default_library", diff --git a/service/stovepipe/server/main.go b/service/stovepipe/server/main.go index 1e412c4be..bacab144d 100644 --- a/service/stovepipe/server/main.go +++ b/service/stovepipe/server/main.go @@ -43,6 +43,7 @@ import ( extqueue "github.com/uber/submitqueue/platform/extension/messagequeue" queueMySQL "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" platformhook "github.com/uber/submitqueue/platform/hook" + servicemq "github.com/uber/submitqueue/service/messagequeue" "github.com/uber/submitqueue/service/stovepipe/server/mapper" "github.com/uber/submitqueue/stovepipe/controller" "github.com/uber/submitqueue/stovepipe/controller/build" @@ -254,11 +255,17 @@ func run() error { } defer queueDB.Close() + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) + if err != nil { + return fmt.Errorf("failed to configure queue subscribers: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -340,6 +347,7 @@ func run() error { storageFty, materializer, registry, + tenants, ) srv := &StovepipeServer{ pingController: pingController, diff --git a/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/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/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/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 fe781b28a..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" @@ -156,12 +157,25 @@ 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) + } + tenants, err := servicemq.ParseRequiredTenants(os.Getenv("MQ_TENANTS")) + if err != nil { + return fmt.Errorf("failed to configure queue subscribers: %w", err) + } + mysqlQueue, err := queueMySQL.NewQueue(queueMySQL.Params{ DB: queueDB, Logger: logger, LogLevel: os.Getenv("QUEUE_LOG_LEVEL"), MetricsScope: scope.SubScope("queue"), + Tenants: tenants, }) if err != nil { return fmt.Errorf("failed to create queue: %w", err) @@ -176,14 +190,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) 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..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, publish.IntentID(buildID), payload, buildID) + return publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuildSignal, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(buildID), + Payload: payload, + PartitionKey: buildID, + }) } // Name returns the controller name for logging and metrics. diff --git a/stovepipe/controller/build/build_test.go b/stovepipe/controller/build/build_test.go index e5c9b7bc9..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..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, publish.IntentID(requestID), payload, requestID) + return publish.Message(ctx, c.registry, stovepipemq.TopicKeyRecord, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(requestID), + Payload: payload, + PartitionKey: requestID, + }) } // Name returns the controller name for logging and metrics. diff --git a/stovepipe/controller/buildsignal/buildsignal_test.go b/stovepipe/controller/buildsignal/buildsignal_test.go index d20ee7289..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 1c9022f0a..2ffd6cb81 100644 --- a/stovepipe/controller/ingest.go +++ b/stovepipe/controller/ingest.go @@ -20,7 +20,6 @@ import ( "fmt" "github.com/uber-go/tally" - entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/extension/counter" @@ -56,13 +55,14 @@ func IsInvalidRequest(err error) bool { // the request ID onto the process stage. Ingestion is idempotent: a re-reported head resolves to // the already-minted request and republishes it only while it remains accepted. type IngestController struct { - logger *zap.SugaredLogger - metricsScope tally.Scope - counters counter.Factory - sourceControl sourcecontrol.Factory - stores storage.Factory - materializer requestlog.Materializer - registry consumer.TopicRegistry + logger *zap.SugaredLogger + metricsScope tally.Scope + counters counter.Factory + sourceControl sourcecontrol.Factory + stores storage.Factory + materializer requestlog.Materializer + registry consumer.TopicRegistry + configuredQueues map[string]struct{} } // NewIngestController creates a new instance of the stovepipe ingest controller. It publishes @@ -75,15 +75,21 @@ func NewIngestController( stores storage.Factory, materializer requestlog.Materializer, registry consumer.TopicRegistry, + configuredQueues []string, ) *IngestController { + queueNames := make(map[string]struct{}, len(configuredQueues)) + for _, queue := range configuredQueues { + queueNames[queue] = struct{}{} + } return &IngestController{ - logger: logger, - metricsScope: scope.SubScope("ingest_controller"), - counters: counters, - sourceControl: sourceControl, - stores: stores, - materializer: materializer, - registry: registry, + logger: logger, + metricsScope: scope.SubScope("ingest_controller"), + counters: counters, + sourceControl: sourceControl, + stores: stores, + materializer: materializer, + registry: registry, + configuredQueues: queueNames, } } @@ -106,6 +112,9 @@ func (c *IngestController) Ingest(ctx context.Context, req entity.IngestRequest) return entity.IngestResult{}, fmt.Errorf("requires the request to have a queue name specified: %w", ErrInvalidRequest) } queue := req.Queue + if _, ok := c.configuredQueues[queue]; !ok { + return entity.IngestResult{}, fmt.Errorf("queue %q is not configured: %w", queue, ErrInvalidRequest) + } store, err := c.stores.For(storage.Config{QueueName: queue}) if err != nil { @@ -312,8 +321,12 @@ func (c *IngestController) publishProcess(ctx context.Context, id, queue string) return fmt.Errorf("failed to serialize process request: %w", err) } - ctx = entityqueue.WithQueueName(ctx, queue) - if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyProcess, publish.IntentID(id), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyProcess, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(id), + Payload: payload, + PartitionKey: queue, + }); err != nil { return fmt.Errorf("failed to publish process request: %w", err) } return nil diff --git a/stovepipe/controller/ingest_test.go b/stovepipe/controller/ingest_test.go index f500110c6..74492d19d 100644 --- a/stovepipe/controller/ingest_test.go +++ b/stovepipe/controller/ingest_test.go @@ -97,7 +97,7 @@ func newIngestController(t *testing.T, ctrl *gomock.Controller) (*IngestControll }) require.NoError(t, err) - c := NewIngestController(zap.NewNop().Sugar(), tally.NewTestScope("test", nil), staticCounterFactory{counter: m.counter}, m.factory, staticStorageFactory{store: store}, m.materializer, registry) + c := NewIngestController(zap.NewNop().Sugar(), tally.NewTestScope("test", nil), staticCounterFactory{counter: m.counter}, m.factory, staticStorageFactory{store: store}, m.materializer, registry, []string{testQueue}) return c, m } @@ -155,6 +155,7 @@ func TestPublishProcessCarriesQueueMetadata(t *testing.T) { require.NoError(t, c.publishProcess(context.Background(), "request/monorepo/main/7", testQueue)) assert.Equal(t, testQueue, got.PartitionKey) + assert.Equal(t, testQueue, got.Tenant) assert.Equal(t, testQueue, got.Metadata[entityqueue.MetadataKeyQueueName]) } @@ -233,6 +234,13 @@ func TestIngestController_Ingest(t *testing.T) { wantErr: true, wantInvalid: true, }, + { + name: "unconfigured queue is invalid", + queue: "monorepo/unconfigured", + setup: func(m ingestMocks) {}, + wantErr: true, + wantInvalid: true, + }, { name: "unknown queue head is invalid", queue: testQueue, diff --git a/stovepipe/controller/process/process.go b/stovepipe/controller/process/process.go index 75e264784..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, publish.IntentID(id), payload, id); err != nil { + if err := publish.Message(ctx, c.registry, stovepipemq.TopicKeyBuild, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(id), + Payload: payload, + PartitionKey: id, + }); err != nil { return fmt.Errorf("failed to publish build request: %w", err) } return nil @@ -498,7 +503,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..94d606719 100644 --- a/submitqueue/core/request/log.go +++ b/submitqueue/core/request/log.go @@ -52,8 +52,12 @@ func PublishLog(ctx context.Context, registry consumer.TopicRegistry, logEntry e cause = append(cause, occurrence) } - if err := publish.Message(ctx, registry, topickey.TopicKeyLog, - publish.IntentID(logEntry.RequestID, cause...), payload, partitionKey); err != nil { + if err := publish.Message(ctx, registry, topickey.TopicKeyLog, publish.MessageParams{ + Tenant: logEntry.Queue, + ID: publish.IntentID(logEntry.RequestID, cause...), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/core/request/log_test.go b/submitqueue/core/request/log_test.go index ea3314795..044aac039 100644 --- a/submitqueue/core/request/log_test.go +++ b/submitqueue/core/request/log_test.go @@ -128,10 +128,12 @@ func TestPublishLog_MessageIDScopedByStatus(t *testing.T) { ctrl := gomock.NewController(t) var ids []string + var tenants []string mockPub := queuemock.NewMockPublisher(ctrl) mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, _ string, msg entityqueue.Message) error { ids = append(ids, msg.ID) + tenants = append(tenants, msg.Tenant) return nil }, ).AnyTimes() @@ -161,6 +163,7 @@ func TestPublishLog_MessageIDScopedByStatus(t *testing.T) { "req/1/cancelled", "req/1/started", }, ids) + require.Equal(t, []string{"req", "req", "req", "req"}, tenants) } // TestPublishLog_MessageIDScopedByOccurrence locks in the recurring half of the diff --git a/submitqueue/core/request/terminate_test.go b/submitqueue/core/request/terminate_test.go index ac98c0098..fac4cfc09 100644 --- a/submitqueue/core/request/terminate_test.go +++ b/submitqueue/core/request/terminate_test.go @@ -124,7 +124,7 @@ func TestTerminateRequest(t *testing.T) { "already in target terminal state republishes log": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - already := entity.Request{ID: requestID, State: entity.RequestStateError, Version: 5} + already := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateError, Version: 5} rs.EXPECT().Get(gomock.Any(), requestID).Return(already, nil) }, wantResult: TerminationResult{ @@ -140,7 +140,7 @@ func TestTerminateRequest(t *testing.T) { "already in target terminal state returns republish error": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - already := entity.Request{ID: requestID, State: entity.RequestStateError, Version: 5} + already := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateError, Version: 5} rs.EXPECT().Get(gomock.Any(), requestID).Return(already, nil) }, publishErr: fmt.Errorf("connection refused"), @@ -150,7 +150,7 @@ func TestTerminateRequest(t *testing.T) { "diverged terminal state is left untouched": { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { - landed := entity.Request{ID: requestID, State: entity.RequestStateLanded, Version: 7} + landed := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateLanded, Version: 7} rs.EXPECT().Get(gomock.Any(), requestID).Return(landed, nil) }, wantResult: TerminationResult{ diff --git a/submitqueue/gateway/controller/cancel.go b/submitqueue/gateway/controller/cancel.go index 644ae3448..cbaafd667 100644 --- a/submitqueue/gateway/controller/cancel.go +++ b/submitqueue/gateway/controller/cancel.go @@ -139,8 +139,12 @@ func (c *cancelController) publishToQueue(ctx context.Context, cancelRequest ent // The request ID is the message ID with no cause: a request is cancelled at // most once, so a second Cancel for one already being cancelled is meant to // dedup rather than enqueue redundant work. - if err := publish.Message(ctx, c.registry, topickey.TopicKeyCancel, - publish.IntentID(cancelRequest.ID), payload, cancelRequest.ID); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyCancel, publish.MessageParams{ + Tenant: cancelRequest.Queue, + ID: publish.IntentID(cancelRequest.ID), + Payload: payload, + PartitionKey: cancelRequest.ID, + }); err != nil { return fmt.Errorf("failed to publish cancel request message: %w", err) } diff --git a/submitqueue/gateway/controller/cancel_test.go b/submitqueue/gateway/controller/cancel_test.go index 45cefe752..554079906 100644 --- a/submitqueue/gateway/controller/cancel_test.go +++ b/submitqueue/gateway/controller/cancel_test.go @@ -153,6 +153,7 @@ func TestCancel_PublishesToQueue(t *testing.T) { assert.Equal(t, "cancel", publishedTopic) assert.Equal(t, "my-queue/7", publishedMessage.ID) + assert.Equal(t, "my-queue", publishedMessage.Tenant) assert.Equal(t, "my-queue/7", publishedMessage.PartitionKey) deserialized, err := entity.CancelRequestFromBytes(publishedMessage.Payload) diff --git a/submitqueue/gateway/controller/land.go b/submitqueue/gateway/controller/land.go index 0a1a3708d..d03b537a8 100644 --- a/submitqueue/gateway/controller/land.go +++ b/submitqueue/gateway/controller/land.go @@ -212,8 +212,12 @@ func (c *landController) publishToQueue(ctx context.Context, landRequest entity. // retry of this same publish dedups instead of enqueuing it twice // - Payload: serialized LandRequest entity // - Partition key: landRequest.Queue (ensures ordering per queue) - if err := publish.Message(ctx, c.registry, topickey.TopicKeyStart, - publish.IntentID(landRequest.ID), payload, landRequest.Queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyStart, publish.MessageParams{ + Tenant: landRequest.Queue, + ID: publish.IntentID(landRequest.ID), + Payload: payload, + PartitionKey: landRequest.Queue, + }); err != nil { return fmt.Errorf("failed to publish land request message: %w", err) } diff --git a/submitqueue/gateway/controller/land_test.go b/submitqueue/gateway/controller/land_test.go index c957ef23b..52b25f8ee 100644 --- a/submitqueue/gateway/controller/land_test.go +++ b/submitqueue/gateway/controller/land_test.go @@ -459,6 +459,7 @@ func TestLand_PublishesToQueue(t *testing.T) { // Verify message was published to the topic registered under TopicKeyStart assert.Equal(t, "start", publishedTopic) assert.Equal(t, "test-queue/123", publishedMessage.ID) + assert.Equal(t, "test-queue", publishedMessage.Tenant) assert.Equal(t, "test-queue", publishedMessage.PartitionKey) // Verify payload can be deserialized diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 0b60b32d6..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, - 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 afa80cc99..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, publish.IntentID(buildID), payload, buildID); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeyBuildSignal, publish.MessageParams{ + Tenant: queue, + ID: publish.IntentID(buildID), + Payload: payload, + PartitionKey: buildID, + }); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to buildsignal: %w", err) } diff --git a/submitqueue/orchestrator/controller/build/build_test.go b/submitqueue/orchestrator/controller/build/build_test.go index dcb9a69cb..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..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, msgID, payload, queue) + return publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: msgID, + Payload: payload, + PartitionKey: queue, + }) } // Name returns the controller name for logging and metrics. diff --git a/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go b/submitqueue/orchestrator/controller/buildsignal/buildsignal_test.go index 9e02cda19..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..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, 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/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..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, 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/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..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, publish.UniqueID(next), payload, queue); err != nil { + if err := publish.Message(ctx, c.registry, topickey.TopicKeySpeculate, publish.MessageParams{ + Tenant: queue, + ID: publish.UniqueID(next), + Payload: payload, + PartitionKey: queue, + }); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to re-trigger speculation for queue %s: %w", queue, err) } diff --git a/submitqueue/orchestrator/controller/dlq/speculate_test.go b/submitqueue/orchestrator/controller/dlq/speculate_test.go index 3cb627fbe..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..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, publish.IntentID(req.GetId()), payload, partitionKey); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: req.GetQueueName(), + ID: publish.IntentID(req.GetId()), + Payload: payload, + PartitionKey: partitionKey, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/merge/merge_test.go b/submitqueue/orchestrator/controller/merge/merge_test.go index 90c63cd6c..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..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, 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 7611dd837..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, msgID, payload, queue, metadata); err != nil { + if err := publish.Message(ctx, c.registry, key, publish.MessageParams{ + Tenant: queue, + ID: msgID, + Payload: payload, + PartitionKey: queue, + Metadata: metadata, + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go index 73b78c3e3..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..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, 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 5dbfbd04f..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, 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 9eacab3cf..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, publish.IntentID(req.GetId()), payload, req.GetQueueName()); err != nil { + if err := publish.Message(ctx, c.registry, c.runwayTopicKey, publish.MessageParams{ + Tenant: req.GetQueueName(), + ID: publish.IntentID(req.GetId()), + Payload: payload, + PartitionKey: req.GetQueueName(), + }); err != nil { return fmt.Errorf("failed to publish message: %w", err) } diff --git a/submitqueue/orchestrator/controller/validate/validate_test.go b/submitqueue/orchestrator/controller/validate/validate_test.go index 4778d408c..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/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..cd92821cf 100644 --- a/test/e2e/submitqueue/harness_test.go +++ b/test/e2e/submitqueue/harness_test.go @@ -35,6 +35,7 @@ import ( changepb "github.com/uber/submitqueue/api/base/change/protopb" mergestrategypb "github.com/uber/submitqueue/api/base/mergestrategy/protopb" gatewaypb "github.com/uber/submitqueue/api/submitqueue/gateway/protopb" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/extension/consumergate" queuemysql "github.com/uber/submitqueue/platform/extension/messagequeue/mysql" @@ -286,6 +287,7 @@ func (s *E2EIntegrationSuite) redeliverBatchMessage(req request) { DB: s.queueDB, Logger: zap.NewNop(), MetricsScope: tally.NoopScope, + Tenants: []string{req.queue}, }) require.NoError(t, err, "failed to open the queue for a manual publish") defer func() { require.NoError(t, queue.Close()) }() @@ -298,8 +300,12 @@ func (s *E2EIntegrationSuite) redeliverBatchMessage(req request) { payload, err := entity.RequestID{ID: req.sqid, Queue: req.queue}.ToBytes() require.NoError(t, err) - require.NoError(t, publish.Message(s.ctx, registry, topickey.TopicKeyBatch, - publish.UniqueID(req.sqid), payload, req.queue), "failed to redeliver the batch message") + require.NoError(t, publish.Message(entityqueue.WithQueueName(s.ctx, req.queue), registry, topickey.TopicKeyBatch, publish.MessageParams{ + Tenant: req.queue, + ID: publish.UniqueID(req.sqid), + Payload: payload, + PartitionKey: req.queue, + }), "failed to redeliver the batch message") s.log.Logf("Redelivered the batch message for %s", req.sqid) } diff --git a/test/integration/extension/messagequeue/mysql/BUILD.bazel b/test/integration/extension/messagequeue/mysql/BUILD.bazel index e6803cf72..942aeff28 100644 --- a/test/integration/extension/messagequeue/mysql/BUILD.bazel +++ b/test/integration/extension/messagequeue/mysql/BUILD.bazel @@ -2,7 +2,10 @@ load("@rules_go//go:def.bzl", "go_test") go_test( name = "go_default_test", - srcs = ["queue_test.go"], + srcs = [ + "queue_test.go", + "tenant_isolation_test.go", + ], data = [ "docker-compose.yml", "//platform/extension/messagequeue/mysql/schema", diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index a87503019..0f54ef176 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -38,6 +38,8 @@ import ( "github.com/uber/submitqueue/test/testutil" ) +const testTenant = "test-tenant" + type SQLQueueIntegrationSuite struct { suite.Suite ctx context.Context @@ -50,6 +52,25 @@ func TestSQLQueueIntegration(t *testing.T) { suite.Run(t, new(SQLQueueIntegrationSuite)) } +func (s *SQLQueueIntegrationSuite) testQueueParams(t *testing.T, extra func(*queueMySQL.Params)) queueMySQL.Params { + p := queueMySQL.Params{ + DB: s.db, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + Tenants: []string{testTenant}, + } + if extra != nil { + extra(&p) + } + return p +} + +func testMessage(id string, payload []byte, partitionKey string, metadata map[string]string) entityqueue.Message { + msg := entityqueue.NewMessage(id, payload, partitionKey, metadata) + msg.Tenant = testTenant + return msg +} + func (s *SQLQueueIntegrationSuite) SetupSuite() { t := s.T() s.ctx = context.Background() @@ -99,6 +120,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 @@ -322,7 +379,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 +402,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 +418,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 +469,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 +486,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 +521,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 +537,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 +570,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 +590,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 +613,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 +632,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 +667,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 +692,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 +708,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 +724,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 +742,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 +778,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 +795,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 +803,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 +818,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 +842,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 +871,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 +887,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 +901,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 +929,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 +958,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 +1002,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 +1033,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 +1068,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 +1079,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 +1104,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 +1133,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 +1152,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 +1233,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 +1252,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 +1281,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 +1293,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 +1330,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 +1353,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 +1370,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 +1384,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 +1413,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 +1450,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 +1478,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 +1502,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 +1528,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 +1566,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 +1575,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 +1592,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 +1615,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 +1625,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 +1646,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 +1656,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 +1672,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 +1692,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 +1700,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 +1714,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 +1728,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 +1751,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 +1780,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 +1805,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 +1832,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 +1878,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 +1894,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 +1947,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 +1960,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 +2004,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 +2017,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 +2061,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 +2129,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 +2145,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 +2185,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 +2199,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 +2211,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 +2221,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 +2243,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 +2254,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 +2267,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 +2285,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 +2303,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 +2321,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 +2330,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 +2356,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 +2378,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 +2389,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 +2423,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 +2438,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 +2480,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 +2499,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 +2525,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 +2549,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 +2597,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 +2639,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 +2678,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 +2698,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 +2757,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 +2794,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 +2804,7 @@ func (s *SQLQueueIntegrationSuite) TestWatermarkAdvancesContiguously() { // Publish 5 messages to the same partition for i := 1; i <= 5; i++ { - msg := entityqueue.NewMessage( + msg := testMessage( fmt.Sprintf("wm-msg-%d", i), []byte(fmt.Sprintf("payload-%d", i)), "wm-part", diff --git a/test/integration/extension/messagequeue/mysql/tenant_isolation_test.go b/test/integration/extension/messagequeue/mysql/tenant_isolation_test.go new file mode 100644 index 000000000..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/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..d900bc11d 100644 --- a/tool/linter/queueshard/main.go +++ b/tool/linter/queueshard/main.go @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Command queueshard checks that every domain table is shardable by queue: its -// primary key must lead with the queue column, and no secondary index may span -// queues. +// Command queueshard checks that every table is shardable by its owning key: +// its primary key must lead with the shard column, and no secondary index may +// span shards. // // A table is shardable by queue when one queue's rows are unreachable through // another queue's binding. That holds exactly when the queue is the leading @@ -32,23 +32,20 @@ import ( "strings" ) -// queueColumns are the column names that identify the owning queue. Most tables -// call it "queue"; a table whose rows *are* queues (stovepipe's queue table) -// names it "name" because the row's identity is the queue itself. -var queueColumns = map[string]bool{ - "queue": true, - "name": true, +// schemaShardColumns identifies the allowed shard columns for each schema root. +var schemaShardColumns = map[string]map[string]bool{ + "submitqueue/extension/storage/mysql/schema": {"queue": true, "name": true}, + "stovepipe/extension/storage/mysql/schema": {"queue": true, "name": true}, + "platform/extension/counter/mysql/schema": {"queue": true, "name": true}, + "platform/extension/messagequeue/mysql/schema": {"tenant": true}, } // schemaRoots are the directories scanned for table definitions. -// -// platform/extension/messagequeue is deliberately absent: it is a message-queue -// backend keyed by (consumer_group, topic, partition_key), not a domain table -// set, and sharding it is tracked separately. var schemaRoots = []string{ "submitqueue/extension/storage/mysql/schema", "stovepipe/extension/storage/mysql/schema", "platform/extension/counter/mysql/schema", + "platform/extension/messagequeue/mysql/schema", } var ( @@ -76,6 +73,7 @@ func main() { var violations []violation var checked int for _, schemaRoot := range schemaRoots { + shardColumns := schemaShardColumns[schemaRoot] files, err := filepath.Glob(filepath.Join(root, schemaRoot, "*.sql")) if err != nil { fmt.Fprintf(os.Stderr, "error globbing %s: %v\n", schemaRoot, err) @@ -95,28 +93,27 @@ func main() { if relErr != nil { rel = file } - found, tableViolations := check(rel, string(content)) + found, tableViolations := check(rel, string(content), shardColumns) checked += found violations = append(violations, tableViolations...) } } if len(violations) > 0 { - fmt.Fprintf(os.Stderr, "%d table(s) are not shardable by queue:\n\n", len(violations)) + fmt.Fprintf(os.Stderr, "%d table(s) are not shardable:\n\n", len(violations)) for _, v := range violations { fmt.Fprintf(os.Stderr, " %s: table %q %s\n", v.file, v.table, v.problem) } - fmt.Fprintf(os.Stderr, "\nEvery table's primary key must lead with the queue column, and no\n") - fmt.Fprintf(os.Stderr, "secondary index may span queues, so that one queue's rows are\n") - fmt.Fprintf(os.Stderr, "unreachable through another queue's binding.\n") + fmt.Fprintf(os.Stderr, "\nEvery table's primary key must lead with its schema's shard column, and no\n") + fmt.Fprintf(os.Stderr, "secondary index may span shards.\n") os.Exit(1) } - fmt.Printf("All %d tables are shardable by queue.\n", checked) + fmt.Printf("All %d tables are shardable.\n", checked) } // check returns the number of tables found in content and any violations. -func check(file, content string) (int, []violation) { +func check(file, content string, shardColumns map[string]bool) (int, []violation) { var violations []violation matches := createTableRe.FindAllStringSubmatch(content, -1) for _, match := range matches { @@ -132,23 +129,29 @@ func check(file, content string) (int, []violation) { violations = append(violations, violation{file, table, "has an empty PRIMARY KEY"}) continue } - if !queueColumns[columns[0]] { + if !shardColumns[columns[0]] { violations = append(violations, violation{ file, table, - fmt.Sprintf("leads its PRIMARY KEY with %q, not the queue column", columns[0]), + fmt.Sprintf("leads its PRIMARY KEY with %q, not a shard column", columns[0]), }) } for _, idx := range indexRe.FindAllStringSubmatch(body, -1) { idxColumns := splitColumns(idx[2]) - if len(idxColumns) == 0 || !queueColumns[idxColumns[0]] { + if table == "queue_messages" && idx[1] == "idx_offset" && + len(idxColumns) == 1 && idxColumns[0] == "offset" { + // InnoDB requires AUTO_INCREMENT to be the leading column of some + // index; queue_messages keeps a dedicated offset-only index for that. + continue + } + if len(idxColumns) == 0 || !shardColumns[idxColumns[0]] { lead := "(empty)" if len(idxColumns) > 0 { lead = idxColumns[0] } violations = append(violations, violation{ file, table, - fmt.Sprintf("has index %q leading with %q, which spans queues", idx[1], lead), + fmt.Sprintf("has index %q leading with %q, which spans shards", idx[1], lead), }) } } diff --git a/tool/linter/queueshard/main_test.go b/tool/linter/queueshard/main_test.go index 2b5887b4f..551d93d2c 100644 --- a/tool/linter/queueshard/main_test.go +++ b/tool/linter/queueshard/main_test.go @@ -47,7 +47,7 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, wantViolations: 1, - wantProblem: `leads its PRIMARY KEY with "request_id", not the queue column`, + wantProblem: `leads its PRIMARY KEY with "request_id", not a shard column`, }, { name: "queue present but not leading is rejected", @@ -87,7 +87,7 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, wantViolations: 1, - wantProblem: `has index "idx_status" leading with "status", which spans queues`, + wantProblem: `has index "idx_status" leading with "status", which spans shards`, }, { name: "a queue-leading secondary index passes", @@ -100,6 +100,40 @@ func TestCheck(t *testing.T) { ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", wantTables: 1, }, + { + name: "queue_messages offset index passes", + schema: "CREATE TABLE IF NOT EXISTS queue_messages (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + }, + { + name: "offset index on another table is rejected", + schema: "CREATE TABLE IF NOT EXISTS queue_offsets (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + wantViolations: 1, + wantProblem: `has index "idx_offset" leading with "offset", which spans shards`, + }, + { + name: "differently named queue_messages offset index is rejected", + schema: "CREATE TABLE IF NOT EXISTS queue_messages (\n" + + " tenant VARCHAR(191) NOT NULL,\n" + + " offset BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,\n" + + " PRIMARY KEY (tenant, offset),\n" + + " KEY idx_unscoped_offset (offset)\n" + + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;", + wantTables: 1, + wantViolations: 1, + wantProblem: `has index "idx_unscoped_offset" leading with "offset", which spans shards`, + }, { name: "a table with no primary key is rejected", schema: "CREATE TABLE IF NOT EXISTS loose (\n" + @@ -118,7 +152,7 @@ func TestCheck(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - tables, violations := check("test.sql", tt.schema) + tables, violations := check("test.sql", tt.schema, map[string]bool{"queue": true, "name": true, "tenant": true}) assert.Equal(t, tt.wantTables, tables) require.Len(t, violations, tt.wantViolations) if tt.wantProblem != "" {