[improve][client][broker] PIP-491: Prevent Delivery Stalls by Making the Client Return Exactly the Permits Used by the Broker

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

@@ -320,15 +320,15 @@ public boolean readCompacted() {
         return readCompacted;
     }
 
-    public Future<Void> sendMessages(final List<? extends Entry> entries, EntryBatchSizes batchSizes,
+    public SendMessageResult sendMessages(final List<? extends Entry> entries, EntryBatchSizes batchSizes,
                                      EntryBatchIndexesAcks batchIndexesAcks,
                                      int totalMessages, long totalBytes, long totalChunkedMessages,
                                      RedeliveryTracker redeliveryTracker) {
         return sendMessages(entries, batchSizes, batchIndexesAcks, totalMessages, totalBytes,
                 totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH);
     }
 
-    public Future<Void> sendMessages(final List<? extends Entry> entries, EntryBatchSizes batchSizes,
+    public SendMessageResult sendMessages(final List<? extends Entry> entries, EntryBatchSizes batchSizes,
                                      EntryBatchIndexesAcks batchIndexesAcks,
                                      int totalMessages, long totalBytes, long totalChunkedMessages,
                                      RedeliveryTracker redeliveryTracker, long epoch) {
@@ -340,9 +340,9 @@ public Future<Void> sendMessages(final List<? extends Entry> entries, EntryBatch
      * Dispatch a list of entries to the consumer. <br/>
      * <b>It is also responsible to release entries data and recycle entries object.</b>
      *
-     * @return a SendMessageInfo object that contains the detail of what was sent to consumer
+     * @return the finalized per-entry permit accounting and the asynchronous write result
      */
-    public Future<Void> sendMessages(final List<? extends Entry> entries,
+    public SendMessageResult sendMessages(final List<? extends Entry> entries,
                                      final List<Integer> stickyKeyHashes,
                                      EntryBatchSizes batchSizes,
                                      EntryBatchIndexesAcks batchIndexesAcks,
@@ -353,6 +353,8 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                                      long epoch) {
         this.lastConsumedTimestamp = System.currentTimeMillis();
 
+        SendMessageResult sendMessageResult = new SendMessageResult(entries.size());
+
         if (entries.isEmpty() || totalMessages == 0) {
             log.debug("List of messages is empty, triggering write future immediately");
             batchSizes.recyle();
@@ -361,20 +363,26 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
             }
             final Promise<Void> writePromise = cnx.newPromise();
             writePromise.setSuccess(null);
-            return writePromise;
+            sendMessageResult.setSendFuture(writePromise);
+            return sendMessageResult;
         }
-        int unackedMessages = totalMessages;
         int totalEntries = 0;
 
         for (int i = 0; i < entries.size(); i++) {
             Entry entry = entries.get(i);
             if (entry != null) {
-                totalEntries++;
+                int batchSize = batchSizes.getBatchSize(i);
+                int messagePermits = batchIndexesAcks == null
+                        ? batchSize : batchIndexesAcks.getUnackedIndexCount(i, batchSize);
+                if (messagePermits == 0) {
+                    entries.set(i, null);
+                    entry.release();
+                    continue;
+                }
                 // Note
                 // Must ensure that the message is written to the pendingAcks before sent is first,
                 // because this consumer is possible to disconnect at this time.
                 if (pendingAcks != null) {
-                    int batchSize = batchSizes.getBatchSize(i);
                     int stickyKeyHash;
                     if (stickyKeyHashes == null) {
                         if (entry instanceof EntryAndMetadata entryAndMetadata) {
@@ -386,22 +394,13 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                         stickyKeyHash = stickyKeyHashes.get(i);
                     }
                     boolean sendingAllowed;
-                    long[] ackSet = batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i);
-                    int remainingUnacked;
-                    if (ackSet != null) {
-                        remainingUnacked = BitSet.valueOf(ackSet).cardinality();
-                        unackedMessages -= (batchSize - remainingUnacked);
-                    } else {
-                        remainingUnacked = batchSize;
-                    }
                     sendingAllowed =
                             pendingAcks.addPendingAckIfAllowed(entry.getLedgerId(), entry.getEntryId(),
-                                    remainingUnacked, stickyKeyHash);
+                                    messagePermits, stickyKeyHash);
                     if (!sendingAllowed) {
                         // sending isn't allowed when pending acks doesn't accept adding the entry
                         // this happens when Key_Shared draining hashes contains the stickyKeyHash
                         // because of race conditions, it might be resolved at the time of sending
-                        totalEntries--;
                         entries.set(i, null);
                         entry.release();
                         log.debug()
@@ -415,32 +414,38 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                                 .attr("entryId", entry.getEntryId())
                                 .attr("batchSize", batchSize)
                                 .log("Added entry to pendingAcks");
+                        totalEntries++;
+                        sendMessageResult.recordMessagePermits(i, messagePermits);
                     }
+                } else {
+                    totalEntries++;
+                    sendMessageResult.recordMessagePermits(i, messagePermits);
                 }
             }
         }
 
         // calculate avg message per entry
-        if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1
+        int sentMessagePermits = sendMessageResult.getTotalMessagePermits();
+        if (totalEntries > 0 && avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1
             // set init value.
-            avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries);
-        } else {
+            avgMessagesPerEntry.set(1.0 * sentMessagePermits / totalEntries);
+        } else if (totalEntries > 0) {
             avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent
-                    + (1 - avgPercent) * totalMessages / totalEntries);
+                    + (1 - avgPercent) * sentMessagePermits / totalEntries);
         }
 
-        // reduce permit and increment unackedMsg count with total number of messages in batch-msgs
-        int ackedCount = batchIndexesAcks == null ? 0 : batchIndexesAcks.getTotalAckedIndexCount();
-        MESSAGE_PERMITS_UPDATER.addAndGet(this, ackedCount - totalMessages);
+        // Reduce permits by the message count represented by the commands that will actually be sent.
+        MESSAGE_PERMITS_UPDATER.addAndGet(this, -sentMessagePermits);
         log.debug()
-                .attr("ackedCount", ackedCount)
+                .attr("sentMessagePermits", sentMessagePermits)
                 .attr("totalMessages", totalMessages)
                 .attr("avgMessagesPerEntry", avgMessagesPerEntry.get())
                 .log("Added minus messages to MESSAGE_PERMITS_UPDATER");
-        incrementUnackedMessages(unackedMessages);
+        incrementUnackedMessages(sentMessagePermits);
         Future<Void> writeAndFlushPromise =
                 cnx.getCommandSender().sendMessagesToConsumer(consumerId, topicName, subscription, partitionIdx,
-                        entries, batchSizes, batchIndexesAcks, redeliveryTracker, epoch);
+                        entries, batchSizes, batchIndexesAcks, redeliveryTracker, epoch, sendMessageResult);
+        sendMessageResult.setSendFuture(writeAndFlushPromise);
         writeAndFlushPromise.addListener(status -> {
             // only increment counters after the messages have been successfully written to the TCP/IP connection
             if (status.isSuccess()) {
@@ -457,7 +462,7 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                         .log("Sent messages to client failed by IO exception, closing the connection");
             }
         });
