[improve][broker] PIP-379: Key_Shared Draining Hashes for Improved Message Ordering

distribution/server/src/assemble/LICENSE.bin.txt

@@ -262,6 +262,7 @@ The Apache Software License, Version 2.0
      - com.fasterxml.jackson.module-jackson-module-parameter-names-2.17.2.jar
  * Caffeine -- com.github.ben-manes.caffeine-caffeine-2.9.1.jar
  * Conscrypt -- org.conscrypt-conscrypt-openjdk-uber-2.5.2.jar
+ * Fastutil -- it.unimi.dsi-fastutil-8.5.14.jar
  * Proto Google Common Protos -- com.google.api.grpc-proto-google-common-protos-2.17.0.jar
  * Bitbucket -- org.bitbucket.b_c-jose4j-0.9.4.jar
  * Gson

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/ManagedCursor.java

@@ -330,6 +330,12 @@ default void asyncReadEntriesWithSkipOrWait(int maxEntries, long maxSizeBytes, R
         asyncReadEntriesOrWait(maxEntries, maxSizeBytes, callback, ctx, maxPosition);
     }
 
+    /**
+     * Returns whether there is a pending read operation in progress.
+     * @return true if there is a pending read operation in progress
+     */
+    boolean hasPendingReadRequest();
+
     /**
      * Cancel a previously scheduled asyncReadEntriesOrWait operation.
      *
@@ -900,4 +906,14 @@ default ManagedCursorAttributes getManagedCursorAttributes() {
     int applyMaxSizeCap(int maxEntries, long maxSizeBytes);
 
     void updateReadStats(int readEntriesCount, long readEntriesSize);
+
+    /**
+     * Get the number of read operations performed on this cursor.
+     * This is used in Pulsar's Key_Shared dispatcher to determine whether a pending read needs to be cancelled
+     * in order to process entries in the replay queue when a hash is unblocked.
+     * @return the number of read operations performed on this cursor
+     */
+    default long getReadOpCount() {
+        return 0;
+    }
 }

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/ManagedCursorImpl.java

