From 7436364a356586f2b9c319f012da000ab82649b8 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 6 Aug 2026 09:35:10 -0700 Subject: [PATCH] fix(messagequeue): release leases on drained partitions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? A lease is held forever once acquired: `ReleaseLease` is only called on shutdown and fair-share shedding, renewal is unconditional, and workers are reconciled from the lease set — so when a partition drains (every message acked and garbage-collected, hence absent from discovery), its holder keeps a goroutine polling an empty partition at ~4-5 queries per poll tick, permanently, plus a lease row and a `queue_offsets` row that nothing ever deletes. On topics with bounded partition keys (queue names) this is harmless stickiness; on topics with unbounded short-lived keys it is a leak proportional to cumulative traffic: buildsignal partitions per batch ID, and the gateway log and cancel topics partition per request ID — every request ever processed would leave a ghost worker behind. The crash variant is worse: a stale lease on a drained partition is never even stealable, because acquisition only probes discovered partitions. ### What? - Idle-lease release: the discovery tick tracks, per owned partition absent from discovery, when it was first observed drained (a partition with any in-flight, postponed, or unacked message still has stored rows and is always discovered — only fully-consumed partitions qualify). After a grace of 2x `LeaseDurationMs` (60s at defaults) the subscriber deletes its consumer group's offsets row (while still holding the lease, so no concurrent initialization is possible), releases the lease, and reconciliation stops the worker. If a message arrives later, the partition reappears in discovery and is reacquired like any new partition — `Initialize` recreates the offsets row, and since drained meant zero stored rows there is nothing to replay. - Stale-lease purge: the lease tick deletes lease rows not renewed within 10x `LeaseDurationMs`, covering holders that crashed while owning a drained partition. Deleting a stale row is equivalent to expiry — a concurrent renewal refreshes the row and the age predicate skips it. - Rebalance integration tests pin `Retry.MaxAttempts` high: they publish messages they never ack, and dead-lettering mid-test would now drain the partitions and dissolve the lease distribution their assertions wait on. ## Test Plan - ✅ New integration test `TestIdleLeaseRelease` covers the full lifecycle against real MySQL: consume, wait out GC + grace, assert the lease and offsets rows are gone, then republish to the same partition key and assert delivery resumes through normal discovery. --- .../messagequeue/mysql/mock_stores.go | 28 ++++ .../messagequeue/mysql/offset_store.go | 20 +++ .../messagequeue/mysql/offset_store_test.go | 51 +++++++ .../mysql/partition_lease_store.go | 30 ++++ .../mysql/partition_lease_store_test.go | 52 +++++++ .../extension/messagequeue/mysql/stores.go | 14 ++ .../messagequeue/mysql/subscriber.go | 133 ++++++++++++++++++ .../messagequeue/mysql/subscriber_test.go | 62 ++++++++ .../messagequeue/mysql/queue_test.go | 91 ++++++++++-- 9 files changed, 471 insertions(+), 10 deletions(-) diff --git a/platform/extension/messagequeue/mysql/mock_stores.go b/platform/extension/messagequeue/mysql/mock_stores.go index 87417b0c..a411e53c 100644 --- a/platform/extension/messagequeue/mysql/mock_stores.go +++ b/platform/extension/messagequeue/mysql/mock_stores.go @@ -152,6 +152,20 @@ func (m *MockoffsetStore) EXPECT() *MockoffsetStoreMockRecorder { return m.recorder } +// DeleteOffset mocks base method. +func (m *MockoffsetStore) DeleteOffset(ctx context.Context, topic, partitionKey, consumerGroup string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteOffset", ctx, 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 { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOffset", reflect.TypeOf((*MockoffsetStore)(nil).DeleteOffset), ctx, topic, partitionKey, consumerGroup) +} + // GetAckedOffset mocks base method. func (m *MockoffsetStore) GetAckedOffset(ctx context.Context, topic, partitionKey, consumerGroup string) (int64, error) { m.ctrl.T.Helper() @@ -281,6 +295,20 @@ func (mr *MockpartitionLeaseStoreMockRecorder) GetLeasedPartitions(ctx, topic, s return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLeasedPartitions", reflect.TypeOf((*MockpartitionLeaseStore)(nil).GetLeasedPartitions), ctx, topic, subscriberName, consumerGroup) } +// PurgeStale mocks base method. +func (m *MockpartitionLeaseStore) PurgeStale(ctx context.Context, topic, consumerGroup string, olderThanMs int64) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PurgeStale", ctx, 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 { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PurgeStale", reflect.TypeOf((*MockpartitionLeaseStore)(nil).PurgeStale), ctx, topic, consumerGroup, olderThanMs) +} + // ReleaseLease mocks base method. func (m *MockpartitionLeaseStore) ReleaseLease(ctx context.Context, topic, partitionKey, subscriberName, consumerGroup string) error { m.ctrl.T.Helper() diff --git a/platform/extension/messagequeue/mysql/offset_store.go b/platform/extension/messagequeue/mysql/offset_store.go index 60963a81..f26b3fb1 100644 --- a/platform/extension/messagequeue/mysql/offset_store.go +++ b/platform/extension/messagequeue/mysql/offset_store.go @@ -132,3 +132,23 @@ func (s *sqloffsetStore) GetMinAckedOffset(ctx context.Context, topic string, pa return minOffset, true, nil } + +// 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) { + op := metrics.Begin(s.scope, "delete_offset", metrics.StorageLatencyBuckets, + metrics.NewTag("topic", topic), + metrics.NewTag("partition_key", partitionKey), + 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) + + if err != nil { + return fmt.Errorf("delete offset topic=%s partition=%s: %w", 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 c145c7c5..112067c8 100644 --- a/platform/extension/messagequeue/mysql/offset_store_test.go +++ b/platform/extension/messagequeue/mysql/offset_store_test.go @@ -188,3 +188,54 @@ func TestOffsetStore_GetMinAckedOffset(t *testing.T) { }) } } + +func TestOffsetStore_DeleteOffset(t *testing.T) { + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + wantErr bool + }{ + { + 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"). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "idempotent - row already gone", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_offsets"). + WithArgs(testConsumerGroup, "test_topic", "part-1"). + WillReturnResult(sqlmock.NewResult(0, 0)) + }, + }, + { + name: "database error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_offsets"). + WithArgs(testConsumerGroup, "test_topic", "part-1"). + WillReturnError(fmt.Errorf("db error")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupoffsetStoreTest(t) + defer db.Close() + + tt.setup(mock) + + err := store.DeleteOffset(context.Background(), "test_topic", "part-1", testConsumerGroup) + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/platform/extension/messagequeue/mysql/partition_lease_store.go b/platform/extension/messagequeue/mysql/partition_lease_store.go index 8d8287c6..d13bcb96 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store.go @@ -228,6 +228,36 @@ func (s *sqlpartitionLeaseStore) GetAllLeases(ctx context.Context, topic string, return leases, nil } +// 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) { + op := metrics.Begin(s.scope, "purge_stale", metrics.StorageLatencyBuckets, metrics.NewTag("topic", topic)) + defer func() { op.Complete(retErr) }() + + threshold := currentTimeMillis() - olderThanMs + + result, err := s.db.ExecContext(ctx, fmt.Sprintf(` + DELETE FROM %s + WHERE consumer_group = ? AND topic = ? AND lease_renewed_at < ? + `, PartitionLeasesTableName), consumerGroup, topic, threshold) + + if err != nil { + return fmt.Errorf("failed to purge stale leases: %w", err) + } + + // RowsAffected error is swallowed because the DELETE itself succeeded; + // the count is for observability only. + 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", + logTopic, topic, + "deleted", deleted, + ) + } + + return 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) diff --git a/platform/extension/messagequeue/mysql/partition_lease_store_test.go b/platform/extension/messagequeue/mysql/partition_lease_store_test.go index 386fc094..1a6a83e9 100644 --- a/platform/extension/messagequeue/mysql/partition_lease_store_test.go +++ b/platform/extension/messagequeue/mysql/partition_lease_store_test.go @@ -17,6 +17,7 @@ package mysql import ( "context" "database/sql" + "fmt" "testing" "time" @@ -412,3 +413,54 @@ func TestPartitionLeaseStore_DiscoverAndAcquirePartitions(t *testing.T) { }) } } + +func TestPartitionLeaseStore_PurgeStale(t *testing.T) { + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + wantErr bool + }{ + { + name: "deletes rows older than threshold", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(0, 2)) + }, + }, + { + 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()). + WillReturnResult(sqlmock.NewResult(0, 0)) + }, + }, + { + name: "database error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_partition_leases"). + WithArgs(testConsumerGroup, "test_topic", sqlmock.AnyArg()). + WillReturnError(fmt.Errorf("db error")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setuppartitionLeaseStoreTest(t) + defer db.Close() + + tt.setup(mock) + + err := store.PurgeStale(context.Background(), "test_topic", testConsumerGroup, 300_000) + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/platform/extension/messagequeue/mysql/stores.go b/platform/extension/messagequeue/mysql/stores.go index c739f2fb..d2c43b49 100644 --- a/platform/extension/messagequeue/mysql/stores.go +++ b/platform/extension/messagequeue/mysql/stores.go @@ -101,6 +101,12 @@ type offsetStore interface { // 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) + + // 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 } // leaseInfo describes one partition's current lease row (internal use only) @@ -137,6 +143,14 @@ type partitionLeaseStore interface { // 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 + // 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) diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index 12791588..e0f1cef4 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -61,6 +61,23 @@ const ( // Purging a live-but-stalled subscriber's row is harmless: its next // heartbeat re-inserts it. heartbeatPurgeAfterLeaseDurations = 10 + + // idleLeaseReleaseAfterLeaseDurations sets how long an owned partition + // must stay drained (no stored messages, hence absent from discovery) + // before its lease is released, as a multiple of LeaseDurationMs (2x = + // 60s at defaults). Long enough to keep ownership sticky across brief + // quiet spells on stable partition keys; short enough that short-lived + // keys (one batch, one request) don't hold leases, offset rows, and + // polling workers forever. A released partition rejoins through normal + // discovery as soon as a message arrives. + idleLeaseReleaseAfterLeaseDurations = 2 + + // leasePurgeAfterLeaseDurations sets the age threshold for purging + // abandoned lease rows, as a multiple of LeaseDurationMs. Covers holders + // that crashed while owning a drained partition: acquisition only probes + // discovered partitions, so nothing would ever steal (and thereby + // refresh or remove) a stale lease on a partition with no messages. + leasePurgeAfterLeaseDurations = 10 ) // HookSignal identifies the type of subscriber lifecycle event. @@ -124,6 +141,14 @@ type subscription struct { // DiscoverAndAcquirePartitions call. 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 map[string]time.Time } // partitionWorker handles polling and delivering messages for a single partition. @@ -580,6 +605,13 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { 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)...) + } s.emitSignal(SignalPartitionUpdate) case <-discoveryTicker.C: @@ -650,10 +682,111 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc return fmt.Errorf("get leased partitions after acquire: %w", err) } + // Idle-lease release: an owned partition absent from discovery has no + // stored messages left — everything was consumed and garbage-collected. + // Held past the grace period, such a lease buys nothing (a worker + // polling an empty partition forever) and on topics with short-lived + // partition keys it leaks a lease row, an offsets row, and a goroutine + // per key ever used. Release drops the partition entirely: reconcile + // stops its worker, and if a message arrives later the partition + // reappears in discovery and is reacquired like any new partition. + grace := time.Duration(idleLeaseReleaseAfterLeaseDurations*cfg.LeaseDurationMs) * time.Millisecond + var expired []string + sub.drainedSince, expired = updateDrainedTracking(sub.drainedSince, leasedPartitions, discoveredPartitions, grace, time.Now()) + 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. + s.logger.Errorw("delete offsets for drained partition failed", + "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. + s.logger.Errorw("release lease for drained partition failed", + "topic", sub.topic, + "partition_key", pk, + "error", err, + ) + continue + } + released[pk] = struct{}{} + delete(sub.drainedSince, pk) + + // 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) + + metrics.NamedCounter(s.scope, "idle_lease", "released", 1, metrics.NewTag("topic", sub.topic)) + s.logger.Infow("released idle partition lease", + "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) + } + } + leasedPartitions = kept + } + } + s.reconcilePartitionWorkers(ctx, sub, leasedPartitions) return nil } +// updateDrainedTracking recomputes, for every owned partition absent from +// this tick's discovery, when it was first observed drained. A partition is +// drained only when zero of its messages remain stored — in-flight, +// postponed, and unacked messages all keep rows in the messages table, so a +// partition with any outstanding work is always discovered. Partitions that +// reappear in discovery (or are no longer owned) are dropped from tracking; +// first-seen times carry over so the clock accumulates across ticks. Returns +// the updated tracking map and the partitions drained for at least grace +// (release candidates), sorted for determinism. +func updateDrainedTracking(prev map[string]time.Time, owned []string, discovered []string, grace time.Duration, now time.Time) (map[string]time.Time, []string) { + discoveredSet := make(map[string]struct{}, len(discovered)) + for _, pk := range discovered { + discoveredSet[pk] = struct{}{} + } + + next := make(map[string]time.Time) + var expired []string + for _, pk := range owned { + if _, live := discoveredSet[pk]; live { + continue + } + since, tracked := prev[pk] + if !tracked { + since = now + } + next[pk] = since + if now.Sub(since) >= grace { + expired = append(expired, pk) + } + } + sort.Strings(expired) + return next, expired +} + // reconcilePartitionWorkers diffs the current set of workers against the current // set of leases and starts/stops workers to match. This is the core of the // supervisor's control loop. diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index 1949a645..263d80d9 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -1008,3 +1008,65 @@ func TestSubscriber_RebalanceUnderCapReleasesNothing(t *testing.T) { require.NoError(t, err) assert.Empty(t, released) } + +func TestUpdateDrainedTracking(t *testing.T) { + now := time.UnixMilli(1_000_000) + earlier := now.Add(-time.Minute) + grace := 30 * time.Second + + tests := []struct { + name string + prev map[string]time.Time + owned []string + discovered []string + wantTracked map[string]time.Time + wantExpired []string + }{ + { + name: "owned and discovered is not tracked", + owned: []string{"p1"}, + discovered: []string{"p1"}, + wantTracked: map[string]time.Time{}, + }, + { + name: "freshly drained starts tracking now", + owned: []string{"p1"}, + discovered: nil, + wantTracked: map[string]time.Time{"p1": now}, + }, + { + name: "already tracked keeps original since time", + prev: map[string]time.Time{"p1": now.Add(-10 * time.Second)}, + owned: []string{"p1"}, + wantTracked: map[string]time.Time{"p1": now.Add(-10 * time.Second)}, + }, + { + name: "drained past grace expires sorted", + prev: map[string]time.Time{"p-b": earlier, "p-a": earlier, "p-young": now.Add(-time.Second)}, + owned: []string{"p-b", "p-a", "p-young"}, + wantTracked: map[string]time.Time{"p-a": earlier, "p-b": earlier, "p-young": now.Add(-time.Second)}, + wantExpired: []string{"p-a", "p-b"}, + }, + { + name: "rediscovered partition resets the clock", + prev: map[string]time.Time{"p1": earlier}, + owned: []string{"p1"}, + discovered: []string{"p1"}, + wantTracked: map[string]time.Time{}, + }, + { + name: "no longer owned is dropped from tracking", + prev: map[string]time.Time{"gone": earlier}, + owned: []string{"p1"}, + wantTracked: map[string]time.Time{"p1": now}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tracked, expired := updateDrainedTracking(tt.prev, tt.owned, tt.discovered, grace, now) + assert.Equal(t, tt.wantTracked, tracked) + assert.Equal(t, tt.wantExpired, expired) + }) + } +} diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index 5e9c15d6..740fdf4e 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -114,6 +114,17 @@ func testSubConfig(subscriberName, consumerGroup string) extqueue.SubscriptionCo return cfg } +// rebalanceTestConfig pins a high retry budget on top of testSubConfig. +// Rebalance tests publish messages they never ack; under the default budget +// those messages dead-letter mid-test, draining the partitions — and a +// drained partition's lease is released after the idle grace, dissolving the +// exact lease distribution the assertions wait on. +func rebalanceTestConfig(subscriberName, consumerGroup string) extqueue.SubscriptionConfig { + cfg := testSubConfig(subscriberName, consumerGroup) + cfg.Retry.MaxAttempts = 1000 + return cfg +} + // receive waits for a single delivery. Bazel's test timeout is the safety net // if the delivery never arrives. func receive(t *testing.T, deliveryChan <-chan extqueue.Delivery) extqueue.Delivery { @@ -1732,7 +1743,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_EvenDistribution() { require.NoError(t, err) defer q1.Close() - _, err = q1.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s1", consumerGroup)) + _, err = q1.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s1", consumerGroup)) require.NoError(t, err) waitForCondition(t, signalCh, func() bool { @@ -1748,7 +1759,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_EvenDistribution() { require.NoError(t, err) defer q2.Close() - _, err = q2.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s2", consumerGroup)) + _, err = q2.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s2", consumerGroup)) require.NoError(t, err) waitForCondition(t, signalCh, func() bool { @@ -1795,9 +1806,9 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_SubscriberLeaves() { require.NoError(t, err) // no defer close — we close explicitly below - _, err = q1.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s1", consumerGroup)) + _, err = q1.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s1", consumerGroup)) require.NoError(t, err) - _, err = q2.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s2", consumerGroup)) + _, err = q2.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s2", consumerGroup)) require.NoError(t, err) waitForCondition(t, signalCh, func() bool { @@ -1860,9 +1871,9 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OddPartitions() { require.NoError(t, err) defer q2.Close() - _, err = q1.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s1", consumerGroup)) + _, err = q1.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s1", consumerGroup)) require.NoError(t, err) - _, err = q2.Subscriber().Subscribe(s.ctx, topic, testSubConfig("s2", consumerGroup)) + _, err = q2.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("s2", consumerGroup)) require.NoError(t, err) // maxPart = ceil(5/2) = 3. One gets 3, the other gets 2. @@ -1913,7 +1924,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_NoOrphans() { }) require.NoError(t, err) queues[i] = q - _, err = q.Subscriber().Subscribe(s.ctx, topic, testSubConfig(name, consumerGroup)) + _, err = q.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig(name, consumerGroup)) require.NoError(t, err) } defer queues[0].Close() @@ -1973,7 +1984,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_MoreSubscribersThanPartitions() }) require.NoError(t, err) queues = append(queues, q) - _, err = q.Subscriber().Subscribe(s.ctx, topic, testSubConfig(name, consumerGroup)) + _, err = q.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig(name, consumerGroup)) require.NoError(t, err) } defer func() { @@ -2037,7 +2048,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_NoStarvation_UnevenSplit() { // Nothing is acked in this test; a high retry budget keeps the // messages out of the DLQ so partitions stay discoverable while the // group converges. - cfg := testSubConfig(name, consumerGroup) + cfg := rebalanceTestConfig(name, consumerGroup) cfg.Retry.MaxAttempts = 1000 _, err = q.Subscriber().Subscribe(s.ctx, topic, cfg) require.NoError(t, err) @@ -2108,7 +2119,7 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OrphanSweep() { require.NoError(t, err) } - deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, testSubConfig("worker-real", consumerGroup)) + deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, rebalanceTestConfig("worker-real", consumerGroup)) require.NoError(t, err) // All three messages must arrive: one via the normal capped acquisition, @@ -2128,6 +2139,66 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_OrphanSweep() { t.Logf("Orphan sweep verified: all 3 partitions processed despite a fair-share cap of 1") } +// TestIdleLeaseRelease verifies the full lifecycle of a short-lived partition +// key: consume everything, wait out garbage collection and the idle grace, +// and the subscriber must release the lease and delete its offsets row (no +// ghost worker, no leaked rows). A message published afterwards must +// resurrect the partition through normal discovery and be delivered. +func (s *SQLQueueIntegrationSuite) TestIdleLeaseRelease() { + t := s.T() + + topic := "idle_release_topic" + consumerGroup := "idle-release-cg" + 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, + }) + require.NoError(t, err) + defer q.Close() + + // Fast poll so the idle GC pass (every 100 idle ticks) fires quickly; + // idle grace is 2x the 3s test lease duration. + cfg := testSubConfig("worker-idle", consumerGroup) + cfg.PollIntervalMs = 50 + + deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, cfg) + require.NoError(t, err) + + msg := entityqueue.NewMessage("idle-1", []byte("x"), partition, nil) + require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) + + delivery := receive(t, deliveryChan) + require.Equal(t, "idle-1", delivery.Message().ID) + require.NoError(t, delivery.Ack(s.ctx)) + + // After GC removes the acked message and the drained grace elapses, the + // lease row and this group's offsets row must both be gone. + 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)) + return n + } + waitForCondition(t, signalCh, func() bool { + return rowCount("queue_partition_leases") == 0 && rowCount("queue_offsets") == 0 + }, "drained partition's lease and offsets rows should be released after the idle grace") + + // 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) + require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg2)) + + delivery2 := receive(t, deliveryChan) + require.Equal(t, "idle-2", delivery2.Message().ID) + require.NoError(t, delivery2.Ack(s.ctx)) + + t.Logf("Idle lease released and partition resurrected on new traffic") +} + // TestNackDoesNotBlockOtherMessages verifies that nacking a message does not // block delivery of subsequent messages in the same partition. The nacked // message should be skipped (invisible) while later messages are delivered.