[2/4 messagequeue] Shard MySQL queues by tenant identity

platform/base/messagequeue/BUILD.bazel

@@ -4,7 +4,9 @@ go_library(
     name = "go_default_library",
     srcs = [
         "context.go",
+        "identity.go",
         "message.go",
+        "partition.go",
     ],
     importpath = "github.com/uber/submitqueue/platform/base/messagequeue",
     visibility = ["//visibility:public"],
@@ -14,8 +16,12 @@ go_test(
     name = "go_default_test",
     srcs = [
         "context_test.go",
+        "identity_test.go",
         "message_test.go",
     ],
     embed = [":go_default_library"],
-    deps = ["@com_github_stretchr_testify//assert:go_default_library"],
+    deps = [
+        "@com_github_stretchr_testify//assert:go_default_library",
+        "@com_github_stretchr_testify//require:go_default_library",
+    ],
 )

platform/base/messagequeue/identity.go

@@ -0,0 +1,45 @@
+// Copyright (c) 2026 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package messagequeue
+
+import "fmt"
+
+// ValidateTenantMetadata verifies that the persisted shard identity agrees
+// with queue metadata when metadata is present.
+func ValidateTenantMetadata(message Message) error {
+	if message.Tenant == "" {
+		return fmt.Errorf("message tenant is required")
+	}
+	queueName := message.Metadata[MetadataKeyQueueName]
+	if queueName != "" && queueName != message.Tenant {
+		return fmt.Errorf("message tenant %q does not match queue metadata %q", message.Tenant, queueName)
+	}
+	return nil
+}
+
+// ValidatePayloadQueue verifies that a queue-bearing payload belongs to the
+// persisted message tenant.
+func ValidatePayloadQueue(message Message, payloadQueue string) error {
+	if err := ValidateTenantMetadata(message); err != nil {
+		return err
+	}
+	if payloadQueue == "" {
+		return fmt.Errorf("payload queue is required")
+	}
+	if payloadQueue != message.Tenant {
+		return fmt.Errorf("message tenant %q does not match payload queue %q", message.Tenant, payloadQueue)
+	}
+	return nil
+}

platform/base/messagequeue/identity_test.go

@@ -0,0 +1,70 @@
+// Copyright (c) 2026 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package messagequeue
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidateTenantMetadata(t *testing.T) {
+	tests := []struct {
+		name    string
+		message Message
+		wantErr bool
+	}{
+		{name: "tenant without metadata", message: Message{Tenant: "queue-a"}},
+		{name: "matching metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-a"}}},
+		{name: "empty tenant", message: Message{}, wantErr: true},
+		{name: "conflicting metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-b"}}, wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := ValidateTenantMetadata(tt.message)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}
+
+func TestValidatePayloadQueue(t *testing.T) {
+	tests := []struct {
+		name         string
+		message      Message
+		payloadQueue string
+		wantErr      bool
+	}{
+		{name: "matching queue", message: Message{Tenant: "queue-a"}, payloadQueue: "queue-a"},
+		{name: "empty payload queue", message: Message{Tenant: "queue-a"}, wantErr: true},
+		{name: "conflicting payload queue", message: Message{Tenant: "queue-a"}, payloadQueue: "queue-b", wantErr: true},
+		{name: "conflicting metadata", message: Message{Tenant: "queue-a", Metadata: map[string]string{MetadataKeyQueueName: "queue-b"}}, payloadQueue: "queue-a", wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := ValidatePayloadQueue(tt.message, tt.payloadQueue)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}

platform/base/messagequeue/message.go

@@ -43,11 +43,15 @@ type Message struct {
 	// Use for trace IDs, request IDs, and cross-service metadata.
 	Metadata map[string]string
 
-	// PartitionKey determines which partition/shard this message goes to.
-	// Messages with the same PartitionKey are guaranteed ordered delivery.
+	// PartitionKey determines the ordering partition within Tenant.
+	// Messages with the same Tenant and PartitionKey are delivered in order.
 	// Optional - if empty, backend may use round-robin distribution.
 	PartitionKey string
 
+	// Tenant is the shard isolation identity persisted by the backend.
+	// Domains map their queue name onto this field at publish time.
+	Tenant string
+
 	// PublishedAt is when the message was published (Unix milliseconds).
 	PublishedAt int64
 }
@@ -79,6 +83,7 @@ func (m Message) Copy() Message {
 		Payload:      payloadCopy,
 		Metadata:     maps.Clone(m.Metadata),
 		PartitionKey: m.PartitionKey,
+		Tenant:       m.Tenant,
 		PublishedAt:  m.PublishedAt,
 	}
 }

platform/base/messagequeue/partition.go

@@ -0,0 +1,32 @@
+// Copyright (c) 2026 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package messagequeue
+
+// PartitionIdentity uniquely identifies a partition within a tenant.
+// It is comparable and may be used as a map key.
+type PartitionIdentity struct {
+	// Tenant is the shard isolation identity.
+	Tenant string
+	// PartitionKey is the partition key within the tenant.
+	PartitionKey string
+}
+
+// PartitionIdentity returns the message's tenant-scoped partition identity.
+func (m Message) PartitionIdentity() PartitionIdentity {
+	return PartitionIdentity{
+		Tenant:       m.Tenant,
+		PartitionKey: m.PartitionKey,
+	}
+}

platform/consumer/consumer.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/uber-go/tally"
+	"github.com/uber/submitqueue/platform/base/failure"
 	entityqueue "github.com/uber/submitqueue/platform/base/messagequeue"
 	"github.com/uber/submitqueue/platform/errs"
 	"github.com/uber/submitqueue/platform/extension/consumergate"
@@ -282,10 +283,10 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic
 		"topic_key", topicKey,
 	)
 
