diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index a7f296536881e..868eea1b4294a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -320,7 +320,7 @@ public boolean readCompacted() { return readCompacted; } - public Future sendMessages(final List entries, EntryBatchSizes batchSizes, + public SendMessageResult sendMessages(final List entries, EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker) { @@ -328,7 +328,7 @@ public Future sendMessages(final List entries, EntryBatch totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH); } - public Future sendMessages(final List entries, EntryBatchSizes batchSizes, + public SendMessageResult sendMessages(final List entries, EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { @@ -340,9 +340,9 @@ public Future sendMessages(final List entries, EntryBatch * Dispatch a list of entries to the consumer.
* It is also responsible to release entries data and recycle entries object. * - * @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 sendMessages(final List entries, + public SendMessageResult sendMessages(final List entries, final List stickyKeyHashes, EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks, @@ -353,6 +353,8 @@ public Future sendMessages(final List 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 sendMessages(final List entries, } final Promise 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 sendMessages(final List 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 sendMessages(final List 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 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 sendMessages(final List entries, .log("Sent messages to client failed by IO exception, closing the connection"); } }); - return writeAndFlushPromise; + return sendMessageResult; } private void incrementUnackedMessages(int unackedMessages) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcks.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcks.java index 3126756b51469..9afd38ab4b7a3 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcks.java +++ b/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 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 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; diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java index d98fc35858327..d365cc68c6561 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java @@ -86,7 +86,8 @@ void sendLookupResponse(String brokerServiceUrl, String brokerServiceUrlTls, boo Future sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription, int partitionIdx, List entries, EntryBatchSizes batchSizes, EntryBatchIndexesAcks batchIndexesAcks, - RedeliveryTracker redeliveryTracker, long epoch); + RedeliveryTracker redeliveryTracker, long epoch, + SendMessageResult sendMessageResult); void sendTcClientConnectResponse(long requestId, ServerError error, String message); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java index f499818399992..e368e02bee758 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java +++ b/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 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); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessageResult.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessageResult.java new file mode 100644 index 0000000000000..3184418e26b9a --- /dev/null +++ b/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}. + * + *

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 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 sendFuture) { + this.sendFuture = sendFuture; + } + + public int getTotalMessagePermits() { + return totalMessagePermits; + } + + public Future getSendFuture() { + return sendFuture; + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java index 5e8422af870f4..e3040731ac986 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java +++ b/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 { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentDispatcherMultipleConsumers.java index d70cf5baef14d..2de3ac13f8927 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentDispatcherMultipleConsumers.java +++ b/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 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); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentStickyKeyDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentStickyKeyDispatcherMultipleConsumers.java index 657a6fcf4d179..be52e974e26df 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentStickyKeyDispatcherMultipleConsumers.java +++ b/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 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); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index b659f6e2200d8..0ae26fbb106d7 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -68,6 +68,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.SharedConsumerAssignor; import org.apache.pulsar.broker.service.Subscription; import org.apache.pulsar.broker.transaction.exception.buffer.TransactionBufferException; @@ -869,18 +870,18 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis readType == ReadType.Replay, c); totalEntriesProcessed += entriesForThisConsumer.size(); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), - sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); + SendMessageResult sendResult = c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, + sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), + sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); int msgSent = sendMessageInfo.getTotalMessages(); remainingMessages -= msgSent; start += maxEntriesInThisBatch; entriesToDispatch -= maxEntriesInThisBatch; - TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, - -(msgSent - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits()); log.debug() .attr("msgSent", msgSent) - .attr("totalAckedIndexCount", () -> batchIndexesAcks.getTotalAckedIndexCount()) + .attr("messagePermits", sendResult.getTotalMessagePermits()) .log("Added permits to TOTAL_AVAILABLE_PERMITS_UPDATER"); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); @@ -976,17 +977,17 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); totalEntriesProcessed += entryAndMetadataList.size(); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() - ).addListener(future -> { + ); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntriesAsync(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -sendResult.getTotalMessagePermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 3de50042b592d..76ff397946d11 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -65,6 +65,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.SharedConsumerAssignor; import org.apache.pulsar.broker.service.StickyKeyConsumerSelector; import org.apache.pulsar.broker.service.Subscription; @@ -759,18 +760,18 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), - sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); + SendMessageResult sendResult = c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, + sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), + sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); int msgSent = sendMessageInfo.getTotalMessages(); remainingMessages -= msgSent; start += messagesForC; entriesToDispatch -= messagesForC; - TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, - -(msgSent - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits()); log.debug() .attr("msgSent", msgSent) - .attr("totalAckedIndexCount", batchIndexesAcks.getTotalAckedIndexCount()) + .attr("messagePermits", sendResult.getTotalMessagePermits()) .log("Added -( minus) permits to TOTAL_AVAILABLE_PERMITS_UPDATER in " + "PersistentDispatcherMultipleConsumers"); totalMessagesSent += sendMessageInfo.getTotalMessages(); @@ -836,17 +837,17 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() - ).addListener(future -> { + ); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntries(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -sendResult.getTotalMessagePermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherSingleActiveConsumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherSingleActiveConsumer.java index 8c2b1d820520b..ffeec30fae136 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherSingleActiveConsumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherSingleActiveConsumer.java @@ -250,7 +250,7 @@ protected void dispatchEntriesToConsumer(Consumer currentConsumer, List e .sendMessages(entries, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker, epoch) - .addListener(future -> { + .getSendFuture().addListener(future -> { if (future.isSuccess()) { acquirePermitsForDeliveredMessages(topic, cursor, entries.size(), sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes()); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java index 07177e039b118..0757ea00b2418 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java @@ -51,6 +51,7 @@ import org.apache.pulsar.broker.service.ImpactedConsumersResult; import org.apache.pulsar.broker.service.PendingAcksMap; 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.StickyKeyDispatcher; import org.apache.pulsar.broker.service.Subscription; @@ -336,17 +337,17 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis totalEntries += filterEntriesForConsumer(entriesForConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); totalEntriesProcessed += entriesForConsumer.size(); - consumer.sendMessages(entriesForConsumer, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entriesForConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), - getRedeliveryTracker()).addListener(future -> { + getRedeliveryTracker()); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && remainingConsumersToFinishSending.decrementAndGet() == 0) { readMoreEntriesAsync(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -sendResult.getTotalMessagePermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersClassic.java index f7b614d60af32..b38a4f6d056b9 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersClassic.java @@ -51,6 +51,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.StickyKeyDispatcher; import org.apache.pulsar.broker.service.Subscription; @@ -389,17 +390,17 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(messagesForC); totalEntries += filterEntriesForConsumer(entriesWithSameKey, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); - consumer.sendMessages(entriesWithSameKey, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entriesWithSameKey, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), - getRedeliveryTracker()).addListener(future -> { + getRedeliveryTracker()); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && keyNumbers.decrementAndGet() == 0) { readMoreEntries(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -sendResult.getTotalMessagePermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } else { diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerTest.java index a974655b06c27..b5c282faa385f 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerTest.java @@ -21,14 +21,28 @@ import static java.util.Collections.emptyMap; import static org.apache.pulsar.client.api.MessageId.latest; import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Exclusive; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Shared; import static org.apache.pulsar.common.api.proto.KeySharedMode.AUTO_SPLIT; import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import io.netty.util.concurrent.ImmediateEventExecutor; import java.net.SocketAddress; +import java.util.ArrayList; +import java.util.List; +import org.apache.bookkeeper.mledger.Entry; +import org.apache.commons.lang3.tuple.Pair; import org.apache.pulsar.broker.PulsarService; import org.apache.pulsar.broker.ServiceConfiguration; +import org.apache.pulsar.broker.service.persistent.PersistentTopic; import org.apache.pulsar.common.api.proto.KeySharedMeta; import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; import org.testng.annotations.BeforeMethod; @@ -37,14 +51,17 @@ @Test(groups = "broker") public class ConsumerTest { private Consumer consumer; + private Subscription subscription; + private ServerCnx cnx; + private Topic topic; private final ConsumerStatsImpl stats = new ConsumerStatsImpl(); @BeforeMethod public void beforeMethod() { - Subscription subscription = mock(Subscription.class); - ServerCnx cnx = mock(ServerCnx.class); + subscription = mock(Subscription.class); + cnx = mock(ServerCnx.class); SocketAddress address = mock(SocketAddress.class); - Topic topic = mock(Topic.class); + topic = mock(PersistentTopic.class, RETURNS_DEEP_STUBS); BrokerService brokerService = mock(BrokerService.class); PulsarService pulsarService = mock(PulsarService.class); ServiceConfiguration serviceConfiguration = mock(ServiceConfiguration.class); @@ -73,4 +90,54 @@ public void testGetBytesOutCounter() { consumer.updateStats(stats); assertEquals(consumer.getBytesOutCounter(), 1L); } + + @Test + public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { + when(topic.getHierarchyTopicPolicies().getMaxUnackedMessagesOnConsumer().get()).thenReturn(100); + Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", true, + cnx, "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, + DEFAULT_CONSUMER_EPOCH); + sharedConsumer.setPendingAcksAddHandler((ignored, ledgerId, entryId, stickyKeyHash) -> ledgerId != 2); + sharedConsumer.flowPermits(100); + + Entry partialBatch = mock(Entry.class); + when(partialBatch.getLedgerId()).thenReturn(1L); + when(partialBatch.getEntryId()).thenReturn(1L); + Entry rejectedBatch = mock(Entry.class); + when(rejectedBatch.getLedgerId()).thenReturn(2L); + when(rejectedBatch.getEntryId()).thenReturn(2L); + Entry emptyPartialBatch = mock(Entry.class); + when(emptyPartialBatch.getLedgerId()).thenReturn(3L); + when(emptyPartialBatch.getEntryId()).thenReturn(3L); + List entries = new ArrayList<>(List.of(partialBatch, rejectedBatch, emptyPartialBatch)); + + EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); + batchSizes.setBatchSize(0, 10); + batchSizes.setBatchSize(1, 10); + batchSizes.setBatchSize(2, 10); + EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); + batchIndexesAcks.setIndexesAcks(0, Pair.of(10, new long[] {0b100101L})); + batchIndexesAcks.setIndexesAcks(2, Pair.of(10, new long[] {0L})); + + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); + when(cnx.getCommandSender()).thenReturn(commandSender); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), + any(), anyLong(), any())).thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); + + try { + SendMessageResult result = sharedConsumer.sendMessages(entries, batchSizes, batchIndexesAcks, + 23, 0, 0, mock(RedeliveryTracker.class)); + + assertEquals(result.getTotalMessagePermits(), 3); + assertEquals(sharedConsumer.getAvailablePermits(), 97); + assertEquals(sharedConsumer.getUnackedMessages(), 3); + assertNull(entries.get(1)); + assertNull(entries.get(2)); + verify(rejectedBatch).release(); + verify(emptyPartialBatch).release(); + } finally { + batchSizes.recyle(); + batchIndexesAcks.recycle(); + } + } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcksTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcksTest.java index 040653d95ac27..61f4a128beb6a 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcksTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/EntryBatchIndexesAcksTest.java @@ -19,12 +19,38 @@ package org.apache.pulsar.broker.service; import static org.testng.Assert.assertEquals; +import java.util.BitSet; +import java.util.SplittableRandom; import org.apache.commons.lang3.tuple.Pair; import org.apache.pulsar.common.util.collections.BitSetRecyclable; import org.testng.annotations.Test; public class EntryBatchIndexesAcksTest { + private static final long RANDOM_SEED = 0x491B0A7EL; + + @Test + void shouldCalculateUnackedIndexesForEachEntry() { + BitSetRecyclable bitSet = BitSetRecyclable.create(); + bitSet.set(1); + bitSet.set(4); + bitSet.set(8); + EntryBatchIndexesAcks acks = EntryBatchIndexesAcks.get(2); + try { + acks.setIndexesAcks(0, Pair.of(10, bitSet.toLongArray())); + + assertEquals(acks.getUnackedIndexCount(0, 10), 3); + assertEquals(acks.getUnackedIndexCount(1, 7), 7); + + bitSet.set(63); + acks.setIndexesAcks(0, Pair.of(10, bitSet.toLongArray())); + assertEquals(acks.getUnackedIndexCount(0, 10), 3); + } finally { + acks.recycle(); + bitSet.recycle(); + } + } + @Test void shouldResetStateBeforeReusing() { // given @@ -50,4 +76,30 @@ void shouldResetStateBeforeReusing() { assertEquals(acks.getTotalAckedIndexCount(), 0); } -} \ No newline at end of file + @Test + void shouldMatchBoundedBitSetCardinalityForRandomAckSets() { + SplittableRandom random = new SplittableRandom(RANDOM_SEED); + EntryBatchIndexesAcks acks = EntryBatchIndexesAcks.get(2); + try { + for (int testCase = 0; testCase < 10_000; testCase++) { + int batchSize = random.nextInt(1, 1025); + int requiredWords = (batchSize + Long.SIZE - 1) / Long.SIZE; + long[] ackSet = new long[requiredWords + random.nextInt(3)]; + for (int i = 0; i < ackSet.length; i++) { + ackSet[i] = random.nextLong(); + } + + BitSet expected = BitSet.valueOf(ackSet); + expected.clear(batchSize, Math.max(batchSize, expected.length())); + acks.setIndexesAcks(0, Pair.of(batchSize, ackSet)); + + assertEquals(acks.getUnackedIndexCount(0, batchSize), expected.cardinality(), + "seed=" + RANDOM_SEED + ", case=" + testCase + ", batchSize=" + batchSize); + assertEquals(acks.getUnackedIndexCount(1, batchSize), batchSize); + } + } finally { + acks.recycle(); + } + } + +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PulsarCommandSenderImplTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PulsarCommandSenderImplTest.java new file mode 100644 index 0000000000000..7b7f96983cdc7 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PulsarCommandSenderImplTest.java @@ -0,0 +1,128 @@ +/* + * 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 static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPromise; +import io.netty.channel.DefaultChannelPromise; +import io.netty.channel.EventLoop; +import io.netty.util.concurrent.ImmediateEventExecutor; +import java.util.ArrayList; +import java.util.List; +import org.apache.bookkeeper.mledger.Entry; +import org.apache.pulsar.broker.PulsarService; +import org.apache.pulsar.broker.ServiceConfiguration; +import org.apache.pulsar.common.api.proto.ProtocolVersion; +import org.apache.pulsar.common.protocol.ByteBufPair; +import org.testng.annotations.Test; + +@Test(groups = "broker") +public class PulsarCommandSenderImplTest { + + @Test + public void testSendMessagesMapsFinalizedPerEntryPermitsToWireCommands() { + ServerCnx cnx = mock(ServerCnx.class); + ChannelHandlerContext ctx = mock(ChannelHandlerContext.class); + Channel channel = mock(Channel.class); + EventLoop eventLoop = mock(EventLoop.class); + ChannelPromise writePromise = new DefaultChannelPromise(channel, ImmediateEventExecutor.INSTANCE); + ChannelPromise voidPromise = mock(ChannelPromise.class); + BrokerService brokerService = mock(BrokerService.class); + PulsarService pulsarService = mock(PulsarService.class); + ServiceConfiguration configuration = mock(ServiceConfiguration.class); + List serializedPermits = new ArrayList<>(); + + when(cnx.ctx()).thenReturn(ctx); + when(ctx.channel()).thenReturn(channel); + when(channel.eventLoop()).thenReturn(eventLoop); + doAnswer(invocation -> { + invocation.getArgument(0).run(); + return null; + }).when(eventLoop).execute(any(Runnable.class)); + when(ctx.newPromise()).thenReturn(writePromise); + when(ctx.voidPromise()).thenReturn(voidPromise); + doAnswer(invocation -> { + invocation.getArgument(1).trySuccess(); + return invocation.getArgument(1); + }).when(ctx).writeAndFlush(any(), eq(writePromise)); + when(cnx.isBatchMessageCompatibleVersion()).thenReturn(true); + when(cnx.getRemoteEndpointProtocolVersion()).thenReturn(ProtocolVersion.v18.getValue()); + when(cnx.supportBrokerMetadata()).thenReturn(true); + when(cnx.getBrokerService()).thenReturn(brokerService); + when(brokerService.getPulsar()).thenReturn(pulsarService); + when(pulsarService.getConfig()).thenReturn(configuration); + when(configuration.isExposingBrokerEntryMetadataToClientEnabled()).thenReturn(true); + when(cnx.newMessageAndIntercept(anyLong(), anyLong(), anyLong(), anyInt(), anyInt(), any(ByteBuf.class), + any(), anyString(), anyLong(), anyInt())).thenAnswer(invocation -> { + serializedPermits.add(invocation.getArgument(9)); + return mock(ByteBufPair.class); + }); + + Entry first = mock(Entry.class); + Entry second = mock(Entry.class); + ByteBuf firstPayload = Unpooled.buffer(0); + ByteBuf secondPayload = Unpooled.buffer(0); + when(first.getLedgerId()).thenReturn(1L); + when(first.getEntryId()).thenReturn(2L); + when(first.getDataBuffer()).thenReturn(firstPayload); + when(second.getLedgerId()).thenReturn(3L); + when(second.getEntryId()).thenReturn(4L); + when(second.getDataBuffer()).thenReturn(secondPayload); + List entries = new ArrayList<>(); + entries.add(first); + entries.add(null); + entries.add(second); + EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); + batchSizes.setBatchSize(0, 10); + batchSizes.setBatchSize(2, 4); + SendMessageResult sendResult = new SendMessageResult(entries.size()); + sendResult.recordMessagePermits(0, 3); + sendResult.recordMessagePermits(2, 4); + + try { + PulsarCommandSenderImpl sender = new PulsarCommandSenderImpl(null, cnx, null); + sender.sendMessagesToConsumer(7, "topic", mock(Subscription.class), -1, entries, batchSizes, null, + mock(RedeliveryTracker.class), 11, sendResult); + + assertEquals(serializedPermits, List.of(3, 4)); + verify(first).release(); + verify(second).release(); + } finally { + while (firstPayload.refCnt() > 0) { + firstPayload.release(); + } + while (secondPayload.refCnt() > 0) { + secondPayload.release(); + } + } + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/SendMessageResultTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/SendMessageResultTest.java new file mode 100644 index 0000000000000..023b9e8546a8e --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/SendMessageResultTest.java @@ -0,0 +1,42 @@ +/* + * 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 static org.testng.Assert.assertEquals; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class SendMessageResultTest { + + @Test + public void testFinalizedPerEntryPermitsAndCheckedSum() { + SendMessageResult result = new SendMessageResult(3); + result.recordMessagePermits(0, 2); + result.recordMessagePermits(2, 3); + + assertEquals(result.getMessagePermits(0), 2); + assertEquals(result.getMessagePermits(1), 0); + assertEquals(result.getMessagePermits(2), 3); + assertEquals(result.getTotalMessagePermits(), 5); + + SendMessageResult overflowing = new SendMessageResult(2); + overflowing.recordMessagePermits(0, Integer.MAX_VALUE); + Assert.assertThrows(ArithmeticException.class, () -> overflowing.recordMessagePermits(1, 1)); + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java new file mode 100644 index 0000000000000..4c0afcaa09f75 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java @@ -0,0 +1,135 @@ +/* + * 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.persistent; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import org.apache.pulsar.broker.service.Consumer; +import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +import org.apache.pulsar.client.api.Message; +import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.api.SubscriptionType; +import org.apache.pulsar.common.util.FutureUtil; +import org.awaitility.Awaitility; +import org.testng.Assert; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +@Test(groups = "broker-impl") +public class PersistentSharedPermitAccountingTest extends SharedPulsarBaseTest { + + private static final int BATCH_SIZE = 10; + + @DataProvider(name = "sharedDispatcherImplementations") + public Object[][] sharedDispatcherImplementations() { + return new Object[][] {{false}, {true}}; + } + + @Test(timeOut = 30000, dataProvider = "sharedDispatcherImplementations") + public void testBatchDebtIsSharedByConsumerAndDispatcherAndRemovedWithConsumer(boolean classicDispatcher) + throws Exception { + boolean originalSetting = getConfig().isSubscriptionSharedUseClassicPersistentImplementation(); + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(classicDispatcher); + String topicName = newTopicName(); + String subscriptionName = "shared-subscription"; + + try (org.apache.pulsar.client.api.Consumer clientConsumer = + pulsarClient.newConsumer(Schema.INT32) + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .acknowledgmentGroupTime(0, TimeUnit.MILLISECONDS) + .subscribe(); + Producer producer = pulsarClient.newProducer(Schema.INT32) + .topic(topicName) + .enableBatching(true) + .batchingMaxMessages(BATCH_SIZE) + .batchingMaxPublishDelay(1, TimeUnit.HOURS) + .create()) { + PersistentSubscription subscription = getSubscription(topicName, subscriptionName); + + sendBatch(producer, 0); + awaitPermitBalances(subscription, 1 - BATCH_SIZE, 1 - BATCH_SIZE, BATCH_SIZE); + + for (int i = 0; i < BATCH_SIZE; i++) { + Message message = clientConsumer.receive(10, TimeUnit.SECONDS); + Assert.assertNotNull(message); + try { + clientConsumer.acknowledge(message); + } finally { + message.release(); + } + } + awaitPermitBalances(subscription, 1, 1, 0); + + sendBatch(producer, BATCH_SIZE); + awaitPermitBalances(subscription, 1 - BATCH_SIZE, 1 - BATCH_SIZE, BATCH_SIZE); + + clientConsumer.close(); + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + Assert.assertTrue(subscription.getConsumers().isEmpty()); + Assert.assertEquals(getTotalAvailablePermits(subscription), 0); + }); + } finally { + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(originalSetting); + } + } + + private PersistentSubscription getSubscription(String topicName, String subscriptionName) throws Exception { + PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get(10, TimeUnit.SECONDS) + .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); + return topic.getSubscription(subscriptionName); + } + + private static void sendBatch(Producer producer, int firstValue) throws Exception { + List> sends = new ArrayList<>(BATCH_SIZE); + for (int i = 0; i < BATCH_SIZE; i++) { + sends.add(producer.sendAsync(firstValue + i)); + } + producer.flush(); + FutureUtil.waitForAll(sends).get(10, TimeUnit.SECONDS); + } + + private static void awaitPermitBalances(PersistentSubscription subscription, int consumerPermits, + int dispatcherPermits, int unackedMessages) { + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + Assert.assertEquals(subscription.getConsumers().size(), 1); + Consumer brokerConsumer = subscription.getConsumers().get(0); + Assert.assertEquals(brokerConsumer.getAvailablePermits(), consumerPermits); + Assert.assertEquals(brokerConsumer.getUnackedMessages(), unackedMessages); + Assert.assertEquals(getTotalAvailablePermits(subscription), dispatcherPermits); + }); + } + + private static int getTotalAvailablePermits(PersistentSubscription subscription) { + if (subscription.getDispatcher() instanceof PersistentDispatcherMultipleConsumers dispatcher) { + return dispatcher.totalAvailablePermits; + } + if (subscription.getDispatcher() instanceof PersistentDispatcherMultipleConsumersClassic dispatcher) { + return dispatcher.totalAvailablePermits; + } + throw new AssertionError("Unexpected dispatcher " + subscription.getDispatcher()); + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/BatchMessageIndexAckTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/BatchMessageIndexAckTest.java index 34be98f35646d..0ca1c14a80237 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/BatchMessageIndexAckTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/BatchMessageIndexAckTest.java @@ -20,10 +20,14 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doReturn; +import java.time.Duration; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.NavigableMap; +import java.util.Set; +import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -32,6 +36,8 @@ import org.apache.bookkeeper.client.api.DigestType; import org.apache.bookkeeper.client.api.LedgerMetadata; import org.apache.bookkeeper.net.BookieId; +import org.apache.pulsar.broker.service.persistent.PersistentSubscription; +import org.apache.pulsar.broker.service.persistent.PersistentTopic; import org.apache.pulsar.client.api.Consumer; import org.apache.pulsar.client.api.Message; import org.apache.pulsar.client.api.MessageId; @@ -42,6 +48,7 @@ import org.apache.pulsar.client.api.SubscriptionType; import org.apache.pulsar.common.policies.data.PersistentTopicInternalStats; import org.apache.pulsar.common.util.FutureUtil; +import org.awaitility.Awaitility; import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; @@ -161,6 +168,122 @@ public Object[][] ackReceiptEnabled() { return new Object[][] { { true }, { false } }; } + @DataProvider(name = "sharedDispatcherAndRemainingMessages") + public Object[][] sharedDispatcherAndRemainingMessages() { + return new Object[][] { + {false, 3}, + {false, 1}, + {true, 3}, + {true, 1} + }; + } + + @Test(timeOut = 30000, dataProvider = "sharedDispatcherAndRemainingMessages") + public void testPartialBatchRedeliveryKeepsSharedPermitBalance(boolean classicDispatcher, + int remainingMessages) throws Exception { + final int batchSize = 10; + final String topic = "persistent://my-property/my-ns/partial-redelivery-" + UUID.randomUUID(); + final String subscriptionName = "shared-sub"; + conf.setSubscriptionSharedUseClassicPersistentImplementation(classicDispatcher); + + @Cleanup + Producer producer = pulsarClient.newProducer(Schema.INT32) + .topic(topic) + .enableBatching(true) + .batchingMaxMessages(batchSize) + .batchingMaxPublishDelay(1, TimeUnit.HOURS) + .create(); + @Cleanup + Consumer firstConsumer = pulsarClient.newConsumer(Schema.INT32) + .topic(topic) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .enableBatchIndexAcknowledgment(true) + .isAckReceiptEnabled(true) + .acknowledgmentGroupTime(0, TimeUnit.MILLISECONDS) + .subscribe(); + + List> sendFutures = new ArrayList<>(batchSize); + for (int i = 0; i < batchSize; i++) { + sendFutures.add(producer.sendAsync(i)); + } + producer.flush(); + FutureUtil.waitForAll(sendFutures).get(10, TimeUnit.SECONDS); + + List> firstDelivery = new ArrayList<>(batchSize); + for (int i = 0; i < batchSize; i++) { + Message message = firstConsumer.receive(10, TimeUnit.SECONDS); + Assert.assertNotNull(message); + firstDelivery.add(message); + } + int acknowledgedMessages = batchSize - remainingMessages; + for (int i = 0; i < acknowledgedMessages; i++) { + firstConsumer.acknowledge(firstDelivery.get(i)); + } + firstDelivery.forEach(Message::release); + firstConsumer.close(); + + @Cleanup + Consumer replacementConsumer = pulsarClient.newConsumer(Schema.INT32) + .topic(topic) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .enableBatchIndexAcknowledgment(true) + .isAckReceiptEnabled(true) + .acknowledgmentGroupTime(0, TimeUnit.MILLISECONDS) + .subscribe(); + + PersistentTopic topicRef = (PersistentTopic) pulsar.getBrokerService().getTopicReference(topic) + .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); + PersistentSubscription subscription = topicRef.getSubscription(subscriptionName); + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + Assert.assertEquals(subscription.getConsumers().size(), 1); + org.apache.pulsar.broker.service.Consumer brokerConsumer = subscription.getConsumers().get(0); + // A batch is an indivisible entry. Its exact permit debit can therefore exceed the one permit that + // admitted it, and the resulting negative balance must remain as debt until the messages are consumed. + Assert.assertEquals(brokerConsumer.getAvailablePermits(), 1 - remainingMessages); + Assert.assertEquals(brokerConsumer.getUnackedMessages(), remainingMessages); + }); + + Set expectedValues = new HashSet<>(); + for (int i = acknowledgedMessages; i < batchSize; i++) { + expectedValues.add(i); + } + for (int i = 0; i < remainingMessages; i++) { + Message message = replacementConsumer.receive(10, TimeUnit.SECONDS); + Assert.assertNotNull(message); + try { + Assert.assertTrue(expectedValues.remove(message.getValue())); + replacementConsumer.acknowledge(message); + } finally { + message.release(); + } + } + Assert.assertTrue(expectedValues.isEmpty()); + Assert.assertNull(replacementConsumer.receive(1, TimeUnit.SECONDS)); + + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + Assert.assertEquals(subscription.getConsumers().size(), 1); + org.apache.pulsar.broker.service.Consumer brokerConsumer = subscription.getConsumers().get(0); + Assert.assertEquals(brokerConsumer.getAvailablePermits(), 1); + Assert.assertEquals(brokerConsumer.getUnackedMessages(), 0); + }); + + CompletableFuture nextMessage = producer.sendAsync(batchSize); + producer.flush(); + nextMessage.get(10, TimeUnit.SECONDS); + Message message = replacementConsumer.receive(10, TimeUnit.SECONDS); + Assert.assertNotNull(message); + try { + Assert.assertEquals(message.getValue(), Integer.valueOf(batchSize)); + replacementConsumer.acknowledge(message); + } finally { + message.release(); + } + } + @Test(dataProvider = "ackReceiptEnabled") public void testBatchMessageIndexAckForSharedSubscription(boolean ackReceiptEnabled) throws Exception { final String topic = "testBatchMessageIndexAckForSharedSubscription"; diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CompactedOutBatchMessageTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CompactedOutBatchMessageTest.java index 56bf88c64299f..c2f9d4fc4b994 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CompactedOutBatchMessageTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CompactedOutBatchMessageTest.java @@ -19,9 +19,12 @@ package org.apache.pulsar.client.impl; import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH; +import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; +import java.util.concurrent.TimeUnit; import org.apache.pulsar.broker.service.SharedPulsarBaseTest; import org.apache.pulsar.client.api.Message; import org.apache.pulsar.common.api.proto.BrokerEntryMetadata; @@ -34,6 +37,77 @@ @Test(groups = "broker-impl") public class CompactedOutBatchMessageTest extends SharedPulsarBaseTest { + @Test + public void testDoesNotReturnBatchPermitsToDifferentConnection() throws Exception { + final int batchSize = 2; + final String topic = newTopicName(); + MessageMetadata metadata = new MessageMetadata() + .setProducerName("foobar") + .setSequenceId(1) + .setPublishTime(1) + .setNumMessagesInBatch(batchSize); + ByteBuf compactedBatch = Unpooled.buffer(1000); + for (int i = 0; i < batchSize; i++) { + Commands.serializeSingleMessageInBatchWithPayload( + new SingleMessageMetadata().setCompactedOut(true), Unpooled.EMPTY_BUFFER, compactedBatch); + } + + try (ConsumerImpl consumer = + (ConsumerImpl) pulsarClient.newConsumer().topic(topic) + .subscriptionName("old-connection-subscription") + .receiverQueueSize(20) + .subscribe()) { + consumer.pause(); + consumer.getPermitState().availablePermits.set(20); + int permitsBefore = consumer.getAvailablePermits(); + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, compactedBatch, + new MessageIdData().setLedgerId(1234).setEntryId(567), mock(ClientCnx.class), + DEFAULT_CONSUMER_EPOCH, false, batchSize); + + assertEquals(consumer.getAvailablePermits(), permitsBefore); + } finally { + compactedBatch.release(); + } + } + + @Test + public void testReturnOnlyRemainingPermitsAfterPartialBatchDeserialization() throws Exception { + final int batchSize = 5; + final int parsedMessages = 2; + final String topic = newTopicName(); + MessageMetadata metadata = new MessageMetadata() + .setProducerName("foobar") + .setSequenceId(1) + .setPublishTime(1) + .setNumMessagesInBatch(batchSize); + ByteBuf truncatedBatch = Unpooled.buffer(1000); + for (int i = 0; i < parsedMessages; i++) { + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {(byte) i}); + Commands.serializeSingleMessageInBatchWithPayload( + new SingleMessageMetadata(), payload, truncatedBatch); + payload.release(); + } + + try (ConsumerImpl consumer = + (ConsumerImpl) pulsarClient.newConsumer().topic(topic) + .subscriptionName("partial-batch-subscription") + .receiverQueueSize(20) + .subscribe()) { + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, truncatedBatch, + new MessageIdData().setLedgerId(1234).setEntryId(567), consumer.cnx(), + DEFAULT_CONSUMER_EPOCH, false, batchSize); + + for (int i = 0; i < parsedMessages; i++) { + Message message = consumer.receive(5, TimeUnit.SECONDS); + assertNotNull(message); + message.release(); + } + assertEquals(consumer.getAvailablePermits(), batchSize); + } finally { + truncatedBatch.release(); + } + } + @Test public void testCompactedOutMessages() throws Exception { final String topic1 = newTopicName(); @@ -63,7 +137,7 @@ public void testCompactedOutMessages() throws Exception { try (ConsumerImpl consumer = (ConsumerImpl) pulsarClient.newConsumer().topic(topic1) - .subscriptionName("my-subscriber-name").subscribe()) { + .subscriptionName("my-subscriber-name").receiverQueueSize(20).subscribe()) { // shove it in the sideways consumer.receiveIndividualMessagesFromBatch(brokerEntryMetadata, metadata, 0, null, batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), @@ -75,6 +149,9 @@ batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), assertEquals(m.getKey(), "key3"); assertEquals(consumer.numMessagesInQueue(), 0); + assertEquals(consumer.getAvailablePermits(), metadata.getNumMessagesInBatch()); + } finally { + batchBuffer.release(); } } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CorruptedBatchMessagePermitTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CorruptedBatchMessagePermitTest.java new file mode 100644 index 0000000000000..d451abd710e4a --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CorruptedBatchMessagePermitTest.java @@ -0,0 +1,99 @@ +/* + * 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.client.impl; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import lombok.Cleanup; +import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +import org.apache.pulsar.broker.service.persistent.PersistentTopic; +import org.apache.pulsar.client.api.Consumer; +import org.apache.pulsar.client.api.Message; +import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.api.SubscriptionType; +import org.apache.pulsar.common.api.proto.CompressionType; +import org.apache.pulsar.common.api.proto.MessageMetadata; +import org.apache.pulsar.common.protocol.Commands; +import org.testng.annotations.Test; + +@Test(groups = "broker-impl") +public class CorruptedBatchMessagePermitTest extends SharedPulsarBaseTest { + + @Test(timeOut = 30000) + public void testSharedConsumerReturnsPermitsForCorruptedBatch() throws Exception { + final int batchSize = 10; + final String topic = newTopicName(); + admin.namespaces().setDeduplicationStatus(getNamespace(), false); + + @Cleanup + Consumer consumer = pulsarClient.newConsumer(Schema.STRING) + .topic(topic) + .subscriptionName("shared-subscription") + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(batchSize) + .subscribe(); + @Cleanup + Producer producer = pulsarClient.newProducer(Schema.STRING) + .topic(topic) + .enableBatching(false) + .create(); + + PersistentTopic persistentTopic = (PersistentTopic) getTopicIfExists(topic).get() + .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); + publishCorruptedCompressedBatch(persistentTopic, batchSize).get(10, TimeUnit.SECONDS); + producer.send("message-after-corrupted-batch"); + + Message message = consumer.receive(10, TimeUnit.SECONDS); + assertNotNull(message); + try { + assertEquals(message.getValue(), "message-after-corrupted-batch"); + } finally { + message.release(); + } + } + + private static CompletableFuture publishCorruptedCompressedBatch(PersistentTopic topic, int batchSize) { + MessageMetadata metadata = new MessageMetadata() + .setProducerName("corrupted-batch-producer") + .setSequenceId(0) + .setPublishTime(System.currentTimeMillis()) + .setNumMessagesInBatch(batchSize) + .setCompression(CompressionType.LZ4) + .setUncompressedSize(1024); + ByteBuf invalidCompressedPayload = Unpooled.wrappedBuffer(new byte[] {1, 2, 3, 4}); + ByteBuf entry = Commands.serializeMetadataAndPayload( + Commands.ChecksumType.Crc32c, metadata, invalidCompressedPayload); + invalidCompressedPayload.release(); + + CompletableFuture result = new CompletableFuture<>(); + topic.publishMessage(entry, (error, ledgerId, entryId) -> { + if (error == null) { + result.complete(null); + } else { + result.completeExceptionally(error); + } + }); + return result; + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/processor/MessagePayloadProcessorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/processor/MessagePayloadProcessorTest.java index 37f557fe0d11f..2669299029d06 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/processor/MessagePayloadProcessorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/processor/MessagePayloadProcessorTest.java @@ -36,6 +36,8 @@ import org.apache.pulsar.client.api.Schema; import org.apache.pulsar.client.api.SubscriptionInitialPosition; import org.apache.pulsar.client.api.TopicMetadata; +import org.apache.pulsar.client.impl.ConsumerImpl; +import org.awaitility.Awaitility; import org.testng.Assert; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -176,6 +178,7 @@ public void testCustomProcessor(final int numMessages, final int batchingMaxMess .topic(topic) .subscriptionName("sub") .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest) + .receiverQueueSize(100) .messagePayloadProcessor(new CustomBatchPayloadProcessor()) .subscribe(); @@ -196,5 +199,8 @@ public void testCustomProcessor(final int numMessages, final int batchingMaxMess Assert.assertEquals(message.getValue(), messagePrefix + i); consumer.acknowledge(message.getMessageId()); } + + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> + Assert.assertEquals(((ConsumerImpl) consumer).getAvailablePermits(), numMessages)); } } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerBase.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerBase.java index ddf0d1a219da0..e6b93a601427c 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerBase.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerBase.java @@ -1001,6 +1001,7 @@ protected boolean enqueueMessageAndCheckBatchReceive(Message message) { updateAutoScaleReceiverQueueHint(); } else { INCOMING_MESSAGES_SIZE_UPDATER.addAndGet(this, -messageSize); + messageDiscarded(message); } } } finally { @@ -1009,6 +1010,10 @@ protected boolean enqueueMessageAndCheckBatchReceive(Message message) { return hasEnoughMessagesForBatchReceive(); } + protected void messageDiscarded(Message message) { + message.release(); + } + protected abstract void updateAutoScaleReceiverQueueHint(); protected boolean hasEnoughMessagesForBatchReceive() { @@ -1244,17 +1249,20 @@ private void executeKeySharedMessageListener(Message message, Runnable runnab protected void callMessageListener(Message msg) { try { + ConsumerImpl receivedConsumer = (msg instanceof TopicMessageImpl) + ? ((TopicMessageImpl) msg).receivedByconsumer : (ConsumerImpl) this; + MessageImpl innerMessage = (MessageImpl) (msg instanceof TopicMessageImpl + ? ((TopicMessageImpl) msg).getMessage() : msg); State state = getState(); if (state == State.Closing || state == State.Closed) { log.debug().attr("messageId", msg.getMessageId()) .log("Consumer has been closed. Skipping message."); + receivedConsumer.increaseAvailablePermits(innerMessage); msg.release(); return; } log.debug().attr("messageId", msg.getMessageId()) .log("Calling message listener for message"); - ConsumerImpl receivedConsumer = (msg instanceof TopicMessageImpl) - ? ((TopicMessageImpl) msg).receivedByconsumer : (ConsumerImpl) this; // check the internal consumer state if (receivedConsumer != this) { @@ -1262,6 +1270,7 @@ protected void callMessageListener(Message msg) { if (receivedByConsumerState == State.Closing || receivedByConsumerState == State.Closed) { log.debug().attr("messageId", msg.getMessageId()) .log("Consumer that received the message has been closed. Skipping message."); + receivedConsumer.increaseAvailablePermits(innerMessage); msg.release(); return; } @@ -1272,11 +1281,10 @@ protected void callMessageListener(Message msg) { receivedConsumer.increaseAvailablePermits((MessageImpl) (msg instanceof TopicMessageImpl ? ((TopicMessageImpl) msg).getMessage() : msg)); - MessageImpl innerMessage = (MessageImpl) (msg instanceof TopicMessageImpl - ? ((TopicMessageImpl) msg).getMessage() : msg); if (!receivedConsumer.isValidConsumerEpoch(innerMessage)) { log.debug().attr("messageId", msg.getMessageId()) .log("Skipping processing message since the consumer epoch is not valid."); + msg.release(); return; } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java index fa25343d244a7..161338c7a07ad 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java @@ -35,6 +35,7 @@ import io.netty.util.Recycler.Handle; import io.netty.util.ReferenceCountUtil; import io.netty.util.Timeout; +import io.netty.util.concurrent.EventExecutor; import io.netty.util.concurrent.FastThreadLocal; import io.opentelemetry.api.common.Attributes; import java.io.IOException; @@ -63,7 +64,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicLongFieldUpdater; import java.util.concurrent.atomic.AtomicReference; @@ -149,13 +149,24 @@ public class ConsumerImpl extends ConsumerBase implements ConnectionHandle final long consumerId; - // Number of messages that have delivered to the application. Every once in a while, this number will be sent to the - // broker to notify that we are ready to get (and store in the incoming messages queue) more messages - @SuppressWarnings("rawtypes") - private static final AtomicIntegerFieldUpdater AVAILABLE_PERMITS_UPDATER = AtomicIntegerFieldUpdater - .newUpdater(ConsumerImpl.class, "availablePermits"); - @SuppressWarnings("unused") - private volatile int availablePermits = 0; + /** + * Permit accounting for one broker-side consumer incarnation. The identity of this object, rather than the + * {@link ClientCnx}, is the incarnation token because a broker can recreate a consumer on the same connection. + */ + static final class ConsumerPermitState { + final ClientCnx cnx; + final AtomicInteger availablePermits = new AtomicInteger(); + final AtomicBoolean flowTaskScheduled = new AtomicBoolean(); + volatile boolean active = true; + volatile boolean flowEnabled; + + ConsumerPermitState(ClientCnx cnx) { + this.cnx = cnx; + } + } + + private final AtomicReference permitState = new AtomicReference<>(); + private final Object permitStateLock = new Object(); protected volatile MessageId lastDequeuedMessageId = MessageId.earliest; private volatile MessageId lastMessageIdInBroker = MessageId.earliest; @@ -331,7 +342,6 @@ protected ConsumerImpl(PulsarClientImpl client, String topic, ConsumerConfigurat } this.initialStartMessageId = this.startMessageId; this.startMessageRollbackDurationInSec = startMessageRollbackDurationInSec; - AVAILABLE_PERMITS_UPDATER.set(this, 0); this.lookupDeadline = System.currentTimeMillis() + client.getConfiguration().getLookupTimeoutMs(); this.partitionIndex = partitionIndex; this.hasParentConsumer = hasParentConsumer; @@ -896,13 +906,15 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { + client.getConfiguration().getOperationTimeoutMs()); int currentSize; + final ConsumerPermitState newPermitState; synchronized (this) { currentSize = incomingMessages.size(); + newPermitState = replacePermitState(cnx); setClientCnx(cnx); - clearReceiverQueue(true); if (possibleSendToDeadLetterTopicMessages != null) { - possibleSendToDeadLetterTopicMessages.clear(); + clearPossibleDeadLetterMessages(); } + clearReceiverQueue(true); } boolean isDurable = subscriptionMode == SubscriptionMode.Durable; @@ -947,8 +959,12 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { cnx.sendRequestWithId(request, requestId).thenRun(() -> { synchronized (ConsumerImpl.this) { - if (changeToReadyState()) { + if (permitState.get() != newPermitState) { + future.complete(null); + return; + } else if (changeToReadyState()) { consumerIsReconnectedToBroker(cnx, currentSize); + enablePermitState(newPermitState); } else { // Consumer was closed while reconnecting, close the connection to make sure the broker // drops the consumer on its side @@ -967,10 +983,15 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { // if the consumer is not partitioned or is re-connected and is partitioned, we send the flow // command to receive messages. if (!(firstTimeConnect && hasParentConsumer) && getCurrentReceiverQueueSize() != 0) { - increaseAvailablePermits(cnx, getCurrentReceiverQueueSize()); + increaseAvailablePermits(newPermitState, getCurrentReceiverQueueSize()); } future.complete(null); }).exceptionally((e) -> { + if (permitState.get() != newPermitState) { + future.complete(null); + return null; + } + deactivatePermitState(newPermitState); deregisterFromClientCnx(); if (getState() == State.Closing || getState() == State.Closed) { // Consumer was closed while reconnecting, close the connection to make sure the broker @@ -1062,8 +1083,6 @@ protected void closeWhenReceivedUnrecoverableError(Throwable t, ClientCnx cnx) { protected void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize) { log.info("Subscribed to topic"); - - AVAILABLE_PERMITS_UPDATER.set(this, 0); } /** @@ -1114,14 +1133,87 @@ private void clearReceiverQueue(boolean updateStartMessageId) { } // else: No message was received or dequeued by this consumer. Next message would still be the startMessageId } - /** - * send the flow command to have the broker start pushing messages. - */ - private void sendFlowPermitsToBroker(ClientCnx cnx, int numMessages) { - if (cnx != null && numMessages > 0) { - log.debug().attr("addingPermits", numMessages).log("Adding additional permits"); - cnx.ctx().writeAndFlush(Commands.newFlow(consumerId, numMessages), cnx.ctx().voidPromise()); + private ConsumerPermitState replacePermitState(ClientCnx cnx) { + ConsumerPermitState newState = new ConsumerPermitState(cnx); + synchronized (permitStateLock) { + ConsumerPermitState previousState = permitState.getAndSet(newState); + if (previousState != null) { + previousState.active = false; + previousState.flowEnabled = false; + previousState.availablePermits.set(0); + } + } + return newState; + } + + private void enablePermitState(ConsumerPermitState state) { + synchronized (permitStateLock) { + if (permitState.get() == state && state.active) { + state.flowEnabled = true; + } + } + scheduleFlowIfNeeded(state); + } + + protected final void deactivatePermitState(ConsumerPermitState state) { + if (state == null) { + return; } + synchronized (permitStateLock) { + state.active = false; + state.flowEnabled = false; + state.availablePermits.set(0); + permitState.compareAndSet(state, null); + } + } + + private ConsumerPermitState getPermitState(ClientCnx cnx) { + ConsumerPermitState state = permitState.get(); + return state != null && state.cnx == cnx && state.active ? state : null; + } + + @VisibleForTesting + ConsumerPermitState getPermitState() { + return permitState.get(); + } + + private void scheduleFlowIfNeeded(ConsumerPermitState state) { + if (state == null || !state.active || !state.flowEnabled || paused || state.availablePermits.get() <= 0 + || state.availablePermits.get() < getCurrentReceiverQueueSize() / 2 + || !state.flowTaskScheduled.compareAndSet(false, true)) { + return; + } + EventExecutor eventExecutor = state.cnx.ctx().executor(); + if (eventExecutor == null) { + // Some test connections use a minimal mocked context. A real Netty context always has an executor. + state.flowTaskScheduled.set(false); + return; + } + eventExecutor.execute(() -> flushPermits(state)); + } + + private void flushPermits(ConsumerPermitState state) { + int permits = 0; + synchronized (permitStateLock) { + if (permitState.get() == state && state.active && state.flowEnabled && !paused + && state.availablePermits.get() >= getCurrentReceiverQueueSize() / 2) { + permits = state.availablePermits.getAndSet(0); + if (permits > 0) { + log.debug().attr("available", permits) + .log("Sending permit-cmd to broker with available permits ="); + state.cnx.ctx().writeAndFlush(Commands.newFlow(consumerId, permits), + state.cnx.ctx().voidPromise()); + } + } + state.flowTaskScheduled.set(false); + } + scheduleFlowIfNeeded(state); + } + + private void closePermitSourceOnOverflow(ConsumerPermitState state) { + deactivatePermitState(state); + log.error("Closing consumer connection because the returned permit count overflowed"); + state.cnx.channel().close(); } @Override @@ -1250,7 +1342,7 @@ private void cleanupAtClose(CompletableFuture closeFuture, Throwable excep private void closeConsumerTasks() { unAckedMessageTracker.close(); if (possibleSendToDeadLetterTopicMessages != null) { - possibleSendToDeadLetterTopicMessages.clear(); + clearPossibleDeadLetterMessages(); } acknowledgmentsGroupingTracker.close(); if (batchReceiveTimeout != null) { @@ -1298,6 +1390,26 @@ protected MessageImpl newSingleMessage(final int index, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { + return newSingleMessage(index, numMessages, brokerEntryMetadata, msgMetadata, singleMessageMetadata, + payload, messageId, schema, containMetadata, ackBitSet, ackSetInMessageId, redeliveryCount, + consumerEpoch, isEncrypted, cnx()); + } + + protected MessageImpl newSingleMessage(final int index, + final int numMessages, + final BrokerEntryMetadata brokerEntryMetadata, + final MessageMetadata msgMetadata, + final SingleMessageMetadata singleMessageMetadata, + final ByteBuf payload, + final MessageIdImpl messageId, + final Schema schema, + final boolean containMetadata, + final BitSetRecyclable ackBitSet, + final BitSet ackSetInMessageId, + final int redeliveryCount, + final long consumerEpoch, + final boolean isEncrypted, + final ClientCnx messageCnx) { log.debug().attr("index", index) .log("processing message num - in batch"); @@ -1332,11 +1444,11 @@ protected MessageImpl newSingleMessage(final int index, final ByteBuf payloadBuffer = (singleMessagePayload != null) ? singleMessagePayload : payload; final MessageImpl message = MessageImpl.create(topicName.toString(), batchMessageIdImpl, msgMetadata, singleMessageMetadata, payloadBuffer, - createEncryptionContext(msgMetadata, isEncrypted), cnx(), schema, redeliveryCount, + createEncryptionContext(msgMetadata, isEncrypted), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); return message; - } catch (IOException | IllegalStateException e) { + } catch (IOException | RuntimeException e) { throw new IllegalStateException(e); } finally { if (singleMessagePayload != null) { @@ -1364,8 +1476,21 @@ protected MessageImpl newMessage(final MessageIdImpl messageId, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { + return newMessage(messageId, brokerEntryMetadata, messageMetadata, payload, schema, redeliveryCount, + consumerEpoch, isEncrypted, cnx()); + } + + protected MessageImpl newMessage(final MessageIdImpl messageId, + final BrokerEntryMetadata brokerEntryMetadata, + final MessageMetadata messageMetadata, + final ByteBuf payload, + final Schema schema, + final int redeliveryCount, + final long consumerEpoch, + final boolean isEncrypted, + final ClientCnx messageCnx) { final MessageImpl message = MessageImpl.create(topicName.toString(), messageId, messageMetadata, payload, - createEncryptionContext(messageMetadata, isEncrypted), cnx(), schema, redeliveryCount, + createEncryptionContext(messageMetadata, isEncrypted), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); return message; @@ -1378,18 +1503,69 @@ private void executeNotifyCallback(final MessageImpl message) { // Enqueue the message so that it can be retrieved when application calls receive() // if the conf.getReceiverQueueSize() is 0 then discard message if no one is waiting for it. // if asyncReceive is waiting then notify callback without adding to incomingMessages queue - internalPinnedExecutor.execute(() -> { - if (!isValidConsumerEpoch(message)) { - increaseAvailablePermits(cnx()); - return; - } - Message interceptMsg = onArrival(message); - if (hasNextPendingReceive()) { - notifyPendingReceivedCallback(interceptMsg, null); - } else if (enqueueMessageAndCheckBatchReceive(interceptMsg) && hasPendingBatchReceive()) { - notifyPendingBatchReceivedCallBack(); - } - }); + try { + internalPinnedExecutor.execute(() -> { + if (message.getPermitState() != permitState.get() || !isValidConsumerEpoch(message)) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + increaseAvailablePermits(message); + message.release(); + return; + } + final Message interceptMsg; + try { + interceptMsg = onArrival(message); + } catch (Throwable t) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + increaseAvailablePermits(message); + message.release(); + log.error().exception(t).log("Message arrival interceptor failed"); + return; + } + if (hasNextPendingReceive()) { + notifyPendingReceivedCallback(interceptMsg, null); + } else if (enqueueMessageAndCheckBatchReceive(interceptMsg) && hasPendingBatchReceive()) { + notifyPendingBatchReceivedCallBack(); + } + }); + } catch (Throwable t) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + throw t; + } + } + + private void executeNotifyCallback(final MessageImpl message, MessagePermitBudget permitBudget) { + try { + int permitCost = permitBudget.claim(); + message.setFlowPermitOwnership(permitBudget.sourceState, permitCost); + executeNotifyCallback(message); + } catch (Throwable t) { + permitBudget.restore(message.takeFlowPermitCost()); + message.release(); + throw t; + } + } + + private void executeNotifyCallbackWithTransferredPermit(final MessageImpl message) { + try { + executeNotifyCallback(message); + } catch (Throwable t) { + increaseAvailablePermits(message); + message.release(); + throw t; + } + } + + private void executeNotifyCallbackForPayloadProcessor(final MessageImpl message) { + try { + executeNotifyCallback(message); + } catch (Throwable t) { + message.takeFlowPermitCost(); + message.release(); + throw t; + } } protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, @@ -1400,14 +1576,48 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe final int redeliveryCount, final List ackSet, long consumerEpoch) { + int messagePermits = messageMetadata.hasNumMessagesInBatch() + ? messageMetadata.getNumMessagesInBatch() : 1; + processPayloadByProcessor(brokerEntryMetadata, messageMetadata, byteBuf, messageId, schema, redeliveryCount, + ackSet, consumerEpoch, messagePermits, getPermitState(cnx())); + } + + protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, + final MessageMetadata messageMetadata, + final ByteBuf byteBuf, + final MessageIdImpl messageId, + final Schema schema, + final int redeliveryCount, + final List ackSet, + long consumerEpoch, + int messagePermits) { + processPayloadByProcessor(brokerEntryMetadata, messageMetadata, byteBuf, messageId, schema, redeliveryCount, + ackSet, consumerEpoch, messagePermits, getPermitState(cnx())); + } + + protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, + final MessageMetadata messageMetadata, + final ByteBuf byteBuf, + final MessageIdImpl messageId, + final Schema schema, + final int redeliveryCount, + final List ackSet, + long consumerEpoch, + int messagePermits, + ConsumerPermitState sourceState) { final MessagePayloadImpl payload = MessagePayloadImpl.create(byteBuf); final MessagePayloadContextImpl entryContext = MessagePayloadContextImpl.get( - brokerEntryMetadata, messageMetadata, messageId, this, redeliveryCount, ackSet, consumerEpoch); - final AtomicInteger skippedMessages = new AtomicInteger(0); + brokerEntryMetadata, messageMetadata, messageId, this, redeliveryCount, ackSet, consumerEpoch, + sourceState); + final AtomicInteger skippedMessages = new AtomicInteger(); try { conf.getPayloadProcessor().process(payload, entryContext, schema, message -> { if (message != null) { - executeNotifyCallback((MessageImpl) message); + MessageImpl messageImpl = (MessageImpl) message; + if (messageImpl.getPermitState() == null) { + messageImpl.setFlowPermitOwnership(sourceState, 1); + } + executeNotifyCallbackForPayloadProcessor(messageImpl); } else { skippedMessages.incrementAndGet(); } @@ -1415,20 +1625,31 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe } catch (Throwable throwable) { log.warn().exception(throwable) .log("unable to obtain message in batch"); - discardCorruptedMessage(messageId, cnx(), ValidationError.BatchDeSerializeError); + discardCorruptedMessage(messageId, sourceState, ValidationError.BatchDeSerializeError, 1); } finally { entryContext.recycle(); payload.release(); // byteBuf.release() is called in this method } if (skippedMessages.get() > 0) { - increaseAvailablePermits(cnx(), skippedMessages.get()); + increaseAvailablePermitsForReceivedMessage(sourceState, skippedMessages.get()); } tryTriggerListener(); } void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, ClientCnx cnx) { + ConsumerPermitState sourceState = getPermitState(cnx); + if (sourceState == null) { + return; + } + final int earlyMessagePermits; + try { + earlyMessagePermits = getEarlyMessagePermits(cmdMessage); + } catch (PermitProtocolException e) { + closeMalformedPermitCommand(sourceState, e.getMessage()); + return; + } long[] ackSet = EMPTY_ACK_SET; if (cmdMessage.getAckSetsCount() > 0) { ackSet = new long[cmdMessage.getAckSetsCount()]; @@ -1448,7 +1669,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien if (!verifyChecksum(headersAndPayload, messageId)) { // discard message with checksum error - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch); + discardCorruptedMessage(messageId, sourceState, ValidationError.ChecksumMismatch, earlyMessagePermits); return; } @@ -1458,11 +1679,18 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien brokerEntryMetadata = Commands.parseBrokerEntryMetadataIfExist(headersAndPayload); msgMetadata = Commands.parseMessageMetadata(headersAndPayload); } catch (Throwable t) { - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch); + discardCorruptedMessage(messageId, sourceState, ValidationError.ChecksumMismatch, earlyMessagePermits); return; } final int numMessages = msgMetadata.getNumMessagesInBatch(); + final int messagePermits; + try { + messagePermits = validateAndResolveMessagePermits(cmdMessage, numMessages, earlyMessagePermits); + } catch (PermitProtocolException e) { + closeMalformedPermitCommand(sourceState, e.getMessage()); + return; + } final int numChunks = msgMetadata.hasNumChunksFromMsg() ? msgMetadata.getNumChunksFromMsg() : 0; final boolean isChunkedMessage = numChunks > 1; MessageIdImpl msgId = new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()); @@ -1471,12 +1699,12 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien log.debug().attr("messageId", msgId) .log("Ignoring message as it was already being acked earlier by same consumer"); - increaseAvailablePermits(cnx, numMessages); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); return; } DecryptResult decryptResult = decryptPayloadIfNeeded(messageId, redeliveryCount, msgMetadata, headersAndPayload, - cnx); + sourceState); if (decryptResult.shouldDiscard()) { // Message was discarded or CryptoKeyReader isn't implemented @@ -1489,7 +1717,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien // uncompress decryptedPayload and release decryptedPayload-ByteBuf ByteBuf uncompressedPayload = (isMessageUndecryptable || isChunkedMessage) ? decryptedPayload.retain() - : uncompressPayloadIfNeeded(messageId, msgMetadata, decryptedPayload, cnx, true); + : uncompressPayloadIfNeeded(messageId, msgMetadata, decryptedPayload, sourceState, true, + messagePermits); decryptedPayload.release(); if (uncompressedPayload == null) { // Message was discarded on decompression error @@ -1499,7 +1728,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien if (conf.getPayloadProcessor() != null) { // uncompressedPayload is released in this method so we don't need to call release() again processPayloadByProcessor(brokerEntryMetadata, msgMetadata, - uncompressedPayload, msgId, schema, redeliveryCount, longArrayToList(ackSet), consumerEpoch); + uncompressedPayload, msgId, schema, redeliveryCount, longArrayToList(ackSet), consumerEpoch, + messagePermits, sourceState); return; } @@ -1508,7 +1738,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien if (isMessageUndecryptable || (numMessages == 1 && !msgMetadata.hasNumMessagesInBatch())) { if (isChunkedMessage) { - uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx); + uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, + sourceState); if (uncompressedPayload == null) { return; } @@ -1539,32 +1770,55 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien .log("Ignoring message from before the startMessageId"); uncompressedPayload.release(); + increaseAvailablePermitsForReceivedMessage(sourceState, 1); return; } + MessagePermitBudget permitBudget = isMessageUndecryptable || isChunkedMessage + ? null : new MessagePermitBudget(sourceState, messagePermits); final MessageImpl message = newMessage(msgId, brokerEntryMetadata, msgMetadata, uncompressedPayload, - schema, redeliveryCount, consumerEpoch, isMessageUndecryptable); + schema, redeliveryCount, consumerEpoch, isMessageUndecryptable, cnx); uncompressedPayload.release(); if (deadLetterPolicy != null && possibleSendToDeadLetterTopicMessages != null) { if (redeliveryCount >= deadLetterPolicy.getMaxRedeliverCount()) { - possibleSendToDeadLetterTopicMessages.put((MessageIdImpl) message.getMessageId(), + retainPossibleDeadLetterMessages((MessageIdImpl) message.getMessageId(), Collections.singletonList(message)); if (redeliveryCount > deadLetterPolicy.getMaxRedeliverCount()) { redeliverUnacknowledgedMessages(Collections.singleton(message.getMessageId())); - // The message is skipped due to reaching the max redelivery count, - // so we need to increase the available permits - increaseAvailablePermits(cnx); + if (permitBudget != null) { + permitBudget.terminate(); + } else { + increaseAvailablePermitsForReceivedMessage(sourceState, 1); + } return; } + message.retain(); } } - executeNotifyCallback(message); + if (permitBudget != null) { + boolean malformed = false; + try { + executeNotifyCallback(message, permitBudget); + } catch (PermitProtocolException e) { + malformed = true; + permitBudget.terminateWithoutCredit(); + closeMalformedPermitCommand(sourceState, e.getMessage()); + } finally { + if (!malformed) { + permitBudget.terminate(); + } + } + } else { + message.setFlowPermitOwnership(sourceState, 1); + executeNotifyCallbackWithTransferredPermit(message); + } } else { // handle batch message enqueuing; uncompressed payload has all messages in batch receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, - uncompressedPayload, messageId, cnx, consumerEpoch, isMessageUndecryptable); + uncompressedPayload, messageId, sourceState, consumerEpoch, isMessageUndecryptable, + messagePermits); uncompressedPayload.release(); } @@ -1572,10 +1826,136 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien } + @VisibleForTesting + static int getMessagePermits(CommandMessage cmdMessage, int batchSize) { + int earlyPermits = getEarlyMessagePermits(cmdMessage); + return validateAndResolveMessagePermits(cmdMessage, batchSize, earlyPermits); + } + + private static int getEarlyMessagePermits(CommandMessage cmdMessage) { + if (cmdMessage.hasMessagePermits()) { + long unsignedPermits = Integer.toUnsignedLong(cmdMessage.getMessagePermits()); + if (unsignedPermits == 0 || unsignedPermits > Integer.MAX_VALUE) { + throw new PermitProtocolException("message_permits is outside the supported unsigned range"); + } + return (int) unsignedPermits; + } + if (cmdMessage.getAckSetsCount() == 0) { + return 1; + } + int permits = getAckSetCardinality(cmdMessage); + if (permits == 0) { + throw new PermitProtocolException("ack_set must contain at least one deliverable index"); + } + return permits; + } + + private static int validateAndResolveMessagePermits(CommandMessage cmdMessage, int batchSize, + int earlyPermits) { + if (batchSize <= 0) { + throw new PermitProtocolException("num_messages_in_batch must be positive"); + } + int expectedPermits; + if (cmdMessage.getAckSetsCount() > 0) { + expectedPermits = validateAckSet(cmdMessage, batchSize); + } else { + expectedPermits = batchSize; + } + if (cmdMessage.hasMessagePermits() && earlyPermits != expectedPermits) { + throw new PermitProtocolException("message_permits does not match the native payload cardinality"); + } + return cmdMessage.hasMessagePermits() ? earlyPermits : expectedPermits; + } + + private static int validateAckSet(CommandMessage cmdMessage, int batchSize) { + int cardinality = 0; + for (int i = 0; i < cmdMessage.getAckSetsCount(); i++) { + long word = cmdMessage.getAckSetAt(i); + long firstIndex = (long) i * Long.SIZE; + if (firstIndex >= batchSize) { + if (word != 0) { + throw new PermitProtocolException("ack_set contains an index outside the native batch"); + } + } else { + int validBits = (int) Math.min(Long.SIZE, batchSize - firstIndex); + if (validBits < Long.SIZE && (word >>> validBits) != 0) { + throw new PermitProtocolException("ack_set contains an index outside the native batch"); + } + } + try { + cardinality = Math.addExact(cardinality, Long.bitCount(word)); + } catch (ArithmeticException e) { + throw new PermitProtocolException("ack_set cardinality exceeds the supported range"); + } + } + if (cardinality == 0) { + throw new PermitProtocolException("ack_set must contain at least one deliverable index"); + } + return cardinality; + } + + private static int getAckSetCardinality(CommandMessage cmdMessage) { + int cardinality = 0; + for (int i = 0; i < cmdMessage.getAckSetsCount(); i++) { + try { + cardinality = Math.addExact(cardinality, Long.bitCount(cmdMessage.getAckSetAt(i))); + } catch (ArithmeticException e) { + throw new PermitProtocolException("ack_set cardinality exceeds the supported range"); + } + } + return cardinality; + } + + private static final class PermitProtocolException extends RuntimeException { + private PermitProtocolException(String message) { + super(message); + } + } + + private final class MessagePermitBudget { + private final ConsumerPermitState sourceState; + private int remainingPermits; + private boolean terminated; + + private MessagePermitBudget(ConsumerPermitState sourceState, int messagePermits) { + this.sourceState = sourceState; + this.remainingPermits = messagePermits; + } + + private int claim() { + if (terminated || remainingPermits == 0) { + throw new PermitProtocolException("native payload expansion exceeded message_permits"); + } + remainingPermits--; + return 1; + } + + private void restore(int permitCost) { + if (!terminated) { + remainingPermits = Math.addExact(remainingPermits, permitCost); + } + } + + private void terminate() { + if (terminated) { + return; + } + terminated = true; + int permits = remainingPermits; + remainingPermits = 0; + increaseAvailablePermitsForReceivedMessage(sourceState, permits); + } + + private void terminateWithoutCredit() { + terminated = true; + remainingPermits = 0; + } + } + private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId, - MessageIdData messageId, ClientCnx cnx) { + MessageIdData messageId, ConsumerPermitState sourceState) { if (msgMetadata.getChunkId() != (msgMetadata.getNumChunksFromMsg() - 1)) { - increaseAvailablePermits(cnx); + increaseAvailablePermitsForReceivedMessage(sourceState, 1); } // Lazy task scheduling to expire incomplete chunk message if (expireTimeOfIncompleteChunkedMessageMillis > 0 && expireChunkMessageTaskScheduled.compareAndSet(false, @@ -1713,7 +2093,8 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m compressedPayload.release(); compressedPayload = chunkedMsgCtx.chunkedMsgBuffer; - ByteBuf uncompressedPayload = uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, cnx, false); + ByteBuf uncompressedPayload = + uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, sourceState, false, 1); compressedPayload.release(); return uncompressedPayload; } @@ -1785,6 +2166,25 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, boolean isEncrypted) { + int messagePermits = ackSet == null || ackSet.length == 0 + ? msgMetadata.getNumMessagesInBatch() : BitSet.valueOf(ackSet).cardinality(); + receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, + uncompressedPayload, messageId, getPermitState(cnx), consumerEpoch, isEncrypted, messagePermits); + } + + void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, + int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, + MessageIdData messageId, ClientCnx cnx, long consumerEpoch, + boolean isEncrypted, int messagePermits) { + receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, + uncompressedPayload, messageId, getPermitState(cnx), consumerEpoch, isEncrypted, messagePermits); + } + + void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, + int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, + MessageIdData messageId, ConsumerPermitState sourceState, + long consumerEpoch, + boolean isEncrypted, int messagePermits) { int batchSize = msgMetadata.getNumMessagesInBatch(); // create ack tracker for entry aka batch @@ -1803,50 +2203,64 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } SingleMessageMetadata singleMessageMetadata = new SingleMessageMetadata(); - int skippedMessages = 0; + MessagePermitBudget permitBudget = new MessagePermitBudget(sourceState, messagePermits); + boolean malformed = false; try { for (int i = 0; i < batchSize; ++i) { final MessageImpl message = newSingleMessage(i, batchSize, brokerEntryMetadata, msgMetadata, singleMessageMetadata, uncompressedPayload, batchMessage, schema, true, - ackBitSet, ackSetInMessageId, redeliveryCount, consumerEpoch, isEncrypted); + ackBitSet, ackSetInMessageId, redeliveryCount, consumerEpoch, isEncrypted, + sourceState == null ? null : sourceState.cnx); if (message == null) { // If it is not in ackBitSet, it means Broker does not want to deliver it to the client, and // did not decrease the permits in the broker-side. // So do not acquire more permits for this message. // Why not skip this single message in the first line of for-loop block? We need call // "newSingleMessage" to move "payload.readerIndex" to a correct value to get the correct data. - if (!isSingleMessageAcked(ackBitSet, i)) { - skippedMessages++; - } continue; } if (possibleToDeadLetter != null) { possibleToDeadLetter.add(message); // Skip the message which reaches the max redelivery count. if (redeliveryCount > deadLetterPolicy.getMaxRedeliverCount()) { - skippedMessages++; continue; } } if (acknowledgmentsGroupingTracker.isDuplicate(message.getMessageId())) { - skippedMessages++; + if (possibleToDeadLetter == null) { + message.release(); + } continue; } - executeNotifyCallback(message); - } - if (ackBitSet != null) { - ackBitSet.recycle(); + if (possibleToDeadLetter != null) { + message.retain(); + } + try { + executeNotifyCallback(message, permitBudget); + } catch (PermitProtocolException e) { + throw e; + } } + } catch (PermitProtocolException e) { + malformed = true; + permitBudget.terminateWithoutCredit(); + closeMalformedPermitCommand(sourceState, e.getMessage()); } catch (IllegalStateException e) { log.warn().exception(e) .log("unable to obtain message in batch"); - discardCorruptedMessage(messageId, cnx, ValidationError.BatchDeSerializeError); + discardCorruptedMessage(messageId, sourceState, ValidationError.BatchDeSerializeError, 0); + } finally { + if (ackBitSet != null) { + ackBitSet.recycle(); + } + if (!malformed) { + permitBudget.terminate(); + } } if (deadLetterPolicy != null && possibleSendToDeadLetterTopicMessages != null) { if (redeliveryCount >= deadLetterPolicy.getMaxRedeliverCount()) { - possibleSendToDeadLetterTopicMessages.put(batchMessage, - possibleToDeadLetter); + retainPossibleDeadLetterMessages(batchMessage, possibleToDeadLetter); if (redeliveryCount > deadLetterPolicy.getMaxRedeliverCount()) { redeliverUnacknowledgedMessages(Collections.singleton(batchMessage)); } @@ -1857,9 +2271,6 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, .attr("remainingCapacity", incomingMessages.remainingCapacity()) .log("enqueued messages in batch. queue size -, available queue size"); - if (skippedMessages > 0) { - increaseAvailablePermits(cnx, skippedMessages); - } } private boolean isPriorEntryIndex(long idx) { @@ -1883,8 +2294,9 @@ private boolean isSameEntry(MessageIdImpl messageId) { */ @Override protected synchronized void messageProcessed(Message msg) { - ClientCnx currentCnx = cnx(); - ClientCnx msgCnx = ((MessageImpl) msg).getCnx(); + ConsumerPermitState currentState = permitState.get(); + ConsumerPermitState messageState = msg instanceof MessageImpl messageImpl + ? messageImpl.getPermitState() : currentState; lastDequeuedMessageId = msg.getMessageId(); messagesPrefetchedGauge.decrement(); @@ -1893,11 +2305,13 @@ protected synchronized void messageProcessed(Message msg) { bytesPrefetchedGauge.subtract(msg.size()); bytesReceivedCounter.add(msg.size()); - if (msgCnx != currentCnx) { + if (messageState != currentState) { // The processed message did belong to the old queue that was cleared after reconnection. } else { if (listener == null && !parentConsumerHasListener) { - increaseAvailablePermits(currentCnx); + if (msg instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); + } } stats.updateNumMsgsReceived(msg); @@ -1930,11 +2344,21 @@ protected void trackMessage(MessageId messageId, int redeliveryCount) { } void increaseAvailablePermits(MessageImpl msg) { - ClientCnx currentCnx = cnx(); - ClientCnx msgCnx = msg.getCnx(); - if (msgCnx == currentCnx) { - increaseAvailablePermits(currentCnx); + increaseAvailablePermits(msg.getPermitState(), msg.takeFlowPermitCost()); + } + + private void increaseAvailablePermitsForReceivedMessage(ConsumerPermitState sourceState, int permits) { + increaseAvailablePermits(sourceState, permits); + } + + @Override + protected void messageDiscarded(Message message) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + if (message instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); } + message.release(); } void increaseAvailablePermits(ClientCnx currentCnx) { @@ -1942,21 +2366,37 @@ void increaseAvailablePermits(ClientCnx currentCnx) { } protected void increaseAvailablePermits(ClientCnx currentCnx, int delta) { - int available = AVAILABLE_PERMITS_UPDATER.addAndGet(this, delta); - while (available >= getCurrentReceiverQueueSize() / 2 && !paused) { - if (AVAILABLE_PERMITS_UPDATER.compareAndSet(this, available, 0)) { - log.debug().attr("available", available) - .log("Sending permit-cmd to broker with available permits ="); - sendFlowPermitsToBroker(currentCnx, available); - break; - } else { - available = AVAILABLE_PERMITS_UPDATER.get(this); + increaseAvailablePermits(getPermitState(currentCnx), delta); + } + + private void increaseAvailablePermits(ConsumerPermitState state, int delta) { + if (state == null || delta < 0 || (delta == 0 && !state.active)) { + return; + } + if (delta > 0) { + while (state.active && permitState.get() == state) { + int current = state.availablePermits.get(); + final int updated; + try { + updated = Math.addExact(current, delta); + } catch (ArithmeticException e) { + closePermitSourceOnOverflow(state); + return; + } + if (state.availablePermits.compareAndSet(current, updated)) { + if (!state.active || permitState.get() != state) { + state.availablePermits.set(0); + return; + } + break; + } } } + scheduleFlowIfNeeded(state); } public void increaseAvailablePermits(int delta) { - increaseAvailablePermits(cnx(), delta); + increaseAvailablePermits(permitState.get(), delta); } @Override @@ -1973,14 +2413,22 @@ protected void setCurrentReceiverQueueSize(int newSize) { @Override public void pause() { - paused = true; + synchronized (permitStateLock) { + paused = true; + } } @Override public void resume() { - if (paused) { - paused = false; - increaseAvailablePermits(cnx(), 0); + ConsumerPermitState state = null; + synchronized (permitStateLock) { + if (paused) { + paused = false; + state = permitState.get(); + } + } + if (state != null) { + increaseAvailablePermits(state, 0); } } @@ -2033,7 +2481,7 @@ public static DecryptResult discard() { @SuppressWarnings({"unchecked", "rawtypes"}) private DecryptResult decryptPayloadIfNeeded(MessageIdData messageId, int redeliveryCount, MessageMetadata msgMetadata, - ByteBuf payload, ClientCnx currentCnx) { + ByteBuf payload, ConsumerPermitState sourceState) { if (msgMetadata.getEncryptionKeysCount() == 0) { return DecryptResult.success(payload.retain()); @@ -2041,7 +2489,7 @@ private DecryptResult decryptPayloadIfNeeded(MessageIdData messageId, int redeli int batchSize = msgMetadata.getNumMessagesInBatch(); // If KeyReader is not configured throw exception based on config param if (conf.getCryptoKeyReader() == null) { - return handleCryptoFailure(payload, messageId, currentCnx, redeliveryCount, batchSize, true); + return handleCryptoFailure(payload, messageId, sourceState, redeliveryCount, batchSize, true); } int maxDecryptedSize = msgCrypto.getMaxOutputSize(payload.readableBytes()); @@ -2055,10 +2503,11 @@ private DecryptResult decryptPayloadIfNeeded(MessageIdData messageId, int redeli decryptedData.release(); - return handleCryptoFailure(payload, messageId, currentCnx, redeliveryCount, batchSize, false); + return handleCryptoFailure(payload, messageId, sourceState, redeliveryCount, batchSize, false); } - private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData messageId, ClientCnx currentCnx, + private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData messageId, + ConsumerPermitState sourceState, int redeliveryCount, int batchSize, boolean cryptoReaderNotExist) { if (conf.getDecryptFailListener() != null) { @@ -2087,7 +2536,7 @@ private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData message .attr("batchIndex", messageId.getBatchIndex()) .log("[--] Discarding message since decryption failed " + "and config is set to discard"); } - discardMessage(messageId, currentCnx, ValidationError.DecryptionError, batchSize); + discardMessage(messageId, sourceState, ValidationError.DecryptionError, batchSize); return DecryptResult.discard(); case FAIL: if (cryptoReaderNotExist) { @@ -2110,7 +2559,7 @@ private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData message } private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetadata msgMetadata, ByteBuf payload, - ClientCnx currentCnx, boolean checkMaxMessageSize) { + ConsumerPermitState sourceState, boolean checkMaxMessageSize, int messagePermits) { CompressionType compressionType = msgMetadata.getCompression(); CompressionCodec codec = CompressionCodecProvider.getCompressionCodec(compressionType); int uncompressedSize = msgMetadata.getUncompressedSize(); @@ -2120,19 +2569,20 @@ private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetada log.error().attr("size", payloadSize) .attr("messageId", messageId) .log("Got corrupted payload message size at"); - discardCorruptedMessage(messageId, currentCnx, ValidationError.UncompressedSizeCorruption); + discardCorruptedMessage(messageId, sourceState, ValidationError.UncompressedSizeCorruption, + messagePermits); return null; } try { ByteBuf uncompressedPayload = codec.decode(payload, uncompressedSize); return uncompressedPayload; - } catch (IOException e) { + } catch (IOException | RuntimeException e) { log.error().attr("compressionType", compressionType) .attr("messageId", messageId) .exceptionMessage(e) .exception(e) .log("Failed to decompress message with at"); - discardCorruptedMessage(messageId, currentCnx, ValidationError.DecompressionError); + discardCorruptedMessage(messageId, sourceState, ValidationError.DecompressionError, messagePermits); return null; } } @@ -2154,33 +2604,40 @@ private boolean verifyChecksum(ByteBuf headersAndPayload, MessageIdData messageI return true; } - private void discardCorruptedMessage(MessageIdImpl messageId, ClientCnx currentCnx, - ValidationError validationError) { + private void discardCorruptedMessage(MessageIdImpl messageId, ConsumerPermitState sourceState, + ValidationError validationError, int messagePermits) { log.error().attr("messageId", messageId) .log("Discarding corrupted message"); ByteBuf cmd = Commands.newAck(consumerId, messageId.getLedgerId(), messageId.getEntryId(), null, AckType.Individual, validationError, Collections.emptyMap(), -1); - currentCnx.ctx().writeAndFlush(cmd, currentCnx.ctx().voidPromise()); - increaseAvailablePermits(currentCnx); + sourceState.cnx.ctx().writeAndFlush(cmd, sourceState.cnx.ctx().voidPromise()); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); stats.incrementNumReceiveFailed(); } - private void discardCorruptedMessage(MessageIdData messageId, ClientCnx currentCnx, - ValidationError validationError) { + private void discardCorruptedMessage(MessageIdData messageId, ConsumerPermitState sourceState, + ValidationError validationError, int messagePermits) { log.error().attr("messageId", messageId) .log("Discarding corrupted message"); - discardMessage(messageId, currentCnx, validationError, 1); + discardMessage(messageId, sourceState, validationError, messagePermits); } - private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, ValidationError validationError, - int batchMessages) { + private void discardMessage(MessageIdData messageId, ConsumerPermitState sourceState, + ValidationError validationError, + int messagePermits) { ByteBuf cmd = Commands.newAck(consumerId, messageId.getLedgerId(), messageId.getEntryId(), null, AckType.Individual, validationError, Collections.emptyMap(), -1); - currentCnx.ctx().writeAndFlush(cmd, currentCnx.ctx().voidPromise()); - increaseAvailablePermits(currentCnx, batchMessages); + sourceState.cnx.ctx().writeAndFlush(cmd, sourceState.cnx.ctx().voidPromise()); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); stats.incrementNumReceiveFailed(); } + private void closeMalformedPermitCommand(ConsumerPermitState sourceState, String reason) { + log.error().attr("reason", reason).log("Closing connection after receiving a malformed message command"); + deactivatePermitState(sourceState); + sourceState.cnx.channel().close(); + } + @Override String getHandlerName() { return subscription; @@ -2201,7 +2658,8 @@ int getPartitionIndex() { @Override public int getAvailablePermits() { - return AVAILABLE_PERMITS_UPDATER.get(this); + ConsumerPermitState state = permitState.get(); + return state == null ? 0 : state.availablePermits.get(); } @Override @@ -2234,7 +2692,7 @@ public void redeliverUnacknowledgedMessages() { } // clear local message - int currentSize; + List> messagesFromQueue; incomingQueueLock.lock(); try { // we should increase epoch every time, because MultiTopicsConsumerImpl also increase it, @@ -2245,8 +2703,9 @@ public void redeliverUnacknowledgedMessages() { } // clear local message - currentSize = incomingMessages.size(); - clearIncomingMessages(); + messagesFromQueue = new ArrayList<>(incomingMessages.size()); + incomingMessages.drainTo(messagesFromQueue); + resetIncomingMessageSize(); unAckedMessageTracker.clear(); } finally { incomingQueueLock.unlock(); @@ -2256,12 +2715,16 @@ public void redeliverUnacknowledgedMessages() { if (cnx != null && isConnected(cnx)) { cnx.ctx().writeAndFlush(Commands.newRedeliverUnacknowledgedMessages( consumerId, CONSUMER_EPOCH.get(this)), cnx.ctx().voidPromise()); - if (currentSize > 0) { - increaseAvailablePermits(cnx, currentSize); - } - log.debug().attr("send", currentSize) + messagesFromQueue.forEach(message -> { + if (message instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); + } + message.release(); + }); + log.debug().attr("messages", messagesFromQueue.size()) .log("Redeliver unacked messages and send permits"); } else { + messagesFromQueue.forEach(Message::release); log.warn("Send redeliver messages command but the" + " client is reconnect or close, so don't" + " need to send redeliver command to" @@ -2284,7 +2747,7 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { } ClientCnx cnx = cnx(); if (isConnected() && cnx.getRemoteEndpointProtocolVersion() >= ProtocolVersion.v2.getValue()) { - int messagesFromQueue = removeExpiredMessagesFromQueue(messageIds); + List> messagesFromQueue = removeExpiredMessagesFromQueue(messageIds); Iterables.partition(messageIds, MAX_REDELIVER_UNACKNOWLEDGED).forEach(ids -> { getRedeliveryMessageIdData(ids).thenAccept(messageIdData -> { if (!messageIdData.isEmpty()) { @@ -2293,10 +2756,13 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { } }); }); - if (messagesFromQueue > 0) { - increaseAvailablePermits(cnx, messagesFromQueue); - } - log.debug().attr("increase", messagesFromQueue) + messagesFromQueue.forEach(message -> { + if (message instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); + } + message.release(); + }); + log.debug().attr("increase", messagesFromQueue.size()) .log("Redeliver unacked messages and increase permits"); return; } @@ -2350,79 +2816,87 @@ private CompletableFuture> getRedeliveryMessageIdData(List processPossibleToDLQ(MessageIdAdv messageId) { - List> deadLetterMessages = null; - if (possibleSendToDeadLetterTopicMessages != null) { - deadLetterMessages = possibleSendToDeadLetterTopicMessages.get(MessageIdAdvUtils.discardBatch(messageId)); - } - CompletableFuture result = new CompletableFuture<>(); - if (deadLetterMessages != null) { - List> finalDeadLetterMessages = deadLetterMessages; - initDeadLetterProducerIfNeeded().thenAcceptAsync(producerDLQ -> { - for (MessageImpl message : finalDeadLetterMessages) { - try { - String originMessageIdStr = message.getMessageId().toString(); - String originTopicNameStr = getOriginTopicNameStr(message); - AutoProduceBytesSchema deadLetterMessageSchema = + if (possibleSendToDeadLetterTopicMessages == null) { + return CompletableFuture.completedFuture(false); + } + MessageIdAdv deadLetterKey = MessageIdAdvUtils.discardBatch(messageId); + final List> deadLetterMessages; + synchronized (possibleSendToDeadLetterTopicMessages) { + deadLetterMessages = possibleSendToDeadLetterTopicMessages.get(deadLetterKey); + if (deadLetterMessages == null) { + return CompletableFuture.completedFuture(false); + } + deadLetterMessages.forEach(MessageImpl::retain); + } + + CompletableFuture sendFuture = initDeadLetterProducerIfNeeded().thenComposeAsync(producerDLQ -> { + List> sendFutures = new ArrayList<>(deadLetterMessages.size()); + try { + for (MessageImpl message : deadLetterMessages) { + String originMessageIdStr = message.getMessageId().toString(); + String originTopicNameStr = getOriginTopicNameStr(message); + AutoProduceBytesSchema deadLetterMessageSchema = (AutoProduceBytesSchema) Schema.AUTO_PRODUCE_BYTES(message.getReaderSchema().get()); - deadLetterMessageSchema.setRequireSchemaValidation(false); - TypedMessageBuilder typedMessageBuilderNew = - producerDLQ.newMessage(deadLetterMessageSchema) - .value(message.getData()) - .properties(getPropertiesMap(message, originMessageIdStr, originTopicNameStr)); - copyMessageKeysIfNeeded(message, typedMessageBuilderNew); - copyMessageEventTime(message, typedMessageBuilderNew); - typedMessageBuilderNew.sendAsync() - .thenAccept(messageIdInDLQ -> { - possibleSendToDeadLetterTopicMessages.remove(messageId); - acknowledgeAsync(messageId).whenComplete((v, ex) -> { - if (ex != null) { - log.warn().attr("messageId", messageId) - .exception(ex) - .log("Failed to acknowledge the" - + " message of the original topic" - + " but send to the DLQ" - + " successfully."); - result.complete(false); - } else { - result.complete(true); - } - }); - }).exceptionally(ex -> { - if (ex instanceof PulsarClientException.ProducerQueueIsFullError) { - log.warn().attr("deadLetterTopic", deadLetterPolicy.getDeadLetterTopic()) - .attr("messageId", messageId) - .exceptionMessage(ex) - .log("Failed to send DLQ message to" - + " with ProducerQueueIsFullError" - + " for message id"); - } else { - log.warn().attr("deadLetterTopic", deadLetterPolicy.getDeadLetterTopic()) - .attr("messageId", messageId) - .exception(ex) - .log("Failed to send DLQ message to for message id"); - } - result.complete(false); - return null; - }); - } catch (Exception e) { - log.warn().attr("deadLetterTopic", deadLetterPolicy.getDeadLetterTopic()) - .attr("messageId", messageId) - .exception(e) - .log("Failed to process DLQ message to for message id"); - result.complete(false); - } + deadLetterMessageSchema.setRequireSchemaValidation(false); + TypedMessageBuilder typedMessageBuilderNew = + producerDLQ.newMessage(deadLetterMessageSchema) + .value(message.getData()) + .properties(getPropertiesMap(message, originMessageIdStr, originTopicNameStr)); + copyMessageKeysIfNeeded(message, typedMessageBuilderNew); + copyMessageEventTime(message, typedMessageBuilderNew); + sendFutures.add(typedMessageBuilderNew.sendAsync()); } - }, internalPinnedExecutor).exceptionally(ex -> { - log.error().attr("topic", deadLetterPolicy.getDeadLetterTopic()) - .exception(ex) - .log("Dead letter producer exception with topic"); - result.complete(false); - return null; + } catch (Exception e) { + return FutureUtil.failedFuture(e); + } + return FutureUtil.waitForAll(sendFutures); + }, internalPinnedExecutor); + + return sendFuture.thenCompose(ignored -> { + releasePossibleDeadLetterMessages(deadLetterKey, deadLetterMessages); + return acknowledgeAsync(messageId).handle((v, ex) -> { + if (ex != null) { + log.warn().attr("messageId", messageId) + .exception(ex) + .log("Failed to acknowledge the message of the original topic" + + " but send to the DLQ successfully."); + return false; + } + return true; }); - } else { - result.complete(false); + }).exceptionally(ex -> { + Throwable cause = FutureUtil.unwrapCompletionException(ex); + log.warn().attr("deadLetterTopic", deadLetterPolicy.getDeadLetterTopic()) + .attr("messageId", messageId) + .exception(cause) + .log("Failed to send DLQ message for message id"); + return false; + }).whenComplete((ignored, ex) -> deadLetterMessages.forEach(MessageImpl::release)); + } + + private void retainPossibleDeadLetterMessages(MessageIdAdv messageId, List> messages) { + synchronized (possibleSendToDeadLetterTopicMessages) { + List> previous = possibleSendToDeadLetterTopicMessages.put(messageId, messages); + if (previous != null && previous != messages) { + previous.forEach(MessageImpl::release); + } + } + } + + private void releasePossibleDeadLetterMessages(MessageIdAdv messageId, List> messages) { + synchronized (possibleSendToDeadLetterTopicMessages) { + if (possibleSendToDeadLetterTopicMessages.remove(messageId, messages)) { + messages.forEach(MessageImpl::release); + } + } + } + + private void clearPossibleDeadLetterMessages() { + synchronized (possibleSendToDeadLetterTopicMessages) { + possibleSendToDeadLetterTopicMessages.values() + .forEach(messages -> messages.forEach(MessageImpl::release)); + possibleSendToDeadLetterTopicMessages.clear(); } - return result; } private void customizeDeadLetterProducerBuilder(DeadLetterProducerBuilderCustomizer customizer, @@ -3008,8 +3482,8 @@ private Optional createEncryptionContext(MessageMetadata msgM return Optional.ofNullable(encryptionCtx); } - private int removeExpiredMessagesFromQueue(Set messageIds) { - int messagesFromQueue = 0; + private List> removeExpiredMessagesFromQueue(Set messageIds) { + List> messagesFromQueue = new ArrayList<>(); Message message; while (true) { message = incomingMessages.pollIf(msg -> { @@ -3020,8 +3494,7 @@ private int removeExpiredMessagesFromQueue(Set messageIds) { break; } decreaseIncomingMessageSize(message); - messagesFromQueue++; - message.release(); + messagesFromQueue.add(message); } return messagesFromQueue; } @@ -3072,6 +3545,7 @@ void resetBackoff() { } void connectionClosed(ClientCnx cnx, Optional initialConnectionDelayMs, Optional hostUrl) { + deactivatePermitState(getPermitState(cnx)); this.connectionHandler.connectionClosed(cnx, initialConnectionDelayMs, hostUrl); } @@ -3081,6 +3555,9 @@ public ClientCnx getClientCnx() { void setClientCnx(ClientCnx clientCnx) { if (clientCnx != null) { + if (getPermitState(clientCnx) == null) { + replacePermitState(clientCnx); + } this.connectionHandler.setClientCnx(clientCnx); clientCnx.registerConsumer(consumerId, this); if (conf.isAckReceiptEnabled() @@ -3096,6 +3573,7 @@ void setClientCnx(ClientCnx clientCnx) { } void deregisterFromClientCnx() { + deactivatePermitState(permitState.get()); setClientCnx(null); } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessageImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessageImpl.java index 4821842f9bb77..bd29083daac89 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessageImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessageImpl.java @@ -36,6 +36,7 @@ import java.util.Optional; import java.util.TreeMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import java.util.stream.Collectors; import lombok.Getter; import org.apache.pulsar.client.api.CompressionType; @@ -85,6 +86,17 @@ public class MessageImpl implements TraceableMessage, Message { private boolean poolMessage; @Getter private long consumerEpoch; + @SuppressWarnings("rawtypes") + private static final AtomicIntegerFieldUpdater FLOW_PERMIT_COST_UPDATER = + AtomicIntegerFieldUpdater.newUpdater(MessageImpl.class, "flowPermitCost"); + @SuppressWarnings("unused") + private volatile int flowPermitCost; + private ConsumerImpl.ConsumerPermitState permitState; + @SuppressWarnings("rawtypes") + private static final AtomicIntegerFieldUpdater REFERENCES_UPDATER = + AtomicIntegerFieldUpdater.newUpdater(MessageImpl.class, "references"); + @SuppressWarnings("unused") + private volatile int references; /** * OpenTelemetry tracing span associated with this message. @@ -107,6 +119,9 @@ public static MessageImpl create(MessageMetadata msgMetadata, ByteBuffer msg.schema = schema; msg.schemaHash = SchemaHash.of(schema); msg.uncompressedSize = payload.remaining(); + msg.permitState = null; + FLOW_PERMIT_COST_UPDATER.set(msg, 0); + REFERENCES_UPDATER.set(msg, 1); return msg; } @@ -198,6 +213,9 @@ private static void init(MessageImpl msg, String topic, BatchMessageIdImp msg.encryptionCtx = encryptionCtx; msg.schema = schema; msg.consumerEpoch = consumerEpoch; + msg.permitState = null; + FLOW_PERMIT_COST_UPDATER.set(msg, 0); + REFERENCES_UPDATER.set(msg, 1); msg.poolMessage = poolMessage; // If it's not pool message then need to make a copy since the passed payload is @@ -721,6 +739,40 @@ public ClientCnx getCnx() { return cnx; } + int getFlowPermitCost() { + return FLOW_PERMIT_COST_UPDATER.get(this); + } + + int takeFlowPermitCost() { + return FLOW_PERMIT_COST_UPDATER.getAndSet(this, 0); + } + + ConsumerImpl.ConsumerPermitState getPermitState() { + return permitState; + } + + void setFlowPermitOwnership(ConsumerImpl.ConsumerPermitState permitState, int flowPermitCost) { + if (flowPermitCost < 0 || (flowPermitCost > 0 && permitState == null)) { + throw new IllegalArgumentException("Permit ownership requires a source state and a non-negative cost"); + } + this.permitState = permitState; + FLOW_PERMIT_COST_UPDATER.set(this, flowPermitCost); + } + + void retain() { + if (!poolMessage) { + return; + } + int current = REFERENCES_UPDATER.get(this); + while (current > 0) { + if (REFERENCES_UPDATER.compareAndSet(this, current, current + 1)) { + return; + } + current = REFERENCES_UPDATER.get(this); + } + throw new IllegalStateException("Cannot retain a released message"); + } + public void recycle() { if (msgMetadata != null) { msgMetadata.clear(); @@ -740,6 +792,9 @@ public void recycle() { schemaState = SchemaState.None; poolMessage = false; consumerEpoch = DEFAULT_CONSUMER_EPOCH; + permitState = null; + FLOW_PERMIT_COST_UPDATER.set(this, 0); + REFERENCES_UPDATER.set(this, 0); if (recyclerHandle != null) { recyclerHandle.recycle(this); @@ -749,8 +804,17 @@ public void recycle() { @Override public void release() { if (poolMessage) { - ReferenceCountUtil.safeRelease(payload); - recycle(); + int current = REFERENCES_UPDATER.get(this); + while (current > 0) { + if (REFERENCES_UPDATER.compareAndSet(this, current, current - 1)) { + if (current == 1) { + ReferenceCountUtil.safeRelease(payload); + recycle(); + } + return; + } + current = REFERENCES_UPDATER.get(this); + } } } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePayloadContextImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePayloadContextImpl.java index 0220f1154d163..64b84268cebe6 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePayloadContextImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePayloadContextImpl.java @@ -50,6 +50,7 @@ protected MessagePayloadContextImpl newObject(Handle private SingleMessageMetadata singleMessageMetadata; private MessageIdImpl messageId; private ConsumerImpl consumer; + private ConsumerImpl.ConsumerPermitState permitState; private int redeliveryCount; private BitSet ackSetInMessageId; private BitSetRecyclable ackBitSet; @@ -66,6 +67,18 @@ public static MessagePayloadContextImpl get(final BrokerEntryMetadata brokerEntr final int redeliveryCount, final List ackSet, final long consumerEpoch) { + return get(brokerEntryMetadata, messageMetadata, messageId, consumer, redeliveryCount, ackSet, + consumerEpoch, consumer.getPermitState()); + } + + public static MessagePayloadContextImpl get(final BrokerEntryMetadata brokerEntryMetadata, + @NonNull final MessageMetadata messageMetadata, + @NonNull final MessageIdImpl messageId, + @NonNull final ConsumerImpl consumer, + final int redeliveryCount, + final List ackSet, + final long consumerEpoch, + final ConsumerImpl.ConsumerPermitState permitState) { final MessagePayloadContextImpl context = RECYCLER.get(); context.consumerEpoch = consumerEpoch; context.brokerEntryMetadata = brokerEntryMetadata; @@ -73,6 +86,7 @@ public static MessagePayloadContextImpl get(final BrokerEntryMetadata brokerEntr context.singleMessageMetadata = new SingleMessageMetadata(); context.messageId = messageId; context.consumer = consumer; + context.permitState = permitState; context.redeliveryCount = redeliveryCount; context.ackSetInMessageId = BatchMessageIdImpl.newAckSet(context.getNumMessages()); boolean isAckSetNotEmpty = ackSet != null && ackSet.size() > 0; @@ -92,6 +106,7 @@ public void recycle() { singleMessageMetadata = null; messageId = null; consumer = null; + permitState = null; redeliveryCount = 0; consumerEpoch = DEFAULT_CONSUMER_EPOCH; ackSetInMessageId = null; @@ -130,7 +145,7 @@ public Message getMessageAt(int index, Schema schema) { final ByteBuf payloadBuffer = MessagePayloadUtils.convertToByteBuf(payload); try { - return consumer.newSingleMessage(index, + MessageImpl message = consumer.newSingleMessage(index, numMessages, brokerEntryMetadata, messageMetadata, @@ -143,7 +158,12 @@ public Message getMessageAt(int index, ackSetInMessageId, redeliveryCount, consumerEpoch, - false); + false, + permitState == null ? null : permitState.cnx); + if (message != null) { + message.setFlowPermitOwnership(permitState, 1); + } + return message; } finally { payloadBuffer.release(); } @@ -153,8 +173,11 @@ public Message getMessageAt(int index, public Message asSingleMessage(MessagePayload payload, Schema schema) { final ByteBuf payloadBuffer = MessagePayloadUtils.convertToByteBuf(payload); try { - return consumer.newMessage(messageId, brokerEntryMetadata, - messageMetadata, payloadBuffer, schema, redeliveryCount, consumerEpoch); + MessageImpl message = consumer.newMessage(messageId, brokerEntryMetadata, + messageMetadata, payloadBuffer, schema, redeliveryCount, consumerEpoch, false, + permitState == null ? null : permitState.cnx); + message.setFlowPermitOwnership(permitState, 1); + return message; } finally { payloadBuffer.release(); } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImpl.java index dad3285b53d0a..35c3ce33649c6 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImpl.java @@ -274,21 +274,28 @@ private void receiveMessageFromConsumer(ConsumerImpl consumer, boolean batchR .log("Receive message from sub consumer"); // Stop to process the remaining message after the consumer is closed. if (getState() == State.Closed) { + messages.forEach(message -> { + MessageImpl messageImpl = (MessageImpl) message; + consumer.increaseAvailablePermits(messageImpl); + message.release(); + }); return; } // Process the message, add to the queue and trigger listener or async callback messages.forEach(msg -> { final boolean skipDueToSeek = duringSeek; MessageImpl msgImpl = (MessageImpl) msg; - ClientCnx cnx = msgImpl.getCnx(); boolean isValidEpoch = isValidConsumerEpoch(msgImpl); if (isValidEpoch && !skipDueToSeek) { messageReceived(consumer, msg); } else if (!isValidEpoch) { - consumer.increaseAvailablePermits(cnx); + consumer.increaseAvailablePermits(msgImpl); + msg.release(); } else if (skipDueToSeek) { log.info().attr("messageId", msg.getMessageId()) .log("Skip processing message received during seek"); + consumer.increaseAvailablePermits(msgImpl); + msg.release(); } }); @@ -351,6 +358,12 @@ protected synchronized void messageProcessed(Message msg) { decreaseIncomingMessageSize(msg); } + @Override + protected void messageDiscarded(Message message) { + returnPermitForParentListener(message); + message.release(); + } + private void resumeReceivingFromPausedConsumersIfNeeded() { if (incomingMessages.size() <= getCurrentReceiverQueueSize() / 2 && !pausedConsumers.isEmpty()) { while (true) { @@ -730,7 +743,7 @@ public void redeliverUnacknowledgedMessages() { consumer.redeliverUnacknowledgedMessages(); consumer.unAckedChunkedMessageIdSequenceMap.clear(); }); - clearIncomingMessages(); + clearIncomingMessagesForRedelivery(); unAckedMessageTracker.clear(); resumeReceivingFromPausedConsumersIfNeeded(); } finally { @@ -952,16 +965,40 @@ private void removeExpiredMessagesFromQueue(Set messageIds) { while (message != null) { decreaseIncomingMessageSize(message); MessageId messageId = message.getMessageId(); + boolean shouldStop = false; if (!messageIds.contains(messageId)) { messageIds.add(messageId); - break; + shouldStop = true; } + returnPermitForParentListener(message); message.release(); + if (shouldStop) { + break; + } message = incomingMessages.poll(); } } } + private void clearIncomingMessagesForRedelivery() { + List> messages = new ArrayList<>(incomingMessages.size()); + incomingMessages.drainTo(messages); + messages.forEach(message -> { + returnPermitForParentListener(message); + message.release(); + }); + resetIncomingMessageSize(); + } + + private void returnPermitForParentListener(Message message) { + if (message instanceof TopicMessageImpl topicMessage) { + Message innerMessage = topicMessage.getMessage(); + if (innerMessage instanceof MessageImpl messageImpl) { + topicMessage.receivedByconsumer.increaseAvailablePermits(messageImpl); + } + } + } + private TopicName getTopicName(String topic) { try { return TopicName.get(topic); diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ZeroQueueConsumerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ZeroQueueConsumerImpl.java index dbbeed2973be2..dcaf360858439 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ZeroQueueConsumerImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ZeroQueueConsumerImpl.java @@ -74,6 +74,7 @@ protected Message internalReceive() throws PulsarClientException { try { Message msg = fetchSingleMessageFromBroker(); trackMessage(msg); + increaseAvailablePermits((MessageImpl) msg); return beforeConsume(msg); } finally { zeroQueueLock.unlock(); @@ -95,7 +96,10 @@ private Message fetchSingleMessageFromBroker() throws PulsarClientException { // Just being cautious if (incomingMessages.size() > 0) { log.error("The incoming message queue should never be greater than 0 when Queue size is 0"); - incomingMessages.forEach(Message::release); + incomingMessages.forEach(message -> { + increaseAvailablePermits((MessageImpl) message); + message.release(); + }); incomingMessages.clear(); } @@ -111,16 +115,17 @@ private Message fetchSingleMessageFromBroker() throws PulsarClientException { do { message = incomingMessages.take(); lastDequeuedMessageId = message.getMessageId(); - ClientCnx msgCnx = ((MessageImpl) message).getCnx(); + ConsumerPermitState messageState = ((MessageImpl) message).getPermitState(); // synchronized need to prevent race between connectionOpened and the check "msgCnx == cnx()" synchronized (this) { // if message received due to an old flow - discard it and wait for the message from the // latest flow command - if (msgCnx == cnx()) { + if (messageState == getPermitState()) { waitingOnReceiveForZeroQueueSize = false; break; } } + message.release(); } while (true); stats.updateNumMsgsReceived(message); @@ -132,7 +137,11 @@ private Message fetchSingleMessageFromBroker() throws PulsarClientException { // Finally blocked is invoked in case the block on incomingMessages is interrupted waitingOnReceiveForZeroQueueSize = false; // Clearing the queue in case there was a race with messageReceived - incomingMessages.clear(); + Message queuedMessage; + while ((queuedMessage = incomingMessages.poll()) != null) { + increaseAvailablePermits((MessageImpl) queuedMessage); + queuedMessage.release(); + } } } @@ -159,6 +168,14 @@ protected boolean canEnqueueMessage(Message message) { } } + @Override + void notifyPendingReceivedCallback(final Message message, Exception exception) { + super.notifyPendingReceivedCallback(message, exception); + if (message instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); + } + } + private void triggerZeroQueueSizeListener(final Message message) { Objects.requireNonNull(listener, "listener can't be null"); Objects.requireNonNull(message, "unqueued message can't be null"); @@ -185,7 +202,7 @@ private void triggerZeroQueueSizeListener(final Message message) { .exception(t) .log("Message listener error in processing unqueued message"); } - increaseAvailablePermits(cnx()); + increaseAvailablePermits((MessageImpl) message); waitingOnListenerForZeroQueueSize = false; }); } @@ -198,11 +215,12 @@ protected void tryTriggerListener() { @Override void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, - MessageIdData messageId, ClientCnx cnx, long consumerEpoch, - boolean isEncrypted) { + MessageIdData messageId, ConsumerPermitState sourceState, + long consumerEpoch, + boolean isEncrypted, int messagePermits) { rejectBatchMessageByClosingConsumer( - new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()) + new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()), sourceState ); } @@ -216,18 +234,24 @@ protected void setCurrentReceiverQueueSize(int newSize) { protected void processPayloadByProcessor(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata messageMetadata, ByteBuf byteBuf, MessageIdImpl messageId, Schema schema, - int redeliveryCount, List ackSet, long consumerEpoch) { + int redeliveryCount, List ackSet, long consumerEpoch, + int messagePermits, ConsumerPermitState sourceState) { if (this.isBatch(messageMetadata)) { - rejectBatchMessageByClosingConsumer(messageId); + try { + rejectBatchMessageByClosingConsumer(messageId, sourceState); + } finally { + byteBuf.release(); + } } else { super.processPayloadByProcessor(brokerEntryMetadata, messageMetadata, byteBuf, messageId, schema, - redeliveryCount, ackSet, consumerEpoch); + redeliveryCount, ackSet, consumerEpoch, messagePermits, sourceState); } } - private void rejectBatchMessageByClosingConsumer(MessageIdImpl messageId) { + private void rejectBatchMessageByClosingConsumer(MessageIdImpl messageId, ConsumerPermitState sourceState) { log.warn().attr("messageId", messageId) .log("Closing consumer - due to unsupported received batch-message with zero receiver queue size"); + deactivatePermitState(sourceState); // close connection closeAsync().handle((ok, e) -> { // notify callback with failure result diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java index 33732e56a5a44..acbbfff6ab4fa 100644 --- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java +++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java @@ -23,6 +23,7 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.any; import static org.mockito.Mockito.atLeast; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; @@ -33,9 +34,16 @@ import static org.testng.Assert.assertTrue; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPromise; +import io.netty.util.concurrent.EventExecutor; +import java.util.ArrayList; import java.util.Arrays; import java.util.BitSet; import java.util.List; +import java.util.Optional; +import java.util.SplittableRandom; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; @@ -56,23 +64,35 @@ import org.apache.pulsar.client.impl.conf.ClientConfigurationData; import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData; import org.apache.pulsar.client.impl.conf.TopicConsumerConfigurationData; +import org.apache.pulsar.client.impl.metrics.InstrumentProvider; +import org.apache.pulsar.client.impl.metrics.UpDownCounter; import org.apache.pulsar.client.util.ExecutorProvider; import org.apache.pulsar.client.util.ScheduledExecutorProvider; +import org.apache.pulsar.common.api.proto.BaseCommand; +import org.apache.pulsar.common.api.proto.CommandMessage; +import org.apache.pulsar.common.api.proto.MessageIdData; import org.apache.pulsar.common.api.proto.MessageMetadata; +import org.apache.pulsar.common.api.proto.SingleMessageMetadata; +import org.apache.pulsar.common.protocol.Commands; import org.apache.pulsar.common.util.Backoff; import org.awaitility.Awaitility; +import org.mockito.ArgumentCaptor; import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class ConsumerImplTest { + private static final long RANDOM_PERMIT_SEED = 0x491C11E17L; + private final String topic = "non-persistent://tenant/ns1/my-topic"; private ExecutorProvider executorProvider; private ExecutorService internalExecutor; private ConsumerImpl consumer; private ConsumerConfigurationData consumerConf; + private UpDownCounter messagesPrefetchedGauge; + private UpDownCounter bytesPrefetchedGauge; @BeforeMethod(alwaysRun = true) public void setUp() { @@ -86,6 +106,13 @@ private void createConsumer(ConsumerConfigurationData consumerConf) { internalExecutor = Executors.newSingleThreadScheduledExecutor(); PulsarClientImpl client = ClientTestFixtures.createPulsarClientMock(executorProvider, internalExecutor); + InstrumentProvider instrumentProvider = spy(InstrumentProvider.NOOP); + messagesPrefetchedGauge = mock(UpDownCounter.class); + bytesPrefetchedGauge = mock(UpDownCounter.class); + doAnswer(invocation -> invocation.getArgument(0).endsWith(".count") + ? messagesPrefetchedGauge : bytesPrefetchedGauge) + .when(instrumentProvider).newUpDownCounter(any(), any(), any(), any(), any()); + when(client.instrumentProvider()).thenReturn(instrumentProvider); ClientConfigurationData clientConf = client.getConfiguration(); clientConf.setOperationTimeoutMs(100); clientConf.setStatsIntervalSeconds(0); @@ -115,6 +142,315 @@ public void testNotifyPendingReceivedCallback_EmptyQueueNotThrowsException() { consumer.notifyPendingReceivedCallback(null, null); } + @Test + public void testGetMessagePermitsUsesCommandValueWhenPresent() { + CommandMessage command = new CommandMessage().setMessagePermits(3); + command.addAckSet(0b10101L); + + Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 3); + } + + @Test + public void testGetMessagePermitsFallsBackToAckSetForOldBroker() { + CommandMessage command = new CommandMessage(); + command.addAckSet(0b101101L); + + Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 4); + } + + @Test + public void testGetMessagePermitsFallsBackToBatchSizeForOldBroker() { + Assert.assertEquals(ConsumerImpl.getMessagePermits(new CommandMessage(), 10), 10); + } + + @Test + public void testGetMessagePermitsValidatesNativeCommand() { + Assert.assertEquals(ConsumerImpl.getMessagePermits(new CommandMessage().setMessagePermits(1), 1), 1); + Assert.assertThrows(RuntimeException.class, + () -> ConsumerImpl.getMessagePermits(new CommandMessage().setMessagePermits(0), 1)); + Assert.assertThrows(RuntimeException.class, + () -> ConsumerImpl.getMessagePermits(new CommandMessage().setMessagePermits(-1), 1)); + Assert.assertThrows(RuntimeException.class, + () -> ConsumerImpl.getMessagePermits(new CommandMessage(), 0)); + Assert.assertThrows(RuntimeException.class, + () -> { + CommandMessage command = new CommandMessage(); + command.addAckSet(0); + ConsumerImpl.getMessagePermits(command, 1); + }); + Assert.assertThrows(RuntimeException.class, + () -> { + CommandMessage command = new CommandMessage(); + command.addAckSet(0b10); + ConsumerImpl.getMessagePermits(command, 1); + }); + Assert.assertThrows(RuntimeException.class, + () -> { + CommandMessage command = new CommandMessage().setMessagePermits(2); + command.addAckSet(0b1); + ConsumerImpl.getMessagePermits(command, 2); + }); + } + + @Test + public void testChecksumFailureReturnsCommandPermits() { + final int messagePermits = 10; + CommandMessage command = new CommandMessage() + .setConsumerId(consumer.consumerId) + .setMessagePermits(messagePermits); + command.setMessageId().setLedgerId(1).setEntryId(2); + MessageMetadata metadata = new MessageMetadata() + .setProducerName("producer") + .setSequenceId(0) + .setPublishTime(System.currentTimeMillis()) + .setNumMessagesInBatch(messagePermits); + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); + ByteBuf metadataAndPayload = Commands.serializeMetadataAndPayload( + Commands.ChecksumType.Crc32c, metadata, payload); + payload.release(); + metadataAndPayload.setByte(metadataAndPayload.writerIndex() - 1, + metadataAndPayload.getByte(metadataAndPayload.writerIndex() - 1) ^ 1); + + ClientCnx cnx = mock(ClientCnx.class); + ChannelHandlerContext context = mock(ChannelHandlerContext.class); + ChannelPromise promise = mock(ChannelPromise.class); + when(cnx.ctx()).thenReturn(context); + when(context.voidPromise()).thenReturn(promise); + consumer.setClientCnx(cnx); + int permitsBefore = consumer.getAvailablePermits(); + try { + consumer.messageReceived(command, metadataAndPayload, cnx); + Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + messagePermits); + } finally { + metadataAndPayload.release(); + } + } + + @Test + public void testPermitReturnIsBoundToConsumerIncarnationOnSameConnection() { + ClientCnx cnx = mock(ClientCnx.class); + consumer.setClientCnx(cnx); + ConsumerImpl.ConsumerPermitState oldState = consumer.getPermitState(); + consumer.increaseAvailablePermits(cnx, 1); + Assert.assertEquals(oldState.availablePermits.get(), 1); + + consumer.deactivatePermitState(oldState); + consumer.setClientCnx(cnx); + ConsumerImpl.ConsumerPermitState replacementState = consumer.getPermitState(); + Assert.assertNotSame(replacementState, oldState); + Assert.assertSame(replacementState.cnx, oldState.cnx); + + MessageImpl oldMessage = new MessageImpl<>(topic, new MessageIdImpl(1, 1, -1), + new MessageMetadata(), Unpooled.EMPTY_BUFFER, cnx, Schema.BYTES); + oldMessage.setFlowPermitOwnership(oldState, 1); + consumer.increaseAvailablePermits(oldMessage); + Assert.assertEquals(replacementState.availablePermits.get(), 0); + + MessageImpl currentMessage = new MessageImpl<>(topic, new MessageIdImpl(1, 2, -1), + new MessageMetadata(), Unpooled.EMPTY_BUFFER, cnx, Schema.BYTES); + currentMessage.setFlowPermitOwnership(replacementState, 1); + consumer.increaseAvailablePermits(currentMessage); + consumer.increaseAvailablePermits(currentMessage); + Assert.assertEquals(replacementState.availablePermits.get(), 1, + "A terminal message path must return its permit at most once"); + } + + @Test + public void testMalformedExplicitPermitsCloseSourceIncarnationWithoutCredit() { + ClientCnx cnx = mock(ClientCnx.class); + Channel channel = mock(Channel.class); + when(cnx.channel()).thenReturn(channel); + consumer.setClientCnx(cnx); + CommandMessage command = new CommandMessage().setMessagePermits(0); + command.setMessageId().setLedgerId(1).setEntryId(2); + + consumer.messageReceived(command, Unpooled.EMPTY_BUFFER, cnx); + + verify(channel).close(); + Assert.assertNull(consumer.getPermitState()); + Assert.assertEquals(consumer.getAvailablePermits(), 0); + } + + @Test + public void testPermitAccumulatorOverflowClosesSourceIncarnation() { + ClientCnx cnx = mock(ClientCnx.class); + Channel channel = mock(Channel.class); + when(cnx.channel()).thenReturn(channel); + consumer.setClientCnx(cnx); + ConsumerImpl.ConsumerPermitState state = consumer.getPermitState(); + state.availablePermits.set(Integer.MAX_VALUE); + MessageImpl message = new MessageImpl<>(topic, new MessageIdImpl(1, 1, -1), + new MessageMetadata(), Unpooled.EMPTY_BUFFER, cnx, Schema.BYTES); + message.setFlowPermitOwnership(state, 1); + + consumer.increaseAvailablePermits(message); + + verify(channel).close(); + Assert.assertNull(consumer.getPermitState()); + } + + @Test + public void testPooledMessageCanBeRetainedByDeadLetterAndMessageLifecycles() { + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); + MessageImpl message = MessageImpl.create(topic, new MessageIdImpl(1, 1, -1), + new MessageMetadata(), payload, Optional.empty(), mock(ClientCnx.class), Schema.BYTES, + 0, true, DEFAULT_CONSUMER_EPOCH); + try { + Assert.assertEquals(payload.refCnt(), 2); + message.retain(); + message.release(); + Assert.assertEquals(payload.refCnt(), 2, + "Releasing one owner must not recycle a message retained by dead-letter handling"); + message.release(); + Assert.assertEquals(payload.refCnt(), 1); + } finally { + payload.release(); + } + } + + @Test + public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetchGauges() throws Exception { + consumer.setCurrentReceiverQueueSize(2); + ClientCnx messageCnx = setCurrentConnectionWithFlowEnabled(); + ChannelHandlerContext context = messageCnx.ctx(); + List flowPermits = new ArrayList<>(); + doAnswer(invocation -> { + flowPermits.add(parseFlowPermits(invocation.getArgument(0))); + return null; + }).when(context).writeAndFlush(any(ByteBuf.class), any(ChannelPromise.class)); + ConsumerBase.CONSUMER_EPOCH.set(consumer, 2); + MessageMetadata metadata = new MessageMetadata() + .setProducerName("producer") + .setSequenceId(1) + .setPublishTime(1) + .setNumMessagesInBatch(1); + ByteBuf batch = Unpooled.buffer(); + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); + Commands.serializeSingleMessageInBatchWithPayload(new SingleMessageMetadata(), payload, batch); + payload.release(); + + try { + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, batch, + new MessageIdData().setLedgerId(1).setEntryId(2), messageCnx, 1, false, 1); + consumer.internalPinnedExecutor.submit( + () -> Assert.assertEquals(consumer.numMessagesInQueue(), 0)).get(5, TimeUnit.SECONDS); + + Assert.assertEquals(consumer.numMessagesInQueue(), 0); + Assert.assertEquals(consumer.getAvailablePermits(), 0); + Assert.assertEquals(flowPermits, List.of(1)); + verify(messagesPrefetchedGauge).increment(); + verify(messagesPrefetchedGauge).decrement(); + ArgumentCaptor addedBytes = ArgumentCaptor.forClass(Long.class); + ArgumentCaptor subtractedBytes = ArgumentCaptor.forClass(Long.class); + verify(bytesPrefetchedGauge).add(addedBytes.capture()); + verify(bytesPrefetchedGauge).subtract(subtractedBytes.capture()); + Assert.assertEquals(subtractedBytes.getValue(), addedBytes.getValue()); + } finally { + batch.release(); + } + } + + @Test + public void testRandomizedBatchDecodeWritesExactlyTheReturnedPermitsToFlow() throws Exception { + consumer.setCurrentReceiverQueueSize(2); + ClientCnx messageCnx = setCurrentConnectionWithFlowEnabled(); + ChannelHandlerContext context = messageCnx.ctx(); + List flowPermits = new ArrayList<>(); + doAnswer(invocation -> { + flowPermits.add(parseFlowPermits(invocation.getArgument(0))); + return null; + }).when(context).writeAndFlush(any(ByteBuf.class), any(ChannelPromise.class)); + SplittableRandom random = new SplittableRandom(RANDOM_PERMIT_SEED); + int expectedReturnedPermits = 0; + + for (int testCase = 0; testCase < 250; testCase++) { + int batchSize = random.nextInt(1, 65); + int requiredWords = (batchSize + Long.SIZE - 1) / Long.SIZE; + long[] ackSet = new long[requiredWords + random.nextInt(3)]; + for (int word = 0; word < ackSet.length; word++) { + ackSet[word] = random.nextLong(); + } + int requiredIndex = random.nextInt(batchSize); + ackSet[requiredIndex / Long.SIZE] |= 1L << (requiredIndex % Long.SIZE); + BitSet deliveredIndexes = BitSet.valueOf(ackSet); + deliveredIndexes.clear(batchSize, Math.max(batchSize, deliveredIndexes.length())); + int messagePermits = deliveredIndexes.cardinality(); + int expectedQueuedMessages = 0; + MessageMetadata metadata = new MessageMetadata() + .setProducerName("producer") + .setSequenceId(testCase) + .setPublishTime(1) + .setNumMessagesInBatch(batchSize); + ByteBuf batch = Unpooled.buffer(); + for (int index = 0; index < batchSize; index++) { + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {(byte) index}); + boolean compactedOut = deliveredIndexes.get(index) && random.nextInt(5) == 0; + if (deliveredIndexes.get(index) && !compactedOut) { + expectedQueuedMessages++; + } + Commands.serializeSingleMessageInBatchWithPayload( + new SingleMessageMetadata().setCompactedOut(compactedOut), payload, batch); + payload.release(); + } + + try { + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, ackSet, batch, + new MessageIdData().setLedgerId(1).setEntryId(testCase), messageCnx, + DEFAULT_CONSUMER_EPOCH, false, messagePermits); + int queuedMessages = consumer.internalPinnedExecutor.submit(consumer::numMessagesInQueue) + .get(5, TimeUnit.SECONDS); + + Assert.assertEquals(queuedMessages, expectedQueuedMessages, + "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); + for (int permit = 0; permit < expectedQueuedMessages; permit++) { + Message message = consumer.incomingMessages.poll(); + Assert.assertNotNull(message); + consumer.messageProcessed(message); + message.release(); + } + expectedReturnedPermits += messagePermits; + Assert.assertEquals(consumer.getAvailablePermits(), 0, + "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); + Assert.assertEquals(flowPermits.stream().mapToInt(Integer::intValue).sum(), expectedReturnedPermits, + "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); + } finally { + batch.release(); + } + } + } + + private ClientCnx setCurrentConnectionWithFlowEnabled() { + ClientCnx messageCnx = mock(ClientCnx.class); + ChannelHandlerContext context = mock(ChannelHandlerContext.class); + Channel channel = mock(Channel.class); + EventExecutor eventExecutor = mock(EventExecutor.class); + when(context.voidPromise()).thenReturn(mock(ChannelPromise.class)); + when(context.channel()).thenReturn(channel); + when(context.executor()).thenReturn(eventExecutor); + doAnswer(invocation -> { + invocation.getArgument(0).run(); + return null; + }).when(eventExecutor).execute(any(Runnable.class)); + when(messageCnx.ctx()).thenReturn(context); + when(messageCnx.channel()).thenReturn(channel); + consumer.setClientCnx(messageCnx); + consumer.getPermitState().flowEnabled = true; + return messageCnx; + } + + private static int parseFlowPermits(ByteBuf frame) { + try { + frame.skipBytes(Integer.BYTES); + int commandSize = (int) frame.readUnsignedInt(); + BaseCommand command = new BaseCommand(); + command.parseFrom(frame, commandSize); + Assert.assertEquals(command.getType(), BaseCommand.Type.FLOW); + return command.getFlow().getMessagePermits(); + } finally { + frame.release(); + } + } + @Test(invocationTimeOut = 500) public void testCorrectBackoffConfiguration() { final Backoff backoff = consumer.getConnectionHandler().backoff; @@ -266,6 +602,7 @@ public void testCreateConsumerWhenSchemaIsNull() throws PulsarClientException { @Test public void testMaxReceiverQueueSize() { + consumer.setClientCnx(mock(ClientCnx.class)); int size = consumer.getCurrentReceiverQueueSize(); int permits = consumer.getAvailablePermits(); consumer.setCurrentReceiverQueueSize(size + 100); @@ -375,16 +712,23 @@ public void testGetMessageAtSyncsAckSetInMessageIdWithBrokerAckSet() { List ackSet = Arrays.stream(brokerAckSet.toLongArray()).boxed().collect(Collectors.toList()); MessageIdImpl messageId = new MessageIdImpl(1L, 2L, -1); + ClientCnx messageCnx = mock(ClientCnx.class); + ConsumerImpl.ConsumerPermitState permitState = new ConsumerImpl.ConsumerPermitState(messageCnx); MessagePayloadContextImpl context = MessagePayloadContextImpl.get( - null, messageMetadata, messageId, consumer, 0, ackSet, DEFAULT_CONSUMER_EPOCH); + null, messageMetadata, messageId, consumer, 0, ackSet, DEFAULT_CONSUMER_EPOCH, permitState); MessagePayload payload0 = MessagePayloadImpl.create(Unpooled.wrappedBuffer(new byte[]{0})); MessagePayload payload1 = MessagePayloadImpl.create(Unpooled.wrappedBuffer(new byte[]{1})); + Message message1 = null; try { // Index 0 is already acked per the broker, so it must not be redelivered to the app. Assert.assertNull(context.getMessageAt(0, batchSize, payload0, false, Schema.BYTES)); - Message message1 = context.getMessageAt(1, batchSize, payload1, false, Schema.BYTES); + message1 = context.getMessageAt(1, batchSize, payload1, false, Schema.BYTES); Assert.assertNotNull(message1); + Assert.assertSame(((MessageImpl) message1).getCnx(), messageCnx, + "The message must retain the connection that delivered its command"); + Assert.assertSame(((MessageImpl) message1).getPermitState(), permitState); + Assert.assertEquals(((MessageImpl) message1).getFlowPermitCost(), 1); BitSet ackSetInMessageId = ((MessageIdAdv) message1.getMessageId()).getAckSet(); Assert.assertFalse(ackSetInMessageId.get(0), @@ -393,6 +737,9 @@ public void testGetMessageAtSyncsAckSetInMessageIdWithBrokerAckSet() { Assert.assertTrue(ackSetInMessageId.get(1), "index 1 is still outstanding"); Assert.assertTrue(ackSetInMessageId.get(2), "index 2 is still outstanding"); } finally { + if (message1 != null) { + message1.release(); + } payload0.release(); payload1.release(); context.recycle(); diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImplTest.java index 7192f0317b5d5..97ac6b5fcaa9f 100644 --- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImplTest.java +++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImplTest.java @@ -40,6 +40,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashSet; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutionException; @@ -50,9 +51,11 @@ import lombok.Cleanup; import org.apache.pulsar.client.api.Consumer; import org.apache.pulsar.client.api.Message; +import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.Messages; import org.apache.pulsar.client.api.PulsarClientException; import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.api.SubscriptionType; import org.apache.pulsar.client.impl.conf.ClientConfigurationData; import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData; import org.apache.pulsar.client.util.ExecutorProvider; @@ -193,6 +196,36 @@ public void testBatchReceiveAsyncCanBeCancelled() { assertFalse(consumer.hasPendingBatchReceive()); } + @Test + public void testRedeliveryReturnsInnerMessagePermitsForParentListener() { + String topic = "persistent://public/default/redelivery-topic"; + ConsumerConfigurationData consumerConfData = new ConsumerConfigurationData<>(); + consumerConfData.setSubscriptionName("subscriptionName"); + consumerConfData.setSubscriptionType(SubscriptionType.Shared); + consumerConfData.setMessageListener((consumer, message) -> { }); + MultiTopicsConsumerImpl consumer = createMultiTopicsConsumer(consumerConfData); + + @SuppressWarnings("unchecked") + ConsumerImpl internalConsumer = mock(ConsumerImpl.class); + @SuppressWarnings("unchecked") + MessageImpl innerMessage = mock(MessageImpl.class); + when(innerMessage.getMessageId()).thenReturn(new MessageIdImpl(1, 2, -1)); + when(innerMessage.size()).thenReturn(1); + TopicMessageImpl topicMessage = new TopicMessageImpl<>(topic, innerMessage, internalConsumer); + consumer.consumers.put(topic, internalConsumer); + consumer.increaseIncomingMessageSize(topicMessage); + consumer.incomingMessages.add(topicMessage); + + Set messageIds = new HashSet<>(); + messageIds.add(topicMessage.getMessageId()); + consumer.redeliverUnacknowledgedMessages(messageIds); + + verify(internalConsumer).increaseAvailablePermits(innerMessage); + verify(internalConsumer).redeliverUnacknowledgedMessages(any()); + verify(innerMessage).release(); + assertTrue(consumer.incomingMessages.isEmpty()); + } + @Test(expectedExceptions = {IllegalArgumentException.class}) public void testValidTopicNames() { ConsumerConfigurationData consumerConfData = new ConsumerConfigurationData<>(); diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java index 575664e8b0930..cf99fb734f756 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java @@ -537,6 +537,21 @@ public static long getEntryTimestamp(ByteBuf headersAndPayloadWithBrokerEntryMet public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long entryId, int partition, int redeliveryCount, long[] ackSet, long consumerEpoch) { + return newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, consumerEpoch, + null); + } + + public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long entryId, int partition, + int redeliveryCount, long[] ackSet, long consumerEpoch, int messagePermits) { + if (messagePermits <= 0) { + throw new IllegalArgumentException("Message permits must be positive"); + } + return newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, consumerEpoch, + Integer.valueOf(messagePermits)); + } + + private static BaseCommand newMessageCommand(long consumerId, long ledgerId, long entryId, int partition, + int redeliveryCount, long[] ackSet, long consumerEpoch, Integer messagePermits) { BaseCommand cmd = localCmd(Type.MESSAGE); CommandMessage msg = cmd.setMessage() .setConsumerId(consumerId); @@ -552,6 +567,9 @@ public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long if (redeliveryCount > 0) { msg.setRedeliveryCount(redeliveryCount); } + if (messagePermits != null) { + msg.setMessagePermits(messagePermits); + } if (ackSet != null) { for (int i = 0; i < ackSet.length; i++) { msg.addAckSet(ackSet[i]); diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto index 5d30972d93791..77a58e3939926 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -581,6 +581,10 @@ message CommandMessage { optional uint32 redelivery_count = 3 [default = 0]; repeated int64 ack_set = 4; optional uint64 consumer_epoch = 5; + + // Number of permits the broker deducted from the consumer for this command. For a batch message with + // batch-index acknowledgments, this excludes the indexes that were already acknowledged. + optional uint32 message_permits = 6; } message CommandAck { diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/compression/CommandsTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/compression/CommandsTest.java index 77978506c2b2f..f8ba819b916b7 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/compression/CommandsTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/compression/CommandsTest.java @@ -30,6 +30,7 @@ import java.io.IOException; import java.util.Base64; import org.apache.pulsar.common.allocator.PulsarByteBufAllocator; +import org.apache.pulsar.common.api.proto.BaseCommand; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.protocol.ByteBufPair; import org.apache.pulsar.common.protocol.Commands; @@ -39,6 +40,25 @@ public class CommandsTest { + @Test + public void testMessageCommandCarriesPermitCount() { + BaseCommand command = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 10); + + assertTrue(command.getMessage().hasMessagePermits()); + assertEquals(command.getMessage().getMessagePermits(), 10); + + BaseCommand singleMessageCommand = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 1); + assertTrue(singleMessageCommand.getMessage().hasMessagePermits()); + assertEquals(singleMessageCommand.getMessage().getMessagePermits(), 1); + + Assert.assertThrows(IllegalArgumentException.class, + () -> Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 0)); + + BaseCommand legacyCommand = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1); + Assert.assertFalse(legacyCommand.getMessage().hasMessagePermits()); + assertEquals(legacyCommand.getMessage().getMessagePermits(), 0); + } + @Test public void testChecksumSendCommand() throws Exception {