[improve][broker] Enable Consumer Throttling and Accurate Unacknowledged Message Tracking for Exclusive and Failover Subscriptions

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

@@ -41,6 +41,7 @@
  * running.
  */
 public class PendingAcksMap {
+
     /**
      * Callback interface for handling the addition of pending acknowledgments.
      */
@@ -227,6 +228,15 @@ public void forEachAndClose(PendingAcksConsumer processor) {
         }
     }
 
+    public void clear() {
+        try {
+            writeLock.lock();
+            pendingAcks.clear();
+        } finally {
+            writeLock.unlock();
+        }
+    }
+
     /**
      * Check if the map contains a pending ack for the given ledger ID and entry ID.
      *
@@ -330,9 +340,10 @@ public boolean remove(long ledgerId, long entryId) {
      *
      * @param markDeleteLedgerId the ledger ID up to which to remove pending acks
      * @param markDeleteEntryId the entry ID up to which to remove pending acks
+     * @return the sum batchSize of removed pending acks
      */
-    public void removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId) {
-        internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, false);
+    public int removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId) {
+        return internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, false);
     }
 
     /**
@@ -345,15 +356,18 @@ public void removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId) {
      * @param markDeleteLedgerId the ledger ID up to which to remove pending acks
      * @param markDeleteEntryId the entry ID up to which to remove pending acks
      * @param useWriteLock true if the method should use a write lock, false otherwise
+     * @return the sum batchSize of removed pending acks
      */
-    private void internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntryId, boolean useWriteLock) {
+    private int internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntryId, boolean useWriteLock) {
         PendingAcksRemoveHandler pendingAcksRemoveHandler = pendingAcksRemoveHandlerSupplier.get();
         // track if the write lock was acquired
         boolean acquiredWriteLock = false;
         // track if a batch was started
         boolean batchStarted = false;
         // track if the method should retry with a write lock
         boolean retryWithWriteLock = false;
+        // the number of removed pending acks
+        int removeCount = 0;
         try {
             if (useWriteLock) {
                 writeLock.lock();
@@ -378,9 +392,10 @@ private void internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntry
                 while (entryMapIterator.hasNext()) {
                     Long2ObjectMap.Entry<IntIntPair> intIntPairEntry = entryMapIterator.next();
                     long entryId = intIntPairEntry.getLongKey();
+                    int batchSize = intIntPairEntry.getValue().leftInt();
                     if (!acquiredWriteLock) {
                         retryWithWriteLock = true;
-                        return;
+                        break;
                     }
                     if (pendingAcksRemoveHandler != null) {
                         if (!batchStarted) {
@@ -390,12 +405,16 @@ private void internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntry
                         int stickyKeyHash = intIntPairEntry.getValue().rightInt();
                         pendingAcksRemoveHandler.handleRemoving(consumer, ledgerId, entryId, stickyKeyHash, closed);
                     }
+                    removeCount += batchSize;
                     entryMapIterator.remove();
                 }
+                if (retryWithWriteLock) {
+                    break;
+                }
                 if (ledgerMap.isEmpty()) {
                     if (!acquiredWriteLock) {
                         retryWithWriteLock = true;
-                        return;
+                        break;
                     }
                     ledgerMapIterator.remove();
                 }
@@ -409,10 +428,11 @@ private void internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntry
             } else {
                 readLock.unlock();
                 if (retryWithWriteLock) {
-                    internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, true);
+                    removeCount = internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, true);
                 }
             }
         }
+        return removeCount;
     }
 
     private void handleRemovePendingAck(long ledgerId, long entryId, int stickyKeyHash) {

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

@@ -89,9 +89,9 @@ public PersistentDispatcherSingleActiveConsumer(ManagedCursor cursor, SubType su
                 : ""/* NonDurableCursor doesn't have name */);
         this.readBatchSize = serviceConfig.getDispatcherMaxReadBatchSize();
         this.readFailureBackoff = new Backoff(serviceConfig.getDispatcherReadFailureBackoffInitialTimeInMs(),
-            TimeUnit.MILLISECONDS, serviceConfig.getDispatcherReadFailureBackoffMaxTimeInMs(),
-            TimeUnit.MILLISECONDS, serviceConfig.getDispatcherReadFailureBackoffMandatoryStopTimeInMs(),
-            TimeUnit.MILLISECONDS);
+                TimeUnit.MILLISECONDS, serviceConfig.getDispatcherReadFailureBackoffMaxTimeInMs(),
+                TimeUnit.MILLISECONDS, serviceConfig.getDispatcherReadFailureBackoffMandatoryStopTimeInMs(),
+                TimeUnit.MILLISECONDS);
         this.redeliveryTracker = RedeliveryTrackerDisabled.REDELIVERY_TRACKER_DISABLED;
         this.initializeDispatchRateLimiterIfNeeded();
     }
