From d88fca9605ecc294a9cfda77b523e578d7cfecc5 Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Fri, 17 Jul 2026 18:36:53 -0700 Subject: [PATCH 1/5] [improve][broker] PIP-486 (PR5): extract subclass seams from the sticky-key dispatcher Preparation for the broker-side entry-bucket handoff, which lives in a dispatcher subclass so the existing Key_Shared paths stay untouched: - getStickyKeyHash reverts to its pre-PIP-486 form (the entry-bucket routing moves into the subclass), and the entryBucketDispatch field is gone from this class. - Protected seams, no behavior change: a constructor overload taking a pre-built selector and draining mode; handleConsumerAdded/handleConsumerRemoved hooks around the selector membership changes (the AUTO_SPLIT per-hash draining wiring is their default body); and shouldBlockDispatch, the single dispatch-hold-back predicate consulted by both the dispatch path and the replay filter. Assisted-by: Claude Code (Fable 5) --- ...tStickyKeyDispatcherMultipleConsumers.java | 145 ++++++++++-------- 1 file changed, 82 insertions(+), 63 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java index d761bc92594b7..ad62e5c0c1407 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java @@ -69,9 +69,6 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi private final boolean allowOutOfOrderDelivery; private final StickyKeyConsumerSelector selector; private final boolean drainingHashesRequired; - // PIP-486: dispatch whole entries by their producer-stamped entry-bucket hash range instead of - // hashing each message's key. Set only by scalable-topic consumers sharing a segment by bucket. - private final boolean entryBucketDispatch; private boolean skipNextReplayToTriggerLookAhead = false; private final KeySharedMode keySharedMode; @@ -82,34 +79,47 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi PersistentStickyKeyDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor, Subscription subscription, ServiceConfiguration conf, KeySharedMeta ksm) { + this(topic, cursor, subscription, conf, ksm, createSelector(ksm, conf), + // recent joined consumer tracking is required only for AUTO_SPLIT mode when + // out-of-order delivery is disabled + ksm.getKeySharedMode() == KeySharedMode.AUTO_SPLIT && !ksm.isAllowOutOfOrderDelivery()); + } + + /** + * Seam for subclasses (PIP-486 entry-bucket dispatch): supply a custom consumer selector and + * opt out of the per-hash draining tracker (a subclass brings its own handoff tracking). + */ + protected PersistentStickyKeyDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor, + Subscription subscription, ServiceConfiguration conf, KeySharedMeta ksm, + StickyKeyConsumerSelector selector, boolean drainingHashesRequired) { super(topic, cursor, subscription, ksm.isAllowOutOfOrderDelivery()); this.log = LOG.with().ctx(super.log).build(); this.allowOutOfOrderDelivery = ksm.isAllowOutOfOrderDelivery(); this.keySharedMode = ksm.getKeySharedMode(); - this.entryBucketDispatch = ksm.isEntryBucketDispatch(); - // recent joined consumer tracking is required only for AUTO_SPLIT mode when out-of-order delivery is disabled - this.drainingHashesRequired = - keySharedMode == KeySharedMode.AUTO_SPLIT && !allowOutOfOrderDelivery; + this.drainingHashesRequired = drainingHashesRequired; this.drainingHashesTracker = drainingHashesRequired ? new DrainingHashesTracker(this.getName(), this::stickyKeyHashUnblocked) : null; this.rescheduleReadHandler = new RescheduleReadHandler(conf::getKeySharedUnblockingIntervalMs, topic.getBrokerService().executor(), this::cancelPendingRead, () -> reScheduleReadInMs(0), () -> havePendingRead, this::getReadMoreEntriesCallCount, () -> !redeliveryMessages.isEmpty()); - switch (this.keySharedMode) { + this.selector = selector; + } + + private static StickyKeyConsumerSelector createSelector(KeySharedMeta ksm, ServiceConfiguration conf) { + boolean drainingHashesRequired = + ksm.getKeySharedMode() == KeySharedMode.AUTO_SPLIT && !ksm.isAllowOutOfOrderDelivery(); + switch (ksm.getKeySharedMode()) { case AUTO_SPLIT: if (conf.isSubscriptionKeySharedUseConsistentHashing()) { - selector = new ConsistentHashingStickyKeyConsumerSelector( + return new ConsistentHashingStickyKeyConsumerSelector( conf.getSubscriptionKeySharedConsistentHashingReplicaPoints(), drainingHashesRequired); - } else { - selector = new HashRangeAutoSplitStickyKeyConsumerSelector(drainingHashesRequired); } - break; + return new HashRangeAutoSplitStickyKeyConsumerSelector(drainingHashesRequired); case STICKY: - this.selector = new HashRangeExclusiveStickyKeyConsumerSelector(); - break; + return new HashRangeExclusiveStickyKeyConsumerSelector(); default: - throw new IllegalArgumentException("Invalid key-shared mode: " + keySharedMode); + throw new IllegalArgumentException("Invalid key-shared mode: " + ksm.getKeySharedMode()); } } @@ -143,38 +153,48 @@ public synchronized CompletableFuture addConsumer(Consumer consumer) { return CompletableFuture.completedFuture(null); } return super.addConsumer(consumer).thenCompose(__ -> selector.addConsumer(consumer)) - .thenAccept(impactedConsumers -> { + .thenAccept(impactedConsumers -> // TODO: Add some way to prevent changes in between the time the consumer is added and the // time the draining hashes are applied. It might be fine for ConsistentHashingStickyKeyConsumerSelector // since it's not really asynchronous, although it returns a CompletableFuture - if (drainingHashesRequired) { - consumer.setPendingAcksAddHandler(this::handleAddingPendingAck); - consumer.setPendingAcksRemoveHandler(new PendingAcksMap.PendingAcksRemoveHandler() { - @Override - public void handleRemoving(Consumer consumer, long ledgerId, long entryId, int stickyKeyHash, - boolean closing) { - drainingHashesTracker.reduceRefCount(consumer, stickyKeyHash, closing); - } - - @Override - public void startBatch() { - drainingHashesTracker.startBatch(); - } - - @Override - public void endBatch() { - drainingHashesTracker.endBatch(); - } - }); - consumer.setDrainingHashesConsumerStatsUpdater(drainingHashesTracker::updateConsumerStats); - registerDrainingHashes(consumer, impactedConsumers.orElseThrow()); - } - }).exceptionally(ex -> { + handleConsumerAdded(consumer, impactedConsumers) + ).exceptionally(ex -> { internalRemoveConsumer(consumer); throw FutureUtil.wrapToCompletionException(ex); }); } + /** + * Hook invoked (under the dispatcher monitor) after a consumer was added to the selector, with + * the ranges the addition moved between consumers. The base wires the AUTO_SPLIT per-hash + * draining tracker; the PIP-486 entry-bucket subclass replaces this with per-bucket tracking. + */ + protected synchronized void handleConsumerAdded(Consumer consumer, + Optional impactedConsumers) { + if (drainingHashesRequired) { + consumer.setPendingAcksAddHandler(this::handleAddingPendingAck); + consumer.setPendingAcksRemoveHandler(new PendingAcksMap.PendingAcksRemoveHandler() { + @Override + public void handleRemoving(Consumer consumer, long ledgerId, long entryId, int stickyKeyHash, + boolean closing) { + drainingHashesTracker.reduceRefCount(consumer, stickyKeyHash, closing); + } + + @Override + public void startBatch() { + drainingHashesTracker.startBatch(); + } + + @Override + public void endBatch() { + drainingHashesTracker.endBatch(); + } + }); + consumer.setDrainingHashesConsumerStatsUpdater(drainingHashesTracker::updateConsumerStats); + registerDrainingHashes(consumer, impactedConsumers.orElseThrow()); + } + } + private synchronized void registerDrainingHashes(Consumer skipConsumer, ImpactedConsumersResult impactedConsumers) { impactedConsumers.processUpdatedHashRanges((c, updatedHashRanges, opType) -> { @@ -210,6 +230,16 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE // The consumer must be removed from the selector before calling the superclass removeConsumer method. Optional impactedConsumers = selector.removeConsumer(consumer); super.removeConsumer(consumer); + handleConsumerRemoved(consumer, impactedConsumers); + } + + /** + * Hook invoked (under the dispatcher monitor) after a consumer was removed from the selector, + * with the ranges the removal moved between the remaining consumers. See + * {@link #handleConsumerAdded}. + */ + protected synchronized void handleConsumerRemoved(Consumer consumer, + Optional impactedConsumers) { if (drainingHashesRequired) { // register draining hashes for the impacted consumers and ranges, in case a hash switched from one // consumer to another. This will handle the case where a hash gets switched from an existing @@ -552,15 +582,22 @@ private boolean canDispatchEntry(Consumer consumer, Entry entry, if (readType == ReadType.Normal && redeliveryMessages.containsStickyKeyHash(stickyKeyHash)) { return false; } - if (drainingHashesRequired) { - // If the hash is draining, do not send the message - if (drainingHashesTracker.shouldBlockStickyKeyHash(consumer, stickyKeyHash)) { - return false; - } + if (shouldBlockDispatch(consumer, stickyKeyHash)) { + return false; } return true; } + /** + * Whether dispatching an entry with the given sticky key hash to the given consumer must be + * held back for an in-progress handoff. The base blocks per-hash while AUTO_SPLIT draining is + * active; the PIP-486 entry-bucket subclass blocks per moving bucket. + */ + protected boolean shouldBlockDispatch(Consumer consumer, int stickyKeyHash) { + // If the hash is draining, do not send the message + return drainingHashesRequired && drainingHashesTracker.shouldBlockStickyKeyHash(consumer, stickyKeyHash); + } + /** * Creates a filter for replaying messages. The filter is stateful and shouldn't be cached or reused. * @see PersistentDispatcherMultipleConsumers#createFilterForReplay() @@ -623,8 +660,7 @@ public boolean test(Position position) { return false; } - if (drainingHashesRequired - && drainingHashesTracker.shouldBlockStickyKeyHash(consumer, stickyKeyHash.intValue())) { + if (shouldBlockDispatch(consumer, stickyKeyHash.intValue())) { // the hash is draining and the consumer is not the draining consumer alreadyBlockedHashes.add(stickyKeyHash); return false; @@ -638,23 +674,6 @@ public boolean test(Position position) { @Override protected int getStickyKeyHash(Entry entry) { if (entry instanceof EntryAndMetadata entryAndMetadata) { - // PIP-486: an entry-bucket subscription routes each whole entry by its producer-stamped - // entry-bucket hash range, so a batch holding one bucket's keys goes to the bucket's owner - // with no per-key hashing. The stamp shares the 16-bit key-hash space, so it feeds the - // selector directly. Cached as THE entry's sticky-key hash so pending acks, redelivery and - // draining all see the value dispatch used. Unstamped entries (non-batched messages) fall - // through to the message's sticky-key hash, which is the same low-16 Murmur value the - // producer would have stamped. - if (entryBucketDispatch) { - var metadata = entryAndMetadata.getMetadata(); - if (metadata != null && metadata.hasEntryHashMin()) { - return entryAndMetadata.getOrUpdateCachedStickyKeyHash(stickyKey -> { - int bucketHash = metadata.getEntryHashMin(); - // 0 is reserved as "hash not set"; nudge to 1, which is still inside bucket 0. - return bucketHash == STICKY_KEY_HASH_NOT_SET ? 1 : bucketHash; - }); - } - } // use the cached sticky key hash if available, otherwise calculate the sticky key hash and cache it return entryAndMetadata.getOrUpdateCachedStickyKeyHash(selector::makeStickyKeyHash); } From ed668827b7b235ce7635c4670d79580be055d8b2 Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Fri, 17 Jul 2026 18:51:33 -0700 Subject: [PATCH 2/5] [improve][broker] PIP-486 (PR5): broker-side entry-bucket handoff dispatcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New PersistentEntryBucketDispatcherMultipleConsumers, selected by PersistentSubscription when the consumers' KeySharedMeta carries entryBucketDispatch, so all PIP-486 logic lives in one subclass and the existing Key_Shared paths stay untouched. - EntryBucketConsumerSelector: the segment's bucket boundaries (declared identically by every consumer at subscribe — a segment's bucketing is immutable) spread deterministically over the connected consumers as contiguous slices; membership changes produce the standard ImpactedConsumersResult diff via ConsumerHashAssignmentsSnapshot. - The dispatcher normalizes every entry's hash to its bucket's canonical value — the stamped entry_hash_min for batched entries, the key's low-16 Murmur for unstamped ones, both landing in the same bucket. One hash value per bucket makes the inherited per-hash machinery operate at bucket granularity automatically: pending acks, replay-order blocking, and the draining tracker. A membership change hands buckets over exactly like AUTO_SPLIT hands over hash ranges — the moving bucket is blocked until the previous owner's outstanding entries are acked, then flows to the new owner in order. No client re-subscribes, no rejected subscribes, no client-side protocol: the handoff is entirely broker-side. - Boundary validation: contiguous ascending tiling of the 16-bit ring at create and for every joining consumer; the AbstractSubscription compatibility check keeps plain Key_Shared consumers and entry-bucket consumers off each other's subscriptions. Assisted-by: Claude Code (Fable 5) --- .../broker/service/AbstractSubscription.java | 9 + .../service/EntryBucketConsumerSelector.java | 162 ++++++++++++++++++ ...ntryBucketDispatcherMultipleConsumers.java | 150 ++++++++++++++++ ...tStickyKeyDispatcherMultipleConsumers.java | 4 +- .../persistent/PersistentSubscription.java | 7 +- 5 files changed, 330 insertions(+), 2 deletions(-) create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelector.java create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumers.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractSubscription.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractSubscription.java index ae108646532dc..f16643914cd68 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractSubscription.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractSubscription.java @@ -73,6 +73,15 @@ protected Optional> checkForConsumerCompatibilityErrorWi + "is different than the connecting consumer's key_shared mode '%s'.", dispatcherKsm.getKeySharedMode(), consumerKsm.getKeySharedMode())))); } + if (dispatcherKsm.isEntryBucketDispatch() != consumerKsm.isEntryBucketDispatch()) { + return Optional.of(FutureUtil.failedFuture(new BrokerServiceException.SubscriptionBusyException( + String.format("Subscription is of different type. %s", + dispatcherKsm.isEntryBucketDispatch() + ? "Active subscription dispatches by entry-bucket while the connecting " + + "consumer does not." : + "Active subscription does not dispatch by entry-bucket while the connecting " + + "consumer requests it.")))); + } if (dispatcherKsm.isAllowOutOfOrderDelivery() != consumerKsm.isAllowOutOfOrderDelivery()) { return Optional.of(FutureUtil.failedFuture(new BrokerServiceException.SubscriptionBusyException( String.format("Subscription is of different type. %s", diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelector.java new file mode 100644 index 0000000000000..5f0d05b88218f --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelector.java @@ -0,0 +1,162 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import org.apache.pulsar.client.api.Range; + +/** + * PIP-486: assigns a scalable-topic segment's entry-buckets to the connected consumers. + * + *

