[fix][broker] Fix persistent throughput degradation caused by permit loss during frequent reconnects on Shared subscriptions

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

@@ -32,7 +32,8 @@
 public abstract class AbstractDispatcherMultipleConsumers extends AbstractBaseDispatcher {
 
     protected final CopyOnWriteArrayList<Consumer> consumerList = new CopyOnWriteArrayList<>();
-    protected final ObjectSet<Consumer> consumerSet = new ObjectHashSet<>();
+    private final ObjectHashSet<Consumer> consumerSetImpl = new ObjectHashSet<>();
+    protected final ObjectSet<Consumer> consumerSet = consumerSetImpl;
     protected volatile int currentConsumerRoundRobinIndex = 0;
 
     protected static final int FALSE = 0;
@@ -58,6 +59,18 @@ public synchronized boolean canUnsubscribe(Consumer consumer) {
         return consumerList.size() == 1 && consumerSet.contains(consumer);
     }
 
+    /**
+     * Checks whether the exact Consumer instance is still connected.
+     *
+     * <p>This differs from {@link ObjectSet#contains(Object)}, which uses {@link Consumer#equals(Object)} and can
+     * match a replacement Consumer that reuses the same protocol identity.
+     * The caller must hold the dispatcher monitor while checking membership and acting on the result.
+     */
+    protected final boolean containsConsumerInstance(Consumer consumer) {
+        int index = consumerSetImpl.indexOf(consumer);
+        return consumerSetImpl.indexExists(index) && consumerSetImpl.indexGet(index) == consumer;
+    }
+
     public boolean isClosed() {
         return isClosed == TRUE;
     }

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

@@ -117,6 +117,18 @@ public class Consumer {
     private static final AtomicIntegerFieldUpdater<Consumer> MESSAGE_PERMITS_UPDATER =
             AtomicIntegerFieldUpdater.newUpdater(Consumer.class, "messagePermits");
     private volatile int messagePermits = 0;
+    /**
+     * Guards the Flow-side compound update of {@link #messagePermits} and
+     * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before the corresponding
+     * task runs on the dispatcher's message thread. Consumer removal can happen while that task is queued or waiting
+     * for the dispatcher monitor, so both values must be observed consistently when calculating how many permits are
+     * already included in the dispatcher total.
+     *
+     * <p>The dispatcher callback is invoked only after this lock is released. This avoids holding the lock while
+     * calling into the subscription and preserves the lock order used by dispatcher flow processing and removal.
+     */
+    private final Object flowPermitAccountingLock = new Object();
+    private int pendingDispatcherFlowPermits = 0;
     // It starts keep tracking of messagePermits once consumer gets blocked, as consumer needs two separate counts:
     // messagePermits (1) before and (2) after being blocked: to dispatch only blockedPermit number of messages at the
     // time of redelivery
@@ -916,7 +928,7 @@ public void flowPermits(int additionalNumberOfMessages) {
         }
         int oldPermits;
         if (!blockedConsumerOnUnackedMsgs) {
-            oldPermits = MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfMessages);
+            oldPermits = addPermitsPendingDispatcherUpdate(additionalNumberOfMessages);
             log.debug()
                     .attr("additionalNumberOfMessages", additionalNumberOfMessages)
                     .log("Added message permits before updating dispatcher");
@@ -943,7 +955,7 @@ public void flowPermits(int additionalNumberOfMessages) {
     void flowConsumerBlockedPermits(Consumer consumer) {
         int additionalNumberOfPermits = PERMITS_RECEIVED_WHILE_CONSUMER_BLOCKED_UPDATER.getAndSet(consumer, 0);
         // add newly flow permits to actual consumer.messagePermits
-        MESSAGE_PERMITS_UPDATER.getAndAdd(consumer, additionalNumberOfPermits);
+        consumer.addPermitsPendingDispatcherUpdate(additionalNumberOfPermits);
         log.debug()
                 .attr("additionalNumberOfPermits", additionalNumberOfPermits)
                 .log("Added blocked permits");
@@ -955,6 +967,57 @@ public int getAvailablePermits() {
         return MESSAGE_PERMITS_UPDATER.get(this);
     }
 
+    /**
+     * Adds permits after a Flow command is accepted and immediately before notifying the dispatcher. The pending
+     * count covers the interval until the dispatcher's asynchronous Flow task starts processing the same permits.
+     */
+    private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) {
+        if (!shouldTrackPendingDispatcherFlowPermits()) {
+            return MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfPermits);
+        }
+        synchronized (flowPermitAccountingLock) {
+            pendingDispatcherFlowPermits += additionalNumberOfPermits;
+            return MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfPermits);
+        }
+    }
+
+    private boolean shouldTrackPendingDispatcherFlowPermits() {
+        return isPersistentTopic && Subscription.isIndividualAckMode(subType);
+    }
+
+    /**
+     * Called at the start of the dispatcher's asynchronous Flow task, before checking whether this consumer is still
+     * connected. At this point the Flow update is no longer pending: the dispatcher will either add the permits to
+     * its total or ignore them because the consumer has already been removed.
+     */
+    public void completePendingDispatcherFlow(int additionalNumberOfPermits) {
+        synchronized (flowPermitAccountingLock) {
+            // Preserve the accounting delta across signed int wrap, matching the other permit counters.
+            pendingDispatcherFlowPermits -= additionalNumberOfPermits;
+        }
+    }
+
+    /**
+     * Called while the dispatcher removes this consumer. Permits belonging to Flow tasks that have not started yet
+     * are excluded because those permits have not been added to the dispatcher total and must not be subtracted from
+     * it during removal.
+     *
+     * <p>This accounting is enabled for persistent Shared and Key_Shared dispatchers. It requires every dispatcher
+     * Flow task to call {@link #completePendingDispatcherFlow(int)} before applying or ignoring it. For these
+     * dispatchers, when observed under the dispatcher monitor, the total available permits equal the sum of this
+     * balance over all connected consumers.
+     *
+     * <p>The returned balance can be negative. A pending Flow makes permits visible on the consumer before its
+     * asynchronous dispatcher update runs, so the dispatcher can consume those permits while they are still counted
+     * as pending. Subtracting the negative balance during removal is required to restore the dispatcher total; callers
+     * must not clamp it to zero.
+     */
+    public int getAvailablePermitsForDispatcherRemoval() {
+        synchronized (flowPermitAccountingLock) {
+            return MESSAGE_PERMITS_UPDATER.get(this) - pendingDispatcherFlowPermits;
+        }
+    }
+
     /**
      * return 0 if there is no entry dispatched yet.
      */

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

@@ -36,6 +36,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -258,9 +259,12 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE
                         notifyAddedToReplay.setTrue();
                     }
                 });
