[improve][broker] PIP-486 (PR4): bucket-shared consumption for scale-up fan-out (#26173)

Author: merlimat

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java

@@ -69,6 +69,9 @@ 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;
@@ -84,6 +87,7 @@ public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDi
 
         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;
@@ -634,6 +638,23 @@ 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);
         }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java

@@ -21,6 +21,7 @@
 import io.github.merlimat.slog.Logger;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashSet;
@@ -37,6 +38,7 @@
 import org.apache.pulsar.broker.resources.ScalableTopicResources;
 import org.apache.pulsar.broker.service.TransportCnx;
 import org.apache.pulsar.common.naming.TopicName;
+import org.apache.pulsar.common.scalable.HashRange;
 import org.apache.pulsar.common.scalable.SegmentInfo;
 import org.apache.pulsar.common.scalable.SegmentTopicName;
 import org.apache.pulsar.common.util.Backoff;
@@ -533,9 +535,9 @@ private void evictExpiredConsumer(String consumerName) {
      * Compute a balanced assignment of segments to consumers.
      *
      * <p>Strategy: sort segments by hash range, then segment id (tiebreak), sort consumers by
-     * name, then round-robin. Deterministic: the same inputs always produce the same output,
-     * so a new leader recomputing assignments after failover gets the same result as the old
-     * leader.
+     * name, then "segments first, entry-buckets absorb the surplus" (see the inline comment).
+     * Deterministic: the same inputs always produce the same output, so a new leader recomputing
+     * assignments after failover gets the same result as the old leader.
      *
      * <p><b>DAG replay.</b> The assignment includes every <em>sealed</em> segment in the
      * DAG. A fresh EARLIEST subscription needs to read messages produced before it joined,
@@ -577,20 +579,69 @@ Map<ConsumerSession, ConsumerAssignment> computeAssignment(
             assignmentLists.put(consumer, new ArrayList<>());
         }
 
-        int consumerIndex = 0;
-        for (SegmentInfo segment : sortedSegments) {
-            TopicName segmentTopic = SegmentTopicName.fromParent(topicName, segment.hashRange(),
-                    segment.segmentId());
-            // PIP-486: assign each whole segment to a single consumer for efficient single-active
-            // (Exclusive) dispatch — no per-bucket pending tracking. A segment's entry-buckets let it be
-            // *shared* across multiple consumers, but fanning a segment out into Key_Shared bucket
-            // ownership is a controller-driven scale-up action handled separately; by default one
-            // consumer owns the whole segment. Empty bucketRanges signals the client to subscribe
-            // Exclusive.
-            ConsumerSession consumer = sortedConsumers.get(consumerIndex % sortedConsumers.size());
-            assignmentLists.get(consumer).add(new ConsumerAssignment.AssignedSegment(
-                    segment.segmentId(), segment.hashRange(), segmentTopic.toString(), List.of()));
-            consumerIndex++;
+        // PIP-486: "segments first, entry-buckets absorb the surplus". While consumers don't outnumber
+        // segments, each whole segment goes to a single consumer for efficient single-active
+        // (Exclusive) dispatch — no per-bucket pending tracking; empty bucketRanges signals the client
+        // to subscribe Exclusive. Only when consumers outnumber segments does the controller fan
+        // segments out by entry-bucket: each owner of a shared segment takes a contiguous slice of its
+        // buckets and subscribes Key_Shared STICKY declaring exactly those ranges. A segment absorbs at
+        // most bucketCount() consumers; consumers beyond the topic's total bucket capacity stay idle.
+        int segmentCount = sortedSegments.size();
+        int consumerCount = sortedConsumers.size();
+        if (consumerCount <= segmentCount) {
+            int consumerIndex = 0;
+            for (SegmentInfo segment : sortedSegments) {
+                TopicName segmentTopic = SegmentTopicName.fromParent(topicName, segment.hashRange(),
+                        segment.segmentId());
+                ConsumerSession consumer = sortedConsumers.get(consumerIndex % sortedConsumers.size());
+                assignmentLists.get(consumer).add(new ConsumerAssignment.AssignedSegment(
+                        segment.segmentId(), segment.hashRange(), segmentTopic.toString(), List.of()));
+                consumerIndex++;
+            }
+        } else {
+            // Per-segment owner counts: one each, then hand the surplus to segments that still have
+            // bucket capacity, round-robin in segment order.
+            int[] owners = new int[segmentCount];
+            Arrays.fill(owners, 1);
+            int surplus = consumerCount - segmentCount;
+            boolean anyCapacityLeft = true;
+            while (surplus > 0 && anyCapacityLeft) {
+                anyCapacityLeft = false;
+                for (int i = 0; i < segmentCount && surplus > 0; i++) {
+                    if (owners[i] < sortedSegments.get(i).bucketCount()) {
+                        owners[i]++;
+                        surplus--;
+                        anyCapacityLeft = true;
+                    }
+                }
+            }
+            int consumerIndex = 0;
+            for (int i = 0; i < segmentCount; i++) {
+                SegmentInfo segment = sortedSegments.get(i);
+                TopicName segmentTopic = SegmentTopicName.fromParent(topicName, segment.hashRange(),
+                        segment.segmentId());
+                int k = owners[i];
+                if (k == 1) {
+                    ConsumerSession consumer = sortedConsumers.get(consumerIndex++);
+                    assignmentLists.get(consumer).add(new ConsumerAssignment.AssignedSegment(
+                            segment.segmentId(), segment.hashRange(), segmentTopic.toString(), List.of()));
+                } else {
+                    List<HashRange> buckets = EntryBucketSplits.ranges(segment.entryBucketSplits());
+                    int base = buckets.size() / k;
+                    int extra = buckets.size() % k;
+                    int from = 0;
+                    for (int c = 0; c < k; c++) {
+                        int size = base + (c < extra ? 1 : 0);
+                        List<HashRange> 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));
+                    }
+                }
+            }
+            // Consumers past consumerIndex found no segment with spare bucket capacity: they keep an
+            // empty assignment (idle) until the layout or the group changes.
         }
 
         Map<ConsumerSession, ConsumerAssignment> result = new LinkedHashMap<>();

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java

