[improve][broker] Add limits for Key_Shared Subscription look ahead in dispatching

conf/broker.conf

@@ -355,6 +355,32 @@ maxUnackedMessagesPerBroker=0
 # limit/2 messages
 maxUnackedMessagesPerSubscriptionOnBrokerBlocked=0.16
 
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *
+# connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription).
+# Setting this value to 0 will disable the limit calculated per consumer.
+keySharedLookAheadMsgInReplayThresholdPerConsumer=1000
+
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *
+# connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription).
+# This value should be set to a value less than 2 * managedLedgerMaxUnackedRangesToPersist.
+# Setting this value to 0 will disable the limit calculated per subscription.
+keySharedLookAheadMsgInReplayThresholdPerSubscription=10000
+
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# This setting controls whether look ahead is enabled when recently joined consumers are present.
+keySharedLookAheadEnabledWhenRecentlyJoinedConsumersExist=false
+
 # Broker periodically checks if subscription is stuck and unblock if flag is enabled. (Default is disabled)
 unblockStuckSubscriptionEnabled=false
 

conf/standalone.conf

@@ -232,6 +232,32 @@ maxUnackedMessagesPerBroker=0
 # limit/2 messages
 maxUnackedMessagesPerSubscriptionOnBrokerBlocked=0.16
 
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *
+# connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription).
+# Setting this value to 0 will disable the limit calculated per consumer.
+keySharedLookAheadMsgInReplayThresholdPerConsumer=1000
+
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *
+# connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription).
+# This value should be set to a value less than 2 * managedLedgerMaxUnackedRangesToPersist.
+# Setting this value to 0 will disable the limit calculated per subscription.
+keySharedLookAheadMsgInReplayThresholdPerSubscription=10000
+
+# For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer
+# or a blocked key hash (because of ordering constraints), the broker will continue reading more
+# messages from the backlog and attempt to dispatch them to consumers until the number of replay
+# messages reaches the calculated threshold.
+# This setting controls whether look ahead is enabled when recently joined consumers are present.
+keySharedLookAheadEnabledWhenRecentlyJoinedConsumersExist=false
+
 # Tick time to schedule task that checks topic publish rate limiting across all topics
 # Reducing to lower value can give more accuracy while throttling publish but
 # it uses more CPU to perform frequent check. (Disable publish throttling with value 0)

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

@@ -949,6 +949,48 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
             + " back and unack count reaches to `limit/2`. Using a value of 0, is disabling unackedMessage-limit"
             + " check and broker doesn't block dispatchers")
     private int maxUnackedMessagesPerBroker = 0;
