You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by lh...@apache.org on 2026/09/14 14:55:20 UTC
(pulsar) branch branch-4.0 updated (a1a814b34cd8 -> 70f53623631f)
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a change to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/pulsar.git
from a1a814b34cd8 [fix][txn] Prevent queued snapshot writes from stalling (#26565)
new 8c406a6b2864 [improve][client] Coalesce message listener drain scheduling (#26531)
new cb104f784736 [improve][ci] Configure SSH access action default and update Upterm (#26575)
new 70f53623631f [fix][client] Fix buffer ownership on the send failure paths (#26455)
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.github/actions/ssh-access/action.yml | 31 +-
.github/workflows/pulsar-ci-flaky.yaml | 1 +
.github/workflows/pulsar-ci.yaml | 1 +
.../impl/ListenerTaskSchedulingBenchmark.java | 102 +++
.../client/impl/RawBatchMessageContainerImpl.java | 97 ++-
.../compaction/StrategicTwoPhaseCompactor.java | 5 +-
.../impl/RawBatchMessageContainerImplTest.java | 98 +++
.../common/protocol/ProducerBatchSendTest.java | 116 ++--
.../client/impl/BatchMessageContainerImpl.java | 122 +++-
.../client/impl/BatchMessageKeyBasedContainer.java | 56 +-
.../apache/pulsar/client/impl/ConsumerBase.java | 51 +-
.../pulsar/client/impl/ListenerTaskScheduler.java | 73 ++
.../apache/pulsar/client/impl/ProducerImpl.java | 204 +++++-
.../client/impl/BatchMessageContainerImplTest.java | 549 +++++++++++++++
.../impl/ConsumerListenerSchedulingTest.java | 172 +++++
.../client/impl/ListenerTaskSchedulerTest.java | 245 +++++++
.../pulsar/client/impl/ProducerImplTest.java | 735 +++++++++++++++++++++
17 files changed, 2490 insertions(+), 168 deletions(-)
create mode 100644 microbench/src/main/java/org/apache/pulsar/client/impl/ListenerTaskSchedulingBenchmark.java
create mode 100644 pulsar-client/src/main/java/org/apache/pulsar/client/impl/ListenerTaskScheduler.java
create mode 100644 pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerListenerSchedulingTest.java
create mode 100644 pulsar-client/src/test/java/org/apache/pulsar/client/impl/ListenerTaskSchedulerTest.java
(pulsar) 02/03: [improve][ci] Configure SSH access action default and update Upterm (#26575)
Posted by lh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit cb104f7847366fbe4e8dcf2c7a57046436b8bdb4
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Mon Sep 14 10:11:32 2026 +0300
[improve][ci] Configure SSH access action default and update Upterm (#26575)
(cherry picked from commit b214a36d6831bc3db4cc99169887ee133771b58e)
(cherry picked from commit 678ab800ec8901b195de7df158cb84a6c5d4c128)
---
.github/actions/ssh-access/action.yml | 31 +++++++++++++++++++++++++++----
.github/workflows/pulsar-ci-flaky.yaml | 1 +
.github/workflows/pulsar-ci.yaml | 1 +
3 files changed, 29 insertions(+), 4 deletions(-)
diff --git a/.github/actions/ssh-access/action.yml b/.github/actions/ssh-access/action.yml
index 89f9b74e6040..11c10ceb3d7e 100644
--- a/.github/actions/ssh-access/action.yml
+++ b/.github/actions/ssh-access/action.yml
@@ -17,6 +17,13 @@
# under the License.
#
+# CI_ENABLE_SSH=true or false explicitly enables or disables both start and wait.
+# In particular, false disables SSH even for public repositories.
+# Without an override, public workflow repositories are enabled; private, internal,
+# and unknown visibility are disabled. The calling workflows map the repository
+# variable CI_ENABLE_SSH into the environment; job/step environment overrides win.
+# Callers still control when this action runs and which GitHub users' keys are allowed.
+
name: ssh access
description: Sets up SSH access to build VM with upterm
inputs:
@@ -47,27 +54,43 @@ inputs:
runs:
using: composite
steps:
- - run: |
+ - name: Configure SSH access
+ env:
+ SSH_ACCESS_ENABLED: ${{ env.CI_ENABLE_SSH || (github.event.repository.visibility == 'public' && 'true' || 'false') }}
+ run: |
+ case "$SSH_ACCESS_ENABLED" in
+ false)
+ echo "SSH access is disabled. Set CI_ENABLE_SSH=true to enable it."
+ exit 0
+ ;;
+ true) ;;
+ *)
+ echo "::error::CI_ENABLE_SSH must be true or false."
+ exit 1
+ ;;
+ esac
if [[ "${{ inputs.action }}" == "start" ]]; then
echo "::group::Installing upterm & tmux"
if [[ "$OSTYPE" == "linux-gnu"* ]]; then
# install upterm
- curl -sL https://github.com/owenthereal/upterm/releases/download/v0.20.0/upterm_linux_amd64.tar.gz | tar zxvf - -C /tmp upterm && sudo install /tmp/upterm /usr/local/bin/ && rm -rf /tmp/upterm
+ curl -fsSL https://github.com/owenthereal/upterm/releases/download/v0.27.0/upterm_linux_amd64.tar.gz | tar zxvf - -C /tmp upterm && sudo install /tmp/upterm /usr/local/bin/ && rm -rf /tmp/upterm
# install tmux if it's not present
if ! command -v tmux &>/dev/null; then
sudo apt-get -y install tmux
fi
elif [[ "$OSTYPE" == "darwin"* ]]; then
- brew install --cask owenthereal/upterm/upterm
+ # Fully qualifying the cask grants trust to this item; suppress installation prompts in CI.
+ HOMEBREW_NO_ASK=1 brew install --cask owenthereal/upterm/upterm
# install tmux if it's not present
if ! command -v tmux &>/dev/null; then
- brew install tmux
+ HOMEBREW_NO_ASK=1 brew install tmux
fi
else
echo "Unsupported $OSTYPE"
exit 0
fi
+ upterm version
echo '::endgroup::'
echo "::group::Configuring ssh and ssh keys"
# generate ssh key
diff --git a/.github/workflows/pulsar-ci-flaky.yaml b/.github/workflows/pulsar-ci-flaky.yaml
index 9d53232df852..696a0ef50278 100644
--- a/.github/workflows/pulsar-ci-flaky.yaml
+++ b/.github/workflows/pulsar-ci-flaky.yaml
@@ -76,6 +76,7 @@ concurrency:
cancel-in-progress: true
env:
+ CI_ENABLE_SSH: ${{ vars.CI_ENABLE_SSH }}
MAVEN_OPTS: -Xss1500k -Xmx2048m -XX:+UnlockDiagnosticVMOptions -XX:+IgnoreUnrecognizedVMOptions -XX:GCLockerRetryAllocationCount=100 -Daether.connector.http.reuseConnections=false -Daether.connector.requestTimeout=60000 -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false -Dmaven.wagon.http.retryHandler.class=standard -Dmaven.wagon.http.retryHandler.count=3 -Dmaven.wagon.http.retryHandler.requestSentEnabled=true -Dmaven.wagon.http.serviceUnavailableRetryStrategy.class=standard -Dmaven. [...]
# defines the retention period for the intermediate build artifacts needed for rerunning a failed build job
# it's possible to rerun individual failed jobs when the build artifacts are available
diff --git a/.github/workflows/pulsar-ci.yaml b/.github/workflows/pulsar-ci.yaml
index 8ee4de81e396..6eee771dcd94 100644
--- a/.github/workflows/pulsar-ci.yaml
+++ b/.github/workflows/pulsar-ci.yaml
@@ -76,6 +76,7 @@ concurrency:
cancel-in-progress: true
env:
+ CI_ENABLE_SSH: ${{ vars.CI_ENABLE_SSH }}
MAVEN_OPTS: -Xss1500k -Xmx2048m -XX:+UnlockDiagnosticVMOptions -XX:+IgnoreUnrecognizedVMOptions -XX:GCLockerRetryAllocationCount=100 -Daether.connector.http.reuseConnections=false -Daether.connector.requestTimeout=60000 -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false -Dmaven.wagon.http.retryHandler.class=standard -Dmaven.wagon.http.retryHandler.count=3 -Dmaven.wagon.http.retryHandler.requestSentEnabled=true -Dmaven.wagon.http.serviceUnavailableRetryStrategy.class=standard -Dmaven. [...]
# defines the retention period for the intermediate build artifacts needed for rerunning a failed build job
# it's possible to rerun individual failed jobs when the build artifacts are available
(pulsar) 03/03: [fix][client] Fix buffer ownership on the send failure paths (#26455)
Posted by lh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 70f53623631f806540ec2e169d4ed0ac7cba131a
Author: Zixuan Liu <no...@gmail.com>
AuthorDate: Mon Sep 14 16:58:42 2026 +0800
[fix][client] Fix buffer ownership on the send failure paths (#26455)
(cherry picked from commit b886d7790984a1ff1dfffbec901bba7128c51803)
(cherry picked from commit 1e6590542930cbeb7750195643067ad3f397f78b)
---
.../client/impl/RawBatchMessageContainerImpl.java | 97 ++-
.../compaction/StrategicTwoPhaseCompactor.java | 5 +-
.../impl/RawBatchMessageContainerImplTest.java | 98 +++
.../common/protocol/ProducerBatchSendTest.java | 116 ++--
.../client/impl/BatchMessageContainerImpl.java | 122 +++-
.../client/impl/BatchMessageKeyBasedContainer.java | 56 +-
.../apache/pulsar/client/impl/ProducerImpl.java | 204 +++++-
.../client/impl/BatchMessageContainerImplTest.java | 549 +++++++++++++++
.../pulsar/client/impl/ProducerImplTest.java | 735 +++++++++++++++++++++
9 files changed, 1837 insertions(+), 145 deletions(-)
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImpl.java b/pulsar-broker/src/main/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImpl.java
index 374f1e30c0a8..06ba2d3b1a49 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImpl.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImpl.java
@@ -18,7 +18,10 @@
*/
package org.apache.pulsar.client.impl;
+import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.util.ReferenceCountUtil;
import java.nio.ByteBuffer;
import java.util.Set;
import org.apache.pulsar.client.api.CryptoKeyReader;
@@ -56,26 +59,45 @@ public class RawBatchMessageContainerImpl extends BatchMessageContainerImpl {
this.compressor = new CompressionCodecNone();
}
+ /** This constructor is for testing only, to track the buffers the container allocates. */
+ @VisibleForTesting
+ RawBatchMessageContainerImpl(ByteBufAllocator allocator) {
+ super(allocator);
+ this.compressionType = CompressionType.NONE;
+ this.compressor = new CompressionCodecNone();
+ }
+
private ByteBuf encrypt(ByteBuf compressedPayload) {
if (msgCrypto == null) {
return compressedPayload;
}
- int maxSize = msgCrypto.getMaxOutputSize(compressedPayload.readableBytes());
- ByteBuf encryptedPayload = allocator.buffer(maxSize);
- ByteBuffer targetBuffer = encryptedPayload.nioBuffer(0, maxSize);
-
+ ByteBuf encryptedPayload = null;
try {
+ int maxSize = msgCrypto.getMaxOutputSize(compressedPayload.readableBytes());
+ encryptedPayload = allocator.buffer(maxSize);
+ ByteBuffer targetBuffer = encryptedPayload.nioBuffer(0, maxSize);
msgCrypto.encrypt(encryptionKeys, cryptoKeyReader, () -> messageMetadata,
compressedPayload.nioBuffer(), targetBuffer);
- } catch (PulsarClientException e) {
- encryptedPayload.release();
+ encryptedPayload.writerIndex(targetBuffer.remaining());
compressedPayload.release();
+ return encryptedPayload;
+ } catch (PulsarClientException e) {
+ // Release the compressed payload and any partially built encrypted buffer before failing the batch.
+ ReferenceCountUtil.safeRelease(encryptedPayload);
+ ReferenceCountUtil.safeRelease(compressedPayload);
discard(e);
throw new RuntimeException("Failed to encrypt payload", e);
+ } catch (Throwable t) {
+ // Never orphan the compressed payload or a partially built encrypted buffer when encryption fails,
+ // whatever the failure is (e.g. an OOM while allocating the encrypted buffer or an unexpected
+ // runtime exception from the crypto provider). Unlike the PulsarClientException branch, the batch is
+ // deliberately not discarded here: this failure escapes toByteBuf() before its serialization try
+ // (whose finally clears the container), so the caller owns the recovery — StrategicTwoPhaseCompactor
+ // discards the container on it.
+ ReferenceCountUtil.safeRelease(encryptedPayload);
+ ReferenceCountUtil.safeRelease(compressedPayload);
+ throw t;
}
- encryptedPayload.writerIndex(targetBuffer.remaining());
- compressedPayload.release();
- return encryptedPayload;
}
@Override
@@ -91,6 +113,11 @@ public class RawBatchMessageContainerImpl extends BatchMessageContainerImpl {
this.cryptoKeyReader = cryptoKeyReader;
}
+ @VisibleForTesting
+ void setMsgCryptoForTesting(MessageCrypto<MessageMetadata, MessageMetadata> msgCrypto) {
+ this.msgCrypto = msgCrypto;
+ }
+
@Override
public boolean add(MessageImpl<?> msg, SendCallback callback) {
this.lastAddedMessageId = (MessageIdAdv) msg.getMessageId();
@@ -168,29 +195,35 @@ public class RawBatchMessageContainerImpl extends BatchMessageContainerImpl {
}
ByteBuf encryptedPayload = encrypt(getCompressedBatchMetadataAndPayload());
- updateAndReserveBatchAllocatedSize(encryptedPayload.capacity());
- ByteBuf metadataAndPayload = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c,
- messageMetadata, encryptedPayload);
-
- MessageIdData idData = new MessageIdData();
- idData.setLedgerId(lastMessageId.getLedgerId());
- idData.setEntryId(lastMessageId.getEntryId());
- idData.setPartition(lastMessageId.getPartitionIndex());
-
- // Format: [IdSize][Id][metadataAndPayloadSize][metadataAndPayload]
- // Following RawMessage.serialize() format as the compacted messages will be parsed as RawMessage in broker
- int idSize = idData.getSerializedSize();
- int headerSize = 4 /* IdSize */ + idSize + 4 /* metadataAndPayloadSize */;
- int totalSize = headerSize + metadataAndPayload.readableBytes();
- ByteBuf buf = PulsarByteBufAllocator.DEFAULT.buffer(totalSize);
- buf.writeInt(idSize);
- idData.writeTo(buf);
- buf.writeInt(metadataAndPayload.readableBytes());
- buf.writeBytes(metadataAndPayload);
- metadataAndPayload.release();
- encryptedPayload.release();
- clear();
- return buf;
+ ByteBuf metadataAndPayload = null;
+ try {
+ updateAndReserveBatchAllocatedSize(encryptedPayload.capacity());
+ metadataAndPayload = Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c,
+ messageMetadata, encryptedPayload);
+
+ MessageIdData idData = new MessageIdData();
+ idData.setLedgerId(lastMessageId.getLedgerId());
+ idData.setEntryId(lastMessageId.getEntryId());
+ idData.setPartition(lastMessageId.getPartitionIndex());
+
+ // Format: [IdSize][Id][metadataAndPayloadSize][metadataAndPayload]
+ // Following RawMessage.serialize() format as the compacted messages will be parsed as RawMessage in broker
+ int idSize = idData.getSerializedSize();
+ int headerSize = 4 /* IdSize */ + idSize + 4 /* metadataAndPayloadSize */;
+ int totalSize = headerSize + metadataAndPayload.readableBytes();
+ ByteBuf buf = PulsarByteBufAllocator.DEFAULT.buffer(totalSize);
+ buf.writeInt(idSize);
+ idData.writeTo(buf);
+ buf.writeInt(metadataAndPayload.readableBytes());
+ buf.writeBytes(metadataAndPayload);
+ return buf;
+ } finally {
+ // Release everything allocated for this serialization on both success and failure, so a failure after
+ // the batch buffer was built (e.g. an OOM) cannot orphan it.
+ ReferenceCountUtil.safeRelease(metadataAndPayload);
+ ReferenceCountUtil.safeRelease(encryptedPayload);
+ clear();
+ }
}
@Override
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/compaction/StrategicTwoPhaseCompactor.java b/pulsar-broker/src/main/java/org/apache/pulsar/compaction/StrategicTwoPhaseCompactor.java
index 3c6749757421..c0cfd67ce4db 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/compaction/StrategicTwoPhaseCompactor.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/compaction/StrategicTwoPhaseCompactor.java
@@ -38,6 +38,7 @@ import org.apache.pulsar.client.api.CryptoKeyReader;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Reader;
import org.apache.pulsar.client.impl.BatchMessageIdImpl;
import org.apache.pulsar.client.impl.CompactionReaderImpl;
@@ -457,7 +458,9 @@ public class StrategicTwoPhaseCompactor extends PublishingOrderCompactor {
} catch (Throwable t) {
log.error("Failed to add entry", t);
- batchMessageContainer.discard((Exception) t);
+ // discard(Exception) cannot take an Error: wrap non-Exception Throwables so the batch container is
+ // always cleared here and stays reusable after a failed flush.
+ batchMessageContainer.discard(t instanceof Exception ? (Exception) t : new PulsarClientException(t));
bkf.completeExceptionally(t);
return bkf;
}
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImplTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImplTest.java
index e8a91a4a194f..96a69ca61616 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImplTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/RawBatchMessageContainerImplTest.java
@@ -21,19 +21,30 @@ package org.apache.pulsar.client.impl;
import static org.apache.pulsar.common.api.proto.CompressionType.NONE;
import static org.apache.pulsar.common.api.proto.CompressionType.ZSTD;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anySet;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
import static org.testng.AssertJUnit.assertFalse;
import static org.testng.AssertJUnit.assertTrue;
import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.Unpooled;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.apache.pulsar.client.api.CryptoKeyReader;
import org.apache.pulsar.client.api.EncryptionKeyInfo;
import org.apache.pulsar.client.api.MessageCrypto;
+import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.impl.crypto.MessageCryptoBc;
import org.apache.pulsar.common.api.EncryptionContext;
@@ -306,4 +317,91 @@ public class RawBatchMessageContainerImplTest {
Assert.assertEquals(container.getNumMessagesInBatch(), 0);
Assert.assertEquals(container.batchedMessageMetadataAndPayload, null);
}
+
+ /**
+ * A crypto provider that fails with an unexpected (non-{@link PulsarClientException}) error after the batch
+ * payload was built must not orphan the compressed payload or the partially built encrypted buffer.
+ */
+ @Test
+ public void testToByteBufReleasesPayloadWhenEncryptionFailsUnexpectedly() throws Exception {
+ setEncryptionAndCompression(true, false);
+ // Track every buffer the container allocates, so the partially built encrypted output buffer is
+ // asserted as well, not just the batch payload it hands over.
+ List<ByteBuf> allocated = new ArrayList<>();
+ ByteBufAllocator trackingAllocator = mock(ByteBufAllocator.class);
+ doAnswer(invocation -> {
+ ByteBuf buffer = Unpooled.buffer(invocation.getArgument(0));
+ allocated.add(buffer);
+ return buffer;
+ }).when(trackingAllocator).buffer(anyInt());
+ RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(trackingAllocator);
+ container.setCryptoKeyReader(cryptoKeyReader);
+ container.add(createMessage("my-topic", "hi-1", 0), null);
+
+ // Replace the real crypto with one whose encrypt() throws an unexpected RuntimeException, so the batch
+ // payload is built (getCompressedBatchMetadataAndPayload) and then encryption fails outside the
+ // PulsarClientException contract.
+ MessageCrypto<MessageMetadata, MessageMetadata> crypto = mock(MessageCrypto.class);
+ when(crypto.getMaxOutputSize(anyInt())).thenReturn(128);
+ doThrow(new RuntimeException("mocked crypto failure"))
+ .when(crypto).encrypt(anySet(), any(), any(), any(), any());
+ container.setMsgCryptoForTesting(crypto);
+
+ Throwable e = null;
+ try {
+ container.toByteBuf();
+ } catch (Throwable ex) {
+ e = ex;
+ }
+ Assert.assertEquals(e.getClass(), RuntimeException.class);
+ Assert.assertTrue(e.getMessage().contains("mocked crypto failure"));
+ // The compressed batch payload must have been released instead of leaked; the container keeps its
+ // (now released) buffer reference until the caller recovers, mirroring the producer path.
+ Assert.assertEquals(container.batchedMessageMetadataAndPayload.refCnt(), 0);
+ // The partially built encrypted output buffer must have been released as well, not only the source.
+ for (ByteBuf buffer : allocated) {
+ Assert.assertEquals(buffer.refCnt(), 0);
+ }
+
+ container.discard(null);
+ }
+
+ /**
+ * A crypto provider failing with a {@link PulsarClientException} must release the compressed batch payload
+ * and the partially built encrypted buffer and discard the batch, so the compactor can reuse the container.
+ */
+ @Test
+ public void testToByteBufReleasesPayloadAndDiscardsWhenEncryptionFailsWithClientException() throws Exception {
+ setEncryptionAndCompression(true, false);
+ List<ByteBuf> allocated = new ArrayList<>();
+ ByteBufAllocator trackingAllocator = mock(ByteBufAllocator.class);
+ doAnswer(invocation -> {
+ ByteBuf buffer = Unpooled.buffer(invocation.getArgument(0));
+ allocated.add(buffer);
+ return buffer;
+ }).when(trackingAllocator).buffer(anyInt());
+ RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(trackingAllocator);
+ container.setCryptoKeyReader(cryptoKeyReader);
+ container.add(createMessage("my-topic", "hi-1", 0), null);
+
+ MessageCrypto<MessageMetadata, MessageMetadata> crypto = mock(MessageCrypto.class);
+ when(crypto.getMaxOutputSize(anyInt())).thenReturn(128);
+ doThrow(new PulsarClientException("mocked crypto failure"))
+ .when(crypto).encrypt(anySet(), any(), any(), any(), any());
+ container.setMsgCryptoForTesting(crypto);
+
+ try {
+ container.toByteBuf();
+ Assert.fail("expected the encryption failure to propagate");
+ } catch (RuntimeException e) {
+ Assert.assertTrue(e.getMessage().contains("Failed to encrypt payload"));
+ Assert.assertTrue(e.getCause() instanceof PulsarClientException);
+ }
+ // Unlike the unexpected-Throwable branch, the PulsarClientException branch discards the batch so the
+ // container is empty and reusable for the next flush.
+ Assert.assertEquals(container.getNumMessagesInBatch(), 0);
+ for (ByteBuf buffer : allocated) {
+ Assert.assertEquals(buffer.refCnt(), 0);
+ }
+ }
}
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/common/protocol/ProducerBatchSendTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/common/protocol/ProducerBatchSendTest.java
index 8f24abad0058..972bea448e0a 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/common/protocol/ProducerBatchSendTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/common/protocol/ProducerBatchSendTest.java
@@ -18,6 +18,7 @@
*/
package org.apache.pulsar.common.protocol;
+import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.spy;
import static org.testng.Assert.assertEquals;
@@ -30,9 +31,11 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.broker.BrokerTestUtil;
+import org.apache.pulsar.client.api.CompressionType;
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.ProducerBuilder;
import org.apache.pulsar.client.api.ProducerConsumerBase;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.impl.ConsumerImpl;
@@ -63,30 +66,42 @@ public class ProducerBatchSendTest extends ProducerConsumerBase {
@DataProvider
public Object[][] flushSend() {
return new Object[][] {
- {Collections.emptyList()},
- {Arrays.asList(1)},
- {Arrays.asList(2)},
- {Arrays.asList(3)},
- {Arrays.asList(1, 2)},
- {Arrays.asList(2, 3)},
- {Arrays.asList(1, 2, 3)},
+ {Collections.emptyList(), CompressionType.NONE},
+ {Arrays.asList(1), CompressionType.NONE},
+ {Arrays.asList(2), CompressionType.NONE},
+ {Arrays.asList(3), CompressionType.NONE},
+ {Arrays.asList(1, 2), CompressionType.NONE},
+ {Arrays.asList(2, 3), CompressionType.NONE},
+ {Arrays.asList(1, 2, 3), CompressionType.NONE},
+ {Collections.emptyList(), CompressionType.ZLIB},
+ {Arrays.asList(1), CompressionType.ZLIB},
+ {Arrays.asList(2), CompressionType.ZLIB},
+ {Arrays.asList(3), CompressionType.ZLIB},
+ {Arrays.asList(1, 2), CompressionType.ZLIB},
+ {Arrays.asList(2, 3), CompressionType.ZLIB},
+ {Arrays.asList(1, 2, 3), CompressionType.ZLIB},
};
}
+ /**
+ * The batch container's {@code createOpSendMsg()} may fail after the batch payload was already built, e.g.
+ * when the command buffer allocation fails in {@link Commands#serializeCommandSendWithSize}. With
+ * compression enabled, the batch buffer has already been released at that point, so the recovery must not
+ * reuse it when the batch is retried.
+ */
@Test(timeOut = 30_000, dataProvider = "flushSend")
- public void testNoEnoughMemSend(List<Integer> flushSend) throws Exception {
+ public void testNoEnoughMemSend(List<Integer> flushSend, CompressionType compressionType) throws Exception {
final String topic = BrokerTestUtil.newUniqueName("persistent://public/default/tp");
final String subscription = "s1";
admin.topics().createNonPartitionedTopic(topic);
admin.topics().createSubscription(topic, subscription, MessageId.earliest);
- Producer<String> producer = pulsarClient.newProducer(Schema.STRING).topic(topic).enableBatching(true)
- .batchingMaxMessages(Integer.MAX_VALUE).batchingMaxPublishDelay(1, TimeUnit.HOURS).create();
+ ProducerBuilder<String> builder = pulsarClient.newProducer(Schema.STRING).topic(topic).enableBatching(true)
+ .batchingMaxMessages(Integer.MAX_VALUE).batchingMaxPublishDelay(1, TimeUnit.HOURS);
+ if (compressionType != CompressionType.NONE) {
+ builder.compressionType(compressionType);
+ }
+ Producer<String> producer = builder.create();
- /**
- * The method {@link org.apache.pulsar.client.impl.BatchMessageContainerImpl#createOpSendMsg} may fail due to
- * many errors, such like allocate more memory failed when calling
- * {@link Commands#serializeCommandSendWithSize}. We mock an error here.
- */
AtomicBoolean failure = new AtomicBoolean(true);
BaseCommand threadLocalBaseCommand = Commands.LOCAL_BASE_COMMAND.get();
BaseCommand spyBaseCommand = spy(threadLocalBaseCommand);
@@ -99,42 +114,47 @@ public class ProducerBatchSendTest extends ProducerConsumerBase {
}).when(spyBaseCommand).setSend();
Commands.LOCAL_BASE_COMMAND.set(spyBaseCommand);
- // Failed sending 3 times.
- producer.sendAsync("1");
- if (flushSend.contains(1)) {
- producer.flushAsync();
- }
- producer.sendAsync("2");
- if (flushSend.contains(2)) {
- producer.flushAsync();
+ // 6 KB payloads give the ZLIB cases a batch that is really compressed (and stay above the
+ // compressMinMsgBodySize threshold on branches that have one) even when a single message is flushed.
+ StringBuilder sb = new StringBuilder();
+ for (int i = 0; i < 750; i++) {
+ sb.append("abcdefgh");
}
- producer.sendAsync("3");
- if (flushSend.contains(3)) {
- producer.flushAsync();
- }
- // Publishing is finished eventually.
- failure.set(false);
- producer.flush();
- Awaitility.await().untilAsserted(() -> {
- assertTrue(admin.topics().getStats(topic).getSubscriptions().get(subscription).getMsgBacklog() > 0);
- });
- // Verify: all messages can be consumed.
- ConsumerImpl<String> consumer = (ConsumerImpl<String>) pulsarClient.newConsumer(Schema.STRING).topic(topic)
- .subscriptionName(subscription).subscribe();
- Message<String> msg1 = consumer.receive(2, TimeUnit.SECONDS);
- assertNotNull(msg1);
- assertEquals(msg1.getValue(), "1");
- Message<String> msg2 = consumer.receive(2, TimeUnit.SECONDS);
- assertNotNull(msg2);
- assertEquals(msg2.getValue(), "2");
- Message<String> msg3 = consumer.receive(2, TimeUnit.SECONDS);
- assertNotNull(msg3);
- assertEquals(msg3.getValue(), "3");
+ try {
+ // Failed sending 3 times.
+ producer.sendAsync(sb + "-1");
+ if (flushSend.contains(1)) {
+ producer.flushAsync();
+ }
+ producer.sendAsync(sb + "-2");
+ if (flushSend.contains(2)) {
+ producer.flushAsync();
+ }
+ producer.sendAsync(sb + "-3");
+ if (flushSend.contains(3)) {
+ producer.flushAsync();
+ }
+ // Publishing is finished eventually.
+ failure.set(false);
+ assertThat(producer.flushAsync()).succeedsWithin(10, TimeUnit.SECONDS);
+ Awaitility.await().untilAsserted(() -> {
+ assertTrue(admin.topics().getStats(topic).getSubscriptions().get(subscription).getMsgBacklog() > 0);
+ });
- // cleanup.
- consumer.close();
- producer.close();
+ // Verify: all messages can be consumed.
+ ConsumerImpl<String> consumer = (ConsumerImpl<String>) pulsarClient.newConsumer(Schema.STRING)
+ .topic(topic).subscriptionName(subscription).subscribe();
+ for (int i = 1; i <= 3; i++) {
+ Message<String> msg = consumer.receive(2, TimeUnit.SECONDS);
+ assertNotNull(msg, "message " + i + " lost");
+ assertEquals(msg.getValue(), sb + "-" + i);
+ }
+ consumer.close();
+ } finally {
+ Commands.LOCAL_BASE_COMMAND.set(threadLocalBaseCommand);
+ producer.close();
+ }
admin.topics().delete(topic, false);
}
}
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageContainerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageContainerImpl.java
index 323727e929df..d3182a7a333c 100644
--- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageContainerImpl.java
+++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageContainerImpl.java
@@ -18,6 +18,7 @@
*/
package org.apache.pulsar.client.impl;
+import static com.google.common.base.Preconditions.checkState;
import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
@@ -58,6 +59,9 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
@Setter
protected long highestSequenceId = -1L;
protected ByteBuf batchedMessageMetadataAndPayload;
+ // Whether the container still owns batchedMessageMetadataAndPayload. A build may release or hand the buffer
+ // over (compression / encryption); failure recovery then reallocates instead of reusing that buffer.
+ private boolean batchPayloadOwned;
protected List<MessageImpl<?>> messages = new ArrayList<>(maxMessagesNum);
protected SendCallback previousCallback = null;
// keep track of callbacks for individual messages being published in a batch
@@ -102,6 +106,7 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
this.firstCallback = callback;
batchedMessageMetadataAndPayload = allocator.buffer(
Math.min(maxBatchSize, getMaxMessageSize()));
+ batchPayloadOwned = true;
updateAndReserveBatchAllocatedSize(batchedMessageMetadataAndPayload.capacity());
if (msg.getMessageBuilder().hasTxnidMostBits() && currentTxnidMostBits == -1) {
currentTxnidMostBits = msg.getMessageBuilder().getTxnidMostBits();
@@ -142,6 +147,11 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
}
protected ByteBuf getCompressedBatchMetadataAndPayload() {
+ // A build may only start while the container owns its batch buffer, allocated by add() or reallocated by
+ // resetPayloadAfterFailedPublishing(). After compression/encryption released it, a build must not be
+ // re-entered without going through resetPayloadAfterFailedPublishing() first.
+ checkState(batchPayloadOwned,
+ "Batch payload buffer is missing or not owned when building a non-empty batch");
int batchWriteIndex = batchedMessageMetadataAndPayload.writerIndex();
int batchReadIndex = batchedMessageMetadataAndPayload.readerIndex();
@@ -151,7 +161,7 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
try {
if (n == 1) {
batchedMessageMetadataAndPayload.writeBytes(msg.getDataBuffer());
- } else {
+ } else {
batchedMessageMetadataAndPayload = Commands.serializeSingleMessageInBatchWithPayload(
msg.getMessageBuilder(), msg.getDataBuffer(), batchedMessageMetadataAndPayload);
}
@@ -169,6 +179,7 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
int uncompressedSize = batchedMessageMetadataAndPayload.readableBytes();
ByteBuf compressedPayload = compressor.encode(batchedMessageMetadataAndPayload);
batchedMessageMetadataAndPayload.release();
+ batchPayloadOwned = false;
if (compressionType != CompressionType.NONE) {
messageMetadata.setCompression(compressionType);
messageMetadata.setUncompressedSize(uncompressedSize);
@@ -181,6 +192,56 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
return compressedPayload;
}
+ /**
+ * Builds the batch payload (compressing it when configured) and encrypts it, returning the buffer that now owns
+ * the data. When compression or encryption replaced the container buffer, {@link #batchPayloadOwned} is cleared
+ * so failure recovery reallocates instead of reusing released memory; otherwise the container keeps ownership.
+ * If encryption fails, a payload the container no longer owns is released before rethrowing.
+ */
+ private ByteBuf buildAndEncryptBatchPayload() throws PulsarClientException {
+ ByteBuf compressedPayload = getCompressedBatchMetadataAndPayload();
+ try {
+ ByteBuf encryptedPayload = producer.encryptMessage(messageMetadata, compressedPayload);
+ if (encryptedPayload != compressedPayload) {
+ // Encryption allocated a new buffer and released the source payload.
+ batchPayloadOwned = false;
+ }
+ return encryptedPayload;
+ } catch (Throwable t) {
+ // Encryption failed: release the payload unless the container still owns it. A buffer the container
+ // owns is left for resetPayloadAfterFailedPublishing() to reuse.
+ if (!batchPayloadOwned) {
+ compressedPayload.release();
+ }
+ throw t;
+ }
+ }
+
+ /**
+ * Releases a payload whose ownership left the container when the command serialization failed and no op ever
+ * took it. A buffer the container still owns is left for {@link #resetPayloadAfterFailedPublishing()} to reuse.
+ */
+ private void releasePayloadIfOrphaned(ByteBuf payload) {
+ if (!batchPayloadOwned) {
+ payload.release();
+ }
+ }
+
+ /**
+ * Releases the command of an operation that was built from this container but never reached the send queue:
+ * in multi-batch mode a later sub-batch can fail to build after this one already produced its operation.
+ * The payload claim is shared with the container when its ownership never left (no compression or
+ * encryption): take it out of the pair before releasing, so the container keeps the buffer and
+ * {@link #resetPayloadAfterFailedPublishing()} reuses it when the messages are retried.
+ */
+ void releaseOrphanedOpCmd(ProducerImpl.OpSendMsg op) {
+ if (batchPayloadOwned) {
+ op.cmd.getSecond().retain();
+ }
+ op.cmd.release();
+ op.recycle();
+ }
+
void updateMaxBatchSize(int uncompressedSize) {
if (uncompressedSize > maxBatchSize) {
maxBatchSize = uncompressedSize;
@@ -212,6 +273,7 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
lowestSequenceId = -1L;
highestSequenceId = -1L;
batchedMessageMetadataAndPayload = null;
+ batchPayloadOwned = false;
currentTxnidMostBits = -1L;
currentTxnidLeastBits = -1L;
batchAllocatedSizeBytes = 0;
@@ -229,15 +291,18 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
if (firstCallback != null) {
firstCallback.sendComplete(ex, null);
}
- if (batchedMessageMetadataAndPayload != null) {
- ReferenceCountUtil.safeRelease(batchedMessageMetadataAndPayload);
- batchedMessageMetadataAndPayload = null;
- }
} catch (Throwable t) {
log.warn("[{}] [{}] Got exception while completing the callback for msg {}:", topicName,
producer.getProducerName(), lowestSequenceId, t);
+ } finally {
+ // In the finally: the callback is application code, and an exception escaping it must not
+ // skip the release - clear() below drops the reference without freeing the buffer.
+ if (batchPayloadOwned && batchedMessageMetadataAndPayload != null) {
+ ReferenceCountUtil.safeRelease(batchedMessageMetadataAndPayload);
+ }
+ batchPayloadOwned = false;
+ clear();
}
- clear();
}
@Override
@@ -250,10 +315,16 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
if (messages.size() == 1) {
messageMetadata.clear();
messageMetadata.copyFrom(messages.get(0).getMessageBuilder());
- ByteBuf encryptedPayload = producer.encryptMessage(messageMetadata, getCompressedBatchMetadataAndPayload());
+ ByteBuf encryptedPayload = buildAndEncryptBatchPayload();
updateAndReserveBatchAllocatedSize(encryptedPayload.capacity());
- ByteBufPair cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(),
- 1, null, messageMetadata, encryptedPayload);
+ ByteBufPair cmd;
+ try {
+ cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(),
+ 1, null, messageMetadata, encryptedPayload);
+ } catch (Throwable t) {
+ releasePayloadIfOrphaned(encryptedPayload);
+ throw t;
+ }
final OpSendMsg op;
// Shouldn't call create(MessageImpl<?> msg, ByteBufPair cmd, long sequenceId, SendCallback callback),
@@ -271,6 +342,9 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
// handle mgs size check as non-batched in `ProducerImpl.isMessageSizeExceeded`
if (op.getMessageHeaderAndPayloadSize() > getMaxMessageSize()) {
+ // The pair took the payload without retaining it, so cmd.release() also releases the
+ // container's claim: clear the ownership flag first, or discard() would double-release.
+ batchPayloadOwned = false;
cmd.release();
producer.semaphoreRelease(1);
producer.client.getMemoryLimitController().releaseMemory(
@@ -282,10 +356,12 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
lowestSequenceId = -1L;
return op;
}
- ByteBuf encryptedPayload = producer.encryptMessage(messageMetadata, getCompressedBatchMetadataAndPayload());
+ ByteBuf encryptedPayload = buildAndEncryptBatchPayload();
updateAndReserveBatchAllocatedSize(encryptedPayload.capacity());
if (encryptedPayload.readableBytes() > getMaxMessageSize()) {
- encryptedPayload.release();
+ // Release only the payload whose ownership left the container (compression or encryption
+ // replaced the batch buffer); the container's own buffer is released by discard() below.
+ releasePayloadIfOrphaned(encryptedPayload);
producer.semaphoreRelease(messages.size());
messages.forEach(msg -> producer.client.getMemoryLimitController()
.releaseMemory(msg.getUncompressedSize()));
@@ -303,8 +379,14 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
if (currentTxnidLeastBits != -1) {
messageMetadata.setTxnidLeastBits(currentTxnidLeastBits);
}
- ByteBufPair cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(),
- messageMetadata.getHighestSequenceId(), numMessagesInBatch, messageMetadata, encryptedPayload);
+ ByteBufPair cmd;
+ try {
+ cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(),
+ messageMetadata.getHighestSequenceId(), numMessagesInBatch, messageMetadata, encryptedPayload);
+ } catch (Throwable t) {
+ releasePayloadIfOrphaned(encryptedPayload);
+ throw t;
+ }
if (log.isDebugEnabled()) {
log.debug("[{}] [{}] Build batch msg seq:{}, highest-seq:{}, numMessagesInBatch: {}, uncompressedSize: {},"
@@ -325,7 +407,19 @@ class BatchMessageContainerImpl extends AbstractBatchMessageContainer {
@Override
public void resetPayloadAfterFailedPublishing() {
- if (batchedMessageMetadataAndPayload != null) {
+ if (messages.isEmpty()) {
+ // Nothing to rebuild: add() allocates a fresh buffer when the next batch starts.
+ return;
+ }
+ if (!batchPayloadOwned) {
+ // The failed build released the batch buffer (compression or encryption): reallocate instead of
+ // reusing memory that may already have been handed to another buffer.
+ batchedMessageMetadataAndPayload = allocator.buffer(
+ Math.min(maxBatchSize, getMaxMessageSize()));
+ batchPayloadOwned = true;
+ updateAndReserveBatchAllocatedSize(batchedMessageMetadataAndPayload.capacity());
+ } else {
+ // The container still owns the buffer: discard the partially written content.
batchedMessageMetadataAndPayload.readerIndex(0);
batchedMessageMetadataAndPayload.writerIndex(0);
}
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageKeyBasedContainer.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageKeyBasedContainer.java
index 96074b43466f..25477406ae81 100644
--- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageKeyBasedContainer.java
+++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/BatchMessageKeyBasedContainer.java
@@ -19,6 +19,7 @@
package org.apache.pulsar.client.impl;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
@@ -108,36 +109,39 @@ class BatchMessageKeyBasedContainer extends AbstractBatchMessageContainer {
@Override
public List<ProducerImpl.OpSendMsg> createOpSendMsgs() throws IOException {
+ // In key based batching, the sequence ids might not be ordered, for example,
+ // | key | sequence id list |
+ // | :-- | :--------------- |
+ // | A | 0, 3, 4 |
+ // | B | 1, 2 |
+ // The message order should be 1, 2, 0, 3, 4 so that a message with a sequence id <= 4 should be dropped.
+ // However, for a MessageMetadata with both `sequence_id` and `highest_sequence_id` fields, the broker will
+ // expect a strict order so that the batch of key "A" (0, 3, 4) will be dropped.
+ // Therefore, we should update the `sequence_id` field to the highest sequence id and remove the
+ // `highest_sequence_id` field to allow the weak order.
+ batches.values().forEach(batchMessageContainer -> {
+ batchMessageContainer.setLowestSequenceId(batchMessageContainer.getHighestSequenceId());
+ });
+ List<BatchMessageContainerImpl> sorted = batches.values().stream().sorted((o1, o2) ->
+ (int) (o1.getLowestSequenceId() - o2.getLowestSequenceId())
+ ).collect(Collectors.toList());
+ // Build the sub-batches in order: when a later sub-batch fails to build, the operations already built
+ // never reach the send queue and must be released here, or their commands leak.
+ List<ProducerImpl.OpSendMsg> ops = new ArrayList<>(sorted.size());
try {
- // In key based batching, the sequence ids might not be ordered, for example,
- // | key | sequence id list |
- // | :-- | :--------------- |
- // | A | 0, 3, 4 |
- // | B | 1, 2 |
- // The message order should be 1, 2, 0, 3, 4 so that a message with a sequence id <= 4 should be dropped.
- // However, for a MessageMetadata with both `sequence_id` and `highest_sequence_id` fields, the broker will
- // expect a strict order so that the batch of key "A" (0, 3, 4) will be dropped.
- // Therefore, we should update the `sequence_id` field to the highest sequence id and remove the
- // `highest_sequence_id` field to allow the weak order.
- batches.values().forEach(batchMessageContainer -> {
- batchMessageContainer.setLowestSequenceId(batchMessageContainer.getHighestSequenceId());
- });
- return batches.values().stream().sorted((o1, o2) ->
- (int) (o1.getLowestSequenceId() - o2.getLowestSequenceId())
- ).map(batchMessageContainer -> {
- try {
- return batchMessageContainer.createOpSendMsg();
- } catch (IOException e) {
- throw new IllegalStateException(e);
+ for (BatchMessageContainerImpl batchMessageContainer : sorted) {
+ ops.add(batchMessageContainer.createOpSendMsg());
+ }
+ } catch (Throwable t) {
+ for (int i = 0; i < ops.size(); i++) {
+ ProducerImpl.OpSendMsg op = ops.get(i);
+ if (op != null) {
+ sorted.get(i).releaseOrphanedOpCmd(op);
}
- }).collect(Collectors.toList());
- } catch (IllegalStateException e) {
- if (e.getCause() instanceof IOException) {
- throw (IOException) e.getCause();
- } else {
- throw e;
}
+ throw t;
}
+ return ops;
}
@Override
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java
index 31037bd44a69..c913698a5977 100644
--- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java
+++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java
@@ -515,7 +515,8 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
* @param payload
* @return a new payload
*/
- private ByteBuf applyCompression(ByteBuf payload) {
+ @VisibleForTesting
+ public ByteBuf applyCompression(ByteBuf payload) {
ByteBuf compressedPayload = compressor.encode(payload);
payload.release();
return compressedPayload;
@@ -571,7 +572,15 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
// Batch will be compressed when closed
// If a message has a delayed delivery time, we'll always send it individually
if (!isBatchMessagingEnabled() || msgMetadata.hasDeliverAtTime()) {
- compressedPayload = applyCompression(payload);
+ try {
+ compressedPayload = applyCompressionOrReleaseSource(payload);
+ } catch (Throwable t) {
+ // canEnqueueRequest has already acquired the send permit and reserved the memory, and no
+ // op exists to release them through the send lifecycle.
+ completeCallbackAndReleaseSemaphore(uncompressedSize, callback,
+ new PulsarClientException(t, msg.getSequenceId()));
+ return;
+ }
compressed = true;
// validate msg-size (For batching this will be check at the batch completion size)
@@ -671,8 +680,11 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
if (chunkId > 0 && conf.isBlockIfQueueFull() && !canEnqueueRequest(callback,
message.getSequenceId(), 0 /* The memory was already reserved */)) {
compressedPayload.release();
- client.getMemoryLimitController().releaseMemory(uncompressedSize - readStartIndex);
- semaphoreRelease(totalChunks - chunkId);
+ // In blocking mode canEnqueueRequest only fails on interruption, before taking this
+ // chunk's permit (and the memory was reserved once for the whole message, the ops of the
+ // earlier chunks carry no share): release the full reservation here and no permits -
+ // the permits of the already-built chunks are released by their own operations.
+ client.getMemoryLimitController().releaseMemory(uncompressedSize);
return;
}
synchronized (this) {
@@ -681,9 +693,33 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
final long sequenceId = updateMessageMetadataSequenceId(msgMetadata);
String uuid = totalChunks > 1 ? String.format("%s-%d", producerName, sequenceId) : null;
- serializeAndSendMessage(msg, payload, sequenceId, uuid, chunkId, totalChunks,
- readStartIndex, payloadChunkSize, compressedPayload, compressed,
- compressedPayload.readableBytes(), callback, chunkedMessageCtx, messageId);
+ try {
+ serializeAndSendMessage(msg, payload, sequenceId, uuid, chunkId, totalChunks,
+ readStartIndex, payloadChunkSize, compressedPayload, compressed,
+ compressedPayload.readableBytes(), callback, chunkedMessageCtx, messageId);
+ } catch (Throwable t) {
+ // For chunked persistent messages the last chunk's slice is never retained, so the base
+ // payload's own ref-count claim is the vehicle that releases it. A failure on any earlier
+ // chunk orphans that claim (the failing chunk's retained slice is already released by the
+ // send-path helper); release it here. Earlier chunks' retained slices keep the memory
+ // alive until their operations complete.
+ if (totalChunks > 1 && chunkId != totalChunks - 1 && TopicName.get(topic).isPersistent()) {
+ ReferenceCountUtil.safeRelease(compressedPayload);
+ }
+ if (totalChunks > 1) {
+ // The chunks after the failing one will never be built: return their pre-acquired
+ // permits (non-blocking mode acquires them up front) and their claims on the chunked
+ // context, which only created ops release. The failing chunk's own permit goes through
+ // the outer catch, the earlier chunks' through their own operations.
+ if (!conf.isBlockIfQueueFull()) {
+ semaphoreRelease(totalChunks - chunkId - 1);
+ }
+ for (int i = chunkId; i < totalChunks; i++) {
+ ReferenceCountUtil.safeRelease(chunkedMessageCtx);
+ }
+ }
+ throw t;
+ }
readStartIndex = ((chunkId + 1) * payloadChunkSize);
}
}
@@ -808,9 +844,9 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
// in this case compression has not been applied by the caller
// but we have to compress the payload if compression is configured
if (!compressed) {
- chunkPayload = applyCompression(chunkPayload);
+ chunkPayload = applyCompressionOrReleaseSource(chunkPayload);
}
- ByteBuf encryptedPayload = encryptMessage(msgMetadata, chunkPayload);
+ ByteBuf encryptedPayload = encryptMessageOrReleaseSource(msgMetadata, chunkPayload);
// When publishing during replication, we need to set the correct number of message in batch
// This is only used in tracking the publish rate stats
@@ -819,28 +855,29 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
: 1;
final OpSendMsg op;
if (msg.getSchemaState() == MessageImpl.SchemaState.Ready) {
- ByteBufPair cmd = sendMessage(producerId, sequenceId, numMessages, messageId, msgMetadata,
- encryptedPayload);
+ ByteBufPair cmd = sendMessageOrReleasePayload(producerId, sequenceId, numMessages, messageId,
+ msgMetadata, encryptedPayload);
op = OpSendMsg.create(rpcLatencyHistogram, msg, cmd, sequenceId, callback);
} else {
op = OpSendMsg.create(rpcLatencyHistogram, msg, null, sequenceId, callback);
+ // Hold on to the payload until the deferred command is built; if the op is failed before
+ // that happens, recycle() releases it instead of orphaning the buffer.
+ op.pendingPayload = encryptedPayload;
final MessageMetadata finalMsgMetadata = msgMetadata;
- op.rePopulate = () -> {
- if (msgMetadata.hasChunkId()) {
- // The message metadata is shared between all chunks in a large message
- // We need to reset the chunk id for each call of this method
- // It's safe to do that because there is only 1 thread to manipulate this message metadata
- finalMsgMetadata.setChunkId(chunkId);
- }
- op.cmd = sendMessage(producerId, sequenceId, numMessages, messageId, finalMsgMetadata,
- encryptedPayload);
- };
+ op.rePopulate = () -> buildDeferredCommand(op, finalMsgMetadata, producerId, sequenceId,
+ numMessages, messageId, chunkId);
}
op.setNumMessagesInBatch(numMessages);
op.setBatchSizeByte(encryptedPayload.readableBytes());
if (totalChunks > 1) {
op.totalChunks = totalChunks;
op.chunkId = chunkId;
+ if (chunkId != totalChunks - 1) {
+ // The message's memory was reserved once by canEnqueueRequest, and every release site
+ // (per-op ack, failPendingMessages, terminal state) releases op.uncompressedSize: only
+ // the last chunk's op may carry the size, or each chunk ack would release it again.
+ op.uncompressedSize = 0;
+ }
}
op.chunkedMessageCtx = chunkedMessageCtx;
lastSendFuture = callback.getFuture();
@@ -965,9 +1002,12 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
return compressedPayload;
}
+ // Nulled out on a successful hand-off: the finally-block releases the encrypted buffer whenever it was
+ // not returned to the caller (a crypto failure of any kind, or a failure while building it).
+ ByteBuf encryptedPayload = null;
try {
int maxSize = msgCrypto.getMaxOutputSize(compressedPayload.readableBytes());
- ByteBuf encryptedPayload = PulsarByteBufAllocator.DEFAULT.buffer(maxSize);
+ encryptedPayload = allocateEncryptedBuffer(maxSize);
ByteBuffer targetBuffer = encryptedPayload.nioBuffer(0, maxSize);
msgCrypto.encrypt(conf.getEncryptionKeys(), conf.getCryptoKeyReader(), () -> msgMetadata,
@@ -975,7 +1015,9 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
encryptedPayload.writerIndex(targetBuffer.remaining());
compressedPayload.release();
- return encryptedPayload;
+ ByteBuf result = encryptedPayload;
+ encryptedPayload = null;
+ return result;
} catch (PulsarClientException e) {
// Unless config is set to explicitly publish un-encrypted message upon failure, fail the request
if (conf.getCryptoFailureAction() == ProducerCryptoFailureAction.SEND) {
@@ -984,9 +1026,77 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
return compressedPayload;
}
throw e;
+ } finally {
+ ReferenceCountUtil.safeRelease(encryptedPayload);
}
}
+ ByteBuf allocateEncryptedBuffer(int maxSize) {
+ return PulsarByteBufAllocator.DEFAULT.buffer(maxSize);
+ }
+
+ /**
+ * Builds the send command and, when the serialization fails, releases the payload instead of orphaning
+ * it. For chunked messages this also returns the slice's claim on the shared payload buffer.
+ */
+ ByteBufPair sendMessageOrReleasePayload(long producerId, long sequenceId, int numMessages,
+ MessageId messageId, MessageMetadata msgMetadata, ByteBuf payload) {
+ try {
+ return sendMessage(producerId, sequenceId, numMessages, messageId, msgMetadata, payload);
+ } catch (Throwable t) {
+ ReferenceCountUtil.safeRelease(payload);
+ throw t;
+ }
+ }
+
+ /**
+ * Applies compression and, when the codec fails, releases the source payload instead of orphaning it:
+ * applyCompression() releases its input only after the codec succeeds. For chunked messages this also
+ * returns the slice's claim on the shared payload buffer.
+ */
+ ByteBuf applyCompressionOrReleaseSource(ByteBuf source) {
+ try {
+ return applyCompression(source);
+ } catch (Throwable t) {
+ ReferenceCountUtil.safeRelease(source);
+ throw t;
+ }
+ }
+
+ /**
+ * Applies encryption and, when the crypto fails, releases the source payload instead of orphaning it:
+ * encryptMessage() leaves the source with the caller on failure (its internal partial output buffer is
+ * released inside). For chunked messages this also returns the slice's claim on the shared payload
+ * buffer.
+ */
+ ByteBuf encryptMessageOrReleaseSource(MessageMetadata msgMetadata, ByteBuf source)
+ throws PulsarClientException {
+ try {
+ return encryptMessage(msgMetadata, source);
+ } catch (Throwable t) {
+ ReferenceCountUtil.safeRelease(source);
+ throw t;
+ }
+ }
+
+ /**
+ * Builds the command of an op whose serialization was deferred until the schema registration completed.
+ * The deferred payload stays owned by the op ({@code pendingPayload}) until the command is built: a
+ * failed construction must not release it, because the op remains pending and the next resend rebuilds
+ * from the same buffer; {@code recycle()} releases it when the op is failed instead.
+ */
+ void buildDeferredCommand(OpSendMsg op, MessageMetadata msgMetadata, long producerId, long sequenceId,
+ int numMessages, MessageId messageId, int chunkId) {
+ if (msgMetadata.hasChunkId()) {
+ // The message metadata is shared between all chunks in a large message. We need to reset the
+ // chunk id for each call of this method. It's safe to do that because there is only 1 thread
+ // to manipulate this message metadata.
+ msgMetadata.setChunkId(chunkId);
+ }
+ op.cmd = sendMessage(producerId, sequenceId, numMessages, messageId, msgMetadata, op.pendingPayload);
+ op.pendingPayload = null;
+ }
+
protected ByteBufPair sendMessage(long producerId, long sequenceId, int numMessages,
MessageId messageId, MessageMetadata msgMetadata,
ByteBuf compressedPayload) {
@@ -1413,6 +1523,16 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
callback.sendComplete(exception, null);
}
+ /**
+ * The send permits currently available on the producer queue ({@code Integer.MAX_VALUE} when the queue is
+ * unbounded). Visible for testing: the queue-accounting assertions on the send failure paths read it to
+ * confirm the permits acquired by {@code canEnqueueRequest} return to their starting point.
+ */
+ @VisibleForTesting
+ int availableSendPermitsForTesting() {
+ return semaphore.map(Semaphore::availablePermits).orElse(Integer.MAX_VALUE);
+ }
+
/**
* Checks message checksum to retry if message was corrupted while sending to broker. Recomputes checksum of the
* message header-payload again.
@@ -1573,6 +1693,7 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
// Volatile so a stale WriteInEventLoopCallback on an older connection's event loop can never observe
// cmd == this after the disposal has released it; the disposal also clears this field before releasing.
volatile ByteBufPair cmd;
+ ByteBuf pendingPayload;
SendCallback callback;
Runnable rePopulate;
ChunkedMessageCtx chunkedMessageCtx;
@@ -1602,6 +1723,7 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
msg = null;
msgs = null;
cmd = null;
+ pendingPayload = null;
callback = null;
rePopulate = null;
writeEventLoop = null;
@@ -1734,6 +1856,7 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
void recycle() {
ReferenceCountUtil.safeRelease(chunkedMessageCtx);
+ ReferenceCountUtil.safeRelease(pendingPayload);
initialize();
recyclerHandle.recycle(this);
}
@@ -1857,6 +1980,17 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
return op;
}
+ public boolean remove(OpSendMsg o) {
+ boolean removed = delegate.remove(o);
+ if (!removed && postponedOpSendMgs != null) {
+ removed = postponedOpSendMgs.remove(o);
+ }
+ if (removed) {
+ messagesCount.addAndGet(-o.numMessagesInBatch);
+ }
+ return removed;
+ }
+
public OpSendMsg peek() {
return delegate.peek();
}
@@ -2587,9 +2721,26 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
}
}
} catch (Throwable t) {
+ // The op may already sit in pendingMessages (e.g. the event loop rejected the write task on
+ // shutdown, after the op was queued): take it back out, or the recycled op would dangle in the
+ // queue and the resend would operate on a pooled instance.
+ pendingMessages.remove(op);
releaseSemaphoreForSendOp(op);
log.warn("[{}] [{}] error while closing out batch -- {}", topic, producerName, t);
- op.sendComplete(new PulsarClientException(t, op.sequenceId));
+ try {
+ // The callback runs application future handlers and can throw (every other sendComplete call
+ // site guards against this); the releases below must still run.
+ op.sendComplete(new PulsarClientException(t, op.sequenceId));
+ } catch (Throwable callbackEx) {
+ log.warn("[{}] [{}] Got exception while completing the callback for msg {}:", topic,
+ producerName, op.sequenceId, callbackEx);
+ }
+ if (op.writeEventLoop != null) {
+ // The cmd was retained for a write that never got queued: that reference has no owner
+ // anymore, so drop it on top of the op's own release.
+ ReferenceCountUtil.safeRelease(op.cmd);
+ }
+ releaseOpCmdAndRecycle(op);
}
}
@@ -2715,6 +2866,11 @@ public class ProducerImpl<T> extends ProducerBase<T> implements TimerTask, Conne
checkState(op.rePopulate != null);
op.rePopulate.run();
if (isMessageSizeExceeded(op)) {
+ // isMessageSizeExceeded already released the accounting and completed the callback:
+ // take the op out of the queue and release its freshly built command too, or it stays
+ // pending until the send timeout fails it - releasing the accounting a second time.
+ msgIterator.remove();
+ releaseOpCmdAndRecycle(op);
continue;
}
}
diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
index 4b1aa23ac4f9..270e19ac84b2 100644
--- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
@@ -18,28 +18,53 @@
*/
package org.apache.pulsar.client.impl;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;
+import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.UnpooledByteBufAllocator;
+import io.netty.buffer.UnpooledHeapByteBuf;
import io.netty.util.ReferenceCountUtil;
+import io.netty.util.Timeout;
+import io.netty.util.Timer;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Random;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import org.apache.pulsar.client.api.CompressionType;
+import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
import org.apache.pulsar.client.impl.conf.ProducerConfigurationData;
+import org.apache.pulsar.client.impl.metrics.InstrumentProvider;
+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.compression.CompressionCodec;
+import org.apache.pulsar.common.protocol.ByteBufPair;
+import org.apache.pulsar.common.protocol.Commands;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
public class BatchMessageContainerImplTest {
@@ -219,6 +244,530 @@ public class BatchMessageContainerImplTest {
return producer;
}
+ @DataProvider
+ public Object[][] compressionTypes() {
+ return new Object[][] {
+ {CompressionType.NONE},
+ {CompressionType.ZLIB},
+ };
+ }
+
+ /**
+ * A failure after the batch payload was built must not break the retry on the next flush, with or without
+ * compression: the container must not reuse a buffer released by the compression path.
+ */
+ @Test(dataProvider = "compressionTypes")
+ public void testRecoveryAfterBatchBuildFailure(CompressionType compressionType) throws Exception {
+ ProducerImpl<?> producer = createTestProducer(compressionType);
+
+ AtomicBoolean fail = new AtomicBoolean(true);
+ doAnswer(invocation -> {
+ if (fail.get()) {
+ throw new RuntimeException("mocked encryption failure");
+ }
+ return invocation.getArgument(1);
+ }).when(producer).encryptMessage(any(), any());
+ doAnswer(invocation -> {
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any());
+
+ BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer);
+ RecordingCompressionCodec compressor = installRecordingCompressor(batchMessageContainer, compressionType);
+ List<MessageImpl<?>> messages = addMessages(batchMessageContainer, 2);
+
+ // First build fails after the batch payload was produced; ProducerImpl.batchMessageAndSend() then resets.
+ assertThatThrownBy(batchMessageContainer::createOpSendMsg)
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked");
+ // The container keeps its buffer reference while messages remain.
+ assertNotNull(batchMessageContainer.batchedMessageMetadataAndPayload);
+ if (compressionType != CompressionType.NONE) {
+ // The compressed payload must not leak when encryption fails before anything took ownership of it.
+ assertEquals(compressor.lastEncoded.get().refCnt(), 0);
+ }
+ // Every build hands the batch buffer to the codec on this branch (the NONE codec returns it retained and
+ // the container drops its own claim), so re-entering the build without reset must fail fast instead of
+ // writing into released memory, with or without compression.
+ assertThatThrownBy(batchMessageContainer::createOpSendMsg)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("not owned");
+ batchMessageContainer.resetPayloadAfterFailedPublishing();
+
+ // The retry must succeed and produce a valid command instead of reusing a released buffer.
+ fail.set(false);
+ ProducerImpl.OpSendMsg op = batchMessageContainer.createOpSendMsg();
+ assertNotNull(op);
+ assertNotNull(op.cmd);
+ if (compressionType != CompressionType.NONE) {
+ // Guard against a vacuous green: the compression branch must actually have run.
+ assertTrue(compressor.encodeCalls.get() > 0, "the compression branch must actually have run");
+ }
+ op.cmd.release();
+ batchMessageContainer.clear();
+ messages.forEach(ReferenceCountUtil::safeRelease);
+ }
+
+ /**
+ * Encryption behaves like compression: encryptMessage() releases the source payload and returns a new buffer.
+ * A later failure must not make the retry reuse the released batch buffer.
+ */
+ @Test
+ public void testRecoveryAfterEncryptionFailure() throws Exception {
+ ProducerImpl<?> producer = createTestProducer(CompressionType.NONE);
+
+ AtomicBoolean failSend = new AtomicBoolean(true);
+ AtomicReference<ByteBuf> firstEncryptedRef = new AtomicReference<>();
+ doAnswer(invocation -> {
+ // Real encryption allocates a new buffer and releases the source payload.
+ ByteBuf source = invocation.getArgument(1);
+ ByteBuf encrypted = PulsarByteBufAllocator.DEFAULT.buffer(source.readableBytes());
+ encrypted.writeBytes(source);
+ source.release();
+ if (failSend.get()) {
+ firstEncryptedRef.set(encrypted);
+ }
+ return encrypted;
+ }).when(producer).encryptMessage(any(), any());
+ doAnswer(invocation -> {
+ if (failSend.getAndSet(false)) {
+ throw new RuntimeException("mocked send failure");
+ }
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any());
+
+ BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer);
+ List<MessageImpl<?>> messages = addMessages(batchMessageContainer, 2);
+
+ assertThatThrownBy(batchMessageContainer::createOpSendMsg).isInstanceOf(RuntimeException.class);
+ // The encrypted payload is orphaned once sendMessage fails; it must be released, not leaked.
+ assertEquals(firstEncryptedRef.get().refCnt(), 0);
+ batchMessageContainer.resetPayloadAfterFailedPublishing();
+
+ ProducerImpl.OpSendMsg op = batchMessageContainer.createOpSendMsg();
+ assertNotNull(op);
+ assertNotNull(op.cmd);
+ op.cmd.release();
+ batchMessageContainer.clear();
+ messages.forEach(ReferenceCountUtil::safeRelease);
+ }
+
+ /**
+ * Without compression or encryption, a failed build must still leave a retry that produces a well-formed,
+ * parseable SEND frame with balanced ref-counts.
+ */
+ @Test
+ public void testNoCompressionBuildFailureProducesValidFrame() throws Exception {
+ assertValidSendFrameAfterFailure(true, false); // failure in encryptMessage
+ assertValidSendFrameAfterFailure(false, true); // failure in sendMessage
+ }
+
+ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean failAtSend) throws Exception {
+ ProducerImpl<?> producer = createTestProducer(CompressionType.NONE);
+
+ AtomicBoolean failOnce = new AtomicBoolean(true);
+ if (failAtEncrypt) {
+ doAnswer(invocation -> {
+ if (failOnce.getAndSet(false)) {
+ throw new RuntimeException("mocked encryption failure");
+ }
+ return invocation.getArgument(1);
+ }).when(producer).encryptMessage(any(), any());
+ } else {
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ }
+ doAnswer(invocation -> {
+ if (failAtSend && failOnce.getAndSet(false)) {
+ throw new RuntimeException("mocked send failure");
+ }
+ MessageMetadata metadata = invocation.getArgument(4);
+ ByteBuf payload = invocation.getArgument(5);
+ return Commands.newSend(0L, metadata.hasSequenceId() ? metadata.getSequenceId() : 0L, 1,
+ Commands.ChecksumType.Crc32c, metadata, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any());
+
+ BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer);
+ List<MessageImpl<?>> messages = addMessages(batchMessageContainer, 3);
+
+ // First build fails after the batch payload was produced; ProducerImpl.batchMessageAndSend() then resets.
+ assertThatThrownBy(batchMessageContainer::createOpSendMsg).isInstanceOf(RuntimeException.class);
+ batchMessageContainer.resetPayloadAfterFailedPublishing();
+
+ // The retry must succeed and produce a well-formed SEND frame.
+ ProducerImpl.OpSendMsg op = batchMessageContainer.createOpSendMsg();
+ assertNotNull(op);
+ assertNotNull(op.cmd);
+
+ ByteBufPair cmd = op.cmd;
+ ByteBuf header = cmd.getFirst();
+ ByteBuf payloadBuf = cmd.getSecond();
+ int totalSize = header.getInt(0);
+ int cmdSize = header.getInt(4);
+ // The total-size field must equal the number of bytes that follow it.
+ assertEquals(totalSize, cmd.readableBytes() - 4,
+ "TOTAL_SIZE must equal the number of bytes following the total-size field");
+ // The command must parse cleanly as a SEND command. Skip both length fields: TOTAL_SIZE and
+ // CMD_SIZE (getInt above reads absolutely and does not move the reader index).
+ BaseCommand parsed = new BaseCommand();
+ header.markReaderIndex();
+ header.skipBytes(8);
+ parsed.parseFrom(header, cmdSize);
+ assertEquals(parsed.getType(), BaseCommand.Type.SEND);
+ header.resetReaderIndex();
+
+ // Ref-counts must be balanced: the op owns the batch buffer exactly once, and it is freed once.
+ assertEquals(payloadBuf.refCnt(), 1);
+ cmd.release();
+ assertEquals(payloadBuf.refCnt(), 0);
+
+ batchMessageContainer.clear();
+ messages.forEach(ReferenceCountUtil::safeRelease);
+ }
+
+ /**
+ * In multi-batch mode, a later sub-batch can fail to build after an earlier one already produced its
+ * operation. The already-built operations never reach the send queue, so their commands must be released
+ * by the container — otherwise every failed flush leaks command buffers, and repeated retries grow the
+ * direct memory usage.
+ */
+ @Test(dataProvider = "compressionTypes")
+ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType compressionType)
+ throws Exception {
+ ProducerImpl<?> producer = createTestProducer(compressionType);
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ List<ByteBufPair> builtPairs = new ArrayList<>();
+ // The pair clears its component references when released, so track the buffers at build time.
+ List<ByteBuf> builtHeaders = new ArrayList<>();
+ List<ByteBuf> builtPayloads = new ArrayList<>();
+ AtomicInteger sendCalls = new AtomicInteger();
+ doAnswer(invocation -> {
+ if (sendCalls.incrementAndGet() == 2) {
+ throw new RuntimeException("mocked second sub-batch failure");
+ }
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ ByteBufPair pair = ByteBufPair.get(header, payload);
+ builtPairs.add(pair);
+ builtHeaders.add(header);
+ builtPayloads.add(payload);
+ return pair;
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any());
+
+ BatchMessageKeyBasedContainer container = new BatchMessageKeyBasedContainer();
+ container.setProducer(producer);
+ List<MessageImpl<?>> messages = new ArrayList<>();
+ try {
+ for (int i = 0; i < 4; i++) {
+ MessageMetadata messageMetadata = new MessageMetadata();
+ messageMetadata.setSequenceId(i);
+ messageMetadata.setProducerName("producer");
+ messageMetadata.setPublishTime(System.currentTimeMillis());
+ messageMetadata.setPartitionKey(i < 2 ? "a" : "b");
+ ByteBuffer payload = ByteBuffer.wrap(("payload-" + i).getBytes(StandardCharsets.UTF_8));
+ MessageImpl<?> message = MessageImpl.create(messageMetadata, payload, Schema.BYTES, null);
+ messages.add(message);
+ container.add(message, null);
+ }
+
+ // Sub-batch "a" builds its operation, sub-batch "b" fails: the built command must not leak.
+ assertThatThrownBy(container::createOpSendMsgs)
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked second");
+ assertEquals(builtPairs.size(), 1);
+ // The pair itself is released (returned to its recycler) and the op recycled, not just the
+ // components freed.
+ assertEquals(builtPairs.get(0).refCnt(), 0);
+ assertEquals(builtHeaders.get(0).refCnt(), 0);
+ // The build handed the payload over to the command (with compression a new buffer; without, the
+ // NONE codec's retained batch buffer whose container claim was dropped), so it must be released too.
+ assertEquals(builtPayloads.get(0).refCnt(), 0);
+
+ // All messages stay in their sub-batches and the retry produces a complete batch again.
+ assertEquals(container.getNumMessagesInBatch(), 4);
+ container.resetPayloadAfterFailedPublishing();
+ List<ProducerImpl.OpSendMsg> ops = container.createOpSendMsgs();
+ assertEquals(ops.size(), 2);
+ ops.forEach(op -> op.cmd.release());
+ assertEquals(builtPairs.get(1).refCnt(), 0);
+ assertEquals(builtPairs.get(2).refCnt(), 0);
+ container.clear();
+ } finally {
+ messages.forEach(ReferenceCountUtil::safeRelease);
+ }
+ }
+
+ /**
+ * discard() invokes the application callback before releasing the batch buffer. The callback is application
+ * code (future handlers run inside sendComplete), so it can throw; the release must not be skipped when it
+ * does, or clear() drops the buffer's reference without freeing it and every failed batch with a throwing
+ * handler leaks its off-heap memory.
+ */
+ @Test
+ public void discardReleasesTheBatchBufferWhenTheCallbackThrows() throws Exception {
+ ProducerImpl<?> producer = createTestProducer(CompressionType.NONE);
+ List<ReleaseCountingByteBuf> containerBuffers = new ArrayList<>();
+ ByteBufAllocator recordingAllocator = mock(ByteBufAllocator.class);
+ doAnswer(invocation -> {
+ ReleaseCountingByteBuf buffer = new ReleaseCountingByteBuf(invocation.getArgument(0));
+ containerBuffers.add(buffer);
+ return buffer;
+ }).when(recordingAllocator).buffer(anyInt());
+
+ BatchMessageContainerImpl container = new BatchMessageContainerImpl(recordingAllocator);
+ container.setProducer(producer);
+ SendCallback throwingCallback = mock(SendCallback.class);
+ doThrow(new RuntimeException("mocked application callback failure"))
+ .when(throwingCallback).sendComplete(any(), any());
+
+ MessageImpl<?> message = createMessage(0);
+ try {
+ container.add(message, throwingCallback);
+ container.discard(new PulsarClientException("mocked discard"));
+
+ assertEquals(containerBuffers.size(), 1, "the batch buffer is the container's only allocation");
+ assertEquals(containerBuffers.get(0).refCnt(), 0,
+ "the batch buffer must be released even when the callback throws");
+ assertTrue(container.isEmpty(), "the container must still be cleared");
+ } finally {
+ ReferenceCountUtil.safeRelease(message);
+ }
+ }
+
+ /**
+ * A single message whose header + payload exceed the max message size is rejected after the command was
+ * already built. Without compression or encryption the command's payload IS the container's batch buffer:
+ * the NONE codec returns it retained and the container drops its own claim, so releasing the command is the
+ * last release of the buffer. The container used to keep reporting ownership after the command release, so
+ * the internal discard() released the already-freed buffer a third time — the failure was swallowed by
+ * safeRelease, but if the buffer had been recycled in between, the extra release corrupted a live buffer.
+ * The batch buffer must end up freed with release() invoked exactly once per claim: the codec hand-off and
+ * the command release.
+ */
+ @Test
+ public void oversizedSingleMessageReleasesTheBatchBufferExactlyOnce() throws Exception {
+ ProducerImpl<?> producer = createTestProducer(CompressionType.NONE);
+ doAnswer(invocation -> {
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any());
+
+ List<ReleaseCountingByteBuf> containerBuffers = new ArrayList<>();
+ ByteBufAllocator recordingAllocator = mock(ByteBufAllocator.class);
+ doAnswer(invocation -> {
+ ReleaseCountingByteBuf buffer = new ReleaseCountingByteBuf(invocation.getArgument(0));
+ containerBuffers.add(buffer);
+ return buffer;
+ }).when(recordingAllocator).buffer(anyInt());
+
+ BatchMessageContainerImpl container = new BatchMessageContainerImpl(recordingAllocator);
+ container.setProducer(producer);
+
+ MessageImpl<?> message = createMessage(0, Commands.DEFAULT_MAX_MESSAGE_SIZE + 2048);
+ try {
+ container.add(message, null);
+ assertNull(container.createOpSendMsg(), "an oversized single message must not produce an op");
+
+ assertEquals(containerBuffers.size(), 1, "the batch buffer is the container's only allocation");
+ ReleaseCountingByteBuf batchBuffer = containerBuffers.get(0);
+ assertEquals(batchBuffer.refCnt(), 0, "the batch buffer must be freed");
+ assertEquals(batchBuffer.releases(), 2,
+ "the batch buffer must be released once per claim (codec hand-off, then the command "
+ + "release), so discard() must not release it again");
+ } finally {
+ ReferenceCountUtil.safeRelease(message);
+ }
+ }
+
+ /**
+ * The multi-message variant of the same defect: the oversized check released the payload, and discard()
+ * then released the batch buffer again although its ownership had already left the container. Without
+ * compression that payload is the container's own batch buffer (returned retained by the NONE codec), so
+ * it must see exactly one release per claim and none from discard().
+ *
+ * <p>The ZLIB variant pins the other arm of that release: with compression the payload is a new buffer,
+ * so the oversized check itself must release it (dropping the orphan release would leak it — its
+ * ref-count stays at 1) while the batch buffer is released once, by the codec hand-off.
+ */
+ @Test(dataProvider = "compressionTypes")
+ public void oversizedBatchReleasesThePayloadExactlyOnce(CompressionType compressionType) throws Exception {
+ ProducerImpl<?> producer = createTestProducer(compressionType);
+
+ List<ReleaseCountingByteBuf> containerBuffers = new ArrayList<>();
+ ByteBufAllocator recordingAllocator = mock(ByteBufAllocator.class);
+ doAnswer(invocation -> {
+ ReleaseCountingByteBuf buffer = new ReleaseCountingByteBuf(invocation.getArgument(0));
+ containerBuffers.add(buffer);
+ return buffer;
+ }).when(recordingAllocator).buffer(anyInt());
+
+ BatchMessageContainerImpl container = new BatchMessageContainerImpl(recordingAllocator);
+ container.setProducer(producer);
+ RecordingCompressionCodec compressor = installRecordingCompressor(container, compressionType);
+
+ // Incompressible payload, so the compressed size still exceeds the 5MB limit in the ZLIB variant.
+ byte[] incompressible = new byte[Commands.DEFAULT_MAX_MESSAGE_SIZE + 2048];
+ new Random(42).nextBytes(incompressible);
+ int half = incompressible.length / 2;
+ MessageImpl<?> first = createMessage(0, incompressible, 0, half);
+ MessageImpl<?> second = createMessage(1, incompressible, half, incompressible.length - half);
+ try {
+ container.add(first, null);
+ container.add(second, null);
+ assertNull(container.createOpSendMsg(), "an oversized batch must not produce an op");
+
+ assertEquals(containerBuffers.size(), 1, "the batch buffer is the container's only allocation");
+ ReleaseCountingByteBuf batchBuffer = containerBuffers.get(0);
+ assertEquals(batchBuffer.refCnt(), 0, "the batch buffer must be freed");
+ if (compressionType != CompressionType.NONE) {
+ assertEquals(batchBuffer.releases(), 1,
+ "the batch buffer must be released exactly once, by the compression hand-off, never "
+ + "by discard()");
+ assertEquals(compressor.lastEncoded.get().refCnt(), 0,
+ "the compressed payload has left the container, so the oversized check must release it");
+ } else {
+ assertEquals(batchBuffer.releases(), 2,
+ "the batch buffer must be released once per claim (the codec hand-off, then the "
+ + "oversized check on the retained payload), never a third time by discard()");
+ }
+ } finally {
+ ReferenceCountUtil.safeRelease(first);
+ ReferenceCountUtil.safeRelease(second);
+ }
+ }
+
+ private MessageImpl<?> createMessage(long sequenceId) {
+ byte[] payload = "payload".getBytes(StandardCharsets.UTF_8);
+ return createMessage(sequenceId, payload, 0, payload.length);
+ }
+
+ private MessageImpl<?> createMessage(long sequenceId, int payloadSize) {
+ return createMessage(sequenceId, new byte[payloadSize], 0, payloadSize);
+ }
+
+ private MessageImpl<?> createMessage(long sequenceId, byte[] payloadArray, int offset, int length) {
+ MessageMetadata messageMetadata = new MessageMetadata();
+ messageMetadata.setSequenceId(sequenceId);
+ messageMetadata.setProducerName("producer");
+ messageMetadata.setPublishTime(System.currentTimeMillis());
+ ByteBuffer payload = ByteBuffer.wrap(payloadArray, offset, length);
+ return MessageImpl.create(messageMetadata, payload, Schema.BYTES, null);
+ }
+
+ /**
+ * A plain heap buffer that counts {@code release()} invocations, making a swallowed double release visible.
+ * (Netty 4.1's {@code WrappedByteBuf} is package-private, so the count is taken on the buffer itself.)
+ */
+ private static final class ReleaseCountingByteBuf extends UnpooledHeapByteBuf {
+
+ private int releases;
+
+ ReleaseCountingByteBuf(int initialCapacity) {
+ super(UnpooledByteBufAllocator.DEFAULT, initialCapacity, Integer.MAX_VALUE);
+ }
+
+ @Override
+ public boolean release() {
+ releases++;
+ return super.release();
+ }
+
+ @Override
+ public boolean release(int decrement) {
+ releases++;
+ return super.release(decrement);
+ }
+
+ int releases() {
+ return releases;
+ }
+ }
+
+ private ProducerImpl<?> createTestProducer(CompressionType compressionType) throws Exception {
+ ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData();
+ producerConfigurationData.setCompressionType(compressionType);
+ PulsarClientImpl pulsarClient = mock(PulsarClientImpl.class);
+ when(pulsarClient.newProducerId()).thenReturn(1L);
+ when(pulsarClient.getCnxPool()).thenReturn(mock(ConnectionPool.class));
+ when(pulsarClient.getMemoryLimitController()).thenReturn(mock(MemoryLimitController.class));
+ Timer timer = mock(Timer.class);
+ when(timer.newTimeout(any(), anyLong(), any())).thenReturn(mock(Timeout.class));
+ when(pulsarClient.timer()).thenReturn(timer);
+ ClientConfigurationData clientConfigurationData = new ClientConfigurationData();
+ clientConfigurationData.setStatsIntervalSeconds(0);
+ when(pulsarClient.getConfiguration()).thenReturn(clientConfigurationData);
+ when(pulsarClient.instrumentProvider()).thenReturn(InstrumentProvider.NOOP);
+
+ ProducerImpl<?> producer = mock(ProducerImpl.class, withSettings()
+ .useConstructor(pulsarClient, "persistent://public/default/batch-container-test",
+ producerConfigurationData, new CompletableFuture<>(), 0, Schema.BYTES,
+ null, Optional.empty())
+ .defaultAnswer(CALLS_REAL_METHODS));
+ return producer;
+ }
+
+ /**
+ * The container compresses through its own {@code compressor}, chosen from the producer's compression type
+ * in {@code setProducer()}. For a compressing type, swap in a codec that encodes into a new buffer and
+ * records it, so a test can assert on the compressed payload's fate; the container releases the source
+ * itself right after {@code encode()}, as with the real codecs. NONE keeps the real codec, which returns
+ * the batch buffer retained.
+ */
+ private static RecordingCompressionCodec installRecordingCompressor(AbstractBatchMessageContainer container,
+ CompressionType compressionType) {
+ if (compressionType == CompressionType.NONE) {
+ return null;
+ }
+ RecordingCompressionCodec compressor = new RecordingCompressionCodec();
+ container.compressor = compressor;
+ return compressor;
+ }
+
+ private static final class RecordingCompressionCodec implements CompressionCodec {
+ private final AtomicReference<ByteBuf> lastEncoded = new AtomicReference<>();
+ private final AtomicInteger encodeCalls = new AtomicInteger();
+
+ @Override
+ public ByteBuf encode(ByteBuf raw) {
+ ByteBuf compressed = PulsarByteBufAllocator.DEFAULT.buffer(raw.readableBytes());
+ compressed.writeBytes(raw, raw.readerIndex(), raw.readableBytes());
+ lastEncoded.set(compressed);
+ encodeCalls.incrementAndGet();
+ return compressed;
+ }
+
+ @Override
+ public ByteBuf decode(ByteBuf encoded, int uncompressedSize) {
+ throw new UnsupportedOperationException("decode is not used by the batch container");
+ }
+ }
+
+ private List<MessageImpl<?>> addMessages(BatchMessageContainerImpl batchMessageContainer, int count) {
+ List<MessageImpl<?>> messages = new ArrayList<>();
+ for (int i = 0; i < count; i++) {
+ MessageMetadata messageMetadata = new MessageMetadata();
+ messageMetadata.setSequenceId(i);
+ messageMetadata.setProducerName("producer");
+ messageMetadata.setPublishTime(System.currentTimeMillis());
+ ByteBuffer payload = ByteBuffer.wrap(("payload-" + i).getBytes(StandardCharsets.UTF_8));
+ MessageImpl<?> message = MessageImpl.create(messageMetadata, payload, Schema.BYTES, null);
+ messages.add(message);
+ batchMessageContainer.add(message, null);
+ }
+ return messages;
+ }
+
/**
* A batch carries one transaction id in its metadata, so every message in it inherits that transaction.
* Mixing a plain message into a transactional batch therefore silently enrolls it in the transaction: it
diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java
index a12b3a8b89af..ecf8601051bb 100644
--- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java
@@ -18,32 +18,56 @@
*/
package org.apache.pulsar.client.impl;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.anyLong;
import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.withSettings;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;
+import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.EventLoop;
+import io.netty.channel.EventLoopGroup;
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.Timeout;
+import io.netty.util.Timer;
+import io.netty.util.concurrent.ScheduledFuture;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
+import java.util.Optional;
+import java.util.TreeSet;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import org.apache.commons.lang3.reflect.FieldUtils;
+import org.apache.pulsar.client.api.CryptoKeyReader;
+import org.apache.pulsar.client.api.MessageCrypto;
+import org.apache.pulsar.client.api.ProducerCryptoFailureAction;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.impl.ProducerImpl.OpSendMsg;
import org.apache.pulsar.client.impl.ProducerImpl.OpSendMsgQueue;
+import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
+import org.apache.pulsar.client.impl.conf.ProducerConfigurationData;
+import org.apache.pulsar.client.impl.metrics.InstrumentProvider;
import org.apache.pulsar.client.impl.metrics.LatencyHistogram;
import org.apache.pulsar.common.api.proto.MessageMetadata;
import org.apache.pulsar.common.protocol.ByteBufPair;
@@ -257,6 +281,609 @@ public class ProducerImplTest {
}
}
+ private ProducerConfigurationData encryptedProducerConf() {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setEncryptionKeys(new TreeSet<>(Collections.singleton("key")));
+ conf.setCryptoKeyReader(mock(CryptoKeyReader.class));
+ return conf;
+ }
+
+ /** A client mock whose stubs satisfy the {@link ProducerImpl} constructor. */
+ private static PulsarClientImpl mockedPulsarClient() {
+ PulsarClientImpl client = mock(PulsarClientImpl.class);
+ when(client.newProducerId()).thenReturn(1L);
+ when(client.getCnxPool()).thenReturn(mock(ConnectionPool.class));
+ Timer timer = mock(Timer.class);
+ when(timer.newTimeout(any(), anyLong(), any())).thenReturn(mock(Timeout.class));
+ when(client.timer()).thenReturn(timer);
+ ClientConfigurationData clientConfigurationData = new ClientConfigurationData();
+ clientConfigurationData.setStatsIntervalSeconds(0);
+ when(client.getConfiguration()).thenReturn(clientConfigurationData);
+ when(client.instrumentProvider()).thenReturn(InstrumentProvider.NOOP);
+ EventLoopGroup eventLoopGroup = mock(EventLoopGroup.class);
+ when(eventLoopGroup.scheduleWithFixedDelay(any(Runnable.class), anyLong(), anyLong(), any()))
+ .thenReturn(mock(ScheduledFuture.class));
+ when(client.eventLoopGroup()).thenReturn(eventLoopGroup);
+ return client;
+ }
+
+ /**
+ * A producer mock built through the real constructor, so the {@code conf}, {@code log}, {@code client} and
+ * {@code msgCrypto} fields hold real values instead of needing reflection: {@code conf.setMessageCrypto()}
+ * installs a test crypto through the same constructor branch production uses.
+ */
+ @SuppressWarnings("unchecked")
+ private static ProducerImpl<byte[]> constructProducer(PulsarClientImpl client,
+ ProducerConfigurationData conf) {
+ return mock(ProducerImpl.class, withSettings()
+ .useConstructor(client, "persistent://public/default/producer-impl-test", conf,
+ new CompletableFuture<>(), 0, Schema.BYTES, null, Optional.empty())
+ .defaultAnswer(CALLS_REAL_METHODS));
+ }
+
+ /**
+ * When encryption fails in any way, the partially built encrypted buffer must be released instead of
+ * leaking; the source payload stays with the caller (the batch container decides its fate).
+ */
+ @Test
+ public void testEncryptMessageReleasesPartialBufferOnFailure() throws Exception {
+ MessageCrypto<?, ?> msgCrypto = mock(MessageCrypto.class);
+ when(msgCrypto.getMaxOutputSize(anyInt())).thenReturn(64);
+ doThrow(new RuntimeException("mocked encryption failure"))
+ .when(msgCrypto).encrypt(any(), any(), any(), any(), any());
+ ProducerConfigurationData conf = encryptedProducerConf();
+ conf.setMessageCrypto(msgCrypto);
+ ProducerImpl<byte[]> producer = constructProducer(mockedPulsarClient(), conf);
+
+ ByteBuf partial = Unpooled.buffer(64);
+ doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt());
+
+ ByteBuf source = Unpooled.buffer(8);
+ assertThatThrownBy(() -> producer.encryptMessage(new MessageMetadata(), source))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked encryption failure");
+ assertEquals(partial.refCnt(), 0, "the partially built encrypted buffer must not leak");
+ assertEquals(source.refCnt(), 1, "the source payload stays with the caller");
+ source.release();
+ }
+
+ /** The SEND crypto-failure action returns the unencrypted source; the partial buffer must not leak. */
+ @Test
+ public void testEncryptMessageCryptoFailureActionSendReleasesPartialBuffer() throws Exception {
+ MessageCrypto<?, ?> msgCrypto = mock(MessageCrypto.class);
+ when(msgCrypto.getMaxOutputSize(anyInt())).thenReturn(64);
+ doThrow(new PulsarClientException("mocked encryption failure"))
+ .when(msgCrypto).encrypt(any(), any(), any(), any(), any());
+ ProducerConfigurationData conf = encryptedProducerConf();
+ conf.setMessageCrypto(msgCrypto);
+ conf.setCryptoFailureAction(ProducerCryptoFailureAction.SEND);
+ ProducerImpl<byte[]> producer = constructProducer(mockedPulsarClient(), conf);
+
+ ByteBuf partial = Unpooled.buffer(64);
+ doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt());
+
+ ByteBuf source = Unpooled.buffer(8);
+ assertSame(producer.encryptMessage(new MessageMetadata(), source), source);
+ assertEquals(partial.refCnt(), 0, "the partially built encrypted buffer must not leak");
+ source.release();
+ }
+
+ /** A failed command serialization must release the payload instead of orphaning it. */
+ @Test
+ public void testSendMessageFailureReleasesPayload() throws Exception {
+ ProducerImpl<byte[]> producer = mock(ProducerImpl.class, CALLS_REAL_METHODS);
+ doThrow(new RuntimeException("mocked serialization failure"))
+ .when(producer)
+ .sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ ByteBuf payload = Unpooled.buffer(8);
+ assertThatThrownBy(() -> producer.sendMessageOrReleasePayload(
+ 1, 1, 1, null, new MessageMetadata(), payload))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked serialization failure");
+ assertEquals(payload.refCnt(), 0, "the payload must be released on a failed serialization");
+ }
+
+ /** A failing compression stage must release the source payload the codec left with the caller. */
+ @Test
+ public void testApplyCompressionFailureReleasesSource() throws Exception {
+ ProducerImpl<byte[]> producer = mock(ProducerImpl.class, CALLS_REAL_METHODS);
+ doThrow(new RuntimeException("mocked compression failure"))
+ .when(producer)
+ .applyCompression(any());
+
+ ByteBuf source = Unpooled.buffer(8);
+ assertThatThrownBy(() -> producer.applyCompressionOrReleaseSource(source))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked compression failure");
+ assertEquals(source.refCnt(), 0, "the source payload must be released on a failed compression");
+ }
+
+ /** A failing encryption stage must release the source payload encryptMessage() left with the caller. */
+ @Test
+ public void testEncryptMessageFailureReleasesSource() throws Exception {
+ ProducerImpl<byte[]> producer = mock(ProducerImpl.class, CALLS_REAL_METHODS);
+ doThrow(new PulsarClientException("mocked encryption failure"))
+ .when(producer)
+ .encryptMessage(any(), any());
+
+ ByteBuf source = Unpooled.buffer(8);
+ assertThatThrownBy(() -> producer.encryptMessageOrReleaseSource(new MessageMetadata(), source))
+ .isInstanceOf(PulsarClientException.class)
+ .hasMessageContaining("mocked encryption failure");
+ assertEquals(source.refCnt(), 0, "the source payload must be released on a failed encryption");
+ }
+
+ /**
+ * An op whose command is deferred until the schema is registered holds its payload; when the op is
+ * failed before the command was built (send timeout, producer close), recycle() must release it.
+ */
+ @Test
+ public void testPendingSchemaOpPayloadReleasedOnFailure() {
+ ByteBuf payload = Unpooled.buffer(8);
+ ProducerImpl.OpSendMsg op = ProducerImpl.OpSendMsg.create(
+ mock(LatencyHistogram.class),
+ mock(MessageImpl.class),
+ null,
+ 1L,
+ mock(SendCallback.class));
+ op.pendingPayload = payload;
+
+ op.recycle();
+
+ assertEquals(payload.refCnt(), 0, "the deferred payload must be released when the op is recycled");
+ }
+
+ /**
+ * A deferred command whose first construction fails must keep the payload alive: the op stays pending
+ * and the next resend (reconnect) rebuilds the command from the same buffer instead of touching a
+ * released one.
+ */
+ @Test
+ public void testDeferredCommandConstructionFailureThenRecovery() throws Exception {
+ ProducerImpl<byte[]> producer = mock(ProducerImpl.class, CALLS_REAL_METHODS);
+ AtomicInteger sendCalls = new AtomicInteger();
+ ByteBufPair builtCmd = mock(ByteBufPair.class);
+ doAnswer(invocation -> {
+ if (sendCalls.incrementAndGet() == 1) {
+ throw new RuntimeException("mocked header allocation failure");
+ }
+ return builtCmd;
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ ByteBuf payload = Unpooled.buffer(8);
+ ProducerImpl.OpSendMsg op = ProducerImpl.OpSendMsg.create(
+ mock(LatencyHistogram.class),
+ mock(MessageImpl.class),
+ null,
+ 1L,
+ mock(SendCallback.class));
+ op.pendingPayload = payload;
+
+ // First construction fails: the payload must stay with the op for the retry.
+ assertThatThrownBy(() -> producer.buildDeferredCommand(op, new MessageMetadata(), 1, 1, 1, null, -1))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked header allocation failure");
+ assertEquals(payload.refCnt(), 1, "the payload must stay alive for the next resend");
+ assertEquals(op.pendingPayload, payload, "the op keeps owning the deferred payload");
+ assertNull(op.cmd);
+
+ // The resend after reconnect rebuilds from the same buffer and hands it to the command.
+ producer.buildDeferredCommand(op, new MessageMetadata(), 1, 1, 1, null, -1);
+ assertEquals(op.cmd, builtCmd, "the retry must rebuild the command");
+ assertNull(op.pendingPayload, "the payload's ownership moved into the command");
+ // Recycling the op now must not release the payload again.
+ op.recycle();
+ assertEquals(payload.refCnt(), 1, "the payload belongs to the command now, not to the op");
+ payload.release();
+ }
+
+ /**
+ * Exercises the send-path wiring: {@code sendAsync()} itself must route through the stage helpers, so a
+ * failing compression or serialization stage releases the payload buffers instead of orphaning them. The
+ * helper-level tests above cannot detect a call site reverting to the bare method.
+ */
+ @Test
+ public void testSendPathFailureReleasesPayloadThroughTheStageHelpers() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ PulsarClientImpl client = mockedPulsarClient();
+ when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+
+ // A failing compression stage must release the message payload and complete the callback: the stage
+ // runs after canEnqueueRequest has acquired the send permit and reserved the memory, but before an op
+ // exists to release them, so both the buffer release and the callback are part of the wiring under
+ // test (leaving the future incomplete is the alternative this pins away).
+ doThrow(new RuntimeException("mocked compression failure"))
+ .when(producer).applyCompression(any());
+ MessageImpl<byte[]> first = newMessage("first");
+ SendCallback firstCallback = mock(SendCallback.class);
+ producer.sendAsync(first, firstCallback);
+ ArgumentCaptor<Throwable> compressionFailure = ArgumentCaptor.forClass(Throwable.class);
+ verify(firstCallback).sendComplete(compressionFailure.capture(), any());
+ assertTrue(String.valueOf(compressionFailure.getValue()).contains("mocked compression failure"),
+ "the callback must be failed with the compression failure");
+ assertEquals(first.getDataBuffer().refCnt(), 0,
+ "the payload must be released by the compression stage");
+
+ // A failing command serialization must release the compressed payload handed to it.
+ ByteBuf compressed = Unpooled.buffer(8);
+ doAnswer(invocation -> {
+ ByteBuf source = invocation.getArgument(0);
+ source.release();
+ return compressed;
+ }).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ doThrow(new RuntimeException("mocked serialization failure"))
+ .when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ MessageImpl<byte[]> second = newMessage("second");
+ SendCallback secondCallback = mock(SendCallback.class);
+ producer.sendAsync(second, secondCallback);
+ verify(secondCallback).sendComplete(any(), any());
+ assertEquals(compressed.refCnt(), 0,
+ "the compressed payload must be released by the serialization stage");
+ }
+
+ /**
+ * Exercises the chunked send path: when a chunk's command serialization fails, the chunk slice's claim on
+ * the shared payload buffer must be released by {@code sendMessageOrReleasePayload}. This cannot be detected
+ * by the helper-level tests alone, because the call site is only reachable with chunking enabled and a
+ * payload large enough to split.
+ */
+ @Test
+ public void testChunkedSendFailureReleasesChunkSliceClaim() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ conf.setChunkingEnabled(true);
+ PulsarClientImpl client = mockedPulsarClient();
+ when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ // A small max message size forces the 10 KB payload into several chunks.
+ ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+ when(connectionHandler.getMaxMessageSize()).thenReturn(1024);
+ doReturn(connectionHandler).when(producer).getConnectionHandler();
+ // No-op compression, so the compressed payload is the message payload itself.
+ doAnswer(invocation -> invocation.getArgument(0)).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ doThrow(new RuntimeException("mocked chunk serialization failure"))
+ .when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ MessageImpl<byte[]> message = newMessage(new byte[10 * 1024]);
+ ByteBuf payload = message.getDataBuffer();
+ SendCallback callback = mock(SendCallback.class);
+ producer.sendAsync(message, callback);
+
+ ArgumentCaptor<Throwable> throwableCaptor = ArgumentCaptor.forClass(Throwable.class);
+ verify(callback).sendComplete(throwableCaptor.capture(), any());
+ assertTrue(throwableCaptor.getValue().getMessage().contains("mocked chunk serialization failure"));
+ // The failing chunk's retained slice is released by the send-path helper and the base payload's own
+ // claim by the chunk-loop failure handling. With no compression the base is the message payload
+ // itself, so nothing may remain after the failure.
+ assertEquals(payload.refCnt(), 0, "the failing chunk slice's claim must be released");
+ }
+
+ /**
+ * When the whole payload is compressed before chunking, a failure on a middle chunk must release not only
+ * the failing chunk's retained slice but also the base payload's own claim — the one the last chunk would
+ * have carried. Otherwise the compressed base buffer leaks on every failed chunked send.
+ */
+ @Test
+ public void testChunkedSendFailureReleasesCompressedBasePayload() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ conf.setChunkingEnabled(true);
+ PulsarClientImpl client = mockedPulsarClient();
+ when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+ when(connectionHandler.getMaxMessageSize()).thenReturn(1024);
+ doReturn(connectionHandler).when(producer).getConnectionHandler();
+
+ ByteBuf compressed = Unpooled.buffer(10 * 1024);
+ compressed.writeBytes(new byte[10 * 1024]);
+ doAnswer(invocation -> {
+ ByteBuf source = invocation.getArgument(0);
+ source.release();
+ return compressed;
+ }).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ doThrow(new RuntimeException("mocked chunk serialization failure"))
+ .when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ SendCallback callback = mock(SendCallback.class);
+ producer.sendAsync(newMessage(new byte[10 * 1024]), callback);
+ verify(callback).sendComplete(any(), any());
+ // The failing chunk's retained slice and the base payload's own claim must both be released.
+ assertEquals(compressed.refCnt(), 0, "the compressed base payload must be released on a chunk failure");
+ }
+
+ /**
+ * A compression failure on the non-batch path happens after canEnqueueRequest has already acquired the
+ * send permit and reserved the message size in the client memory controller, but before any op exists to
+ * release them through the send lifecycle. The failure must still complete the callback (otherwise the
+ * returned future hangs) and return the accounting to its starting point, or repeated compression
+ * failures exhaust the producer queue and the client memory limit.
+ */
+ @Test
+ public void compressionFailureBeforeAnOpExistsReleasesPermitAndMemory() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ // The default of 0 means an unbounded queue with no permit accounting to assert on.
+ conf.setMaxPendingMessages(100);
+ PulsarClientImpl client = mockedPulsarClient();
+ MemoryLimitController memoryLimitController = new MemoryLimitController(1024 * 1024);
+ when(client.getMemoryLimitController()).thenReturn(memoryLimitController);
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ doThrow(new RuntimeException("mocked compression failure"))
+ .when(producer).applyCompression(any());
+
+ int maxPermits = producer.availableSendPermitsForTesting();
+ SendCallback callback = mock(SendCallback.class);
+ producer.sendAsync(newMessage(new byte[16]), callback);
+
+ ArgumentCaptor<Throwable> throwableCaptor = ArgumentCaptor.forClass(Throwable.class);
+ verify(callback).sendComplete(throwableCaptor.capture(), any());
+ assertTrue(String.valueOf(throwableCaptor.getValue()).contains("mocked compression failure"),
+ "the callback must be failed with the compression failure");
+ assertEquals(memoryLimitController.currentUsage(), 0,
+ "the memory reserved by canEnqueueRequest must be released");
+ assertEquals(producer.availableSendPermitsForTesting(), maxPermits,
+ "the send permit acquired by canEnqueueRequest must be released");
+ }
+
+ /**
+ * In non-blocking mode (the default) the chunk loop pre-acquires a send permit for every chunk of the
+ * message before building any of them. A chunk-build failure used to strand the permits of the chunks
+ * that were never going to be built — and their claims on the chunked-message context — permanently
+ * shrinking the producer queue (failing the first of N chunks cost N - 1 permits). The failure handling
+ * must return every pre-acquired permit: the failing chunk's own permit is released by the outer catch,
+ * and the never-built chunks' permits here.
+ */
+ @Test
+ public void chunkedNonBlockingFailureReturnsThePreAcquiredPermits() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ conf.setChunkingEnabled(true);
+ // The default of 0 means an unbounded queue: the permits this test asserts on only exist with a bound.
+ conf.setMaxPendingMessages(100);
+ PulsarClientImpl client = mockedPulsarClient();
+ MemoryLimitController memoryLimitController = new MemoryLimitController(1024 * 1024);
+ when(client.getMemoryLimitController()).thenReturn(memoryLimitController);
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ // A small max message size forces the 10 KB payload into several chunks.
+ ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+ when(connectionHandler.getMaxMessageSize()).thenReturn(1024);
+ doReturn(connectionHandler).when(producer).getConnectionHandler();
+ doAnswer(invocation -> invocation.getArgument(0)).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ doThrow(new RuntimeException("mocked chunk serialization failure"))
+ .when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ int maxPermits = producer.availableSendPermitsForTesting();
+ SendCallback callback = mock(SendCallback.class);
+ producer.sendAsync(newMessage(new byte[10 * 1024]), callback);
+
+ ArgumentCaptor<Throwable> chunkFailure = ArgumentCaptor.forClass(Throwable.class);
+ verify(callback).sendComplete(chunkFailure.capture(), any());
+ assertTrue(String.valueOf(chunkFailure.getValue()).contains("mocked chunk serialization failure"),
+ "the chunk-build path must be the one that failed: " + chunkFailure.getValue());
+ assertEquals(memoryLimitController.currentUsage(), 0,
+ "the memory reserved for the message must be released");
+ assertEquals(producer.availableSendPermitsForTesting(), maxPermits,
+ "the permits pre-acquired for the never-built chunks must be released");
+ }
+
+ /**
+ * The client memory for a chunked message is reserved once, by canEnqueueRequest before the chunk loop,
+ * but released per chunk operation (every release site releases op.uncompressedSize). Every chunk op used
+ * to carry the full message size, so each of the N chunk acks released the whole reservation again and the
+ * memory gauge drifted low by (N - 1) sizes - silently disabling the client memory limit for chunked
+ * producers. Only the last chunk's op may carry the size: the ops of one message must carry the
+ * reservation exactly once in total.
+ */
+ @Test
+ public void chunkedOpsCarryTheMemoryReservationExactlyOnce() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ conf.setChunkingEnabled(true);
+ conf.setMaxPendingMessages(100);
+ PulsarClientImpl client = mockedPulsarClient();
+ when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ // A small max message size forces the 10 KB payload into several chunks.
+ ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+ when(connectionHandler.getMaxMessageSize()).thenReturn(1024);
+ doReturn(connectionHandler).when(producer).getConnectionHandler();
+ doAnswer(invocation -> invocation.getArgument(0)).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ doAnswer(invocation -> {
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = Unpooled.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+ List<ProducerImpl.OpSendMsg> ops = new ArrayList<>();
+ doAnswer(invocation -> {
+ ops.add(invocation.getArgument(0));
+ return null;
+ }).when(producer).processOpSendMsg(any());
+
+ int payloadSize = 10 * 1024;
+ producer.sendAsync(newMessage(new byte[payloadSize]), mock(SendCallback.class));
+
+ assertTrue(ops.size() > 1, "the payload must actually be split into chunks");
+ for (int i = 0; i < ops.size(); i++) {
+ ProducerImpl.OpSendMsg op = ops.get(i);
+ boolean isLastChunk = op.chunkId == op.totalChunks - 1;
+ assertEquals(op.uncompressedSize, isLastChunk ? payloadSize : 0,
+ "chunk " + op.chunkId + " of " + op.totalChunks
+ + " must carry the reservation only if it is the last chunk");
+ }
+ assertEquals(ops.stream().mapToLong(op -> op.uncompressedSize).sum(), payloadSize,
+ "the chunk ops must carry the message's reservation exactly once in total");
+ ops.forEach(op -> ReferenceCountUtil.safeRelease(op.cmd));
+ }
+
+ /**
+ * In blocking mode the per-chunk canEnqueueRequest only fails on interruption, before taking that chunk's
+ * permit and after the whole message's memory was reserved once. The failure branch used to release only
+ * uncompressedSize - readStartIndex (leaking the already-chunked bytes on the memory gauge, since the
+ * earlier chunk ops carry no share) and N - chunkId permits that were never acquired in blocking mode.
+ * Triggered by setting the interrupt flag while chunk 0's command is being built: Semaphore.acquire()
+ * throws immediately on an interrupt flag even with permits available.
+ */
+ @Test
+ public void chunkedBlockingInterruptionReleasesTheWholeReservationAndNoForeignPermits() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ conf.setChunkingEnabled(true);
+ conf.setMaxPendingMessages(100);
+ conf.setBlockIfQueueFull(true);
+ PulsarClientImpl client = mockedPulsarClient();
+ MemoryLimitController memoryLimitController = new MemoryLimitController(1024 * 1024);
+ when(client.getMemoryLimitController()).thenReturn(memoryLimitController);
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ // A small max message size forces the 10 KB payload into several chunks.
+ ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+ when(connectionHandler.getMaxMessageSize()).thenReturn(1024);
+ doReturn(connectionHandler).when(producer).getConnectionHandler();
+ doAnswer(invocation -> invocation.getArgument(0)).when(producer).applyCompression(any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ AtomicInteger sendCalls = new AtomicInteger();
+ doAnswer(invocation -> {
+ if (sendCalls.incrementAndGet() == 1) {
+ // Chunk 0 is built normally; chunk 1's enqueue must observe the interrupt flag.
+ Thread.currentThread().interrupt();
+ }
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = Unpooled.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+ List<ProducerImpl.OpSendMsg> ops = new ArrayList<>();
+ doAnswer(invocation -> {
+ ops.add(invocation.getArgument(0));
+ return null;
+ }).when(producer).processOpSendMsg(any());
+
+ int maxPermits = producer.availableSendPermitsForTesting();
+ SendCallback callback = mock(SendCallback.class);
+ try {
+ producer.sendAsync(newMessage(new byte[10 * 1024]), callback);
+ } finally {
+ // canEnqueueRequest's interrupt handler restores the interrupt status; don't leak it to other tests.
+ Thread.interrupted();
+ }
+
+ assertEquals(ops.size(), 1, "chunk 0 must have been built before the interruption failed chunk 1");
+ verify(callback).sendComplete(any(), any());
+ assertEquals(memoryLimitController.currentUsage(), 0,
+ "the whole memory reservation must be released, including the already-chunked bytes");
+ assertEquals(producer.availableSendPermitsForTesting(), maxPermits - 1,
+ "only the built chunk's permit may still be held (by its own operation)");
+ producer.releaseSemaphoreForSendOp(ops.get(0));
+ assertEquals(producer.availableSendPermitsForTesting(), maxPermits,
+ "completing the built chunk's op must return the queue to its starting point");
+ ops.forEach(op -> ReferenceCountUtil.safeRelease(op.cmd));
+ }
+
+ /**
+ * A message whose schema is not yet registered goes through the deferred-command branch of the real send
+ * path: the op holds the encrypted payload (pendingPayload) instead of a command and builds it later via
+ * rePopulate(). A failed first build must keep the payload with the op for the resend, and failing the op
+ * before the command was built must release it.
+ */
+ @Test
+ public void testDeferredSchemaOpWiringThroughSendPath() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ PulsarClientImpl client = mockedPulsarClient();
+ when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+ // Keep the schema state non-Ready: the real populateMessageSchema would mark it Ready for the
+ // matching schema, so the deferred branch would not run.
+ doAnswer(invocation -> true).when(producer).populateMessageSchema(any(), any());
+ doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any());
+ AtomicReference<ProducerImpl.OpSendMsg> captured = new AtomicReference<>();
+ doAnswer(invocation -> {
+ captured.set(invocation.getArgument(0));
+ return null;
+ }).when(producer).processOpSendMsg(any());
+ AtomicInteger sendCalls = new AtomicInteger();
+ doAnswer(invocation -> {
+ if (sendCalls.incrementAndGet() == 1) {
+ throw new RuntimeException("mocked deferred header allocation failure");
+ }
+ ByteBuf payload = invocation.getArgument(5);
+ ByteBuf header = Unpooled.buffer();
+ header.writeInt(4 + 4 + payload.readableBytes());
+ header.writeInt(0);
+ return ByteBufPair.get(header, payload);
+ }).when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any());
+
+ MessageImpl<byte[]> message = newMessage("deferred");
+ ByteBuf payload = message.getDataBuffer();
+ producer.sendAsync(message, mock(SendCallback.class));
+
+ ProducerImpl.OpSendMsg op = captured.get();
+ assertNotNull(op);
+ assertNull(op.cmd, "the command must be deferred until the schema is registered");
+ assertSame(op.pendingPayload, payload, "the op must own the deferred payload");
+ assertNotNull(op.rePopulate);
+
+ // First build fails: the payload stays with the op for the next resend.
+ assertThatThrownBy(() -> op.rePopulate.run())
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("mocked deferred");
+ assertEquals(payload.refCnt(), 1, "the payload must stay alive for the next resend");
+ assertSame(op.pendingPayload, payload, "the op keeps owning the deferred payload");
+ assertNull(op.cmd);
+
+ // The resend after reconnect rebuilds from the same buffer and hands it to the command.
+ op.rePopulate.run();
+ assertNotNull(op.cmd);
+ assertNull(op.pendingPayload, "the payload's ownership moved into the command");
+ op.cmd.release();
+ assertEquals(payload.refCnt(), 0, "the payload belongs to the command now, not to the op");
+ op.recycle();
+
+ // A second deferred op failed before the schema was registered (e.g. send timeout): recycle() must
+ // release the deferred payload instead of orphaning it.
+ MessageImpl<byte[]> secondMessage = newMessage("deferred-2");
+ ByteBuf secondPayload = secondMessage.getDataBuffer();
+ producer.sendAsync(secondMessage, mock(SendCallback.class));
+ ProducerImpl.OpSendMsg secondOp = captured.get();
+ assertSame(secondOp.pendingPayload, secondPayload);
+ secondOp.recycle();
+ assertEquals(secondPayload.refCnt(), 0,
+ "the deferred payload must be released when the op is failed");
+ }
+
+ private static MessageImpl<byte[]> newMessage(String content) {
+ MessageMetadata metadata = new MessageMetadata();
+ metadata.setPublishTime(System.currentTimeMillis());
+ return MessageImpl.create(metadata,
+ ByteBuffer.wrap(content.getBytes(StandardCharsets.UTF_8)), Schema.BYTES, null);
+ }
+
+ private static MessageImpl<byte[]> newMessage(byte[] content) {
+ MessageMetadata metadata = new MessageMetadata();
+ metadata.setPublishTime(System.currentTimeMillis());
+ return MessageImpl.create(metadata, ByteBuffer.wrap(content), Schema.BYTES, null);
+ }
+
/**
* Regression test for the send-timeout vs in-flight-write race (scenario B):
*
@@ -427,6 +1054,114 @@ public class ProducerImplTest {
verify(stats).updateNumMsgsSent(1, 0);
}
+ /**
+ * If the connection's event loop rejects the write task (shutting down during a producer close or
+ * reconnect), the catch of processOpSendMsg used to release only the permit and the memory: the cmd -
+ * retained for a write that never got queued - leaked together with the chunked-message context claim,
+ * and the op stayed in pendingMessages unrecycled. The catch must take the op back out of the queue,
+ * drop the orphaned write reference, and release the op's own command.
+ */
+ @Test
+ @SuppressWarnings("unchecked")
+ public void processOpSendMsgCatchReleasesCmdAndRemovesTheOpFromTheQueue() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ PulsarClientImpl client = mockedPulsarClient();
+ Mockito.when(client.getMemoryLimitController())
+ .thenReturn(Mockito.mock(MemoryLimitController.class));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+
+ // The event loop rejects every task: the deferred release falls back to inline.
+ EventLoop eventLoop = Mockito.mock(EventLoop.class);
+ Mockito.doThrow(new RejectedExecutionException("mocked event loop shutdown"))
+ .when(eventLoop).execute(any(Runnable.class));
+ Channel channel = Mockito.mock(Channel.class);
+ Mockito.when(channel.eventLoop()).thenReturn(eventLoop);
+ ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class);
+ Mockito.when(ctx.channel()).thenReturn(channel);
+ ClientCnx cnx = Mockito.mock(ClientCnx.class);
+ Mockito.when(cnx.ctx()).thenReturn(ctx);
+ Mockito.doReturn(cnx).when(producer).getCnxIfReady();
+
+ OpSendMsgQueue pendingQueue = producer.pendingMessages;
+
+ ByteBufPair cmd = ByteBufPair.get(
+ Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)),
+ Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8)));
+ MessageImpl<?> msg = Mockito.mock(MessageImpl.class);
+ Mockito.when(msg.getUncompressedSize()).thenReturn(10);
+ SendCallback callback = Mockito.mock(SendCallback.class);
+ OpSendMsg op = OpSendMsg.create(
+ Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, callback);
+ op.totalChunks = 1;
+ op.chunkId = 0;
+ op.numMessagesInBatch = 1;
+ // Bypass the message-dependent paths (batch scheduling / schema registration / size checks) so the test
+ // focuses on the write hand-off.
+ op.msg = null;
+
+ producer.processOpSendMsg(op);
+
+ assertEquals(pendingQueue.size(), 0, "the failed op must be taken back out of the pending queue");
+ assertEquals(pendingQueue.messagesCount(), 0, "the queue's message accounting must follow");
+ assertEquals(cmd.refCnt(), 0,
+ "both the op's reference and the orphaned write reference must be released");
+ verify(callback).sendComplete(any(), any());
+ }
+
+ /**
+ * The same failure path with a callback that throws: sendComplete runs application future handlers, and
+ * the exception escaping it must not skip the cmd/op cleanup that follows - by then the op is already out
+ * of the queue with its accounting released, so a skip would strand the buffers for good.
+ */
+ @Test
+ @SuppressWarnings("unchecked")
+ public void processOpSendMsgCatchSurvivesAThrowingCallback() throws Exception {
+ ProducerConfigurationData conf = new ProducerConfigurationData();
+ conf.setBatchingEnabled(false);
+ PulsarClientImpl client = mockedPulsarClient();
+ Mockito.when(client.getMemoryLimitController())
+ .thenReturn(Mockito.mock(MemoryLimitController.class));
+ ProducerImpl<byte[]> producer = constructProducer(client, conf);
+ producer.setState(ProducerImpl.State.Ready);
+
+ EventLoop eventLoop = Mockito.mock(EventLoop.class);
+ Mockito.doThrow(new RejectedExecutionException("mocked event loop shutdown"))
+ .when(eventLoop).execute(any(Runnable.class));
+ Channel channel = Mockito.mock(Channel.class);
+ Mockito.when(channel.eventLoop()).thenReturn(eventLoop);
+ ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class);
+ Mockito.when(ctx.channel()).thenReturn(channel);
+ ClientCnx cnx = Mockito.mock(ClientCnx.class);
+ Mockito.when(cnx.ctx()).thenReturn(ctx);
+ Mockito.doReturn(cnx).when(producer).getCnxIfReady();
+
+ OpSendMsgQueue pendingQueue = producer.pendingMessages;
+
+ ByteBufPair cmd = ByteBufPair.get(
+ Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)),
+ Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8)));
+ MessageImpl<?> msg = Mockito.mock(MessageImpl.class);
+ Mockito.when(msg.getUncompressedSize()).thenReturn(10);
+ SendCallback callback = Mockito.mock(SendCallback.class);
+ Mockito.doThrow(new RuntimeException("mocked application callback failure"))
+ .when(callback).sendComplete(any(), any());
+ OpSendMsg op = OpSendMsg.create(
+ Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, callback);
+ op.totalChunks = 1;
+ op.chunkId = 0;
+ op.numMessagesInBatch = 1;
+ op.msg = null;
+
+ producer.processOpSendMsg(op);
+
+ assertEquals(pendingQueue.size(), 0, "the failed op must be taken back out of the pending queue");
+ assertEquals(cmd.refCnt(), 0,
+ "the cleanup must still run after the throwing callback");
+ verify(callback).sendComplete(any(), any());
+ }
+
/**
* Regression test for the stale write-callback window: an op is handed to connection A's event loop, then
* re-sent on connection B after a reconnect (so {@code op.writeEventLoop} now points to B's loop), and the
(pulsar) 01/03: [improve][client] Coalesce message listener drain scheduling (#26531)
Posted by lh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 8c406a6b2864b59e99e03b2756ba841d7d2f9dd3
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Mon Sep 14 09:29:22 2026 +0300
[improve][client] Coalesce message listener drain scheduling (#26531)
(cherry picked from commit a6278f292533259274a0246cc2a08c8f463ae8e0)
(cherry picked from commit c377a7853a3c05d3c4accafe5d247277195fd060)
---
.../impl/ListenerTaskSchedulingBenchmark.java | 102 +++++++++
.../apache/pulsar/client/impl/ConsumerBase.java | 51 +++--
.../pulsar/client/impl/ListenerTaskScheduler.java | 73 ++++++
.../impl/ConsumerListenerSchedulingTest.java | 172 +++++++++++++++
.../client/impl/ListenerTaskSchedulerTest.java | 245 +++++++++++++++++++++
5 files changed, 624 insertions(+), 19 deletions(-)
diff --git a/microbench/src/main/java/org/apache/pulsar/client/impl/ListenerTaskSchedulingBenchmark.java b/microbench/src/main/java/org/apache/pulsar/client/impl/ListenerTaskSchedulingBenchmark.java
new file mode 100644
index 000000000000..282e631f4573
--- /dev/null
+++ b/microbench/src/main/java/org/apache/pulsar/client/impl/ListenerTaskSchedulingBenchmark.java
@@ -0,0 +1,102 @@
+/*
+ * 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 java.util.Queue;
+import java.util.concurrent.Executor;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+/** Models queued arrival/drain scheduling, excluding message decoding, listener calls and OS thread handoff. */
+@State(Scope.Thread)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@Fork(2)
+public class ListenerTaskSchedulingBenchmark {
+ @Param({"1", "16", "1024"})
+ public int notifications;
+
+ @Param({"fresh", "reused", "coalesced"})
+ public String scheduling;
+
+ private QueuedExecutor executor;
+ private ListenerTaskScheduler scheduler;
+ private Runnable arrival;
+ private Runnable trigger;
+ private int incoming;
+ private int received;
+
+ @Setup
+ public void setup() {
+ executor = new QueuedExecutor();
+ scheduler = new ListenerTaskScheduler(executor, this::drain);
+ arrival = () -> incoming++;
+ Runnable reusedDrain = this::drain;
+ trigger = switch (scheduling) {
+ case "fresh" -> () -> executor.execute(() -> drain());
+ case "reused" -> () -> executor.execute(reusedDrain);
+ case "coalesced" -> scheduler::trigger;
+ default -> throw new IllegalArgumentException(scheduling);
+ };
+ }
+
+ @Benchmark
+ public int notifyAndDrain() {
+ received = 0;
+ for (int i = 0; i < notifications; i++) {
+ executor.execute(arrival);
+ trigger.run();
+ }
+ Runnable task;
+ while ((task = executor.tasks.poll()) != null) {
+ task.run();
+ }
+ if (received != notifications || incoming != 0) {
+ throw new IllegalStateException("A notification was lost");
+ }
+ return received;
+ }
+
+ private void drain() {
+ received += incoming;
+ incoming = 0;
+ }
+
+ private static final class QueuedExecutor implements Executor {
+ private final Queue<Runnable> tasks = new LinkedBlockingQueue<>();
+
+ @Override
+ public void execute(Runnable command) {
+ tasks.add(command);
+ }
+ }
+}
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 65cdbf07d5b5..e9053d33a894 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
@@ -86,6 +86,7 @@ public abstract class ConsumerBase<T> extends HandlerState implements Consumer<T
protected final MessageListenerExecutor messageListenerExecutor;
protected final ExecutorService externalPinnedExecutor;
protected final ExecutorService internalPinnedExecutor;
+ private final ListenerTaskScheduler listenerTaskScheduler;
protected final UnAckedMessageTracker unAckedMessageTracker;
final GrowableArrayBlockingQueue<Message<T>> incomingMessages;
protected Map<MessageIdAdv, MessageIdImpl[]> unAckedChunkedMessageIdSequenceMap = new ConcurrentHashMap<>();
@@ -147,6 +148,8 @@ public abstract class ConsumerBase<T> extends HandlerState implements Consumer<T
: conf.getMessageListenerExecutor();
this.externalPinnedExecutor = executorProvider.getExecutor();
this.internalPinnedExecutor = client.getInternalExecutorService();
+ this.listenerTaskScheduler = listener == null
+ ? null : new ListenerTaskScheduler(internalPinnedExecutor, this::drainListener);
this.pendingReceives = Queues.newConcurrentLinkedQueue();
this.pendingBatchReceives = Queues.newConcurrentLinkedQueue();
this.schema = schema;
@@ -1189,27 +1192,37 @@ public abstract class ConsumerBase<T> extends HandlerState implements Consumer<T
// The messages are added into the receiver queue by the internal pinned executor,
// so need to use internal pinned executor to avoid race condition which message
// might be added into the receiver queue but not able to read here.
- internalPinnedExecutor.execute(() -> {
- try {
- Message<T> msg;
- do {
- msg = internalReceive(0, TimeUnit.MILLISECONDS);
- if (msg != null) {
- // Trigger the notification on the message listener in a separate thread to avoid blocking the
- // internal pinned executor thread while the message processing happens
- final Message<T> finalMsg = msg;
- MESSAGE_LISTENER_QUEUE_SIZE_UPDATER.incrementAndGet(this);
+ listenerTaskScheduler.trigger();
+ }
+
+ private void drainListener() {
+ try {
+ Message<T> msg;
+ do {
+ msg = internalReceive(0, TimeUnit.MILLISECONDS);
+ if (msg != null) {
+ // Trigger the notification on the message listener in a separate thread to avoid blocking the
+ // internal pinned executor thread while the message processing happens
+ final Message<T> finalMsg = msg;
+ MESSAGE_LISTENER_QUEUE_SIZE_UPDATER.incrementAndGet(this);
+ try {
messageListenerExecutor.execute(msg, () -> callMessageListener(finalMsg));
- } else {
- if (log.isDebugEnabled()) {
- log.debug("[{}] [{}] Message has been cleared from the queue", topic, subscription);
- }
+ } catch (RuntimeException | Error error) {
+ // A failed submission has dequeued a message, but may leave more messages to drain.
+ // Preserve a later turn even when earlier failures consumed the coalesced notification.
+ // Retrying here makes progress; retrying a failed dequeue could loop indefinitely.
+ listenerTaskScheduler.trigger();
+ throw error;
}
- } while (msg != null);
- } catch (PulsarClientException e) {
- log.warn("[{}] [{}] Failed to dequeue the message for listener", topic, subscription, e);
- }
- });
+ } else {
+ if (log.isDebugEnabled()) {
+ log.debug("[{}] [{}] Message has been cleared from the queue", topic, subscription);
+ }
+ }
+ } while (msg != null);
+ } catch (PulsarClientException e) {
+ log.warn("[{}] [{}] Failed to dequeue the message for listener", topic, subscription, e);
+ }
}
private void executeMessageListener(Message<?> message, Runnable runnable) {
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ListenerTaskScheduler.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ListenerTaskScheduler.java
new file mode 100644
index 000000000000..821ae1693621
--- /dev/null
+++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ListenerTaskScheduler.java
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
+/**
+ * Coalesces listener-drain requests on an asynchronous, serial FIFO executor.
+ * An additional trigger requires a later turn: its arrival task may be queued behind the current drain.
+ * Arrival tasks must be enqueued before their corresponding triggers. Rejected submissions reset the state
+ * for a later notification to retry; delivery is not guaranteed while the executor rejects work.
+ */
+final class ListenerTaskScheduler implements Runnable {
+ private static final int IDLE = 0;
+ private static final int SCHEDULED = 1;
+ private static final int TRIGGERED_AGAIN = 2;
+ private static final AtomicIntegerFieldUpdater<ListenerTaskScheduler> STATE_UPDATER =
+ AtomicIntegerFieldUpdater.newUpdater(ListenerTaskScheduler.class, "state");
+
+ private final Executor executor;
+ private final Runnable drain;
+ private volatile int state;
+
+ ListenerTaskScheduler(Executor executor, Runnable drain) {
+ this.executor = executor;
+ this.drain = drain;
+ }
+
+ void trigger() {
+ int previous = STATE_UPDATER.getAndUpdate(this,
+ current -> current == IDLE ? SCHEDULED : TRIGGERED_AGAIN);
+ if (previous == IDLE) {
+ try {
+ executor.execute(this);
+ } catch (RejectedExecutionException error) {
+ // A later notification can retry if the executor becomes available again.
+ STATE_UPDATER.set(this, IDLE);
+ throw error;
+ }
+ }
+ }
+
+ @Override
+ public void run() {
+ try {
+ drain.run();
+ } finally {
+ if (STATE_UPDATER.getAndSet(this, IDLE) == TRIGGERED_AGAIN) {
+ // Enqueue after pending arrival tasks, rather than draining again in this turn.
+ // A concurrent trigger after the reset either schedules this turn or marks it for a later turn.
+ trigger();
+ }
+ }
+ }
+}
diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerListenerSchedulingTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerListenerSchedulingTest.java
new file mode 100644
index 000000000000..59004099bed4
--- /dev/null
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerListenerSchedulingTest.java
@@ -0,0 +1,172 @@
+/*
+ * 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.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.apache.pulsar.client.api.Message;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData;
+import org.apache.pulsar.client.util.ExecutorProvider;
+import org.testng.annotations.Test;
+
+public class ConsumerListenerSchedulingTest {
+ @Test
+ public void consecutiveRejectionsDoNotStrandCoalescedMessages() {
+ Fixture fixture = new Fixture(2);
+ Message<byte[]> first = newMessage();
+ Message<byte[]> second = newMessage();
+ Message<byte[]> third = newMessage();
+ fixture.addAndTrigger(first);
+ fixture.addAndTrigger(second);
+ fixture.addAndTrigger(third);
+
+ assertEquals(fixture.tasks.size(), 1);
+ expectThrows(RejectedExecutionException.class, fixture::runNext);
+ expectThrows(RejectedExecutionException.class, fixture::runNext);
+ fixture.runRemaining();
+
+ assertEquals(fixture.attempted, Arrays.asList(first, second, third));
+ assertEquals(fixture.accepted, Collections.singletonList(third));
+ assertTrue(fixture.consumer.incomingMessages.isEmpty());
+ }
+
+ @Test
+ public void rejectingEveryMessageFinishesAndAllowsLaterNotifications() {
+ Fixture fixture = new Fixture(3);
+ for (int i = 0; i < 3; i++) {
+ fixture.addAndTrigger(newMessage());
+ }
+ for (int i = 0; i < 3; i++) {
+ expectThrows(RejectedExecutionException.class, fixture::runNext);
+ }
+ fixture.runRemaining();
+ assertEquals(fixture.attempted.size(), 3);
+ assertTrue(fixture.accepted.isEmpty());
+ assertTrue(fixture.consumer.incomingMessages.isEmpty());
+
+ Message<byte[]> next = newMessage();
+ fixture.addAndTrigger(next);
+ fixture.runRemaining();
+ assertEquals(fixture.accepted, Collections.singletonList(next));
+ }
+
+ @Test
+ public void rejectionRetryRunsAfterPendingArrivals() {
+ Fixture fixture = new Fixture(2);
+ Message<byte[]> first = newMessage();
+ Message<byte[]> second = newMessage();
+ Message<byte[]> third = newMessage();
+ for (Message<byte[]> message : Arrays.asList(first, second, third)) {
+ fixture.tasks.add(() -> fixture.consumer.incomingMessages.add(message));
+ fixture.consumer.tryTriggerListener();
+ }
+
+ fixture.runNext(); // First arrival.
+ expectThrows(RejectedExecutionException.class, fixture::runNext);
+ assertEquals(fixture.attempted, Collections.singletonList(first));
+ assertTrue(fixture.consumer.incomingMessages.isEmpty());
+ fixture.runNext(); // Second arrival, ahead of the retry.
+ fixture.runNext(); // Third arrival, ahead of the retry.
+ assertEquals(fixture.consumer.incomingMessages.size(), 2);
+ expectThrows(RejectedExecutionException.class, fixture::runNext);
+ fixture.runRemaining();
+
+ assertEquals(fixture.attempted, Arrays.asList(first, second, third));
+ assertEquals(fixture.accepted, Collections.singletonList(third));
+ assertTrue(fixture.consumer.incomingMessages.isEmpty());
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Message<byte[]> newMessage() {
+ return mock(Message.class);
+ }
+
+ private static final class Fixture {
+ private final Queue<Runnable> tasks = new ArrayDeque<>();
+ private final List<Message<?>> attempted = new ArrayList<>();
+ private final List<Message<?>> accepted = new ArrayList<>();
+ private final ConsumerImpl<byte[]> consumer;
+
+ private Fixture(int rejectedSubmissions) {
+ ExecutorService executor = mock(ExecutorService.class);
+ doAnswer(invocation -> {
+ tasks.add(invocation.getArgument(0));
+ return null;
+ }).when(executor).execute(any(Runnable.class));
+ ExecutorProvider executorProvider = mock(ExecutorProvider.class);
+ when(executorProvider.getExecutor()).thenReturn(executor);
+ PulsarClientImpl client = ClientTestFixtures.createPulsarClientMock(executorProvider, executor);
+ client.getConfiguration().setStatsIntervalSeconds(0);
+ ConsumerConfigurationData<byte[]> conf = new ConsumerConfigurationData<>();
+ conf.setSubscriptionName("test-sub");
+ conf.setMessageListener((ignored, message) -> { });
+ conf.setMessageListenerExecutor((message, runnable) -> {
+ attempted.add(message);
+ if (attempted.size() <= rejectedSubmissions) {
+ throw new RejectedExecutionException("listener submission rejected");
+ }
+ accepted.add(message);
+ });
+ consumer = new ConsumerImpl<byte[]>(client, "non-persistent://tenant/ns/listener-scheduling", conf,
+ executorProvider, -1, false, false, new CompletableFuture<>(), null, 0, Schema.BYTES,
+ null, true) {
+ @Override
+ protected Message<byte[]> internalReceive(long timeout, TimeUnit unit) {
+ // Keep the real listener drain and scheduler, without message decoding or broker state.
+ return incomingMessages.poll();
+ }
+ };
+ consumer.setState(HandlerState.State.Ready);
+ assertTrue(tasks.isEmpty());
+ }
+
+ private void addAndTrigger(Message<byte[]> message) {
+ consumer.incomingMessages.add(message);
+ consumer.tryTriggerListener();
+ }
+
+ private void runNext() {
+ tasks.remove().run();
+ }
+
+ private void runRemaining() {
+ int runs = 0;
+ while (!tasks.isEmpty()) {
+ assertTrue(++runs <= 10, "Listener drain must not retry indefinitely");
+ runNext();
+ }
+ }
+ }
+}
diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ListenerTaskSchedulerTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ListenerTaskSchedulerTest.java
new file mode 100644
index 000000000000..8ef1950480e4
--- /dev/null
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ListenerTaskSchedulerTest.java
@@ -0,0 +1,245 @@
+/*
+ * 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.assertTrue;
+import static org.testng.Assert.expectThrows;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.BitSet;
+import java.util.Collections;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.testng.annotations.Test;
+
+public class ListenerTaskSchedulerTest {
+ @Test
+ public void drainsArrivalQueuedBehindAnAlreadyScheduledDrain() {
+ QueuedExecutor executor = new QueuedExecutor();
+ Queue<Integer> arrivals = new ArrayDeque<>();
+ List<Integer> received = new ArrayList<>();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, () -> {
+ runs.incrementAndGet();
+ Integer arrival;
+ while ((arrival = arrivals.poll()) != null) {
+ received.add(arrival);
+ }
+ });
+ executor.execute(() -> arrivals.add(1));
+ scheduler.trigger();
+ executor.execute(() -> arrivals.add(2));
+ scheduler.trigger();
+ assertEquals(executor.tasks.size(), 3);
+ executor.runNext();
+ executor.runNext();
+ assertEquals(received, Collections.singletonList(1));
+ executor.runAll();
+ assertEquals(received, Arrays.asList(1, 2));
+ assertEquals(runs.get(), 2);
+ }
+
+ @Test
+ public void coalescesABurstAndReturnsToIdle() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, runs::incrementAndGet);
+ for (int i = 0; i < 1000; i++) {
+ scheduler.trigger();
+ }
+ assertEquals(executor.tasks.size(), 1);
+ executor.runAll();
+ assertEquals(runs.get(), 2);
+ scheduler.trigger();
+ executor.runAll();
+ assertEquals(runs.get(), 3);
+ }
+
+ @Test
+ public void triggerDuringDrainSchedulesAnotherTurn() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger runs = new AtomicInteger();
+ AtomicReference<ListenerTaskScheduler> reference = new AtomicReference<>();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, () -> {
+ if (runs.incrementAndGet() == 1) {
+ reference.get().trigger();
+ }
+ });
+ reference.set(scheduler);
+ scheduler.trigger();
+ executor.runAll();
+ assertEquals(runs.get(), 2);
+ }
+
+ @Test
+ public void drainFailurePreservesLaterTriggers() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, () -> {
+ if (runs.incrementAndGet() == 1) {
+ throw new IllegalStateException("drain failed");
+ }
+ });
+ scheduler.trigger();
+ scheduler.trigger();
+ expectThrows(IllegalStateException.class, executor::runNext);
+ executor.runAll();
+ assertEquals(runs.get(), 2);
+ }
+
+ @Test
+ public void rejectedSubmissionAllowsLaterRetry() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger submissions = new AtomicInteger();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(task -> {
+ if (submissions.incrementAndGet() == 1) {
+ throw new RejectedExecutionException("first submission rejected");
+ }
+ executor.execute(task);
+ }, runs::incrementAndGet);
+ expectThrows(RejectedExecutionException.class, scheduler::trigger);
+ scheduler.trigger();
+ executor.runAll();
+ assertEquals(runs.get(), 1);
+ }
+
+ @Test
+ public void rejectedFollowUpAllowsLaterRetry() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger submissions = new AtomicInteger();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(task -> {
+ if (submissions.incrementAndGet() == 2) {
+ throw new RejectedExecutionException("follow-up submission rejected");
+ }
+ executor.execute(task);
+ }, runs::incrementAndGet);
+ scheduler.trigger();
+ scheduler.trigger();
+ expectThrows(RejectedExecutionException.class, executor::runNext);
+ assertEquals(runs.get(), 1);
+ assertTrue(executor.tasks.isEmpty());
+ scheduler.trigger();
+ executor.runAll();
+ assertEquals(runs.get(), 2);
+ }
+
+ @Test
+ public void drainFailureReturnsToIdle() {
+ QueuedExecutor executor = new QueuedExecutor();
+ AtomicInteger runs = new AtomicInteger();
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, () -> {
+ if (runs.incrementAndGet() == 1) {
+ throw new IllegalStateException("drain failed");
+ }
+ });
+ scheduler.trigger();
+ expectThrows(IllegalStateException.class, executor::runNext);
+ assertTrue(executor.tasks.isEmpty());
+ scheduler.trigger();
+ executor.runAll();
+ assertEquals(runs.get(), 2);
+ }
+
+ @Test(invocationCount = 5, timeOut = 30000)
+ public void concurrentNotificationsDoNotLoseArrivals() throws Exception {
+ int producerCount = 4;
+ int perProducer = 5000;
+ int total = producerCount * perProducer;
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ ExecutorService producers = Executors.newFixedThreadPool(producerCount);
+ Queue<Integer> arrivals = new ArrayDeque<>();
+ BitSet received = new BitSet(total);
+ AtomicInteger duplicates = new AtomicInteger();
+ CountDownLatch complete = new CountDownLatch(total);
+ CountDownLatch start = new CountDownLatch(1);
+ ListenerTaskScheduler scheduler = new ListenerTaskScheduler(executor, () -> {
+ Integer id;
+ while ((id = arrivals.poll()) != null) {
+ if (received.get(id)) {
+ duplicates.incrementAndGet();
+ }
+ received.set(id);
+ complete.countDown();
+ }
+ });
+ try {
+ List<Future<?>> jobs = new ArrayList<>();
+ for (int producer = 0; producer < producerCount; producer++) {
+ int first = producer * perProducer;
+ jobs.add(producers.submit(() -> {
+ start.await();
+ for (int i = 0; i < perProducer; i++) {
+ int id = first + i;
+ executor.execute(() -> arrivals.add(id));
+ scheduler.trigger();
+ }
+ return null;
+ }));
+ }
+ start.countDown();
+ for (Future<?> job : jobs) {
+ job.get(10, TimeUnit.SECONDS);
+ }
+ assertTrue(complete.await(10, TimeUnit.SECONDS), "An arrival was left without a drain");
+ executor.submit(() -> {
+ assertEquals(received.cardinality(), total);
+ assertEquals(duplicates.get(), 0);
+ assertTrue(arrivals.isEmpty());
+ }).get(10, TimeUnit.SECONDS);
+ } finally {
+ start.countDown();
+ producers.shutdownNow();
+ executor.shutdownNow();
+ assertTrue(producers.awaitTermination(5, TimeUnit.SECONDS));
+ assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ private static final class QueuedExecutor implements Executor {
+ private final Queue<Runnable> tasks = new ArrayDeque<>();
+
+ @Override
+ public void execute(Runnable command) {
+ tasks.add(command);
+ }
+
+ void runNext() {
+ tasks.remove().run();
+ }
+
+ void runAll() {
+ while (!tasks.isEmpty()) {
+ runNext();
+ }
+ }
+ }
+}