@@ -71,6 +71,7 @@
 import org.apache.pulsar.broker.ServiceConfiguration;
 import org.apache.pulsar.broker.service.BrokerService;
 import org.apache.pulsar.broker.service.Consumer;
+import org.apache.pulsar.broker.service.EntryAndMetadata;
 import org.apache.pulsar.broker.service.EntryBatchIndexesAcks;
 import org.apache.pulsar.broker.service.EntryBatchSizes;
 import org.apache.pulsar.broker.service.PendingAcksMap;
@@ -847,6 +848,59 @@ 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));
+        EntryImpl entry = createEntry(1, 1, "msg", 1, "some-key");
+        try {
+            MessageMetadata stamped = new MessageMetadata()
+                    .setProducerName("p").setSequenceId(1).setPublishTime(1)
+                    .setEntryHashMin(0x1234).setEntryHashMax(0x5678);
+            assertEquals(bucketDispatcher.getStickyKeyHash(
+                    EntryAndMetadata.create(entry, stamped)), 0x1234);
+
+            // entry_hash_min == 0 collides with the reserved "hash not set" sentinel, so it is nudged
+            // to 1, which is still inside bucket 0.
+            MessageMetadata zero = new MessageMetadata()
+                    .setProducerName("p").setSequenceId(1).setPublishTime(1)
+                    .setEntryHashMin(0).setEntryHashMax(0);
+            assertEquals(bucketDispatcher.getStickyKeyHash(
+                    EntryAndMetadata.create(entry, zero)), 1);
+
+            // Unstamped entries (non-batched messages) fall back to the message's sticky-key hash.
+            MessageMetadata unstamped = new MessageMetadata()
+                    .setProducerName("p").setSequenceId(1).setPublishTime(1).setPartitionKey("some-key");
+            assertEquals(bucketDispatcher.getStickyKeyHash(
+                            EntryAndMetadata.create(entry, unstamped)),
+                    bucketDispatcher.getSelector().makeStickyKeyHash("some-key".getBytes(UTF_8)));
+        } finally {
+            entry.release();
+        }
+    }
+
+    @Test
+    public void testStampedEntryStillRoutesByKeyWithoutEntryBucketDispatch() {
+        // Without the entryBucketDispatch flag (any plain Key_Shared subscription), a stamped entry
+        // keeps dispatching by the message key — the stamp is ignored.
+        EntryImpl entry = createEntry(1, 1, "msg", 1, "some-key");
+        try {
+            MessageMetadata stamped = new MessageMetadata()
+                    .setProducerName("p").setSequenceId(1).setPublishTime(1)
+                    .setPartitionKey("some-key")
+                    .setEntryHashMin(0x1234).setEntryHashMax(0x5678);
+            assertEquals(persistentDispatcher.getStickyKeyHash(
+                            EntryAndMetadata.create(entry, stamped)),
+                    persistentDispatcher.getSelector().makeStickyKeyHash("some-key".getBytes(UTF_8)));
+        } finally {
+            entry.release();
+        }
+    }
+
     private EntryImpl createEntry(long ledgerId, long entryId, String message, long sequenceId) {
         return createEntry(ledgerId, entryId, message, sequenceId, "testKey");
     }

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java

