[improve][broker] PIP-486 (PR5): broker-side entry-bucket handoff for orderly consumer scaling

[merlimat]

Motivation

Continues PIP-486 after #26114, #26115, #26118, #26129, #26173. PR4 activated bucket-shared consumption for stable membership and left the hard part open: preserving per-key order across an ownership change of a live bucket (a consumer joining, leaving, or crashing mid-traffic).

The design follows three constraints: every language SDK has to implement whatever protocol we pick, so the client surface must stay minimal; transitions should be orderly, with no transient errors and retries when a consumer joins; and a consumer must not re-subscribe for buckets it keeps. The answer is to move the handoff into the broker, at entry-bucket granularity — exactly how Key_Shared AUTO_SPLIT hands over hash ranges today, but per whole bucket instead of per key. Consumers keep consuming throughout a membership change; a bucket move is observable only as a short pause on that bucket.

The mechanism that makes this nearly free is canonical-hash normalization: the dispatcher normalizes every entry's hash to its bucket's canonical value (the bucket's start hash, nudged to 1 for bucket 0 since 0 is the reserved "not set" sentinel). One hash value per bucket means the existing per-hash Key_Shared machinery — pending acks, replay-order blocking, DrainingHashesTracker — operates at exactly bucket granularity, with no new tracking code.

Modifications

Broker — dispatcher seams

• PersistentStickyKeyDispatcherMultipleConsumers gains a protected constructor taking a StickyKeyConsumerSelector and drainingHashesRequired, plus handleConsumerAdded/handleConsumerRemoved and shouldBlockDispatch hooks (defaults preserve AUTO_SPLIT behavior exactly).
• PR4's entryBucketDispatch field and getStickyKeyHash hook are reverted out of the parent: plain Key_Shared dispatch is back to its pre-PR4 form, and bucket logic lives only in the subclass. hasSameKeySharedPolicy rejects a flag mismatch.

Broker — entry-bucket dispatcher

• EntryBucketConsumerSelector: the segment's bucket boundaries are fixed at construction; the only moving part is membership. The N buckets are spread deterministically over the connected consumers sorted by (name, id) — consumer j of k owns the contiguous slice [j*N/k, (j+1)*N/k) — so a membership change moves as few whole buckets as possible, and membership diffs come from the existing ConsumerHashAssignmentsSnapshot.resolveImpactedConsumers.
• PersistentEntryBucketDispatcherMultipleConsumers (selected by PersistentSubscription when the consumer's KeySharedMeta carries entryBucketDispatch): the bucket boundaries arrive in the subscribe itself (KeySharedMeta.hashRanges = the segment's full boundary list, immutable for the segment's life; validated as ascending, contiguous, tiling the 16-bit ring). A joiner declaring different boundaries is rejected (ConsumerAssignException) — a mismatch is a stale layout, not a race. getStickyKeyHash routes a stamped entry by canonical(entry_hash_min) written through the cached-hash path, and unstamped entries by the key's hash normalized to the same canonical value. drainingHashesRequired=true gives the per-bucket handoff: a moving bucket is withheld from its new owner until every message the prior owner has in flight for it is acked, then flows in order.

Controller

• SubscriptionCoordinator: every sharer of a bucketed segment now receives the segment's full boundary list in bucket_ranges (identical for all sharers) instead of a per-consumer slice; empty still means sole owner → Exclusive. The controller only decides which consumers share a segment; the bucket→consumer spread is computed broker-side from live membership, so a join/leave needs no controller round-trip and no re-subscribes.

Client — v5 stream consumer

• The only client-side transition left is the mode flip (Exclusive ⇄ Key_Shared) when a segment's bucket_ranges flips between empty and non-empty, or a segment drop. Before closing the per-segment consumer the client now drains: delivery of that segment is paused and the close waits until every already-delivered message is acked (shared segments: the individual-ack translation queue is empty; whole segments: the cumulative high-water mark reached the latest delivered id). Completion is ack-driven, not time-based, so the flip cannot discard or reorder in-flight messages.
• The bucket_ranges proto comment is refreshed to the boundaries-for-all semantics. No wire-format changes: field shapes from PR3 and the PR4 flag are unchanged, and no new commands are introduced.

Verifications

• EntryBucketConsumerSelectorTest: bucket-index/canonical-hash math; deterministic, add-order-independent spread; more consumers than buckets; key-hash normalization; impacted-bucket reporting on join/leave; merged assignment snapshot.
• PersistentEntryBucketDispatcherMultipleConsumersTest: boundary-list validation accept/reject cases.
• PersistentStickyKeyDispatcherMultipleConsumersTest: the stamped-range routing test now exercises the subclass and asserts canonical values; a stamped entry on a plain Key_Shared subscription still routes by key.
• V5EntryBucketDispatchTest: the PR4 e2e tests, plus testConsumerJoiningMidTrafficPreservesPerKeyOrder — a second stream consumer joins while a keyed stream is in flight (per-message acks), and per-key order is asserted across the broker-side handoff with nothing lost or duplicated.
• Regression sweep: all v5 client suites, the scalable broker suites, the sticky-key dispatcher suite, and checkstyle pass.

[void-ptr974]

Thanks for the detailed write-up. The broker-side handoff approach looks really nice to me, especially the canonical bucket hash trick to reuse the existing Key_Shared draining path.

One thing I wanted to sanity-check: drainAndClose() waits for messages already delivered to the app to be acked before closing the old v4 consumer. But in subscribeAssigned(), it looks like the old segment future is removed/replaced in segmentConsumers before that drain completes, while ackSegmentUpTo() uses segmentConsumers.get(segmentId) to decide where to send the ack.

Could an app ack during the drain miss the old v4 consumer, or in the resubscribe case land on the new one instead? I may be missing another guard here, but otherwise the drain seems to depend on an ack path that no longer points at the consumer being drained.

[merlimat]

Good catch — confirmed, and it was worse than a missed ack: on the drop path the segmentConsumers entry was removed before the drain started, so ackSegmentUpTo() early-returned and the acks the drain depends on could never land — the release deadlocked until the whole consumer closed. On the re-subscribe path the slot already held the new consumer's chain, so mid-drain acks were deferred onto the new consumer — and on an Exclusive→Key_Shared flip those are cumulative acks against a Key_Shared consumer, which the client rejects: silently lost, and the messages redelivered after the flip.

Fixed in 90573d7: a drainingConsumers map keeps the ack path pointed at the consumer being drained (drainAndClose registers it there before vacating the slot, and now owns that removal so the two stay ordered; ackSegmentUpTo checks it first). Two new e2e tests pin both paths — segment release with pending acks, and the join flip with a slow-acking app — and both fail without the fix.

[void-ptr974]

Thanks for the quick fix. I took another pass through the ack/drain handoff paths and the added coverage; this looks good to me.

[lhotari]

I ran a local combined Claude Code Fable / Codex (gpt-5.6-sol) code review on this PR (head 90573d78280). Sharing the suspected bugs below (quality-level notes omitted). The first finding was flagged independently by both reviewers.

1. Stale latestDelivered / lastCumulativeAcked across mode flips can wedge a later drain — deadlock on a single-segment topic (found by both reviewers)

ScalableStreamConsumer.subscribeAssigned's re-subscribe (mode-flip) branch never resets latestDelivered / lastCumulativeAcked — only the drop path and the sealed-segment path clear them — and lastCumulativeAcked is only advanced by the Exclusive branch of ackSegmentUpTo; during a Key_Shared phase, acks take the sharedSegmentUnacked queue path and never touch it.

1. A owns a segment Exclusive and acks up to X (lastCumulativeAcked = X).
2. B joins → A flips to Key_Shared; shared-phase traffic advances latestDelivered to Y > X, individually acked (queue drains).
3. B leaves → A flips back to Exclusive (the queue-based drain passes instantly). The maps now hold delivered = Y, acked = X < Y, and the application holds nothing unacked.
4. B rejoins → the new drain finds no shared queue, and isReleaseDrained() is false forever: there is nothing left to ack, and the drain pauses the segment so no new message can ever be delivered or acked to advance the watermark.

On a single-segment topic this is a permanent wedge: the joiner's subscribe retries until operation timeout and the incumbent stops receiving, with only the "Still waiting for the application to ack" warning firing. Multi-segment topics can self-heal accidentally when a position-vector ack from another segment re-merges the stale position. The converse also holds: a stale cumulative high-water from an earlier Exclusive generation can make a later drain pass prematurely. A fix would clear both maps once drained completes in the re-subscribe branch (mirroring the drop path — safe because a completed drain proves everything delivered was acked), or more generally scope the release state to the segment-consumer generation. None of the new e2e tests exercise a second flip cycle, which is why this survives them.

2. Drain barrier counts acks as applied before the broker accepts them — including transactional acks (Codex)

In ackSegmentUpTo, lastCumulativeAcked is advanced (and ids are removed from sharedSegmentUnacked) before acknowledgeCumulativeAsync() / acknowledgeAsync() is invoked, and the returned ack futures are ignored. isReleaseDrained() can therefore pass — and drainAndClose close the old consumer — while acks are still pending or have failed; a failed ack does not restore the drain state, and the broker then redelivers messages the old application already processed to the new owner (duplicates across the handoff). The transactional variant is stronger: an ack added to a transaction counts toward the barrier immediately (lastCumulativeAcked.merge runs before the v4Txn branch), so the handoff can complete while the application's transaction is still open or later aborts. The barrier should reflect successfully applied acks, and transactional acks during a handoff need explicit semantics (wait for commit, or reject them while draining).

3. Pause does not fence an in-flight receive (both, Codex with the sharper framing)

pausedSegments only stops the receive loop from re-arming after messageSink.accept() completes. A receiveAsync() armed just before drainAndClose pauses the segment can complete after the drain predicate has already passed: it unconditionally updates latestDelivered / the unacked queue and hands one more message to the application after the barrier succeeded, so the new owner can attach while the old application still holds that message (which the broker then redelivers — a duplicate the exact-sequence e2e assertions would only catch under unlucky timing). This needs a generation/fence check when the receive completes, or the drain should wait for the outstanding receive operation to settle.

4. Classic dispatcher can be silently reused for an entry-bucket consumer (Fable)

PersistentStickyKeyDispatcherMultipleConsumersClassic.hasSameKeySharedPolicy still compares only mode + allowOutOfOrderDelivery, while the modern dispatcher now rejects an entryBucketDispatch mismatch in both directions. With subscriptionKeySharedUseClassicPersistentImplementation=true, if a segment-topic subscription is left with a consumer-less classic STICKY dispatcher (e.g. a plain v4 Key_Shared consumer subscribed and left), a later entryBucketDispatch consumer matches it and reuseOrCreateDispatcher silently reuses the classic dispatcher: per-key dispatch, no boundary validation, no bucket handoff. The new AbstractSubscription compatibility check only covers the consumers-connected case, so it doesn't close this hole. The classic implementation needs the same !ksm.isEntryBucketDispatch() check.

5. closeAsync() can complete without awaiting a concurrently-registered draining consumer (Codex)

closeAsync() builds its close set via Stream.concat(segmentConsumers.values().stream(), drainingConsumers.values().stream()); both CHM spliterators are constructed eagerly at the concat call and are only weakly consistent. A reconcile's drainAndClose running concurrently can have its drainingConsumers.put land after spliterator construction (not guaranteed to be traversed) while its segmentConsumers.remove is reflected — so that consumer's close future is excluded from the allOf, and closeAsync() can report completion while a per-segment consumer is still closing (the drain loop does still close it via the closed flag, just un-awaited). A stable snapshot coordinated with drain registration, or a common lifecycle registry, would close this.

---

Assisted-by: Claude Code (Fable 5) + Codex CLI (gpt-5.6-sol); findings manually reviewed.

[lhotari]

Please check the Claude Code Fable / Codex gpt-5.6-sol code review comments.