+
+    @FieldContext(
+            category = CATEGORY_POLICIES,
+            doc = "For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer"
+                    + " or a blocked key hash (because of ordering constraints), the broker will continue reading more"
+                    + " messages from the backlog and attempt to dispatch them to consumers until the number of replay"
+                    + " messages reaches the calculated threshold.\n"
+                    + "Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *"
+                    + " connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription)"
+                    + ".\n"
+                    + "Setting this value to 0 will disable the limit calculated per consumer.",
+            dynamic = true
+    )
+    private int keySharedLookAheadMsgInReplayThresholdPerConsumer = 1000;
+
+    @FieldContext(
+            category = CATEGORY_POLICIES,
+            doc = "For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer"
+                    + " or a blocked key hash (because of ordering constraints), the broker will continue reading more"
+                    + " messages from the backlog and attempt to dispatch them to consumers until the number of replay"
+                    + " messages reaches the calculated threshold.\n"
+                    + "Formula: threshold = max(keySharedLookAheadMsgInReplayThresholdPerConsumer *"
+                    + " connected consumer count, keySharedLookAheadMsgInReplayThresholdPerSubscription)"
+                    + ".\n"
+                    + "This value should be set to a value less than 2 * managedLedgerMaxUnackedRangesToPersist.\n"
+                    + "Setting this value to 0 will disable the limit calculated per subscription.\n",
+            dynamic = true
+    )
+    private int keySharedLookAheadMsgInReplayThresholdPerSubscription = 10000;
+
+
+    @FieldContext(
+            category = CATEGORY_POLICIES,
+            doc = "For Key_Shared subscriptions, if messages cannot be dispatched to consumers due to a slow consumer"
+                    + " or a blocked key hash (because of ordering constraints), the broker will continue reading more"
+                    + " messages from the backlog and attempt to dispatch them to consumers until the number of replay"
+                    + " messages reaches the calculated threshold.\n"
+                    + "This setting controls whether look ahead is enabled when recently joined consumers are present.",
+            dynamic = true
+    )
+    private boolean keySharedLookAheadEnabledWhenRecentlyJoinedConsumersExist = false;
+
     @FieldContext(
         category = CATEGORY_POLICIES,
         doc = "Once broker reaches maxUnackedMessagesPerBroker limit, it blocks subscriptions which has higher "

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

@@ -22,7 +22,10 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.NavigableSet;
+import java.util.Optional;
 import java.util.Set;
+import java.util.TreeSet;
+import java.util.function.Predicate;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.bookkeeper.mledger.Position;
 import org.apache.bookkeeper.mledger.PositionFactory;
@@ -146,8 +149,33 @@ public boolean containsStickyKeyHashes(Set<Integer> stickyKeyHashes) {
         return false;
     }
 
-    public NavigableSet<Position> getMessagesToReplayNow(int maxMessagesToRead) {
-        return messagesToRedeliver.items(maxMessagesToRead, PositionFactory::create);
+    public boolean containsStickyKeyHash(int stickyKeyHash) {
+        return !allowOutOfOrderDelivery && hashesRefCount.containsKey(stickyKeyHash);
+    }
+
+    public Optional<Position> getFirstPositionInReplay() {
+        NavigableSet<Position> items = messagesToRedeliver.items(1, PositionFactory::create);
+        return items.isEmpty() ? Optional.empty() : Optional.of(items.first());
+    }
+
+    /**
+     * Get the messages to replay now.
+     *
+     * @param maxMessagesToRead
+     *            the max messages to read
+     * @param filter
+     *            the filter to use to select the messages to replay
+     * @return the messages to replay now
+     */
+    public NavigableSet<Position> getMessagesToReplayNow(int maxMessagesToRead, Predicate<Position> filter) {
+        NavigableSet<Position> items = new TreeSet<>();
+        messagesToRedeliver.processItems(PositionFactory::create, item -> {
+            if (filter.test(item)) {
+                items.add(item);
+            }
+            return items.size() < maxMessagesToRead;
+        });
+        return items;
     }
 
     /**

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

@@ -345,26 +345,24 @@ public synchronized void readMoreEntries() {
                 return;
             }
 
-            NavigableSet<Position> messagesToReplayNow = getMessagesToReplayNow(messagesToRead);
-            NavigableSet<Position> messagesToReplayFiltered = filterOutEntriesWillBeDiscarded(messagesToReplayNow);
-            if (!messagesToReplayFiltered.isEmpty()) {
+            Set<Position> messagesToReplayNow =
+                    canReplayMessages() ? getMessagesToReplayNow(messagesToRead) : Collections.emptySet();
+            if (!messagesToReplayNow.isEmpty()) {
                 if (log.isDebugEnabled()) {
                     log.debug("[{}] Schedule replay of {} messages for {} consumers", name,
-                            messagesToReplayFiltered.size(), consumerList.size());
+                            messagesToReplayNow.size(), consumerList.size());
                 }
-
                 havePendingReplayRead = true;
-                minReplayedPosition = messagesToReplayNow.first();
+                updateMinReplayedPosition();
                 Set<? extends Position> deletedMessages = topic.isDelayedDeliveryEnabled()
-                        ? asyncReplayEntriesInOrder(messagesToReplayFiltered)
-                        : asyncReplayEntries(messagesToReplayFiltered);
+                        ? asyncReplayEntriesInOrder(messagesToReplayNow)
+                        : asyncReplayEntries(messagesToReplayNow);
                 // clear already acked positions from replay bucket
-
                 deletedMessages.forEach(position -> redeliveryMessages.remove(position.getLedgerId(),
                         position.getEntryId()));
                 // if all the entries are acked-entries and cleared up from redeliveryMessages, try to read
                 // next entries as readCompletedEntries-callback was never called
-                if ((messagesToReplayFiltered.size() - deletedMessages.size()) == 0) {
+                if ((messagesToReplayNow.size() - deletedMessages.size()) == 0) {
                     havePendingReplayRead = false;
                     readMoreEntriesAsync();
                 }
@@ -386,13 +384,7 @@ public synchronized void readMoreEntries() {
                             consumerList.size());
                 }
                 havePendingRead = true;
-                NavigableSet<Position> toReplay = getMessagesToReplayNow(1);
-                if (!toReplay.isEmpty()) {
-                    minReplayedPosition = toReplay.first();
-                    redeliveryMessages.add(minReplayedPosition.getLedgerId(), minReplayedPosition.getEntryId());
-                } else {
-                    minReplayedPosition = null;
-                }
+                updateMinReplayedPosition();
 
                 // Filter out and skip read delayed messages exist in DelayedDeliveryTracker
                 if (delayedDeliveryTracker.isPresent()) {
@@ -427,6 +419,19 @@ public synchronized void readMoreEntries() {
         }
     }
 
+    /**
+     * Controls whether replaying entries is currently enabled.
+     * Subclasses can override this method to temporarily disable replaying entries.
+     * @return true if replaying entries is currently enabled
+     */
+    protected boolean canReplayMessages() {
+        return true;
+    }
+
+    private void updateMinReplayedPosition() {
+        minReplayedPosition = getFirstPositionInReplay().orElse(null);
+    }
+
     private boolean shouldPauseOnAckStatePersist(ReadType readType) {
         // Allows new consumers to consume redelivered messages caused by the just-closed consumer.
         if (readType != ReadType.Normal) {
@@ -751,8 +756,9 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis
         int remainingMessages = 0;
         boolean hasChunk = false;
         for (int i = 0; i < metadataArray.length; i++) {
-            final MessageMetadata metadata = Commands.peekAndCopyMessageMetadata(
-                    entries.get(i).getDataBuffer(), subscription.toString(), -1);
+            Entry entry = entries.get(i);
+            MessageMetadata metadata = entry instanceof EntryAndMetadata ? ((EntryAndMetadata) entry).getMetadata()
+                    : Commands.peekAndCopyMessageMetadata(entry.getDataBuffer(), subscription.toString(), -1);
             if (metadata != null) {
                 remainingMessages += metadata.getNumMessagesInBatch();
                 if (!hasChunk && metadata.hasUuid()) {
@@ -1232,25 +1238,30 @@ protected synchronized NavigableSet<Position> getMessagesToReplayNow(int maxMess
                     delayedDeliveryTracker.get().getScheduledMessages(maxMessagesToRead);
             messagesAvailableNow.forEach(p -> redeliveryMessages.add(p.getLedgerId(), p.getEntryId()));
         }
-
         if (!redeliveryMessages.isEmpty()) {
-            return redeliveryMessages.getMessagesToReplayNow(maxMessagesToRead);
+            return redeliveryMessages.getMessagesToReplayNow(maxMessagesToRead, createFilterForReplay());
         } else {
             return Collections.emptyNavigableSet();
         }
     }
 
+    protected Optional<Position> getFirstPositionInReplay() {
+        return redeliveryMessages.getFirstPositionInReplay();
+    }
+
     /**
-     * This is a mode method designed for Key_Shared mode.
+     * Creates a stateful filter for filtering replay positions.
+     * This is only used for Key_Shared mode to skip replaying certain entries.
      * Filter out the entries that will be discarded due to the order guarantee mechanism of Key_Shared mode.
      * This method is in order to avoid the scenario below:
      * - Get positions from the Replay queue.
      * - Read entries from BK.
      * - The order guarantee mechanism of Key_Shared mode filtered out all the entries.
      * - Delivery non entry to the client, but we did a BK read.
      */
-    protected NavigableSet<Position> filterOutEntriesWillBeDiscarded(NavigableSet<Position> src) {
-        return src;
+    protected Predicate<Position> createFilterForReplay() {
+        // pick all positions from the replay
+        return position -> true;
     }
 
     /**