diff --git a/platform/extension/messagequeue/mysql/subscriber.go b/platform/extension/messagequeue/mysql/subscriber.go index a63b289e..10d9f497 100644 --- a/platform/extension/messagequeue/mysql/subscriber.go +++ b/platform/extension/messagequeue/mysql/subscriber.go @@ -485,6 +485,13 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { leaseTicker := time.NewTicker(time.Duration(cfg.LeaseRenewalIntervalMs) * time.Millisecond) defer leaseTicker.Stop() + // Orphan sweep pacing: an uncapped acquisition pass runs every two lease + // durations. Two lease durations is past every transient window in the + // protocol — an expiring lease, a crashed peer's heartbeat going stale — + // so anything still unleased at sweep time is genuinely unclaimed. + orphanSweepInterval := 2 * time.Duration(cfg.LeaseDurationMs) * time.Millisecond + lastOrphanSweep := time.Now() + // Send initial heartbeat so this subscriber is immediately visible to // ActiveSubscribers. Without this, other subscribers compute incorrect // fair shares until the first leaseTicker fires. @@ -533,10 +540,26 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { // Rebalance, renew, and heartbeat are independent operations. // Each can fail without affecting the others — the next tick retries. - if err := s.rebalance(ctx, sub, leasedPartitions); err != nil { + // Renewal covers only the partitions kept after shedding; renewing + // a just-released lease would spuriously fail with ErrLeaseExpired. + released, err := s.rebalance(ctx, sub, leasedPartitions) + if err != nil { s.logger.Errorw("rebalance failed", append(logFields, "error", err)...) } - if err := s.renewLeases(ctx, sub, leasedPartitions); err != nil { + kept := leasedPartitions + if len(released) > 0 { + releasedSet := make(map[string]struct{}, len(released)) + for _, pk := range released { + releasedSet[pk] = struct{}{} + } + kept = make([]string, 0, len(leasedPartitions)) + for _, pk := range leasedPartitions { + if _, ok := releasedSet[pk]; !ok { + kept = append(kept, pk) + } + } + } + if err := s.renewLeases(ctx, sub, kept); err != nil { s.logger.Errorw("lease renewal failed", append(logFields, "error", err)...) } if err := s.sendHeartbeat(ctx, sub); err != nil { @@ -545,7 +568,18 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { s.emitSignal(SignalPartitionUpdate) case <-discoveryTicker.C: - if err := s.discoverAndReconcileWorkers(ctx, sub); err != nil { + // Orphan sweep: periodically run acquisition with no cap. In a + // healthy group the sweep is a no-op — TryAcquireLease cannot + // steal a valid lease — but a partition left unleased for any + // reason the cap arithmetic missed (divergent heartbeat views, a + // subscriber that heartbeats without acquiring) is picked up by + // whichever subscriber sweeps first. An over-cap grab is shed at + // the next rebalance once a peer has spare capacity to take it. + uncapped := time.Since(lastOrphanSweep) >= orphanSweepInterval + if uncapped { + lastOrphanSweep = time.Now() + } + if err := s.discoverAndReconcileWorkers(ctx, sub, uncapped); err != nil { s.logger.Errorw("partition discovery failed, will retry on next tick", append(logFields, "error", err)...) } s.emitSignal(SignalPartitionUpdate) @@ -554,8 +588,9 @@ func (s *subscriber) managePartitions(ctx context.Context, sub *subscription) { } // discoverAndReconcileWorkers discovers new partitions and reconciles workers. -// Uses load-based fair share to limit how many partitions this subscriber acquires. -func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subscription) error { +// Uses fair share to limit how many partitions this subscriber acquires; +// uncapped skips the fair-share cap entirely (the orphan sweep). +func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subscription, uncapped bool) error { cfg := sub.config // Get current leased partitions for fair share computation. @@ -565,15 +600,21 @@ func (s *subscriber) discoverAndReconcileWorkers(ctx context.Context, sub *subsc } // Use cached discovered partitions from last tick for fair share cap. - // On the first tick, lastDiscoveredPartitions is nil → fairShareCap uses - // only owned partitions, which gives unlimited cap for new subscribers. + // On the first tick, lastDiscoveredPartitions is nil → fairShareCap sees + // only owned partitions, so a joiner's first-tick cap floors at 1 and + // ramps once discovery is cached. sub.workersMu.Lock() cachedDiscovered := sub.lastDiscoveredPartitions sub.workersMu.Unlock() - maxPartitions, err := s.fairShareCap(ctx, sub, leasedPartitions, cachedDiscovered) - if err != nil { - return fmt.Errorf("compute fair share cap: %w", err) + // maxPartitions == 0 means unlimited (the orphan sweep, or an + // uncontended single subscriber via fairShareCap). + maxPartitions := 0 + if !uncapped { + maxPartitions, err = s.fairShareCap(ctx, sub, leasedPartitions, cachedDiscovered) + if err != nil { + return fmt.Errorf("compute fair share cap: %w", err) + } } // Discover and try to acquire leases for new partitions. @@ -1026,8 +1067,12 @@ func (s *subscriber) deregisterHeartbeat(ctx context.Context, sub *subscription) } // rebalance checks if this subscriber holds more partitions than its fair share -// and releases extras so other subscribers can pick them up. -func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []string) error { +// and releases extras so other subscribers can pick them up. Returns the +// partitions actually released so the caller renews only the remainder — +// renewing a just-released lease would spuriously fail with ErrLeaseExpired. +// The owned slice is never mutated (the caller shares it with lease renewal). +// On error, partitions released before the failure are still returned. +func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []string) (released []string, retErr error) { cfg := sub.config // Use cached discovered partitions from the most recent discovery tick. @@ -1037,20 +1082,24 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []s maxPart, err := s.fairShareCap(ctx, sub, owned, discoveredPartitions) if err != nil { - return fmt.Errorf("compute fair share cap: %w", err) + return nil, fmt.Errorf("compute fair share cap: %w", err) } if maxPart == 0 || len(owned) <= maxPart { - return nil + return nil, nil } - // Sort deterministically so the same partitions are released across runs. - sort.Strings(owned) + // Sort a copy deterministically so the same partitions are released + // across runs without reordering the caller's slice. + sortedOwned := make([]string, len(owned)) + copy(sortedOwned, owned) + sort.Strings(sortedOwned) // Release excess partitions - for _, pk := range owned[maxPart:] { + for _, pk := range sortedOwned[maxPart:] { if err := s.leaseStore.ReleaseLease(ctx, sub.topic, pk, cfg.SubscriberName, cfg.ConsumerGroup); err != nil { - return fmt.Errorf("release partition %s during rebalance: %w", pk, err) + return released, fmt.Errorf("release partition %s during rebalance: %w", pk, err) } + released = append(released, pk) // Stop the worker immediately to prevent duplicate processing. s.stopPartitionWorker(sub, pk) @@ -1063,7 +1112,7 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []s "max_partitions", maxPart, ) } - return nil + return released, nil } // fairShareCap computes the max partitions this subscriber should own. @@ -1071,6 +1120,16 @@ func (s *subscriber) rebalance(ctx context.Context, sub *subscription, owned []s // owned is the caller-provided list of leased partitions. // discoveredPartitions is an optional pre-fetched list of all known partitions; // if nil, only owned partitions are used for fair share computation. +// +// The cap is remainder-aware: subscribers rank themselves in the sorted +// active list, the first (P mod N) ranks get floor(P/N)+1, and the rest get +// floor(P/N), so per-rank caps sum to exactly P. Independent ceil(P/N) caps +// sum to more than P and admit stable starvation states — e.g. P=12, N=5 +// could settle at 3/3/3/3/0 with every subscriber at cap and nobody obliged +// to shed for the empty one. With caps summing to P, a subscriber over its +// cap implies another under its cap (rebalance sheds, the peer acquires), +// and an unleased partition implies a subscriber with spare cap to claim it +// — neither a starved subscriber nor a leftover partition is a stable state. func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, owned []string, discoveredPartitions []string) (int, error) { cfg := sub.config @@ -1082,8 +1141,6 @@ func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, owned return 0, nil } - activeSubscribers := len(active) - // Count all known partitions as the union of owned + discovered. // Using max(owned, discovered) would undercount when some partitions // have leases but no messages, or vice versa. @@ -1098,8 +1155,31 @@ func (s *subscriber) fairShareCap(ctx context.Context, sub *subscription, owned } totalPartitions := len(partitionSet) - // ceil(totalPartitions / activeSubscribers) - maxPart := (totalPartitions + activeSubscribers - 1) / activeSubscribers + // Rank in the sorted active list. ActiveSubscribers row order is not + // guaranteed, so sorting is what lets every subscriber derive the same + // ranking from the same set without coordination. + sort.Strings(active) + n := len(active) + rank := -1 + for i, name := range active { + if name == cfg.SubscriberName { + rank = i + break + } + } + + var maxPart int + if rank < 0 { + // Own heartbeat not visible this tick (e.g. the write failed): fall + // back to a conservative ceil over n+1 contenders instead of + // claiming a rank that may belong to another subscriber. + maxPart = (totalPartitions + n) / (n + 1) + } else { + maxPart = totalPartitions / n + if rank < totalPartitions%n { + maxPart++ + } + } if maxPart < 1 { maxPart = 1 } diff --git a/platform/extension/messagequeue/mysql/subscriber_test.go b/platform/extension/messagequeue/mysql/subscriber_test.go index c021397a..02112b2e 100644 --- a/platform/extension/messagequeue/mysql/subscriber_test.go +++ b/platform/extension/messagequeue/mysql/subscriber_test.go @@ -793,3 +793,217 @@ func TestSubscriber_StopAllWorkers(t *testing.T) { <-doneCh } } + +func TestSubscriber_FairShareCap(t *testing.T) { + tests := []struct { + name string + self string + active []string // as returned by the heartbeat store, deliberately unsorted + owned []string + discovered []string + want int + }{ + { + // The starvation case: P=12, N=5. Independent ceil caps were 3 + // for every rank (sum 15), so 3/3/3/3/0 was stable. Remainder + // caps are 3,3,2,2,2 (sum 12): rank 0 gets the remainder… + name: "uneven split first rank gets remainder", + self: "s1", + active: []string{"s3", "s1", "s5", "s2", "s4"}, + discovered: partitionKeysN(12), + want: 3, + }, + { + // …and the last rank gets the floor, not zero-forever. + name: "uneven split last rank gets floor", + self: "s5", + active: []string{"s3", "s1", "s5", "s2", "s4"}, + discovered: partitionKeysN(12), + want: 2, + }, + { + name: "even split", + self: "s2", + active: []string{"s2", "s1"}, + discovered: partitionKeysN(4), + want: 2, + }, + { + name: "single subscriber is unlimited", + self: "s1", + active: []string{"s1"}, + discovered: partitionKeysN(4), + want: 0, + }, + { + // P < N: the remainder share is 0 for high ranks, but the cap + // keeps the historical minimum of 1 — with fewer partitions than + // subscribers somebody idles regardless, and the floor preserves + // the maxPart=0-means-unlimited contract. + name: "fewer partitions than subscribers floors at one", + self: "s4", + active: []string{"s1", "s2", "s3", "s4"}, + discovered: partitionKeysN(2), + want: 1, + }, + { + // Own heartbeat missing from the active list (write failed this + // interval): conservative ceil over n+1 contenders, never + // unlimited. + name: "missing own heartbeat falls back to ceil", + self: "s-missing", + active: []string{"s1", "s2"}, + discovered: partitionKeysN(9), + want: 3, + }, + { + name: "owned and discovered are unioned", + self: "s1", + active: []string{"s1", "s2"}, + owned: []string{"pk-00", "pk-extra"}, + discovered: []string{"pk-00", "pk-01", "pk-02"}, + want: 2, // union = 4 partitions, rank 0 of 2 -> 2 + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + mockHB := NewMocksubscriberHeartbeatStore(ctrl) + mockHB.EXPECT(). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(tt.active, nil). + AnyTimes() + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), mockHB, + NewMockdeliveryStateStore(ctrl), + ) + sub := &subscription{ + topic: "test-topic", + config: extqueue.DefaultSubscriptionConfig(tt.self, "test-cg"), + } + + got, err := s.fairShareCap(context.Background(), sub, tt.owned, tt.discovered) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } + + // The anti-starvation invariant: whenever P >= N, per-rank caps sum to + // exactly P, so no combination of at-cap subscribers can leave a + // subscriber starved or a partition unclaimed. + t.Run("caps sum to partition total", func(t *testing.T) { + for n := 2; n <= 6; n++ { + for p := n; p <= 13; p++ { + active := make([]string, n) + for i := range active { + active[i] = fmt.Sprintf("s%d", i) + } + + ctrl := gomock.NewController(t) + mockHB := NewMocksubscriberHeartbeatStore(ctrl) + mockHB.EXPECT(). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(active, nil). + AnyTimes() + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), mockHB, + NewMockdeliveryStateStore(ctrl), + ) + + sum := 0 + for _, self := range active { + sub := &subscription{ + topic: "test-topic", + config: extqueue.DefaultSubscriptionConfig(self, "test-cg"), + } + cap, err := s.fairShareCap(context.Background(), sub, nil, partitionKeysN(p)) + require.NoError(t, err) + sum += cap + } + require.Equal(t, p, sum, "n=%d p=%d", n, p) + } + } + }) +} + +// partitionKeysN generates n distinct partition keys. +func partitionKeysN(n int) []string { + keys := make([]string, n) + for i := range keys { + keys[i] = fmt.Sprintf("pk-%02d", i) + } + return keys +} + +func TestSubscriber_RebalanceReleasesExcess(t *testing.T) { + ctrl := gomock.NewController(t) + + // Two active subscribers, four partitions: self is rank 0 -> cap 2. + mockHB := NewMocksubscriberHeartbeatStore(ctrl) + mockHB.EXPECT(). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return([]string{"s1", "s2"}, nil) + + // The lexicographically largest partitions beyond the cap are released. + mockLease := NewMockpartitionLeaseStore(ctrl) + mockLease.EXPECT(). + ReleaseLease(gomock.Any(), "test-topic", "pk-c", "s1", "test-cg"). + Return(nil) + mockLease.EXPECT(). + ReleaseLease(gomock.Any(), "test-topic", "pk-d", "s1", "test-cg"). + Return(nil) + + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + mockLease, mockHB, NewMockdeliveryStateStore(ctrl), + ) + sub := &subscription{ + topic: "test-topic", + config: extqueue.DefaultSubscriptionConfig("s1", "test-cg"), + workers: make(map[string]*partitionWorker), + } + + owned := []string{"pk-d", "pk-a", "pk-c", "pk-b"} + released, err := s.rebalance(context.Background(), sub, owned) + require.NoError(t, err) + assert.Equal(t, []string{"pk-c", "pk-d"}, released) + // The caller's slice is shared with lease renewal and must not be + // reordered (regression: rebalance used to sort it in place, making the + // subsequent renewal hit the released tail and log ErrLeaseExpired). + assert.Equal(t, []string{"pk-d", "pk-a", "pk-c", "pk-b"}, owned) +} + +func TestSubscriber_RebalanceUnderCapReleasesNothing(t *testing.T) { + ctrl := gomock.NewController(t) + + mockHB := NewMocksubscriberHeartbeatStore(ctrl) + mockHB.EXPECT(). + ActiveSubscribers(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return([]string{"s1", "s2"}, nil) + + // No ReleaseLease expectations: owning exactly the cap sheds nothing. + s := NewSubscriber( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, + NewMockmessageStore(ctrl), NewMockoffsetStore(ctrl), + NewMockpartitionLeaseStore(ctrl), mockHB, NewMockdeliveryStateStore(ctrl), + ) + sub := &subscription{ + topic: "test-topic", + config: extqueue.DefaultSubscriptionConfig("s1", "test-cg"), + workers: make(map[string]*partitionWorker), + // Four known partitions across two subscribers -> rank-0 cap is 2: + // owning exactly the cap must shed nothing. + lastDiscoveredPartitions: []string{"pk-a", "pk-b", "pk-c", "pk-d"}, + } + + released, err := s.rebalance(context.Background(), sub, []string{"pk-a", "pk-b"}) + require.NoError(t, err) + assert.Empty(t, released) +} diff --git a/test/integration/extension/messagequeue/mysql/queue_test.go b/test/integration/extension/messagequeue/mysql/queue_test.go index fd913905..1c1d6c24 100644 --- a/test/integration/extension/messagequeue/mysql/queue_test.go +++ b/test/integration/extension/messagequeue/mysql/queue_test.go @@ -1967,6 +1967,136 @@ func (s *SQLQueueIntegrationSuite) TestRebalance_MoreSubscribersThanPartitions() t.Logf("More subscribers than partitions verified: 2 partitions, 4 subscribers, max 1 each") } +// TestRebalance_NoStarvation_UnevenSplit reproduces the starvation case the +// old independent ceil(P/N) caps admitted: with 12 partitions and 5 +// subscribers every cap was 3 (sum 15), so 3/3/3/3/0 was a stable end state +// with nobody obliged to shed for the empty subscriber. Remainder-aware caps +// are 3+3+2+2+2 (sum 12), so every subscriber must converge to at least +// floor(12/5)=2 partitions. +func (s *SQLQueueIntegrationSuite) TestRebalance_NoStarvation_UnevenSplit() { + t := s.T() + + topic := "rebalance_starvation_topic" + consumerGroup := "rebalance-starvation-cg" + + signalCh := make(chan queueMySQL.HookSignal, 100) + + pubQ, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, + }) + require.NoError(t, err) + defer pubQ.Close() + + const partitionCount = 12 + for i := 0; i < partitionCount; i++ { + pk := fmt.Sprintf("pk-%02d", i) + msg := entityqueue.NewMessage(fmt.Sprintf("rb-starve-%d", i), []byte("x"), pk, nil) + require.NoError(t, pubQ.Publisher().Publish(s.ctx, topic, msg)) + } + + subNames := []string{"s1", "s2", "s3", "s4", "s5"} + var queues []extqueue.Queue + for _, name := range subNames { + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, + OnSignal: signalCh, + }) + require.NoError(t, err) + queues = append(queues, q) + // 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.Retry.MaxAttempts = 1000 + _, err = q.Subscriber().Subscribe(s.ctx, topic, cfg) + require.NoError(t, err) + } + defer func() { + for _, q := range queues { + q.Close() + } + }() + + waitForCondition(t, signalCh, func() bool { + leases, _ := getPartitionLeases(s.db, topic, consumerGroup) + total := 0 + minOwned := partitionCount + maxOwned := 0 + for _, name := range subNames { + owned := len(leases[name]) + total += owned + if owned < minOwned { + minOwned = owned + } + if owned > maxOwned { + maxOwned = owned + } + } + return total == partitionCount && minOwned >= 2 && maxOwned <= 3 + }, "12 partitions across 5 subscribers must split 3+3+2+2+2 — no subscriber starved") + + t.Logf("No starvation: 12 partitions split with every subscriber owning 2-3") +} + +// TestRebalance_OrphanSweep verifies the guarantee that no partition is left +// unprocessed even when the fair-share arithmetic refuses to assign it. +// Phantom heartbeat rows (never acquiring anything) inflate the active count +// so the one real subscriber's cap is 1 with 3 partitions published — the +// normal acquisition path claims one partition and stops. The periodic +// uncapped orphan sweep (every 2x LeaseDurationMs) must pick up the other +// two anyway, proven by every partition's message being delivered and acked. +func (s *SQLQueueIntegrationSuite) TestRebalance_OrphanSweep() { + t := s.T() + + topic := "rebalance_sweep_topic" + consumerGroup := "rebalance-sweep-cg" + partitions := []string{"pk-a", "pk-b", "pk-c"} + + q, err := queueMySQL.NewQueue(queueMySQL.Params{ + DB: s.db, Logger: zaptest.NewLogger(t), MetricsScope: tally.NoopScope, + }) + require.NoError(t, err) + defer q.Close() + + for i, pk := range partitions { + msg := entityqueue.NewMessage(fmt.Sprintf("sweep-%d", i), []byte("x"), pk, nil) + require.NoError(t, q.Publisher().Publish(s.ctx, topic, msg)) + } + + // Two phantom subscribers that heartbeat but never acquire. Their + // heartbeat_at is stamped in the future so they stay "active" for the + // whole test without a refresh loop. Sorted, the real subscriber ranks + // last of 3, so its remainder-aware cap is 3/3 = 1. + futureMs := time.Now().Add(10 * time.Minute).UnixMilli() + for i := 0; i < 2; i++ { + _, err := s.db.ExecContext(s.ctx, ` + INSERT INTO queue_subscriber_heartbeats (consumer_group, topic, subscriber_name, heartbeat_at, deregistered_at) + VALUES (?, ?, ?, ?, 0) + ON DUPLICATE KEY UPDATE heartbeat_at = VALUES(heartbeat_at), deregistered_at = 0 + `, consumerGroup, topic, fmt.Sprintf("phantom-%d", i), futureMs) + require.NoError(t, err) + } + + deliveryChan, err := q.Subscriber().Subscribe(s.ctx, topic, testSubConfig("worker-real", consumerGroup)) + require.NoError(t, err) + + // All three messages must arrive: one via the normal capped acquisition, + // the other two only after the sweep bypasses the cap (~2x the 3s test + // lease duration). Acking promptly proves processing, which is the + // guarantee — lease ownership may churn afterwards as rebalance sheds + // the over-cap sweep grabs. + received := make(map[string]bool) + receiveN(t, deliveryChan, len(partitions), func(delivery extqueue.Delivery, _ int) { + received[delivery.Message().PartitionKey] = true + require.NoError(t, delivery.Ack(s.ctx)) + }) + for _, pk := range partitions { + assert.True(t, received[pk], "partition %s must have been processed", pk) + } + + t.Logf("Orphan sweep verified: all 3 partitions processed despite a fair-share cap of 1") +} + // 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.