From 66537f64c0e1ea99f05bd40071bd13a2fac87742 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 6 Aug 2026 09:00:03 -0700 Subject: [PATCH 1/2] feat(messagequeue): split partition discovery cadence from polling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Partition discovery, lease acquisition, and worker reconciliation all run on the message-poll ticker (`PollIntervalMs`, 100ms default). Discovery drives topic-wide work every tick — a `DISTINCT partition_key` scan, an active-subscriber read, self-lease reads, and lease-acquisition probes — whose query volume multiplies with subscribers × topics at 10x/sec, even though its outcome only changes when membership or the partition set changes. Message polling needs 100ms latency; discovery does not. ### What? New `PartitionDiscoveryIntervalMs` subscription config (default 1s) drives the supervisor's discovery ticker; `PollIntervalMs` continues to drive per-partition message polling unchanged. This cuts discovery-driven query volume ~10x at the default settings. The accepted trade-off is that a brand-new partition's first message now waits up to the discovery interval (1s) before a worker picks it up; messages on already-owned partitions are unaffected. Integration test configs pin discovery to 100ms so lease-handoff and rebalance convergence assertions stay fast. ## Test Plan - ✅ Full Docker integration suite (`bazel test //test/integration/extension/messagequeue/...`) — discovery-latency-sensitive tests (empty-topic wake-up, rebalance convergence, crash recovery) pass with the new cadence. --- .../extension/messagequeue/mysql/README.md | 1 + .../messagequeue/mysql/subscriber.go | 2 +- .../messagequeue/subscription_config.go | 23 +++++++++----- .../messagequeue/mysql/queue_test.go | 30 +++++++++++++++++++ 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/platform/extension/messagequeue/mysql/README.md b/platform/extension/messagequeue/mysql/README.md index a19765e4..c9fa8c69 100644 --- a/platform/extension/messagequeue/mysql/README.md +++ b/platform/extension/messagequeue/mysql/README.md @@ -69,6 +69,7 @@ subConfig.DLQ.TopicSuffix = "_dlq" // DLQ topic suffix | `SubscriberName` | Unique worker identifier for partition leasing (e.g., hostname, pod name) | | `ConsumerGroup` | Consumer group for independent offset tracking | | `PollIntervalMs` | How often to poll for new messages | +| `PartitionDiscoveryIntervalMs` | How often to discover partitions, attempt lease acquisition, and reconcile workers | | `BatchSize` | Maximum messages to fetch per poll. Set to `1` for strict serialization | | `VisibilityTimeoutMs` | How long messages are invisible after fetch. Must exceed max processing time for `BatchSize=1` | | `LeaseRenewalIntervalMs` | How often to renew partition leases | diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index 10d9f497..6277c7fe 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -479,7 +479,7 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { "subscriber_name", cfg.SubscriberName, } - discoveryTicker := time.NewTicker(time.Duration(cfg.PollIntervalMs) * time.Millisecond) + discoveryTicker := time.NewTicker(time.Duration(cfg.PartitionDiscoveryIntervalMs) * time.Millisecond) defer discoveryTicker.Stop() leaseTicker := time.NewTicker(time.Duration(cfg.LeaseRenewalIntervalMs) * time.Millisecond) diff --git a/platform/extension/messagequeue/subscription_config.go b/platform/extension/messagequeue/subscription_config.go index fad4a7cb..b23c884c 100644 --- a/platform/extension/messagequeue/subscription_config.go +++ b/platform/extension/messagequeue/subscription_config.go @@ -30,6 +30,14 @@ type SubscriptionConfig struct { // PollIntervalMs is how often to poll for new messages (in milliseconds). PollIntervalMs int64 + // PartitionDiscoveryIntervalMs is how often to discover partitions, + // attempt lease acquisition, and reconcile partition workers (in + // milliseconds). Separate from PollIntervalMs: message polling needs low + // latency, while discovery drives topic-wide queries whose volume + // multiplies with subscribers and topics and whose outcome only changes + // on membership or partition changes. + PartitionDiscoveryIntervalMs int64 + // BatchSize is the maximum number of messages to fetch per poll. BatchSize int @@ -98,13 +106,14 @@ func DLQSubscriptionConfig(subscriberName, consumerGroup string) SubscriptionCon // DefaultSubscriptionConfig returns a SubscriptionConfig with sensible defaults. func DefaultSubscriptionConfig(subscriberName, consumerGroup string) SubscriptionConfig { return SubscriptionConfig{ - SubscriberName: subscriberName, - ConsumerGroup: consumerGroup, - PollIntervalMs: 100, // 100ms - BatchSize: 10, - VisibilityTimeoutMs: 60000, // 60s - LeaseRenewalIntervalMs: 10000, // 10s - LeaseDurationMs: 30000, // 30s + SubscriberName: subscriberName, + ConsumerGroup: consumerGroup, + PollIntervalMs: 100, // 100ms + PartitionDiscoveryIntervalMs: 1000, // 1s + BatchSize: 10, + VisibilityTimeoutMs: 60000, // 60s + LeaseRenewalIntervalMs: 10000, // 10s + LeaseDurationMs: 30000, // 30s Retry: RetryConfig{ MaxAttempts: 3, InitialBackoffMs: 1000, // 1s diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index 1c1d6c24..310c1805 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -102,8 +102,12 @@ func (s *SQLQueueIntegrationSuite) TearDownSuite() { // 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 // subscriber can't find invisible messages until the DB timeout expires. +// Partition discovery is likewise pinned to 100ms so initial lease +// acquisition and rebalance convergence stay fast under the 1s production +// default. func testSubConfig(subscriberName, consumerGroup string) extqueue.SubscriptionConfig { cfg := extqueue.DefaultSubscriptionConfig(subscriberName, consumerGroup) + cfg.PartitionDiscoveryIntervalMs = 100 cfg.VisibilityTimeoutMs = 2000 cfg.LeaseDurationMs = 3000 cfg.LeaseRenewalIntervalMs = 1000 @@ -344,6 +348,7 @@ func (s *SQLQueueIntegrationSuite) TestPublishAndSubscribe() { // Subscribe first with config subConfig := extqueue.DefaultSubscriptionConfig("test-worker-1", "test-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -414,6 +419,7 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPerPartitionIsolation() { // Subscribe with short poll interval for fast test subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "isolation-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -481,6 +487,7 @@ func (s *SQLQueueIntegrationSuite) TestSubscriberPartitionOrderPreserved() { // Subscribe and receive all subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "order-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -521,6 +528,7 @@ func (s *SQLQueueIntegrationSuite) TestMultiplePartitions() { // Subscribe subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "multi-partition-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -650,6 +658,7 @@ func (s *SQLQueueIntegrationSuite) TestIdempotentPublish() { // Subscribe subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "idempotent-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -696,6 +705,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentPublishers() { // Subscribe subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "concurrent-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -825,10 +835,12 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroups() { // Subscribe both groups subConfig1 := extqueue.DefaultSubscriptionConfig("worker-1", "group-A") + subConfig1.PartitionDiscoveryIntervalMs = 100 deliveryChan1, err := subscriber1.Subscribe(s.ctx, topic, subConfig1) require.NoError(t, err) subConfig2 := extqueue.DefaultSubscriptionConfig("worker-1", "group-B") + subConfig2.PartitionDiscoveryIntervalMs = 100 deliveryChan2, err := subscriber2.Subscribe(s.ctx, topic, subConfig2) require.NoError(t, err) @@ -904,10 +916,12 @@ func (s *SQLQueueIntegrationSuite) TestMultipleWorkersInConsumerGroup() { // Subscribe both workers subConfig1 := extqueue.DefaultSubscriptionConfig("worker-1", consumerGroup) + subConfig1.PartitionDiscoveryIntervalMs = 100 deliveryChan1, err := subscriber1.Subscribe(s.ctx, topic, subConfig1) require.NoError(t, err) subConfig2 := extqueue.DefaultSubscriptionConfig("worker-2", consumerGroup) + subConfig2.PartitionDiscoveryIntervalMs = 100 deliveryChan2, err := subscriber2.Subscribe(s.ctx, topic, subConfig2) require.NoError(t, err) @@ -979,6 +993,7 @@ func (s *SQLQueueIntegrationSuite) TestConcurrentSubscribers() { subscriber := q.Subscriber() subConfig := extqueue.DefaultSubscriptionConfig(fmt.Sprintf("worker-%d", i), consumerGroup) + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) deliveryChans = append(deliveryChans, deliveryChan) @@ -1078,6 +1093,7 @@ func (s *SQLQueueIntegrationSuite) TestDeadLetterQueue() { t.Logf("Subscribing to DLQ topic: %s", dlqTopic) dlqConfig := extqueue.DefaultSubscriptionConfig("worker-1", "dlq-consumer") + dlqConfig.PartitionDiscoveryIntervalMs = 100 dlqDeliveryChan, err := subscriber.Subscribe(s.ctx, dlqTopic, dlqConfig) require.NoError(t, err) @@ -1129,6 +1145,7 @@ func (s *SQLQueueIntegrationSuite) TestMessageOrderingWithinPartition() { // Subscribe first subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "ordering-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -1191,6 +1208,7 @@ func (s *SQLQueueIntegrationSuite) TestLateSubscriber() { // Now subscribe (late subscriber) subscriber := q.Subscriber() subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "late-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) t.Logf("Late subscriber joined after messages published") @@ -1232,6 +1250,7 @@ func (s *SQLQueueIntegrationSuite) TestEmptyTopicSubscribe() { // Subscribe to empty topic (no messages published yet) subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "empty-consumer") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 // 100 milliseconds deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -1490,6 +1509,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ConsumerLagAfterPartialAck() { // Subscribe and ack only 2 subConfig := extqueue.DefaultSubscriptionConfig("worker-1", consumerGroup) + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -1539,6 +1559,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_LeasesAndOffsets() { require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("lo-1", []byte("a"), "p1", nil))) subConfig := extqueue.DefaultSubscriptionConfig("admin-worker-1", consumerGroup) + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -1615,6 +1636,7 @@ func (s *SQLQueueIntegrationSuite) TestAdmin_ResetOffsetAndReleaseLease() { require.NoError(t, publisher.Publish(s.ctx, topic, entityqueue.NewMessage("r1", []byte("a"), "rp1", nil))) subConfig := extqueue.DefaultSubscriptionConfig("reset-worker", consumerGroup) + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 deliveryChan, err := subscriber.Subscribe(s.ctx, topic, subConfig) require.NoError(t, err) @@ -2115,6 +2137,7 @@ func (s *SQLQueueIntegrationSuite) TestInFlightMessageDoesNotBlockOtherMessages( // Subscribe with batch=10 to fetch multiple messages per poll. The default // 60s visibility timeout keeps msg-1 invisible for the whole test. subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "nack-nb-cg") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 50 subConfig.BatchSize = 10 deliveryCh, err := q.Subscriber().Subscribe(s.ctx, topic, subConfig) @@ -2172,6 +2195,7 @@ func (s *SQLQueueIntegrationSuite) TestPostponeBlocksPartitionUntilDue() { // Subscribe with batch=10 so the barrier — not the batch size — is what // keeps later offsets back. subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "postpone-cg") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 50 subConfig.BatchSize = 10 deliveryCh, err := q.Subscriber().Subscribe(s.ctx, topic, subConfig) @@ -2268,6 +2292,7 @@ func (s *SQLQueueIntegrationSuite) TestPostponeResetsRetryBudget() { dlqTopic := topic + subConfig.DLQ.TopicSuffix dlqConfig := extqueue.DefaultSubscriptionConfig("worker-1", "postpone-budget-cg") + dlqConfig.PartitionDiscoveryIntervalMs = 100 dlqDeliveryChan, err := q.Subscriber().Subscribe(s.ctx, dlqTopic, dlqConfig) require.NoError(t, err) @@ -2299,6 +2324,7 @@ func (s *SQLQueueIntegrationSuite) TestBatchSizeOneStrictSerialization() { // Subscribe with batchSize=1 for strict serialization subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "serial-cg") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 50 subConfig.BatchSize = 1 deliveryCh, err := q.Subscriber().Subscribe(s.ctx, topic, subConfig) @@ -2346,8 +2372,10 @@ func (s *SQLQueueIntegrationSuite) TestMultipleConsumerGroupsIndependentState() // Two consumer groups subscribing to the same topic cfg1 := extqueue.DefaultSubscriptionConfig("worker-1", "cg-alpha") + cfg1.PartitionDiscoveryIntervalMs = 100 cfg1.PollIntervalMs = 50 cfg2 := extqueue.DefaultSubscriptionConfig("worker-2", "cg-beta") + cfg2.PartitionDiscoveryIntervalMs = 100 cfg2.PollIntervalMs = 50 ch1, err := q.Subscriber().Subscribe(s.ctx, topic, cfg1) @@ -2480,6 +2508,7 @@ func (s *SQLQueueIntegrationSuite) TestCrashAfterRejectDoesNotLoseMessages() { // Verify DLQ contains msg-B dlqTopic := topic + subConfig.DLQ.TopicSuffix dlqConfig := extqueue.DefaultSubscriptionConfig("worker-2", "crash-reject-cg") + dlqConfig.PartitionDiscoveryIntervalMs = 100 dlqConfig.PollIntervalMs = 100 dlqChan, err := q2.Subscriber().Subscribe(s.ctx, dlqTopic, dlqConfig) require.NoError(t, err) @@ -2637,6 +2666,7 @@ func (s *SQLQueueIntegrationSuite) TestWatermarkAdvancesContiguously() { } subConfig := extqueue.DefaultSubscriptionConfig("worker-1", "watermark-cg") + subConfig.PartitionDiscoveryIntervalMs = 100 subConfig.PollIntervalMs = 100 subConfig.VisibilityTimeoutMs = 30000 // long visibility so nothing re-delivers subConfig.BatchSize = 10 From d27bb4dc36714d95793e8c72626f8a16e57dc63f Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 6 Aug 2026 09:06:30 -0700 Subject: [PATCH 2/2] perf(messagequeue): lease-aware partition acquisition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Every discovery tick, `DiscoverAndAcquirePartitions` attempts `TryAcquireLease` on every discovered partition — including partitions this subscriber already owns and partitions validly leased by peers. Each attempt is an `INSERT … ON DUPLICATE KEY UPDATE` (a write that takes a row lock even when the steal condition fails) plus an ownership SELECT. At steady state that is ~2 × P×(N−1)/N queries per subscriber per tick of contended writes that can never win, concentrated on the same lease rows from every node in the group. ### What? `partitionLeaseStore` gains `GetAllLeases` — a single PK-prefix read of every lease row for `(consumer_group, topic)`. `DiscoverAndAcquirePartitions` uses it to classify discovered partitions: self-owned rows count toward the cap but are not re-probed (renewal is the lease tick's job), rows validly held by another subscriber are skipped without any write, and only unleased or stale (stealable) partitions are attempted. The classification is advisory — `TryAcquireLease` remains the atomic arbiter, so races between the read and the write resolve exactly as before. The read also replaces the previous pre-loop `GetLeasedPartitions` query. A `lease_aware_skipped` counter records how many probes each tick avoids. ## Test Plan - ✅ Full Docker integration suite — rebalance, crash-recovery (stale-lease steal), and orphan-sweep tests exercise the new classification against real MySQL. --- .../messagequeue/mysql/mock_stores.go | 15 ++ .../mysql/partition_lease_store.go | 93 +++++-- .../mysql/partition_lease_store_test.go | 249 +++++++++++------- .../extension/messagequeue/mysql/stores.go | 17 ++ 4 files changed, 256 insertions(+), 118 deletions(-) diff --git a/platform/extension/messagequeue/mysql/mock_stores.go b/platform/extension/messagequeue/mysql/mock_stores.go index f8d7d105..28c446bf 100644 --- a/platform/extension/messagequeue/mysql/mock_stores.go +++ b/platform/extension/messagequeue/mysql/mock_stores.go @@ -251,6 +251,21 @@ func (mr *MockpartitionLeaseStoreMockRecorder) DiscoverAndAcquirePartitions(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverAndAcquirePartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).DiscoverAndAcquirePartitions), ctx, topic, subscriberName, consumerGroup, leaseDurationMs, maxPartitions) } +// GetAllLeases mocks base method. +func (m *MockpartitionLeaseStore) GetAllLeases(ctx context.Context, topic, consumerGroup string) ([]leaseInfo, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAllLeases", ctx, 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 { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllLeases", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetAllLeases), ctx, topic, consumerGroup) +} + // GetLeasedPartitions mocks base method. func (m *MockpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic, subscriberName, consumerGroup string) ([]string, error) { m.ctrl.T.Helper() diff --git a/platform/extension/messagequeue/mysql/partition_lease_store.go b/platform/extension/messagequeue/mysql/partition_lease_store.go index 6e2fb227..8d8287c6 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store.go @@ -196,9 +196,51 @@ func (s *sqlpartitionLeaseStore) GetLeasedPartitions(ctx context.Context, topic return partitions, nil } +// 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) { + 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) + + if err != nil { + return nil, fmt.Errorf("get all leases topic=%s: %w", topic, err) + } + defer rows.Close() + + var leases []leaseInfo + 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) + } + leases = append(leases, lease) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("row iteration topic=%s: %w", topic, err) + } + + return leases, nil +} + // 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. // maxPartitions limits how many total partitions this subscriber can own (0 = unlimited) +// +// Acquisition is lease-aware: one GetAllLeases read classifies every +// discovered partition, and TryAcquireLease is attempted only for partitions +// that are unleased or whose lease is stale (stealable). Partitions already +// owned by this subscriber are counted against the cap but not re-probed +// (renewal is the lease tick's job), and partitions validly held by another +// subscriber are skipped entirely — probing them is a guaranteed-futile +// 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) { op := metrics.Begin(s.scope, "discover_and_acquire", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) defer func() { op.Complete(retErr) }() @@ -234,28 +276,44 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex "count", len(partitions), ) - // Query owned partitions once before the loop to avoid N+1 queries. - // Build a set of already-owned partition keys so we can distinguish - // re-acquiring an already-owned partition from acquiring a new one. + // 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) + if err != nil { + return 0, nil, fmt.Errorf("get all leases for acquisition topic=%s: %w", topic, err) + } + staleThreshold := currentTimeMillis() - leaseDurationMs ownedCount := 0 ownedSet := make(map[string]struct{}) - if maxPartitions > 0 { - owned, err := s.GetLeasedPartitions(ctx, topic, subscriberName, consumerGroup) - if err != nil { - return 0, nil, fmt.Errorf("get owned partitions for cap check topic=%s: %w", topic, err) - } - ownedCount = len(owned) - for _, pk := range owned { - ownedSet[pk] = struct{}{} + heldByOther := make(map[string]struct{}) + for _, lease := range allLeases { + switch { + case lease.LeasedBy == subscriberName: + // Self-owned, fresh or stale: a stale self-lease means our own + // renewals are lagging, not that ownership moved. + ownedSet[lease.PartitionKey] = struct{}{} + ownedCount++ + case lease.LeaseRenewedAt >= staleThreshold: + heldByOther[lease.PartitionKey] = struct{}{} } } // Sort partitions deterministically sort.Strings(partitions) - // Try to acquire leases for discovered partitions + // Try to acquire leases for unleased or stale discovered partitions acquiredCount := 0 + skippedCount := 0 for _, partitionKey := range partitions { + if _, owned := ownedSet[partitionKey]; owned { + continue + } + if _, held := heldByOther[partitionKey]; held { + skippedCount++ + continue + } + // Enforce maxPartitions cap using local count if maxPartitions > 0 && ownedCount >= maxPartitions { s.logger.Debugw("reached max partitions cap, stopping acquisition", @@ -279,22 +337,19 @@ func (s *sqlpartitionLeaseStore) DiscoverAndAcquirePartitions(ctx context.Contex continue } if acquired { - // Only count as newly acquired if not already owned. - // TryAcquireLease returns true for already-owned partitions (renew), - // so we must not double-count them against the maxPartitions cap. - if _, alreadyOwned := ownedSet[partitionKey]; !alreadyOwned { - acquiredCount++ - ownedCount++ - } + acquiredCount++ + ownedCount++ } } metrics.NamedCounter(s.scope, "discover_and_acquire", "partitions_discovered", int64(len(partitions)), metrics.NewTag("topic", topic)) 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", logTopic, topic, "discovered_count", len(partitions), "acquired_count", acquiredCount, + "skipped_held_by_other", skippedCount, ) return acquiredCount, partitions, nil diff --git a/platform/extension/messagequeue/mysql/partition_lease_store_test.go b/platform/extension/messagequeue/mysql/partition_lease_store_test.go index 6d405833..386fc094 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store_test.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store_test.go @@ -18,6 +18,7 @@ import ( "context" "database/sql" "testing" + "time" "github.com/DATA-DOG/go-sqlmock" "github.com/stretchr/testify/require" @@ -216,37 +217,124 @@ func TestPartitionLeaseStore_GetLeasedPartitions(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestPartitionLeaseStore_GetAllLeases(t *testing.T) { + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + want []leaseInfo + }{ + { + name: "returns leases held by any subscriber", + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"partition_key", "leased_by", "lease_renewed_at"}). + 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"). + WillReturnRows(rows) + }, + want: []leaseInfo{ + {PartitionKey: "part1", LeasedBy: testSubscriberName, LeaseRenewedAt: 1000}, + {PartitionKey: "part2", LeasedBy: "other-worker", LeaseRenewedAt: 2000}, + }, + }, + { + 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"). + WillReturnRows(sqlmock.NewRows([]string{"partition_key", "leased_by", "lease_renewed_at"})) + }, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setuppartitionLeaseStoreTest(t) + defer db.Close() + + tt.setup(mock) + + leases, err := store.GetAllLeases(context.Background(), "test_topic", testConsumerGroup) + require.NoError(t, err) + require.Equal(t, tt.want, leases) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { + leaseColumns := []string{"partition_key", "leased_by", "lease_renewed_at"} + freshMs := time.Now().UnixMilli() + staleMs := freshMs - testLeaseDurationMs - 60_000 + + // expectDiscover mocks the DISTINCT partition scan. + expectDiscover := func(mock sqlmock.Sqlmock, partitions ...string) { + rows := sqlmock.NewRows([]string{"partition_key"}) + for _, pk := range partitions { + rows.AddRow(pk) + } + mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). + WithArgs("test_topic"). + WillReturnRows(rows) + } + + // expectAcquire mocks one TryAcquireLease attempt whose ownership check + // reports the given owner. + expectAcquire := func(mock sqlmock.Sqlmock, owner string) { + mock.ExpectExec("INSERT INTO queue_partition_leases"). + WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). + WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(owner)) + } + tests := []struct { name string maxPartitions int setup func(mock sqlmock.Sqlmock) wantAcquired int - wantErr bool }{ { - name: "unlimited - acquires all available", + name: "acquires unleased, skips fresh lease held by other", maxPartitions: 0, setup: func(mock sqlmock.Sqlmock) { - // Discover partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Acquire part1 - success - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // Acquire part2 - taken by other worker - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow("other-worker")) + expectDiscover(mock, "part1", "part2") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part2", "other-worker", freshMs)) + // Only unleased part1 is attempted; part2's fresh lease is + // never write-probed. + expectAcquire(mock, testSubscriberName) + }, + wantAcquired: 1, + }, + { + name: "stale lease held by other is stealable", + maxPartitions: 0, + 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"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part1", "other-worker", staleMs)) + expectAcquire(mock, testSubscriberName) + }, + wantAcquired: 1, + }, + { + name: "self-owned partitions are not re-probed", + maxPartitions: 0, + 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"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("part1", testSubscriberName, freshMs)) + // Only part2 is attempted; renewal of part1 is the lease + // tick's job. + expectAcquire(mock, testSubscriberName) }, wantAcquired: 1, }, @@ -254,33 +342,13 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { name: "stops acquiring when cap reached", maxPartitions: 2, setup: func(mock sqlmock.Sqlmock) { - // Discover 3 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2"). - AddRow("part3") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: owns 0 partitions - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"})) - - // Acquire part1 - success - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // Acquire part2 - success (now at cap of 2, stops) - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // part3 is never attempted because ownedCount (2) >= maxPartitions (2) + expectDiscover(mock, "part1", "part2", "part3") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns)) + // part1 and part2 acquired; part3 never attempted at the cap. + expectAcquire(mock, testSubscriberName) + expectAcquire(mock, testSubscriberName) }, wantAcquired: 2, }, @@ -288,52 +356,42 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { name: "pre-owned partitions count toward cap", maxPartitions: 3, setup: func(mock sqlmock.Sqlmock) { - // Discover 3 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2"). - AddRow("part3") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: already owns 2 partitions - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"}). - AddRow("existing1"). - AddRow("existing2")) - - // Acquire part1 - success (now at 3, cap reached) - mock.ExpectExec("INSERT INTO queue_partition_leases"). - WillReturnResult(sqlmock.NewResult(1, 1)) - mock.ExpectQuery("SELECT leased_by FROM queue_partition_leases"). - WillReturnRows(sqlmock.NewRows([]string{"leased_by"}).AddRow(testSubscriberName)) - - // part2, part3 never attempted because ownedCount (3) >= maxPartitions (3) + expectDiscover(mock, "part1", "part2", "part3") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("existing1", testSubscriberName, freshMs). + AddRow("existing2", testSubscriberName, freshMs)) + // One acquisition reaches the cap of 3; part2/part3 skipped. + expectAcquire(mock, testSubscriberName) }, wantAcquired: 1, }, { - name: "already at cap - acquires nothing", + name: "already at cap acquires nothing", maxPartitions: 2, setup: func(mock sqlmock.Sqlmock) { - // Discover 2 partitions - rows := sqlmock.NewRows([]string{"partition_key"}). - AddRow("part1"). - AddRow("part2") - mock.ExpectQuery("SELECT DISTINCT partition_key FROM queue_messages"). - WithArgs("test_topic"). - WillReturnRows(rows) - - // Pre-loop GetLeasedPartitions: already owns 2 partitions (at cap) - mock.ExpectQuery("SELECT partition_key FROM queue_partition_leases"). - WithArgs(testConsumerGroup, "test_topic", testSubscriberName). - WillReturnRows(sqlmock.NewRows([]string{"partition_key"}). - AddRow("existing1"). - AddRow("existing2")) - - // No acquire attempts - immediately breaks + expectDiscover(mock, "part1", "part2") + mock.ExpectQuery("SELECT partition_key, leased_by, lease_renewed_at FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic"). + WillReturnRows(sqlmock.NewRows(leaseColumns). + AddRow("existing1", testSubscriberName, freshMs). + AddRow("existing2", testSubscriberName, freshMs)) + // No acquire attempts. + }, + wantAcquired: 0, + }, + { + name: "lost race counts nothing", + maxPartitions: 0, + 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"). + WillReturnRows(sqlmock.NewRows(leaseColumns)) + // Attempted while unleased, but another subscriber won the + // atomic acquire between the read and the write. + expectAcquire(mock, "other-worker") }, wantAcquired: 0, }, @@ -344,19 +402,12 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { db, mock, store := setuppartitionLeaseStoreTest(t) defer db.Close() - ctx := context.Background() - topic := "test_topic" - tt.setup(mock) - acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(ctx, topic, testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) - if tt.wantErr { - require.Error(t, err) - } else { - require.NoError(t, err) - require.Equal(t, tt.wantAcquired, acquired) - require.NotNil(t, discoveredPartitions) - } + acquired, discoveredPartitions, err := store.DiscoverAndAcquirePartitions(context.Background(), "test_topic", testSubscriberName, testConsumerGroup, testLeaseDurationMs, tt.maxPartitions) + require.NoError(t, err) + require.Equal(t, tt.wantAcquired, acquired) + require.NotNil(t, discoveredPartitions) require.NoError(t, mock.ExpectationsWereMet()) }) } diff --git a/platform/extension/messagequeue/mysql/stores.go b/platform/extension/messagequeue/mysql/stores.go index 8c39070f..9b28c26d 100644 --- a/platform/extension/messagequeue/mysql/stores.go +++ b/platform/extension/messagequeue/mysql/stores.go @@ -103,6 +103,17 @@ type offsetStore interface { GetMinAckedOffset(ctx context.Context, topic string, partitionKey string) (offset int64, found bool, err error) } +// leaseInfo describes one partition's current lease row (internal use only) +type leaseInfo struct { + // PartitionKey is the partition this lease covers + PartitionKey string + // LeasedBy is the subscriber name currently holding the lease + LeasedBy string + // LeaseRenewedAt is the epoch milliseconds of the last renewal; a lease + // is stale (stealable) once this is older than the lease duration + LeaseRenewedAt int64 +} + // partitionLeaseStore handles partition lease operations (internal use only) type partitionLeaseStore interface { // TryAcquireLease attempts to acquire or renew a lease for a partition @@ -120,6 +131,12 @@ type partitionLeaseStore interface { // GetLeasedPartitions returns all partitions currently leased by this worker GetLeasedPartitions(ctx context.Context, 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) + // 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)