@@ -27,6 +27,8 @@
 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;
@@ -40,6 +42,7 @@
 import org.apache.pulsar.broker.resources.ScalableTopicResources;
 import org.apache.pulsar.broker.service.TransportCnx;
 import org.apache.pulsar.common.naming.TopicName;
+import org.apache.pulsar.common.scalable.HashRange;
 import org.awaitility.Awaitility;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
@@ -444,27 +447,74 @@ public void testSingleBucketSegmentsHaveNoBucketRanges() throws Exception {
     }
 
     @Test
-    public void testBucketedSegmentIsAssignedWholeToOneConsumer() throws Exception {
-        // One segment with N=4 entry-buckets (budget 4 / 1 segment). Even with several consumers, the
-        // controller assigns the whole segment to a single consumer with empty bucketRanges (efficient
-        // single-active / Exclusive dispatch); fanning it out into per-bucket Key_Shared ownership is a
-        // separate controller-driven scale-up action.
-        SubscriptionCoordinator c = new SubscriptionCoordinator("test-sub", topicName,
-                SegmentLayout.fromMetadata(ScalableTopicController.createInitialMetadata(1, 4, Map.of())),
-                resources, scheduler, Duration.ofMillis(200));
+    public void testLoneConsumerOwnsBucketedSegmentWhole() throws Exception {
+        // One segment with N=4 entry-buckets, one consumer: no surplus, so the segment stays whole
+        // (empty bucketRanges -> Exclusive single-active dispatch).
+        SubscriptionCoordinator c = bucketedCoordinator();
+        Map<ConsumerSession, ConsumerAssignment> result =
+                c.registerConsumer("consumer-1", 1L, mock(TransportCnx.class)).get();
+
+        ConsumerAssignment assignment = findByName(result, "consumer-1");
+        assertEquals(assignment.assignedSegments().size(), 1);
+        assertTrue(assignment.assignedSegments().get(0).bucketRanges().isEmpty());
+    }
+
+    @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.
+        SubscriptionCoordinator c = bucketedCoordinator();
         c.registerConsumer("consumer-1", 1L, mock(TransportCnx.class)).get();
         Map<ConsumerSession, ConsumerAssignment> result =
                 c.registerConsumer("consumer-2", 2L, mock(TransportCnx.class)).get();
 
-        int owners = 0;
+        assertEquals(result.size(), 2);
+        List<HashRange> owned = new ArrayList<>();
         for (ConsumerAssignment assignment : result.values()) {
-            for (ConsumerAssignment.AssignedSegment seg : assignment.assignedSegments()) {
-                assertEquals(seg.segmentId(), 0);
-                assertTrue(seg.bucketRanges().isEmpty());
-                owners++;
+            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(
+                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).
+        SubscriptionCoordinator c = bucketedCoordinator();
+        for (int i = 1; i <= 4; i++) {
+            c.registerConsumer("consumer-" + i, i, mock(TransportCnx.class)).get();
+        }
+        Map<ConsumerSession, ConsumerAssignment> result =
+                c.registerConsumer("consumer-5", 5L, mock(TransportCnx.class)).get();
+
+        assertEquals(result.size(), 5);
+        List<HashRange> owned = new ArrayList<>();
+        int idle = 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());
         }
-        assertEquals(owners, 1);
+        assertEquals(idle, 1);
+        assertEquals(owned.size(), 4);
+    }
+
+    private SubscriptionCoordinator bucketedCoordinator() {
+        // One segment carrying the whole default budget: N = 4 entry-buckets.
+        return new SubscriptionCoordinator("test-sub", topicName,
+                SegmentLayout.fromMetadata(ScalableTopicController.createInitialMetadata(1, 4, Map.of())),
+                resources, scheduler, Duration.ofMillis(200));
     }
 
     // --- Helpers ---