@@ -231,18 +231,19 @@ protected void dispatchEntriesToConsumer(Consumer currentConsumer, List<Entry> e
                                              EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks,
                                              SendMessageInfo sendMessageInfo, long epoch) {
         currentConsumer
-            .sendMessages(entries, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(),
-                    sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(),
-                    redeliveryTracker, epoch)
-            .addListener(future -> {
-                if (future.isSuccess()) {
-                    acquirePermitsForDeliveredMessages(topic, cursor, entries.size(),
-                            sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes());
-
-                    // Schedule a new read batch operation only after the previous batch has been written to the socket.
-                    executor.execute(() -> readMoreEntries(getActiveConsumer()));
-                }
-            });
+                .sendMessages(entries, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(),
+                        sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(),
+                        redeliveryTracker, epoch)
+                .addListener(future -> {
+                    if (future.isSuccess()) {
+                        acquirePermitsForDeliveredMessages(topic, cursor, entries.size(),
+                                sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes());
+
+                        // Schedule a new read batch operation only after the previous batch has been written to the
+                        // socket.
+                        executor.execute(() -> readMoreEntries(getActiveConsumer()));
+                    }
+                });
     }
 
     @Override
@@ -435,6 +436,11 @@ protected Pair<Integer, Long> calculateToRead(Consumer consumer) {
             int avgMessagesPerEntry = Math.max(1, consumer.getAvgMessagesPerEntry());
             messagesToRead = Math.min((int) Math.ceil(availablePermits * 1.0 / avgMessagesPerEntry), readBatchSize);
         }
+        // We need to ensure that we don't exceed the max unacked messages.
+        int maxUnackedMessages = consumer.getMaxUnackedMessages();
+        if (maxUnackedMessages > 0) {
+            messagesToRead = Math.min(messagesToRead, maxUnackedMessages - consumer.getUnackedMessages());
+        }
 
         // throttle only if: (1) cursor is not active (or flag for throttle-nonBacklogConsumer is enabled) bcz
         // active-cursor reads message from cache rather from bookkeeper (2) if topic has reached message-rate
@@ -585,6 +591,11 @@ public boolean checkAndUnblockIfStuck() {
         return false;
     }
 