-        return writeAndFlushPromise;
+        return sendMessageResult;
     }
 
     private void incrementUnackedMessages(int unackedMessages) {

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

@@ -20,7 +20,6 @@
 
 
 import io.netty.util.Recycler;
-import java.util.BitSet;
 import org.apache.commons.lang3.tuple.Pair;
 
 @SuppressWarnings("unchecked")
@@ -37,17 +36,36 @@ public long[] getAckSet(int entryIdx) {
         return pair == null ? null : pair.getRight();
     }
 
+    public int getUnackedIndexCount(int entryIdx, int batchSize) {
+        Pair<Integer, long[]> pair = indexesAcks[entryIdx];
+        return pair == null ? batchSize : getCardinality(pair.getRight(), batchSize);
+    }
+
     public int getTotalAckedIndexCount() {
         int count = 0;
         for (int i = 0; i < size; i++) {
             Pair<Integer, long[]> pair = indexesAcks[i];
             if (pair != null) {
-                count += pair.getLeft() - BitSet.valueOf(pair.getRight()).cardinality();
+                count += pair.getLeft() - getUnackedIndexCount(i, pair.getLeft());
             }
         }
         return count;
     }
 
+    private static int getCardinality(long[] ackSet, int batchSize) {
+        int cardinality = 0;
+        int completeWords = Math.min(batchSize >>> 6, ackSet.length);
+        for (int i = 0; i < completeWords; i++) {
+            cardinality += Long.bitCount(ackSet[i]);
+        }
+        int remainingBits = batchSize & 63;
+        if (remainingBits > 0 && completeWords < ackSet.length) {
+            long mask = -1L >>> (Long.SIZE - remainingBits);
+            cardinality += Long.bitCount(ackSet[completeWords] & mask);
+        }
+        return cardinality;
+    }
+
     public void recycle() {
         for (int i = 0; i < size; i++) {
             indexesAcks[i] = null;

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

@@ -86,7 +86,8 @@ void sendLookupResponse(String brokerServiceUrl, String brokerServiceUrlTls, boo
     Future<Void> sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription,
                                         int partitionIdx, List<? extends Entry> entries, EntryBatchSizes batchSizes,
                                         EntryBatchIndexesAcks batchIndexesAcks,
-                                        RedeliveryTracker redeliveryTracker, long epoch);
+                                        RedeliveryTracker redeliveryTracker, long epoch,
+                                        SendMessageResult sendMessageResult);
 
     void sendTcClientConnectResponse(long requestId, ServerError error, String message);
 

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

@@ -235,7 +235,8 @@ public boolean sendTopicMigrated(ResourceType type, long resourceId, String brok
     public ChannelPromise sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription,
                                                  int partitionIdx, List<? extends Entry> entries,
                                                  EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks,
-                                                 RedeliveryTracker redeliveryTracker, long epoch) {
+                                                 RedeliveryTracker redeliveryTracker, long epoch,
+                                                 SendMessageResult sendMessageResult) {
         final ChannelHandlerContext ctx = cnx.ctx();
         final ChannelPromise writePromise = ctx.newPromise();
         ctx.channel().eventLoop().execute(() -> {
@@ -294,10 +295,12 @@ public ChannelPromise sendMessagesToConsumer(long consumerId, String topicName,
                 int redeliveryCount = redeliveryTracker
                         .getRedeliveryCount(entry.getLedgerId(), entry.getEntryId());
 
+                long[] ackSet = batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i);
+                int messagePermits = sendMessageResult.getMessagePermits(i);
+
                 ctx.write(
                         cnx.newMessageAndIntercept(consumerId, entry.getLedgerId(), entry.getEntryId(), partitionIdx,
-                                redeliveryCount, metadataAndPayload,
-                                batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i), topicName, epoch),
+                                redeliveryCount, metadataAndPayload, ackSet, topicName, epoch, messagePermits),
                         ctx.voidPromise());
                 entriesToRelease.add(entry);
             }

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

@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.broker.service;
+
+import io.netty.util.concurrent.Future;
+
+/**
+ * Finalized permit accounting for one call to {@link Consumer#sendMessages}.
+ *
+ * <p>The per-entry values are finalized after admission and remain available to the asynchronous command sender.
+ * Dispatchers use the total instead of reconstructing it from inputs that the sender can recycle asynchronously.
+ */
+public final class SendMessageResult {
+    private final int[] messagePermits;
+    private int totalMessagePermits;
+    private Future<Void> sendFuture;
+
+    SendMessageResult(int entries) {
+        this.messagePermits = new int[entries];
+    }
+
+    void recordMessagePermits(int entryIndex, int permits) {
+        if (permits <= 0) {
+            throw new IllegalArgumentException("Message permits must be positive");
+        }
+        messagePermits[entryIndex] = permits;
+        totalMessagePermits = Math.addExact(totalMessagePermits, permits);
+    }
+
+    int getMessagePermits(int entryIndex) {
+        return messagePermits[entryIndex];
+    }
+
+    void setSendFuture(Future<Void> sendFuture) {
+        this.sendFuture = sendFuture;
+    }
+
+    public int getTotalMessagePermits() {
+        return totalMessagePermits;
+    }
+
+    public Future<Void> getSendFuture() {
+        return sendFuture;
+    }
+}

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

@@ -4489,6 +4489,19 @@ public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long e
             int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet, String topic, long epoch) {
         BaseCommand command = Commands.newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount,
                 ackSet, epoch);
+        return newMessageAndIntercept(consumerId, ledgerId, entryId, metadataAndPayload, topic, command);
+    }
+
+    public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId, int partition,
+            int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet, String topic, long epoch,
+            int messagePermits) {
+        BaseCommand command = Commands.newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount,
+                ackSet, epoch, messagePermits);
+        return newMessageAndIntercept(consumerId, ledgerId, entryId, metadataAndPayload, topic, command);
+    }
+
+    private ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId,
+            ByteBuf metadataAndPayload, String topic, BaseCommand command) {
         ByteBufPair res = Commands.serializeCommandMessageWithSize(command, metadataAndPayload);
         if (brokerInterceptor != null) {
             try {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentDispatcherMultipleConsumers.java

@@ -35,6 +35,7 @@
 import org.apache.pulsar.broker.service.RedeliveryTracker;
 import org.apache.pulsar.broker.service.RedeliveryTrackerDisabled;
 import org.apache.pulsar.broker.service.SendMessageInfo;
+import org.apache.pulsar.broker.service.SendMessageResult;
 import org.apache.pulsar.broker.service.Subscription;
 import org.apache.pulsar.broker.service.persistent.DispatchRateLimiter;
 import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType;
@@ -204,10 +205,11 @@ public synchronized void sendMessages(List<Entry> entries) {
             SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal();
             EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size());
             filterEntriesForConsumer(entries, batchSizes, sendMessageInfo, null, null, false, consumer);
-            consumer.sendMessages(entries, batchSizes, null, sendMessageInfo.getTotalMessages(),
+            SendMessageResult sendResult = consumer.sendMessages(entries, batchSizes, null,
+                    sendMessageInfo.getTotalMessages(),
                     sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker());
 
-            TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendMessageInfo.getTotalMessages());
+            TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits());
         } else {
             entries.forEach(entry -> {
                 int totalMsgs = getNumberOfMessagesInBatch(entry);

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentStickyKeyDispatcherMultipleConsumers.java

@@ -35,6 +35,7 @@
 import org.apache.pulsar.broker.service.HashRangeAutoSplitStickyKeyConsumerSelector;
 import org.apache.pulsar.broker.service.HashRangeExclusiveStickyKeyConsumerSelector;
 import org.apache.pulsar.broker.service.SendMessageInfo;
+import org.apache.pulsar.broker.service.SendMessageResult;
 import org.apache.pulsar.broker.service.StickyKeyConsumerSelector;
 import org.apache.pulsar.broker.service.Subscription;
 import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType;
@@ -178,11 +179,12 @@ public synchronized void sendMessages(List<Entry> entries) {
             filterEntriesForConsumer(entriesForConsumer, batchSizes, sendMessageInfo, null, null, false, consumer);
 
             if (consumer.getAvailablePermits() > 0 && consumer.isWritable()) {
-                consumer.sendMessages(entriesForConsumer, stickyKeysForConsumer, batchSizes,
+                SendMessageResult sendResult = consumer.sendMessages(entriesForConsumer, stickyKeysForConsumer,
+                        batchSizes,
                         null, sendMessageInfo.getTotalMessages(),
                         sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(),
                         getRedeliveryTracker(), Commands.DEFAULT_CONSUMER_EPOCH);
-                TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendMessageInfo.getTotalMessages());
+                TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits());
             } else {
                 entriesForConsumer.forEach(e -> {
                     int totalMsgs = getNumberOfMessagesInBatch(e);