@@ -216,6 +216,9 @@ public class ManagedCursorImpl implements ManagedCursor {
     private int individualDeletedMessagesSerializedSize;
     private static final String COMPACTION_CURSOR_NAME = "__compaction";
     private volatile boolean cacheReadEntry = false;
+    private static final AtomicLongFieldUpdater<ManagedCursorImpl> READ_OP_COUNT_UPDATER =
+            AtomicLongFieldUpdater.newUpdater(ManagedCursorImpl.class, "readOpCount");
+    private volatile long readOpCount;
 
     // active state cache in ManagedCursor. It should be in sync with the state in activeCursors in ManagedLedger.
     private volatile boolean isActive = false;
@@ -571,6 +574,7 @@ protected void recoverFromLedger(final ManagedCursorInfo info, final VoidCallbac
                 return;
             }
 
+            READ_OP_COUNT_UPDATER.incrementAndGet(ManagedCursorImpl.this);
             lh.asyncReadEntries(lastEntryInLedger, lastEntryInLedger, (rc1, lh1, seq, ctx1) -> {
                 if (log.isDebugEnabled()) {
                     log.debug("[{}} readComplete rc={} entryId={}", ledger.getName(), rc1, lh1.getLastAddConfirmed());
@@ -1102,6 +1106,7 @@ public boolean cancelPendingReadRequest() {
         return op != null;
     }
 
+    @Override
     public boolean hasPendingReadRequest() {
         return WAITING_READ_OP_UPDATER.get(this) != null;
     }
@@ -3804,13 +3809,19 @@ public ManagedCursorAttributes getManagedCursorAttributes() {
         return ATTRIBUTES_UPDATER.updateAndGet(this, old -> old != null ? old : new ManagedCursorAttributes(this));
     }
 
+    @Override
+    public long getReadOpCount() {
+        return readOpCount;
+    }
+
     @Override
     public ManagedLedgerInternalStats.CursorStats getCursorStats() {
         ManagedLedgerInternalStats.CursorStats cs = new ManagedLedgerInternalStats.CursorStats();
         cs.markDeletePosition = getMarkDeletedPosition().toString();
         cs.readPosition = getReadPosition().toString();
         cs.waitingReadOp = hasPendingReadRequest();
         cs.pendingReadOps = getPendingReadOpsCount();
+        cs.readOpCount = getReadOpCount();
         cs.messagesConsumedCounter = getMessagesConsumedCounter();
         cs.cursorLedger = getCursorLedger();
         cs.cursorLedgerLastEntry = getCursorLedgerLastEntry();

managed-ledger/src/test/java/org/apache/bookkeeper/mledger/impl/ManagedCursorContainerTest.java

@@ -315,6 +315,11 @@ public void asyncReadEntriesOrWait(int maxEntries, long maxSizeBytes, ReadEntrie
 
         }
 
+        @Override
+        public boolean hasPendingReadRequest() {
+            return false;
+        }
+
         @Override
         public boolean cancelPendingReadRequest() {
             return true;

pom.xml

@@ -162,6 +162,7 @@ flexible messaging model and an intuitive client API.</description>
     <bouncycastle.bcpkix-fips.version>1.0.7</bouncycastle.bcpkix-fips.version>
     <bouncycastle.bc-fips.version>1.0.2.5</bouncycastle.bc-fips.version>
     <jackson.version>2.17.2</jackson.version>
+    <fastutil.version>8.5.14</fastutil.version>
     <reflections.version>0.10.2</reflections.version>
     <swagger.version>1.6.2</swagger.version>
     <puppycrawl.checkstyle.version>10.14.2</puppycrawl.checkstyle.version>
@@ -911,6 +912,12 @@ flexible messaging model and an intuitive client API.</description>
         <scope>import</scope>
       </dependency>
 
+      <dependency>
+        <groupId>it.unimi.dsi</groupId>
+        <artifactId>fastutil</artifactId>
+        <version>${fastutil.version}</version>
+      </dependency>
+
       <dependency>
         <groupId>org.codehaus.jettison</groupId>
         <artifactId>jettison</artifactId>

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java

@@ -979,6 +979,16 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
     )
     private int keySharedLookAheadMsgInReplayThresholdPerSubscription = 20000;
 
+    @FieldContext(
+            category = CATEGORY_POLICIES,
+            doc = "For Key_Shared subscriptions, when a blocked key hash gets unblocked,"
+                    + " a redelivery will be attempted after a delay. This setting controls the delay."
+                    + " The reason to have the delay is to batch multiple unblocking events instead of triggering"
+                    + " redelivery for each unblocking event.",
+            dynamic = true
+    )
+    private long keySharedUnblockingIntervalMs = 10L;
+
     @FieldContext(
         category = CATEGORY_POLICIES,
         doc = "Once broker reaches maxUnackedMessagesPerBroker limit, it blocks subscriptions which has higher "

pulsar-broker/pom.xml

@@ -63,6 +63,11 @@
       <artifactId>protobuf-java</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>it.unimi.dsi</groupId>
+      <artifactId>fastutil</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>${project.groupId}</groupId>
       <artifactId>pulsar-client-original</artifactId>

pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/AbstractDelayedDeliveryTracker.java

@@ -146,7 +146,7 @@ public void run(Timeout timeout) throws Exception {
             lastTickRun = clock.millis();
             currentTimeoutTarget = -1;
             this.timeout = null;
-            dispatcher.readMoreEntries();
+            dispatcher.readMoreEntriesAsync();
         }
     }
 

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerService.java

@@ -3351,7 +3351,7 @@ public void unblockDispatchersOnUnAckMessages(List<PersistentDispatcherMultipleC
         try {
             dispatcherList.forEach(dispatcher -> {
                 dispatcher.unBlockDispatcherOnUnackedMsgs();
-                executor().execute(() -> dispatcher.readMoreEntries());
+                dispatcher.readMoreEntriesAsync();
                 log.info("[{}] Dispatcher is unblocked", dispatcher.getName());
                 blockedDispatchers.remove(dispatcher);
             });

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java

@@ -18,23 +18,18 @@
  */
 package org.apache.pulsar.broker.service;
 
-import java.util.ArrayList;
-import java.util.IdentityHashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.NavigableMap;
+import java.util.SortedMap;
 import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import org.apache.pulsar.client.api.Range;
-import org.apache.pulsar.common.util.Murmur3_32Hash;
 
 /**
- * This is a consumer selector based fixed hash range.
- *
- * The implementation uses consistent hashing to evenly split, the
- * number of keys assigned to each consumer.
+ * This is a consumer selector using consistent hashing to evenly split
+ * the number of keys assigned to each consumer.
  */
 public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyConsumerSelector {
     // use NUL character as field separator for hash key calculation
@@ -47,14 +42,22 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons
     private final ConsumerNameIndexTracker consumerNameIndexTracker = new ConsumerNameIndexTracker();
 
     private final int numberOfPoints;
+    private final Range keyHashRange;
+    private ConsumerHashAssignmentsSnapshot consumerHashAssignmentsSnapshot;
 
     public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) {
+        this(numberOfPoints, DEFAULT_RANGE_SIZE - 1);
+    }
+
+    public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints, int rangeMaxValue) {
         this.hashRing = new TreeMap<>();
         this.numberOfPoints = numberOfPoints;
+        this.keyHashRange = Range.of(STICKY_KEY_HASH_NOT_SET + 1, rangeMaxValue);
+        this.consumerHashAssignmentsSnapshot = ConsumerHashAssignmentsSnapshot.empty();
     }
 
     @Override
-    public CompletableFuture<Void> addConsumer(Consumer consumer) {
+    public CompletableFuture<ImpactedConsumersResult> addConsumer(Consumer consumer) {
         rwLock.writeLock().lock();
         try {
             ConsumerIdentityWrapper consumerIdentityWrapper = new ConsumerIdentityWrapper(consumer);
@@ -72,7 +75,11 @@ public CompletableFuture<Void> addConsumer(Consumer consumer) {
                     consumerNameIndexTracker.decreaseConsumerRefCount(removed);
                 }
             }
-            return CompletableFuture.completedFuture(null);
+            ConsumerHashAssignmentsSnapshot assignmentsAfter = internalGetConsumerHashAssignmentsSnapshot();
+            ImpactedConsumersResult impactedConsumers =
+                    consumerHashAssignmentsSnapshot.resolveImpactedConsumers(assignmentsAfter);
+            consumerHashAssignmentsSnapshot = assignmentsAfter;
+            return CompletableFuture.completedFuture(impactedConsumers);
         } finally {
             rwLock.writeLock().unlock();
         }
@@ -88,14 +95,14 @@ public CompletableFuture<Void> addConsumer(Consumer consumer) {
      * @param hashRingPointIndex the index of the hash ring point
      * @return the hash value
      */
-    private static int calculateHashForConsumerAndIndex(Consumer consumer, int consumerNameIndex,
+    private int calculateHashForConsumerAndIndex(Consumer consumer, int consumerNameIndex,
                                                         int hashRingPointIndex) {
         String key = consumer.consumerName() + KEY_SEPARATOR + consumerNameIndex + KEY_SEPARATOR + hashRingPointIndex;
-        return Murmur3_32Hash.getInstance().makeHash(key.getBytes());
+        return makeStickyKeyHash(key.getBytes());
     }
 
     @Override
-    public void removeConsumer(Consumer consumer) {
+    public ImpactedConsumersResult removeConsumer(Consumer consumer) {
         rwLock.writeLock().lock();
         try {
             ConsumerIdentityWrapper consumerIdentityWrapper = new ConsumerIdentityWrapper(consumer);
@@ -109,6 +116,11 @@ public void removeConsumer(Consumer consumer) {
                     }
                 }
             }
+            ConsumerHashAssignmentsSnapshot assignmentsAfter = internalGetConsumerHashAssignmentsSnapshot();
+            ImpactedConsumersResult impactedConsumers =
+                    consumerHashAssignmentsSnapshot.resolveImpactedConsumers(assignmentsAfter);
+            consumerHashAssignmentsSnapshot = assignmentsAfter;
+            return impactedConsumers;
         } finally {
             rwLock.writeLock().unlock();
         }
@@ -134,32 +146,58 @@ public Consumer select(int hash) {
     }
 
     @Override
-    public Map<Consumer, List<Range>> getConsumerKeyHashRanges() {
-        Map<Consumer, List<Range>> result = new IdentityHashMap<>();
+    public Range getKeyHashRange() {
+        return keyHashRange;
+    }
+
+    @Override
+    public ConsumerHashAssignmentsSnapshot getConsumerHashAssignmentsSnapshot() {
         rwLock.readLock().lock();
         try {
-            if (hashRing.isEmpty()) {
-                return result;
-            }
-            int start = 0;
-            int lastKey = 0;
-            for (Map.Entry<Integer, ConsumerIdentityWrapper> entry: hashRing.entrySet()) {
-                Consumer consumer = entry.getValue().consumer;
-                result.computeIfAbsent(consumer, key -> new ArrayList<>())
-                        .add(Range.of(start, entry.getKey()));
-                lastKey = entry.getKey();
-                start = lastKey + 1;
-            }
-            // Handle wrap-around in the hash ring, the first consumer will also contain the range from the last key
-            // to the maximum value of the hash range
-            Consumer firstConsumer = hashRing.firstEntry().getValue().consumer;
-            List<Range> ranges = result.get(firstConsumer);
-            if (lastKey != Integer.MAX_VALUE - 1) {
-                ranges.add(Range.of(lastKey + 1, Integer.MAX_VALUE - 1));
-            }
+            return consumerHashAssignmentsSnapshot;
         } finally {
             rwLock.readLock().unlock();
         }
-        return result;
+    }
+
+    private ConsumerHashAssignmentsSnapshot internalGetConsumerHashAssignmentsSnapshot() {
+        if (hashRing.isEmpty()) {
+            return ConsumerHashAssignmentsSnapshot.empty();
+        }
+        SortedMap<Range, Consumer> result = new TreeMap<>();
+        int start = getKeyHashRange().getStart();
+        int lastKey = -1;
+        Consumer previousConsumer = null;
+        Range previousRange = null;
+        for (Map.Entry<Integer, ConsumerIdentityWrapper> entry: hashRing.entrySet()) {
+            Consumer consumer = entry.getValue().consumer;
+            Range range;
+            if (consumer == previousConsumer) {
+                // join ranges
+                result.remove(previousRange);
+                range = Range.of(previousRange.getStart(), entry.getKey());
+            } else {
+                range = Range.of(start, entry.getKey());
+            }
+            result.put(range, consumer);
+            lastKey = entry.getKey();
+            start = lastKey + 1;
+            previousConsumer = consumer;
+            previousRange = range;
+        }
+        // Handle wrap-around
+        Consumer firstConsumer = hashRing.firstEntry().getValue().consumer;
+        if (lastKey != getKeyHashRange().getEnd()) {
+            Range range;
+            if (firstConsumer == previousConsumer && previousRange.getEnd() == lastKey) {
+                // join ranges
+                result.remove(previousRange);
+                range = Range.of(previousRange.getStart(), getKeyHashRange().getEnd());
+            } else {
+                range = Range.of(lastKey + 1, getKeyHashRange().getEnd());
+            }
+            result.put(range, firstConsumer);
+        }
+        return ConsumerHashAssignmentsSnapshot.of(result);
     }
 }