+    public int getUnackedMessages() {
+        Consumer activeConsumer = getActiveConsumer();
+        return activeConsumer != null ? activeConsumer.getUnackedMessages() : 0;
+    }
+
     private static final Logger log = LoggerFactory.getLogger(PersistentDispatcherSingleActiveConsumer.class);
 
     public static class ReadEntriesCtx {
@@ -597,13 +608,14 @@ public static class ReadEntriesCtx {
         private ReadEntriesCtx(Recycler.Handle<ReadEntriesCtx> recyclerHandle) {
             this.recyclerHandle = recyclerHandle;
         }
+
         private static final Recycler<ReadEntriesCtx> RECYCLER =
                 new Recycler<ReadEntriesCtx>() {
-            @Override
-            protected ReadEntriesCtx newObject(Recycler.Handle<ReadEntriesCtx> recyclerHandle) {
-                return new ReadEntriesCtx(recyclerHandle);
-            }
-        };
+                    @Override
+                    protected ReadEntriesCtx newObject(Recycler.Handle<ReadEntriesCtx> recyclerHandle) {
+                        return new ReadEntriesCtx(recyclerHandle);
+                    }
+                };
 
         public static ReadEntriesCtx create(Consumer consumer, long epoch) {
             ReadEntriesCtx readEntriesCtx = RECYCLER.get();

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

@@ -1315,7 +1315,9 @@ public CompletableFuture<SubscriptionStatsImpl> getStatsAsync(GetStatsOptions ge
         SubType subType = getType();
         subStats.type = getTypeString();
         if (dispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
-            Consumer activeConsumer = ((PersistentDispatcherSingleActiveConsumer) dispatcher).getActiveConsumer();
+            PersistentDispatcherSingleActiveConsumer d = ((PersistentDispatcherSingleActiveConsumer) dispatcher);
+            subStats.unackedMessages = d.getUnackedMessages();
+            Consumer activeConsumer = d.getActiveConsumer();
             if (activeConsumer != null) {
                 subStats.activeConsumerName = activeConsumer.consumerName();
             }

pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/MaxUnackedMessagesTest.java

@@ -252,12 +252,12 @@ public void testMaxUnackedMessagesOnSubApplied() throws Exception {
     }
 
     @Test(timeOut = 30000)
-    public void testMaxUnackedMessagesOnConsumer() throws Exception {
+    public void testMaxUnackedMessagesOnSharedConsumer() throws Exception {
         final String topicName = testTopic + System.currentTimeMillis();
         final String subscriberName = "test-sub" + System.currentTimeMillis();
         final int unackMsgAllowed = 100;
         final int receiverQueueSize = 10;
-        final int totalProducedMsgs = 300;
+        final int totalProducedMsgs = 400;
 
         ConsumerBuilder<String> consumerBuilder = pulsarClient.newConsumer(Schema.STRING).topic(topicName)
                 .subscriptionName(subscriberName).receiverQueueSize(receiverQueueSize)
@@ -305,16 +305,183 @@ public void testMaxUnackedMessagesOnConsumer() throws Exception {
             String message = "my-message-" + i;
             producer.send(message);
         }
+        AtomicInteger consumer1Counter = new AtomicInteger(0);
         AtomicInteger consumer2Counter = new AtomicInteger(0);
         AtomicInteger consumer3Counter = new AtomicInteger(0);
-        CountDownLatch countDownLatch = new CountDownLatch(2);
+        CountDownLatch countDownLatch = new CountDownLatch(3);
+        startConsumer(consumer1, consumer1Counter, countDownLatch);
         startConsumer(consumer2, consumer2Counter, countDownLatch);
         startConsumer(consumer3, consumer3Counter, countDownLatch);
         countDownLatch.await(10, TimeUnit.SECONDS);
+        assertEquals(consumer1Counter.get(), unackMsgAllowed);
         assertEquals(consumer2Counter.get(), unackMsgAllowed);
         assertEquals(consumer3Counter.get(), unackMsgAllowed);
     }
 
+    @Test(timeOut = 30000)
+    public void testMaxUnackedMessagesOnExclusiveConsumer() throws Exception {
+        final String topicName = testTopic + System.currentTimeMillis();
+        final String subscriberName = "test-sub-exclusive" + System.currentTimeMillis();
+        final int unackMsgAllowed = 100;
+        final int receiverQueueSize = 10;
+        final int totalProducedMsgs = 300;
+
+        ConsumerBuilder<String> consumerBuilder = pulsarClient.newConsumer(Schema.STRING).topic(topicName)
+                .subscriptionName(subscriberName).receiverQueueSize(receiverQueueSize)
+                .ackTimeout(1, TimeUnit.MINUTES)
+                .subscriptionType(SubscriptionType.Exclusive);
+        @Cleanup
+        Consumer<String> consumer = consumerBuilder.subscribe();
+        // 1) Produced Messages
+        @Cleanup
+        Producer<String> producer = pulsarClient.newProducer(Schema.STRING).topic(topicName).create();
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            String message = "my-message-" + i;
+            producer.send(message);
+        }
+        // 2) Unlimited, so all messages can be consumed
+        int count = 0;
+        List<Message<String>> list = new ArrayList<>(totalProducedMsgs);
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            Message<String> message = consumer.receive(1, TimeUnit.SECONDS);
+            if (message == null) {
+                break;
+            }
+            count++;
+            list.add(message);
+        }
+        assertEquals(count, totalProducedMsgs);
+        list.forEach(message -> {
+            try {
+                consumer.acknowledge(message);
+            } catch (PulsarClientException e) {
+            }
+        });
+        // 3) Set restrictions, so only part of the data can be consumed
+        waitCacheInit(topicName);
+        admin.topics().setMaxUnackedMessagesOnConsumer(topicName, unackMsgAllowed);
+        Awaitility.await().untilAsserted(()
+                -> assertNotNull(admin.topics().getMaxUnackedMessagesOnConsumer(topicName)));
+        assertEquals(admin.topics().getMaxUnackedMessagesOnConsumer(topicName).intValue(), unackMsgAllowed);
+        // 4) consumer can only consume 100 messages
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            String message = "my-message-" + i;
+            producer.send(message);
+        }
+        int consumerCounter = 0;
+        Message<String> message = null;
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            try {
+                Message<String> msg = consumer.receive(500, TimeUnit.MILLISECONDS);
+                if (msg == null) {
+                    break;
+                }
+                message = msg;
+                ++consumerCounter;
+            } catch (PulsarClientException e) {
+                break;
+            }
+        }
+        assertEquals(consumerCounter, unackMsgAllowed);
+        consumer.acknowledgeCumulative(message.getMessageId());
+        consumerCounter = 0;
+        for (int i = 0; i < totalProducedMsgs - unackMsgAllowed; i++) {
+            try {
+                message = consumer.receive(500, TimeUnit.MILLISECONDS);
+                if (message == null) {
+                    break;
+                }
+                ++consumerCounter;
+            } catch (PulsarClientException e) {
+                break;
+            }
+        }
+        assertEquals(consumerCounter, unackMsgAllowed);
+    }
+
+    @Test(timeOut = 30000)
+    public void testMaxUnackedMessagesOnFailOverConsumer() throws Exception {
+        final String topicName = testTopic + System.currentTimeMillis();
+        final String subscriberName = "test-sub-failover" + System.currentTimeMillis();
+        final int unackMsgAllowed = 100;
+        final int receiverQueueSize = 10;
+        final int totalProducedMsgs = 300;
+
+        ConsumerBuilder<String> consumerBuilder = pulsarClient.newConsumer(Schema.STRING).topic(topicName)
+                .subscriptionName(subscriberName).receiverQueueSize(receiverQueueSize)
+                .ackTimeout(1, TimeUnit.MINUTES)
+                .subscriptionType(SubscriptionType.Failover);
+        @Cleanup
+        Consumer<String> consumer = consumerBuilder.subscribe();
+        // 1) Produced Messages
+        @Cleanup
+        Producer<String> producer = pulsarClient.newProducer(Schema.STRING).topic(topicName).create();
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            String message = "my-message-" + i;
+            producer.send(message);
+        }
+        // 2) Unlimited, so all messages can be consumed
+        int count = 0;
+        List<Message<String>> list = new ArrayList<>(totalProducedMsgs);
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            Message<String> message = consumer.receive(1, TimeUnit.SECONDS);
+            if (message == null) {
+                break;
+            }
+            count++;
+            list.add(message);
+        }
+        assertEquals(count, totalProducedMsgs);
+        list.forEach(message -> {
+            try {
+                consumer.acknowledge(message);
+            } catch (PulsarClientException e) {
+            }
+        });
+        // 3) Set restrictions, so only part of the data can be consumed
+        waitCacheInit(topicName);
+        admin.topics().setMaxUnackedMessagesOnConsumer(topicName, unackMsgAllowed);
+        Awaitility.await().untilAsserted(()
+                -> assertNotNull(admin.topics().getMaxUnackedMessagesOnConsumer(topicName)));
+        assertEquals(admin.topics().getMaxUnackedMessagesOnConsumer(topicName).intValue(), unackMsgAllowed);
+        // 4) consumer can only consume 100 messages
+        for (int i = 0; i < totalProducedMsgs; i++) {
+            String message = "my-message-" + i;
+            producer.send(message);
+        }
+        int consumerCounter = 0;
+        while (true) {
+            try {
+                Message<String> message = consumer.receive(500, TimeUnit.MILLISECONDS);
+                if (message == null) {
+                    break;
+                }
+                ++consumerCounter;
+            } catch (PulsarClientException e) {
+                break;
+            }
+        }
+        assertEquals(consumerCounter, unackMsgAllowed);
+
+        // 5) failover consumer can only consume 100 messages
+        @Cleanup
+        Consumer<String> failoverConsumer = consumerBuilder.subscribe();
+        consumer.close();
+        consumerCounter = 0;
+        while (true) {
+            try {
+                Message<String> message = failoverConsumer.receive(1, TimeUnit.SECONDS);
+                if (message == null) {
+                    break;
+                }
+                ++consumerCounter;
+            } catch (PulsarClientException e) {
+                break;
+            }
+        }
+        assertEquals(consumerCounter, unackMsgAllowed);
+    }
+
     private void startConsumer(Consumer<String> consumer, AtomicInteger consumerCounter,
                                CountDownLatch countDownLatch) {
         new Thread(() -> {