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.