-                totalAvailablePermits -= consumer.getAvailablePermits();
+                // Restore the invariant that the dispatcher total equals the sum of the removal balances of the
+                // remaining consumers. Exclude Flow permits that have not updated the dispatcher total yet.
+                int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval();
+                totalAvailablePermits -= availablePermits;
                 log.debug()
-                        .attr("diffAvailablePermits", consumer.getAvailablePermits())
+                        .attr("availablePermits", availablePermits)
                         .attr("totalAvailablePermits", totalAvailablePermits)
                         .log("Decreased totalAvailablePermits");
                 if (notifyAddedToReplay.booleanValue()) {
@@ -300,27 +304,45 @@ protected synchronized void clearComponentsAfterRemovedAllConsumers() {
 
     @Override
     public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) {
-        topic.getBrokerService().executor().execute(() -> {
-            internalConsumerFlow(consumer, additionalNumberOfMessages);
-        });
+        Runnable flowTask = () -> internalConsumerFlow(consumer, additionalNumberOfMessages);
+        try {
+            dispatchMessagesThread.execute(flowTask);
+        } catch (RejectedExecutionException e) {
+            // Leave the permits pending so removal excludes this unapplied Flow. Never wait for the dispatcher
+            // monitor on the connection EventLoop, including while the broker is shutting down.
+            log.debug()
+                    .attr("consumer", consumer)
+                    .attr("executorShutdown", dispatchMessagesThread.isShutdown())
+                    .exception(e)
+                    .log("Unable to schedule flow control update");
+        }
     }
 
-    private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) {
-        if (!consumerSet.contains(consumer)) {
+    private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) {
+        boolean connected;
+        int updatedTotalAvailablePermits = 0;
+        synchronized (this) {
+            consumer.completePendingDispatcherFlow(additionalNumberOfMessages);
+            connected = containsConsumerInstance(consumer);
+            if (connected) {
+                totalAvailablePermits += additionalNumberOfMessages;
+                updatedTotalAvailablePermits = totalAvailablePermits;
+            }
+        }
+
+        if (!connected) {
             log.debug()
                     .attr("consumer", consumer)
                     .log("Ignoring flow control from disconnected consumer");
             return;
         }
 
-        totalAvailablePermits += additionalNumberOfMessages;
-
         log.debug()
                 .attr("consumer", consumer)
-                .attr("totalAvailablePermits", totalAvailablePermits)
+                .attr("totalAvailablePermits", updatedTotalAvailablePermits)
                 .attr("additionalNumberOfMessages", additionalNumberOfMessages)
                 .log("Trigger new read after receiving flow control message");
-        readMoreEntriesAsync();
+        readMoreEntries();
     }
 
     /**

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

@@ -35,6 +35,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -244,9 +245,12 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE
                 consumer.getPendingAcks().forEach((ledgerId, entryId, batchSize, stickyKeyHash) -> {
                     addMessageToReplay(ledgerId, entryId, stickyKeyHash);
                 });
-                totalAvailablePermits -= consumer.getAvailablePermits();
+                // Restore the invariant that the dispatcher total equals the sum of the removal balances of the
+                // remaining consumers. Exclude Flow permits that have not updated the dispatcher total yet.
+                int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval();
+                totalAvailablePermits -= availablePermits;
                 log.debug()
-                        .attr("availablePermits", consumer.getAvailablePermits())
+                        .attr("availablePermits", availablePermits)
                         .attr("totalAvailablePermits", totalAvailablePermits)
                         .log("Decreased totalAvailablePermits by in PersistentDispatcherMultipleConsumers. "
                                 + "New dispatcher permit count is");
@@ -280,24 +284,42 @@ private synchronized void clearComponentsAfterRemovedAllConsumers() {
 
     @Override
     public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) {
-        topic.getBrokerService().executor().execute(() -> {
-            internalConsumerFlow(consumer, additionalNumberOfMessages);
-        });
+        Runnable flowTask = () -> internalConsumerFlow(consumer, additionalNumberOfMessages);
+        try {
+            dispatchMessagesThread.execute(flowTask);
+        } catch (RejectedExecutionException e) {
+            // Leave the permits pending so removal excludes this unapplied Flow. Never wait for the dispatcher
+            // monitor on the connection EventLoop, including while the broker is shutting down.
+            log.debug()
+                    .attr("consumer", consumer)
+                    .attr("executorShutdown", dispatchMessagesThread.isShutdown())
+                    .exception(e)
+                    .log("Unable to schedule flow control update");
+        }
     }
 
-    private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) {
-        if (!consumerSet.contains(consumer)) {
+    private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) {
+        boolean connected;
+        int updatedTotalAvailablePermits = 0;
+        synchronized (this) {
+            consumer.completePendingDispatcherFlow(additionalNumberOfMessages);
+            connected = containsConsumerInstance(consumer);
+            if (connected) {
+                totalAvailablePermits += additionalNumberOfMessages;
+                updatedTotalAvailablePermits = totalAvailablePermits;
+            }
+        }
+
+        if (!connected) {
             log.debug()
                     .attr("consumer", consumer)
                     .log("Ignoring flow control from disconnected consumer");
             return;
         }
 
-        totalAvailablePermits += additionalNumberOfMessages;
-
         log.debug()
                 .attr("consumer", consumer)
-                .attr("totalAvailablePermits", totalAvailablePermits)
+                .attr("totalAvailablePermits", updatedTotalAvailablePermits)
                 .attr("additionalNumberOfMessages", additionalNumberOfMessages)
                 .log("- Trigger new read after receiving flow control message with permits " + "after adding permits");
         readMoreEntries();
@@ -307,6 +329,7 @@ private synchronized void internalConsumerFlow(Consumer consumer, int additional
      * We should not call readMoreEntries() recursively in the same thread as there is a risk of StackOverflowError.
      *
      */
+    @Override
     public void readMoreEntriesAsync() {
         topic.getBrokerService().executor().execute(this::readMoreEntries);
     }