diff --git a/platform/extension/messagequeue/mysql/vitess/BUILD.bazel b/platform/extension/messagequeue/mysql/vitess/BUILD.bazel new file mode 100644 index 00000000..ccc3f247 --- /dev/null +++ b/platform/extension/messagequeue/mysql/vitess/BUILD.bazel @@ -0,0 +1,5 @@ +filegroup( + name = "vschema", + srcs = ["vschema.json"], + visibility = ["//visibility:public"], +) diff --git a/platform/extension/messagequeue/mysql/vitess/vschema.json b/platform/extension/messagequeue/mysql/vitess/vschema.json new file mode 100644 index 00000000..826bed8c --- /dev/null +++ b/platform/extension/messagequeue/mysql/vitess/vschema.json @@ -0,0 +1,50 @@ +{ + "sharded": true, + "vindexes": { + "tenant_xxhash": { + "type": "xxhash" + } + }, + "tables": { + "queue_delivery_state": { + "column_vindexes": [ + { + "column": "tenant", + "name": "tenant_xxhash" + } + ] + }, + "queue_messages": { + "column_vindexes": [ + { + "column": "tenant", + "name": "tenant_xxhash" + } + ] + }, + "queue_offsets": { + "column_vindexes": [ + { + "column": "tenant", + "name": "tenant_xxhash" + } + ] + }, + "queue_partition_leases": { + "column_vindexes": [ + { + "column": "tenant", + "name": "tenant_xxhash" + } + ] + }, + "queue_subscriber_heartbeats": { + "column_vindexes": [ + { + "column": "tenant", + "name": "tenant_xxhash" + } + ] + } + } +} diff --git a/test/integration/extension/messagequeue/mysql/vitess/BUILD.bazel b/test/integration/extension/messagequeue/mysql/vitess/BUILD.bazel new file mode 100644 index 00000000..5da4d195 --- /dev/null +++ b/test/integration/extension/messagequeue/mysql/vitess/BUILD.bazel @@ -0,0 +1,29 @@ +load("@rules_go//go:def.bzl", "go_test") + +go_test( + name = "go_default_test", + srcs = ["vitess_test.go"], + data = [ + "Dockerfile", + "docker-compose.yml", + "//platform/extension/messagequeue/mysql/schema", + "//platform/extension/messagequeue/mysql/vitess:vschema", + ], + tags = [ + "integration", + "requires-network", + ], + deps = [ + "//platform/base/failure:go_default_library", + "//platform/base/messagequeue:go_default_library", + "//platform/extension/messagequeue:go_default_library", + "//platform/extension/messagequeue/mysql:go_default_library", + "//platform/extension/messagequeue/mysql/ctl/lib:go_default_library", + "//test/testutil:go_default_library", + "@com_github_go_sql_driver_mysql//:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + "@com_github_uber_go_tally//:go_default_library", + "@org_uber_go_zap//zaptest:go_default_library", + ], +) diff --git a/test/integration/extension/messagequeue/mysql/vitess/Dockerfile b/test/integration/extension/messagequeue/mysql/vitess/Dockerfile new file mode 100644 index 00000000..6cc1026c --- /dev/null +++ b/test/integration/extension/messagequeue/mysql/vitess/Dockerfile @@ -0,0 +1,4 @@ +FROM vitess/vttestserver:v23.0.5-mysql80@sha256:79653b39ca8511ebe225018f32090e061b93c0ac5d40596a3c3da5abf921e272 + +COPY platform/extension/messagequeue/mysql/schema/*.sql /vt/test-schema/submitqueue/ +COPY platform/extension/messagequeue/mysql/vitess/vschema.json /vt/test-schema/submitqueue/vschema.json diff --git a/test/integration/extension/messagequeue/mysql/vitess/docker-compose.yml b/test/integration/extension/messagequeue/mysql/vitess/docker-compose.yml new file mode 100644 index 00000000..058930a5 --- /dev/null +++ b/test/integration/extension/messagequeue/mysql/vitess/docker-compose.yml @@ -0,0 +1,22 @@ +services: + vtcombo: + platform: linux/amd64 + build: + context: ${REPO_ROOT} + dockerfile: test/integration/extension/messagequeue/mysql/vitess/Dockerfile + command: + - /vt/bin/vttestserver + - --alsologtostderr + - --keyspaces=submitqueue + - --mysql-bind-host=0.0.0.0 + - --num-shards=2 + - --port=33574 + - --schema-dir=/vt/test-schema + - --vtcombo-bind-host=0.0.0.0 + ports: + - "33577" + healthcheck: + test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-P", "33577"] + interval: 2s + timeout: 2s + retries: 90 diff --git a/test/integration/extension/messagequeue/mysql/vitess/vitess_test.go b/test/integration/extension/messagequeue/mysql/vitess/vitess_test.go new file mode 100644 index 00000000..c543d773 --- /dev/null +++ b/test/integration/extension/messagequeue/mysql/vitess/vitess_test.go @@ -0,0 +1,269 @@ +// 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 vitess + +import ( + "context" + "database/sql" + "fmt" + "testing" + + _ "github.com/go-sql-driver/mysql" + "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" + queueAdmin "github.com/uber/submitqueue/platform/extension/messagequeue/mysql/ctl/lib" + "github.com/uber/submitqueue/test/testutil" + "go.uber.org/zap/zaptest" +) + +const ( + keyspace = "submitqueue" + shardLower = "-80" + shardUpper = "80-" + vtgatePort = 33577 + testTopic = "vitess_tenant_isolation" + partitionKey = "shared-partition" + messageID = "shared-message" +) + +func TestTenantShardingThroughVTGate(t *testing.T) { + ctx := t.Context() + log := testutil.NewTestLogger(t) + stack := testutil.NewComposeStack( + t, + log, + ctx, + "docker-compose.yml", + "ext-messagequeue-vitess", + testutil.WithBuildContext(vitessBuildContext()), + ) + require.NoError(t, stack.Up()) + + vtgate := connectVTGate(t, stack, keyspace) + lowerShard := connectVTGate(t, stack, keyspace+":"+shardLower) + upperShard := connectVTGate(t, stack, keyspace+":"+shardUpper) + shards := map[string]*sql.DB{ + shardLower: lowerShard, + shardUpper: upperShard, + } + + tenants := findTenantsOnDifferentShards(t, ctx, vtgate, shards) + tenantLower := tenants[shardLower] + tenantUpper := tenants[shardUpper] + require.NotEqual(t, tenantLower, tenantUpper) + + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: vtgate, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + Tenants: []string{tenantLower, tenantUpper}, + }) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, q.Close()) + }) + + cfg := extqueue.DefaultSubscriptionConfig("vitess-worker", "vitess-consumer") + cfg.PartitionDiscoveryIntervalMs = 100 + cfg.VisibilityTimeoutMs = cfg.LeaseDurationMs * 10 + deliveries, err := q.Subscriber().Subscribe(ctx, testTopic, cfg) + require.NoError(t, err) + + for _, tenant := range []string{tenantLower, tenantUpper} { + msg := entityqueue.NewMessage(messageID, []byte(tenant), partitionKey, nil) + msg.Tenant = tenant + require.NoError(t, q.Publisher().Publish(ctx, testTopic, msg)) + } + + received := make(map[string]extqueue.Delivery, 2) + for len(received) < 2 { + select { + case <-ctx.Done(): + require.FailNow(t, "timed out waiting for both tenant deliveries", ctx.Err()) + case delivery, ok := <-deliveries: + require.True(t, ok) + require.NotNil(t, delivery) + received[delivery.Message().Tenant] = delivery + } + } + assert.ElementsMatch(t, []string{tenantLower, tenantUpper}, mapKeys(received)) + + for _, table := range []string{ + "queue_messages", + "queue_delivery_state", + "queue_offsets", + "queue_partition_leases", + "queue_subscriber_heartbeats", + } { + assertTenantOnOnlyShard(t, ctx, shards, table, tenantLower, shardLower) + assertTenantOnOnlyShard(t, ctx, shards, table, tenantUpper, shardUpper) + } + + require.NoError(t, received[tenantLower].Reject( + ctx, + failure.New("vitess shard-local DLQ test"), + )) + assertTopicOnOnlyShard(t, ctx, shards, tenantLower, testTopic+"_dlq", shardLower) + assert.Zero(t, tenantTopicCount(t, ctx, lowerShard, tenantLower, testTopic)) + assertTopicOnOnlyShard(t, ctx, shards, tenantUpper, testTopic, shardUpper) + + topics, err := queueAdmin.NewAdminStore(vtgate).ListTopics(ctx, queueAdmin.TenantScope{AllTenants: true}) + require.NoError(t, err) + assert.Contains(t, topics, queueAdmin.TopicInfo{Tenant: tenantLower, Topic: testTopic + "_dlq", MessageCount: 1}) + assert.Contains(t, topics, queueAdmin.TopicInfo{Tenant: tenantUpper, Topic: testTopic, MessageCount: 1}) + + require.NoError(t, received[tenantUpper].Ack(ctx)) +} + +func vitessBuildContext() map[string]string { + const ( + schemaRoot = "platform/extension/messagequeue/mysql/schema/" + testRoot = "test/integration/extension/messagequeue/mysql/vitess/" + ) + files := map[string]string{ + testRoot + "Dockerfile": testRoot + "Dockerfile", + "platform/extension/messagequeue/mysql/vitess/vschema.json": "platform/extension/messagequeue/mysql/vitess/vschema.json", + } + for _, name := range []string{ + "queue_delivery_state.sql", + "queue_messages.sql", + "queue_offsets.sql", + "queue_partition_leases.sql", + "queue_subscriber_heartbeats.sql", + } { + files[schemaRoot+name] = schemaRoot + name + } + return files +} + +func connectVTGate(t *testing.T, stack *testutil.ComposeStack, database string) *sql.DB { + t.Helper() + port, err := stack.ServicePort("vtcombo", vtgatePort) + require.NoError(t, err) + db, err := sql.Open("mysql", fmt.Sprintf("root@tcp(localhost:%d)/%s?parseTime=true&interpolateParams=true", port, database)) + require.NoError(t, err) + require.NoError(t, db.Ping()) + t.Cleanup(func() { + require.NoError(t, db.Close()) + }) + return db +} + +func findTenantsOnDifferentShards( + t *testing.T, + ctx context.Context, + vtgate *sql.DB, + shards map[string]*sql.DB, +) map[string]string { + t.Helper() + const probeTopic = "vitess_routing_probe" + + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: vtgate, + Logger: zaptest.NewLogger(t), + MetricsScope: tally.NoopScope, + }) + require.NoError(t, err) + + tenantsByShard := make(map[string]string, 2) + var publishedTenants []string + for candidate := 0; candidate < 64 && len(tenantsByShard) < len(shards); candidate++ { + tenant := fmt.Sprintf("vitess-tenant-%d", candidate) + msg := entityqueue.NewMessage(messageID, []byte(tenant), partitionKey, nil) + msg.Tenant = tenant + require.NoError(t, q.Publisher().Publish(ctx, probeTopic, msg)) + publishedTenants = append(publishedTenants, tenant) + + for shard, db := range shards { + if tenantTopicCount(t, ctx, db, tenant, probeTopic) == 1 { + tenantsByShard[shard] = tenant + } + } + } + require.NoError(t, q.Close()) + require.Len(t, tenantsByShard, len(shards)) + + for _, tenant := range publishedTenants { + _, err := vtgate.ExecContext(ctx, "DELETE FROM queue_messages WHERE tenant = ? AND topic = ?", tenant, probeTopic) + require.NoError(t, err) + } + return tenantsByShard +} + +func assertTenantOnOnlyShard( + t *testing.T, + ctx context.Context, + shards map[string]*sql.DB, + table string, + tenant string, + expectedShard string, +) { + t.Helper() + for shard, db := range shards { + var count int + err := db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+table+" WHERE tenant = ?", tenant).Scan(&count) + require.NoError(t, err) + if shard == expectedShard { + require.Positive(t, count, "%s should contain %s on shard %s", table, tenant, shard) + } else { + require.Zero(t, count, "%s should not contain %s on shard %s", table, tenant, shard) + } + } +} + +func assertTopicOnOnlyShard( + t *testing.T, + ctx context.Context, + shards map[string]*sql.DB, + tenant string, + topic string, + expectedShard string, +) { + t.Helper() + for shard, db := range shards { + count := tenantTopicCount(t, ctx, db, tenant, topic) + if shard == expectedShard { + require.Equal(t, 1, count) + } else { + require.Zero(t, count) + } + } +} + +func tenantTopicCount(t *testing.T, ctx context.Context, db *sql.DB, tenant string, topic string) int { + t.Helper() + var count int + err := db.QueryRowContext( + ctx, + "SELECT COUNT(*) FROM queue_messages WHERE tenant = ? AND topic = ?", + tenant, + topic, + ).Scan(&count) + require.NoError(t, err) + return count +} + +func mapKeys(deliveries map[string]extqueue.Delivery) []string { + keys := make([]string, 0, len(deliveries)) + for tenant := range deliveries { + keys = append(keys, tenant) + } + return keys +}