-	// partitionChs maps partition keys to per-partition delivery channels.
+	// partitionChs maps tenant-scoped partitions to delivery channels.
 	// Each channel is created lazily on the first message for that partition
 	// and is never removed — partitions are stable for the lifetime of a subscription.
-	partitionChs := make(map[string]chan extqueue.Delivery)
+	partitionChs := make(map[entityqueue.PartitionIdentity]chan extqueue.Delivery)
 	var wg sync.WaitGroup
 
 	for {
@@ -307,15 +308,18 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic
 				m.shutdownPartitions(partitionChs, &wg)
 				return
 			}
+			if m.rejectDeliveryWithInvalidTenant(ctx, controller, delivery) {
+				continue
+			}
 
 			// Route delivery to its partition's channel, creating the channel
 			// and spawning a processPartition goroutine if this is the first
 			// message for that partition.
-			partitionKey := delivery.Message().PartitionKey
-			ch, exists := partitionChs[partitionKey]
+			partition := delivery.Message().PartitionIdentity()
+			ch, exists := partitionChs[partition]
 			if !exists {
 				ch = make(chan extqueue.Delivery, batchSize)
-				partitionChs[partitionKey] = ch
+				partitionChs[partition] = ch
 				wg.Add(1)
 				go func(pCh <-chan extqueue.Delivery) {
 					defer wg.Done()
@@ -336,9 +340,32 @@ func (m *consumer) consumeLoop(ctx context.Context, controller Controller, topic
 	}
 }
 
+func (m *consumer) rejectDeliveryWithInvalidTenant(ctx context.Context, controller Controller, delivery extqueue.Delivery) bool {
+	msg := delivery.Message()
+	if err := entityqueue.ValidateTenantMetadata(msg); err != nil {
+		m.logger.Errorw("rejecting message with inconsistent queue identity",
+			"controller", controller.Name(),
+			"topic_key", controller.TopicKey(),
+			"message_id", msg.ID,
+			"tenant", msg.Tenant,
+			"error", err,
+		)
+		if rejectErr := delivery.Reject(ctx, failure.New(err.Error())); rejectErr != nil {
+			m.logger.Errorw("failed to reject message with inconsistent queue identity",
+				"controller", controller.Name(),
+				"topic_key", controller.TopicKey(),
+				"message_id", msg.ID,
+				"error", rejectErr,
+			)
+		}
+		return true
+	}
+	return false
+}
+
 // shutdownPartitions closes all partition channels to signal processPartition
 // goroutines to exit, then waits for them to finish draining.
-func (m *consumer) shutdownPartitions(partitionChs map[string]chan extqueue.Delivery, wg *sync.WaitGroup) {
+func (m *consumer) shutdownPartitions(partitionChs map[entityqueue.PartitionIdentity]chan extqueue.Delivery, wg *sync.WaitGroup) {
 	for _, ch := range partitionChs {
 		close(ch)
 	}
@@ -373,11 +400,9 @@ func (m *consumer) processDelivery(ctx context.Context, controller Controller, d
 	const opName = "process"
 
 	msg := delivery.Message()
-	queueName := msg.Metadata[entityqueue.MetadataKeyQueueName]
+	queueName := msg.Tenant
 	ctx = entityqueue.WithQueueName(ctx, queueName)
-	if queueName != "" {
-		ctx = metrics.WithContextTags(ctx, metrics.NewTag("queue", queueName))
-	}
+	ctx = metrics.WithContextTags(ctx, metrics.NewTag("queue", queueName))
 
 	// Consumer gate: a delivery whose gate is closed is recorded as parked and
 	// postponed (barrier + re-check on redelivery); a false return also covers
@@ -575,7 +600,10 @@ func (m *consumer) checkGate(ctx context.Context, controller Controller, deliver
 	consumerGroup := controller.ConsumerGroup()
 	topic := controller.TopicKey().String()
 
-	entry, err := m.gate.Enter(ctx, consumergate.Key{ConsumerGroup: consumerGroup, PartitionKey: msg.PartitionKey})
+	entry, err := m.gate.Enter(ctx, consumergate.Key{
+		ConsumerGroup: consumerGroup,
+		Partition:     msg.PartitionIdentity(),
+	})
 	if err != nil {
 		if errors.Is(err, context.Canceled) {
 			// Shutting down: leave the delivery in flight; visibility lapses
@@ -708,10 +736,11 @@ func (m *consumer) unsubscribeAll(timeoutMs int64) error {
 	var timedOutControllers []string
 	for topicKey, sub := range m.subscriptions {
 		start := time.Now()
+		timer := time.NewTimer(remaining)
 		select {
 		case <-sub.done:
-			// Controller stopped gracefully
-		case <-time.After(remaining):
+			timer.Stop()
+		case <-timer.C:
 			m.logger.Errorw("timeout waiting for controller to stop",
 				"controller", sub.controller.Name(),
 				"topic_key", topicKey,