From 87bb02dfc38bf60c8382d9f24e3f99b7edc177a8 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 15 Aug 2026 19:40:34 +0800 Subject: [PATCH 1/8] [improve] PIP-491: Make batch permit accounting explicit Assisted-by: Codex (GPT-5) --- .../pulsar/broker/service/Consumer.java | 82 +++-- .../broker/service/EntryBatchIndexesAcks.java | 22 +- .../broker/service/EntryBatchPermits.java | 56 ++++ .../broker/service/PulsarCommandSender.java | 2 +- .../service/PulsarCommandSenderImpl.java | 6 +- .../pulsar/broker/service/ServerCnx.java | 9 +- ...PersistentDispatcherMultipleConsumers.java | 18 +- ...entDispatcherMultipleConsumersClassic.java | 18 +- .../pulsar/broker/service/ConsumerTest.java | 78 ++++- .../service/EntryBatchIndexesAcksTest.java | 21 +- .../impl/CompactedOutBatchMessageTest.java | 103 ++++++- .../impl/CorruptedBatchMessagePermitTest.java | 107 +++++++ .../pulsar/client/impl/ConsumerImpl.java | 282 +++++++++++++++--- .../client/impl/ZeroQueueConsumerImpl.java | 2 +- .../pulsar/client/impl/ConsumerImplTest.java | 57 ++++ .../pulsar/common/protocol/Commands.java | 11 + pulsar-common/src/main/proto/PulsarApi.proto | 3 + .../common/compression/CommandsTest.java | 15 + 18 files changed, 797 insertions(+), 95 deletions(-) create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CorruptedBatchMessagePermitTest.java 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..84211cebcc134 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 @@ -324,7 +324,15 @@ public Future sendMessages(final List entries, EntryBatch EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker) { - return sendMessages(entries, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + return sendMessages(entries, null, batchSizes, batchIndexesAcks, new EntryBatchPermits(entries.size()), + totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH); + } + + public Future sendMessages(final List entries, EntryBatchSizes batchSizes, + EntryBatchIndexesAcks batchIndexesAcks, EntryBatchPermits batchPermits, + int totalMessages, long totalBytes, long totalChunkedMessages, + RedeliveryTracker redeliveryTracker) { + return sendMessages(entries, null, batchSizes, batchIndexesAcks, batchPermits, totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH); } @@ -332,8 +340,8 @@ public Future sendMessages(final List entries, EntryBatch EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { - return sendMessages(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, - totalChunkedMessages, redeliveryTracker, epoch); + return sendMessages(entries, null, batchSizes, batchIndexesAcks, new EntryBatchPermits(entries.size()), + totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, epoch); } /** @@ -351,6 +359,21 @@ public Future sendMessages(final List entries, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { + return sendMessages(entries, stickyKeyHashes, batchSizes, batchIndexesAcks, + new EntryBatchPermits(entries.size()), totalMessages, totalBytes, totalChunkedMessages, + redeliveryTracker, epoch); + } + + private Future sendMessages(final List entries, + final List stickyKeyHashes, + EntryBatchSizes batchSizes, + EntryBatchIndexesAcks batchIndexesAcks, + EntryBatchPermits batchPermits, + int totalMessages, + long totalBytes, + long totalChunkedMessages, + RedeliveryTracker redeliveryTracker, + long epoch) { this.lastConsumedTimestamp = System.currentTimeMillis(); if (entries.isEmpty() || totalMessages == 0) { @@ -363,18 +386,25 @@ public Future sendMessages(final List entries, writePromise.setSuccess(null); return writePromise; } - 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) { + totalEntries--; + 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) { @@ -385,18 +415,9 @@ public Future sendMessages(final List entries, } else { 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 = + boolean 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 @@ -415,32 +436,37 @@ public Future sendMessages(final List entries, .attr("entryId", entry.getEntryId()) .attr("batchSize", batchSize) .log("Added entry to pendingAcks"); + batchPermits.setPermits(i, messagePermits); } + } else { + batchPermits.setPermits(i, messagePermits); } } } // calculate avg message per entry - if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1 - // set init value. - avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries); - } else { - avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent - + (1 - avgPercent) * totalMessages / totalEntries); + if (totalEntries > 0) { + if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1 + // set init value. + avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries); + } else { + avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent + + (1 - avgPercent) * totalMessages / 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); + int sentMessagePermits = batchPermits.getTotalPermits(); + // Reduce permits by the logical messages represented by commands that survived final admission. + 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, batchPermits, redeliveryTracker, epoch); writeAndFlushPromise.addListener(status -> { // only increment counters after the messages have been successfully written to the TCP/IP connection if (status.isSuccess()) { 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/EntryBatchPermits.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java new file mode 100644 index 0000000000000..8bdd27d7d85c2 --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java @@ -0,0 +1,56 @@ +/* + * 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; + +/** + * Carries the finalized permit debit for each entry in a send operation. + * + *

The broker consumer populates this object synchronously after final send admission. The command sender and + * persistent Shared dispatcher then consume the same values without deriving them again from batch metadata. + */ +public final class EntryBatchPermits { + private final int[] permits; + private int totalPermits; + + public EntryBatchPermits(int entriesListSize) { + if (entriesListSize < 0) { + throw new IllegalArgumentException("entriesListSize must not be negative"); + } + permits = new int[entriesListSize]; + } + + void setPermits(int entryIdx, int messagePermits) { + if (messagePermits <= 0) { + throw new IllegalArgumentException("messagePermits must be positive"); + } + if (permits[entryIdx] != 0) { + throw new IllegalStateException("Permits already finalized for entry " + entryIdx); + } + permits[entryIdx] = messagePermits; + totalPermits = Math.addExact(totalPermits, messagePermits); + } + + int getPermits(int entryIdx) { + return permits[entryIdx]; + } + + public int getTotalPermits() { + return totalPermits; + } +} 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..4e091080d96f5 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 @@ -85,7 +85,7 @@ void sendLookupResponse(String brokerServiceUrl, String brokerServiceUrlTls, boo Future sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription, int partitionIdx, List entries, EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, + EntryBatchIndexesAcks batchIndexesAcks, EntryBatchPermits batchPermits, RedeliveryTracker redeliveryTracker, long epoch); 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..863e4508f3bbc 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) { + EntryBatchPermits batchPermits, RedeliveryTracker redeliveryTracker, + long epoch) { final ChannelHandlerContext ctx = cnx.ctx(); final ChannelPromise writePromise = ctx.newPromise(); ctx.channel().eventLoop().execute(() -> { @@ -297,7 +298,8 @@ public ChannelPromise sendMessagesToConsumer(long consumerId, String topicName, ctx.write( cnx.newMessageAndIntercept(consumerId, entry.getLedgerId(), entry.getEntryId(), partitionIdx, redeliveryCount, metadataAndPayload, - batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i), topicName, epoch), + batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i), topicName, epoch, + batchPermits.getPermits(i)), ctx.voidPromise()); entriesToRelease.add(entry); } 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..851c0e02d0092 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 @@ -4487,8 +4487,15 @@ private TopicName validateTopicName(String topic, long requestId, Object request public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId, int partition, int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet, String topic, long epoch) { + return newMessageAndIntercept(consumerId, ledgerId, entryId, partition, redeliveryCount, + metadataAndPayload, ackSet, topic, epoch, 1); + } + + 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); + ackSet, epoch, messagePermits); ByteBufPair res = Commands.serializeCommandMessageWithSize(command, metadataAndPayload); if (brokerInterceptor != null) { try { 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..1d545d226014f 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 @@ -63,6 +63,7 @@ import org.apache.pulsar.broker.service.Consumer; import org.apache.pulsar.broker.service.EntryAndMetadata; import org.apache.pulsar.broker.service.EntryBatchIndexesAcks; +import org.apache.pulsar.broker.service.EntryBatchPermits; import org.apache.pulsar.broker.service.EntryBatchSizes; import org.apache.pulsar.broker.service.InMemoryRedeliveryTracker; import org.apache.pulsar.broker.service.RedeliveryTracker; @@ -863,24 +864,25 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesForThisConsumer.size()); EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesForThisConsumer.size()); + EntryBatchPermits batchPermits = new EntryBatchPermits(entriesForThisConsumer.size()); totalEntries += filterEntriesForConsumer(metadataArray, start, entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); totalEntriesProcessed += entriesForThisConsumer.size(); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), - sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); + c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, batchPermits, + 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, -batchPermits.getTotalPermits()); log.debug() .attr("msgSent", msgSent) - .attr("totalAckedIndexCount", () -> batchIndexesAcks.getTotalAckedIndexCount()) + .attr("messagePermits", batchPermits.getTotalPermits()) .log("Added permits to TOTAL_AVAILABLE_PERMITS_UPDATER"); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); @@ -972,11 +974,12 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, final SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal(); final EntryBatchSizes batchSizes = EntryBatchSizes.get(messagesForC); final EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(messagesForC); + final EntryBatchPermits batchPermits = new EntryBatchPermits(messagesForC); totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); totalEntriesProcessed += entryAndMetadataList.size(); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, + consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, batchPermits, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() ).addListener(future -> { @@ -985,8 +988,7 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -batchPermits.getTotalPermits()); 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..e27ad26452d2e 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 @@ -60,6 +60,7 @@ import org.apache.pulsar.broker.service.Consumer; import org.apache.pulsar.broker.service.EntryAndMetadata; import org.apache.pulsar.broker.service.EntryBatchIndexesAcks; +import org.apache.pulsar.broker.service.EntryBatchPermits; import org.apache.pulsar.broker.service.EntryBatchSizes; import org.apache.pulsar.broker.service.InMemoryRedeliveryTracker; import org.apache.pulsar.broker.service.RedeliveryTracker; @@ -755,22 +756,23 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesForThisConsumer.size()); EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesForThisConsumer.size()); + EntryBatchPermits batchPermits = new EntryBatchPermits(entriesForThisConsumer.size()); totalEntries += filterEntriesForConsumer(metadataArray, start, entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), - sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); + c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, batchPermits, + 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, -batchPermits.getTotalPermits()); log.debug() .attr("msgSent", msgSent) - .attr("totalAckedIndexCount", batchIndexesAcks.getTotalAckedIndexCount()) + .attr("messagePermits", batchPermits.getTotalPermits()) .log("Added -( minus) permits to TOTAL_AVAILABLE_PERMITS_UPDATER in " + "PersistentDispatcherMultipleConsumers"); totalMessagesSent += sendMessageInfo.getTotalMessages(); @@ -833,10 +835,11 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, final SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal(); final EntryBatchSizes batchSizes = EntryBatchSizes.get(messagesForC); final EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(messagesForC); + final EntryBatchPermits batchPermits = new EntryBatchPermits(messagesForC); totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, + consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, batchPermits, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() ).addListener(future -> { @@ -845,8 +848,7 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, - -(sendMessageInfo.getTotalMessages() - batchIndexesAcks.getTotalAckedIndexCount())); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -batchPermits.getTotalPermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } 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..0d222180f9eee 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,29 @@ 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.ArgumentMatchers.eq; 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 static org.testng.Assert.assertSame; +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,19 +52,22 @@ @Test(groups = "broker") public class ConsumerTest { private Consumer consumer; + private Subscription subscription; + private ServerCnx cnx; 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); + PersistentTopic topic = mock(PersistentTopic.class); BrokerService brokerService = mock(BrokerService.class); PulsarService pulsarService = mock(PulsarService.class); ServiceConfiguration serviceConfiguration = mock(ServiceConfiguration.class); when(cnx.clientAddress()).thenReturn(address); + when(subscription.getName()).thenReturn("subscription"); when(subscription.getTopic()).thenReturn(topic); when(topic.getBrokerService()).thenReturn(brokerService); when(brokerService.getPulsar()).thenReturn(pulsarService); @@ -73,4 +91,58 @@ public void testGetBytesOutCounter() { consumer.updateStats(stats); assertEquals(consumer.getBytesOutCounter(), 1L); } + + @Test + public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { + Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, 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})); + EntryBatchPermits batchPermits = new EntryBatchPermits(entries.size()); + + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); + when(cnx.getCommandSender()).thenReturn(commandSender); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), + any(), any(), anyLong())) + .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); + + try { + sharedConsumer.sendMessages(entries, batchSizes, batchIndexesAcks, batchPermits, 23, 0, 0, + mock(RedeliveryTracker.class)); + + assertEquals(batchPermits.getTotalPermits(), 3); + assertEquals(sharedConsumer.getAvailablePermits(), 97); + assertEquals(sharedConsumer.getUnackedMessages(), 3); + assertNull(entries.get(1)); + assertNull(entries.get(2)); + verify(rejectedBatch).release(); + verify(emptyPartialBatch).release(); + verify(commandSender).sendMessagesToConsumer(eq(2L), eq("topic"), eq(subscription), anyInt(), eq(entries), + eq(batchSizes), eq(batchIndexesAcks), eq(batchPermits), any(), eq(DEFAULT_CONSUMER_EPOCH)); + assertSame(entries.get(0), partialBatch); + } 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..7910e19e43c6f 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 @@ -25,6 +25,25 @@ public class EntryBatchIndexesAcksTest { + @Test + void shouldCalculateUnackedIndexesWithinBatchBounds() { + BitSetRecyclable bitSet = BitSetRecyclable.create(); + bitSet.set(1); + bitSet.set(4); + bitSet.set(8); + bitSet.set(63); + 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); + } finally { + acks.recycle(); + bitSet.recycle(); + } + } + @Test void shouldResetStateBeforeReusing() { // given @@ -50,4 +69,4 @@ void shouldResetStateBeforeReusing() { assertEquals(acks.getTotalAckedIndexCount(), 0); } -} \ No newline at end of file +} 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..45a726618b31e 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,104 @@ @Test(groups = "broker-impl") public class CompactedOutBatchMessageTest extends SharedPulsarBaseTest { + @Test + public void testBatchRemainderIsNotReturnedToReplacementConnection() throws Exception { + final int batchSize = 2; + 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(newTopicName()) + .subscriptionName("old-connection-subscription") + .receiverQueueSize(20) + .subscribe()) { + 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 testStaleEpochDiscardDoesNotReturnPermitToReplacementConnection() throws Exception { + MessageMetadata metadata = new MessageMetadata() + .setProducerName("foobar") + .setSequenceId(1) + .setPublishTime(1) + .setNumMessagesInBatch(1); + ByteBuf batch = Unpooled.buffer(100); + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); + Commands.serializeSingleMessageInBatchWithPayload(new SingleMessageMetadata(), payload, batch); + payload.release(); + + try (ConsumerImpl consumer = + (ConsumerImpl) pulsarClient.newConsumer().topic(newTopicName()) + .subscriptionName("stale-epoch-subscription") + .receiverQueueSize(20) + .subscribe()) { + ConsumerBase.CONSUMER_EPOCH.set(consumer, 2); + int permitsBefore = consumer.getAvailablePermits(); + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, batch, + new MessageIdData().setLedgerId(1234).setEntryId(567), mock(ClientCnx.class), 1, false, 1); + consumer.internalPinnedExecutor.submit(() -> assertEquals(consumer.numMessagesInQueue(), 0)) + .get(5, TimeUnit.SECONDS); + + assertEquals(consumer.getAvailablePermits(), permitsBefore); + } finally { + batch.release(); + } + } + + @Test + public void testPartialBatchDeserializationReturnsOnlyCommandRemainder() throws Exception { + final int batchSize = 5; + final int parsedMessages = 2; + 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(newTopicName()) + .subscriptionName("partial-batch-subscription") + .receiverQueueSize(20) + .subscribe()) { + int permitsBefore = consumer.getAvailablePermits(); + 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(), permitsBefore + batchSize); + } finally { + truncatedBatch.release(); + } + } + @Test public void testCompactedOutMessages() throws Exception { final String topic1 = newTopicName(); @@ -67,7 +168,7 @@ public void testCompactedOutMessages() throws Exception { // shove it in the sideways consumer.receiveIndividualMessagesFromBatch(brokerEntryMetadata, metadata, 0, null, batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), - consumer.cnx(), DEFAULT_CONSUMER_EPOCH, false); + consumer.cnx(), DEFAULT_CONSUMER_EPOCH, false, metadata.getNumMessagesInBatch()); Message m = consumer.receive(); assertEquals(((BatchMessageIdImpl) m.getMessageId()).getLedgerId(), 1234); assertEquals(((BatchMessageIdImpl) m.getMessageId()).getEntryId(), 567); 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..94d8efe40af7b --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/CorruptedBatchMessagePermitTest.java @@ -0,0 +1,107 @@ +/* + * 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.time.Duration; +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.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.Producer; +import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.api.SubscriptionType; +import org.apache.pulsar.common.api.proto.MessageMetadata; +import org.apache.pulsar.common.protocol.Commands; +import org.awaitility.Awaitility; +import org.testng.annotations.Test; + +@Test(groups = "broker-impl") +public class CorruptedBatchMessagePermitTest extends SharedPulsarBaseTest { + + @Test(timeOut = 30000) + public void testSharedConsumerReturnsAllPermitsForCorruptedBatch() throws Exception { + final int batchSize = 10; + final String subscriptionName = "shared-subscription"; + final String topicName = newTopicName(); + admin.namespaces().setDeduplicationStatus(getNamespace(), false); + + @Cleanup + Consumer consumer = pulsarClient.newConsumer(Schema.STRING) + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .subscribe(); + @Cleanup + Producer producer = pulsarClient.newProducer(Schema.STRING) + .topic(topicName) + .enableBatching(false) + .create(); + + PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get() + .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); + PersistentSubscription subscription = topic.getSubscription(subscriptionName); + publishCorruptedBatch(topic, batchSize).get(10, TimeUnit.SECONDS); + + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + assertEquals(subscription.getConsumers().size(), 1); + assertEquals(subscription.getConsumers().get(0).getAvailablePermits(), 1); + }); + + 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 publishCorruptedBatch(PersistentTopic topic, int batchSize) { + MessageMetadata metadata = new MessageMetadata() + .setProducerName("corrupted-batch-producer") + .setSequenceId(0) + .setPublishTime(System.currentTimeMillis()) + .setNumMessagesInBatch(batchSize); + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); + ByteBuf entry = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c, metadata, payload); + payload.release(); + int lastByteIndex = entry.writerIndex() - 1; + entry.setByte(lastByteIndex, entry.getByte(lastByteIndex) ^ 1); + + 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-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..ddc255926828d 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 @@ -147,6 +147,37 @@ public class ConsumerImpl extends ConsumerBase implements ConnectionHandle private static final long[] EMPTY_ACK_SET = new long[0]; private static final int MAX_REDELIVER_UNACKNOWLEDGED = 1000; + private static final class MessagePermitBudget { + private int remainingPermits; + + private MessagePermitBudget(int messagePermits) { + remainingPermits = messagePermits; + } + + private void claim() { + if (remainingPermits == 0) { + throw new InvalidMessagePermitsException("Batch contains more deliverable messages than permits"); + } + remainingPermits--; + } + + private void restore() { + remainingPermits++; + } + + private int drain() { + int permits = remainingPermits; + remainingPermits = 0; + return permits; + } + } + + private static final class InvalidMessagePermitsException extends IllegalStateException { + private InvalidMessagePermitsException(String message) { + super(message); + } + } + final long consumerId; // Number of messages that have delivered to the application. Every once in a while, this number will be sent to the @@ -1298,6 +1329,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()); + } + + private 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,7 +1383,7 @@ 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; @@ -1364,23 +1415,38 @@ 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()); + } + + private 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; } private void executeNotifyCallback(final MessageImpl message) { + ClientCnx messageCnx = message.getCnx(); + int messageSize = message.size(); messagesPrefetchedGauge.increment(); - bytesPrefetchedGauge.add(message.size()); + bytesPrefetchedGauge.add(messageSize); // 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()); + completeStaleEpochDiscard(messageCnx, messageSize); return; } Message interceptMsg = onArrival(message); @@ -1392,6 +1458,12 @@ private void executeNotifyCallback(final MessageImpl message) { }); } + private void completeStaleEpochDiscard(ClientCnx messageCnx, int messageSize) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(messageSize); + increaseAvailablePermitsForReceivedMessage(messageCnx, 1); + } + protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, final MessageMetadata messageMetadata, final ByteBuf byteBuf, @@ -1428,6 +1500,58 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe tryTriggerListener(); } + @VisibleForTesting + static int getMessagePermitsForEarlyFailure(CommandMessage command) { + if (command.hasMessagePermits()) { + return getExplicitMessagePermits(command); + } + return command.getAckSetsCount() > 0 ? getAckSetCardinality(command, -1) : 1; + } + + @VisibleForTesting + static int getMessagePermits(CommandMessage command, int batchSize) { + int expectedPermits = command.getAckSetsCount() > 0 + ? getAckSetCardinality(command, batchSize) : batchSize; + if (!command.hasMessagePermits()) { + return expectedPermits; + } + + int explicitPermits = getExplicitMessagePermits(command); + if (explicitPermits != expectedPermits) { + throw new InvalidMessagePermitsException("Explicit message permits " + explicitPermits + + " do not match the payload and ack set value " + expectedPermits); + } + return explicitPermits; + } + + private static int getExplicitMessagePermits(CommandMessage command) { + long messagePermits = Integer.toUnsignedLong(command.getMessagePermits()); + if (messagePermits == 0 || messagePermits > Integer.MAX_VALUE) { + throw new InvalidMessagePermitsException("Invalid explicit message permits " + messagePermits); + } + return (int) messagePermits; + } + + private static int getAckSetCardinality(CommandMessage command, int batchSize) { + int words = command.getAckSetsCount(); + int completeWords = batchSize < 0 ? words : Math.min(batchSize >>> 6, words); + long cardinality = 0; + for (int i = 0; i < completeWords; i++) { + cardinality += Long.bitCount(command.getAckSetAt(i)); + } + if (batchSize >= 0) { + int remainingBits = batchSize & 63; + if (remainingBits > 0 && completeWords < words) { + long mask = -1L >>> (Long.SIZE - remainingBits); + cardinality += Long.bitCount(command.getAckSetAt(completeWords) & mask); + } + } + if (cardinality > Integer.MAX_VALUE) { + throw new InvalidMessagePermitsException("Ack set message permits exceed the supported range"); + } + return (int) cardinality; + } + void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, ClientCnx cnx) { long[] ackSet = EMPTY_ACK_SET; if (cmdMessage.getAckSetsCount() > 0) { @@ -1443,12 +1567,19 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien if (cmdMessage.hasConsumerEpoch()) { consumerEpoch = cmdMessage.getConsumerEpoch(); } + final int earlyFailurePermits; + try { + earlyFailurePermits = getMessagePermitsForEarlyFailure(cmdMessage); + } catch (InvalidMessagePermitsException e) { + closeConnectionForInvalidMessagePermits(messageId, cnx, e); + return; + } log.debug().attr("messageId", messageId) .log("Received message"); if (!verifyChecksum(headersAndPayload, messageId)) { // discard message with checksum error - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch); + discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); return; } @@ -1458,11 +1589,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, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); return; } final int numMessages = msgMetadata.getNumMessagesInBatch(); + final int messagePermits; + try { + messagePermits = getMessagePermits(cmdMessage, numMessages); + } catch (InvalidMessagePermitsException e) { + closeConnectionForInvalidMessagePermits(messageId, cnx, e); + return; + } final int numChunks = msgMetadata.hasNumChunksFromMsg() ? msgMetadata.getNumChunksFromMsg() : 0; final boolean isChunkedMessage = numChunks > 1; MessageIdImpl msgId = new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()); @@ -1471,7 +1609,7 @@ 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(cnx, messagePermits); return; } @@ -1489,7 +1627,7 @@ 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, cnx, true, messagePermits); decryptedPayload.release(); if (uncompressedPayload == null) { // Message was discarded on decompression error @@ -1544,7 +1682,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien final MessageImpl message = newMessage(msgId, brokerEntryMetadata, msgMetadata, uncompressedPayload, - schema, redeliveryCount, consumerEpoch, isMessageUndecryptable); + schema, redeliveryCount, consumerEpoch, isMessageUndecryptable, cnx); uncompressedPayload.release(); if (deadLetterPolicy != null && possibleSendToDeadLetterTopicMessages != null) { @@ -1555,7 +1693,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien 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); + increaseAvailablePermits(message); return; } } @@ -1564,7 +1702,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien } else { // handle batch message enqueuing; uncompressed payload has all messages in batch receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, - uncompressedPayload, messageId, cnx, consumerEpoch, isMessageUndecryptable); + uncompressedPayload, messageId, cnx, consumerEpoch, isMessageUndecryptable, messagePermits); uncompressedPayload.release(); } @@ -1713,7 +1851,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, cnx, false, 1); compressedPayload.release(); return uncompressedPayload; } @@ -1784,7 +1923,7 @@ private void interceptAndComplete(final Message message, final CompletableFut void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, - boolean isEncrypted) { + boolean isEncrypted, int messagePermits) { int batchSize = msgMetadata.getNumMessagesInBatch(); // create ack tracker for entry aka batch @@ -1803,44 +1942,60 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } SingleMessageMetadata singleMessageMetadata = new SingleMessageMetadata(); - int skippedMessages = 0; + MessagePermitBudget permitBudget = new MessagePermitBudget(messagePermits); + boolean corrupted = 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, 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++; + message.release(); continue; } - executeNotifyCallback(message); + try { + permitBudget.claim(); + } catch (InvalidMessagePermitsException e) { + message.release(); + throw e; + } + try { + executeNotifyCallback(message); + } catch (RuntimeException e) { + permitBudget.restore(); + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + message.release(); + throw e; + } } + } catch (InvalidMessagePermitsException e) { + closeConnectionForInvalidMessagePermits(messageId, cnx, e); + return; + } catch (RuntimeException e) { + log.warn().exception(e) + .log("unable to obtain message in batch"); + discardCorruptedMessage(messageId, cnx, ValidationError.BatchDeSerializeError, permitBudget.drain()); + corrupted = true; + } finally { if (ackBitSet != null) { ackBitSet.recycle(); } - } catch (IllegalStateException e) { - log.warn().exception(e) - .log("unable to obtain message in batch"); - discardCorruptedMessage(messageId, cnx, ValidationError.BatchDeSerializeError); } if (deadLetterPolicy != null && possibleSendToDeadLetterTopicMessages != null) { @@ -1857,8 +2012,8 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, .attr("remainingCapacity", incomingMessages.remainingCapacity()) .log("enqueued messages in batch. queue size -, available queue size"); - if (skippedMessages > 0) { - increaseAvailablePermits(cnx, skippedMessages); + if (!corrupted) { + increaseAvailablePermitsForReceivedMessage(cnx, permitBudget.drain()); } } @@ -1930,10 +2085,48 @@ protected void trackMessage(MessageId messageId, int redeliveryCount) { } void increaseAvailablePermits(MessageImpl msg) { - ClientCnx currentCnx = cnx(); - ClientCnx msgCnx = msg.getCnx(); - if (msgCnx == currentCnx) { - increaseAvailablePermits(currentCnx); + increaseAvailablePermitsForReceivedMessage(msg.getCnx(), 1); + } + + private void increaseAvailablePermitsForReceivedMessage(ClientCnx messageCnx, int delta) { + if (delta <= 0) { + return; + } + + int permitsToSend = 0; + boolean overflow = false; + synchronized (this) { + if (messageCnx == null || messageCnx != cnx()) { + return; + } + + int available = 0; + while (true) { + int current = AVAILABLE_PERMITS_UPDATER.get(this); + try { + available = Math.addExact(current, delta); + } catch (ArithmeticException e) { + overflow = true; + break; + } + if (AVAILABLE_PERMITS_UPDATER.compareAndSet(this, current, available)) { + break; + } + } + while (!overflow && available >= getCurrentReceiverQueueSize() / 2 && !paused) { + if (AVAILABLE_PERMITS_UPDATER.compareAndSet(this, available, 0)) { + permitsToSend = available; + break; + } + available = AVAILABLE_PERMITS_UPDATER.get(this); + } + } + + if (overflow) { + closeConnectionForInvalidMessagePermits(null, messageCnx, + new InvalidMessagePermitsException("Client permit accumulator overflow")); + } else if (permitsToSend > 0) { + sendFlowPermitsToBroker(messageCnx, permitsToSend); } } @@ -2110,7 +2303,7 @@ private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData message } private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetadata msgMetadata, ByteBuf payload, - ClientCnx currentCnx, boolean checkMaxMessageSize) { + ClientCnx currentCnx, boolean checkMaxMessageSize, int messagePermits) { CompressionType compressionType = msgMetadata.getCompression(); CompressionCodec codec = CompressionCodecProvider.getCompressionCodec(compressionType); int uncompressedSize = msgMetadata.getUncompressedSize(); @@ -2120,19 +2313,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, currentCnx, 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, currentCnx, ValidationError.DecompressionError, messagePermits); return null; } } @@ -2166,19 +2360,29 @@ private void discardCorruptedMessage(MessageIdImpl messageId, ClientCnx currentC } private void discardCorruptedMessage(MessageIdData messageId, ClientCnx currentCnx, - ValidationError validationError) { + ValidationError validationError, int messagePermits) { log.error().attr("messageId", messageId) .log("Discarding corrupted message"); - discardMessage(messageId, currentCnx, validationError, 1); + discardMessage(messageId, currentCnx, validationError, messagePermits); } private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, ValidationError validationError, - int batchMessages) { + 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); + increaseAvailablePermitsForReceivedMessage(currentCnx, messagePermits); + stats.incrementNumReceiveFailed(); + } + + private void closeConnectionForInvalidMessagePermits(MessageIdData messageId, ClientCnx messageCnx, + InvalidMessagePermitsException error) { + log.warn() + .attr("messageId", messageId) + .exceptionMessage(error) + .log("Closing source connection because message permit accounting is invalid"); stats.incrementNumReceiveFailed(); + messageCnx.ctx().close(); } @Override 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..01e9c70a2ca8e 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 @@ -199,7 +199,7 @@ protected void tryTriggerListener() { void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, - boolean isEncrypted) { + boolean isEncrypted, int messagePermits) { rejectBatchMessageByClosingConsumer( new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()) 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..06684fa74d45e 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 @@ -33,6 +33,7 @@ import static org.testng.Assert.assertTrue; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; +import io.netty.channel.ChannelHandlerContext; import java.util.Arrays; import java.util.BitSet; import java.util.List; @@ -58,6 +59,7 @@ import org.apache.pulsar.client.impl.conf.TopicConsumerConfigurationData; import org.apache.pulsar.client.util.ExecutorProvider; import org.apache.pulsar.client.util.ScheduledExecutorProvider; +import org.apache.pulsar.common.api.proto.CommandMessage; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.util.Backoff; import org.awaitility.Awaitility; @@ -115,6 +117,61 @@ public void testNotifyPendingReceivedCallback_EmptyQueueNotThrowsException() { consumer.notifyPendingReceivedCallback(null, null); } + @Test + public void testResolveExplicitMessagePermits() { + CommandMessage command = new CommandMessage().setMessagePermits(3); + command.addAckSet(0b100101L); + + Assert.assertEquals(ConsumerImpl.getMessagePermitsForEarlyFailure(command), 3); + Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 3); + } + + @Test + public void testResolveLegacyMessagePermits() { + CommandMessage partialBatch = new CommandMessage(); + partialBatch.addAckSet((1L << 1) | (1L << 4) | (1L << 63)); + + Assert.assertEquals(ConsumerImpl.getMessagePermitsForEarlyFailure(partialBatch), 3); + Assert.assertEquals(ConsumerImpl.getMessagePermits(partialBatch, 10), 2); + Assert.assertEquals(ConsumerImpl.getMessagePermits(new CommandMessage(), 10), 10); + } + + @Test + public void testRejectInvalidExplicitMessagePermits() { + CommandMessage zero = new CommandMessage().setMessagePermits(0); + CommandMessage unsignedOverflow = new CommandMessage().setMessagePermits(-1); + CommandMessage mismatch = new CommandMessage().setMessagePermits(2); + mismatch.addAckSet(0b111L); + + Assert.expectThrows(IllegalStateException.class, + () -> ConsumerImpl.getMessagePermitsForEarlyFailure(zero)); + Assert.expectThrows(IllegalStateException.class, + () -> ConsumerImpl.getMessagePermitsForEarlyFailure(unsignedOverflow)); + Assert.expectThrows(IllegalStateException.class, + () -> ConsumerImpl.getMessagePermits(mismatch, 10)); + } + + @Test + public void testInvalidExplicitMessagePermitsCloseSourceConnection() { + CommandMessage command = new CommandMessage() + .setConsumerId(consumer.consumerId) + .setMessagePermits(0); + command.setMessageId().setLedgerId(1).setEntryId(2); + ClientCnx messageCnx = mock(ClientCnx.class); + ChannelHandlerContext context = mock(ChannelHandlerContext.class); + when(messageCnx.ctx()).thenReturn(context); + ByteBuf emptyPayload = Unpooled.buffer(0); + int permitsBefore = consumer.getAvailablePermits(); + + try { + consumer.messageReceived(command, emptyPayload, messageCnx); + verify(context).close(); + Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore); + } finally { + emptyPayload.release(); + } + } + @Test(invocationTimeOut = 500) public void testCorrectBackoffConfiguration() { final Backoff backoff = consumer.getConnectionHandler().backoff; 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..f1ed8b0e45a4e 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,14 @@ 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, 1); + } + + 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("messagePermits must be positive"); + } BaseCommand cmd = localCmd(Type.MESSAGE); CommandMessage msg = cmd.setMessage() .setConsumerId(consumerId); @@ -552,6 +560,9 @@ public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long if (redeliveryCount > 0) { msg.setRedeliveryCount(redeliveryCount); } + if (messagePermits > 1) { + 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..c5acf55c2dc33 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -581,6 +581,9 @@ message CommandMessage { optional uint32 redelivery_count = 3 [default = 0]; repeated int64 ack_set = 4; optional uint64 consumer_epoch = 5; + + // Number of consumer permits debited by the broker for this command. + optional uint32 message_permits = 6 [default = 1]; } 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..067c284a4bf24 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,20 @@ public class CommandsTest { + @Test + public void testMessageCommandCarriesPermitDebit() { + BaseCommand batched = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 10); + assertTrue(batched.getMessage().hasMessagePermits()); + assertEquals(batched.getMessage().getMessagePermits(), 10); + + BaseCommand single = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 1); + Assert.assertFalse(single.getMessage().hasMessagePermits()); + assertEquals(single.getMessage().getMessagePermits(), 1); + + Assert.expectThrows(IllegalArgumentException.class, + () -> Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 0)); + } + @Test public void testChecksumSendCommand() throws Exception { From be7a65c20e01b09c4defa4c0cb6791f2b449a153 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 15 Aug 2026 21:14:13 +0800 Subject: [PATCH 2/8] [refactor][broker] Encapsulate finalized send permits Keep per-entry permit accounting owned by Consumer and expose a single send result to Shared dispatchers and the command sender.\n\nAssisted-by: OpenAI Codex --- .../pulsar/broker/service/Consumer.java | 68 +++++++++--------- .../broker/service/EntryBatchPermits.java | 56 --------------- .../broker/service/PulsarCommandSender.java | 2 +- .../service/PulsarCommandSenderImpl.java | 4 +- .../broker/service/SendMessagesResult.java | 72 +++++++++++++++++++ ...PersistentDispatcherMultipleConsumers.java | 20 +++--- ...entDispatcherMultipleConsumersClassic.java | 20 +++--- .../pulsar/broker/service/ConsumerTest.java | 13 ++-- 8 files changed, 138 insertions(+), 117 deletions(-) delete mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java 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 84211cebcc134..48e6c726ba4b7 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 @@ -324,15 +324,18 @@ public Future sendMessages(final List entries, EntryBatch EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker) { - return sendMessages(entries, null, batchSizes, batchIndexesAcks, new EntryBatchPermits(entries.size()), - totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH); - } - - public Future sendMessages(final List entries, EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, EntryBatchPermits batchPermits, - int totalMessages, long totalBytes, long totalChunkedMessages, - RedeliveryTracker redeliveryTracker) { - return sendMessages(entries, null, batchSizes, batchIndexesAcks, batchPermits, totalMessages, totalBytes, + return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH).getWriteFuture(); + } + + public SendMessagesResult sendMessagesWithResult(final List entries, + EntryBatchSizes batchSizes, + EntryBatchIndexesAcks batchIndexesAcks, + int totalMessages, + long totalBytes, + long totalChunkedMessages, + RedeliveryTracker redeliveryTracker) { + return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH); } @@ -340,15 +343,15 @@ public Future sendMessages(final List entries, EntryBatch EntryBatchIndexesAcks batchIndexesAcks, int totalMessages, long totalBytes, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { - return sendMessages(entries, null, batchSizes, batchIndexesAcks, new EntryBatchPermits(entries.size()), - totalMessages, totalBytes, totalChunkedMessages, redeliveryTracker, epoch); + return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + totalChunkedMessages, redeliveryTracker, epoch).getWriteFuture(); } /** * 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 a future completed after the entries have been written */ public Future sendMessages(final List entries, final List stickyKeyHashes, @@ -359,21 +362,19 @@ public Future sendMessages(final List entries, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { - return sendMessages(entries, stickyKeyHashes, batchSizes, batchIndexesAcks, - new EntryBatchPermits(entries.size()), totalMessages, totalBytes, totalChunkedMessages, - redeliveryTracker, epoch); - } - - private Future sendMessages(final List entries, - final List stickyKeyHashes, - EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, - EntryBatchPermits batchPermits, - int totalMessages, - long totalBytes, - long totalChunkedMessages, - RedeliveryTracker redeliveryTracker, - long epoch) { + return sendMessagesInternal(entries, stickyKeyHashes, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + totalChunkedMessages, redeliveryTracker, epoch).getWriteFuture(); + } + + private SendMessagesResult sendMessagesInternal(final List entries, + final List stickyKeyHashes, + EntryBatchSizes batchSizes, + EntryBatchIndexesAcks batchIndexesAcks, + int totalMessages, + long totalBytes, + long totalChunkedMessages, + RedeliveryTracker redeliveryTracker, + long epoch) { this.lastConsumedTimestamp = System.currentTimeMillis(); if (entries.isEmpty() || totalMessages == 0) { @@ -384,8 +385,9 @@ private Future sendMessages(final List entries, } final Promise writePromise = cnx.newPromise(); writePromise.setSuccess(null); - return writePromise; + return new SendMessagesResult(0).setWriteFuture(writePromise); } + SendMessagesResult sendResult = new SendMessagesResult(entries.size()); int totalEntries = 0; for (int i = 0; i < entries.size(); i++) { @@ -436,10 +438,10 @@ private Future sendMessages(final List entries, .attr("entryId", entry.getEntryId()) .attr("batchSize", batchSize) .log("Added entry to pendingAcks"); - batchPermits.setPermits(i, messagePermits); + sendResult.setMessagePermits(i, messagePermits); } } else { - batchPermits.setPermits(i, messagePermits); + sendResult.setMessagePermits(i, messagePermits); } } } @@ -455,7 +457,7 @@ private Future sendMessages(final List entries, } } - int sentMessagePermits = batchPermits.getTotalPermits(); + int sentMessagePermits = sendResult.getTotalMessagePermits(); // Reduce permits by the logical messages represented by commands that survived final admission. MESSAGE_PERMITS_UPDATER.addAndGet(this, -sentMessagePermits); log.debug() @@ -466,7 +468,7 @@ private Future sendMessages(final List entries, incrementUnackedMessages(sentMessagePermits); Future writeAndFlushPromise = cnx.getCommandSender().sendMessagesToConsumer(consumerId, topicName, subscription, partitionIdx, - entries, batchSizes, batchIndexesAcks, batchPermits, redeliveryTracker, epoch); + entries, batchSizes, batchIndexesAcks, sendResult, redeliveryTracker, epoch); writeAndFlushPromise.addListener(status -> { // only increment counters after the messages have been successfully written to the TCP/IP connection if (status.isSuccess()) { @@ -483,7 +485,7 @@ private Future sendMessages(final List entries, .log("Sent messages to client failed by IO exception, closing the connection"); } }); - return writeAndFlushPromise; + return sendResult.setWriteFuture(writeAndFlushPromise); } private void incrementUnackedMessages(int unackedMessages) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java deleted file mode 100644 index 8bdd27d7d85c2..0000000000000 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/EntryBatchPermits.java +++ /dev/null @@ -1,56 +0,0 @@ -/* - * 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; - -/** - * Carries the finalized permit debit for each entry in a send operation. - * - *

The broker consumer populates this object synchronously after final send admission. The command sender and - * persistent Shared dispatcher then consume the same values without deriving them again from batch metadata. - */ -public final class EntryBatchPermits { - private final int[] permits; - private int totalPermits; - - public EntryBatchPermits(int entriesListSize) { - if (entriesListSize < 0) { - throw new IllegalArgumentException("entriesListSize must not be negative"); - } - permits = new int[entriesListSize]; - } - - void setPermits(int entryIdx, int messagePermits) { - if (messagePermits <= 0) { - throw new IllegalArgumentException("messagePermits must be positive"); - } - if (permits[entryIdx] != 0) { - throw new IllegalStateException("Permits already finalized for entry " + entryIdx); - } - permits[entryIdx] = messagePermits; - totalPermits = Math.addExact(totalPermits, messagePermits); - } - - int getPermits(int entryIdx) { - return permits[entryIdx]; - } - - public int getTotalPermits() { - return totalPermits; - } -} 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 4e091080d96f5..6de3c31fc459e 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 @@ -85,7 +85,7 @@ void sendLookupResponse(String brokerServiceUrl, String brokerServiceUrlTls, boo Future sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription, int partitionIdx, List entries, EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, EntryBatchPermits batchPermits, + EntryBatchIndexesAcks batchIndexesAcks, SendMessagesResult sendResult, RedeliveryTracker redeliveryTracker, long epoch); 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 863e4508f3bbc..437719c329f86 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,7 @@ 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, - EntryBatchPermits batchPermits, RedeliveryTracker redeliveryTracker, + SendMessagesResult sendResult, RedeliveryTracker redeliveryTracker, long epoch) { final ChannelHandlerContext ctx = cnx.ctx(); final ChannelPromise writePromise = ctx.newPromise(); @@ -299,7 +299,7 @@ public ChannelPromise sendMessagesToConsumer(long consumerId, String topicName, cnx.newMessageAndIntercept(consumerId, entry.getLedgerId(), entry.getEntryId(), partitionIdx, redeliveryCount, metadataAndPayload, batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i), topicName, epoch, - batchPermits.getPermits(i)), + sendResult.getMessagePermits(i)), ctx.voidPromise()); entriesToRelease.add(entry); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java new file mode 100644 index 0000000000000..d70c0ca1306d1 --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java @@ -0,0 +1,72 @@ +/* + * 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; + +/** + * Contains the finalized permit accounting and asynchronous write completion for a send operation. + * + *

The broker consumer populates this object synchronously after final send admission. The command sender reads + * the per-entry values, while the persistent Shared dispatcher reads their sum. This keeps every covered broker + * counter and command serialization tied to the same finalized result. + */ +public final class SendMessagesResult { + private final int[] messagePermits; + private int totalMessagePermits; + private Future writeFuture; + + SendMessagesResult(int entriesListSize) { + if (entriesListSize < 0) { + throw new IllegalArgumentException("entriesListSize must not be negative"); + } + messagePermits = new int[entriesListSize]; + } + + void setMessagePermits(int entryIdx, int permits) { + if (permits <= 0) { + throw new IllegalArgumentException("message permits must be positive"); + } + if (messagePermits[entryIdx] != 0) { + throw new IllegalStateException("Permits already finalized for entry " + entryIdx); + } + messagePermits[entryIdx] = permits; + totalMessagePermits = Math.addExact(totalMessagePermits, permits); + } + + int getMessagePermits(int entryIdx) { + return messagePermits[entryIdx]; + } + + SendMessagesResult setWriteFuture(Future writeFuture) { + if (this.writeFuture != null) { + throw new IllegalStateException("Write future already set"); + } + this.writeFuture = writeFuture; + return this; + } + + public int getTotalMessagePermits() { + return totalMessagePermits; + } + + public Future getWriteFuture() { + return writeFuture; + } +} 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 1d545d226014f..bd021357d0cf7 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 @@ -63,12 +63,12 @@ import org.apache.pulsar.broker.service.Consumer; import org.apache.pulsar.broker.service.EntryAndMetadata; import org.apache.pulsar.broker.service.EntryBatchIndexesAcks; -import org.apache.pulsar.broker.service.EntryBatchPermits; import org.apache.pulsar.broker.service.EntryBatchSizes; import org.apache.pulsar.broker.service.InMemoryRedeliveryTracker; 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.SendMessagesResult; import org.apache.pulsar.broker.service.SharedConsumerAssignor; import org.apache.pulsar.broker.service.Subscription; import org.apache.pulsar.broker.transaction.exception.buffer.TransactionBufferException; @@ -864,14 +864,14 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesForThisConsumer.size()); EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesForThisConsumer.size()); - EntryBatchPermits batchPermits = new EntryBatchPermits(entriesForThisConsumer.size()); totalEntries += filterEntriesForConsumer(metadataArray, start, entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); totalEntriesProcessed += entriesForThisConsumer.size(); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, batchPermits, + SendMessagesResult sendResult = c.sendMessagesWithResult( + entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); @@ -879,10 +879,10 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis remainingMessages -= msgSent; start += maxEntriesInThisBatch; entriesToDispatch -= maxEntriesInThisBatch; - TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -batchPermits.getTotalPermits()); + TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits()); log.debug() .attr("msgSent", msgSent) - .attr("messagePermits", batchPermits.getTotalPermits()) + .attr("messagePermits", sendResult.getTotalMessagePermits()) .log("Added permits to TOTAL_AVAILABLE_PERMITS_UPDATER"); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); @@ -974,21 +974,21 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, final SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal(); final EntryBatchSizes batchSizes = EntryBatchSizes.get(messagesForC); final EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(messagesForC); - final EntryBatchPermits batchPermits = new EntryBatchPermits(messagesForC); totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); totalEntriesProcessed += entryAndMetadataList.size(); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, batchPermits, + SendMessagesResult sendResult = consumer.sendMessagesWithResult( + entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), - sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() - ).addListener(future -> { + sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker()); + sendResult.getWriteFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntriesAsync(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -batchPermits.getTotalPermits()); + 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 e27ad26452d2e..866d6f2cc6b5a 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 @@ -60,12 +60,12 @@ import org.apache.pulsar.broker.service.Consumer; import org.apache.pulsar.broker.service.EntryAndMetadata; import org.apache.pulsar.broker.service.EntryBatchIndexesAcks; -import org.apache.pulsar.broker.service.EntryBatchPermits; import org.apache.pulsar.broker.service.EntryBatchSizes; import org.apache.pulsar.broker.service.InMemoryRedeliveryTracker; 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.SendMessagesResult; import org.apache.pulsar.broker.service.SharedConsumerAssignor; import org.apache.pulsar.broker.service.StickyKeyConsumerSelector; import org.apache.pulsar.broker.service.Subscription; @@ -756,12 +756,12 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesForThisConsumer.size()); EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesForThisConsumer.size()); - EntryBatchPermits batchPermits = new EntryBatchPermits(entriesForThisConsumer.size()); totalEntries += filterEntriesForConsumer(metadataArray, start, entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); - c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, batchPermits, + SendMessagesResult sendResult = c.sendMessagesWithResult( + entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); @@ -769,10 +769,10 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis remainingMessages -= msgSent; start += messagesForC; entriesToDispatch -= messagesForC; - TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -batchPermits.getTotalPermits()); + TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendResult.getTotalMessagePermits()); log.debug() .attr("msgSent", msgSent) - .attr("messagePermits", batchPermits.getTotalPermits()) + .attr("messagePermits", sendResult.getTotalMessagePermits()) .log("Added -( minus) permits to TOTAL_AVAILABLE_PERMITS_UPDATER in " + "PersistentDispatcherMultipleConsumers"); totalMessagesSent += sendMessageInfo.getTotalMessages(); @@ -835,20 +835,20 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, final SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal(); final EntryBatchSizes batchSizes = EntryBatchSizes.get(messagesForC); final EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(messagesForC); - final EntryBatchPermits batchPermits = new EntryBatchPermits(messagesForC); totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); - consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, batchPermits, + SendMessagesResult sendResult = consumer.sendMessagesWithResult( + entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), - sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() - ).addListener(future -> { + sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker()); + sendResult.getWriteFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntries(); } }); - TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -batchPermits.getTotalPermits()); + TOTAL_AVAILABLE_PERMITS_UPDATER.getAndAdd(this, -sendResult.getTotalMessagePermits()); totalMessagesSent += sendMessageInfo.getTotalMessages(); totalBytesSent += sendMessageInfo.getTotalBytes(); } 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 0d222180f9eee..e9a4737c83c55 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 @@ -46,6 +46,7 @@ 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.mockito.ArgumentCaptor; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -118,7 +119,6 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); batchIndexesAcks.setIndexesAcks(0, Pair.of(10, new long[] {0b100101L})); batchIndexesAcks.setIndexesAcks(2, Pair.of(10, new long[] {0L})); - EntryBatchPermits batchPermits = new EntryBatchPermits(entries.size()); PulsarCommandSender commandSender = mock(PulsarCommandSender.class); when(cnx.getCommandSender()).thenReturn(commandSender); @@ -127,18 +127,21 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); try { - sharedConsumer.sendMessages(entries, batchSizes, batchIndexesAcks, batchPermits, 23, 0, 0, - mock(RedeliveryTracker.class)); + SendMessagesResult sendResult = sharedConsumer.sendMessagesWithResult( + entries, batchSizes, batchIndexesAcks, 23, 0, 0, mock(RedeliveryTracker.class)); - assertEquals(batchPermits.getTotalPermits(), 3); + assertEquals(sendResult.getTotalMessagePermits(), 3); assertEquals(sharedConsumer.getAvailablePermits(), 97); assertEquals(sharedConsumer.getUnackedMessages(), 3); assertNull(entries.get(1)); assertNull(entries.get(2)); verify(rejectedBatch).release(); verify(emptyPartialBatch).release(); + ArgumentCaptor sendResultCaptor = ArgumentCaptor.forClass(SendMessagesResult.class); verify(commandSender).sendMessagesToConsumer(eq(2L), eq("topic"), eq(subscription), anyInt(), eq(entries), - eq(batchSizes), eq(batchIndexesAcks), eq(batchPermits), any(), eq(DEFAULT_CONSUMER_EPOCH)); + eq(batchSizes), eq(batchIndexesAcks), sendResultCaptor.capture(), any(), + eq(DEFAULT_CONSUMER_EPOCH)); + assertSame(sendResultCaptor.getValue(), sendResult); assertSame(entries.get(0), partialBatch); } finally { batchSizes.recyle(); From 4d2c250a87b6b6c1808ab9a22dcc909e96262d9a Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 16 Aug 2026 01:58:58 +0800 Subject: [PATCH 3/8] [improve] Complete explicit message permit accounting Assisted-by: OpenAI Codex (GPT-5) --- .../pulsar/broker/service/Consumer.java | 90 ++++++------ .../broker/service/SendMessagesResult.java | 3 +- .../pulsar/broker/service/ServerCnx.java | 10 +- .../pulsar/broker/service/ConsumerTest.java | 17 +++ .../client/impl/BatchMessageIndexAckTest.java | 114 +++++++++++++++ .../impl/CompactedOutBatchMessageTest.java | 42 +----- .../impl/CorruptedBatchMessagePermitTest.java | 101 +++++++++++-- .../pulsar/client/impl/ConsumerImpl.java | 121 +++------------- .../client/impl/MessagePermitAccounting.java | 111 ++++++++++++++ .../pulsar/client/impl/ConsumerImplTest.java | 135 +++++++++++++----- .../impl/MessagePermitAccountingTest.java | 84 +++++++++++ .../pulsar/common/protocol/Commands.java | 22 +-- pulsar-common/src/main/proto/PulsarApi.proto | 2 +- .../common/compression/CommandsTest.java | 5 +- 14 files changed, 611 insertions(+), 246 deletions(-) create mode 100644 pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java create mode 100644 pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java 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 48e6c726ba4b7..8f26477ed115e 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 @@ -392,58 +392,54 @@ private SendMessagesResult sendMessagesInternal(final List entr 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) { - totalEntries--; - 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 stickyKeyHash; - if (stickyKeyHashes == null) { - if (entry instanceof EntryAndMetadata entryAndMetadata) { - stickyKeyHash = entryAndMetadata.getCachedStickyKeyHash(); - } else { - stickyKeyHash = STICKY_KEY_HASH_NOT_SET; - } - } else { - stickyKeyHash = stickyKeyHashes.get(i); - } - boolean sendingAllowed = - pendingAcks.addPendingAckIfAllowed(entry.getLedgerId(), entry.getEntryId(), - 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() - .attr("ledgerId", entry.getLedgerId()) - .attr("entryId", entry.getEntryId()) - .attr("batchSize", batchSize) - .log("Skipping sending of entry since adding to pending acks failed"); + if (entry == null) { + continue; + } + 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 stickyKeyHash; + if (stickyKeyHashes == null) { + if (entry instanceof EntryAndMetadata entryAndMetadata) { + stickyKeyHash = entryAndMetadata.getCachedStickyKeyHash(); } else { - log.debug() - .attr("ledgerId", entry.getLedgerId()) - .attr("entryId", entry.getEntryId()) - .attr("batchSize", batchSize) - .log("Added entry to pendingAcks"); - sendResult.setMessagePermits(i, messagePermits); + stickyKeyHash = STICKY_KEY_HASH_NOT_SET; } } else { - sendResult.setMessagePermits(i, messagePermits); + stickyKeyHash = stickyKeyHashes.get(i); + } + boolean sendingAllowed = pendingAcks.addPendingAckIfAllowed( + entry.getLedgerId(), entry.getEntryId(), 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 + entries.set(i, null); + entry.release(); + log.debug() + .attr("ledgerId", entry.getLedgerId()) + .attr("entryId", entry.getEntryId()) + .attr("batchSize", batchSize) + .log("Skipping sending of entry since adding to pending acks failed"); + continue; } + log.debug() + .attr("ledgerId", entry.getLedgerId()) + .attr("entryId", entry.getEntryId()) + .attr("batchSize", batchSize) + .log("Added entry to pendingAcks"); } + sendResult.setMessagePermits(i, messagePermits); + totalEntries++; } // calculate avg message per entry diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java index d70c0ca1306d1..26089afcf2410 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java @@ -46,8 +46,9 @@ void setMessagePermits(int entryIdx, int permits) { if (messagePermits[entryIdx] != 0) { throw new IllegalStateException("Permits already finalized for entry " + entryIdx); } + int updatedTotalMessagePermits = Math.addExact(totalMessagePermits, permits); messagePermits[entryIdx] = permits; - totalMessagePermits = Math.addExact(totalMessagePermits, permits); + totalMessagePermits = updatedTotalMessagePermits; } int getMessagePermits(int entryIdx) { 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 851c0e02d0092..1411465e4d800 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 @@ -4487,8 +4487,9 @@ private TopicName validateTopicName(String topic, long requestId, Object request public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId, int partition, int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet, String topic, long epoch) { - return newMessageAndIntercept(consumerId, ledgerId, entryId, partition, redeliveryCount, - metadataAndPayload, ackSet, topic, epoch, 1); + BaseCommand command = Commands.newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, + ackSet, epoch); + return serializeMessageAndIntercept(command, consumerId, ledgerId, entryId, metadataAndPayload); } public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId, int partition, @@ -4496,6 +4497,11 @@ public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long e int messagePermits) { BaseCommand command = Commands.newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, epoch, messagePermits); + return serializeMessageAndIntercept(command, consumerId, ledgerId, entryId, metadataAndPayload); + } + + private ByteBufPair serializeMessageAndIntercept(BaseCommand command, long consumerId, long ledgerId, long entryId, + ByteBuf metadataAndPayload) { ByteBufPair res = Commands.serializeCommandMessageWithSize(command, metadataAndPayload); if (brokerInterceptor != null) { try { 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 e9a4737c83c55..c10cfd8c9e084 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 @@ -35,6 +35,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertSame; +import static org.testng.Assert.expectThrows; import io.netty.util.concurrent.ImmediateEventExecutor; import java.net.SocketAddress; import java.util.ArrayList; @@ -131,6 +132,9 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { entries, batchSizes, batchIndexesAcks, 23, 0, 0, mock(RedeliveryTracker.class)); assertEquals(sendResult.getTotalMessagePermits(), 3); + assertEquals(sendResult.getMessagePermits(0), 3); + assertEquals(sendResult.getMessagePermits(1), 0); + assertEquals(sendResult.getMessagePermits(2), 0); assertEquals(sharedConsumer.getAvailablePermits(), 97); assertEquals(sharedConsumer.getUnackedMessages(), 3); assertNull(entries.get(1)); @@ -148,4 +152,17 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { batchIndexesAcks.recycle(); } } + + @Test + public void testSendMessagesResultRejectsPartialFinalization() { + SendMessagesResult sendResult = new SendMessagesResult(2); + sendResult.setMessagePermits(0, Integer.MAX_VALUE); + + expectThrows(ArithmeticException.class, () -> sendResult.setMessagePermits(1, 1)); + assertEquals(sendResult.getMessagePermits(0), Integer.MAX_VALUE); + assertEquals(sendResult.getMessagePermits(1), 0); + assertEquals(sendResult.getTotalMessagePermits(), Integer.MAX_VALUE); + expectThrows(IllegalStateException.class, () -> sendResult.setMessagePermits(0, 1)); + expectThrows(IllegalArgumentException.class, () -> new SendMessagesResult(1).setMessagePermits(0, 0)); + } } 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..2b755e50ddfd8 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,113 @@ 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(); + + 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)); + + 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); + 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 45a726618b31e..ba4d86ce1c5d4 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 @@ -21,7 +21,6 @@ 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; @@ -97,44 +96,6 @@ public void testStaleEpochDiscardDoesNotReturnPermitToReplacementConnection() th } } - @Test - public void testPartialBatchDeserializationReturnsOnlyCommandRemainder() throws Exception { - final int batchSize = 5; - final int parsedMessages = 2; - 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(newTopicName()) - .subscriptionName("partial-batch-subscription") - .receiverQueueSize(20) - .subscribe()) { - int permitsBefore = consumer.getAvailablePermits(); - 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(), permitsBefore + batchSize); - } finally { - truncatedBatch.release(); - } - } - @Test public void testCompactedOutMessages() throws Exception { final String topic1 = newTopicName(); @@ -165,6 +126,7 @@ public void testCompactedOutMessages() throws Exception { try (ConsumerImpl consumer = (ConsumerImpl) pulsarClient.newConsumer().topic(topic1) .subscriptionName("my-subscriber-name").subscribe()) { + int permitsBefore = consumer.getAvailablePermits(); // shove it in the sideways consumer.receiveIndividualMessagesFromBatch(brokerEntryMetadata, metadata, 0, null, batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), @@ -176,6 +138,8 @@ batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), assertEquals(m.getKey(), "key3"); assertEquals(consumer.numMessagesInQueue(), 0); + assertEquals(consumer.getAvailablePermits(), permitsBefore + metadata.getNumMessagesInBatch()); + m.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 index 94d8efe40af7b..242165d2495b2 100644 --- 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 @@ -18,6 +18,7 @@ */ package org.apache.pulsar.client.impl; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import io.netty.buffer.ByteBuf; @@ -35,6 +36,7 @@ import org.apache.pulsar.client.api.Schema; import org.apache.pulsar.client.api.SubscriptionType; 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.awaitility.Awaitility; import org.testng.annotations.Test; @@ -43,8 +45,9 @@ public class CorruptedBatchMessagePermitTest extends SharedPulsarBaseTest { @Test(timeOut = 30000) - public void testSharedConsumerReturnsAllPermitsForCorruptedBatch() throws Exception { + public void testSharedConsumerReturnsAllPermitsForRepeatedCorruptedBatches() throws Exception { final int batchSize = 10; + final int corruptedBatches = 3; final String subscriptionName = "shared-subscription"; final String topicName = newTopicName(); admin.namespaces().setDeduplicationStatus(getNamespace(), false); @@ -65,27 +68,82 @@ public void testSharedConsumerReturnsAllPermitsForCorruptedBatch() throws Except PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get() .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); PersistentSubscription subscription = topic.getSubscription(subscriptionName); - publishCorruptedBatch(topic, batchSize).get(10, TimeUnit.SECONDS); + for (int i = 1; i <= corruptedBatches; i++) { + publishCorruptedBatch(topic, batchSize, i).get(10, TimeUnit.SECONDS); + awaitConsumerPermitBalance(subscription, (long) batchSize * i); + } - Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { - assertEquals(subscription.getConsumers().size(), 1); - assertEquals(subscription.getConsumers().get(0).getAvailablePermits(), 1); - }); + producer.send("message-after-corrupted-batches"); + Message message = consumer.receive(10, TimeUnit.SECONDS); + assertNotNull(message); + try { + assertEquals(message.getValue(), "message-after-corrupted-batches"); + } finally { + message.release(); + } + } + + @Test(timeOut = 30000) + public void testSharedConsumerReturnsOnlyUnparsedRemainderForMalformedBatch() throws Exception { + final int batchSize = 5; + final int parsedMessages = 2; + final String subscriptionName = "shared-partial-batch-subscription"; + final String topicName = newTopicName(); + admin.namespaces().setDeduplicationStatus(getNamespace(), false); + + @Cleanup + Consumer consumer = pulsarClient.newConsumer(Schema.STRING) + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .subscribe(); + @Cleanup + Producer producer = pulsarClient.newProducer(Schema.STRING) + .topic(topicName) + .enableBatching(false) + .create(); - producer.send("message-after-corrupted-batch"); + PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get() + .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); + PersistentSubscription subscription = topic.getSubscription(subscriptionName); + publishMalformedBatch(topic, batchSize, parsedMessages).get(10, TimeUnit.SECONDS); + + for (int i = 0; i < parsedMessages; i++) { + Message message = consumer.receive(10, TimeUnit.SECONDS); + assertNotNull(message); + try { + assertEquals(message.getValue(), "partial-" + i); + } finally { + message.release(); + } + } + awaitConsumerPermitBalance(subscription, batchSize); + + producer.send("message-after-malformed-batch"); Message message = consumer.receive(10, TimeUnit.SECONDS); assertNotNull(message); try { - assertEquals(message.getValue(), "message-after-corrupted-batch"); + assertEquals(message.getValue(), "message-after-malformed-batch"); } finally { message.release(); } } - private static CompletableFuture publishCorruptedBatch(PersistentTopic topic, int batchSize) { + private static void awaitConsumerPermitBalance(PersistentSubscription subscription, long expectedMessagesOut) { + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + assertEquals(subscription.getConsumers().size(), 1); + org.apache.pulsar.broker.service.Consumer brokerConsumer = subscription.getConsumers().get(0); + assertEquals(brokerConsumer.getMsgOutCounter(), expectedMessagesOut); + assertEquals(brokerConsumer.getAvailablePermits(), 1); + }); + } + + private static CompletableFuture publishCorruptedBatch(PersistentTopic topic, int batchSize, + long sequenceId) { MessageMetadata metadata = new MessageMetadata() .setProducerName("corrupted-batch-producer") - .setSequenceId(0) + .setSequenceId(sequenceId) .setPublishTime(System.currentTimeMillis()) .setNumMessagesInBatch(batchSize); ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); @@ -94,6 +152,29 @@ private static CompletableFuture publishCorruptedBatch(PersistentTopic top int lastByteIndex = entry.writerIndex() - 1; entry.setByte(lastByteIndex, entry.getByte(lastByteIndex) ^ 1); + return publish(topic, entry); + } + + private static CompletableFuture publishMalformedBatch(PersistentTopic topic, int batchSize, + int parsedMessages) { + MessageMetadata metadata = new MessageMetadata() + .setProducerName("malformed-batch-producer") + .setSequenceId(0) + .setPublishTime(System.currentTimeMillis()) + .setNumMessagesInBatch(batchSize); + ByteBuf payload = Unpooled.buffer(); + for (int i = 0; i < parsedMessages; i++) { + ByteBuf singlePayload = Unpooled.copiedBuffer("partial-" + i, UTF_8); + Commands.serializeSingleMessageInBatchWithPayload( + new SingleMessageMetadata(), singlePayload, payload); + singlePayload.release(); + } + ByteBuf entry = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c, metadata, payload); + payload.release(); + return publish(topic, entry); + } + + private static CompletableFuture publish(PersistentTopic topic, ByteBuf entry) { CompletableFuture result = new CompletableFuture<>(); topic.publishMessage(entry, (error, ledgerId, entryId) -> { if (error == null) { 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 ddc255926828d..1ee1b633147e8 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 @@ -98,6 +98,8 @@ import org.apache.pulsar.client.api.TopicMessageId; import org.apache.pulsar.client.api.TypedMessageBuilder; import org.apache.pulsar.client.api.transaction.TxnID; +import org.apache.pulsar.client.impl.MessagePermitAccounting.Budget; +import org.apache.pulsar.client.impl.MessagePermitAccounting.InvalidMessagePermitsException; import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData; import org.apache.pulsar.client.impl.crypto.MessageCryptoBc; import org.apache.pulsar.client.impl.metrics.Counter; @@ -147,37 +149,6 @@ public class ConsumerImpl extends ConsumerBase implements ConnectionHandle private static final long[] EMPTY_ACK_SET = new long[0]; private static final int MAX_REDELIVER_UNACKNOWLEDGED = 1000; - private static final class MessagePermitBudget { - private int remainingPermits; - - private MessagePermitBudget(int messagePermits) { - remainingPermits = messagePermits; - } - - private void claim() { - if (remainingPermits == 0) { - throw new InvalidMessagePermitsException("Batch contains more deliverable messages than permits"); - } - remainingPermits--; - } - - private void restore() { - remainingPermits++; - } - - private int drain() { - int permits = remainingPermits; - remainingPermits = 0; - return permits; - } - } - - private static final class InvalidMessagePermitsException extends IllegalStateException { - private InvalidMessagePermitsException(String message) { - super(message); - } - } - final long consumerId; // Number of messages that have delivered to the application. Every once in a while, this number will be sent to the @@ -1446,7 +1417,7 @@ private void executeNotifyCallback(final MessageImpl message) { // if asyncReceive is waiting then notify callback without adding to incomingMessages queue internalPinnedExecutor.execute(() -> { if (!isValidConsumerEpoch(message)) { - completeStaleEpochDiscard(messageCnx, messageSize); + completeStaleEpochDiscard(message, messageCnx, messageSize); return; } Message interceptMsg = onArrival(message); @@ -1458,10 +1429,11 @@ private void executeNotifyCallback(final MessageImpl message) { }); } - private void completeStaleEpochDiscard(ClientCnx messageCnx, int messageSize) { + private void completeStaleEpochDiscard(MessageImpl message, ClientCnx messageCnx, int messageSize) { messagesPrefetchedGauge.decrement(); bytesPrefetchedGauge.subtract(messageSize); - increaseAvailablePermitsForReceivedMessage(messageCnx, 1); + returnPermitsToSourceConnection(messageCnx, 1); + message.release(); } protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, @@ -1500,58 +1472,6 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe tryTriggerListener(); } - @VisibleForTesting - static int getMessagePermitsForEarlyFailure(CommandMessage command) { - if (command.hasMessagePermits()) { - return getExplicitMessagePermits(command); - } - return command.getAckSetsCount() > 0 ? getAckSetCardinality(command, -1) : 1; - } - - @VisibleForTesting - static int getMessagePermits(CommandMessage command, int batchSize) { - int expectedPermits = command.getAckSetsCount() > 0 - ? getAckSetCardinality(command, batchSize) : batchSize; - if (!command.hasMessagePermits()) { - return expectedPermits; - } - - int explicitPermits = getExplicitMessagePermits(command); - if (explicitPermits != expectedPermits) { - throw new InvalidMessagePermitsException("Explicit message permits " + explicitPermits - + " do not match the payload and ack set value " + expectedPermits); - } - return explicitPermits; - } - - private static int getExplicitMessagePermits(CommandMessage command) { - long messagePermits = Integer.toUnsignedLong(command.getMessagePermits()); - if (messagePermits == 0 || messagePermits > Integer.MAX_VALUE) { - throw new InvalidMessagePermitsException("Invalid explicit message permits " + messagePermits); - } - return (int) messagePermits; - } - - private static int getAckSetCardinality(CommandMessage command, int batchSize) { - int words = command.getAckSetsCount(); - int completeWords = batchSize < 0 ? words : Math.min(batchSize >>> 6, words); - long cardinality = 0; - for (int i = 0; i < completeWords; i++) { - cardinality += Long.bitCount(command.getAckSetAt(i)); - } - if (batchSize >= 0) { - int remainingBits = batchSize & 63; - if (remainingBits > 0 && completeWords < words) { - long mask = -1L >>> (Long.SIZE - remainingBits); - cardinality += Long.bitCount(command.getAckSetAt(completeWords) & mask); - } - } - if (cardinality > Integer.MAX_VALUE) { - throw new InvalidMessagePermitsException("Ack set message permits exceed the supported range"); - } - return (int) cardinality; - } - void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, ClientCnx cnx) { long[] ackSet = EMPTY_ACK_SET; if (cmdMessage.getAckSetsCount() > 0) { @@ -1569,7 +1489,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien } final int earlyFailurePermits; try { - earlyFailurePermits = getMessagePermitsForEarlyFailure(cmdMessage); + earlyFailurePermits = MessagePermitAccounting.resolveForEarlyFailure(cmdMessage); } catch (InvalidMessagePermitsException e) { closeConnectionForInvalidMessagePermits(messageId, cnx, e); return; @@ -1577,8 +1497,13 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien log.debug().attr("messageId", messageId) .log("Received message"); - if (!verifyChecksum(headersAndPayload, messageId)) { - // discard message with checksum error + try { + if (!verifyChecksum(headersAndPayload, messageId)) { + // discard message with checksum error + discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); + return; + } + } catch (Throwable t) { discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); return; } @@ -1596,7 +1521,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien final int numMessages = msgMetadata.getNumMessagesInBatch(); final int messagePermits; try { - messagePermits = getMessagePermits(cmdMessage, numMessages); + messagePermits = MessagePermitAccounting.resolve(cmdMessage, numMessages); } catch (InvalidMessagePermitsException e) { closeConnectionForInvalidMessagePermits(messageId, cnx, e); return; @@ -1609,7 +1534,7 @@ 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"); - increaseAvailablePermitsForReceivedMessage(cnx, messagePermits); + returnPermitsToSourceConnection(cnx, messagePermits); return; } @@ -1942,8 +1867,7 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } SingleMessageMetadata singleMessageMetadata = new SingleMessageMetadata(); - MessagePermitBudget permitBudget = new MessagePermitBudget(messagePermits); - boolean corrupted = false; + Budget permitBudget = new Budget(messagePermits); try { for (int i = 0; i < batchSize; ++i) { final MessageImpl message = newSingleMessage(i, batchSize, brokerEntryMetadata, msgMetadata, @@ -1991,7 +1915,6 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, log.warn().exception(e) .log("unable to obtain message in batch"); discardCorruptedMessage(messageId, cnx, ValidationError.BatchDeSerializeError, permitBudget.drain()); - corrupted = true; } finally { if (ackBitSet != null) { ackBitSet.recycle(); @@ -2012,9 +1935,7 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, .attr("remainingCapacity", incomingMessages.remainingCapacity()) .log("enqueued messages in batch. queue size -, available queue size"); - if (!corrupted) { - increaseAvailablePermitsForReceivedMessage(cnx, permitBudget.drain()); - } + returnPermitsToSourceConnection(cnx, permitBudget.drain()); } private boolean isPriorEntryIndex(long idx) { @@ -2085,10 +2006,10 @@ protected void trackMessage(MessageId messageId, int redeliveryCount) { } void increaseAvailablePermits(MessageImpl msg) { - increaseAvailablePermitsForReceivedMessage(msg.getCnx(), 1); + returnPermitsToSourceConnection(msg.getCnx(), 1); } - private void increaseAvailablePermitsForReceivedMessage(ClientCnx messageCnx, int delta) { + private void returnPermitsToSourceConnection(ClientCnx messageCnx, int delta) { if (delta <= 0) { return; } @@ -2371,7 +2292,7 @@ private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, Valid ByteBuf cmd = Commands.newAck(consumerId, messageId.getLedgerId(), messageId.getEntryId(), null, AckType.Individual, validationError, Collections.emptyMap(), -1); currentCnx.ctx().writeAndFlush(cmd, currentCnx.ctx().voidPromise()); - increaseAvailablePermitsForReceivedMessage(currentCnx, messagePermits); + returnPermitsToSourceConnection(currentCnx, messagePermits); stats.incrementNumReceiveFailed(); } diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java new file mode 100644 index 0000000000000..416399e0d92ae --- /dev/null +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java @@ -0,0 +1,111 @@ +/* + * 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 org.apache.pulsar.common.api.proto.CommandMessage; + +/** + * Resolves the permit count represented by a message command and tracks that permit budget while a batch is decoded. + */ +final class MessagePermitAccounting { + + private MessagePermitAccounting() { + } + + static int resolveForEarlyFailure(CommandMessage command) { + if (command.hasMessagePermits()) { + return getExplicitMessagePermits(command); + } + return command.getAckSetsCount() > 0 ? getAckSetCardinality(command, -1) : 1; + } + + static int resolve(CommandMessage command, int batchSize) { + int expectedPermits = command.getAckSetsCount() > 0 + ? getAckSetCardinality(command, batchSize) : batchSize; + if (!command.hasMessagePermits()) { + return expectedPermits; + } + + int explicitPermits = getExplicitMessagePermits(command); + if (explicitPermits != expectedPermits) { + throw new InvalidMessagePermitsException("Explicit message permits " + explicitPermits + + " do not match the payload and ack set value " + expectedPermits); + } + return explicitPermits; + } + + private static int getExplicitMessagePermits(CommandMessage command) { + long messagePermits = Integer.toUnsignedLong(command.getMessagePermits()); + if (messagePermits == 0 || messagePermits > Integer.MAX_VALUE) { + throw new InvalidMessagePermitsException("Invalid explicit message permits " + messagePermits); + } + return (int) messagePermits; + } + + private static int getAckSetCardinality(CommandMessage command, int batchSize) { + int words = command.getAckSetsCount(); + int completeWords = batchSize < 0 ? words : Math.min(batchSize >>> 6, words); + long cardinality = 0; + for (int i = 0; i < completeWords; i++) { + cardinality += Long.bitCount(command.getAckSetAt(i)); + } + if (batchSize >= 0) { + int remainingBits = batchSize & 63; + if (remainingBits > 0 && completeWords < words) { + long mask = -1L >>> (Long.SIZE - remainingBits); + cardinality += Long.bitCount(command.getAckSetAt(completeWords) & mask); + } + } + if (cardinality > Integer.MAX_VALUE) { + throw new InvalidMessagePermitsException("Ack set message permits exceed the supported range"); + } + return (int) cardinality; + } + + static final class Budget { + private int remainingPermits; + + Budget(int messagePermits) { + remainingPermits = messagePermits; + } + + void claim() { + if (remainingPermits == 0) { + throw new InvalidMessagePermitsException("Batch contains more deliverable messages than permits"); + } + remainingPermits--; + } + + void restore() { + remainingPermits++; + } + + int drain() { + int permits = remainingPermits; + remainingPermits = 0; + return permits; + } + } + + static final class InvalidMessagePermitsException extends IllegalStateException { + InvalidMessagePermitsException(String message) { + super(message); + } + } +} 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 06684fa74d45e..fd39d3e4093a0 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 @@ -34,6 +34,8 @@ import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPromise; +import io.netty.util.ReferenceCountUtil; import java.util.Arrays; import java.util.BitSet; import java.util.List; @@ -60,9 +62,12 @@ import org.apache.pulsar.client.util.ExecutorProvider; import org.apache.pulsar.client.util.ScheduledExecutorProvider; import org.apache.pulsar.common.api.proto.CommandMessage; +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.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; @@ -117,40 +122,6 @@ public void testNotifyPendingReceivedCallback_EmptyQueueNotThrowsException() { consumer.notifyPendingReceivedCallback(null, null); } - @Test - public void testResolveExplicitMessagePermits() { - CommandMessage command = new CommandMessage().setMessagePermits(3); - command.addAckSet(0b100101L); - - Assert.assertEquals(ConsumerImpl.getMessagePermitsForEarlyFailure(command), 3); - Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 3); - } - - @Test - public void testResolveLegacyMessagePermits() { - CommandMessage partialBatch = new CommandMessage(); - partialBatch.addAckSet((1L << 1) | (1L << 4) | (1L << 63)); - - Assert.assertEquals(ConsumerImpl.getMessagePermitsForEarlyFailure(partialBatch), 3); - Assert.assertEquals(ConsumerImpl.getMessagePermits(partialBatch, 10), 2); - Assert.assertEquals(ConsumerImpl.getMessagePermits(new CommandMessage(), 10), 10); - } - - @Test - public void testRejectInvalidExplicitMessagePermits() { - CommandMessage zero = new CommandMessage().setMessagePermits(0); - CommandMessage unsignedOverflow = new CommandMessage().setMessagePermits(-1); - CommandMessage mismatch = new CommandMessage().setMessagePermits(2); - mismatch.addAckSet(0b111L); - - Assert.expectThrows(IllegalStateException.class, - () -> ConsumerImpl.getMessagePermitsForEarlyFailure(zero)); - Assert.expectThrows(IllegalStateException.class, - () -> ConsumerImpl.getMessagePermitsForEarlyFailure(unsignedOverflow)); - Assert.expectThrows(IllegalStateException.class, - () -> ConsumerImpl.getMessagePermits(mismatch, 10)); - } - @Test public void testInvalidExplicitMessagePermitsCloseSourceConnection() { CommandMessage command = new CommandMessage() @@ -172,6 +143,102 @@ public void testInvalidExplicitMessagePermitsCloseSourceConnection() { } } + @Test + public void testTruncatedFrameAndMetadataFailureReturnExplicitMessagePermits() { + ClientCnx messageCnx = setCurrentConnection(); + ChannelHandlerContext context = messageCnx.ctx(); + int permitsBefore = consumer.getAvailablePermits(); + ByteBuf[] malformedFrames = { + Unpooled.wrappedBuffer(new byte[] {1}), + Unpooled.wrappedBuffer(new byte[] {0, 0, 0, 10}) + }; + + try { + for (ByteBuf malformedFrame : malformedFrames) { + consumer.messageReceived(newCommandMessage(5), malformedFrame, messageCnx); + } + + Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + 10); + releaseValidationCommands(context, 2); + } finally { + Arrays.stream(malformedFrames).forEach(ByteBuf::release); + } + } + + @Test + public void testDecompressionFailureReturnsExplicitMessagePermits() { + ClientCnx messageCnx = setCurrentConnection(); + ChannelHandlerContext context = messageCnx.ctx(); + MessageMetadata metadata = new MessageMetadata() + .setProducerName("producer") + .setSequenceId(1) + .setPublishTime(1) + .setNumMessagesInBatch(5) + .setCompression(CompressionType.ZLIB) + .setUncompressedSize(100); + ByteBuf invalidCompressedPayload = Unpooled.wrappedBuffer(new byte[] {1}); + ByteBuf metadataAndPayload = Commands.serializeMetadataAndPayload( + Commands.ChecksumType.Crc32c, metadata, invalidCompressedPayload); + invalidCompressedPayload.release(); + int permitsBefore = consumer.getAvailablePermits(); + + try { + consumer.messageReceived(newCommandMessage(5), metadataAndPayload, messageCnx); + + Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + 5); + releaseValidationCommands(context, 1); + } finally { + metadataAndPayload.release(); + } + } + + @Test + public void testPermitAccumulatorOverflowClosesSourceConnection() { + ClientCnx messageCnx = setCurrentConnection(); + ChannelHandlerContext context = messageCnx.ctx(); + consumer.paused = true; + ByteBuf firstMalformedMetadata = Unpooled.wrappedBuffer(new byte[] {1}); + ByteBuf secondMalformedMetadata = Unpooled.wrappedBuffer(new byte[] {1}); + + try { + consumer.messageReceived(newCommandMessage(Integer.MAX_VALUE), firstMalformedMetadata, messageCnx); + Assert.assertEquals(consumer.getAvailablePermits(), Integer.MAX_VALUE); + + consumer.messageReceived(newCommandMessage(1), secondMalformedMetadata, messageCnx); + + Assert.assertEquals(consumer.getAvailablePermits(), Integer.MAX_VALUE); + verify(context).close(); + releaseValidationCommands(context, 2); + } finally { + firstMalformedMetadata.release(); + secondMalformedMetadata.release(); + } + } + + private ClientCnx setCurrentConnection() { + ClientCnx messageCnx = mock(ClientCnx.class); + ChannelHandlerContext context = mock(ChannelHandlerContext.class); + when(context.voidPromise()).thenReturn(mock(ChannelPromise.class)); + when(messageCnx.ctx()).thenReturn(context); + consumer.setClientCnx(messageCnx); + return messageCnx; + } + + private CommandMessage newCommandMessage(int messagePermits) { + CommandMessage command = new CommandMessage() + .setConsumerId(consumer.consumerId) + .setMessagePermits(messagePermits); + command.setMessageId().setLedgerId(1).setEntryId(2); + return command; + } + + private static void releaseValidationCommands(ChannelHandlerContext context, int expectedCommands) { + ArgumentCaptor commandCaptor = ArgumentCaptor.forClass(Object.class); + verify(context, times(expectedCommands)).writeAndFlush( + commandCaptor.capture(), any(ChannelPromise.class)); + commandCaptor.getAllValues().forEach(ReferenceCountUtil::release); + } + @Test(invocationTimeOut = 500) public void testCorrectBackoffConfiguration() { final Backoff backoff = consumer.getConnectionHandler().backoff; diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java new file mode 100644 index 0000000000000..a73e4e6f3390e --- /dev/null +++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java @@ -0,0 +1,84 @@ +/* + * 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.expectThrows; +import org.apache.pulsar.client.impl.MessagePermitAccounting.Budget; +import org.apache.pulsar.common.api.proto.CommandMessage; +import org.testng.annotations.Test; + +public class MessagePermitAccountingTest { + + @Test + public void testResolveExplicitMessagePermits() { + CommandMessage command = new CommandMessage().setMessagePermits(3); + command.addAckSet(0b100101L); + + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(command), 3); + assertEquals(MessagePermitAccounting.resolve(command, 10), 3); + + CommandMessage partialBatchWithOnePermit = new CommandMessage().setMessagePermits(1); + partialBatchWithOnePermit.addAckSet(0b100L); + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(partialBatchWithOnePermit), 1); + assertEquals(MessagePermitAccounting.resolve(partialBatchWithOnePermit, 10), 1); + + CommandMessage singleMessage = new CommandMessage().setMessagePermits(1); + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(singleMessage), 1); + assertEquals(MessagePermitAccounting.resolve(singleMessage, 1), 1); + } + + @Test + public void testResolveLegacyMessagePermits() { + CommandMessage partialBatch = new CommandMessage(); + partialBatch.addAckSet((1L << 1) | (1L << 4) | (1L << 63)); + + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(new CommandMessage()), 1); + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(partialBatch), 3); + assertEquals(MessagePermitAccounting.resolve(partialBatch, 10), 2); + assertEquals(MessagePermitAccounting.resolve(new CommandMessage(), 10), 10); + } + + @Test + public void testRejectInvalidExplicitMessagePermits() { + CommandMessage zero = new CommandMessage().setMessagePermits(0); + CommandMessage unsignedOverflow = new CommandMessage().setMessagePermits(-1); + CommandMessage mismatch = new CommandMessage().setMessagePermits(2); + mismatch.addAckSet(0b111L); + CommandMessage onePermitMismatch = new CommandMessage().setMessagePermits(1); + + expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolveForEarlyFailure(zero)); + expectThrows(IllegalStateException.class, + () -> MessagePermitAccounting.resolveForEarlyFailure(unsignedOverflow)); + expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolve(mismatch, 10)); + expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolve(onePermitMismatch, 10)); + } + + @Test + public void testBudgetTracksClaimsAndRestores() { + Budget budget = new Budget(2); + budget.claim(); + budget.restore(); + budget.claim(); + + assertEquals(budget.drain(), 1); + assertEquals(budget.drain(), 0); + expectThrows(IllegalStateException.class, budget::claim); + } +} 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 f1ed8b0e45a4e..e1baa2c027e9b 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,14 +537,6 @@ 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, 1); - } - - 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("messagePermits must be positive"); - } BaseCommand cmd = localCmd(Type.MESSAGE); CommandMessage msg = cmd.setMessage() .setConsumerId(consumerId); @@ -560,9 +552,6 @@ public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long if (redeliveryCount > 0) { msg.setRedeliveryCount(redeliveryCount); } - if (messagePermits > 1) { - msg.setMessagePermits(messagePermits); - } if (ackSet != null) { for (int i = 0; i < ackSet.length; i++) { msg.addAckSet(ackSet[i]); @@ -571,6 +560,17 @@ public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long return cmd; } + 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("messagePermits must be positive"); + } + BaseCommand command = newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, + consumerEpoch); + command.getMessage().setMessagePermits(messagePermits); + return command; + } + public static ByteBufPair newMessage(long consumerId, long ledgerId, long entryId, int partition, int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet) { return serializeCommandMessageWithSize( diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto index c5acf55c2dc33..ae4f06747716f 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -583,7 +583,7 @@ message CommandMessage { optional uint64 consumer_epoch = 5; // Number of consumer permits debited by the broker for this command. - optional uint32 message_permits = 6 [default = 1]; + 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 067c284a4bf24..97b05ec09856a 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 @@ -47,9 +47,12 @@ public void testMessageCommandCarriesPermitDebit() { assertEquals(batched.getMessage().getMessagePermits(), 10); BaseCommand single = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 1); - Assert.assertFalse(single.getMessage().hasMessagePermits()); + assertTrue(single.getMessage().hasMessagePermits()); assertEquals(single.getMessage().getMessagePermits(), 1); + BaseCommand legacy = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1); + Assert.assertFalse(legacy.getMessage().hasMessagePermits()); + Assert.expectThrows(IllegalArgumentException.class, () -> Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 0)); } From 169fbb6b375b2dbcc0245d99627ea6d4fad48aa6 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 16 Aug 2026 12:56:45 +0800 Subject: [PATCH 4/8] [fix][client][broker] Harden explicit permit accounting lifecycle Assisted-by: Codex/GPT-5.6 --- .../pulsar/broker/service/Consumer.java | 3 +- .../broker/service/EntryBatchIndexesAcks.java | 3 +- .../pulsar/broker/service/ConsumerTest.java | 34 ++ .../pulsar/client/impl/ConsumerImpl.java | 334 ++++++++++++------ .../pulsar/client/impl/MessageImpl.java | 12 + .../client/impl/MultiTopicsConsumerImpl.java | 4 +- .../client/impl/ZeroQueueConsumerImpl.java | 25 +- .../pulsar/client/impl/ConsumerImplTest.java | 126 +++++++ 8 files changed, 423 insertions(+), 118 deletions(-) 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 8f26477ed115e..a8f098bd12313 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 @@ -478,7 +478,8 @@ private SendMessagesResult sendMessagesInternal(final List entr } else { log.debug() .exceptionMessage(status.cause()) - .log("Sent messages to client failed by IO exception, closing the connection"); + .log("Sent messages to client failed by IO exception, disconnecting the consumer"); + disconnect(); } }); return sendResult.setWriteFuture(writeAndFlushPromise); 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 9afd38ab4b7a3..c38b0569af4fe 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,6 +20,7 @@ import io.netty.util.Recycler; +import java.util.BitSet; import org.apache.commons.lang3.tuple.Pair; @SuppressWarnings("unchecked") @@ -46,7 +47,7 @@ public int getTotalAckedIndexCount() { for (int i = 0; i < size; i++) { Pair pair = indexesAcks[i]; if (pair != null) { - count += pair.getLeft() - getUnackedIndexCount(i, pair.getLeft()); + count += pair.getLeft() - BitSet.valueOf(pair.getRight()).cardinality(); } } return count; 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 c10cfd8c9e084..02a529ecedb2b 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 @@ -37,9 +37,11 @@ import static org.testng.Assert.assertSame; import static org.testng.Assert.expectThrows; import io.netty.util.concurrent.ImmediateEventExecutor; +import java.io.IOException; import java.net.SocketAddress; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import org.apache.bookkeeper.mledger.Entry; import org.apache.commons.lang3.tuple.Pair; import org.apache.pulsar.broker.PulsarService; @@ -153,6 +155,38 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { } } + @Test + public void testSendMessagesWriteFailureDisconnectsConsumer() throws BrokerServiceException { + Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, + "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, + DEFAULT_CONSUMER_EPOCH); + sharedConsumer.setPendingAcksAddHandler((ignored, ledgerId, entryId, stickyKeyHash) -> true); + sharedConsumer.flowPermits(100); + + Entry entry = mock(Entry.class); + when(entry.getLedgerId()).thenReturn(1L); + when(entry.getEntryId()).thenReturn(1L); + List entries = new ArrayList<>(List.of(entry)); + EntryBatchSizes batchSizes = EntryBatchSizes.get(1); + batchSizes.setBatchSize(0, 1); + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); + when(cnx.getCommandSender()).thenReturn(commandSender); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), + any(), any(), anyLong())) + .thenReturn(ImmediateEventExecutor.INSTANCE.newFailedFuture(new IOException("write failed"))); + + try { + sharedConsumer.sendMessagesWithResult( + entries, batchSizes, null, 1, 0, 0, mock(RedeliveryTracker.class)); + + verify(cnx).closeConsumer(sharedConsumer, Optional.empty()); + verify(subscription).removeConsumer(sharedConsumer, false); + verify(cnx).removedConsumer(sharedConsumer); + } finally { + batchSizes.recyle(); + } + } + @Test public void testSendMessagesResultRejectsPartialFinalization() { SendMessagesResult sendResult = new SendMessagesResult(2); 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 1ee1b633147e8..49c11590744b7 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 @@ -151,13 +151,23 @@ 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 + // Permit returns are scoped to one local broker-consumer incarnation. ClientCnx cannot be used as the identity + // because a pooled physical connection can outlive a broker-side consumer and then be reused by its replacement. + private volatile ConsumerPermitState permitState = new ConsumerPermitState(false); + @SuppressWarnings("rawtypes") - private static final AtomicIntegerFieldUpdater AVAILABLE_PERMITS_UPDATER = AtomicIntegerFieldUpdater - .newUpdater(ConsumerImpl.class, "availablePermits"); - @SuppressWarnings("unused") - private volatile int availablePermits = 0; + private static final AtomicIntegerFieldUpdater AVAILABLE_PERMITS_UPDATER = + AtomicIntegerFieldUpdater.newUpdater(ConsumerPermitState.class, "availablePermits"); + + static final class ConsumerPermitState { + private volatile boolean flowEnabled; + @SuppressWarnings("unused") + private volatile int availablePermits; + + private ConsumerPermitState(boolean flowEnabled) { + this.flowEnabled = flowEnabled; + } + } protected volatile MessageId lastDequeuedMessageId = MessageId.earliest; private volatile MessageId lastMessageIdInBroker = MessageId.earliest; @@ -333,7 +343,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; @@ -898,9 +907,11 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { + client.getConfiguration().getOperationTimeoutMs()); int currentSize; + ConsumerPermitState connectionPermitState; synchronized (this) { currentSize = incomingMessages.size(); setClientCnx(cnx); + connectionPermitState = permitState; clearReceiverQueue(true); if (possibleSendToDeadLetterTopicMessages != null) { possibleSendToDeadLetterTopicMessages.clear(); @@ -950,7 +961,7 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { cnx.sendRequestWithId(request, requestId).thenRun(() -> { synchronized (ConsumerImpl.this) { if (changeToReadyState()) { - consumerIsReconnectedToBroker(cnx, currentSize); + consumerIsReconnectedToBroker(cnx, currentSize, connectionPermitState); } else { // Consumer was closed while reconnecting, close the connection to make sure the broker // drops the consumer on its side @@ -969,7 +980,7 @@ 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()); + updateAvailablePermits(connectionPermitState, getCurrentReceiverQueueSize()); } future.complete(null); }).exceptionally((e) -> { @@ -1062,10 +1073,22 @@ protected void closeWhenReceivedUnrecoverableError(Throwable t, ClientCnx cnx) { }); } - protected void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize) { + protected synchronized void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize) { log.info("Subscribed to topic"); + ConsumerPermitState connectionPermitState = permitState; + if (cnx() == cnx) { + AVAILABLE_PERMITS_UPDATER.set(connectionPermitState, 0); + // Flow becomes valid only after Subscribe succeeds. This prevents a receiver-queue update racing with + // connection setup from writing Flow before the broker has created this consumer incarnation. + connectionPermitState.flowEnabled = true; + } + } - AVAILABLE_PERMITS_UPDATER.set(this, 0); + private synchronized void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize, + ConsumerPermitState connectionPermitState) { + if (permitState == connectionPermitState) { + consumerIsReconnectedToBroker(cnx, currentQueueSize); + } } /** @@ -1300,9 +1323,10 @@ protected MessageImpl newSingleMessage(final int index, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { + ClientCnx messageCnx = cnx(); return newSingleMessage(index, numMessages, brokerEntryMetadata, msgMetadata, singleMessageMetadata, payload, messageId, schema, containMetadata, ackBitSet, ackSetInMessageId, redeliveryCount, consumerEpoch, - isEncrypted, cnx()); + isEncrypted, messageCnx, getPermitState(messageCnx)); } private MessageImpl newSingleMessage(final int index, @@ -1319,7 +1343,8 @@ private MessageImpl newSingleMessage(final int index, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted, - final ClientCnx messageCnx) { + final ClientCnx messageCnx, + final ConsumerPermitState messagePermitState) { log.debug().attr("index", index) .log("processing message num - in batch"); @@ -1357,6 +1382,7 @@ private MessageImpl newSingleMessage(final int index, createEncryptionContext(msgMetadata, isEncrypted), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); + message.setPermitState(messagePermitState); return message; } catch (IOException | IllegalStateException e) { throw new IllegalStateException(e); @@ -1386,8 +1412,9 @@ protected MessageImpl newMessage(final MessageIdImpl messageId, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { + ClientCnx messageCnx = cnx(); return newMessage(messageId, brokerEntryMetadata, messageMetadata, payload, schema, redeliveryCount, - consumerEpoch, isEncrypted, cnx()); + consumerEpoch, isEncrypted, messageCnx, getPermitState(messageCnx)); } private MessageImpl newMessage(final MessageIdImpl messageId, @@ -1398,16 +1425,18 @@ private MessageImpl newMessage(final MessageIdImpl messageId, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted, - final ClientCnx messageCnx) { + final ClientCnx messageCnx, + final ConsumerPermitState messagePermitState) { final MessageImpl message = MessageImpl.create(topicName.toString(), messageId, messageMetadata, payload, createEncryptionContext(messageMetadata, isEncrypted), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); + message.setPermitState(messagePermitState); return message; } private void executeNotifyCallback(final MessageImpl message) { - ClientCnx messageCnx = message.getCnx(); + ConsumerPermitState messagePermitState = message.getPermitState(); int messageSize = message.size(); messagesPrefetchedGauge.increment(); bytesPrefetchedGauge.add(messageSize); @@ -1417,7 +1446,8 @@ private void executeNotifyCallback(final MessageImpl message) { // if asyncReceive is waiting then notify callback without adding to incomingMessages queue internalPinnedExecutor.execute(() -> { if (!isValidConsumerEpoch(message)) { - completeStaleEpochDiscard(message, messageCnx, messageSize); + // isValidConsumerEpoch has already released the stale message. + completeStaleEpochDiscardAccounting(messagePermitState, messageSize); return; } Message interceptMsg = onArrival(message); @@ -1429,11 +1459,23 @@ private void executeNotifyCallback(final MessageImpl message) { }); } - private void completeStaleEpochDiscard(MessageImpl message, ClientCnx messageCnx, int messageSize) { + private void completeStaleEpochDiscardAccounting(ConsumerPermitState messagePermitState, int messageSize) { messagesPrefetchedGauge.decrement(); bytesPrefetchedGauge.subtract(messageSize); - returnPermitsToSourceConnection(messageCnx, 1); - message.release(); + returnPermitsToSourceConsumer(messagePermitState, 1); + } + + 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) { + ClientCnx messageCnx = cnx(); + processPayloadByProcessor(brokerEntryMetadata, messageMetadata, byteBuf, messageId, schema, redeliveryCount, + ackSet, consumerEpoch, getPermitState(messageCnx)); } protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, @@ -1443,7 +1485,8 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe final Schema schema, final int redeliveryCount, final List ackSet, - long consumerEpoch) { + long consumerEpoch, + final ConsumerPermitState messagePermitState) { final MessagePayloadImpl payload = MessagePayloadImpl.create(byteBuf); final MessagePayloadContextImpl entryContext = MessagePayloadContextImpl.get( brokerEntryMetadata, messageMetadata, messageId, this, redeliveryCount, ackSet, consumerEpoch); @@ -1451,7 +1494,9 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe try { conf.getPayloadProcessor().process(payload, entryContext, schema, message -> { if (message != null) { - executeNotifyCallback((MessageImpl) message); + MessageImpl messageImpl = (MessageImpl) message; + messageImpl.setPermitState(messagePermitState); + executeNotifyCallback(messageImpl); } else { skippedMessages.incrementAndGet(); } @@ -1459,20 +1504,21 @@ 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, messagePermitState, ValidationError.BatchDeSerializeError); } finally { entryContext.recycle(); payload.release(); // byteBuf.release() is called in this method } if (skippedMessages.get() > 0) { - increaseAvailablePermits(cnx(), skippedMessages.get()); + returnPermitsToSourceConsumer(messagePermitState, skippedMessages.get()); } tryTriggerListener(); } void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, ClientCnx cnx) { + final ConsumerPermitState messagePermitState = getPermitState(cnx); long[] ackSet = EMPTY_ACK_SET; if (cmdMessage.getAckSetsCount() > 0) { ackSet = new long[cmdMessage.getAckSetsCount()]; @@ -1500,11 +1546,13 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien try { if (!verifyChecksum(headersAndPayload, messageId)) { // discard message with checksum error - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); + discardCorruptedMessage(messageId, cnx, messagePermitState, + ValidationError.ChecksumMismatch, earlyFailurePermits); return; } } catch (Throwable t) { - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); + discardCorruptedMessage(messageId, cnx, messagePermitState, + ValidationError.ChecksumMismatch, earlyFailurePermits); return; } @@ -1514,7 +1562,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien brokerEntryMetadata = Commands.parseBrokerEntryMetadataIfExist(headersAndPayload); msgMetadata = Commands.parseMessageMetadata(headersAndPayload); } catch (Throwable t) { - discardCorruptedMessage(messageId, cnx, ValidationError.ChecksumMismatch, earlyFailurePermits); + discardCorruptedMessage(messageId, cnx, messagePermitState, + ValidationError.ChecksumMismatch, earlyFailurePermits); return; } @@ -1534,7 +1583,7 @@ 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"); - returnPermitsToSourceConnection(cnx, messagePermits); + returnPermitsToSourceConsumer(messagePermitState, messagePermits); return; } @@ -1552,7 +1601,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, messagePermits); + : uncompressPayloadIfNeeded(messageId, msgMetadata, decryptedPayload, cnx, messagePermitState, + true, messagePermits); decryptedPayload.release(); if (uncompressedPayload == null) { // Message was discarded on decompression error @@ -1562,7 +1612,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, + messagePermitState); return; } @@ -1571,7 +1622,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, cnx, + messagePermitState); if (uncompressedPayload == null) { return; } @@ -1607,7 +1659,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien final MessageImpl message = newMessage(msgId, brokerEntryMetadata, msgMetadata, uncompressedPayload, - schema, redeliveryCount, consumerEpoch, isMessageUndecryptable, cnx); + schema, redeliveryCount, consumerEpoch, isMessageUndecryptable, cnx, + messagePermitState); uncompressedPayload.release(); if (deadLetterPolicy != null && possibleSendToDeadLetterTopicMessages != null) { @@ -1627,7 +1680,8 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien } else { // handle batch message enqueuing; uncompressed payload has all messages in batch receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, - uncompressedPayload, messageId, cnx, consumerEpoch, isMessageUndecryptable, messagePermits); + uncompressedPayload, messageId, cnx, consumerEpoch, isMessageUndecryptable, messagePermits, + messagePermitState); uncompressedPayload.release(); } @@ -1636,9 +1690,9 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien } private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId, - MessageIdData messageId, ClientCnx cnx) { + MessageIdData messageId, ClientCnx cnx, ConsumerPermitState messagePermitState) { if (msgMetadata.getChunkId() != (msgMetadata.getNumChunksFromMsg() - 1)) { - increaseAvailablePermits(cnx); + returnPermitsToSourceConsumer(messagePermitState, 1); } // Lazy task scheduling to expire incomplete chunk message if (expireTimeOfIncompleteChunkedMessageMillis > 0 && expireChunkMessageTaskScheduled.compareAndSet(false, @@ -1777,7 +1831,8 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m compressedPayload.release(); compressedPayload = chunkedMsgCtx.chunkedMsgBuffer; ByteBuf uncompressedPayload = - uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, cnx, false, 1); + uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, cnx, messagePermitState, + false, 1); compressedPayload.release(); return uncompressedPayload; } @@ -1849,6 +1904,16 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, boolean isEncrypted, int messagePermits) { + receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, + uncompressedPayload, messageId, cnx, consumerEpoch, isEncrypted, messagePermits, + getPermitState(cnx)); + } + + void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, + int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, + MessageIdData messageId, ClientCnx cnx, long consumerEpoch, + boolean isEncrypted, int messagePermits, + ConsumerPermitState messagePermitState) { int batchSize = msgMetadata.getNumMessagesInBatch(); // create ack tracker for entry aka batch @@ -1872,7 +1937,8 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, 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, cnx); + ackBitSet, ackSetInMessageId, redeliveryCount, consumerEpoch, isEncrypted, cnx, + messagePermitState); 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. @@ -1914,7 +1980,8 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } catch (RuntimeException e) { log.warn().exception(e) .log("unable to obtain message in batch"); - discardCorruptedMessage(messageId, cnx, ValidationError.BatchDeSerializeError, permitBudget.drain()); + discardCorruptedMessage(messageId, cnx, messagePermitState, + ValidationError.BatchDeSerializeError, permitBudget.drain()); } finally { if (ackBitSet != null) { ackBitSet.recycle(); @@ -1935,7 +2002,7 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, .attr("remainingCapacity", incomingMessages.remainingCapacity()) .log("enqueued messages in batch. queue size -, available queue size"); - returnPermitsToSourceConnection(cnx, permitBudget.drain()); + returnPermitsToSourceConsumer(messagePermitState, permitBudget.drain()); } private boolean isPriorEntryIndex(long idx) { @@ -1959,8 +2026,7 @@ private boolean isSameEntry(MessageIdImpl messageId) { */ @Override protected synchronized void messageProcessed(Message msg) { - ClientCnx currentCnx = cnx(); - ClientCnx msgCnx = ((MessageImpl) msg).getCnx(); + ConsumerPermitState messagePermitState = ((MessageImpl) msg).getPermitState(); lastDequeuedMessageId = msg.getMessageId(); messagesPrefetchedGauge.decrement(); @@ -1969,11 +2035,11 @@ protected synchronized void messageProcessed(Message msg) { bytesPrefetchedGauge.subtract(msg.size()); bytesReceivedCounter.add(msg.size()); - if (msgCnx != currentCnx) { + if (messagePermitState != permitState) { // The processed message did belong to the old queue that was cleared after reconnection. } else { if (listener == null && !parentConsumerHasListener) { - increaseAvailablePermits(currentCnx); + updateAvailablePermits(messagePermitState, 1); } stats.updateNumMsgsReceived(msg); @@ -2006,49 +2072,79 @@ protected void trackMessage(MessageId messageId, int redeliveryCount) { } void increaseAvailablePermits(MessageImpl msg) { - returnPermitsToSourceConnection(msg.getCnx(), 1); + increaseAvailablePermits(msg.getPermitState()); + } + + void increaseAvailablePermits(ConsumerPermitState messagePermitState) { + returnPermitsToSourceConsumer(messagePermitState, 1); } - private void returnPermitsToSourceConnection(ClientCnx messageCnx, int delta) { + private void returnPermitsToSourceConsumer(ConsumerPermitState messagePermitState, int delta) { if (delta <= 0) { return; } + updateAvailablePermits(messagePermitState, delta); + } - int permitsToSend = 0; - boolean overflow = false; - synchronized (this) { - if (messageCnx == null || messageCnx != cnx()) { + protected void updateAvailablePermits(ConsumerPermitState messagePermitState, int delta) { + if (messagePermitState == null || messagePermitState != permitState) { + return; + } + + int available; + while (true) { + int current = AVAILABLE_PERMITS_UPDATER.get(messagePermitState); + try { + available = Math.addExact(current, delta); + } catch (ArithmeticException e) { + scheduleCloseForInvalidPermitAccumulator(messagePermitState, + new InvalidMessagePermitsException("Client permit accumulator overflow")); return; } - - int available = 0; - while (true) { - int current = AVAILABLE_PERMITS_UPDATER.get(this); - try { - available = Math.addExact(current, delta); - } catch (ArithmeticException e) { - overflow = true; - break; - } - if (AVAILABLE_PERMITS_UPDATER.compareAndSet(this, current, available)) { - break; - } + if (AVAILABLE_PERMITS_UPDATER.compareAndSet(messagePermitState, current, available)) { + break; } - while (!overflow && available >= getCurrentReceiverQueueSize() / 2 && !paused) { - if (AVAILABLE_PERMITS_UPDATER.compareAndSet(this, available, 0)) { - permitsToSend = available; - break; + } + + while (messagePermitState == permitState && available >= getCurrentReceiverQueueSize() / 2 && !paused + && messagePermitState.flowEnabled) { + if (AVAILABLE_PERMITS_UPDATER.compareAndSet(messagePermitState, available, 0)) { + ClientCnx flowCnx = cnx(); + if (flowCnx != null) { + scheduleFlowPermitsToBroker(messagePermitState, flowCnx, available); } - available = AVAILABLE_PERMITS_UPDATER.get(this); + return; } + available = AVAILABLE_PERMITS_UPDATER.get(messagePermitState); } + } - if (overflow) { - closeConnectionForInvalidMessagePermits(null, messageCnx, - new InvalidMessagePermitsException("Client permit accumulator overflow")); - } else if (permitsToSend > 0) { - sendFlowPermitsToBroker(messageCnx, permitsToSend); + private void scheduleCloseForInvalidPermitAccumulator(ConsumerPermitState sourcePermitState, + InvalidMessagePermitsException error) { + ClientCnx sourceCnx = cnx(); + if (sourceCnx == null) { + return; } + sourceCnx.ctx().channel().eventLoop().execute(() -> { + if (sourcePermitState == permitState) { + closeConnectionForInvalidMessagePermits(null, sourceCnx, error); + } + }); + } + + private void scheduleFlowPermitsToBroker(ConsumerPermitState sourcePermitState, ClientCnx sourceCnx, + int permits) { + sourceCnx.ctx().channel().eventLoop().execute(() -> { + // The event loop is also the final ordering point for Subscribe and Flow on this connection. If this task + // wins, Flow is written before a later replacement Subscribe. If replacement wins, its new state makes + // this task stale and the old Flow is dropped, including when the physical ClientCnx is reused. + if (sourcePermitState != permitState) { + return; + } + log.debug().attr("available", permits) + .log("Sending permit-cmd to broker with available permits ="); + sendFlowPermitsToBroker(sourceCnx, permits); + }); } void increaseAvailablePermits(ClientCnx currentCnx) { @@ -2056,17 +2152,7 @@ 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); - } - } + updateAvailablePermits(getPermitState(currentCnx), delta); } public void increaseAvailablePermits(int delta) { @@ -2201,7 +2287,8 @@ 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, currentCnx, getPermitState(currentCnx), + ValidationError.DecryptionError, batchSize); return DecryptResult.discard(); case FAIL: if (cryptoReaderNotExist) { @@ -2224,7 +2311,8 @@ private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData message } private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetadata msgMetadata, ByteBuf payload, - ClientCnx currentCnx, boolean checkMaxMessageSize, int messagePermits) { + ClientCnx currentCnx, ConsumerPermitState messagePermitState, + boolean checkMaxMessageSize, int messagePermits) { CompressionType compressionType = msgMetadata.getCompression(); CompressionCodec codec = CompressionCodecProvider.getCompressionCodec(compressionType); int uncompressedSize = msgMetadata.getUncompressedSize(); @@ -2234,8 +2322,8 @@ 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, - messagePermits); + discardCorruptedMessage(messageId, currentCnx, messagePermitState, + ValidationError.UncompressedSizeCorruption, messagePermits); return null; } try { @@ -2247,7 +2335,8 @@ private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetada .exceptionMessage(e) .exception(e) .log("Failed to decompress message with at"); - discardCorruptedMessage(messageId, currentCnx, ValidationError.DecompressionError, messagePermits); + discardCorruptedMessage(messageId, currentCnx, messagePermitState, + ValidationError.DecompressionError, messagePermits); return null; } } @@ -2269,30 +2358,35 @@ private boolean verifyChecksum(ByteBuf headersAndPayload, MessageIdData messageI return true; } - private void discardCorruptedMessage(MessageIdImpl messageId, ClientCnx currentCnx, + private void discardCorruptedMessage(MessageIdImpl messageId, ConsumerPermitState messagePermitState, ValidationError validationError) { 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); + ClientCnx currentCnx = cnx(); + if (currentCnx != null) { + currentCnx.ctx().writeAndFlush(cmd, currentCnx.ctx().voidPromise()); + } else { + cmd.release(); + } + returnPermitsToSourceConsumer(messagePermitState, 1); stats.incrementNumReceiveFailed(); } private void discardCorruptedMessage(MessageIdData messageId, ClientCnx currentCnx, - ValidationError validationError, int messagePermits) { + ConsumerPermitState messagePermitState, ValidationError validationError, int messagePermits) { log.error().attr("messageId", messageId) .log("Discarding corrupted message"); - discardMessage(messageId, currentCnx, validationError, messagePermits); + discardMessage(messageId, currentCnx, messagePermitState, validationError, messagePermits); } - private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, ValidationError validationError, - int messagePermits) { + private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, + ConsumerPermitState messagePermitState, 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()); - returnPermitsToSourceConnection(currentCnx, messagePermits); + returnPermitsToSourceConsumer(messagePermitState, messagePermits); stats.incrementNumReceiveFailed(); } @@ -2303,7 +2397,9 @@ private void closeConnectionForInvalidMessagePermits(MessageIdData messageId, Cl .exceptionMessage(error) .log("Closing source connection because message permit accounting is invalid"); stats.incrementNumReceiveFailed(); - messageCnx.ctx().close(); + if (messageCnx != null) { + messageCnx.ctx().close(); + } } @Override @@ -2326,7 +2422,7 @@ int getPartitionIndex() { @Override public int getAvailablePermits() { - return AVAILABLE_PERMITS_UPDATER.get(this); + return AVAILABLE_PERMITS_UPDATER.get(permitState); } @Override @@ -2408,6 +2504,7 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { return; } ClientCnx cnx = cnx(); + ConsumerPermitState redeliveryPermitState = getPermitState(cnx); if (isConnected() && cnx.getRemoteEndpointProtocolVersion() >= ProtocolVersion.v2.getValue()) { int messagesFromQueue = removeExpiredMessagesFromQueue(messageIds); Iterables.partition(messageIds, MAX_REDELIVER_UNACKNOWLEDGED).forEach(ids -> { @@ -2419,7 +2516,7 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { }); }); if (messagesFromQueue > 0) { - increaseAvailablePermits(cnx, messagesFromQueue); + updateAvailablePermits(redeliveryPermitState, messagesFromQueue); } log.debug().attr("increase", messagesFromQueue) .log("Redeliver unacked messages and increase permits"); @@ -3196,8 +3293,12 @@ void resetBackoff() { this.connectionHandler.resetBackoff(); } - void connectionClosed(ClientCnx cnx, Optional initialConnectionDelayMs, Optional hostUrl) { + synchronized void connectionClosed(ClientCnx cnx, Optional initialConnectionDelayMs, Optional hostUrl) { + boolean wasCurrentConnection = this.connectionHandler.cnx() == cnx; this.connectionHandler.connectionClosed(cnx, initialConnectionDelayMs, hostUrl); + if (wasCurrentConnection && this.connectionHandler.cnx() != cnx) { + permitState = new ConsumerPermitState(false); + } } public ClientCnx getClientCnx() { @@ -3206,13 +3307,21 @@ public ClientCnx getClientCnx() { void setClientCnx(ClientCnx clientCnx) { if (clientCnx != null) { - this.connectionHandler.setClientCnx(clientCnx); + synchronized (this) { + this.connectionHandler.setClientCnx(clientCnx); + // Always create new permit state, including when the pooled ClientCnx instance is unchanged. + permitState = new ConsumerPermitState(false); + } clientCnx.registerConsumer(consumerId, this); if (conf.isAckReceiptEnabled() && !Commands.peerSupportsAckReceipt(clientCnx.getRemoteEndpointProtocolVersion())) { log.warn().attr("nowversion", clientCnx.getRemoteEndpointProtocolVersion()) .log("Server don't support ack for receipt! " + "ProtoVersion >=17 support! nowVersion"); } + } else { + synchronized (this) { + permitState = new ConsumerPermitState(false); + } } ClientCnx previousClientCnx = clientCnxUsedForConsumerRegistration.getAndSet(clientCnx); if (previousClientCnx != null && previousClientCnx != clientCnx) { @@ -3220,6 +3329,29 @@ void setClientCnx(ClientCnx clientCnx) { } } + private ConsumerPermitState getPermitState(ClientCnx clientCnx) { + while (true) { + ConsumerPermitState currentPermitState = permitState; + ClientCnx currentCnx = cnx(); + if (currentPermitState != permitState) { + continue; + } + if (currentPermitState.flowEnabled) { + return currentCnx == clientCnx ? currentPermitState : null; + } + return clientCnx == null ? currentPermitState : null; + } + } + + @VisibleForTesting + ConsumerPermitState getPermitState() { + return permitState; + } + + boolean isCurrentPermitState(ConsumerPermitState messagePermitState) { + return messagePermitState != null && messagePermitState == permitState; + } + void deregisterFromClientCnx() { 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..c5c28bb6f0b96 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 @@ -66,6 +66,7 @@ public class MessageImpl implements TraceableMessage, Message { protected MessageId messageId; private final MessageMetadata msgMetadata; private ClientCnx cnx; + private ConsumerImpl.ConsumerPermitState permitState; private ByteBuf payload; private Schema schema; @@ -102,6 +103,7 @@ public static MessageImpl create(MessageMetadata msgMetadata, ByteBuffer msg.messageId = null; msg.topic = topic; msg.cnx = null; + msg.permitState = null; msg.payload = Unpooled.wrappedBuffer(payload); msg.properties = null; msg.schema = schema; @@ -194,6 +196,7 @@ private static void init(MessageImpl msg, String topic, BatchMessageIdImp msg.messageId = batchMessageIdImpl; msg.topic = topic; msg.cnx = cnx; + msg.permitState = null; msg.redeliveryCount = redeliveryCount; msg.encryptionCtx = encryptionCtx; msg.schema = schema; @@ -721,6 +724,14 @@ public ClientCnx getCnx() { return cnx; } + ConsumerImpl.ConsumerPermitState getPermitState() { + return permitState; + } + + void setPermitState(ConsumerImpl.ConsumerPermitState permitState) { + this.permitState = permitState; + } + public void recycle() { if (msgMetadata != null) { msgMetadata.clear(); @@ -729,6 +740,7 @@ public void recycle() { brokerEntryMetadata.clear(); } cnx = null; + permitState = null; messageId = null; topic = null; payload = null; 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..805b4b07bc9e2 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 @@ -280,12 +280,12 @@ private void receiveMessageFromConsumer(ConsumerImpl consumer, boolean batchR messages.forEach(msg -> { final boolean skipDueToSeek = duringSeek; MessageImpl msgImpl = (MessageImpl) msg; - ClientCnx cnx = msgImpl.getCnx(); + ConsumerImpl.ConsumerPermitState permitState = msgImpl.getPermitState(); boolean isValidEpoch = isValidConsumerEpoch(msgImpl); if (isValidEpoch && !skipDueToSeek) { messageReceived(consumer, msg); } else if (!isValidEpoch) { - consumer.increaseAvailablePermits(cnx); + consumer.increaseAvailablePermits(permitState); } else if (skipDueToSeek) { log.info().attr("messageId", msg.getMessageId()) .log("Skip processing message received during seek"); 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 01e9c70a2ca8e..e0e5ec5b7141e 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 @@ -111,15 +111,12 @@ private Message fetchSingleMessageFromBroker() throws PulsarClientException { do { message = incomingMessages.take(); lastDequeuedMessageId = message.getMessageId(); - ClientCnx msgCnx = ((MessageImpl) message).getCnx(); - // 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()) { - waitingOnReceiveForZeroQueueSize = false; - break; - } + ConsumerPermitState messagePermitState = ((MessageImpl) message).getPermitState(); + // If the message belongs to an old broker-consumer incarnation, discard it and wait for the message + // from the latest Flow command. ClientCnx identity alone is insufficient because it can be reused. + if (isCurrentPermitState(messagePermitState)) { + waitingOnReceiveForZeroQueueSize = false; + break; } } while (true); @@ -185,7 +182,7 @@ private void triggerZeroQueueSizeListener(final Message message) { .exception(t) .log("Message listener error in processing unqueued message"); } - increaseAvailablePermits(cnx()); + increaseAvailablePermits((MessageImpl) message); waitingOnListenerForZeroQueueSize = false; }); } @@ -199,7 +196,8 @@ protected void tryTriggerListener() { void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, - boolean isEncrypted, int messagePermits) { + boolean isEncrypted, int messagePermits, + ConsumerPermitState messagePermitState) { rejectBatchMessageByClosingConsumer( new MessageIdImpl(messageId.getLedgerId(), messageId.getEntryId(), getPartitionIndex()) @@ -216,12 +214,13 @@ 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, + ConsumerPermitState messagePermitState) { if (this.isBatch(messageMetadata)) { rejectBatchMessageByClosingConsumer(messageId); } else { super.processPayloadByProcessor(brokerEntryMetadata, messageMetadata, byteBuf, messageId, schema, - redeliveryCount, ackSet, consumerEpoch); + redeliveryCount, ackSet, consumerEpoch, messagePermitState); } } 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 fd39d3e4093a0..8afa0677f47fc 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,9 +23,11 @@ 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; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -33,14 +35,18 @@ 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.channel.EventLoop; import io.netty.util.ReferenceCountUtil; +import java.util.ArrayList; import java.util.Arrays; import java.util.BitSet; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -215,12 +221,132 @@ public void testPermitAccumulatorOverflowClosesSourceConnection() { } } + @Test + public void testSameClientCnxReuseCreatesNewPermitIncarnation() { + consumer.setCurrentReceiverQueueSize(2); + ClientCnx messageCnx = setCurrentConnection(); + ChannelHandlerContext context = messageCnx.ctx(); + ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); + MessageImpl oldMessage = mock(MessageImpl.class); + when(oldMessage.getPermitState()).thenReturn(oldPermitState); + + // Recreate the broker consumer while reusing the same pooled physical ClientCnx. + consumer.setClientCnx(messageCnx); + Assert.assertNotSame(consumer.getPermitState(), oldPermitState); + consumer.consumerIsReconnectedToBroker(messageCnx, 0); + consumer.increaseAvailablePermits(oldMessage); + + Assert.assertEquals(consumer.getAvailablePermits(), 0); + verify(context, never()).writeAndFlush(any(), any(ChannelPromise.class)); + } + + @Test(invocationTimeOut = 5000) + public void testPermitReturnRaceWithSameClientCnxReconnectDoesNotContaminateNewAccumulator() { + consumer.paused = true; + ClientCnx messageCnx = setCurrentConnection(); + ExecutorService raceExecutor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 100; i++) { + ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); + MessageImpl oldMessage = mock(MessageImpl.class); + when(oldMessage.getPermitState()).thenReturn(oldPermitState); + CountDownLatch start = new CountDownLatch(1); + + CompletableFuture returnPermit = CompletableFuture.runAsync(() -> { + await(start); + consumer.increaseAvailablePermits(oldMessage); + }, raceExecutor); + CompletableFuture reconnect = CompletableFuture.runAsync(() -> { + await(start); + consumer.setClientCnx(messageCnx); + consumer.consumerIsReconnectedToBroker(messageCnx, 0); + }, raceExecutor); + + start.countDown(); + CompletableFuture.allOf(returnPermit, reconnect).join(); + Assert.assertEquals(consumer.getAvailablePermits(), 0); + } + } finally { + raceExecutor.shutdownNow(); + } + } + + @Test(invocationTimeOut = 5000) + public void testConcurrentPermitReturnsStayInCurrentIncarnationAccumulator() { + consumer.paused = true; + setCurrentConnection(); + ConsumerImpl.ConsumerPermitState currentPermitState = consumer.getPermitState(); + int threadCount = 8; + int returnsPerThread = 1000; + ExecutorService returnExecutor = Executors.newFixedThreadPool(threadCount); + CountDownLatch start = new CountDownLatch(1); + try { + List> returns = new ArrayList<>(threadCount); + for (int i = 0; i < threadCount; i++) { + returns.add(CompletableFuture.runAsync(() -> { + await(start); + for (int permit = 0; permit < returnsPerThread; permit++) { + consumer.increaseAvailablePermits(currentPermitState); + } + }, returnExecutor)); + } + start.countDown(); + CompletableFuture.allOf(returns.toArray(CompletableFuture[]::new)).join(); + + Assert.assertEquals(consumer.getAvailablePermits(), threadCount * returnsPerThread); + } finally { + returnExecutor.shutdownNow(); + } + } + + @Test + public void testQueuedFlowFromOldIncarnationIsDroppedAfterSameClientCnxReuse() { + consumer.setCurrentReceiverQueueSize(2); + ClientCnx messageCnx = setCurrentConnection(); + ChannelHandlerContext context = messageCnx.ctx(); + EventLoop eventLoop = context.channel().eventLoop(); + List queuedTasks = new ArrayList<>(); + doAnswer(invocation -> { + queuedTasks.add(invocation.getArgument(0)); + return null; + }).when(eventLoop).execute(any(Runnable.class)); + + ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); + consumer.increaseAvailablePermits(oldPermitState); + Assert.assertEquals(queuedTasks.size(), 1); + + // Recreate the broker consumer before the old Flow task reaches the shared physical connection. + consumer.setClientCnx(messageCnx); + queuedTasks.get(0).run(); + + Assert.assertEquals(consumer.getAvailablePermits(), 0); + verify(context, never()).writeAndFlush(any(), any(ChannelPromise.class)); + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new CompletionException(e); + } + } + private ClientCnx setCurrentConnection() { ClientCnx messageCnx = mock(ClientCnx.class); ChannelHandlerContext context = mock(ChannelHandlerContext.class); + Channel channel = mock(Channel.class); + EventLoop eventLoop = mock(EventLoop.class); when(context.voidPromise()).thenReturn(mock(ChannelPromise.class)); + when(context.channel()).thenReturn(channel); + when(channel.eventLoop()).thenReturn(eventLoop); + doAnswer(invocation -> { + invocation.getArgument(0).run(); + return null; + }).when(eventLoop).execute(any(Runnable.class)); when(messageCnx.ctx()).thenReturn(context); consumer.setClientCnx(messageCnx); + consumer.consumerIsReconnectedToBroker(messageCnx, 0); return messageCnx; } From bd4a56494f3766132ed619ccdf10d9a2b564bfc0 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 16 Aug 2026 14:07:00 +0800 Subject: [PATCH 5/8] [test][broker] Cover finalized permit carrier lifetime Assisted-by: OpenAI Codex --- .../pulsar/broker/service/ConsumerTest.java | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) 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 02a529ecedb2b..860603f2c1dd8 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 @@ -37,6 +37,7 @@ import static org.testng.Assert.assertSame; import static org.testng.Assert.expectThrows; import io.netty.util.concurrent.ImmediateEventExecutor; +import io.netty.util.concurrent.Promise; import java.io.IOException; import java.net.SocketAddress; import java.util.ArrayList; @@ -114,6 +115,8 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { when(emptyPartialBatch.getLedgerId()).thenReturn(3L); when(emptyPartialBatch.getEntryId()).thenReturn(3L); List entries = new ArrayList<>(List.of(partialBatch, rejectedBatch, emptyPartialBatch)); + // A null slot represents an entry removed by dispatcher filtering before final send admission. + entries.add(null); EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); batchSizes.setBatchSize(0, 10); @@ -137,6 +140,7 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { assertEquals(sendResult.getMessagePermits(0), 3); assertEquals(sendResult.getMessagePermits(1), 0); assertEquals(sendResult.getMessagePermits(2), 0); + assertEquals(sendResult.getMessagePermits(3), 0); assertEquals(sharedConsumer.getAvailablePermits(), 97); assertEquals(sharedConsumer.getUnackedMessages(), 3); assertNull(entries.get(1)); @@ -155,6 +159,62 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { } } + @Test + public void testFinalizedPermitsOutliveRecycledSenderInputs() { + Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, + "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, + DEFAULT_CONSUMER_EPOCH); + sharedConsumer.setPendingAcksAddHandler((ignored, ledgerId, entryId, stickyKeyHash) -> true); + sharedConsumer.flowPermits(100); + + Entry partialBatch = mock(Entry.class); + when(partialBatch.getLedgerId()).thenReturn(1L); + when(partialBatch.getEntryId()).thenReturn(1L); + Entry completeBatch = mock(Entry.class); + when(completeBatch.getLedgerId()).thenReturn(2L); + when(completeBatch.getEntryId()).thenReturn(2L); + List entries = new ArrayList<>(List.of(partialBatch, completeBatch)); + + EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); + batchSizes.setBatchSize(0, 10); + batchSizes.setBatchSize(1, 4); + EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); + batchIndexesAcks.setIndexesAcks(0, Pair.of(10, new long[] {0b100101L})); + + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); + when(cnx.getCommandSender()).thenReturn(commandSender); + Promise pendingWrite = ImmediateEventExecutor.INSTANCE.newPromise(); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), + any(), any(), anyLong())).thenReturn(pendingWrite); + + SendMessagesResult sendResult; + try { + sendResult = sharedConsumer.sendMessagesWithResult( + entries, batchSizes, batchIndexesAcks, 14, 0, 0, mock(RedeliveryTracker.class)); + } finally { + // The command sender owns and recycles these inputs before its asynchronous write completes. + batchSizes.recyle(); + batchIndexesAcks.recycle(); + } + + EntryBatchSizes reusedBatchSizes = EntryBatchSizes.get(entries.size()); + EntryBatchIndexesAcks reusedBatchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); + try { + // Mutate the recycled carriers while the write is still pending. The finalized result must be independent. + reusedBatchSizes.setBatchSize(0, 1); + reusedBatchSizes.setBatchSize(1, 1); + reusedBatchIndexesAcks.setIndexesAcks(0, Pair.of(1, new long[] {1L})); + + assertEquals(sendResult.getMessagePermits(0), 3); + assertEquals(sendResult.getMessagePermits(1), 4); + assertEquals(sendResult.getTotalMessagePermits(), 7); + } finally { + pendingWrite.trySuccess(null); + reusedBatchSizes.recyle(); + reusedBatchIndexesAcks.recycle(); + } + } + @Test public void testSendMessagesWriteFailureDisconnectsConsumer() throws BrokerServiceException { Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, From 60c6ccbf451cf7c8a68b36fa4ea450213ed0e997 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 16 Aug 2026 17:41:42 +0800 Subject: [PATCH 6/8] [test][client][broker] Expand permit accounting coverage Assisted-by: OpenAI Codex --- .../pulsar/broker/service/ConsumerTest.java | 89 ++++++++ .../service/EntryBatchIndexesAcksTest.java | 30 +++ .../service/PulsarCommandSenderImplTest.java | 128 +++++++++++ .../PersistentSharedPermitAccountingTest.java | 206 ++++++++++++++++++ .../client/impl/BatchMessageIndexAckTest.java | 15 +- .../pulsar/client/impl/ConsumerImplTest.java | 106 +++++++++ .../impl/MessagePermitAccountingTest.java | 71 ++++++ 7 files changed, 642 insertions(+), 3 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PulsarCommandSenderImplTest.java create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java 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 860603f2c1dd8..b5e7a1793bebf 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 @@ -41,8 +41,10 @@ import java.io.IOException; import java.net.SocketAddress; import java.util.ArrayList; +import java.util.BitSet; import java.util.List; import java.util.Optional; +import java.util.SplittableRandom; import org.apache.bookkeeper.mledger.Entry; import org.apache.commons.lang3.tuple.Pair; import org.apache.pulsar.broker.PulsarService; @@ -56,6 +58,8 @@ @Test(groups = "broker") public class ConsumerTest { + private static final long RANDOM_SEED = 0x491C05E5L; + private Consumer consumer; private Subscription subscription; private ServerCnx cnx; @@ -247,6 +251,91 @@ public void testSendMessagesWriteFailureDisconnectsConsumer() throws BrokerServi } } + @Test + public void testRandomizedFinalAdmissionUsesOnePermitResultForAllBrokerCounters() { + final int entriesCount = 1_000; + final int initialPermits = 1_000_000; + SplittableRandom random = new SplittableRandom(RANDOM_SEED); + boolean[] admissionAllowed = new boolean[entriesCount]; + int[] expectedPermits = new int[entriesCount]; + Entry[] originalEntries = new Entry[entriesCount]; + List entries = new ArrayList<>(entriesCount); + EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesCount); + EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesCount); + int totalMessages = 0; + int expectedTotalPermits = 0; + + for (int i = 0; i < entriesCount; i++) { + int batchSize = random.nextInt(1, 257); + batchSizes.setBatchSize(i, batchSize); + totalMessages += batchSize; + if (random.nextInt(8) == 0) { + entries.add(null); + continue; + } + + Entry entry = mock(Entry.class); + originalEntries[i] = entry; + entries.add(entry); + when(entry.getLedgerId()).thenReturn((long) i); + when(entry.getEntryId()).thenReturn((long) i); + + int messagePermits = batchSize; + if (random.nextInt(4) != 0) { + int requiredWords = (batchSize + Long.SIZE - 1) / Long.SIZE; + long[] ackSet = new long[requiredWords + random.nextInt(2)]; + for (int word = 0; word < ackSet.length; word++) { + ackSet[word] = random.nextLong(); + } + BitSet boundedAckSet = BitSet.valueOf(ackSet); + boundedAckSet.clear(batchSize, Math.max(batchSize, boundedAckSet.length())); + messagePermits = boundedAckSet.cardinality(); + batchIndexesAcks.setIndexesAcks(i, Pair.of(batchSize, ackSet)); + } + + admissionAllowed[i] = random.nextInt(5) != 0; + if (messagePermits > 0 && admissionAllowed[i]) { + expectedPermits[i] = messagePermits; + expectedTotalPermits += messagePermits; + } + } + + Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, + "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, + DEFAULT_CONSUMER_EPOCH); + sharedConsumer.setPendingAcksAddHandler( + (ignored, ledgerId, entryId, stickyKeyHash) -> admissionAllowed[(int) ledgerId]); + sharedConsumer.flowPermits(initialPermits); + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); + when(cnx.getCommandSender()).thenReturn(commandSender); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), + any(), any(), anyLong())) + .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); + + try { + SendMessagesResult sendResult = sharedConsumer.sendMessagesWithResult( + entries, batchSizes, batchIndexesAcks, totalMessages, 0, 0, mock(RedeliveryTracker.class)); + + assertEquals(sendResult.getTotalMessagePermits(), expectedTotalPermits, "seed=" + RANDOM_SEED); + assertEquals(sharedConsumer.getAvailablePermits(), initialPermits - expectedTotalPermits, + "seed=" + RANDOM_SEED); + assertEquals(sharedConsumer.getUnackedMessages(), expectedTotalPermits, "seed=" + RANDOM_SEED); + for (int i = 0; i < entriesCount; i++) { + assertEquals(sendResult.getMessagePermits(i), expectedPermits[i], + "seed=" + RANDOM_SEED + ", entry=" + i); + if (originalEntries[i] != null && expectedPermits[i] == 0) { + assertNull(entries.get(i)); + verify(originalEntries[i]).release(); + } else { + assertSame(entries.get(i), originalEntries[i]); + } + } + } finally { + batchSizes.recyle(); + batchIndexesAcks.recycle(); + } + } + @Test public void testSendMessagesResultRejectsPartialFinalization() { SendMessagesResult sendResult = new SendMessagesResult(2); 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 7910e19e43c6f..58b200e591be0 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,16 @@ 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 shouldCalculateUnackedIndexesWithinBatchBounds() { BitSetRecyclable bitSet = BitSetRecyclable.create(); @@ -69,4 +73,30 @@ void shouldResetStateBeforeReusing() { assertEquals(acks.getTotalAckedIndexCount(), 0); } + @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..b218653900fbf --- /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); + SendMessagesResult sendResult = new SendMessagesResult(entries.size()); + sendResult.setMessagePermits(0, 3); + sendResult.setMessagePermits(2, 4); + + try { + PulsarCommandSenderImpl sender = new PulsarCommandSenderImpl(null, cnx, null); + sender.sendMessagesToConsumer(7, "topic", mock(Subscription.class), -1, entries, batchSizes, null, + sendResult, mock(RedeliveryTracker.class), 11); + + 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/persistent/PersistentSharedPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java new file mode 100644 index 0000000000000..b5ecb2aba71a2 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentSharedPermitAccountingTest.java @@ -0,0 +1,206 @@ +/* + * 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 io.netty.buffer.ByteBuf; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelOutboundHandlerAdapter; +import io.netty.channel.ChannelPromise; +import io.netty.util.ReferenceCountUtil; +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.pulsar.broker.service.Consumer; +import org.apache.pulsar.broker.service.ServerCnx; +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); + } + } + + @Test(timeOut = 30000, dataProvider = "sharedDispatcherImplementations") + public void testWriteFailureRemovesDebitedConsumer(boolean classicDispatcher) throws Exception { + boolean originalSetting = getConfig().isSubscriptionSharedUseClassicPersistentImplementation(); + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(classicDispatcher); + String topicName = newTopicName(); + String subscriptionName = "write-failure-subscription"; + + try (org.apache.pulsar.client.api.Consumer clientConsumer = + pulsarClient.newConsumer(Schema.INT32) + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .receiverQueueSize(1) + .subscribe(); + Producer producer = pulsarClient.newProducer(Schema.INT32) + .topic(topicName) + .enableBatching(true) + .batchingMaxMessages(BATCH_SIZE) + .batchingMaxPublishDelay(1, TimeUnit.HOURS) + .create()) { + PersistentSubscription subscription = getSubscription(topicName, subscriptionName); + Consumer originalConsumer = subscription.getConsumers().get(0); + ServerCnx serverCnx = (ServerCnx) originalConsumer.cnx(); + AtomicBoolean writeFailed = new AtomicBoolean(); + failNextFlush(serverCnx, writeFailed); + + sendBatch(producer, 0); + + Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { + Assert.assertTrue(writeFailed.get()); + Assert.assertFalse(subscription.getConsumers().contains(originalConsumer)); + CompletableFuture registeredConsumer = + serverCnx.getConsumers().get(originalConsumer.consumerId()); + Assert.assertTrue(registeredConsumer == null || registeredConsumer.getNow(null) != originalConsumer); + }); + + 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()); + } + + private static void failNextFlush(ServerCnx serverCnx, AtomicBoolean writeFailed) throws Exception { + serverCnx.ctx().channel().eventLoop().submit(() -> serverCnx.ctx().pipeline().addFirst( + new ChannelOutboundHandlerAdapter() { + @Override + public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception { + if (msg instanceof ByteBuf buffer && !buffer.isReadable() + && writeFailed.compareAndSet(false, true)) { + ctx.pipeline().remove(this); + ReferenceCountUtil.safeRelease(msg); + promise.setFailure(new IOException("intentional message write failure")); + return; + } + ctx.write(msg, promise); + } + })).sync(); + } +} 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 2b755e50ddfd8..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 @@ -235,6 +235,18 @@ public void testPartialBatchRedeliveryKeepsSharedPermitBalance(boolean classicDi .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); @@ -252,9 +264,6 @@ public void testPartialBatchRedeliveryKeepsSharedPermitBalance(boolean classicDi Assert.assertTrue(expectedValues.isEmpty()); Assert.assertNull(replacementConsumer.receive(1, TimeUnit.SECONDS)); - 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); 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 8afa0677f47fc..23e8f655a8312 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 @@ -44,6 +44,7 @@ import java.util.Arrays; import java.util.BitSet; import java.util.List; +import java.util.SplittableRandom; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; @@ -65,11 +66,15 @@ 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.CommandMessage; import org.apache.pulsar.common.api.proto.CompressionType; +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; @@ -80,12 +85,16 @@ 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() { @@ -99,6 +108,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); @@ -323,6 +339,96 @@ public void testQueuedFlowFromOldIncarnationIsDroppedAfterSameClientCnxReuse() { verify(context, never()).writeAndFlush(any(), any(ChannelPromise.class)); } + @Test + public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetchGauges() throws Exception { + consumer.paused = true; + ClientCnx messageCnx = setCurrentConnection(); + 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(), 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 testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws Exception { + consumer.paused = true; + ClientCnx messageCnx = setCurrentConnection(); + 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(); + 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}); + Commands.serializeSingleMessageInBatchWithPayload(new SingleMessageMetadata(), 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, messagePermits, + "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); + for (int permit = 0; permit < messagePermits; permit++) { + Message message = consumer.incomingMessages.poll(); + Assert.assertNotNull(message); + consumer.messageProcessed(message); + message.release(); + } + expectedReturnedPermits += messagePermits; + Assert.assertEquals(consumer.getAvailablePermits(), expectedReturnedPermits, + "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); + } finally { + batch.release(); + } + } + } + private static void await(CountDownLatch latch) { try { latch.await(); diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java index a73e4e6f3390e..cc8fb9dcf8752 100644 --- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java +++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java @@ -20,12 +20,16 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.expectThrows; +import java.util.BitSet; +import java.util.SplittableRandom; import org.apache.pulsar.client.impl.MessagePermitAccounting.Budget; import org.apache.pulsar.common.api.proto.CommandMessage; import org.testng.annotations.Test; public class MessagePermitAccountingTest { + private static final long RANDOM_SEED = 0x4915A17E5L; + @Test public void testResolveExplicitMessagePermits() { CommandMessage command = new CommandMessage().setMessagePermits(3); @@ -81,4 +85,71 @@ public void testBudgetTracksClaimsAndRestores() { assertEquals(budget.drain(), 0); expectThrows(IllegalStateException.class, budget::claim); } + + @Test + public void testRandomizedAckSetResolutionMatchesBoundedBitSetCardinality() { + SplittableRandom random = new SplittableRandom(RANDOM_SEED); + for (int testCase = 0; testCase < 10_000; testCase++) { + int batchSize = random.nextInt(1, 513); + 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(); + } + // Explicit message permits must be positive. Also guarantees coverage when random data happens to be zero. + int requiredIndex = random.nextInt(batchSize); + ackSet[requiredIndex / Long.SIZE] |= 1L << (requiredIndex % Long.SIZE); + + BitSet boundedAckSet = BitSet.valueOf(ackSet); + boundedAckSet.clear(batchSize, Math.max(batchSize, boundedAckSet.length())); + int expectedPermits = boundedAckSet.cardinality(); + int allAckSetPermits = BitSet.valueOf(ackSet).cardinality(); + String description = "seed=" + RANDOM_SEED + ", case=" + testCase + ", batchSize=" + batchSize; + + CommandMessage explicit = commandWithAckSet(ackSet).setMessagePermits(expectedPermits); + assertEquals(MessagePermitAccounting.resolve(explicit, batchSize), expectedPermits, description); + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(explicit), expectedPermits, description); + + CommandMessage legacy = commandWithAckSet(ackSet); + assertEquals(MessagePermitAccounting.resolve(legacy, batchSize), expectedPermits, description); + assertEquals(MessagePermitAccounting.resolveForEarlyFailure(legacy), allAckSetPermits, description); + } + } + + @Test + public void testRandomizedBudgetConservesPermitsAcrossClaimsAndRestores() { + SplittableRandom random = new SplittableRandom(RANDOM_SEED); + for (int testCase = 0; testCase < 10_000; testCase++) { + int initialPermits = random.nextInt(1, 257); + int modeledRemaining = initialPermits; + int permanentlyClaimed = 0; + Budget budget = new Budget(initialPermits); + + int operations = random.nextInt(1, initialPermits * 3 + 1); + for (int operation = 0; operation < operations && modeledRemaining > 0; operation++) { + if (random.nextBoolean()) { + budget.claim(); + modeledRemaining--; + permanentlyClaimed++; + } else { + budget.claim(); + budget.restore(); + } + } + + int returnedPermits = budget.drain(); + String description = "seed=" + RANDOM_SEED + ", case=" + testCase; + assertEquals(returnedPermits, modeledRemaining, description); + assertEquals(permanentlyClaimed + returnedPermits, initialPermits, description); + assertEquals(budget.drain(), 0, description); + } + } + + private static CommandMessage commandWithAckSet(long[] ackSet) { + CommandMessage command = new CommandMessage(); + for (long word : ackSet) { + command.addAckSet(word); + } + return command; + } } From a1bddb33dab18f7c16e589532b3e7f67db46d8f8 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 16 Aug 2026 22:30:35 +0800 Subject: [PATCH 7/8] [test][client] Verify returned permits on CommandFlow Assisted-by: OpenAI Codex --- .../pulsar/client/impl/ConsumerImplTest.java | 51 ++++++++++++++++--- 1 file changed, 43 insertions(+), 8 deletions(-) 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 23e8f655a8312..a04dd638825ba 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 @@ -70,6 +70,7 @@ 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.CompressionType; import org.apache.pulsar.common.api.proto.MessageIdData; @@ -341,8 +342,14 @@ public void testQueuedFlowFromOldIncarnationIsDroppedAfterSameClientCnxReuse() { @Test public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetchGauges() throws Exception { - consumer.paused = true; + consumer.setCurrentReceiverQueueSize(2); ClientCnx messageCnx = setCurrentConnection(); + 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") @@ -361,7 +368,8 @@ public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetc () -> Assert.assertEquals(consumer.numMessagesInQueue(), 0)).get(5, TimeUnit.SECONDS); Assert.assertEquals(consumer.numMessagesInQueue(), 0); - Assert.assertEquals(consumer.getAvailablePermits(), 1); + Assert.assertEquals(consumer.getAvailablePermits(), 0); + Assert.assertEquals(flowPermits, List.of(1)); verify(messagesPrefetchedGauge).increment(); verify(messagesPrefetchedGauge).decrement(); ArgumentCaptor addedBytes = ArgumentCaptor.forClass(Long.class); @@ -375,9 +383,15 @@ public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetc } @Test - public void testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws Exception { - consumer.paused = true; + public void testRandomizedBatchDecodeWritesExactlyTheReturnedPermitsToFlow() throws Exception { + consumer.setCurrentReceiverQueueSize(2); ClientCnx messageCnx = setCurrentConnection(); + 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; @@ -393,6 +407,7 @@ public void testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws 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) @@ -401,7 +416,12 @@ public void testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws ByteBuf batch = Unpooled.buffer(); for (int index = 0; index < batchSize; index++) { ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {(byte) index}); - Commands.serializeSingleMessageInBatchWithPayload(new SingleMessageMetadata(), payload, batch); + 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(); } @@ -412,16 +432,18 @@ public void testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws int queuedMessages = consumer.internalPinnedExecutor.submit(consumer::numMessagesInQueue) .get(5, TimeUnit.SECONDS); - Assert.assertEquals(queuedMessages, messagePermits, + Assert.assertEquals(queuedMessages, expectedQueuedMessages, "seed=" + RANDOM_PERMIT_SEED + ", case=" + testCase); - for (int permit = 0; permit < messagePermits; permit++) { + 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(), expectedReturnedPermits, + 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(); @@ -429,6 +451,19 @@ public void testRandomizedBatchDecodeReturnsExactlyTheDeliveredPermits() throws } } + 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(); + } + } + private static void await(CountDownLatch latch) { try { latch.await(); From ba15c790e37d099f343c7749aee17bafb9c2b7b6 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 22 Aug 2026 10:50:49 +0800 Subject: [PATCH 8/8] [fix][client][broker] Align PIP-491 implementation with revised design Assisted-by: OpenAI Codex (GPT-5.6) --- .../pulsar/broker/service/Consumer.java | 160 ++- .../broker/service/EntryBatchIndexesAcks.java | 3 +- .../broker/service/PulsarCommandSender.java | 5 +- .../service/PulsarCommandSenderImpl.java | 11 +- .../broker/service/SendMessageResult.java | 61 + .../broker/service/SendMessagesResult.java | 73 -- .../pulsar/broker/service/ServerCnx.java | 8 +- ...PersistentDispatcherMultipleConsumers.java | 6 +- ...tStickyKeyDispatcherMultipleConsumers.java | 6 +- ...PersistentDispatcherMultipleConsumers.java | 13 +- ...entDispatcherMultipleConsumersClassic.java | 13 +- ...sistentDispatcherSingleActiveConsumer.java | 2 +- ...tStickyKeyDispatcherMultipleConsumers.java | 9 +- ...KeyDispatcherMultipleConsumersClassic.java | 9 +- .../pulsar/broker/service/ConsumerTest.java | 228 +--- .../service/EntryBatchIndexesAcksTest.java | 7 +- .../service/PulsarCommandSenderImplTest.java | 8 +- .../broker/service/SendMessageResultTest.java | 42 + .../PersistentSharedPermitAccountingTest.java | 71 -- .../impl/CompactedOutBatchMessageTest.java | 58 +- .../impl/CorruptedBatchMessagePermitTest.java | 127 +- .../MessagePayloadProcessorTest.java | 6 + .../pulsar/client/impl/ConsumerBase.java | 16 +- .../pulsar/client/impl/ConsumerImpl.java | 1081 ++++++++++------- .../pulsar/client/impl/MessageImpl.java | 66 +- .../impl/MessagePayloadContextImpl.java | 31 +- .../client/impl/MessagePermitAccounting.java | 111 -- .../client/impl/MultiTopicsConsumerImpl.java | 45 +- .../client/impl/ZeroQueueConsumerImpl.java | 57 +- .../pulsar/client/impl/ConsumerImplTest.java | 366 +++--- .../impl/MessagePermitAccountingTest.java | 155 --- .../impl/MultiTopicsConsumerImplTest.java | 33 + .../pulsar/common/protocol/Commands.java | 29 +- pulsar-common/src/main/proto/PulsarApi.proto | 3 +- .../common/compression/CommandsTest.java | 22 +- 35 files changed, 1354 insertions(+), 1587 deletions(-) create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessageResult.java delete mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/SendMessageResultTest.java delete mode 100644 pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java delete mode 100644 pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java 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 a8f098bd12313..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,40 +320,29 @@ 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) { - return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, - totalChunkedMessages, redeliveryTracker, DEFAULT_CONSUMER_EPOCH).getWriteFuture(); - } - - public SendMessagesResult sendMessagesWithResult(final List entries, - EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, - int totalMessages, - long totalBytes, - long totalChunkedMessages, - RedeliveryTracker redeliveryTracker) { - return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + return sendMessages(entries, batchSizes, batchIndexesAcks, totalMessages, totalBytes, 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) { - return sendMessagesInternal(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, - totalChunkedMessages, redeliveryTracker, epoch).getWriteFuture(); + return sendMessages(entries, null, batchSizes, batchIndexesAcks, totalMessages, totalBytes, + totalChunkedMessages, redeliveryTracker, epoch); } /** * Dispatch a list of entries to the consumer.
* It is also responsible to release entries data and recycle entries object. * - * @return a future completed after the entries have been written + * @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, @@ -362,21 +351,10 @@ public Future sendMessages(final List entries, long totalChunkedMessages, RedeliveryTracker redeliveryTracker, long epoch) { - return sendMessagesInternal(entries, stickyKeyHashes, batchSizes, batchIndexesAcks, totalMessages, totalBytes, - totalChunkedMessages, redeliveryTracker, epoch).getWriteFuture(); - } - - private SendMessagesResult sendMessagesInternal(final List entries, - final List stickyKeyHashes, - EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, - int totalMessages, - long totalBytes, - long totalChunkedMessages, - RedeliveryTracker redeliveryTracker, - 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(); @@ -385,76 +363,78 @@ private SendMessagesResult sendMessagesInternal(final List entr } final Promise writePromise = cnx.newPromise(); writePromise.setSuccess(null); - return new SendMessagesResult(0).setWriteFuture(writePromise); + sendMessageResult.setSendFuture(writePromise); + return sendMessageResult; } - SendMessagesResult sendResult = new SendMessagesResult(entries.size()); int totalEntries = 0; for (int i = 0; i < entries.size(); i++) { Entry entry = entries.get(i); - if (entry == null) { - continue; - } - 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 stickyKeyHash; - if (stickyKeyHashes == null) { - if (entry instanceof EntryAndMetadata entryAndMetadata) { - stickyKeyHash = entryAndMetadata.getCachedStickyKeyHash(); - } else { - stickyKeyHash = STICKY_KEY_HASH_NOT_SET; - } - } else { - stickyKeyHash = stickyKeyHashes.get(i); - } - boolean sendingAllowed = pendingAcks.addPendingAckIfAllowed( - entry.getLedgerId(), entry.getEntryId(), 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 + if (entry != null) { + int batchSize = batchSizes.getBatchSize(i); + int messagePermits = batchIndexesAcks == null + ? batchSize : batchIndexesAcks.getUnackedIndexCount(i, batchSize); + if (messagePermits == 0) { entries.set(i, null); entry.release(); - log.debug() - .attr("ledgerId", entry.getLedgerId()) - .attr("entryId", entry.getEntryId()) - .attr("batchSize", batchSize) - .log("Skipping sending of entry since adding to pending acks failed"); continue; } - log.debug() - .attr("ledgerId", entry.getLedgerId()) - .attr("entryId", entry.getEntryId()) - .attr("batchSize", batchSize) - .log("Added entry to pendingAcks"); + // 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 stickyKeyHash; + if (stickyKeyHashes == null) { + if (entry instanceof EntryAndMetadata entryAndMetadata) { + stickyKeyHash = entryAndMetadata.getCachedStickyKeyHash(); + } else { + stickyKeyHash = STICKY_KEY_HASH_NOT_SET; + } + } else { + stickyKeyHash = stickyKeyHashes.get(i); + } + boolean sendingAllowed; + sendingAllowed = + pendingAcks.addPendingAckIfAllowed(entry.getLedgerId(), entry.getEntryId(), + 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 + entries.set(i, null); + entry.release(); + log.debug() + .attr("ledgerId", entry.getLedgerId()) + .attr("entryId", entry.getEntryId()) + .attr("batchSize", batchSize) + .log("Skipping sending of entry since adding to pending acks failed"); + } else { + log.debug() + .attr("ledgerId", entry.getLedgerId()) + .attr("entryId", entry.getEntryId()) + .attr("batchSize", batchSize) + .log("Added entry to pendingAcks"); + totalEntries++; + sendMessageResult.recordMessagePermits(i, messagePermits); + } + } else { + totalEntries++; + sendMessageResult.recordMessagePermits(i, messagePermits); + } } - sendResult.setMessagePermits(i, messagePermits); - totalEntries++; } // calculate avg message per entry - if (totalEntries > 0) { - if (avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1 - // set init value. - avgMessagesPerEntry.set(1.0 * totalMessages / totalEntries); - } else { - avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent - + (1 - avgPercent) * totalMessages / totalEntries); - } + int sentMessagePermits = sendMessageResult.getTotalMessagePermits(); + if (totalEntries > 0 && avgMessagesPerEntry.get() < 1) { //valid avgMessagesPerEntry should always >= 1 + // set init value. + avgMessagesPerEntry.set(1.0 * sentMessagePermits / totalEntries); + } else if (totalEntries > 0) { + avgMessagesPerEntry.set(avgMessagesPerEntry.get() * avgPercent + + (1 - avgPercent) * sentMessagePermits / totalEntries); } - int sentMessagePermits = sendResult.getTotalMessagePermits(); - // Reduce permits by the logical messages represented by commands that survived final admission. + // Reduce permits by the message count represented by the commands that will actually be sent. MESSAGE_PERMITS_UPDATER.addAndGet(this, -sentMessagePermits); log.debug() .attr("sentMessagePermits", sentMessagePermits) @@ -464,7 +444,8 @@ private SendMessagesResult sendMessagesInternal(final List entr incrementUnackedMessages(sentMessagePermits); Future writeAndFlushPromise = cnx.getCommandSender().sendMessagesToConsumer(consumerId, topicName, subscription, partitionIdx, - entries, batchSizes, batchIndexesAcks, sendResult, 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()) { @@ -478,11 +459,10 @@ private SendMessagesResult sendMessagesInternal(final List entr } else { log.debug() .exceptionMessage(status.cause()) - .log("Sent messages to client failed by IO exception, disconnecting the consumer"); - disconnect(); + .log("Sent messages to client failed by IO exception, closing the connection"); } }); - return sendResult.setWriteFuture(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 c38b0569af4fe..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") @@ -47,7 +46,7 @@ public int getTotalAckedIndexCount() { 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; 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 6de3c31fc459e..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 @@ -85,8 +85,9 @@ void sendLookupResponse(String brokerServiceUrl, String brokerServiceUrlTls, boo Future sendMessagesToConsumer(long consumerId, String topicName, Subscription subscription, int partitionIdx, List entries, EntryBatchSizes batchSizes, - EntryBatchIndexesAcks batchIndexesAcks, SendMessagesResult sendResult, - RedeliveryTracker redeliveryTracker, long epoch); + EntryBatchIndexesAcks batchIndexesAcks, + 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 437719c329f86..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,8 +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, - SendMessagesResult sendResult, RedeliveryTracker redeliveryTracker, - long epoch) { + RedeliveryTracker redeliveryTracker, long epoch, + SendMessageResult sendMessageResult) { final ChannelHandlerContext ctx = cnx.ctx(); final ChannelPromise writePromise = ctx.newPromise(); ctx.channel().eventLoop().execute(() -> { @@ -295,11 +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, - sendResult.getMessagePermits(i)), + 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/SendMessagesResult.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java deleted file mode 100644 index 26089afcf2410..0000000000000 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SendMessagesResult.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * 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; - -/** - * Contains the finalized permit accounting and asynchronous write completion for a send operation. - * - *

The broker consumer populates this object synchronously after final send admission. The command sender reads - * the per-entry values, while the persistent Shared dispatcher reads their sum. This keeps every covered broker - * counter and command serialization tied to the same finalized result. - */ -public final class SendMessagesResult { - private final int[] messagePermits; - private int totalMessagePermits; - private Future writeFuture; - - SendMessagesResult(int entriesListSize) { - if (entriesListSize < 0) { - throw new IllegalArgumentException("entriesListSize must not be negative"); - } - messagePermits = new int[entriesListSize]; - } - - void setMessagePermits(int entryIdx, int permits) { - if (permits <= 0) { - throw new IllegalArgumentException("message permits must be positive"); - } - if (messagePermits[entryIdx] != 0) { - throw new IllegalStateException("Permits already finalized for entry " + entryIdx); - } - int updatedTotalMessagePermits = Math.addExact(totalMessagePermits, permits); - messagePermits[entryIdx] = permits; - totalMessagePermits = updatedTotalMessagePermits; - } - - int getMessagePermits(int entryIdx) { - return messagePermits[entryIdx]; - } - - SendMessagesResult setWriteFuture(Future writeFuture) { - if (this.writeFuture != null) { - throw new IllegalStateException("Write future already set"); - } - this.writeFuture = writeFuture; - return this; - } - - public int getTotalMessagePermits() { - return totalMessagePermits; - } - - public Future getWriteFuture() { - return writeFuture; - } -} 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 1411465e4d800..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,7 +4489,7 @@ 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 serializeMessageAndIntercept(command, consumerId, ledgerId, entryId, metadataAndPayload); + return newMessageAndIntercept(consumerId, ledgerId, entryId, metadataAndPayload, topic, command); } public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long entryId, int partition, @@ -4497,11 +4497,11 @@ public ByteBufPair newMessageAndIntercept(long consumerId, long ledgerId, long e int messagePermits) { BaseCommand command = Commands.newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, epoch, messagePermits); - return serializeMessageAndIntercept(command, consumerId, ledgerId, entryId, metadataAndPayload); + return newMessageAndIntercept(consumerId, ledgerId, entryId, metadataAndPayload, topic, command); } - private ByteBufPair serializeMessageAndIntercept(BaseCommand command, long consumerId, long ledgerId, long entryId, - ByteBuf metadataAndPayload) { + 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 bd021357d0cf7..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,7 +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.SendMessagesResult; +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; @@ -870,8 +870,7 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis readType == ReadType.Replay, c); totalEntriesProcessed += entriesForThisConsumer.size(); - SendMessagesResult sendResult = c.sendMessagesWithResult( - entriesForThisConsumer, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); @@ -978,11 +977,11 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); totalEntriesProcessed += entryAndMetadataList.size(); - SendMessagesResult sendResult = consumer.sendMessagesWithResult( - entryAndMetadataList, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), - sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker()); - sendResult.getWriteFuture().addListener(future -> { + sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() + ); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntriesAsync(); } 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 866d6f2cc6b5a..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,7 +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.SendMessagesResult; +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; @@ -760,8 +760,7 @@ protected synchronized boolean trySendMessagesToConsumers(ReadType readType, Lis entriesForThisConsumer, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, c); - SendMessagesResult sendResult = c.sendMessagesWithResult( - entriesForThisConsumer, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = c.sendMessages(entriesForThisConsumer, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), sendMessageInfo.getTotalChunkedMessages(), redeliveryTracker); @@ -838,11 +837,11 @@ private boolean sendChunkedMessagesToConsumers(ReadType readType, totalEntries += filterEntriesForConsumer(entryAndMetadataList, batchSizes, sendMessageInfo, batchIndexesAcks, cursor, readType == ReadType.Replay, consumer); - SendMessagesResult sendResult = consumer.sendMessagesWithResult( - entryAndMetadataList, batchSizes, batchIndexesAcks, + SendMessageResult sendResult = consumer.sendMessages(entryAndMetadataList, batchSizes, batchIndexesAcks, sendMessageInfo.getTotalMessages(), sendMessageInfo.getTotalBytes(), - sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker()); - sendResult.getWriteFuture().addListener(future -> { + sendMessageInfo.getTotalChunkedMessages(), getRedeliveryTracker() + ); + sendResult.getSendFuture().addListener(future -> { if (future.isDone() && numConsumers.decrementAndGet() == 0) { readMoreEntries(); } 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 b5e7a1793bebf..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 @@ -28,23 +28,16 @@ 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.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 static org.testng.Assert.assertSame; -import static org.testng.Assert.expectThrows; import io.netty.util.concurrent.ImmediateEventExecutor; -import io.netty.util.concurrent.Promise; -import java.io.IOException; import java.net.SocketAddress; import java.util.ArrayList; -import java.util.BitSet; import java.util.List; -import java.util.Optional; -import java.util.SplittableRandom; import org.apache.bookkeeper.mledger.Entry; import org.apache.commons.lang3.tuple.Pair; import org.apache.pulsar.broker.PulsarService; @@ -52,17 +45,15 @@ 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.mockito.ArgumentCaptor; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @Test(groups = "broker") public class ConsumerTest { - private static final long RANDOM_SEED = 0x491C05E5L; - private Consumer consumer; private Subscription subscription; private ServerCnx cnx; + private Topic topic; private final ConsumerStatsImpl stats = new ConsumerStatsImpl(); @BeforeMethod @@ -70,13 +61,12 @@ public void beforeMethod() { subscription = mock(Subscription.class); cnx = mock(ServerCnx.class); SocketAddress address = mock(SocketAddress.class); - PersistentTopic topic = mock(PersistentTopic.class); + topic = mock(PersistentTopic.class, RETURNS_DEEP_STUBS); BrokerService brokerService = mock(BrokerService.class); PulsarService pulsarService = mock(PulsarService.class); ServiceConfiguration serviceConfiguration = mock(ServiceConfiguration.class); when(cnx.clientAddress()).thenReturn(address); - when(subscription.getName()).thenReturn("subscription"); when(subscription.getTopic()).thenReturn(topic); when(topic.getBrokerService()).thenReturn(brokerService); when(brokerService.getPulsar()).thenReturn(pulsarService); @@ -103,8 +93,9 @@ public void testGetBytesOutCounter() { @Test public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { - Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, - "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, + 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); @@ -119,8 +110,6 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { when(emptyPartialBatch.getLedgerId()).thenReturn(3L); when(emptyPartialBatch.getEntryId()).thenReturn(3L); List entries = new ArrayList<>(List.of(partialBatch, rejectedBatch, emptyPartialBatch)); - // A null slot represents an entry removed by dispatcher filtering before final send admission. - entries.add(null); EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); batchSizes.setBatchSize(0, 10); @@ -133,219 +122,22 @@ public void testSendMessagesFinalizesPermitsAfterPendingAckAdmission() { PulsarCommandSender commandSender = mock(PulsarCommandSender.class); when(cnx.getCommandSender()).thenReturn(commandSender); when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), - any(), any(), anyLong())) - .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); + any(), anyLong(), any())).thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); try { - SendMessagesResult sendResult = sharedConsumer.sendMessagesWithResult( - entries, batchSizes, batchIndexesAcks, 23, 0, 0, mock(RedeliveryTracker.class)); + SendMessageResult result = sharedConsumer.sendMessages(entries, batchSizes, batchIndexesAcks, + 23, 0, 0, mock(RedeliveryTracker.class)); - assertEquals(sendResult.getTotalMessagePermits(), 3); - assertEquals(sendResult.getMessagePermits(0), 3); - assertEquals(sendResult.getMessagePermits(1), 0); - assertEquals(sendResult.getMessagePermits(2), 0); - assertEquals(sendResult.getMessagePermits(3), 0); + 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(); - ArgumentCaptor sendResultCaptor = ArgumentCaptor.forClass(SendMessagesResult.class); - verify(commandSender).sendMessagesToConsumer(eq(2L), eq("topic"), eq(subscription), anyInt(), eq(entries), - eq(batchSizes), eq(batchIndexesAcks), sendResultCaptor.capture(), any(), - eq(DEFAULT_CONSUMER_EPOCH)); - assertSame(sendResultCaptor.getValue(), sendResult); - assertSame(entries.get(0), partialBatch); } finally { batchSizes.recyle(); batchIndexesAcks.recycle(); } } - - @Test - public void testFinalizedPermitsOutliveRecycledSenderInputs() { - Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, - "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, - DEFAULT_CONSUMER_EPOCH); - sharedConsumer.setPendingAcksAddHandler((ignored, ledgerId, entryId, stickyKeyHash) -> true); - sharedConsumer.flowPermits(100); - - Entry partialBatch = mock(Entry.class); - when(partialBatch.getLedgerId()).thenReturn(1L); - when(partialBatch.getEntryId()).thenReturn(1L); - Entry completeBatch = mock(Entry.class); - when(completeBatch.getLedgerId()).thenReturn(2L); - when(completeBatch.getEntryId()).thenReturn(2L); - List entries = new ArrayList<>(List.of(partialBatch, completeBatch)); - - EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); - batchSizes.setBatchSize(0, 10); - batchSizes.setBatchSize(1, 4); - EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); - batchIndexesAcks.setIndexesAcks(0, Pair.of(10, new long[] {0b100101L})); - - PulsarCommandSender commandSender = mock(PulsarCommandSender.class); - when(cnx.getCommandSender()).thenReturn(commandSender); - Promise pendingWrite = ImmediateEventExecutor.INSTANCE.newPromise(); - when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), - any(), any(), anyLong())).thenReturn(pendingWrite); - - SendMessagesResult sendResult; - try { - sendResult = sharedConsumer.sendMessagesWithResult( - entries, batchSizes, batchIndexesAcks, 14, 0, 0, mock(RedeliveryTracker.class)); - } finally { - // The command sender owns and recycles these inputs before its asynchronous write completes. - batchSizes.recyle(); - batchIndexesAcks.recycle(); - } - - EntryBatchSizes reusedBatchSizes = EntryBatchSizes.get(entries.size()); - EntryBatchIndexesAcks reusedBatchIndexesAcks = EntryBatchIndexesAcks.get(entries.size()); - try { - // Mutate the recycled carriers while the write is still pending. The finalized result must be independent. - reusedBatchSizes.setBatchSize(0, 1); - reusedBatchSizes.setBatchSize(1, 1); - reusedBatchIndexesAcks.setIndexesAcks(0, Pair.of(1, new long[] {1L})); - - assertEquals(sendResult.getMessagePermits(0), 3); - assertEquals(sendResult.getMessagePermits(1), 4); - assertEquals(sendResult.getTotalMessagePermits(), 7); - } finally { - pendingWrite.trySuccess(null); - reusedBatchSizes.recyle(); - reusedBatchIndexesAcks.recycle(); - } - } - - @Test - public void testSendMessagesWriteFailureDisconnectsConsumer() throws BrokerServiceException { - Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, - "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, - DEFAULT_CONSUMER_EPOCH); - sharedConsumer.setPendingAcksAddHandler((ignored, ledgerId, entryId, stickyKeyHash) -> true); - sharedConsumer.flowPermits(100); - - Entry entry = mock(Entry.class); - when(entry.getLedgerId()).thenReturn(1L); - when(entry.getEntryId()).thenReturn(1L); - List entries = new ArrayList<>(List.of(entry)); - EntryBatchSizes batchSizes = EntryBatchSizes.get(1); - batchSizes.setBatchSize(0, 1); - PulsarCommandSender commandSender = mock(PulsarCommandSender.class); - when(cnx.getCommandSender()).thenReturn(commandSender); - when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), - any(), any(), anyLong())) - .thenReturn(ImmediateEventExecutor.INSTANCE.newFailedFuture(new IOException("write failed"))); - - try { - sharedConsumer.sendMessagesWithResult( - entries, batchSizes, null, 1, 0, 0, mock(RedeliveryTracker.class)); - - verify(cnx).closeConsumer(sharedConsumer, Optional.empty()); - verify(subscription).removeConsumer(sharedConsumer, false); - verify(cnx).removedConsumer(sharedConsumer); - } finally { - batchSizes.recyle(); - } - } - - @Test - public void testRandomizedFinalAdmissionUsesOnePermitResultForAllBrokerCounters() { - final int entriesCount = 1_000; - final int initialPermits = 1_000_000; - SplittableRandom random = new SplittableRandom(RANDOM_SEED); - boolean[] admissionAllowed = new boolean[entriesCount]; - int[] expectedPermits = new int[entriesCount]; - Entry[] originalEntries = new Entry[entriesCount]; - List entries = new ArrayList<>(entriesCount); - EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesCount); - EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(entriesCount); - int totalMessages = 0; - int expectedTotalPermits = 0; - - for (int i = 0; i < entriesCount; i++) { - int batchSize = random.nextInt(1, 257); - batchSizes.setBatchSize(i, batchSize); - totalMessages += batchSize; - if (random.nextInt(8) == 0) { - entries.add(null); - continue; - } - - Entry entry = mock(Entry.class); - originalEntries[i] = entry; - entries.add(entry); - when(entry.getLedgerId()).thenReturn((long) i); - when(entry.getEntryId()).thenReturn((long) i); - - int messagePermits = batchSize; - if (random.nextInt(4) != 0) { - int requiredWords = (batchSize + Long.SIZE - 1) / Long.SIZE; - long[] ackSet = new long[requiredWords + random.nextInt(2)]; - for (int word = 0; word < ackSet.length; word++) { - ackSet[word] = random.nextLong(); - } - BitSet boundedAckSet = BitSet.valueOf(ackSet); - boundedAckSet.clear(batchSize, Math.max(batchSize, boundedAckSet.length())); - messagePermits = boundedAckSet.cardinality(); - batchIndexesAcks.setIndexesAcks(i, Pair.of(batchSize, ackSet)); - } - - admissionAllowed[i] = random.nextInt(5) != 0; - if (messagePermits > 0 && admissionAllowed[i]) { - expectedPermits[i] = messagePermits; - expectedTotalPermits += messagePermits; - } - } - - Consumer sharedConsumer = new Consumer(subscription, Shared, "topic", 2, 0, "shared-consumer", false, cnx, - "myrole-1", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), latest, - DEFAULT_CONSUMER_EPOCH); - sharedConsumer.setPendingAcksAddHandler( - (ignored, ledgerId, entryId, stickyKeyHash) -> admissionAllowed[(int) ledgerId]); - sharedConsumer.flowPermits(initialPermits); - PulsarCommandSender commandSender = mock(PulsarCommandSender.class); - when(cnx.getCommandSender()).thenReturn(commandSender); - when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), any(), any(), any(), - any(), any(), anyLong())) - .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); - - try { - SendMessagesResult sendResult = sharedConsumer.sendMessagesWithResult( - entries, batchSizes, batchIndexesAcks, totalMessages, 0, 0, mock(RedeliveryTracker.class)); - - assertEquals(sendResult.getTotalMessagePermits(), expectedTotalPermits, "seed=" + RANDOM_SEED); - assertEquals(sharedConsumer.getAvailablePermits(), initialPermits - expectedTotalPermits, - "seed=" + RANDOM_SEED); - assertEquals(sharedConsumer.getUnackedMessages(), expectedTotalPermits, "seed=" + RANDOM_SEED); - for (int i = 0; i < entriesCount; i++) { - assertEquals(sendResult.getMessagePermits(i), expectedPermits[i], - "seed=" + RANDOM_SEED + ", entry=" + i); - if (originalEntries[i] != null && expectedPermits[i] == 0) { - assertNull(entries.get(i)); - verify(originalEntries[i]).release(); - } else { - assertSame(entries.get(i), originalEntries[i]); - } - } - } finally { - batchSizes.recyle(); - batchIndexesAcks.recycle(); - } - } - - @Test - public void testSendMessagesResultRejectsPartialFinalization() { - SendMessagesResult sendResult = new SendMessagesResult(2); - sendResult.setMessagePermits(0, Integer.MAX_VALUE); - - expectThrows(ArithmeticException.class, () -> sendResult.setMessagePermits(1, 1)); - assertEquals(sendResult.getMessagePermits(0), Integer.MAX_VALUE); - assertEquals(sendResult.getMessagePermits(1), 0); - assertEquals(sendResult.getTotalMessagePermits(), Integer.MAX_VALUE); - expectThrows(IllegalStateException.class, () -> sendResult.setMessagePermits(0, 1)); - expectThrows(IllegalArgumentException.class, () -> new SendMessagesResult(1).setMessagePermits(0, 0)); - } } 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 58b200e591be0..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 @@ -30,18 +30,21 @@ public class EntryBatchIndexesAcksTest { private static final long RANDOM_SEED = 0x491B0A7EL; @Test - void shouldCalculateUnackedIndexesWithinBatchBounds() { + void shouldCalculateUnackedIndexesForEachEntry() { BitSetRecyclable bitSet = BitSetRecyclable.create(); bitSet.set(1); bitSet.set(4); bitSet.set(8); - bitSet.set(63); 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(); 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 index b218653900fbf..7b7f96983cdc7 100644 --- 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 @@ -104,14 +104,14 @@ public void testSendMessagesMapsFinalizedPerEntryPermitsToWireCommands() { EntryBatchSizes batchSizes = EntryBatchSizes.get(entries.size()); batchSizes.setBatchSize(0, 10); batchSizes.setBatchSize(2, 4); - SendMessagesResult sendResult = new SendMessagesResult(entries.size()); - sendResult.setMessagePermits(0, 3); - sendResult.setMessagePermits(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, - sendResult, mock(RedeliveryTracker.class), 11); + mock(RedeliveryTracker.class), 11, sendResult); assertEquals(serializedPermits, List.of(3, 4)); verify(first).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 index b5ecb2aba71a2..4c0afcaa09f75 100644 --- 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 @@ -18,20 +18,12 @@ */ package org.apache.pulsar.broker.service.persistent; -import io.netty.buffer.ByteBuf; -import io.netty.channel.ChannelHandlerContext; -import io.netty.channel.ChannelOutboundHandlerAdapter; -import io.netty.channel.ChannelPromise; -import io.netty.util.ReferenceCountUtil; -import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; import org.apache.pulsar.broker.service.Consumer; -import org.apache.pulsar.broker.service.ServerCnx; import org.apache.pulsar.broker.service.SharedPulsarBaseTest; import org.apache.pulsar.client.api.Message; import org.apache.pulsar.client.api.MessageId; @@ -105,52 +97,6 @@ public void testBatchDebtIsSharedByConsumerAndDispatcherAndRemovedWithConsumer(b } } - @Test(timeOut = 30000, dataProvider = "sharedDispatcherImplementations") - public void testWriteFailureRemovesDebitedConsumer(boolean classicDispatcher) throws Exception { - boolean originalSetting = getConfig().isSubscriptionSharedUseClassicPersistentImplementation(); - getConfig().setSubscriptionSharedUseClassicPersistentImplementation(classicDispatcher); - String topicName = newTopicName(); - String subscriptionName = "write-failure-subscription"; - - try (org.apache.pulsar.client.api.Consumer clientConsumer = - pulsarClient.newConsumer(Schema.INT32) - .topic(topicName) - .subscriptionName(subscriptionName) - .subscriptionType(SubscriptionType.Shared) - .receiverQueueSize(1) - .subscribe(); - Producer producer = pulsarClient.newProducer(Schema.INT32) - .topic(topicName) - .enableBatching(true) - .batchingMaxMessages(BATCH_SIZE) - .batchingMaxPublishDelay(1, TimeUnit.HOURS) - .create()) { - PersistentSubscription subscription = getSubscription(topicName, subscriptionName); - Consumer originalConsumer = subscription.getConsumers().get(0); - ServerCnx serverCnx = (ServerCnx) originalConsumer.cnx(); - AtomicBoolean writeFailed = new AtomicBoolean(); - failNextFlush(serverCnx, writeFailed); - - sendBatch(producer, 0); - - Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { - Assert.assertTrue(writeFailed.get()); - Assert.assertFalse(subscription.getConsumers().contains(originalConsumer)); - CompletableFuture registeredConsumer = - serverCnx.getConsumers().get(originalConsumer.consumerId()); - Assert.assertTrue(registeredConsumer == null || registeredConsumer.getNow(null) != originalConsumer); - }); - - 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")); @@ -186,21 +132,4 @@ private static int getTotalAvailablePermits(PersistentSubscription subscription) } throw new AssertionError("Unexpected dispatcher " + subscription.getDispatcher()); } - - private static void failNextFlush(ServerCnx serverCnx, AtomicBoolean writeFailed) throws Exception { - serverCnx.ctx().channel().eventLoop().submit(() -> serverCnx.ctx().pipeline().addFirst( - new ChannelOutboundHandlerAdapter() { - @Override - public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception { - if (msg instanceof ByteBuf buffer && !buffer.isReadable() - && writeFailed.compareAndSet(false, true)) { - ctx.pipeline().remove(this); - ReferenceCountUtil.safeRelease(msg); - promise.setFailure(new IOException("intentional message write failure")); - return; - } - ctx.write(msg, promise); - } - })).sync(); - } } 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 ba4d86ce1c5d4..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 @@ -21,6 +21,7 @@ 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; @@ -37,8 +38,9 @@ public class CompactedOutBatchMessageTest extends SharedPulsarBaseTest { @Test - public void testBatchRemainderIsNotReturnedToReplacementConnection() throws Exception { + public void testDoesNotReturnBatchPermitsToDifferentConnection() throws Exception { final int batchSize = 2; + final String topic = newTopicName(); MessageMetadata metadata = new MessageMetadata() .setProducerName("foobar") .setSequenceId(1) @@ -51,10 +53,12 @@ public void testBatchRemainderIsNotReturnedToReplacementConnection() throws Exce } try (ConsumerImpl consumer = - (ConsumerImpl) pulsarClient.newConsumer().topic(newTopicName()) + (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), @@ -67,32 +71,40 @@ public void testBatchRemainderIsNotReturnedToReplacementConnection() throws Exce } @Test - public void testStaleEpochDiscardDoesNotReturnPermitToReplacementConnection() throws Exception { + 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(1); - ByteBuf batch = Unpooled.buffer(100); - ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); - Commands.serializeSingleMessageInBatchWithPayload(new SingleMessageMetadata(), payload, batch); - payload.release(); + .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(newTopicName()) - .subscriptionName("stale-epoch-subscription") + (ConsumerImpl) pulsarClient.newConsumer().topic(topic) + .subscriptionName("partial-batch-subscription") .receiverQueueSize(20) .subscribe()) { - ConsumerBase.CONSUMER_EPOCH.set(consumer, 2); - int permitsBefore = consumer.getAvailablePermits(); - consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, batch, - new MessageIdData().setLedgerId(1234).setEntryId(567), mock(ClientCnx.class), 1, false, 1); - consumer.internalPinnedExecutor.submit(() -> assertEquals(consumer.numMessagesInQueue(), 0)) - .get(5, TimeUnit.SECONDS); + consumer.receiveIndividualMessagesFromBatch(null, metadata, 0, null, truncatedBatch, + new MessageIdData().setLedgerId(1234).setEntryId(567), consumer.cnx(), + DEFAULT_CONSUMER_EPOCH, false, batchSize); - assertEquals(consumer.getAvailablePermits(), permitsBefore); + for (int i = 0; i < parsedMessages; i++) { + Message message = consumer.receive(5, TimeUnit.SECONDS); + assertNotNull(message); + message.release(); + } + assertEquals(consumer.getAvailablePermits(), batchSize); } finally { - batch.release(); + truncatedBatch.release(); } } @@ -125,12 +137,11 @@ public void testCompactedOutMessages() throws Exception { try (ConsumerImpl consumer = (ConsumerImpl) pulsarClient.newConsumer().topic(topic1) - .subscriptionName("my-subscriber-name").subscribe()) { - int permitsBefore = consumer.getAvailablePermits(); + .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), - consumer.cnx(), DEFAULT_CONSUMER_EPOCH, false, metadata.getNumMessagesInBatch()); + consumer.cnx(), DEFAULT_CONSUMER_EPOCH, false); Message m = consumer.receive(); assertEquals(((BatchMessageIdImpl) m.getMessageId()).getLedgerId(), 1234); assertEquals(((BatchMessageIdImpl) m.getMessageId()).getEntryId(), 567); @@ -138,8 +149,9 @@ batchBuffer, new MessageIdData().setLedgerId(1234).setEntryId(567), assertEquals(m.getKey(), "key3"); assertEquals(consumer.numMessagesInQueue(), 0); - assertEquals(consumer.getAvailablePermits(), permitsBefore + metadata.getNumMessagesInBatch()); - m.release(); + 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 index 242165d2495b2..d451abd710e4a 100644 --- 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 @@ -18,163 +18,74 @@ */ package org.apache.pulsar.client.impl; -import static java.nio.charset.StandardCharsets.UTF_8; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; -import java.time.Duration; 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.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.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.api.proto.SingleMessageMetadata; import org.apache.pulsar.common.protocol.Commands; -import org.awaitility.Awaitility; import org.testng.annotations.Test; @Test(groups = "broker-impl") public class CorruptedBatchMessagePermitTest extends SharedPulsarBaseTest { @Test(timeOut = 30000) - public void testSharedConsumerReturnsAllPermitsForRepeatedCorruptedBatches() throws Exception { + public void testSharedConsumerReturnsPermitsForCorruptedBatch() throws Exception { final int batchSize = 10; - final int corruptedBatches = 3; - final String subscriptionName = "shared-subscription"; - final String topicName = newTopicName(); + final String topic = newTopicName(); admin.namespaces().setDeduplicationStatus(getNamespace(), false); @Cleanup Consumer consumer = pulsarClient.newConsumer(Schema.STRING) - .topic(topicName) - .subscriptionName(subscriptionName) + .topic(topic) + .subscriptionName("shared-subscription") .subscriptionType(SubscriptionType.Shared) - .receiverQueueSize(1) + .receiverQueueSize(batchSize) .subscribe(); @Cleanup Producer producer = pulsarClient.newProducer(Schema.STRING) - .topic(topicName) + .topic(topic) .enableBatching(false) .create(); - PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get() + PersistentTopic persistentTopic = (PersistentTopic) getTopicIfExists(topic).get() .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); - PersistentSubscription subscription = topic.getSubscription(subscriptionName); - for (int i = 1; i <= corruptedBatches; i++) { - publishCorruptedBatch(topic, batchSize, i).get(10, TimeUnit.SECONDS); - awaitConsumerPermitBalance(subscription, (long) batchSize * i); - } - - producer.send("message-after-corrupted-batches"); - Message message = consumer.receive(10, TimeUnit.SECONDS); - assertNotNull(message); - try { - assertEquals(message.getValue(), "message-after-corrupted-batches"); - } finally { - message.release(); - } - } - - @Test(timeOut = 30000) - public void testSharedConsumerReturnsOnlyUnparsedRemainderForMalformedBatch() throws Exception { - final int batchSize = 5; - final int parsedMessages = 2; - final String subscriptionName = "shared-partial-batch-subscription"; - final String topicName = newTopicName(); - admin.namespaces().setDeduplicationStatus(getNamespace(), false); - - @Cleanup - Consumer consumer = pulsarClient.newConsumer(Schema.STRING) - .topic(topicName) - .subscriptionName(subscriptionName) - .subscriptionType(SubscriptionType.Shared) - .receiverQueueSize(1) - .subscribe(); - @Cleanup - Producer producer = pulsarClient.newProducer(Schema.STRING) - .topic(topicName) - .enableBatching(false) - .create(); - - PersistentTopic topic = (PersistentTopic) getTopicIfExists(topicName).get() - .orElseThrow(() -> new IllegalStateException("Topic was not loaded")); - PersistentSubscription subscription = topic.getSubscription(subscriptionName); - publishMalformedBatch(topic, batchSize, parsedMessages).get(10, TimeUnit.SECONDS); + publishCorruptedCompressedBatch(persistentTopic, batchSize).get(10, TimeUnit.SECONDS); + producer.send("message-after-corrupted-batch"); - for (int i = 0; i < parsedMessages; i++) { - Message message = consumer.receive(10, TimeUnit.SECONDS); - assertNotNull(message); - try { - assertEquals(message.getValue(), "partial-" + i); - } finally { - message.release(); - } - } - awaitConsumerPermitBalance(subscription, batchSize); - - producer.send("message-after-malformed-batch"); Message message = consumer.receive(10, TimeUnit.SECONDS); assertNotNull(message); try { - assertEquals(message.getValue(), "message-after-malformed-batch"); + assertEquals(message.getValue(), "message-after-corrupted-batch"); } finally { message.release(); } } - private static void awaitConsumerPermitBalance(PersistentSubscription subscription, long expectedMessagesOut) { - Awaitility.await().atMost(Duration.ofSeconds(10)).untilAsserted(() -> { - assertEquals(subscription.getConsumers().size(), 1); - org.apache.pulsar.broker.service.Consumer brokerConsumer = subscription.getConsumers().get(0); - assertEquals(brokerConsumer.getMsgOutCounter(), expectedMessagesOut); - assertEquals(brokerConsumer.getAvailablePermits(), 1); - }); - } - - private static CompletableFuture publishCorruptedBatch(PersistentTopic topic, int batchSize, - long sequenceId) { + private static CompletableFuture publishCorruptedCompressedBatch(PersistentTopic topic, int batchSize) { MessageMetadata metadata = new MessageMetadata() .setProducerName("corrupted-batch-producer") - .setSequenceId(sequenceId) - .setPublishTime(System.currentTimeMillis()) - .setNumMessagesInBatch(batchSize); - ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); - ByteBuf entry = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c, metadata, payload); - payload.release(); - int lastByteIndex = entry.writerIndex() - 1; - entry.setByte(lastByteIndex, entry.getByte(lastByteIndex) ^ 1); - - return publish(topic, entry); - } - - private static CompletableFuture publishMalformedBatch(PersistentTopic topic, int batchSize, - int parsedMessages) { - MessageMetadata metadata = new MessageMetadata() - .setProducerName("malformed-batch-producer") .setSequenceId(0) .setPublishTime(System.currentTimeMillis()) - .setNumMessagesInBatch(batchSize); - ByteBuf payload = Unpooled.buffer(); - for (int i = 0; i < parsedMessages; i++) { - ByteBuf singlePayload = Unpooled.copiedBuffer("partial-" + i, UTF_8); - Commands.serializeSingleMessageInBatchWithPayload( - new SingleMessageMetadata(), singlePayload, payload); - singlePayload.release(); - } - ByteBuf entry = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c, metadata, payload); - payload.release(); - return publish(topic, entry); - } + .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(); - private static CompletableFuture publish(PersistentTopic topic, ByteBuf entry) { CompletableFuture result = new CompletableFuture<>(); topic.publishMessage(entry, (error, ledgerId, entryId) -> { if (error == null) { 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 49c11590744b7..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; @@ -98,8 +98,6 @@ import org.apache.pulsar.client.api.TopicMessageId; import org.apache.pulsar.client.api.TypedMessageBuilder; import org.apache.pulsar.client.api.transaction.TxnID; -import org.apache.pulsar.client.impl.MessagePermitAccounting.Budget; -import org.apache.pulsar.client.impl.MessagePermitAccounting.InvalidMessagePermitsException; import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData; import org.apache.pulsar.client.impl.crypto.MessageCryptoBc; import org.apache.pulsar.client.impl.metrics.Counter; @@ -151,24 +149,25 @@ public class ConsumerImpl extends ConsumerBase implements ConnectionHandle final long consumerId; - // Permit returns are scoped to one local broker-consumer incarnation. ClientCnx cannot be used as the identity - // because a pooled physical connection can outlive a broker-side consumer and then be reused by its replacement. - private volatile ConsumerPermitState permitState = new ConsumerPermitState(false); - - @SuppressWarnings("rawtypes") - private static final AtomicIntegerFieldUpdater AVAILABLE_PERMITS_UPDATER = - AtomicIntegerFieldUpdater.newUpdater(ConsumerPermitState.class, "availablePermits"); - + /** + * 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 { - private volatile boolean flowEnabled; - @SuppressWarnings("unused") - private volatile int availablePermits; + final ClientCnx cnx; + final AtomicInteger availablePermits = new AtomicInteger(); + final AtomicBoolean flowTaskScheduled = new AtomicBoolean(); + volatile boolean active = true; + volatile boolean flowEnabled; - private ConsumerPermitState(boolean flowEnabled) { - this.flowEnabled = 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; @@ -907,15 +906,15 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { + client.getConfiguration().getOperationTimeoutMs()); int currentSize; - ConsumerPermitState connectionPermitState; + final ConsumerPermitState newPermitState; synchronized (this) { currentSize = incomingMessages.size(); + newPermitState = replacePermitState(cnx); setClientCnx(cnx); - connectionPermitState = permitState; - clearReceiverQueue(true); if (possibleSendToDeadLetterTopicMessages != null) { - possibleSendToDeadLetterTopicMessages.clear(); + clearPossibleDeadLetterMessages(); } + clearReceiverQueue(true); } boolean isDurable = subscriptionMode == SubscriptionMode.Durable; @@ -960,8 +959,12 @@ public CompletableFuture connectionOpened(final ClientCnx cnx) { cnx.sendRequestWithId(request, requestId).thenRun(() -> { synchronized (ConsumerImpl.this) { - if (changeToReadyState()) { - consumerIsReconnectedToBroker(cnx, currentSize, connectionPermitState); + 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 @@ -980,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) { - updateAvailablePermits(connectionPermitState, 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 @@ -1073,22 +1081,8 @@ protected void closeWhenReceivedUnrecoverableError(Throwable t, ClientCnx cnx) { }); } - protected synchronized void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize) { + protected void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize) { log.info("Subscribed to topic"); - ConsumerPermitState connectionPermitState = permitState; - if (cnx() == cnx) { - AVAILABLE_PERMITS_UPDATER.set(connectionPermitState, 0); - // Flow becomes valid only after Subscribe succeeds. This prevents a receiver-queue update racing with - // connection setup from writing Flow before the broker has created this consumer incarnation. - connectionPermitState.flowEnabled = true; - } - } - - private synchronized void consumerIsReconnectedToBroker(ClientCnx cnx, int currentQueueSize, - ConsumerPermitState connectionPermitState) { - if (permitState == connectionPermitState) { - consumerIsReconnectedToBroker(cnx, currentQueueSize); - } } /** @@ -1139,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 @@ -1275,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) { @@ -1323,28 +1390,26 @@ protected MessageImpl newSingleMessage(final int index, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { - ClientCnx messageCnx = cnx(); - return newSingleMessage(index, numMessages, brokerEntryMetadata, msgMetadata, singleMessageMetadata, payload, - messageId, schema, containMetadata, ackBitSet, ackSetInMessageId, redeliveryCount, consumerEpoch, - isEncrypted, messageCnx, getPermitState(messageCnx)); - } - - private 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, - final ConsumerPermitState messagePermitState) { + 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"); @@ -1382,9 +1447,8 @@ private MessageImpl newSingleMessage(final int index, createEncryptionContext(msgMetadata, isEncrypted), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); - message.setPermitState(messagePermitState); return message; - } catch (IOException | IllegalStateException e) { + } catch (IOException | RuntimeException e) { throw new IllegalStateException(e); } finally { if (singleMessagePayload != null) { @@ -1412,70 +1476,123 @@ protected MessageImpl newMessage(final MessageIdImpl messageId, final int redeliveryCount, final long consumerEpoch, final boolean isEncrypted) { - ClientCnx messageCnx = cnx(); return newMessage(messageId, brokerEntryMetadata, messageMetadata, payload, schema, redeliveryCount, - consumerEpoch, isEncrypted, messageCnx, getPermitState(messageCnx)); - } - - private 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 ConsumerPermitState messagePermitState) { + 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), messageCnx, schema, redeliveryCount, poolMessages, consumerEpoch); message.setBrokerEntryMetadata(brokerEntryMetadata); - message.setPermitState(messagePermitState); return message; } private void executeNotifyCallback(final MessageImpl message) { - ConsumerPermitState messagePermitState = message.getPermitState(); - int messageSize = message.size(); messagesPrefetchedGauge.increment(); - bytesPrefetchedGauge.add(messageSize); + bytesPrefetchedGauge.add(message.size()); // 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)) { - // isValidConsumerEpoch has already released the stale message. - completeStaleEpochDiscardAccounting(messagePermitState, messageSize); - 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 completeStaleEpochDiscardAccounting(ConsumerPermitState messagePermitState, int messageSize) { - messagesPrefetchedGauge.decrement(); - bytesPrefetchedGauge.subtract(messageSize); - returnPermitsToSourceConsumer(messagePermitState, 1); + 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, + final MessageMetadata messageMetadata, + final ByteBuf byteBuf, + final MessageIdImpl messageId, + final Schema schema, + 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) { - ClientCnx messageCnx = cnx(); + 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, getPermitState(messageCnx)); + ackSet, consumerEpoch, messagePermits, getPermitState(cnx())); } protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMetadata, @@ -1486,17 +1603,21 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe final int redeliveryCount, final List ackSet, long consumerEpoch, - final ConsumerPermitState messagePermitState) { + 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) { MessageImpl messageImpl = (MessageImpl) message; - messageImpl.setPermitState(messagePermitState); - executeNotifyCallback(messageImpl); + if (messageImpl.getPermitState() == null) { + messageImpl.setFlowPermitOwnership(sourceState, 1); + } + executeNotifyCallbackForPayloadProcessor(messageImpl); } else { skippedMessages.incrementAndGet(); } @@ -1504,21 +1625,31 @@ protected void processPayloadByProcessor(final BrokerEntryMetadata brokerEntryMe } catch (Throwable throwable) { log.warn().exception(throwable) .log("unable to obtain message in batch"); - discardCorruptedMessage(messageId, messagePermitState, ValidationError.BatchDeSerializeError); + discardCorruptedMessage(messageId, sourceState, ValidationError.BatchDeSerializeError, 1); } finally { entryContext.recycle(); payload.release(); // byteBuf.release() is called in this method } if (skippedMessages.get() > 0) { - returnPermitsToSourceConsumer(messagePermitState, skippedMessages.get()); + increaseAvailablePermitsForReceivedMessage(sourceState, skippedMessages.get()); } tryTriggerListener(); } void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, ClientCnx cnx) { - final ConsumerPermitState messagePermitState = getPermitState(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()]; @@ -1533,26 +1664,12 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien if (cmdMessage.hasConsumerEpoch()) { consumerEpoch = cmdMessage.getConsumerEpoch(); } - final int earlyFailurePermits; - try { - earlyFailurePermits = MessagePermitAccounting.resolveForEarlyFailure(cmdMessage); - } catch (InvalidMessagePermitsException e) { - closeConnectionForInvalidMessagePermits(messageId, cnx, e); - return; - } log.debug().attr("messageId", messageId) .log("Received message"); - try { - if (!verifyChecksum(headersAndPayload, messageId)) { - // discard message with checksum error - discardCorruptedMessage(messageId, cnx, messagePermitState, - ValidationError.ChecksumMismatch, earlyFailurePermits); - return; - } - } catch (Throwable t) { - discardCorruptedMessage(messageId, cnx, messagePermitState, - ValidationError.ChecksumMismatch, earlyFailurePermits); + if (!verifyChecksum(headersAndPayload, messageId)) { + // discard message with checksum error + discardCorruptedMessage(messageId, sourceState, ValidationError.ChecksumMismatch, earlyMessagePermits); return; } @@ -1562,17 +1679,16 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien brokerEntryMetadata = Commands.parseBrokerEntryMetadataIfExist(headersAndPayload); msgMetadata = Commands.parseMessageMetadata(headersAndPayload); } catch (Throwable t) { - discardCorruptedMessage(messageId, cnx, messagePermitState, - ValidationError.ChecksumMismatch, earlyFailurePermits); + discardCorruptedMessage(messageId, sourceState, ValidationError.ChecksumMismatch, earlyMessagePermits); return; } final int numMessages = msgMetadata.getNumMessagesInBatch(); final int messagePermits; try { - messagePermits = MessagePermitAccounting.resolve(cmdMessage, numMessages); - } catch (InvalidMessagePermitsException e) { - closeConnectionForInvalidMessagePermits(messageId, cnx, e); + messagePermits = validateAndResolveMessagePermits(cmdMessage, numMessages, earlyMessagePermits); + } catch (PermitProtocolException e) { + closeMalformedPermitCommand(sourceState, e.getMessage()); return; } final int numChunks = msgMetadata.hasNumChunksFromMsg() ? msgMetadata.getNumChunksFromMsg() : 0; @@ -1583,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"); - returnPermitsToSourceConsumer(messagePermitState, messagePermits); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); return; } DecryptResult decryptResult = decryptPayloadIfNeeded(messageId, redeliveryCount, msgMetadata, headersAndPayload, - cnx); + sourceState); if (decryptResult.shouldDiscard()) { // Message was discarded or CryptoKeyReader isn't implemented @@ -1601,8 +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, messagePermitState, - true, messagePermits); + : uncompressPayloadIfNeeded(messageId, msgMetadata, decryptedPayload, sourceState, true, + messagePermits); decryptedPayload.release(); if (uncompressedPayload == null) { // Message was discarded on decompression error @@ -1613,7 +1729,7 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien // 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, - messagePermitState); + messagePermits, sourceState); return; } @@ -1622,8 +1738,8 @@ uncompressedPayload, msgId, schema, redeliveryCount, longArrayToList(ackSet), co if (isMessageUndecryptable || (numMessages == 1 && !msgMetadata.hasNumMessagesInBatch())) { if (isChunkedMessage) { - uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx, - messagePermitState); + uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, + sourceState); if (uncompressedPayload == null) { return; } @@ -1654,34 +1770,55 @@ uncompressedPayload, msgId, schema, redeliveryCount, longArrayToList(ackSet), co .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, cnx, - messagePermitState); + 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(message); + 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, messagePermits, - messagePermitState); + uncompressedPayload, messageId, sourceState, consumerEpoch, isMessageUndecryptable, + messagePermits); uncompressedPayload.release(); } @@ -1689,10 +1826,136 @@ uncompressedPayload, msgId, schema, redeliveryCount, longArrayToList(ackSet), co } + @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, ConsumerPermitState messagePermitState) { + MessageIdData messageId, ConsumerPermitState sourceState) { if (msgMetadata.getChunkId() != (msgMetadata.getNumChunksFromMsg() - 1)) { - returnPermitsToSourceConsumer(messagePermitState, 1); + increaseAvailablePermitsForReceivedMessage(sourceState, 1); } // Lazy task scheduling to expire incomplete chunk message if (expireTimeOfIncompleteChunkedMessageMillis > 0 && expireChunkMessageTaskScheduled.compareAndSet(false, @@ -1831,8 +2094,7 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m compressedPayload.release(); compressedPayload = chunkedMsgCtx.chunkedMsgBuffer; ByteBuf uncompressedPayload = - uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, cnx, messagePermitState, - false, 1); + uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, sourceState, false, 1); compressedPayload.release(); return uncompressedPayload; } @@ -1903,17 +2165,26 @@ private void interceptAndComplete(final Message message, final CompletableFut void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, MessageMetadata msgMetadata, int redeliveryCount, long[] ackSet, ByteBuf uncompressedPayload, MessageIdData messageId, ClientCnx cnx, long consumerEpoch, - boolean isEncrypted, int messagePermits) { + boolean isEncrypted) { + int messagePermits = ackSet == null || ackSet.length == 0 + ? msgMetadata.getNumMessagesInBatch() : BitSet.valueOf(ackSet).cardinality(); receiveIndividualMessagesFromBatch(brokerEntryMetadata, msgMetadata, redeliveryCount, ackSet, - uncompressedPayload, messageId, cnx, consumerEpoch, isEncrypted, messagePermits, - getPermitState(cnx)); + 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, - ConsumerPermitState messagePermitState) { + 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 @@ -1932,13 +2203,14 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } SingleMessageMetadata singleMessageMetadata = new SingleMessageMetadata(); - Budget permitBudget = new Budget(messagePermits); + 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, cnx, - messagePermitState); + 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. @@ -1955,43 +2227,40 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, } } if (acknowledgmentsGroupingTracker.isDuplicate(message.getMessageId())) { - message.release(); + if (possibleToDeadLetter == null) { + message.release(); + } continue; } - try { - permitBudget.claim(); - } catch (InvalidMessagePermitsException e) { - message.release(); - throw e; + if (possibleToDeadLetter != null) { + message.retain(); } try { - executeNotifyCallback(message); - } catch (RuntimeException e) { - permitBudget.restore(); - messagesPrefetchedGauge.decrement(); - bytesPrefetchedGauge.subtract(message.size()); - message.release(); + executeNotifyCallback(message, permitBudget); + } catch (PermitProtocolException e) { throw e; } } - } catch (InvalidMessagePermitsException e) { - closeConnectionForInvalidMessagePermits(messageId, cnx, e); - return; - } catch (RuntimeException 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, messagePermitState, - ValidationError.BatchDeSerializeError, permitBudget.drain()); + 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)); } @@ -2002,7 +2271,6 @@ void receiveIndividualMessagesFromBatch(BrokerEntryMetadata brokerEntryMetadata, .attr("remainingCapacity", incomingMessages.remainingCapacity()) .log("enqueued messages in batch. queue size -, available queue size"); - returnPermitsToSourceConsumer(messagePermitState, permitBudget.drain()); } private boolean isPriorEntryIndex(long idx) { @@ -2026,7 +2294,9 @@ private boolean isSameEntry(MessageIdImpl messageId) { */ @Override protected synchronized void messageProcessed(Message msg) { - ConsumerPermitState messagePermitState = ((MessageImpl) msg).getPermitState(); + ConsumerPermitState currentState = permitState.get(); + ConsumerPermitState messageState = msg instanceof MessageImpl messageImpl + ? messageImpl.getPermitState() : currentState; lastDequeuedMessageId = msg.getMessageId(); messagesPrefetchedGauge.decrement(); @@ -2035,11 +2305,13 @@ protected synchronized void messageProcessed(Message msg) { bytesPrefetchedGauge.subtract(msg.size()); bytesReceivedCounter.add(msg.size()); - if (messagePermitState != permitState) { + if (messageState != currentState) { // The processed message did belong to the old queue that was cleared after reconnection. } else { if (listener == null && !parentConsumerHasListener) { - updateAvailablePermits(messagePermitState, 1); + if (msg instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); + } } stats.updateNumMsgsReceived(msg); @@ -2072,91 +2344,59 @@ protected void trackMessage(MessageId messageId, int redeliveryCount) { } void increaseAvailablePermits(MessageImpl msg) { - increaseAvailablePermits(msg.getPermitState()); + increaseAvailablePermits(msg.getPermitState(), msg.takeFlowPermitCost()); } - void increaseAvailablePermits(ConsumerPermitState messagePermitState) { - returnPermitsToSourceConsumer(messagePermitState, 1); + private void increaseAvailablePermitsForReceivedMessage(ConsumerPermitState sourceState, int permits) { + increaseAvailablePermits(sourceState, permits); } - private void returnPermitsToSourceConsumer(ConsumerPermitState messagePermitState, int delta) { - if (delta <= 0) { - return; + @Override + protected void messageDiscarded(Message message) { + messagesPrefetchedGauge.decrement(); + bytesPrefetchedGauge.subtract(message.size()); + if (message instanceof MessageImpl messageImpl) { + increaseAvailablePermits(messageImpl); } - updateAvailablePermits(messagePermitState, delta); + message.release(); } - protected void updateAvailablePermits(ConsumerPermitState messagePermitState, int delta) { - if (messagePermitState == null || messagePermitState != permitState) { - return; - } - - int available; - while (true) { - int current = AVAILABLE_PERMITS_UPDATER.get(messagePermitState); - try { - available = Math.addExact(current, delta); - } catch (ArithmeticException e) { - scheduleCloseForInvalidPermitAccumulator(messagePermitState, - new InvalidMessagePermitsException("Client permit accumulator overflow")); - return; - } - if (AVAILABLE_PERMITS_UPDATER.compareAndSet(messagePermitState, current, available)) { - break; - } - } + void increaseAvailablePermits(ClientCnx currentCnx) { + increaseAvailablePermits(currentCnx, 1); + } - while (messagePermitState == permitState && available >= getCurrentReceiverQueueSize() / 2 && !paused - && messagePermitState.flowEnabled) { - if (AVAILABLE_PERMITS_UPDATER.compareAndSet(messagePermitState, available, 0)) { - ClientCnx flowCnx = cnx(); - if (flowCnx != null) { - scheduleFlowPermitsToBroker(messagePermitState, flowCnx, available); - } - return; - } - available = AVAILABLE_PERMITS_UPDATER.get(messagePermitState); - } + protected void increaseAvailablePermits(ClientCnx currentCnx, int delta) { + increaseAvailablePermits(getPermitState(currentCnx), delta); } - private void scheduleCloseForInvalidPermitAccumulator(ConsumerPermitState sourcePermitState, - InvalidMessagePermitsException error) { - ClientCnx sourceCnx = cnx(); - if (sourceCnx == null) { + private void increaseAvailablePermits(ConsumerPermitState state, int delta) { + if (state == null || delta < 0 || (delta == 0 && !state.active)) { return; } - sourceCnx.ctx().channel().eventLoop().execute(() -> { - if (sourcePermitState == permitState) { - closeConnectionForInvalidMessagePermits(null, sourceCnx, error); - } - }); - } - - private void scheduleFlowPermitsToBroker(ConsumerPermitState sourcePermitState, ClientCnx sourceCnx, - int permits) { - sourceCnx.ctx().channel().eventLoop().execute(() -> { - // The event loop is also the final ordering point for Subscribe and Flow on this connection. If this task - // wins, Flow is written before a later replacement Subscribe. If replacement wins, its new state makes - // this task stale and the old Flow is dropped, including when the physical ClientCnx is reused. - if (sourcePermitState != permitState) { - 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; + } } - log.debug().attr("available", permits) - .log("Sending permit-cmd to broker with available permits ="); - sendFlowPermitsToBroker(sourceCnx, permits); - }); - } - - void increaseAvailablePermits(ClientCnx currentCnx) { - increaseAvailablePermits(currentCnx, 1); - } - - protected void increaseAvailablePermits(ClientCnx currentCnx, int delta) { - updateAvailablePermits(getPermitState(currentCnx), delta); + } + scheduleFlowIfNeeded(state); } public void increaseAvailablePermits(int delta) { - increaseAvailablePermits(cnx(), delta); + increaseAvailablePermits(permitState.get(), delta); } @Override @@ -2173,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); } } @@ -2233,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()); @@ -2241,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()); @@ -2255,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) { @@ -2287,8 +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, getPermitState(currentCnx), - ValidationError.DecryptionError, batchSize); + discardMessage(messageId, sourceState, ValidationError.DecryptionError, batchSize); return DecryptResult.discard(); case FAIL: if (cryptoReaderNotExist) { @@ -2311,8 +2559,7 @@ private DecryptResult handleCryptoFailure(ByteBuf payload, MessageIdData message } private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetadata msgMetadata, ByteBuf payload, - ClientCnx currentCnx, ConsumerPermitState messagePermitState, - boolean checkMaxMessageSize, int messagePermits) { + ConsumerPermitState sourceState, boolean checkMaxMessageSize, int messagePermits) { CompressionType compressionType = msgMetadata.getCompression(); CompressionCodec codec = CompressionCodecProvider.getCompressionCodec(compressionType); int uncompressedSize = msgMetadata.getUncompressedSize(); @@ -2322,8 +2569,8 @@ private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetada log.error().attr("size", payloadSize) .attr("messageId", messageId) .log("Got corrupted payload message size at"); - discardCorruptedMessage(messageId, currentCnx, messagePermitState, - ValidationError.UncompressedSizeCorruption, messagePermits); + discardCorruptedMessage(messageId, sourceState, ValidationError.UncompressedSizeCorruption, + messagePermits); return null; } try { @@ -2335,8 +2582,7 @@ private ByteBuf uncompressPayloadIfNeeded(MessageIdData messageId, MessageMetada .exceptionMessage(e) .exception(e) .log("Failed to decompress message with at"); - discardCorruptedMessage(messageId, currentCnx, messagePermitState, - ValidationError.DecompressionError, messagePermits); + discardCorruptedMessage(messageId, sourceState, ValidationError.DecompressionError, messagePermits); return null; } } @@ -2358,48 +2604,38 @@ private boolean verifyChecksum(ByteBuf headersAndPayload, MessageIdData messageI return true; } - private void discardCorruptedMessage(MessageIdImpl messageId, ConsumerPermitState messagePermitState, - 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); - ClientCnx currentCnx = cnx(); - if (currentCnx != null) { - currentCnx.ctx().writeAndFlush(cmd, currentCnx.ctx().voidPromise()); - } else { - cmd.release(); - } - returnPermitsToSourceConsumer(messagePermitState, 1); + sourceState.cnx.ctx().writeAndFlush(cmd, sourceState.cnx.ctx().voidPromise()); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); stats.incrementNumReceiveFailed(); } - private void discardCorruptedMessage(MessageIdData messageId, ClientCnx currentCnx, - ConsumerPermitState messagePermitState, ValidationError validationError, int messagePermits) { + private void discardCorruptedMessage(MessageIdData messageId, ConsumerPermitState sourceState, + ValidationError validationError, int messagePermits) { log.error().attr("messageId", messageId) .log("Discarding corrupted message"); - discardMessage(messageId, currentCnx, messagePermitState, validationError, messagePermits); + discardMessage(messageId, sourceState, validationError, messagePermits); } - private void discardMessage(MessageIdData messageId, ClientCnx currentCnx, - ConsumerPermitState messagePermitState, ValidationError validationError, int messagePermits) { + 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()); - returnPermitsToSourceConsumer(messagePermitState, messagePermits); + sourceState.cnx.ctx().writeAndFlush(cmd, sourceState.cnx.ctx().voidPromise()); + increaseAvailablePermitsForReceivedMessage(sourceState, messagePermits); stats.incrementNumReceiveFailed(); } - private void closeConnectionForInvalidMessagePermits(MessageIdData messageId, ClientCnx messageCnx, - InvalidMessagePermitsException error) { - log.warn() - .attr("messageId", messageId) - .exceptionMessage(error) - .log("Closing source connection because message permit accounting is invalid"); - stats.incrementNumReceiveFailed(); - if (messageCnx != null) { - messageCnx.ctx().close(); - } + 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 @@ -2422,7 +2658,8 @@ int getPartitionIndex() { @Override public int getAvailablePermits() { - return AVAILABLE_PERMITS_UPDATER.get(permitState); + ConsumerPermitState state = permitState.get(); + return state == null ? 0 : state.availablePermits.get(); } @Override @@ -2455,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, @@ -2466,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(); @@ -2477,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" @@ -2504,9 +2746,8 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { return; } ClientCnx cnx = cnx(); - ConsumerPermitState redeliveryPermitState = getPermitState(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()) { @@ -2515,10 +2756,13 @@ public void redeliverUnacknowledgedMessages(Set messageIds) { } }); }); - if (messagesFromQueue > 0) { - updateAvailablePermits(redeliveryPermitState, 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; } @@ -2572,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, @@ -3230,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 -> { @@ -3242,8 +3494,7 @@ private int removeExpiredMessagesFromQueue(Set messageIds) { break; } decreaseIncomingMessageSize(message); - messagesFromQueue++; - message.release(); + messagesFromQueue.add(message); } return messagesFromQueue; } @@ -3293,12 +3544,9 @@ void resetBackoff() { this.connectionHandler.resetBackoff(); } - synchronized void connectionClosed(ClientCnx cnx, Optional initialConnectionDelayMs, Optional hostUrl) { - boolean wasCurrentConnection = this.connectionHandler.cnx() == cnx; + void connectionClosed(ClientCnx cnx, Optional initialConnectionDelayMs, Optional hostUrl) { + deactivatePermitState(getPermitState(cnx)); this.connectionHandler.connectionClosed(cnx, initialConnectionDelayMs, hostUrl); - if (wasCurrentConnection && this.connectionHandler.cnx() != cnx) { - permitState = new ConsumerPermitState(false); - } } public ClientCnx getClientCnx() { @@ -3307,21 +3555,16 @@ public ClientCnx getClientCnx() { void setClientCnx(ClientCnx clientCnx) { if (clientCnx != null) { - synchronized (this) { - this.connectionHandler.setClientCnx(clientCnx); - // Always create new permit state, including when the pooled ClientCnx instance is unchanged. - permitState = new ConsumerPermitState(false); + if (getPermitState(clientCnx) == null) { + replacePermitState(clientCnx); } + this.connectionHandler.setClientCnx(clientCnx); clientCnx.registerConsumer(consumerId, this); if (conf.isAckReceiptEnabled() && !Commands.peerSupportsAckReceipt(clientCnx.getRemoteEndpointProtocolVersion())) { log.warn().attr("nowversion", clientCnx.getRemoteEndpointProtocolVersion()) .log("Server don't support ack for receipt! " + "ProtoVersion >=17 support! nowVersion"); } - } else { - synchronized (this) { - permitState = new ConsumerPermitState(false); - } } ClientCnx previousClientCnx = clientCnxUsedForConsumerRegistration.getAndSet(clientCnx); if (previousClientCnx != null && previousClientCnx != clientCnx) { @@ -3329,30 +3572,8 @@ void setClientCnx(ClientCnx clientCnx) { } } - private ConsumerPermitState getPermitState(ClientCnx clientCnx) { - while (true) { - ConsumerPermitState currentPermitState = permitState; - ClientCnx currentCnx = cnx(); - if (currentPermitState != permitState) { - continue; - } - if (currentPermitState.flowEnabled) { - return currentCnx == clientCnx ? currentPermitState : null; - } - return clientCnx == null ? currentPermitState : null; - } - } - - @VisibleForTesting - ConsumerPermitState getPermitState() { - return permitState; - } - - boolean isCurrentPermitState(ConsumerPermitState messagePermitState) { - return messagePermitState != null && messagePermitState == permitState; - } - 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 c5c28bb6f0b96..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; @@ -66,7 +67,6 @@ public class MessageImpl implements TraceableMessage, Message { protected MessageId messageId; private final MessageMetadata msgMetadata; private ClientCnx cnx; - private ConsumerImpl.ConsumerPermitState permitState; private ByteBuf payload; private Schema schema; @@ -86,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. @@ -103,12 +114,14 @@ public static MessageImpl create(MessageMetadata msgMetadata, ByteBuffer msg.messageId = null; msg.topic = topic; msg.cnx = null; - msg.permitState = null; msg.payload = Unpooled.wrappedBuffer(payload); msg.properties = null; 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; } @@ -196,11 +209,13 @@ private static void init(MessageImpl msg, String topic, BatchMessageIdImp msg.messageId = batchMessageIdImpl; msg.topic = topic; msg.cnx = cnx; - msg.permitState = null; msg.redeliveryCount = redeliveryCount; 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 @@ -724,12 +739,38 @@ 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 setPermitState(ConsumerImpl.ConsumerPermitState 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() { @@ -740,7 +781,6 @@ public void recycle() { brokerEntryMetadata.clear(); } cnx = null; - permitState = null; messageId = null; topic = null; payload = null; @@ -752,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); @@ -761,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/MessagePermitAccounting.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java deleted file mode 100644 index 416399e0d92ae..0000000000000 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/MessagePermitAccounting.java +++ /dev/null @@ -1,111 +0,0 @@ -/* - * 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 org.apache.pulsar.common.api.proto.CommandMessage; - -/** - * Resolves the permit count represented by a message command and tracks that permit budget while a batch is decoded. - */ -final class MessagePermitAccounting { - - private MessagePermitAccounting() { - } - - static int resolveForEarlyFailure(CommandMessage command) { - if (command.hasMessagePermits()) { - return getExplicitMessagePermits(command); - } - return command.getAckSetsCount() > 0 ? getAckSetCardinality(command, -1) : 1; - } - - static int resolve(CommandMessage command, int batchSize) { - int expectedPermits = command.getAckSetsCount() > 0 - ? getAckSetCardinality(command, batchSize) : batchSize; - if (!command.hasMessagePermits()) { - return expectedPermits; - } - - int explicitPermits = getExplicitMessagePermits(command); - if (explicitPermits != expectedPermits) { - throw new InvalidMessagePermitsException("Explicit message permits " + explicitPermits - + " do not match the payload and ack set value " + expectedPermits); - } - return explicitPermits; - } - - private static int getExplicitMessagePermits(CommandMessage command) { - long messagePermits = Integer.toUnsignedLong(command.getMessagePermits()); - if (messagePermits == 0 || messagePermits > Integer.MAX_VALUE) { - throw new InvalidMessagePermitsException("Invalid explicit message permits " + messagePermits); - } - return (int) messagePermits; - } - - private static int getAckSetCardinality(CommandMessage command, int batchSize) { - int words = command.getAckSetsCount(); - int completeWords = batchSize < 0 ? words : Math.min(batchSize >>> 6, words); - long cardinality = 0; - for (int i = 0; i < completeWords; i++) { - cardinality += Long.bitCount(command.getAckSetAt(i)); - } - if (batchSize >= 0) { - int remainingBits = batchSize & 63; - if (remainingBits > 0 && completeWords < words) { - long mask = -1L >>> (Long.SIZE - remainingBits); - cardinality += Long.bitCount(command.getAckSetAt(completeWords) & mask); - } - } - if (cardinality > Integer.MAX_VALUE) { - throw new InvalidMessagePermitsException("Ack set message permits exceed the supported range"); - } - return (int) cardinality; - } - - static final class Budget { - private int remainingPermits; - - Budget(int messagePermits) { - remainingPermits = messagePermits; - } - - void claim() { - if (remainingPermits == 0) { - throw new InvalidMessagePermitsException("Batch contains more deliverable messages than permits"); - } - remainingPermits--; - } - - void restore() { - remainingPermits++; - } - - int drain() { - int permits = remainingPermits; - remainingPermits = 0; - return permits; - } - } - - static final class InvalidMessagePermitsException extends IllegalStateException { - InvalidMessagePermitsException(String message) { - super(message); - } - } -} 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 805b4b07bc9e2..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; - ConsumerImpl.ConsumerPermitState permitState = msgImpl.getPermitState(); boolean isValidEpoch = isValidConsumerEpoch(msgImpl); if (isValidEpoch && !skipDueToSeek) { messageReceived(consumer, msg); } else if (!isValidEpoch) { - consumer.increaseAvailablePermits(permitState); + 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 e0e5ec5b7141e..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,13 +115,17 @@ private Message fetchSingleMessageFromBroker() throws PulsarClientException { do { message = incomingMessages.take(); lastDequeuedMessageId = message.getMessageId(); - ConsumerPermitState messagePermitState = ((MessageImpl) message).getPermitState(); - // If the message belongs to an old broker-consumer incarnation, discard it and wait for the message - // from the latest Flow command. ClientCnx identity alone is insufficient because it can be reused. - if (isCurrentPermitState(messagePermitState)) { - waitingOnReceiveForZeroQueueSize = false; - break; + 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 (messageState == getPermitState()) { + waitingOnReceiveForZeroQueueSize = false; + break; + } } + message.release(); } while (true); stats.updateNumMsgsReceived(message); @@ -129,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(); + } } } @@ -156,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"); @@ -195,12 +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, int messagePermits, - ConsumerPermitState messagePermitState) { + 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 ); } @@ -215,18 +235,23 @@ protected void processPayloadByProcessor(BrokerEntryMetadata brokerEntryMetadata MessageMetadata messageMetadata, ByteBuf byteBuf, MessageIdImpl messageId, Schema schema, int redeliveryCount, List ackSet, long consumerEpoch, - ConsumerPermitState messagePermitState) { + 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, messagePermitState); + 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 a04dd638825ba..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 @@ -27,7 +27,6 @@ import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -38,16 +37,15 @@ import io.netty.channel.Channel; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelPromise; -import io.netty.channel.EventLoop; -import io.netty.util.ReferenceCountUtil; +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.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -72,7 +70,6 @@ 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.CompressionType; import org.apache.pulsar.common.api.proto.MessageIdData; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.api.proto.SingleMessageMetadata; @@ -146,204 +143,175 @@ public void testNotifyPendingReceivedCallback_EmptyQueueNotThrowsException() { } @Test - public void testInvalidExplicitMessagePermitsCloseSourceConnection() { - CommandMessage command = new CommandMessage() - .setConsumerId(consumer.consumerId) - .setMessagePermits(0); - command.setMessageId().setLedgerId(1).setEntryId(2); - ClientCnx messageCnx = mock(ClientCnx.class); - ChannelHandlerContext context = mock(ChannelHandlerContext.class); - when(messageCnx.ctx()).thenReturn(context); - ByteBuf emptyPayload = Unpooled.buffer(0); - int permitsBefore = consumer.getAvailablePermits(); + public void testGetMessagePermitsUsesCommandValueWhenPresent() { + CommandMessage command = new CommandMessage().setMessagePermits(3); + command.addAckSet(0b10101L); - try { - consumer.messageReceived(command, emptyPayload, messageCnx); - verify(context).close(); - Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore); - } finally { - emptyPayload.release(); - } + Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 3); } @Test - public void testTruncatedFrameAndMetadataFailureReturnExplicitMessagePermits() { - ClientCnx messageCnx = setCurrentConnection(); - ChannelHandlerContext context = messageCnx.ctx(); - int permitsBefore = consumer.getAvailablePermits(); - ByteBuf[] malformedFrames = { - Unpooled.wrappedBuffer(new byte[] {1}), - Unpooled.wrappedBuffer(new byte[] {0, 0, 0, 10}) - }; + public void testGetMessagePermitsFallsBackToAckSetForOldBroker() { + CommandMessage command = new CommandMessage(); + command.addAckSet(0b101101L); - try { - for (ByteBuf malformedFrame : malformedFrames) { - consumer.messageReceived(newCommandMessage(5), malformedFrame, messageCnx); - } + Assert.assertEquals(ConsumerImpl.getMessagePermits(command, 10), 4); + } - Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + 10); - releaseValidationCommands(context, 2); - } finally { - Arrays.stream(malformedFrames).forEach(ByteBuf::release); - } + @Test + public void testGetMessagePermitsFallsBackToBatchSizeForOldBroker() { + Assert.assertEquals(ConsumerImpl.getMessagePermits(new CommandMessage(), 10), 10); } @Test - public void testDecompressionFailureReturnsExplicitMessagePermits() { - ClientCnx messageCnx = setCurrentConnection(); - ChannelHandlerContext context = messageCnx.ctx(); + 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(1) - .setPublishTime(1) - .setNumMessagesInBatch(5) - .setCompression(CompressionType.ZLIB) - .setUncompressedSize(100); - ByteBuf invalidCompressedPayload = Unpooled.wrappedBuffer(new byte[] {1}); + .setSequenceId(0) + .setPublishTime(System.currentTimeMillis()) + .setNumMessagesInBatch(messagePermits); + ByteBuf payload = Unpooled.wrappedBuffer(new byte[] {1}); ByteBuf metadataAndPayload = Commands.serializeMetadataAndPayload( - Commands.ChecksumType.Crc32c, metadata, invalidCompressedPayload); - invalidCompressedPayload.release(); - int permitsBefore = consumer.getAvailablePermits(); + 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(newCommandMessage(5), metadataAndPayload, messageCnx); - - Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + 5); - releaseValidationCommands(context, 1); + consumer.messageReceived(command, metadataAndPayload, cnx); + Assert.assertEquals(consumer.getAvailablePermits(), permitsBefore + messagePermits); } finally { metadataAndPayload.release(); } } @Test - public void testPermitAccumulatorOverflowClosesSourceConnection() { - ClientCnx messageCnx = setCurrentConnection(); - ChannelHandlerContext context = messageCnx.ctx(); - consumer.paused = true; - ByteBuf firstMalformedMetadata = Unpooled.wrappedBuffer(new byte[] {1}); - ByteBuf secondMalformedMetadata = Unpooled.wrappedBuffer(new byte[] {1}); + 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); - try { - consumer.messageReceived(newCommandMessage(Integer.MAX_VALUE), firstMalformedMetadata, messageCnx); - Assert.assertEquals(consumer.getAvailablePermits(), Integer.MAX_VALUE); + consumer.deactivatePermitState(oldState); + consumer.setClientCnx(cnx); + ConsumerImpl.ConsumerPermitState replacementState = consumer.getPermitState(); + Assert.assertNotSame(replacementState, oldState); + Assert.assertSame(replacementState.cnx, oldState.cnx); - consumer.messageReceived(newCommandMessage(1), secondMalformedMetadata, messageCnx); + 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); - Assert.assertEquals(consumer.getAvailablePermits(), Integer.MAX_VALUE); - verify(context).close(); - releaseValidationCommands(context, 2); - } finally { - firstMalformedMetadata.release(); - secondMalformedMetadata.release(); - } + 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 testSameClientCnxReuseCreatesNewPermitIncarnation() { - consumer.setCurrentReceiverQueueSize(2); - ClientCnx messageCnx = setCurrentConnection(); - ChannelHandlerContext context = messageCnx.ctx(); - ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); - MessageImpl oldMessage = mock(MessageImpl.class); - when(oldMessage.getPermitState()).thenReturn(oldPermitState); + 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); - // Recreate the broker consumer while reusing the same pooled physical ClientCnx. - consumer.setClientCnx(messageCnx); - Assert.assertNotSame(consumer.getPermitState(), oldPermitState); - consumer.consumerIsReconnectedToBroker(messageCnx, 0); - consumer.increaseAvailablePermits(oldMessage); + consumer.messageReceived(command, Unpooled.EMPTY_BUFFER, cnx); + verify(channel).close(); + Assert.assertNull(consumer.getPermitState()); Assert.assertEquals(consumer.getAvailablePermits(), 0); - verify(context, never()).writeAndFlush(any(), any(ChannelPromise.class)); } - @Test(invocationTimeOut = 5000) - public void testPermitReturnRaceWithSameClientCnxReconnectDoesNotContaminateNewAccumulator() { - consumer.paused = true; - ClientCnx messageCnx = setCurrentConnection(); - ExecutorService raceExecutor = Executors.newFixedThreadPool(2); - try { - for (int i = 0; i < 100; i++) { - ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); - MessageImpl oldMessage = mock(MessageImpl.class); - when(oldMessage.getPermitState()).thenReturn(oldPermitState); - CountDownLatch start = new CountDownLatch(1); - - CompletableFuture returnPermit = CompletableFuture.runAsync(() -> { - await(start); - consumer.increaseAvailablePermits(oldMessage); - }, raceExecutor); - CompletableFuture reconnect = CompletableFuture.runAsync(() -> { - await(start); - consumer.setClientCnx(messageCnx); - consumer.consumerIsReconnectedToBroker(messageCnx, 0); - }, raceExecutor); - - start.countDown(); - CompletableFuture.allOf(returnPermit, reconnect).join(); - Assert.assertEquals(consumer.getAvailablePermits(), 0); - } - } finally { - raceExecutor.shutdownNow(); - } - } + @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); - @Test(invocationTimeOut = 5000) - public void testConcurrentPermitReturnsStayInCurrentIncarnationAccumulator() { - consumer.paused = true; - setCurrentConnection(); - ConsumerImpl.ConsumerPermitState currentPermitState = consumer.getPermitState(); - int threadCount = 8; - int returnsPerThread = 1000; - ExecutorService returnExecutor = Executors.newFixedThreadPool(threadCount); - CountDownLatch start = new CountDownLatch(1); - try { - List> returns = new ArrayList<>(threadCount); - for (int i = 0; i < threadCount; i++) { - returns.add(CompletableFuture.runAsync(() -> { - await(start); - for (int permit = 0; permit < returnsPerThread; permit++) { - consumer.increaseAvailablePermits(currentPermitState); - } - }, returnExecutor)); - } - start.countDown(); - CompletableFuture.allOf(returns.toArray(CompletableFuture[]::new)).join(); + consumer.increaseAvailablePermits(message); - Assert.assertEquals(consumer.getAvailablePermits(), threadCount * returnsPerThread); - } finally { - returnExecutor.shutdownNow(); - } + verify(channel).close(); + Assert.assertNull(consumer.getPermitState()); } @Test - public void testQueuedFlowFromOldIncarnationIsDroppedAfterSameClientCnxReuse() { - consumer.setCurrentReceiverQueueSize(2); - ClientCnx messageCnx = setCurrentConnection(); - ChannelHandlerContext context = messageCnx.ctx(); - EventLoop eventLoop = context.channel().eventLoop(); - List queuedTasks = new ArrayList<>(); - doAnswer(invocation -> { - queuedTasks.add(invocation.getArgument(0)); - return null; - }).when(eventLoop).execute(any(Runnable.class)); - - ConsumerImpl.ConsumerPermitState oldPermitState = consumer.getPermitState(); - consumer.increaseAvailablePermits(oldPermitState); - Assert.assertEquals(queuedTasks.size(), 1); - - // Recreate the broker consumer before the old Flow task reaches the shared physical connection. - consumer.setClientCnx(messageCnx); - queuedTasks.get(0).run(); - - Assert.assertEquals(consumer.getAvailablePermits(), 0); - verify(context, never()).writeAndFlush(any(), any(ChannelPromise.class)); + 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 = setCurrentConnection(); + ClientCnx messageCnx = setCurrentConnectionWithFlowEnabled(); ChannelHandlerContext context = messageCnx.ctx(); List flowPermits = new ArrayList<>(); doAnswer(invocation -> { @@ -385,7 +353,7 @@ public void testStaleEpochBatchOnCurrentIncarnationReturnsPermitAndClosesPrefetc @Test public void testRandomizedBatchDecodeWritesExactlyTheReturnedPermitsToFlow() throws Exception { consumer.setCurrentReceiverQueueSize(2); - ClientCnx messageCnx = setCurrentConnection(); + ClientCnx messageCnx = setCurrentConnectionWithFlowEnabled(); ChannelHandlerContext context = messageCnx.ctx(); List flowPermits = new ArrayList<>(); doAnswer(invocation -> { @@ -451,59 +419,36 @@ public void testRandomizedBatchDecodeWritesExactlyTheReturnedPermitsToFlow() thr } } - 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(); - } - } - - private static void await(CountDownLatch latch) { - try { - latch.await(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new CompletionException(e); - } - } - - private ClientCnx setCurrentConnection() { + private ClientCnx setCurrentConnectionWithFlowEnabled() { ClientCnx messageCnx = mock(ClientCnx.class); ChannelHandlerContext context = mock(ChannelHandlerContext.class); Channel channel = mock(Channel.class); - EventLoop eventLoop = mock(EventLoop.class); + EventExecutor eventExecutor = mock(EventExecutor.class); when(context.voidPromise()).thenReturn(mock(ChannelPromise.class)); when(context.channel()).thenReturn(channel); - when(channel.eventLoop()).thenReturn(eventLoop); + when(context.executor()).thenReturn(eventExecutor); doAnswer(invocation -> { invocation.getArgument(0).run(); return null; - }).when(eventLoop).execute(any(Runnable.class)); + }).when(eventExecutor).execute(any(Runnable.class)); when(messageCnx.ctx()).thenReturn(context); + when(messageCnx.channel()).thenReturn(channel); consumer.setClientCnx(messageCnx); - consumer.consumerIsReconnectedToBroker(messageCnx, 0); + consumer.getPermitState().flowEnabled = true; return messageCnx; } - private CommandMessage newCommandMessage(int messagePermits) { - CommandMessage command = new CommandMessage() - .setConsumerId(consumer.consumerId) - .setMessagePermits(messagePermits); - command.setMessageId().setLedgerId(1).setEntryId(2); - return command; - } - - private static void releaseValidationCommands(ChannelHandlerContext context, int expectedCommands) { - ArgumentCaptor commandCaptor = ArgumentCaptor.forClass(Object.class); - verify(context, times(expectedCommands)).writeAndFlush( - commandCaptor.capture(), any(ChannelPromise.class)); - commandCaptor.getAllValues().forEach(ReferenceCountUtil::release); + 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) @@ -657,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); @@ -766,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), @@ -784,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/MessagePermitAccountingTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java deleted file mode 100644 index cc8fb9dcf8752..0000000000000 --- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/MessagePermitAccountingTest.java +++ /dev/null @@ -1,155 +0,0 @@ -/* - * 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.expectThrows; -import java.util.BitSet; -import java.util.SplittableRandom; -import org.apache.pulsar.client.impl.MessagePermitAccounting.Budget; -import org.apache.pulsar.common.api.proto.CommandMessage; -import org.testng.annotations.Test; - -public class MessagePermitAccountingTest { - - private static final long RANDOM_SEED = 0x4915A17E5L; - - @Test - public void testResolveExplicitMessagePermits() { - CommandMessage command = new CommandMessage().setMessagePermits(3); - command.addAckSet(0b100101L); - - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(command), 3); - assertEquals(MessagePermitAccounting.resolve(command, 10), 3); - - CommandMessage partialBatchWithOnePermit = new CommandMessage().setMessagePermits(1); - partialBatchWithOnePermit.addAckSet(0b100L); - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(partialBatchWithOnePermit), 1); - assertEquals(MessagePermitAccounting.resolve(partialBatchWithOnePermit, 10), 1); - - CommandMessage singleMessage = new CommandMessage().setMessagePermits(1); - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(singleMessage), 1); - assertEquals(MessagePermitAccounting.resolve(singleMessage, 1), 1); - } - - @Test - public void testResolveLegacyMessagePermits() { - CommandMessage partialBatch = new CommandMessage(); - partialBatch.addAckSet((1L << 1) | (1L << 4) | (1L << 63)); - - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(new CommandMessage()), 1); - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(partialBatch), 3); - assertEquals(MessagePermitAccounting.resolve(partialBatch, 10), 2); - assertEquals(MessagePermitAccounting.resolve(new CommandMessage(), 10), 10); - } - - @Test - public void testRejectInvalidExplicitMessagePermits() { - CommandMessage zero = new CommandMessage().setMessagePermits(0); - CommandMessage unsignedOverflow = new CommandMessage().setMessagePermits(-1); - CommandMessage mismatch = new CommandMessage().setMessagePermits(2); - mismatch.addAckSet(0b111L); - CommandMessage onePermitMismatch = new CommandMessage().setMessagePermits(1); - - expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolveForEarlyFailure(zero)); - expectThrows(IllegalStateException.class, - () -> MessagePermitAccounting.resolveForEarlyFailure(unsignedOverflow)); - expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolve(mismatch, 10)); - expectThrows(IllegalStateException.class, () -> MessagePermitAccounting.resolve(onePermitMismatch, 10)); - } - - @Test - public void testBudgetTracksClaimsAndRestores() { - Budget budget = new Budget(2); - budget.claim(); - budget.restore(); - budget.claim(); - - assertEquals(budget.drain(), 1); - assertEquals(budget.drain(), 0); - expectThrows(IllegalStateException.class, budget::claim); - } - - @Test - public void testRandomizedAckSetResolutionMatchesBoundedBitSetCardinality() { - SplittableRandom random = new SplittableRandom(RANDOM_SEED); - for (int testCase = 0; testCase < 10_000; testCase++) { - int batchSize = random.nextInt(1, 513); - 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(); - } - // Explicit message permits must be positive. Also guarantees coverage when random data happens to be zero. - int requiredIndex = random.nextInt(batchSize); - ackSet[requiredIndex / Long.SIZE] |= 1L << (requiredIndex % Long.SIZE); - - BitSet boundedAckSet = BitSet.valueOf(ackSet); - boundedAckSet.clear(batchSize, Math.max(batchSize, boundedAckSet.length())); - int expectedPermits = boundedAckSet.cardinality(); - int allAckSetPermits = BitSet.valueOf(ackSet).cardinality(); - String description = "seed=" + RANDOM_SEED + ", case=" + testCase + ", batchSize=" + batchSize; - - CommandMessage explicit = commandWithAckSet(ackSet).setMessagePermits(expectedPermits); - assertEquals(MessagePermitAccounting.resolve(explicit, batchSize), expectedPermits, description); - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(explicit), expectedPermits, description); - - CommandMessage legacy = commandWithAckSet(ackSet); - assertEquals(MessagePermitAccounting.resolve(legacy, batchSize), expectedPermits, description); - assertEquals(MessagePermitAccounting.resolveForEarlyFailure(legacy), allAckSetPermits, description); - } - } - - @Test - public void testRandomizedBudgetConservesPermitsAcrossClaimsAndRestores() { - SplittableRandom random = new SplittableRandom(RANDOM_SEED); - for (int testCase = 0; testCase < 10_000; testCase++) { - int initialPermits = random.nextInt(1, 257); - int modeledRemaining = initialPermits; - int permanentlyClaimed = 0; - Budget budget = new Budget(initialPermits); - - int operations = random.nextInt(1, initialPermits * 3 + 1); - for (int operation = 0; operation < operations && modeledRemaining > 0; operation++) { - if (random.nextBoolean()) { - budget.claim(); - modeledRemaining--; - permanentlyClaimed++; - } else { - budget.claim(); - budget.restore(); - } - } - - int returnedPermits = budget.drain(); - String description = "seed=" + RANDOM_SEED + ", case=" + testCase; - assertEquals(returnedPermits, modeledRemaining, description); - assertEquals(permanentlyClaimed + returnedPermits, initialPermits, description); - assertEquals(budget.drain(), 0, description); - } - } - - private static CommandMessage commandWithAckSet(long[] ackSet) { - CommandMessage command = new CommandMessage(); - for (long word : ackSet) { - command.addAckSet(word); - } - return command; - } -} 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 e1baa2c027e9b..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]); @@ -560,17 +578,6 @@ public static BaseCommand newMessageCommand(long consumerId, long ledgerId, long return cmd; } - 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("messagePermits must be positive"); - } - BaseCommand command = newMessageCommand(consumerId, ledgerId, entryId, partition, redeliveryCount, ackSet, - consumerEpoch); - command.getMessage().setMessagePermits(messagePermits); - return command; - } - public static ByteBufPair newMessage(long consumerId, long ledgerId, long entryId, int partition, int redeliveryCount, ByteBuf metadataAndPayload, long[] ackSet) { return serializeCommandMessageWithSize( diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto index ae4f06747716f..77a58e3939926 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -582,7 +582,8 @@ message CommandMessage { repeated int64 ack_set = 4; optional uint64 consumer_epoch = 5; - // Number of consumer permits debited by the broker for this command. + // 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; } 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 97b05ec09856a..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 @@ -41,20 +41,22 @@ public class CommandsTest { @Test - public void testMessageCommandCarriesPermitDebit() { - BaseCommand batched = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 10); - assertTrue(batched.getMessage().hasMessagePermits()); - assertEquals(batched.getMessage().getMessagePermits(), 10); + public void testMessageCommandCarriesPermitCount() { + BaseCommand command = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 10); - BaseCommand single = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 1); - assertTrue(single.getMessage().hasMessagePermits()); - assertEquals(single.getMessage().getMessagePermits(), 1); + assertTrue(command.getMessage().hasMessagePermits()); + assertEquals(command.getMessage().getMessagePermits(), 10); - BaseCommand legacy = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1); - Assert.assertFalse(legacy.getMessage().hasMessagePermits()); + BaseCommand singleMessageCommand = Commands.newMessageCommand(1, 2, 3, 4, 0, null, -1, 1); + assertTrue(singleMessageCommand.getMessage().hasMessagePermits()); + assertEquals(singleMessageCommand.getMessage().getMessagePermits(), 1); - Assert.expectThrows(IllegalArgumentException.class, + 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