[fix][broker] PIP-379 Key_Shared implementation racecondition causing out-of-order message delivery

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

@@ -340,7 +340,7 @@ public void reduceRefCount(Consumer consumer, int stickyHash, boolean closing) {
             lock.writeLock().lock();
             try {
                 removed = drainingHashes.remove(stickyHash);
-                if (!closing && removed.isBlocking()) {
+                if (removed.isBlocking()) {
                     if (batchLevel > 0) {
                         unblockedWhileBatching = true;
                     } else {

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

@@ -335,6 +335,13 @@ public synchronized void readMoreEntries() {
             }
             return;
         }
+        if (havePendingReplayRead) {
+            if (log.isDebugEnabled()) {
+                log.debug("[{}] [{}] Skipping read for the topic, Due to replay in-progress.", topic.getName(),
+                        getSubscriptionName());
+            }
+            return;
+        }
         if (isSendInProgress()) {
             // we cannot read more entries while sending the previous batch
             // otherwise we could re-read the same entries and send duplicates
@@ -379,13 +386,23 @@ public synchronized void readMoreEntries() {
             long bytesToRead = calculateResult.getRight();
 
             if (messagesToRead == -1 || bytesToRead == -1) {
-                // Skip read as topic/dispatcher has exceed the dispatch rate or previous pending read hasn't complete.
+                // Skip read as topic/dispatcher has exceed the dispatch rate
                 return;
             }
 
             Set<Position> messagesToReplayNow =
                     canReplayMessages() ? getMessagesToReplayNow(messagesToRead) : Collections.emptySet();
             if (!messagesToReplayNow.isEmpty()) {
+                // before replaying, cancel possible pending read that is waiting for more entries
+                cancelPendingRead();
+                if (havePendingRead) {
+                    // skip read since a pending read is already in progress which cannot be cancelled
+                    if (log.isDebugEnabled()) {
+                        log.debug("[{}] [{}] Skipping replay read for the topic, Due to pending read in-progress.",
+                                topic.getName(), getSubscriptionName());
+                    }
+                    return;
+                }
                 if (log.isDebugEnabled()) {
                     log.debug("[{}] Schedule replay of {} messages for {} consumers", name,
                             messagesToReplayNow.size(), consumerList.size());
@@ -615,13 +632,6 @@ protected Pair<Integer, Long> calculateToRead(int currentTotalAvailablePermits)
             }
         }
 
-        if (havePendingReplayRead) {
-            if (log.isDebugEnabled()) {
-                log.debug("[{}] Skipping replay while awaiting previous read to complete", name);
-            }
-            return Pair.of(-1, -1L);
-        }
-
         // If messagesToRead is 0 or less, correct it to 1 to prevent IllegalArgumentException
         messagesToRead = Math.max(messagesToRead, 1);
         bytesToRead = Math.max(bytesToRead, 1);
@@ -717,6 +727,12 @@ public SubType getType() {
     public final synchronized void readEntriesComplete(List<Entry> entries, Object ctx) {
         ReadType readType = (ReadType) ctx;
         if (readType == ReadType.Normal) {
+            if (!havePendingRead) {
+                log.debug("Discarding read entries as there is no pending read");
+                entries.forEach(Entry::release);
+                readMoreEntriesAsync();
+                return;
+            }
             havePendingRead = false;
         } else {
             havePendingReplayRead = false;
@@ -1422,6 +1438,9 @@ public void cursorIsReset() {
 
     protected boolean addMessageToReplay(long ledgerId, long entryId, long stickyKeyHash) {
         if (checkIfMessageIsUnacked(ledgerId, entryId)) {
+            if (log.isDebugEnabled()) {
+                log.debug("[{}] Adding message to replay for {}:{} hash: {}", name, ledgerId, entryId, stickyKeyHash);
+            }
             redeliveryMessages.add(ledgerId, entryId, stickyKeyHash);
             return true;
         } else {

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

@@ -20,6 +20,8 @@
 
 import static org.apache.pulsar.broker.service.StickyKeyConsumerSelector.STICKY_KEY_HASH_NOT_SET;
 import com.google.common.annotations.VisibleForTesting;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSet;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -407,6 +409,8 @@ private Map<Consumer, List<Entry>> filterAndGroupEntriesForDispatching(List<Entr
         Set<Consumer> blockedByHashConsumers = lookAheadAllowed && readType == ReadType.Normal ? new HashSet<>() : null;
         // in replay read mode, keep track of consumers for entries, used for look-ahead check
         Set<Consumer> consumersForEntriesForLookaheadCheck = lookAheadAllowed ? new HashSet<>() : null;
+        // track already blocked hashes to block any further messages with the same hash
+        IntSet alreadyBlockedHashes = new IntOpenHashSet();
 
         for (Entry inputEntry : entries) {
             EntryAndMetadata entry;
@@ -419,24 +423,29 @@ private Map<Consumer, List<Entry>> filterAndGroupEntriesForDispatching(List<Entr
                         Commands.peekAndCopyMessageMetadata(inputEntry.getDataBuffer(), getSubscriptionName(), -1));
             }
             int stickyKeyHash = getStickyKeyHash(entry);
-            Consumer consumer = selector.select(stickyKeyHash);
+            Consumer consumer = null;
             MutableBoolean blockedByHash = null;
             boolean dispatchEntry = false;
-            if (consumer != null) {
-                if (lookAheadAllowed) {
-                    consumersForEntriesForLookaheadCheck.add(consumer);
-                }
-                blockedByHash = lookAheadAllowed && readType == ReadType.Normal ? new MutableBoolean(false) : null;
-                MutableInt permits =
-                        permitsForConsumer.computeIfAbsent(consumer,
-                                k -> new MutableInt(getAvailablePermits(consumer)));
-                // a consumer was found for the sticky key hash and the entry can be dispatched
-                if (permits.intValue() > 0
-                        && canDispatchEntry(consumer, entry, readType, stickyKeyHash, blockedByHash)) {
-                    // decrement the permits for the consumer
-                    permits.decrement();
-                    // allow the entry to be dispatched
-                    dispatchEntry = true;
+            // check if the hash is already blocked
+            boolean hashIsAlreadyBlocked = alreadyBlockedHashes.contains(stickyKeyHash);
+            if (!hashIsAlreadyBlocked) {
+                consumer = selector.select(stickyKeyHash);
+                if (consumer != null) {
+                    if (lookAheadAllowed) {
+                        consumersForEntriesForLookaheadCheck.add(consumer);
+                    }
+                    blockedByHash = lookAheadAllowed && readType == ReadType.Normal ? new MutableBoolean(false) : null;
+                    MutableInt permits =
+                            permitsForConsumer.computeIfAbsent(consumer,
+                                    k -> new MutableInt(getAvailablePermits(k)));
+                    // a consumer was found for the sticky key hash and the entry can be dispatched
+                    if (permits.intValue() > 0
+                            && canDispatchEntry(consumer, entry, readType, stickyKeyHash, blockedByHash)) {
+                        // decrement the permits for the consumer
+                        permits.decrement();
+                        // allow the entry to be dispatched
+                        dispatchEntry = true;
+                    }
                 }
             }
             if (dispatchEntry) {
@@ -445,6 +454,10 @@ && canDispatchEntry(consumer, entry, readType, stickyKeyHash, blockedByHash)) {
                         entriesGroupedByConsumer.computeIfAbsent(consumer, k -> new ArrayList<>());
                 consumerEntries.add(entry);
             } else {
+                if (!hashIsAlreadyBlocked) {
+                    // the hash is blocked, add it to the set of blocked hashes
+                    alreadyBlockedHashes.add(stickyKeyHash);
+                }
                 if (blockedByHash != null && blockedByHash.isTrue()) {
                     // the entry is blocked by hash, add the consumer to the blocked set
                     blockedByHashConsumers.add(consumer);
@@ -536,6 +549,9 @@ private class ReplayPositionFilter implements Predicate<Position> {
         // tracks the available permits for each consumer for the duration of the filter usage
         // the filter is stateful and shouldn't be shared or reused later
         private final Map<Consumer, MutableInt> availablePermitsMap = new HashMap<>();
+        // tracks the hashes that have been blocked during the filtering
+        // it is necessary to block all later messages after a hash gets blocked so that ordering is preserved
+        private final Set<Long> alreadyBlockedHashes = new HashSet<>();
 
         @Override
         public boolean test(Position position) {
@@ -553,25 +569,34 @@ public boolean test(Position position) {
                 }
                 return true;
             }
+            // check if the hash is already blocked, if so, then replaying of the position should be skipped
+            // to preserve ordering
+            if (alreadyBlockedHashes.contains(stickyKeyHash)) {
+                return false;
+            }
 
             // find the consumer for the sticky key hash
             Consumer consumer = selector.select(stickyKeyHash.intValue());
             // skip replaying the message position if there's no assigned consumer
             if (consumer == null) {
+                alreadyBlockedHashes.add(stickyKeyHash);
                 return false;
             }
+
             // lookup the available permits for the consumer
             MutableInt availablePermits =
                     availablePermitsMap.computeIfAbsent(consumer,
                             k -> new MutableInt(getAvailablePermits(consumer)));
             // skip replaying the message position if the consumer has no available permits
             if (availablePermits.intValue() <= 0) {
+                alreadyBlockedHashes.add(stickyKeyHash);
                 return false;
             }
 
             if (drainingHashesRequired
                     && drainingHashesTracker.shouldBlockStickyKeyHash(consumer, stickyKeyHash.intValue())) {
                 // the hash is draining and the consumer is not the draining consumer
+                alreadyBlockedHashes.add(stickyKeyHash);
                 return false;
             }
 

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerTest.java

@@ -191,23 +191,4 @@ public void unblockingHandler_InvokesStickyKeyHashUnblocked() {
         // then unblocking call should be done
         verify(unblockingHandler).stickyKeyHashUnblocked(1);
     }
-
-    @Test
-    public void unblockingHandler_DoesNotInvokeStickyKeyHashUnblockedWhenClosing() {
-        // given a tracker with unblocking handler
-        UnblockingHandler unblockingHandler = mock(UnblockingHandler.class);
-        DrainingHashesTracker tracker = new DrainingHashesTracker("dispatcher1", unblockingHandler);
-
-        // when a hash is draining
-        Consumer consumer = createMockConsumer("consumer1");
-        tracker.addEntry(consumer, 1);
-        // aand hash gets blocked
-        Consumer consumer2 = createMockConsumer("consumer2");
-        tracker.shouldBlockStickyKeyHash(consumer2, 1);
-        // and hash gets unblocked
-        tracker.reduceRefCount(consumer, 1, true);
-
-        // then unblocking call should be done
-        verify(unblockingHandler, never()).stickyKeyHashUnblocked(anyInt());
-    }
 }