The bucket boundaries are fixed for the selector's life (a segment's bucketing is immutable — + * every consumer of the subscription declares the identical boundary list at subscribe time, see + * {@code KeySharedMeta.entryBucketDispatch}). The selector's only moving part is membership: the + * {@code N} buckets are spread deterministically over the connected consumers — consumers + * sorted by name (id as tiebreak), consumer {@code j} of {@code k} owning the contiguous bucket + * slice {@code [j*N/k, (j+1)*N/k)} — so every broker computes the same spread from the same + * membership, and a membership change moves as few whole buckets as possible. + * + *

Every hash this selector sees or produces is a bucket's canonical hash (the bucket's + * start hash, nudged to 1 for bucket 0 since 0 is the reserved "not set" value). The owning + * dispatcher normalizes each entry's stamped {@code entry_hash_min} — and, for unstamped entries, + * the message key's hash — to the canonical value, which is what makes the existing per-hash + * machinery (pending acks, replay ordering, {@link DrainingHashesTracker} draining) operate at + * exactly bucket granularity: one hash value per bucket. + */ +public class EntryBucketConsumerSelector implements StickyKeyConsumerSelector { + + private final List bucketRanges; + private final int[] bucketStarts; + + /** Connected consumers, sorted by (name, id). Guarded by this. */ + private final List consumers = new ArrayList<>(); + /** Bucket index → owning consumer; empty array while no consumer is connected. */ + private volatile Consumer[] bucketOwners = new Consumer[0]; + private volatile ConsumerHashAssignmentsSnapshot assignmentsSnapshot = + ConsumerHashAssignmentsSnapshot.empty(); + + /** + * @param bucketRanges ascending, inclusive, contiguous ranges tiling {@code [0, 0xFFFF]}; + * already validated by the caller + */ + public EntryBucketConsumerSelector(List bucketRanges) { + this.bucketRanges = List.copyOf(bucketRanges); + this.bucketStarts = new int[bucketRanges.size()]; + for (int i = 0; i < bucketRanges.size(); i++) { + bucketStarts[i] = bucketRanges.get(i).getStart(); + } + } + + public List getBucketRanges() { + return bucketRanges; + } + + /** The bucket index owning the given (16-bit) hash. */ + public int bucketIndexOfHash(int hash) { + int low = 0; + int high = bucketStarts.length - 1; + while (low < high) { + int mid = (low + high + 1) >>> 1; + if (bucketStarts[mid] <= hash) { + low = mid; + } else { + high = mid - 1; + } + } + return low; + } + + /** + * The canonical hash representing a bucket — its start hash, nudged to 1 for bucket 0 + * (0 is the reserved {@link #STICKY_KEY_HASH_NOT_SET} value; 1 is still inside bucket 0). + */ + public int canonicalHashOfBucket(int bucketIndex) { + return Math.max(bucketStarts[bucketIndex], 1); + } + + /** Normalize any 16-bit hash to its bucket's canonical hash. */ + public int canonicalHashOf(int hash) { + return canonicalHashOfBucket(bucketIndexOfHash(hash)); + } + + @Override + public int makeStickyKeyHash(byte[] stickyKey) { + return canonicalHashOf( + StickyKeyConsumerSelectorUtils.makeStickyKeyHash(stickyKey, getKeyHashRange())); + } + + @Override + public synchronized CompletableFuture> addConsumer(Consumer consumer) { + consumers.add(consumer); + consumers.sort(Comparator.comparing(Consumer::consumerName) + .thenComparingLong(Consumer::consumerId)); + return CompletableFuture.completedFuture(Optional.of(recomputeAssignments())); + } + + @Override + public synchronized Optional removeConsumer(Consumer consumer) { + consumers.remove(consumer); + return Optional.of(recomputeAssignments()); + } + + private ImpactedConsumersResult recomputeAssignments() { + int bucketCount = bucketRanges.size(); + Consumer[] owners = new Consumer[consumers.isEmpty() ? 0 : bucketCount]; + List assignments = new ArrayList<>(bucketCount); + if (!consumers.isEmpty()) { + int consumerCount = consumers.size(); + for (int j = 0; j < consumerCount; j++) { + int from = (int) ((long) j * bucketCount / consumerCount); + int to = (int) ((long) (j + 1) * bucketCount / consumerCount); + for (int bucket = from; bucket < to; bucket++) { + owners[bucket] = consumers.get(j); + assignments.add(new HashRangeAssignment(bucketRanges.get(bucket), consumers.get(j))); + } + } + } + ConsumerHashAssignmentsSnapshot after = ConsumerHashAssignmentsSnapshot.of(assignments); + ImpactedConsumersResult impacted = assignmentsSnapshot.resolveImpactedConsumers(after); + bucketOwners = owners; + assignmentsSnapshot = after; + return impacted; + } + + @Override + public Consumer select(int hash) { + Consumer[] owners = bucketOwners; + if (owners.length == 0) { + return null; + } + return owners[bucketIndexOfHash(hash)]; + } + + @Override + public Range getKeyHashRange() { + return Range.of(0, DEFAULT_RANGE_SIZE - 1); + } + + @Override + public ConsumerHashAssignmentsSnapshot getConsumerHashAssignmentsSnapshot() { + return assignmentsSnapshot; + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumers.java new file mode 100644 index 0000000000000..1ecedac3dbeaa --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumers.java @@ -0,0 +1,150 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service.persistent; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import org.apache.bookkeeper.mledger.Entry; +import org.apache.bookkeeper.mledger.ManagedCursor; +import org.apache.pulsar.broker.ServiceConfiguration; +import org.apache.pulsar.broker.service.BrokerServiceException; +import org.apache.pulsar.broker.service.Consumer; +import org.apache.pulsar.broker.service.EntryAndMetadata; +import org.apache.pulsar.broker.service.EntryBucketConsumerSelector; +import org.apache.pulsar.broker.service.Subscription; +import org.apache.pulsar.client.api.Range; +import org.apache.pulsar.common.api.proto.IntRange; +import org.apache.pulsar.common.api.proto.KeySharedMeta; + +/** + * PIP-486: Key_Shared dispatcher for a scalable-topic segment shared by entry-bucket. + * + *

Selected by {@link PersistentSubscription} when the consumers' {@link KeySharedMeta} carries + * {@code entryBucketDispatch}. Producers of a scalable segment batch per entry-bucket and stamp + * each entry's effective bucket hash range ({@code entry_hash_min}); this dispatcher routes each + * whole entry to the consumer owning its bucket — no per-message key hashing, no entry + * filtering, batching intact. + * + *

The mechanics reuse the battle-tested per-hash Key_Shared machinery at bucket granularity by + * normalizing every entry's hash to its bucket's canonical value (one hash value per + * bucket, see {@link EntryBucketConsumerSelector#canonicalHashOf}): pending acks, replay-order + * blocking, and the draining tracker then operate per bucket automatically. A membership change + * therefore hands buckets over exactly like AUTO_SPLIT hands over hash ranges — the moving bucket + * is blocked until the previous owner's outstanding entries are acked, then flows to the new owner + * in order. Clients never re-subscribe, are never rejected, and implement nothing beyond the + * subscribe flag: the handoff is entirely broker-side. + * + *

The bucket boundaries arrive in the subscribe itself ({@code KeySharedMeta.hashRanges}): a + * segment's bucketing is immutable for its life, so every consumer declares the identical + * boundary list; the first consumer's boundaries create the selector and later consumers are + * validated against them. + */ +public class PersistentEntryBucketDispatcherMultipleConsumers + extends PersistentStickyKeyDispatcherMultipleConsumers { + + private final EntryBucketConsumerSelector bucketSelector; + + PersistentEntryBucketDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor, + Subscription subscription, ServiceConfiguration conf, KeySharedMeta ksm) { + // Draining is required: the inherited DrainingHashesTracker tracks the canonical bucket + // hashes, which makes it a per-bucket handoff tracker. + super(topic, cursor, subscription, conf, ksm, + new EntryBucketConsumerSelector(validateBucketBoundaries(ksm)), true); + this.bucketSelector = (EntryBucketConsumerSelector) getSelector(); + } + + /** + * Validate and convert the boundaries declared at subscribe time: ascending, inclusive, + * contiguous ranges tiling the whole 16-bit entry-bucket ring, with bucket 0 wide enough to + * contain the canonical hash 1. + */ + static List validateBucketBoundaries(KeySharedMeta ksm) { + int count = ksm.getHashRangesCount(); + if (count == 0) { + throw new IllegalArgumentException( + "Entry-bucket subscription must declare the segment's bucket boundaries"); + } + List ranges = new ArrayList<>(count); + int expectedStart = 0; + for (int i = 0; i < count; i++) { + IntRange r = ksm.getHashRangeAt(i); + if (r.getStart() != expectedStart || r.getEnd() < r.getStart()) { + throw new IllegalArgumentException("Entry-bucket boundaries must be ascending, " + + "contiguous and start at 0: found [" + r.getStart() + "," + r.getEnd() + + "] where start " + expectedStart + " was expected"); + } + ranges.add(Range.of(r.getStart(), r.getEnd())); + expectedStart = r.getEnd() + 1; + } + if (expectedStart != EntryBucketConsumerSelector.DEFAULT_RANGE_SIZE) { + throw new IllegalArgumentException("Entry-bucket boundaries must tile the 16-bit ring: " + + "last range ends at " + (expectedStart - 1)); + } + if (ranges.get(0).getEnd() < 1) { + throw new IllegalArgumentException( + "Entry-bucket 0 must span at least [0,1] to hold the canonical hash"); + } + return ranges; + } + + @Override + public synchronized CompletableFuture addConsumer(Consumer consumer) { + // A segment's bucketing is immutable, so every consumer must declare the boundaries the + // dispatcher was created with — a mismatch is a client bug or a stale layout, not a race. + try { + List declared = validateBucketBoundaries(consumer.getKeySharedMeta()); + if (!declared.equals(bucketSelector.getBucketRanges())) { + return CompletableFuture.failedFuture(new BrokerServiceException.ConsumerAssignException( + "Consumer declares different entry-bucket boundaries than the subscription: " + + declared + " != " + bucketSelector.getBucketRanges())); + } + } catch (IllegalArgumentException e) { + return CompletableFuture.failedFuture( + new BrokerServiceException.ConsumerAssignException(e.getMessage())); + } + return super.addConsumer(consumer); + } + + @Override + public boolean hasSameKeySharedPolicy(KeySharedMeta ksm) { + return ksm.getKeySharedMode() == getKeySharedMode() + && ksm.isAllowOutOfOrderDelivery() == isAllowOutOfOrderDelivery() + && ksm.isEntryBucketDispatch(); + } + + @Override + protected int getStickyKeyHash(Entry entry) { + if (entry instanceof EntryAndMetadata entryAndMetadata) { + // Cached so pending acks, redelivery and draining all see the same canonical value. + return entryAndMetadata.getOrUpdateCachedStickyKeyHash(stickyKey -> { + var metadata = entryAndMetadata.getMetadata(); + if (metadata != null && metadata.hasEntryHashMin()) { + // The producer stamped the entry's effective bucket range: route the whole + // entry by its bucket. + return bucketSelector.canonicalHashOf(metadata.getEntryHashMin()); + } + // Unstamped (non-batched) entry: the key's low-16 Murmur hash is the same value + // the producer would have stamped, normalized to the same bucket. + return bucketSelector.makeStickyKeyHash(stickyKey); + }); + } + return bucketSelector.makeStickyKeyHash(peekStickyKey(entry)); + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java index ad62e5c0c1407..a522c0859d3ce 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java @@ -806,7 +806,9 @@ public boolean isAllowOutOfOrderDelivery() { public boolean hasSameKeySharedPolicy(KeySharedMeta ksm) { return (ksm.getKeySharedMode() == this.keySharedMode - && ksm.isAllowOutOfOrderDelivery() == this.allowOutOfOrderDelivery); + && ksm.isAllowOutOfOrderDelivery() == this.allowOutOfOrderDelivery + // PIP-486: entry-bucket subscriptions use PersistentEntryBucketDispatcherMultipleConsumers + && !ksm.isEntryBucketDispatch()); } public Map> getConsumerKeyHashRanges() { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentSubscription.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentSubscription.java index fc8227cd153c0..9ba3a4be50e28 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentSubscription.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentSubscription.java @@ -340,7 +340,12 @@ protected Dispatcher reuseOrCreateDispatcher(Dispatcher dispatcher, Consumer con || !((StickyKeyDispatcher) dispatcher) .hasSameKeySharedPolicy(ksm)) { previousDispatcher = dispatcher; - if (config.isSubscriptionKeySharedUseClassicPersistentImplementation()) { + if (ksm.isEntryBucketDispatch()) { + // PIP-486 scalable-topic segment shared by entry-bucket. Always uses the + // modern implementation; the classic dispatcher has no bucket support. + dispatcher = new PersistentEntryBucketDispatcherMultipleConsumers(topic, cursor, + this, config, ksm); + } else if (config.isSubscriptionKeySharedUseClassicPersistentImplementation()) { dispatcher = new PersistentStickyKeyDispatcherMultipleConsumersClassic(topic, cursor, this, config, ksm); From 368d4938f1904776c6bf8ec9c4653ff02a3bc5e0 Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Fri, 17 Jul 2026 19:01:49 -0700 Subject: [PATCH 3/5] [improve][broker] PIP-486 (PR5): controller sends bucket boundaries; client drains before mode flips The controller now sends every sharer of a bucketed segment the segment's full, identical bucket-boundary list in bucket_ranges (instead of slicing buckets across consumers): the bucket-to-consumer spread is computed broker-side by EntryBucketConsumerSelector from live membership, so the controller only decides *which consumers* share a segment, not who owns which bucket. An empty bucket_ranges still means sole ownership (Exclusive). The v5 stream consumer's only transition is the Exclusive/Key_Shared mode flip when bucket_ranges flips between empty and non-empty. Before closing the old per-segment consumer it now drains: delivery of the segment is paused and the close waits until every released message is acked (shared segments: the individual-ack translation queue is empty; whole segments: the cumulative high-water mark reached the latest delivered id), so no prefetched message is dropped or redelivered out of order across the flip. The same drain runs when a segment is dropped from the assignment. Re-adds the mid-traffic e2e (consumer joins while a keyed stream is in flight, per-key order asserted across the handoff), now exercising the broker-side per-bucket draining handoff, with per-message cumulative acks so draining tracks real consumption. Assisted-by: Claude Code (Fable 5) --- .../scalable/SubscriptionCoordinator.java | 16 +-- .../scalable/SubscriptionCoordinatorTest.java | 31 +++-- .../api/v5/V5EntryBucketDispatchTest.java | 88 +++++++++++++- .../impl/v5/ScalableStreamConsumer.java | 113 +++++++++++++++--- 4 files changed, 202 insertions(+), 46 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java index 9b33ff7120bef..22a8ccea0f511 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java @@ -626,17 +626,17 @@ Map computeAssignment( assignmentLists.get(consumer).add(new ConsumerAssignment.AssignedSegment( segment.segmentId(), segment.hashRange(), segmentTopic.toString(), List.of())); } else { - List buckets = EntryBucketSplits.ranges(segment.entryBucketSplits()); - int base = buckets.size() / k; - int extra = buckets.size() % k; - int from = 0; + // Shared segment: every attached consumer subscribes Key_Shared declaring the + // segment's full (immutable) bucket boundary list — WHO owns WHICH bucket is + // decided broker-side by the segment dispatcher's deterministic spread, and + // bucket handoffs on membership changes drain broker-side too. The controller + // only directs membership. + List boundaries = EntryBucketSplits.ranges(segment.entryBucketSplits()); for (int c = 0; c < k; c++) { - int size = base + (c < extra ? 1 : 0); - List slice = List.copyOf(buckets.subList(from, from + size)); - from += size; ConsumerSession consumer = sortedConsumers.get(consumerIndex++); assignmentLists.get(consumer).add(new ConsumerAssignment.AssignedSegment( - segment.segmentId(), segment.hashRange(), segmentTopic.toString(), slice)); + segment.segmentId(), segment.hashRange(), segmentTopic.toString(), + boundaries)); } } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java index 2480d4de09271..b821820e556cc 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java @@ -27,8 +27,6 @@ import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; import java.time.Duration; -import java.util.ArrayList; -import java.util.Comparator; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -461,32 +459,30 @@ public void testLoneConsumerOwnsBucketedSegmentWhole() throws Exception { @Test public void testSurplusConsumersFanOutBucketedSegment() throws Exception { - // One segment with N=4 entry-buckets, two consumers: the surplus consumer fans the segment - // out — each owner takes a contiguous half of the buckets, disjoint and tiling the ring. + // One segment with N=4 entry-buckets, two consumers: the surplus consumer shares the + // segment. Every sharer carries the segment's FULL bucket boundary list (identical for + // all) — which buckets each consumer owns is decided broker-side by the segment + // dispatcher, and handoffs drain broker-side. SubscriptionCoordinator c = bucketedCoordinator(); c.registerConsumer("consumer-1", 1L, mock(TransportCnx.class)).get(); Map result = c.registerConsumer("consumer-2", 2L, mock(TransportCnx.class)).get(); assertEquals(result.size(), 2); - List owned = new ArrayList<>(); for (ConsumerAssignment assignment : result.values()) { assertEquals(assignment.assignedSegments().size(), 1); ConsumerAssignment.AssignedSegment seg = assignment.assignedSegments().get(0); assertEquals(seg.segmentId(), 0); - assertEquals(seg.bucketRanges().size(), 2); - owned.addAll(seg.bucketRanges()); - } - owned.sort(Comparator.comparingInt(HashRange::start)); - assertEquals(owned, List.of( + assertEquals(seg.bucketRanges(), List.of( HashRange.of(0x0000, 0x3FFF), HashRange.of(0x4000, 0x7FFF), HashRange.of(0x8000, 0xBFFF), HashRange.of(0xC000, 0xFFFF))); + } } @Test public void testFanOutCapsAtBucketCountAndLeavesRestIdle() throws Exception { - // One segment with N=4 entry-buckets, five consumers: four owners with one bucket each; the - // fifth consumer exceeds the topic's bucket capacity and stays idle (empty assignment). + // One segment with N=4 entry-buckets, five consumers: four sharers (each carrying the full + // boundary list); the fifth exceeds the segment's bucket capacity and stays idle. SubscriptionCoordinator c = bucketedCoordinator(); for (int i = 1; i <= 4; i++) { c.registerConsumer("consumer-" + i, i, mock(TransportCnx.class)).get(); @@ -495,19 +491,20 @@ public void testFanOutCapsAtBucketCountAndLeavesRestIdle() throws Exception { c.registerConsumer("consumer-5", 5L, mock(TransportCnx.class)).get(); assertEquals(result.size(), 5); - List owned = new ArrayList<>(); int idle = 0; + int sharers = 0; for (ConsumerAssignment assignment : result.values()) { if (assignment.assignedSegments().isEmpty()) { idle++; continue; } - ConsumerAssignment.AssignedSegment seg = assignment.assignedSegments().get(0); - assertEquals(seg.bucketRanges().size(), 1); - owned.addAll(seg.bucketRanges()); + sharers++; + assertEquals(assignment.assignedSegments().get(0).bucketRanges(), List.of( + HashRange.of(0x0000, 0x3FFF), HashRange.of(0x4000, 0x7FFF), + HashRange.of(0x8000, 0xBFFF), HashRange.of(0xC000, 0xFFFF))); } assertEquals(idle, 1); - assertEquals(owned.size(), 4); + assertEquals(sharers, 4); } private SubscriptionCoordinator bucketedCoordinator() { diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java index 6539de7b261fc..1978dc607cfae 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java @@ -180,22 +180,98 @@ public void testTwoConsumersShareBucketedSegmentByEntryBucket() throws Exception assertNull(c.receive(Duration.ofSeconds(3)), "acked messages were redelivered"); } - /** Drains until idle, recording values per key in arrival order, then acks cumulatively. */ + @Test + public void testConsumerJoiningMidTrafficPreservesPerKeyOrder() throws Exception { + String topic = newScalableTopic(1); + // Pin the layout so the second consumer is served by entry-bucket fan-out (see above). + admin.scalableTopics().setAutoScalePolicy(topic, + AutoScalePolicyOverride.builder().enabled(false).build()); + String subscription = "bucket-handoff"; + + @Cleanup + Producer producer = v5Client.newProducer(Schema.string()) + .topic(topic) + .create(); + @Cleanup + StreamConsumer a = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribe(); + + // Continuous keyed traffic; consumer B joins mid-stream, forcing a live two-phase handoff: + // A pauses, drains what it already handed to the application, shrinks from the whole + // segment to half the buckets and confirms; only then is B granted the other half. + List keys = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + keys.add("key-" + i); + } + int perKey = 40; + Map> sent = new HashMap<>(); + for (String k : keys) { + sent.put(k, new ArrayList<>()); + } + + Map> aGot = new ConcurrentHashMap<>(); + Map> bGot = new ConcurrentHashMap<>(); + Thread ta = drainOrdered(a, aGot); + + StreamConsumer b = null; + try { + for (int i = 0; i < perKey; i++) { + if (i == perKey / 3) { + // Mid-traffic join: triggers the handoff while messages flow. + b = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribe(); + } + for (String k : keys) { + String value = k + "-" + i; + producer.newMessage().key(k).value(value).send(); + sent.get(k).add(value); + } + } + + Thread tb = drainOrdered(b, bGot); + ta.join(); + tb.join(); + + assertFalse(aGot.isEmpty(), "consumer A received nothing"); + assertFalse(bGot.isEmpty(), "consumer B received nothing — the handoff did not happen"); + + // The handoff contract: for every key, what A processed followed by what B processed + // is exactly the sent sequence — no loss, no duplicates (a clean drain leaves nothing + // to redeliver), and never a message on B before A finished everything older. + for (String k : keys) { + List combined = new ArrayList<>(aGot.getOrDefault(k, List.of())); + combined.addAll(bGot.getOrDefault(k, List.of())); + assertEquals(combined, sent.get(k), "per-key order across the handoff for key=" + k); + } + } finally { + if (b != null) { + b.close(); + } + } + } + + /** + * Drains until idle, recording values per key in arrival order. Acks every message as it is + * processed — like an order-sensitive application would — so a release drain during a handoff + * can complete promptly. + */ private Thread drainOrdered(StreamConsumer consumer, Map> into) { Thread t = new Thread(() -> { try { - MessageId last = null; while (true) { Message msg = consumer.receive(Duration.ofSeconds(2)); if (msg == null) { - if (last != null) { - consumer.acknowledgeCumulative(last); - } return; } String key = msg.key().orElseThrow(() -> new AssertionError("missing key")); into.computeIfAbsent(key, __ -> new ArrayList<>()).add(msg.value()); - last = msg.id(); + consumer.acknowledgeCumulative(msg.id()); } } catch (Exception ignored) { } diff --git a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java index 14f26de25d771..9d1e062797eae 100644 --- a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java +++ b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java @@ -24,6 +24,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ConcurrentHashMap; @@ -110,6 +111,19 @@ final class ScalableStreamConsumer private final ConcurrentHashMap> sharedSegmentUnacked = new ConcurrentHashMap<>(); + /** + * PIP-486: per whole-segment (Exclusive) consumer, the highest position the application has + * cumulatively acked — the release drain barrier for segments without individual-ack tracking. + */ + private final ConcurrentHashMap lastCumulativeAcked = + new ConcurrentHashMap<>(); + + /** + * PIP-486: segments paused for a release — their receive loops stop re-arming so what was + * already handed to the application can drain before the segment consumer is closed. + */ + private final Set pausedSegments = ConcurrentHashMap.newKeySet(); + /** Latest controller assignment; {@link #reconcile()} converges the segment consumers onto it. */ private volatile List latestAssignment; /** Coalesces concurrent reconcile attempts; only one runs at a time. */ @@ -280,6 +294,8 @@ private void ackSegmentUpTo(long segmentId, org.apache.pulsar.client.api.Message } var unacked = sharedSegmentUnacked.get(segmentId); if (unacked == null) { + lastCumulativeAcked.merge(segmentId, position, + (a, b) -> a.compareTo(b) >= 0 ? a : b); future.thenAccept(c -> { if (v4Txn == null) { c.acknowledgeCumulativeAsync(position); @@ -361,6 +377,7 @@ CompletableFuture closeAsync() { .whenComplete((__, ___) -> { segmentConsumers.clear(); sharedSegmentUnacked.clear(); + lastCumulativeAcked.clear(); }); } @@ -478,20 +495,25 @@ private CompletableFuture subscribeAssigned(List assigned) // consumer): close our v4 consumer so the Exclusive lock is released and // the new owner can attach. Sealed-and-drained segments take a different // path: the receive loop closes them on TopicTerminated. + List> futures = new ArrayList<>(); for (var entry : segmentConsumers.entrySet()) { if (!assignedIds.contains(entry.getKey())) { - log.info().attr("segmentId", entry.getKey()) - .log("Closing consumer for segment removed from assignment"); - entry.getValue().thenAccept(c -> c.closeAsync()); - segmentConsumers.remove(entry.getKey()); - segmentBucketRanges.remove(entry.getKey()); - sharedSegmentUnacked.remove(entry.getKey()); - latestDelivered.remove(entry.getKey()); + long segmentId = entry.getKey(); + var existing = entry.getValue(); + log.info().attr("segmentId", segmentId) + .log("Releasing segment removed from assignment"); + segmentConsumers.remove(segmentId, existing); + segmentBucketRanges.remove(segmentId); + futures.add(drainAndClose(segmentId, existing) + .whenComplete((__, ___) -> { + sharedSegmentUnacked.remove(segmentId); + lastCumulativeAcked.remove(segmentId); + latestDelivered.remove(segmentId); + })); } } // Subscribe to newly-assigned segments. - List> futures = new ArrayList<>(); for (var seg : assigned) { // PIP-486: if the controller changed which entry-buckets we own on a segment we are // already subscribed to (a consumer joined or left), re-subscribe with the new ranges so @@ -503,12 +525,10 @@ private CompletableFuture subscribeAssigned(List assigned) .log("Re-subscribing segment for changed entry-bucket ownership"); segmentConsumers.remove(seg.segmentId(), existing); futures.add(segmentConsumers.computeIfAbsent(seg.segmentId(), id -> - // Await our own close before re-subscribing: the new subscribe (a different - // type, or ranges overlapping the old declaration) is rejected while the old - // consumer is still attached. - existing.handle((c, ex) -> c) - .thenCompose(c -> c != null ? c.closeAsync() - : CompletableFuture.completedFuture(null)) + // Drain what the application already holds, then await our own close before + // re-subscribing with the new mode: the subscription type cannot change + // while the old consumer is still attached. + drainAndClose(seg.segmentId(), existing) .thenCompose(__ -> createSegmentConsumerAsync(seg)))); } else { futures.add(segmentConsumers.computeIfAbsent(seg.segmentId(), @@ -520,6 +540,66 @@ private CompletableFuture subscribeAssigned(List assigned) return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)); } + /** + * PIP-486 release: pause the segment's receive loop (stop handing its messages to the + * application), wait until every message already handed out is acked, then close the v4 + * consumer — anything still prefetched goes back to the broker unread for in-order + * redelivery to the next owner. Deliberately unbounded: ownership must not move while the + * application still holds unacked messages, so a stalled application stalls the release, + * visibly, rather than silently breaking per-key order. Terminates early only when this + * consumer closes. + */ + private CompletableFuture drainAndClose( + long segmentId, CompletableFuture> existing) { + pausedSegments.add(segmentId); + return awaitReleaseDrained(segmentId, 0) + .thenCompose(__ -> existing.handle((c, ex) -> c)) + .thenCompose(c -> c != null ? c.closeAsync() : CompletableFuture.completedFuture(null)) + .whenComplete((__, ___) -> pausedSegments.remove(segmentId)); + } + + /** Poll until the segment's delivered messages are all acked (or this consumer closes). */ + private CompletableFuture awaitReleaseDrained(long segmentId, int pollCount) { + if (closed || isReleaseDrained(segmentId)) { + return CompletableFuture.completedFuture(null); + } + if (pollCount > 0 && pollCount % 200 == 0) { + // ~every 10s at the 50ms poll interval: make a stalled release observable. + log.warn().attr("segmentId", segmentId) + .log("Still waiting for the application to ack the segment's messages " + + "before releasing it"); + } + CompletableFuture result = new CompletableFuture<>(); + scheduler().schedule(() -> awaitReleaseDrained(segmentId, pollCount + 1) + .whenComplete((__, ex) -> { + if (ex != null) { + result.completeExceptionally(ex); + } else { + result.complete(null); + } + }), + 50, TimeUnit.MILLISECONDS); + return result; + } + + /** + * Whether everything this consumer handed to the application for the segment has been acked: + * bucket-shared segments track individual ids; whole (Exclusive) segments compare the + * cumulative-ack high-water mark against the latest delivered position. + */ + private boolean isReleaseDrained(long segmentId) { + var unacked = sharedSegmentUnacked.get(segmentId); + if (unacked != null) { + return unacked.isEmpty(); + } + org.apache.pulsar.client.api.MessageId delivered = latestDelivered.get(segmentId); + if (delivered == null) { + return true; // nothing was ever handed to the application + } + org.apache.pulsar.client.api.MessageId acked = lastCumulativeAcked.get(segmentId); + return acked != null && acked.compareTo(delivered) >= 0; + } + private CompletableFuture> createSegmentConsumerAsync( ActiveSegment segment) { PulsarClientImpl v4Client = client.v4Client(); @@ -604,7 +684,9 @@ private void startReceiveLoop(org.apache.pulsar.client.api.Consumer v4Consume // Re-arm only once the sink has room, so a slow consumer pauses this segment's // receive loop (and the v4 flow-control permits) instead of buffering unboundedly. messageSink.accept(new MessageV5<>(v4Msg, msgId)).thenRun(() -> { - if (!closed) { + // A paused segment is being released (PIP-486): stop pulling so the messages + // already handed to the application can drain before the close. + if (!closed && !pausedSegments.contains(segmentId)) { startReceiveLoop(v4Consumer, segmentId); } }); @@ -627,6 +709,7 @@ private void startReceiveLoop(org.apache.pulsar.client.api.Consumer v4Consume .log("Sealed segment drained, closing v4 consumer"); segmentConsumers.remove(segmentId); sharedSegmentUnacked.remove(segmentId); + lastCumulativeAcked.remove(segmentId); latestDelivered.remove(segmentId); v4Consumer.closeAsync(); return null; From c93f27a9c210077780afd64dcef9250983094db7 Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Fri, 17 Jul 2026 19:09:22 -0700 Subject: [PATCH 4/5] [improve][broker] PIP-486 (PR5): unit tests for bucket selector and dispatcher Adds EntryBucketConsumerSelectorTest (bucket index/canonical-hash math, deterministic add-order-independent spread, more-consumers-than-buckets, key-hash normalization, impacted-bucket reporting on membership changes, merged assignment snapshot) and PersistentEntryBucketDispatcherMultipleConsumersTest (boundary-list validation accept/reject cases). Reworks the PR4-era testEntryBucketDispatchRoutesByStampedRange against the new subclass: entry-bucket routing now lives in PersistentEntryBucketDispatcherMultipleConsumers and normalizes stamps to the bucket's canonical hash, so the test constructs the subclass with a boundary list and asserts canonical values. Also refreshes the stale bucket_ranges comment in PulsarApi.proto to the boundaries-for-all semantics. Assisted-by: Claude Code (Fable 5) --- .../EntryBucketConsumerSelectorTest.java | 189 ++++++++++++++++++ ...BucketDispatcherMultipleConsumersTest.java | 87 ++++++++ ...ckyKeyDispatcherMultipleConsumersTest.java | 29 +-- pulsar-common/src/main/proto/PulsarApi.proto | 10 +- 4 files changed, 299 insertions(+), 16 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelectorTest.java create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumersTest.java diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelectorTest.java new file mode 100644 index 0000000000000..0c89e172e3dd2 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBucketConsumerSelectorTest.java @@ -0,0 +1,189 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pulsar.client.api.Range; +import org.testng.annotations.Test; + +@Test(groups = "broker") +public class EntryBucketConsumerSelectorTest { + + private static final List FOUR_BUCKETS = List.of( + Range.of(0x0000, 0x3FFF), + Range.of(0x4000, 0x7FFF), + Range.of(0x8000, 0xBFFF), + Range.of(0xC000, 0xFFFF)); + + private static Consumer mockConsumer(String name, long id) { + Consumer consumer = mock(Consumer.class); + when(consumer.consumerName()).thenReturn(name); + when(consumer.consumerId()).thenReturn(id); + return consumer; + } + + @Test + public void testBucketIndexAndCanonicalHash() { + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + assertEquals(selector.bucketIndexOfHash(0x0000), 0); + assertEquals(selector.bucketIndexOfHash(0x3FFF), 0); + assertEquals(selector.bucketIndexOfHash(0x4000), 1); + assertEquals(selector.bucketIndexOfHash(0x7FFF), 1); + assertEquals(selector.bucketIndexOfHash(0x8000), 2); + assertEquals(selector.bucketIndexOfHash(0xC000), 3); + assertEquals(selector.bucketIndexOfHash(0xFFFF), 3); + + // Bucket 0's canonical hash is nudged to 1 (0 is the reserved "not set" value). + assertEquals(selector.canonicalHashOfBucket(0), 1); + assertEquals(selector.canonicalHashOfBucket(1), 0x4000); + assertEquals(selector.canonicalHashOfBucket(2), 0x8000); + assertEquals(selector.canonicalHashOfBucket(3), 0xC000); + + assertEquals(selector.canonicalHashOf(0x1234), 1); + assertEquals(selector.canonicalHashOf(0x4000), 0x4000); + assertEquals(selector.canonicalHashOf(0x7FFF), 0x4000); + assertEquals(selector.canonicalHashOf(0xFFFF), 0xC000); + } + + @Test + public void testDeterministicSpreadIsAddOrderIndependent() { + Consumer a = mockConsumer("consumer-a", 1); + Consumer b = mockConsumer("consumer-b", 2); + + EntryBucketConsumerSelector selector1 = new EntryBucketConsumerSelector(FOUR_BUCKETS); + selector1.addConsumer(a); + selector1.addConsumer(b); + + EntryBucketConsumerSelector selector2 = new EntryBucketConsumerSelector(FOUR_BUCKETS); + selector2.addConsumer(b); + selector2.addConsumer(a); + + // Sorted by name: "consumer-a" owns buckets 0-1, "consumer-b" owns buckets 2-3, + // regardless of connection order. + for (int bucket = 0; bucket < 4; bucket++) { + int hash = selector1.canonicalHashOfBucket(bucket); + Consumer expected = bucket < 2 ? a : b; + assertSame(selector1.select(hash), expected); + assertSame(selector2.select(hash), expected); + } + } + + @Test + public void testSpreadWithMoreConsumersThanBuckets() { + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + for (int i = 0; i < 6; i++) { + selector.addConsumer(mockConsumer("consumer-" + i, i)); + } + // Every bucket has exactly one owner; only 4 of the 6 consumers own a bucket. + Set owners = new HashSet<>(); + for (int bucket = 0; bucket < 4; bucket++) { + Consumer owner = selector.select(selector.canonicalHashOfBucket(bucket)); + assertNotNull(owner); + owners.add(owner); + } + assertEquals(owners.size(), 4); + } + + @Test + public void testMakeStickyKeyHashNormalizesToCanonical() { + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + for (int i = 0; i < 100; i++) { + byte[] key = ("key-" + i).getBytes(StandardCharsets.UTF_8); + int rawHash = StickyKeyConsumerSelectorUtils.makeStickyKeyHash(key, selector.getKeyHashRange()); + int hash = selector.makeStickyKeyHash(key); + assertEquals(hash, selector.canonicalHashOf(rawHash)); + assertEquals(selector.bucketIndexOfHash(hash), selector.bucketIndexOfHash(rawHash)); + } + } + + @Test + public void testMembershipChangeReportsImpactedBuckets() { + Consumer a = mockConsumer("consumer-a", 1); + Consumer b = mockConsumer("consumer-b", 2); + + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + selector.addConsumer(a).join(); + + // consumer-b joining takes buckets 2-3 away from consumer-a. + ImpactedConsumersResult impacted = selector.addConsumer(b).join().get(); + Map removed = new HashMap<>(); + impacted.processUpdatedHashRanges((consumer, ranges, op) -> { + if (op == ImpactedConsumersResult.OperationType.REMOVE) { + removed.put(consumer, ranges); + } + }); + assertEquals(removed.keySet(), Set.of(a)); + UpdatedHashRanges lostByA = removed.get(a); + assertFalse(lostByA.containsStickyKey(selector.canonicalHashOfBucket(0))); + assertFalse(lostByA.containsStickyKey(selector.canonicalHashOfBucket(1))); + assertTrue(lostByA.containsStickyKey(selector.canonicalHashOfBucket(2))); + assertTrue(lostByA.containsStickyKey(selector.canonicalHashOfBucket(3))); + + // consumer-b leaving hands buckets 2-3 back: the removal is attributed to consumer-b. + ImpactedConsumersResult afterLeave = selector.removeConsumer(b).get(); + Map removedOnLeave = new HashMap<>(); + afterLeave.processUpdatedHashRanges((consumer, ranges, op) -> { + if (op == ImpactedConsumersResult.OperationType.REMOVE) { + removedOnLeave.put(consumer, ranges); + } + }); + assertEquals(removedOnLeave.keySet(), Set.of(b)); + for (int bucket = 0; bucket < 4; bucket++) { + assertSame(selector.select(selector.canonicalHashOfBucket(bucket)), a); + } + } + + @Test + public void testSelectWithNoConsumers() { + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + assertNull(selector.select(1)); + + Consumer a = mockConsumer("consumer-a", 1); + selector.addConsumer(a).join(); + selector.removeConsumer(a); + assertNull(selector.select(1)); + } + + @Test + public void testConsumerHashAssignmentsSnapshot() { + Consumer a = mockConsumer("consumer-a", 1); + Consumer b = mockConsumer("consumer-b", 2); + EntryBucketConsumerSelector selector = new EntryBucketConsumerSelector(FOUR_BUCKETS); + selector.addConsumer(a).join(); + selector.addConsumer(b).join(); + + // Contiguous bucket slices are merged: a owns buckets 0-1, b owns buckets 2-3. + Map> ranges = selector.getConsumerHashAssignmentsSnapshot().getRangesByConsumer(); + assertEquals(ranges.get(a), List.of(Range.of(0x0000, 0x7FFF))); + assertEquals(ranges.get(b), List.of(Range.of(0x8000, 0xFFFF))); + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumersTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumersTest.java new file mode 100644 index 0000000000000..e785b0ff3adbe --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentEntryBucketDispatcherMultipleConsumersTest.java @@ -0,0 +1,87 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service.persistent; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; +import java.util.List; +import org.apache.pulsar.client.api.Range; +import org.apache.pulsar.common.api.proto.KeySharedMeta; +import org.apache.pulsar.common.api.proto.KeySharedMode; +import org.testng.annotations.Test; + +@Test(groups = "broker") +public class PersistentEntryBucketDispatcherMultipleConsumersTest { + + private static KeySharedMeta ksmWithRanges(int[][] ranges) { + KeySharedMeta ksm = new KeySharedMeta().setKeySharedMode(KeySharedMode.STICKY) + .setEntryBucketDispatch(true); + for (int[] range : ranges) { + ksm.addHashRange().setStart(range[0]).setEnd(range[1]); + } + return ksm; + } + + @Test + public void testValidBoundaries() { + List ranges = PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x3FFF}, {0x4000, 0x7FFF}, {0x8000, 0xBFFF}, {0xC000, 0xFFFF}})); + assertEquals(ranges, List.of( + Range.of(0x0000, 0x3FFF), + Range.of(0x4000, 0x7FFF), + Range.of(0x8000, 0xBFFF), + Range.of(0xC000, 0xFFFF))); + + // A single bucket spanning the whole ring is valid. + assertEquals(PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0xFFFF}})), List.of(Range.of(0x0000, 0xFFFF))); + } + + @Test + public void testInvalidBoundariesAreRejected() { + // No boundaries at all. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{}))); + // Not starting at 0. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0001, 0xFFFF}}))); + // Gap between buckets. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x3FFF}, {0x5000, 0xFFFF}}))); + // Overlapping buckets. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x7FFF}, {0x4000, 0xFFFF}}))); + // End before start. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x3FFF}, {0x4000, 0x2000}}))); + // Not tiling the whole 16-bit ring. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x3FFF}, {0x4000, 0x7FFF}}))); + // Bucket 0 too narrow to hold the canonical hash 1. + assertThrows(IllegalArgumentException.class, () -> + PersistentEntryBucketDispatcherMultipleConsumers.validateBucketBoundaries( + ksmWithRanges(new int[][]{{0x0000, 0x0000}, {0x0001, 0xFFFF}}))); + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java index 8e68500f80ac7..2293fad12285e 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java @@ -850,29 +850,34 @@ public boolean trackDelayedDelivery(long ledgerId, long entryId, MessageMetadata @Test public void testEntryBucketDispatchRoutesByStampedRange() { - // PIP-486: with entryBucketDispatch set on the subscription's KeySharedMeta, a stamped entry - // routes as a whole by its entry-bucket hash (entry_hash_min), not by the message key. - PersistentStickyKeyDispatcherMultipleConsumers bucketDispatcher = - new PersistentStickyKeyDispatcherMultipleConsumers(topicMock, cursorMock, subscriptionMock, - configMock, new KeySharedMeta().setKeySharedMode(KeySharedMode.STICKY) - .setEntryBucketDispatch(true)); + // PIP-486: the entry-bucket dispatcher routes a stamped entry as a whole by its entry-bucket + // (entry_hash_min normalized to the bucket's canonical hash), not by the message key. + KeySharedMeta ksm = new KeySharedMeta().setKeySharedMode(KeySharedMode.STICKY) + .setEntryBucketDispatch(true); + ksm.addHashRange().setStart(0x0000).setEnd(0x3FFF); + ksm.addHashRange().setStart(0x4000).setEnd(0x7FFF); + ksm.addHashRange().setStart(0x8000).setEnd(0xBFFF); + ksm.addHashRange().setStart(0xC000).setEnd(0xFFFF); + PersistentEntryBucketDispatcherMultipleConsumers bucketDispatcher = + new PersistentEntryBucketDispatcherMultipleConsumers(topicMock, cursorMock, subscriptionMock, + configMock, ksm); EntryImpl entry = createEntry(1, 1, "msg", 1, "some-key"); try { MessageMetadata stamped = new MessageMetadata() .setProducerName("p").setSequenceId(1).setPublishTime(1) - .setEntryHashMin(0x1234).setEntryHashMax(0x5678); + .setEntryHashMin(0x4567).setEntryHashMax(0x4FFF); assertEquals(bucketDispatcher.getStickyKeyHash( - EntryAndMetadata.create(entry, stamped)), 0x1234); + EntryAndMetadata.create(entry, stamped)), 0x4000); - // entry_hash_min == 0 collides with the reserved "hash not set" sentinel, so it is nudged - // to 1, which is still inside bucket 0. + // Bucket 0's canonical hash is nudged to 1 (0 is the reserved "hash not set" sentinel). MessageMetadata zero = new MessageMetadata() .setProducerName("p").setSequenceId(1).setPublishTime(1) - .setEntryHashMin(0).setEntryHashMax(0); + .setEntryHashMin(0).setEntryHashMax(0x3FFF); assertEquals(bucketDispatcher.getStickyKeyHash( EntryAndMetadata.create(entry, zero)), 1); - // Unstamped entries (non-batched messages) fall back to the message's sticky-key hash. + // Unstamped entries (non-batched messages) fall back to the message key's hash, + // normalized to the same canonical bucket value. MessageMetadata unstamped = new MessageMetadata() .setProducerName("p").setSequenceId(1).setPublishTime(1).setPartitionKey("some-key"); assertEquals(bucketDispatcher.getStickyKeyHash( diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto index 874ac9713ba56..5d30972d93791 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -950,10 +950,12 @@ message ScalableAssignedSegment { required uint32 hash_end = 3; // Fully-qualified segment:// topic name the consumer should attach to. required string segment_topic = 4; - // PIP-486: the entry-bucket hash ranges (16-bit, inclusive) this consumer owns within the - // segment. Empty means the consumer owns the whole segment (single bucket) and subscribes - // Shared; non-empty means the segment is shared by bucket, and the consumer subscribes - // Key_Shared STICKY declaring exactly these ranges. + // PIP-486: the segment's full entry-bucket boundary list (16-bit inclusive ranges tiling + // 0x0000-0xFFFF). Empty means this consumer is the segment's sole owner and subscribes + // Exclusive; non-empty means the segment is shared by entry-bucket, and the consumer + // subscribes Key_Shared STICKY with entryBucketDispatch, declaring exactly this boundary + // list (identical for every sharer). The bucket-to-consumer spread is computed broker-side + // from the live membership, so handoffs need no client protocol. repeated IntRange bucket_ranges = 5; } From 90573d782806a4448eca949aad97420616db5d9c Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Mon, 20 Jul 2026 17:12:53 -0700 Subject: [PATCH 5/5] [fix][client] Scalable stream consumer: keep ack routing on the draining segment consumer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A release drain completes only when the application acks what it was already handed, but both drain sites broke the ack path they depend on (found in review of the PR): - Dropping a segment removed its segmentConsumers entry before the drain, so ackSegmentUpTo hit the null-entry early-return: the translation queue never emptied, the cumulative high-water never advanced, and the unbounded drain waited forever — the consumer never released the segment and its new owner could never attach. - A mode flip immediately repopulated the slot with the re-subscribe chain, so mid-drain acks were deferred onto the future new consumer: misdirected individual acks, and on a whole-to-shared flip cumulative acks against a Key_Shared consumer, which v4 rejects — the messages were redelivered after the flip and keys saw duplicates. Fix: a drainingConsumers map keeps the ack path pointed at the consumer being drained. drainAndClose registers the consumer there before vacating the segmentConsumers slot (owning the removal so the two stay ordered — also required because the flip path's computeIfAbsent mapping function must not mutate the map reentrantly), and clears it after the close; ackSegmentUpTo consults it first. closeAsync also sweeps draining consumers so close does not complete before they are gone. Two e2e tests pin the behavior, both verified to fail before the fix: - testSegmentReleaseWithPendingAcksHandsOffCleanly: a consumer loses a whole segment while the application holds unacked messages; the ack arrives mid-drain. Before the fix the release deadlocked and the new owner timed out waiting to attach. - testConsumerJoiningWithSlowAcksPreservesPerKeyOrder: the Exclusive to Key_Shared flip with a non-empty drain window. Before the fix the mid-drain ack was lost and redelivery produced per-key duplicates. Assisted-by: Claude Code (Fable 5) --- .../api/v5/V5EntryBucketDispatchTest.java | 187 +++++++++++++++++- .../impl/v5/ScalableStreamConsumer.java | 55 ++++-- 2 files changed, 226 insertions(+), 16 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java index 1978dc607cfae..f6727522c3285 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java @@ -28,7 +28,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; import lombok.Cleanup; import org.apache.pulsar.client.api.v5.config.SubscriptionInitialPosition; import org.apache.pulsar.client.api.v5.schema.Schema; @@ -256,16 +258,199 @@ public void testConsumerJoiningMidTrafficPreservesPerKeyOrder() throws Exception } } + @Test + public void testSegmentReleaseWithPendingAcksHandsOffCleanly() throws Exception { + // The release (drop) path: a consumer loses a whole segment in a rebalance while the + // application still holds delivered-but-unacked messages. The release must wait for those + // acks — and the acks, arriving while the segment is already vacated from the consumer's + // active set, must still reach the draining consumer, or the release never completes and the + // new owner never gets the segment. + String topic = newScalableTopic(2); + admin.scalableTopics().setAutoScalePolicy(topic, + AutoScalePolicyOverride.builder().enabled(false).build()); + String subscription = "segment-release"; + + @Cleanup + Producer producer = v5Client.newProducer(Schema.string()) + .topic(topic) + .create(); + @Cleanup + StreamConsumer a = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribe(); + + // Phase 1: A owns both segments; deliver keyed traffic and hold every ack. + List keys = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + keys.add("key-" + i); + } + int perKey = 10; + Map> sent = new HashMap<>(); + for (String k : keys) { + sent.put(k, new ArrayList<>()); + } + for (int i = 0; i < perKey; i++) { + for (String k : keys) { + String value = k + "-" + i; + producer.newMessage().key(k).value(value).send(); + sent.get(k).add(value); + } + } + Map> aGot = new ConcurrentHashMap<>(); + Message last = null; + for (int i = 0; i < keys.size() * perKey; i++) { + Message msg = a.receive(Duration.ofSeconds(5)); + assertNotNull(msg, "missed message #" + i + " in phase 1"); + String key = msg.key().orElseThrow(() -> new AssertionError("missing key")); + aGot.computeIfAbsent(key, __ -> new ArrayList<>()).add(msg.value()); + last = msg; + } + + // B joins: the controller rebalances to one segment each, so A must release one whole + // segment — but it cannot until the application acks what it already received. Subscribe + // asynchronously: B's attach only completes once A releases, which needs the ack below. + CompletableFuture> bFuture = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribeAsync(); + // Let the assignment update land and the release drain start waiting, so the ack below + // genuinely arrives mid-drain. + Thread.sleep(1500); + a.acknowledgeCumulative(last.id()); + // Bounded: if the release/flip wedges, fail here instead of hanging the suite. + @Cleanup + StreamConsumer b = bFuture.get(30, TimeUnit.SECONDS); + + // Phase 2: more traffic; the moved segment's share must now flow to B. + for (int i = perKey; i < perKey * 2; i++) { + for (String k : keys) { + String value = k + "-" + i; + producer.newMessage().key(k).value(value).send(); + sent.get(k).add(value); + } + } + Map> bGot = new ConcurrentHashMap<>(); + Thread ta = drainOrdered(a, aGot, Duration.ofSeconds(5)); + Thread tb = drainOrdered(b, bGot, Duration.ofSeconds(5)); + ta.join(); + tb.join(); + + assertFalse(bGot.isEmpty(), "consumer B received nothing — the segment was never released"); + // Phase 1 was fully acked before the release completed, so nothing is redelivered: for every + // key, A's deliveries followed by B's are exactly the sent sequence. + for (String k : keys) { + List combined = new ArrayList<>(aGot.getOrDefault(k, List.of())); + combined.addAll(bGot.getOrDefault(k, List.of())); + assertEquals(combined, sent.get(k), "per-key order across the release for key=" + k); + } + } + + @Test + public void testConsumerJoiningWithSlowAcksPreservesPerKeyOrder() throws Exception { + // The mode-flip path with a non-empty drain window: when B joins, A flips the segment from + // Exclusive to Key_Shared. Acks issued while the flip's drain is waiting must reach the old + // consumer being drained — routed anywhere else they are lost (a cumulative ack is invalid on + // the new Key_Shared consumer), the broker redelivers, and keys see duplicates. + String topic = newScalableTopic(1); + admin.scalableTopics().setAutoScalePolicy(topic, + AutoScalePolicyOverride.builder().enabled(false).build()); + String subscription = "slow-ack-join"; + + @Cleanup + Producer producer = v5Client.newProducer(Schema.string()) + .topic(topic) + .create(); + @Cleanup + StreamConsumer a = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribe(); + + // Phase 1: deliver to A and hold every ack. + List keys = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + keys.add("key-" + i); + } + int perKey = 5; + Map> sent = new HashMap<>(); + for (String k : keys) { + sent.put(k, new ArrayList<>()); + } + for (int i = 0; i < perKey; i++) { + for (String k : keys) { + String value = k + "-" + i; + producer.newMessage().key(k).value(value).send(); + sent.get(k).add(value); + } + } + Map> aGot = new ConcurrentHashMap<>(); + Message last = null; + for (int i = 0; i < keys.size() * perKey; i++) { + Message msg = a.receive(Duration.ofSeconds(5)); + assertNotNull(msg, "missed message #" + i + " in phase 1"); + String key = msg.key().orElseThrow(() -> new AssertionError("missing key")); + aGot.computeIfAbsent(key, __ -> new ArrayList<>()).add(msg.value()); + last = msg; + } + + // B joins → A must flip the segment to Key_Shared, which drains first. Subscribe + // asynchronously (B's attach completes only after A's flip) and ack mid-drain. + CompletableFuture> bFuture = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribeAsync(); + Thread.sleep(1500); + a.acknowledgeCumulative(last.id()); + // Bounded: if the release/flip wedges, fail here instead of hanging the suite. + @Cleanup + StreamConsumer b = bFuture.get(30, TimeUnit.SECONDS); + + // Phase 2: post-flip traffic, spread across both owners by bucket. + for (int i = perKey; i < perKey * 2; i++) { + for (String k : keys) { + String value = k + "-" + i; + producer.newMessage().key(k).value(value).send(); + sent.get(k).add(value); + } + } + Map> bGot = new ConcurrentHashMap<>(); + Thread ta = drainOrdered(a, aGot, Duration.ofSeconds(5)); + Thread tb = drainOrdered(b, bGot, Duration.ofSeconds(5)); + ta.join(); + tb.join(); + + assertFalse(bGot.isEmpty(), "consumer B received nothing — the handoff did not happen"); + // The mid-drain ack reached the draining consumer, so phase 1 is acked on the broker and + // never redelivered: A's deliveries followed by B's are exactly the sent sequence — no + // duplicates, no loss, per key. + for (String k : keys) { + List combined = new ArrayList<>(aGot.getOrDefault(k, List.of())); + combined.addAll(bGot.getOrDefault(k, List.of())); + assertEquals(combined, sent.get(k), "per-key order across the slow-ack flip for key=" + k); + } + } + /** * Drains until idle, recording values per key in arrival order. Acks every message as it is * processed — like an order-sensitive application would — so a release drain during a handoff * can complete promptly. */ private Thread drainOrdered(StreamConsumer consumer, Map> into) { + return drainOrdered(consumer, into, Duration.ofSeconds(2)); + } + + /** Variant with a configurable idle timeout, for tests whose handoff takes a few seconds. */ + private Thread drainOrdered(StreamConsumer consumer, Map> into, + Duration idleTimeout) { Thread t = new Thread(() -> { try { while (true) { - Message msg = consumer.receive(Duration.ofSeconds(2)); + Message msg = consumer.receive(idleTimeout); if (msg == null) { return; } diff --git a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java index 9d1e062797eae..bab570b8489c5 100644 --- a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java +++ b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableStreamConsumer.java @@ -33,6 +33,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Stream; import org.apache.pulsar.client.api.KeySharedPolicy; import org.apache.pulsar.client.api.Range; import org.apache.pulsar.client.api.SubscriptionType; @@ -124,6 +125,15 @@ final class ScalableStreamConsumer */ private final Set pausedSegments = ConcurrentHashMap.newKeySet(); + /** + * PIP-486: segment consumers currently draining for a release. A drain completes only when the + * application acks what it was handed, so the ack path must keep reaching the old consumer after + * its {@link #segmentConsumers} slot is vacated (or repopulated with the re-subscribe chain on a + * mode flip) — {@link #ackSegmentUpTo} consults this map first. + */ + private final ConcurrentHashMap>> + drainingConsumers = new ConcurrentHashMap<>(); + /** Latest controller assignment; {@link #reconcile()} converges the segment consumers onto it. */ private volatile List latestAssignment; /** Coalesces concurrent reconcile attempts; only one runs at a time. */ @@ -288,7 +298,11 @@ public void acknowledgeCumulative(MessageId messageId, Transaction txn) { */ private void ackSegmentUpTo(long segmentId, org.apache.pulsar.client.api.MessageId position, org.apache.pulsar.client.api.transaction.Transaction v4Txn) { - var future = segmentConsumers.get(segmentId); + // A draining consumer takes precedence: during a release the segment's slot in + // segmentConsumers is vacated (or already holds the re-subscribe chain), but the acks that + // complete the drain must still reach the consumer being drained. + var draining = drainingConsumers.get(segmentId); + var future = draining != null ? draining : segmentConsumers.get(segmentId); if (future == null) { return; } @@ -366,16 +380,19 @@ CompletableFuture closeAsync() { receiveQueue.close(); session.close(); + // Draining consumers live outside segmentConsumers; their drain loops observe the closed + // flag and finish on their own, but they are closed here too (idempotent) so this future + // does not complete before they are gone. List> futures = new ArrayList<>(); - for (var future : segmentConsumers.values()) { - futures.add(future - .handle((consumer, ex) -> consumer) - .thenCompose(consumer -> consumer != null ? consumer.closeAsync() - : CompletableFuture.completedFuture(null))); - } + Stream.concat(segmentConsumers.values().stream(), drainingConsumers.values().stream()) + .forEach(future -> futures.add(future + .handle((consumer, ex) -> consumer) + .thenCompose(consumer -> consumer != null ? consumer.closeAsync() + : CompletableFuture.completedFuture(null)))); return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)) .whenComplete((__, ___) -> { segmentConsumers.clear(); + drainingConsumers.clear(); sharedSegmentUnacked.clear(); lastCumulativeAcked.clear(); }); @@ -502,7 +519,6 @@ private CompletableFuture subscribeAssigned(List assigned) var existing = entry.getValue(); log.info().attr("segmentId", segmentId) .log("Releasing segment removed from assignment"); - segmentConsumers.remove(segmentId, existing); segmentBucketRanges.remove(segmentId); futures.add(drainAndClose(segmentId, existing) .whenComplete((__, ___) -> { @@ -523,13 +539,14 @@ private CompletableFuture subscribeAssigned(List assigned) && !seg.ownedBucketRanges().equals(segmentBucketRanges.get(seg.segmentId()))) { log.info().attr("segmentId", seg.segmentId()) .log("Re-subscribing segment for changed entry-bucket ownership"); - segmentConsumers.remove(seg.segmentId(), existing); + // Drain what the application already holds, then await our own close before + // re-subscribing with the new mode: the subscription type cannot change while the + // old consumer is still attached. The drain must start outside computeIfAbsent — + // it vacates this segment's slot itself, and ConcurrentHashMap forbids mutating + // the map from within the mapping function. + var drained = drainAndClose(seg.segmentId(), existing); futures.add(segmentConsumers.computeIfAbsent(seg.segmentId(), id -> - // Drain what the application already holds, then await our own close before - // re-subscribing with the new mode: the subscription type cannot change - // while the old consumer is still attached. - drainAndClose(seg.segmentId(), existing) - .thenCompose(__ -> createSegmentConsumerAsync(seg)))); + drained.thenCompose(__ -> createSegmentConsumerAsync(seg)))); } else { futures.add(segmentConsumers.computeIfAbsent(seg.segmentId(), id -> createSegmentConsumerAsync(seg))); @@ -551,11 +568,19 @@ private CompletableFuture subscribeAssigned(List assigned) */ private CompletableFuture drainAndClose( long segmentId, CompletableFuture> existing) { + // Register as draining before vacating the segmentConsumers slot, so at every instant the + // ack path resolves to the consumer being drained — the drain only completes through those + // acks. The slot removal lives here (not at the call sites) to keep that ordering. + drainingConsumers.put(segmentId, existing); + segmentConsumers.remove(segmentId, existing); pausedSegments.add(segmentId); return awaitReleaseDrained(segmentId, 0) .thenCompose(__ -> existing.handle((c, ex) -> c)) .thenCompose(c -> c != null ? c.closeAsync() : CompletableFuture.completedFuture(null)) - .whenComplete((__, ___) -> pausedSegments.remove(segmentId)); + .whenComplete((__, ___) -> { + drainingConsumers.remove(segmentId, existing); + pausedSegments.remove(segmentId); + }); } /** Poll until the segment's delivered messages are all acked (or this consumer closes). */