From 907abb6c63cc521e51909c9f6572dcc62c5e11c1 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Thu, 3 Sep 2026 11:33:07 +0800 Subject: [PATCH 1/8] [fix][client] Fix batch container buffer ownership on the failure-recovery path When compression or encryption releases the batch buffer before a flush that then fails, the old recovery reused the released buffer (use-after-free) or orphaned the compressed payload (leak). Track batch-buffer ownership explicitly instead of nulling the buffer inside the build method: - BatchMessageContainerImpl: add a batchPayloadOwned flag; a build only starts while the container owns its buffer. Compression/encryption that releases it clears the flag, and resetPayloadAfterFailedPublishing() reallocates instead of reusing released memory. Releasing the compressed payload when encryption fails fixes the orphaned-payload leak. - resetPayloadAfterFailedPublishing(): skip reallocation when the container has no messages left, since key-based batching forwards the reset to sub-batches that already succeeded and cleared. - RawBatchMessageContainerImpl: symmetrically release the compressed payload and any partially built encrypted buffer when encryption fails, and release the serialized payloads in toByteBuf() on failure instead of orphaning them. - StrategicTwoPhaseCompactor: always clear the batch container on a failed flush, including when the failure is an Error (discard(Exception) cannot take an Error). - Tests: cover failure recovery with/without compression, encryption-failure buffer release, and the fail-fast guard against re-entering a build without reset. Assisted-by: Claude Code Assisted-by: Codex --- .../impl/RawBatchMessageContainerImpl.java | 87 +++--- .../StrategicTwoPhaseCompactor.java | 5 +- .../RawBatchMessageContainerImplTest.java | 41 +++ .../protocol/ProducerBatchSendTest.java | 118 ++++---- .../impl/BatchMessageContainerImpl.java | 115 ++++++-- .../impl/BatchMessageContainerImplTest.java | 252 ++++++++++++++++++ 6 files changed, 514 insertions(+), 104 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 9b9f79a8ec5ea..3304a37bca1f0 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,9 @@ */ package org.apache.pulsar.client.impl; +import com.google.common.annotations.VisibleForTesting; import io.netty.buffer.ByteBuf; +import io.netty.util.ReferenceCountUtil; import java.nio.ByteBuffer; import java.util.Set; import org.apache.pulsar.client.api.CryptoKeyReader; @@ -60,22 +62,32 @@ 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: the caller owns recovery (StrategicTwoPhaseCompactor discards the + // container on any Throwable from toByteBuf()), so the batch is failed exactly once at the call site. + ReferenceCountUtil.safeRelease(encryptedPayload); + ReferenceCountUtil.safeRelease(compressedPayload); + throw t; } - encryptedPayload.writerIndex(targetBuffer.remaining()); - compressedPayload.release(); - return encryptedPayload; } @Override @@ -91,6 +103,11 @@ public void setCryptoKeyReader(CryptoKeyReader cryptoKeyReader) { this.cryptoKeyReader = cryptoKeyReader; } + @VisibleForTesting + void setMsgCryptoForTesting(MessageCrypto msgCrypto) { + this.msgCrypto = msgCrypto; + } + @Override public boolean add(MessageImpl msg, SendCallback callback) { this.lastAddedMessageId = (MessageIdAdv) msg.getMessageId(); @@ -168,29 +185,35 @@ public ByteBuf toByteBuf() { } ByteBuf encryptedPayload = encrypt(getCompressedBatchMetadataAndPayload(false)); - 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 0569c0f6850f2..2c59153751237 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 @@ -39,6 +39,7 @@ 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; @@ -469,7 +470,9 @@ private CompletableFuture flushBatchMessage(LedgerHandle lh, String top } catch (Throwable t) { log.error().exception(t).log("Failed to add entry"); - 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 f6e06261e87b7..da12e8c223dd2 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,6 +21,12 @@ 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.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; @@ -308,4 +314,39 @@ public void testToByteBufWithEncryptionWithInvalidEncryptKeys() { 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); + RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(); + 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 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); + + container.discard(null); + } } 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 f5ed0d3db788c..8a25c760728ad 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.atomic.AtomicBoolean; import lombok.CustomLog; import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +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.Schema; import org.apache.pulsar.client.impl.ConsumerImpl; import org.apache.pulsar.common.api.proto.BaseCommand; @@ -47,30 +50,42 @@ public class ProducerBatchSendTest extends SharedPulsarBaseTest { @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}, }; } - @Test(timeOut = 30_000, dataProvider = "flushSend") - public void testNoEnoughMemSend(List flushSend) throws Exception { + /** + * {@link org.apache.pulsar.client.impl.BatchMessageContainerImpl#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(dataProvider = "flushSend") + public void testNoEnoughMemSend(List flushSend, CompressionType compressionType) throws Exception { final String topic = newTopicName(); final String subscription = "s1"; admin.topics().createNonPartitionedTopic(topic); admin.topics().createSubscription(topic, subscription, MessageId.earliest); - Producer producer = pulsarClient.newProducer(Schema.STRING).topic(topic).enableBatching(true) - .batchingMaxMessages(Integer.MAX_VALUE).batchingMaxPublishDelay(1, TimeUnit.HOURS).create(); + ProducerBuilder 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 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); @@ -83,41 +98,46 @@ public void testNoEnoughMemSend(List flushSend) throws Exception { }).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 stay above the 4 KB compressMinMsgBodySize threshold, so the ZLIB cases + // really take the compression branch 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 consumer = (ConsumerImpl) pulsarClient.newConsumer(Schema.STRING).topic(topic) - .subscriptionName(subscription).subscribe(); - Message msg1 = consumer.receive(2, TimeUnit.SECONDS); - assertNotNull(msg1); - assertEquals(msg1.getValue(), "1"); - Message msg2 = consumer.receive(2, TimeUnit.SECONDS); - assertNotNull(msg2); - assertEquals(msg2.getValue(), "2"); - Message 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 consumer = (ConsumerImpl) pulsarClient.newConsumer(Schema.STRING) + .topic(topic).subscriptionName(subscription).subscribe(); + for (int i = 1; i <= 3; i++) { + Message 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(); + } } } 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 15072a0ad5fe4..2aade526199cc 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> messages = new ArrayList<>(maxMessagesNum); protected SendCallback previousCallback = null; // keep track of callbacks for individual messages being published in a batch @@ -114,6 +118,7 @@ public boolean add(MessageImpl msg, SendCallback callback) { 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(); @@ -171,6 +176,11 @@ protected ByteBuf getCompressedBatchMetadataAndPayload() { } protected ByteBuf getCompressedBatchMetadataAndPayload(boolean clientOperation) { + // 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(); @@ -180,7 +190,7 @@ protected ByteBuf getCompressedBatchMetadataAndPayload(boolean clientOperation) try { if (n == 1) { batchedMessageMetadataAndPayload.writeBytes(msg.getDataBuffer()); - } else { + } else { batchedMessageMetadataAndPayload = Commands.serializeSingleMessageInBatchWithPayload( msg.getMessageBuilder(), msg.getDataBuffer(), batchedMessageMetadataAndPayload); } @@ -197,10 +207,12 @@ protected ByteBuf getCompressedBatchMetadataAndPayload(boolean clientOperation) int uncompressedSize = batchedMessageMetadataAndPayload.readableBytes(); ByteBuf compressedPayload; - if (clientOperation && producer != null){ + if (clientOperation && producer != null) { if (compressionType != CompressionType.NONE && uncompressedSize > producer.conf.getCompressMinMsgBodySize()) { compressedPayload = producer.applyCompression(batchedMessageMetadataAndPayload); + // applyCompression released the batch buffer: the container no longer owns it. + batchPayloadOwned = false; messageMetadata.setCompression(compressionType); messageMetadata.setUncompressedSize(uncompressedSize); } else { @@ -209,6 +221,7 @@ protected ByteBuf getCompressedBatchMetadataAndPayload(boolean clientOperation) } else { compressedPayload = compressor.encode(batchedMessageMetadataAndPayload); batchedMessageMetadataAndPayload.release(); + batchPayloadOwned = false; if (compressionType != CompressionType.NONE) { messageMetadata.setCompression(compressionType); messageMetadata.setUncompressedSize(uncompressedSize); @@ -222,6 +235,41 @@ protected ByteBuf getCompressedBatchMetadataAndPayload(boolean clientOperation) 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(); + } + } + void updateMaxBatchSize(int uncompressedSize) { if (uncompressedSize > maxBatchSize) { maxBatchSize = uncompressedSize; @@ -255,6 +303,7 @@ public void clear() { minEntryBucketHash = Integer.MAX_VALUE; maxEntryBucketHash = Integer.MIN_VALUE; batchedMessageMetadataAndPayload = null; + batchPayloadOwned = false; currentTxnidMostBits = -1L; currentTxnidLeastBits = -1L; batchAllocatedSizeBytes = 0; @@ -272,10 +321,10 @@ public void discard(Exception ex) { if (firstCallback != null) { firstCallback.sendComplete(ex, null); } - if (batchedMessageMetadataAndPayload != null) { + if (batchPayloadOwned && batchedMessageMetadataAndPayload != null) { ReferenceCountUtil.safeRelease(batchedMessageMetadataAndPayload); - batchedMessageMetadataAndPayload = null; } + batchPayloadOwned = false; } catch (Throwable t) { log.warn().attr("topic", topicName) .attr("producerName", producer.getProducerName()) @@ -297,11 +346,16 @@ public OpSendMsg createOpSendMsg() throws IOException { messageMetadata.clear(); messageMetadata.copyFrom(messages.get(0).getMessageBuilder()); stampEntryBucketRange(); - 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), @@ -330,8 +384,7 @@ public OpSendMsg createOpSendMsg() throws IOException { lowestSequenceId = -1L; return op; } - ByteBuf encryptedPayload = producer.encryptMessage(messageMetadata, - getCompressedBatchMetadataAndPayload()); + ByteBuf encryptedPayload = buildAndEncryptBatchPayload(); updateAndReserveBatchAllocatedSize(encryptedPayload.capacity()); if (encryptedPayload.readableBytes() > getMaxMessageSize()) { encryptedPayload.release(); @@ -353,17 +406,23 @@ public OpSendMsg createOpSendMsg() throws IOException { messageMetadata.setTxnidLeastBits(currentTxnidLeastBits); } stampEntryBucketRange(); - ByteBufPair cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(), - messageMetadata.getHighestSequenceId(), numMessagesInBatch, messageMetadata, encryptedPayload); - log.debug(e -> e.attr("topic", topicName) - .attr("producerName", producer.getProducerName()) - .attr("seq", messageMetadata.getSequenceId()) - .attr("numMessagesInBatch", messageMetadata.getNumMessagesInBatch()) - .attr("highestSeq", messageMetadata.getHighestSequenceId()) - .attr("uncompressedsize", messageMetadata.getUncompressedSize()) - .attr("payloadsize", encryptedPayload.readableBytes()) - .log("Build batch message") - ); + ByteBufPair cmd; + try { + cmd = producer.sendMessage(producer.producerId, messageMetadata.getSequenceId(), + messageMetadata.getHighestSequenceId(), numMessagesInBatch, messageMetadata, encryptedPayload); + } catch (Throwable t) { + releasePayloadIfOrphaned(encryptedPayload); + throw t; + } + log.debug(e -> e.attr("topic", topicName) + .attr("producerName", producer.getProducerName()) + .attr("seq", messageMetadata.getSequenceId()) + .attr("numMessagesInBatch", messageMetadata.getNumMessagesInBatch()) + .attr("highestSeq", messageMetadata.getHighestSequenceId()) + .attr("uncompressedsize", messageMetadata.getUncompressedSize()) + .attr("payloadsize", encryptedPayload.readableBytes()) + .log("Build batch message") + ); OpSendMsg op = OpSendMsg.create(producer.rpcLatencyHistogram, messages, cmd, messageMetadata.getSequenceId(), messageMetadata.getHighestSequenceId(), firstCallback, batchAllocatedSizeBytes); @@ -376,7 +435,19 @@ public OpSendMsg createOpSendMsg() throws IOException { @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/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java index cc40cd27f8e96..97563f208f79f 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,25 +18,38 @@ */ 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.atLeastOnce; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; 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.util.ReferenceCountUtil; 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.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.apache.pulsar.client.api.CompressionType; import org.apache.pulsar.client.api.Schema; import org.apache.pulsar.client.impl.conf.ProducerConfigurationData; +import org.apache.pulsar.common.allocator.PulsarByteBufAllocator; +import org.apache.pulsar.common.api.proto.BaseCommand; import org.apache.pulsar.common.api.proto.MessageMetadata; +import org.apache.pulsar.common.protocol.ByteBufPair; +import org.apache.pulsar.common.protocol.Commands; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; public class BatchMessageContainerImplTest { @@ -230,4 +243,243 @@ private void addMessagesAndCreateOpSendMsg(BatchMessageContainerImpl batchMessag batchMessageContainer.clear(); messages.forEach(ReferenceCountUtil::safeRelease); } + + @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); + + AtomicReference compressedRef = new AtomicReference<>(); + doAnswer(invocation -> { + ByteBuf source = invocation.getArgument(0); + ByteBuf compressed = PulsarByteBufAllocator.DEFAULT.buffer(source.readableBytes()); + compressed.writeBytes(source); + source.release(); + compressedRef.set(compressed); + return compressed; + }).when(producer).applyCompression(any()); + AtomicBoolean fail = new AtomicBoolean(true); + when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> { + if (fail.get()) { + throw new RuntimeException("mocked encryption failure"); + } + return invocation.getArgument(1); + }); + when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(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); + }); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + List> 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"); + if (compressionType != CompressionType.NONE) { + // The container keeps its buffer reference while messages remain, and the compressed payload must not + // leak when encryption fails before anything took ownership of it. + assertNotNull(batchMessageContainer.batchedMessageMetadataAndPayload); + assertEquals(compressedRef.get().refCnt(), 0); + // Re-entering the build without reset must fail fast instead of writing into released memory. + 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. + verify(producer, atLeastOnce()).applyCompression(any()); + } + 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 firstEncryptedRef = new AtomicReference<>(); + when(producer.encryptMessage(any(), any())).thenAnswer(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.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(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); + }); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + List> 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) { + when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> { + if (failOnce.getAndSet(false)) { + throw new RuntimeException("mocked encryption failure"); + } + return invocation.getArgument(1); + }); + } else { + when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> invocation.getArgument(1)); + } + when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(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); + }); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + List> 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. + BaseCommand parsed = new BaseCommand(); + header.markReaderIndex(); + header.skipBytes(4); + 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); + } + + private ProducerImpl createTestProducer(CompressionType compressionType) throws Exception { + ProducerImpl producer = mock(ProducerImpl.class); + ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData(); + producerConfigurationData.setCompressionType(compressionType); + // Force the compression branch even for the tiny payloads used here, so the ZLIB case + // actually compresses instead of silently taking the below-threshold no-compression path. + producerConfigurationData.setCompressMinMsgBodySize(0); + PulsarClientImpl pulsarClient = mock(PulsarClientImpl.class); + when(pulsarClient.getMemoryLimitController()).thenReturn(mock(MemoryLimitController.class)); + try { + Field clientFiled = HandlerState.class.getDeclaredField("client"); + clientFiled.setAccessible(true); + clientFiled.set(producer, pulsarClient); + Field confFiled = ProducerBase.class.getDeclaredField("conf"); + confFiled.setAccessible(true); + confFiled.set(producer, producerConfigurationData); + } catch (Exception e) { + fail(e.getMessage()); + } + when(producer.getConfiguration()).thenReturn(producerConfigurationData); + // Mirror ProducerImpl.applyCompression semantics: encode into a new buffer and release + // the source, so the container's ownership-transfer logic is exercised the real way. + when(producer.applyCompression(any())).thenAnswer(invocation -> { + ByteBuf source = invocation.getArgument(0); + ByteBuf compressed = PulsarByteBufAllocator.DEFAULT.buffer(source.readableBytes()); + compressed.writeBytes(source); + source.release(); + return compressed; + }); + return producer; + } + + private List> addMessages(BatchMessageContainerImpl batchMessageContainer, int count) { + List> 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; + } } From ca49737de8f18129879f3ff63772daacc691fa5c Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Mon, 7 Sep 2026 16:03:27 +0800 Subject: [PATCH 2/8] [fix][client] Release built multi-batch operations when a later sub-batch fails Motivation createOpSendMsgs() built all sub-batches through a stream that collected into a list. When a later sub-batch failed to build, the stream aborted and the operations already built were unreachable: they never reached the send queue, so nothing released their commands. Every failed flush then leaked the command buffers (serialized header plus the transferred batch payload), and repeated retries under a persistent failure cause would keep growing the direct memory usage. Modifications - BatchMessageKeyBasedContainer and EntryBucketBatchContainer build their sub-batches in an explicit loop; on failure they release the commands of the operations already built before rethrowing. - The release is ownership-aware (releaseOrphanedOpCmd): the serialized header is always solely owned by the command, while the payload is only released when its ownership left the container (compression or encryption); otherwise the container keeps the buffer and the retry after resetPayloadAfterFailedPublishing() reuses it. - The messages stay in their sub-batches, so their semaphore permits and memory reservations settle when the retry completes. - Add a regression test for the partial multi-batch failure with and without compression (fails on the previous code with a leaked header). --- .../impl/BatchMessageContainerImpl.java | 14 ++++ .../impl/BatchMessageKeyBasedContainer.java | 56 ++++++++------- .../impl/EntryBucketBatchContainer.java | 39 +++++----- .../impl/BatchMessageContainerImplTest.java | 71 +++++++++++++++++++ 4 files changed, 137 insertions(+), 43 deletions(-) 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 2aade526199cc..0d0ee5985efc3 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 @@ -270,6 +270,20 @@ private void releasePayloadIfOrphaned(ByteBuf payload) { } } + /** + * 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 serialized header is always solely owned by the command; the payload is only released when its + * ownership left the container, otherwise the container keeps the buffer and + * {@link #resetPayloadAfterFailedPublishing()} reuses it when the messages are retried. + */ + void releaseOrphanedOpCmd(ProducerImpl.OpSendMsg op) { + op.cmd.getFirst().release(); + if (!batchPayloadOwned) { + op.cmd.getSecond().release(); + } + } + void updateMaxBatchSize(int uncompressedSize) { if (uncompressedSize > maxBatchSize) { maxBatchSize = uncompressedSize; 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 77e409670ccb8..fa80918f52372 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; @@ -95,36 +96,39 @@ public int getBatchAllocatedSizeBytes() { @Override public List 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 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 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/EntryBucketBatchContainer.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/EntryBucketBatchContainer.java index ce778cefb469c..10b56bbee1f2b 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/EntryBucketBatchContainer.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/EntryBucketBatchContainer.java @@ -20,6 +20,7 @@ import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -129,26 +130,30 @@ public int getBatchAllocatedSizeBytes() { @Override public List createOpSendMsgs() throws IOException { + // As in key-based batching: within a bucket the sequence ids need not be contiguous, so + // collapse to the highest sequence id and drop the highest_sequence_id field to allow the + // broker's weak-order check to pass. + batches.values().forEach(c -> c.setLowestSequenceId(c.getHighestSequenceId())); + List sorted = batches.values().stream() + .sorted((o1, o2) -> (int) (o1.getLowestSequenceId() - o2.getLowestSequenceId())) + .collect(Collectors.toList()); + // Build the buckets in order: when a later bucket fails to build, the operations already built + // never reach the send queue and must be released here, or their commands leak. + List ops = new ArrayList<>(sorted.size()); try { - // As in key-based batching: within a bucket the sequence ids need not be contiguous, so - // collapse to the highest sequence id and drop the highest_sequence_id field to allow the - // broker's weak-order check to pass. - batches.values().forEach(c -> c.setLowestSequenceId(c.getHighestSequenceId())); - return batches.values().stream() - .sorted((o1, o2) -> (int) (o1.getLowestSequenceId() - o2.getLowestSequenceId())) - .map(c -> { - try { - return c.createOpSendMsg(); - } catch (IOException e) { - throw new IllegalStateException(e); - } - }).collect(Collectors.toList()); - } catch (IllegalStateException e) { - if (e.getCause() instanceof IOException) { - throw (IOException) e.getCause(); + for (BatchMessageContainerImpl batchMessageContainer : sorted) { + ops.add(batchMessageContainer.createOpSendMsg()); } - throw e; + } 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); + } + } + throw t; } + return ops; } @Override 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 97563f208f79f..e909304de758a 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 @@ -40,6 +40,7 @@ import java.util.ArrayList; import java.util.List; 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.Schema; @@ -436,6 +437,76 @@ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean fai 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); + when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> invocation.getArgument(1)); + List builtPairs = new ArrayList<>(); + AtomicInteger sendCalls = new AtomicInteger(); + when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(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); + return pair; + }); + + BatchMessageKeyBasedContainer container = new BatchMessageKeyBasedContainer(); + container.setProducer(producer); + List> 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); + assertEquals(builtPairs.get(0).getFirst().refCnt(), 0); + if (compressionType != CompressionType.NONE) { + // Compression handed the payload over to the command, so it must be released as well. + assertEquals(builtPairs.get(0).getSecond().refCnt(), 0); + } else { + // Without compression the container still owns the payload buffer and reuses it on retry. + assertEquals(builtPairs.get(0).getSecond().refCnt(), 1); + } + + // All messages stay in their sub-batches and the retry produces a complete batch again. + assertEquals(container.getNumMessagesInBatch(), 4); + container.resetPayloadAfterFailedPublishing(); + List 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); + } + } + private ProducerImpl createTestProducer(CompressionType compressionType) throws Exception { ProducerImpl producer = mock(ProducerImpl.class); ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData(); From 36b7ab52305e068830773e718d3e2a37da2750bc Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Mon, 7 Sep 2026 17:08:07 +0800 Subject: [PATCH 3/8] [fix][client] Fix encrypted payload leaks on send failure paths Motivation The payload buffer produced by the encryption step must be released when a send fails midway - by whichever component last took its ownership. All three failure points along its lifecycle leaked it instead, so every failed send leaked direct memory, accumulating exactly under the sustained-failure conditions (e.g. direct-memory pressure) that trigger these failures in the first place: 1. During encryption: encryptMessage() leaked the partially built output buffer whenever the crypto failed - on the PulsarClientException rethrow, on the SEND (publish-unencrypted) fallback, and on any RuntimeException or Error escaping the crypto, which had no cleanup at all. The batch container path only released the source payload around this call (previous commit); the output buffer allocated inside was unreachable to it. 2. After encryption: serializeAndSendMessage() orphaned the encrypted payload when building the send command threw (e.g. a header allocation failure). For chunked messages the retained slice also leaked its claim on the shared base buffer, so the base never returned to the pool. 3. Deferred command: an operation whose command is built only after the schema registration completes (rePopulate) captured the encrypted payload in a closure; if the op was failed first (send timeout, producer close), the payload leaked with the unreachable closure. This completes the failure-path ownership fixes for the non-batch send path; the batch flush path is covered by the previous commits. Modifications - encryptMessage() releases the partially built output buffer on every exit that does not hand it to the caller; the source payload remains the caller's responsibility (unchanged contract). - New sendMessageOrReleasePayload() used at both command-building sites in serializeAndSendMessage(): a failed serialization releases the encrypted payload (and the chunk slice's claim on the base buffer) before rethrowing. - OpSendMsg tracks the deferred payload (pendingPayload); rePopulate() clears it once the payload moved into the command, and recycle() releases it when the op is failed before that. - Tests for all three failure points. --- .../pulsar/client/impl/ProducerImpl.java | 46 +++++++- .../pulsar/client/impl/ProducerImplTest.java | 106 ++++++++++++++++++ 2 files changed, 146 insertions(+), 6 deletions(-) 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 052055b329a62..03cd974982149 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 @@ -840,11 +840,14 @@ private void serializeAndSendMessage(MessageImpl msg, : 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 rePopulate() builds the command; 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()) { @@ -853,8 +856,11 @@ private void serializeAndSendMessage(MessageImpl msg, // 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); + // Clear the field before the call: a failed serialization releases the payload inside + // sendMessageOrReleasePayload, a successful one hands it to the command. + op.pendingPayload = null; + op.cmd = sendMessageOrReleasePayload(producerId, sequenceId, numMessages, messageId, + finalMsgMetadata, encryptedPayload); }; } op.setNumMessagesInBatch(numMessages); @@ -990,9 +996,12 @@ protected ByteBuf encryptMessage(MessageMetadata msgMetadata, ByteBuf compressed 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); ((MessageCrypto) msgCrypto).encrypt(conf.getEncryptionKeys(), conf.getCryptoKeyReader(), @@ -1000,7 +1009,9 @@ protected ByteBuf encryptMessage(MessageMetadata msgMetadata, ByteBuf compressed 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) { @@ -1010,6 +1021,26 @@ protected ByteBuf encryptMessage(MessageMetadata msgMetadata, ByteBuf compressed 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; } } @@ -1613,6 +1644,7 @@ protected static final class OpSendMsg { MessageImpl msg; List> msgs; ByteBufPair cmd; + ByteBuf pendingPayload; SendCallback callback; Runnable rePopulate; ChunkedMessageCtx chunkedMessageCtx; @@ -1633,6 +1665,7 @@ void initialize() { msg = null; msgs = null; cmd = null; + pendingPayload = null; callback = null; rePopulate = null; sequenceId = -1L; @@ -1759,6 +1792,7 @@ void sendComplete(final Exception e) { void recycle() { ReferenceCountUtil.safeRelease(chunkedMessageCtx); + ReferenceCountUtil.safeRelease(pendingPayload); initialize(); recyclerHandle.recycle(this); } 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 c35f8a77b515b..4c516a75de85c 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,6 +18,7 @@ */ package org.apache.pulsar.client.impl; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -25,12 +26,20 @@ import static org.mockito.Mockito.withSettings; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; import java.nio.ByteBuffer; import java.util.Collections; +import java.util.TreeSet; 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.conf.ProducerConfigurationData; import org.apache.pulsar.client.impl.metrics.LatencyHistogram; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.protocol.ByteBufPair; @@ -193,4 +202,101 @@ public void testProcessOpSendMsgInTerminalStateReleasesMemoryOnce() throws Excep "The memory reserved for the message must be released exactly once in state " + state); } } + + private ProducerConfigurationData encryptedProducerConf() { + ProducerConfigurationData conf = new ProducerConfigurationData(); + conf.setEncryptionKeys(new TreeSet<>(Collections.singleton("key"))); + conf.setCryptoKeyReader(mock(CryptoKeyReader.class)); + return conf; + } + + /** + * 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 { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + FieldUtils.writeField(producer, "conf", encryptedProducerConf(), true); + + ByteBuf partial = Unpooled.buffer(64); + Mockito.doReturn(partial).when(producer).allocateEncryptedBuffer(Mockito.anyInt()); + MessageCrypto msgCrypto = Mockito.mock(MessageCrypto.class); + Mockito.when(msgCrypto.getMaxOutputSize(Mockito.anyInt())).thenReturn(64); + Mockito.doThrow(new RuntimeException("mocked encryption failure")) + .when(msgCrypto).encrypt(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any()); + FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); + + 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 { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + // Mock instances skip field initializers: provide the logger the SEND fallback branch uses. + FieldUtils.writeField(producer, "log", + Mockito.mock(io.github.merlimat.slog.Logger.class, Mockito.RETURNS_DEEP_STUBS), true); + ProducerConfigurationData conf = encryptedProducerConf(); + conf.setCryptoFailureAction(ProducerCryptoFailureAction.SEND); + FieldUtils.writeField(producer, "conf", conf, true); + + ByteBuf partial = Unpooled.buffer(64); + Mockito.doReturn(partial).when(producer).allocateEncryptedBuffer(Mockito.anyInt()); + MessageCrypto msgCrypto = Mockito.mock(MessageCrypto.class); + Mockito.when(msgCrypto.getMaxOutputSize(Mockito.anyInt())).thenReturn(64); + Mockito.doThrow(new PulsarClientException("mocked encryption failure")) + .when(msgCrypto).encrypt(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.any()); + FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); + + 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 producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doThrow(new RuntimeException("mocked serialization failure")) + .when(producer) + .sendMessage(Mockito.anyLong(), Mockito.anyLong(), Mockito.anyInt(), Mockito.any(), + Mockito.any(), Mockito.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"); + } + + /** + * 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( + Mockito.mock(LatencyHistogram.class), + Mockito.mock(MessageImpl.class), + null, + 1L, + Mockito.mock(SendCallback.class)); + op.pendingPayload = payload; + + op.recycle(); + + assertEquals(payload.refCnt(), 0, "the deferred payload must be released when the op is recycled"); + } } From 669587320b09a0b03b46fba307774c6f3b0fecb5 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Tue, 8 Sep 2026 10:00:49 +0800 Subject: [PATCH 4/8] [fix][client] Release the ByteBufPair and recycle the op on orphan cleanup Motivation releaseOrphanedOpCmd() freed the pair's component buffers but left the ByteBufPair and the OpSendMsg unreleased: both are recycler-based, so on repeated multi-batch build failures they were garbage-collected instead of returned to their pools, adding heap allocation and GC pressure (noted in review). Modifications - When the payload claim is shared with the container (no compression or encryption), take it out of the pair with a retain() before releasing, so resetPayloadAfterFailedPublishing() keeps reusing the buffer. - Release the pair itself (returning it to its recycler) and recycle the orphaned op. - The regression test now also asserts the pair itself is released. --- .../client/impl/BatchMessageContainerImpl.java | 12 +++++++----- .../client/impl/BatchMessageContainerImplTest.java | 14 +++++++++++--- 2 files changed, 18 insertions(+), 8 deletions(-) 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 0d0ee5985efc3..466656b284fd7 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 @@ -273,15 +273,17 @@ private void releasePayloadIfOrphaned(ByteBuf payload) { /** * 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 serialized header is always solely owned by the command; the payload is only released when its - * ownership left the container, otherwise the container keeps the buffer and + * 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) { - op.cmd.getFirst().release(); - if (!batchPayloadOwned) { - op.cmd.getSecond().release(); + if (batchPayloadOwned) { + // Extract the container's payload claim from the pair before releasing it. + op.cmd.getSecond().retain(); } + op.cmd.release(); + op.recycle(); } void updateMaxBatchSize(int uncompressedSize) { 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 e909304de758a..10121be5222b3 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 @@ -449,6 +449,9 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType ProducerImpl producer = createTestProducer(compressionType); when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> invocation.getArgument(1)); List builtPairs = new ArrayList<>(); + // The pair clears its component references when released, so track the buffers at build time. + List builtHeaders = new ArrayList<>(); + List builtPayloads = new ArrayList<>(); AtomicInteger sendCalls = new AtomicInteger(); when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(invocation -> { if (sendCalls.incrementAndGet() == 2) { @@ -460,6 +463,8 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType header.writeInt(0); ByteBufPair pair = ByteBufPair.get(header, payload); builtPairs.add(pair); + builtHeaders.add(header); + builtPayloads.add(payload); return pair; }); @@ -484,13 +489,16 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType .isInstanceOf(RuntimeException.class) .hasMessageContaining("mocked second"); assertEquals(builtPairs.size(), 1); - assertEquals(builtPairs.get(0).getFirst().refCnt(), 0); + // 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); if (compressionType != CompressionType.NONE) { // Compression handed the payload over to the command, so it must be released as well. - assertEquals(builtPairs.get(0).getSecond().refCnt(), 0); + assertEquals(builtPayloads.get(0).refCnt(), 0); } else { // Without compression the container still owns the payload buffer and reuses it on retry. - assertEquals(builtPairs.get(0).getSecond().refCnt(), 1); + assertEquals(builtPayloads.get(0).refCnt(), 1); } // All messages stay in their sub-batches and the retry produces a complete batch again. From e1902e6c5162da0d3abc325dfdb99d04a3098dd1 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Tue, 8 Sep 2026 10:04:06 +0800 Subject: [PATCH 5/8] [fix][client] Release the source payload when compression or encryption fails Motivation Both applyCompression() and encryptMessage() leave their source payload with the caller on failure (they release it only after success), but the non-batch send path had no caller-side cleanup: when either stage threw, the exception propagated out of serializeAndSendMessage() and the source payload was orphaned. For chunked messages the retained slice's claim on the shared base buffer leaked as well, so the base never returned to the pool. Repeated failures therefore leaked direct memory (noted in review). The batch-container path already defends around these calls via its ownership tracking; the non-batch path did not. Modifications - New applyCompressionOrReleaseSource() and encryptMessageOrReleaseSource() used in serializeAndSendMessage(): on failure they release the source payload (and the chunk slice's claim) before rethrowing, mirroring sendMessageOrReleasePayload() on the serialization stage. - Tests for both stages. --- .../pulsar/client/impl/ProducerImpl.java | 36 +++++++++++++++++-- .../pulsar/client/impl/ProducerImplTest.java | 30 ++++++++++++++++ 2 files changed, 63 insertions(+), 3 deletions(-) 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 03cd974982149..ef51b0493db46 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 @@ -585,7 +585,7 @@ public void sendAsync(Message message, SendCallback callback) { // If a message has a delayed delivery time, we'll always send it individually if (!isBatchMessagingEnabled() || msgMetadata.hasDeliverAtTime()) { if (payload.readableBytes() > conf.getCompressMinMsgBodySize()) { - compressedPayload = applyCompression(payload); + compressedPayload = applyCompressionOrReleaseSource(payload); compressed = true; // validate msg-size (For batching this will be check at the batch completion size) @@ -829,9 +829,9 @@ private void serializeAndSendMessage(MessageImpl msg, // 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.readableBytes() > conf.getCompressMinMsgBodySize()) { - 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 @@ -1044,6 +1044,36 @@ ByteBufPair sendMessageOrReleasePayload(long producerId, long sequenceId, int nu } } + /** + * 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; + } + } + protected ByteBufPair sendMessage(long producerId, long sequenceId, int numMessages, MessageId messageId, MessageMetadata msgMetadata, ByteBuf compressedPayload) { 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 4c516a75de85c..c78e504d95088 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 @@ -280,6 +280,36 @@ public void testSendMessageFailureReleasesPayload() throws Exception { 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 producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doThrow(new RuntimeException("mocked compression failure")) + .when(producer) + .applyCompression(Mockito.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 producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doThrow(new PulsarClientException("mocked encryption failure")) + .when(producer) + .encryptMessage(Mockito.any(), Mockito.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. From 7f69f2f086ea34f08d44f4a68a5f2fa686195d37 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Tue, 8 Sep 2026 14:40:07 +0800 Subject: [PATCH 6/8] [fix][client] Keep the deferred payload alive when command construction fails Motivation A deferred command (schema registration pending) released its payload when the construction failed: rePopulate cleared pendingPayload before calling the releasing serialization wrapper. The op stays in pendingMessages after such a failure, and the next resend invokes the closure again - on the released buffer, producing an IllegalReferenceCountException on every reconnect and leaving the send stuck (noted in review). The release-on-throw semantics were wrong for this path: unlike the non-deferred send path, the op itself is the retry owner of the payload. Modifications - Extract the closure body into buildDeferredCommand(): the command is built from op.pendingPayload and the field is cleared only after the construction succeeds. A failed construction keeps the payload with the op for the next resend; recycle() still releases it when the op is failed instead. - Regression test: a failed first construction followed by recovery - the payload survives the failure, the retry rebuilds the command from the same buffer, and recycling the op afterwards does not double-release. --- .../pulsar/client/impl/ProducerImpl.java | 36 +++--- .../pulsar/client/impl/ProducerImplTest.java | 105 +++++++++++++----- 2 files changed, 99 insertions(+), 42 deletions(-) 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 ef51b0493db46..4b5dacbb2e2a3 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 @@ -845,23 +845,12 @@ private void serializeAndSendMessage(MessageImpl msg, op = OpSendMsg.create(rpcLatencyHistogram, msg, cmd, sequenceId, callback); } else { op = OpSendMsg.create(rpcLatencyHistogram, msg, null, sequenceId, callback); - // Hold on to the payload until rePopulate() builds the command; if the op is failed before + // 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); - } - // Clear the field before the call: a failed serialization releases the payload inside - // sendMessageOrReleasePayload, a successful one hands it to the command. - op.pendingPayload = null; - op.cmd = sendMessageOrReleasePayload(producerId, sequenceId, numMessages, messageId, - finalMsgMetadata, encryptedPayload); - }; + op.rePopulate = () -> buildDeferredCommand(op, finalMsgMetadata, producerId, sequenceId, + numMessages, messageId, chunkId); } op.setNumMessagesInBatch(numMessages); op.setBatchSizeByte(encryptedPayload.readableBytes()); @@ -1074,6 +1063,25 @@ ByteBuf encryptMessageOrReleaseSource(MessageMetadata msgMetadata, ByteBuf sourc } } + /** + * 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); + // The payload's ownership moved into the command. + op.pendingPayload = null; + } + protected ByteBufPair sendMessage(long producerId, long sequenceId, int numMessages, MessageId messageId, MessageMetadata msgMetadata, ByteBuf compressedPayload) { 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 c78e504d95088..c1b9c90dd1ceb 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 @@ -19,7 +19,14 @@ 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.RETURNS_DEEP_STUBS; +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; @@ -33,6 +40,7 @@ import java.nio.ByteBuffer; import java.util.Collections; import java.util.TreeSet; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.pulsar.client.api.CryptoKeyReader; import org.apache.pulsar.client.api.MessageCrypto; @@ -216,16 +224,15 @@ private ProducerConfigurationData encryptedProducerConf() { */ @Test public void testEncryptMessageReleasesPartialBufferOnFailure() throws Exception { - ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); FieldUtils.writeField(producer, "conf", encryptedProducerConf(), true); ByteBuf partial = Unpooled.buffer(64); - Mockito.doReturn(partial).when(producer).allocateEncryptedBuffer(Mockito.anyInt()); - MessageCrypto msgCrypto = Mockito.mock(MessageCrypto.class); - Mockito.when(msgCrypto.getMaxOutputSize(Mockito.anyInt())).thenReturn(64); - Mockito.doThrow(new RuntimeException("mocked encryption failure")) - .when(msgCrypto).encrypt(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), - Mockito.any()); + doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); + 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()); FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); ByteBuf source = Unpooled.buffer(8); @@ -240,21 +247,20 @@ public void testEncryptMessageReleasesPartialBufferOnFailure() throws Exception /** The SEND crypto-failure action returns the unencrypted source; the partial buffer must not leak. */ @Test public void testEncryptMessageCryptoFailureActionSendReleasesPartialBuffer() throws Exception { - ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); // Mock instances skip field initializers: provide the logger the SEND fallback branch uses. FieldUtils.writeField(producer, "log", - Mockito.mock(io.github.merlimat.slog.Logger.class, Mockito.RETURNS_DEEP_STUBS), true); + mock(io.github.merlimat.slog.Logger.class, RETURNS_DEEP_STUBS), true); ProducerConfigurationData conf = encryptedProducerConf(); conf.setCryptoFailureAction(ProducerCryptoFailureAction.SEND); FieldUtils.writeField(producer, "conf", conf, true); ByteBuf partial = Unpooled.buffer(64); - Mockito.doReturn(partial).when(producer).allocateEncryptedBuffer(Mockito.anyInt()); - MessageCrypto msgCrypto = Mockito.mock(MessageCrypto.class); - Mockito.when(msgCrypto.getMaxOutputSize(Mockito.anyInt())).thenReturn(64); - Mockito.doThrow(new PulsarClientException("mocked encryption failure")) - .when(msgCrypto).encrypt(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), - Mockito.any()); + doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); + 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()); FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); ByteBuf source = Unpooled.buffer(8); @@ -266,11 +272,10 @@ public void testEncryptMessageCryptoFailureActionSendReleasesPartialBuffer() thr /** A failed command serialization must release the payload instead of orphaning it. */ @Test public void testSendMessageFailureReleasesPayload() throws Exception { - ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); - Mockito.doThrow(new RuntimeException("mocked serialization failure")) + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new RuntimeException("mocked serialization failure")) .when(producer) - .sendMessage(Mockito.anyLong(), Mockito.anyLong(), Mockito.anyInt(), Mockito.any(), - Mockito.any(), Mockito.any()); + .sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any()); ByteBuf payload = Unpooled.buffer(8); assertThatThrownBy(() -> producer.sendMessageOrReleasePayload( @@ -283,10 +288,10 @@ public void testSendMessageFailureReleasesPayload() throws Exception { /** A failing compression stage must release the source payload the codec left with the caller. */ @Test public void testApplyCompressionFailureReleasesSource() throws Exception { - ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); - Mockito.doThrow(new RuntimeException("mocked compression failure")) + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new RuntimeException("mocked compression failure")) .when(producer) - .applyCompression(Mockito.any()); + .applyCompression(any()); ByteBuf source = Unpooled.buffer(8); assertThatThrownBy(() -> producer.applyCompressionOrReleaseSource(source)) @@ -298,10 +303,10 @@ public void testApplyCompressionFailureReleasesSource() throws Exception { /** A failing encryption stage must release the source payload encryptMessage() left with the caller. */ @Test public void testEncryptMessageFailureReleasesSource() throws Exception { - ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); - Mockito.doThrow(new PulsarClientException("mocked encryption failure")) + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new PulsarClientException("mocked encryption failure")) .when(producer) - .encryptMessage(Mockito.any(), Mockito.any()); + .encryptMessage(any(), any()); ByteBuf source = Unpooled.buffer(8); assertThatThrownBy(() -> producer.encryptMessageOrReleaseSource(new MessageMetadata(), source)) @@ -318,15 +323,59 @@ public void testEncryptMessageFailureReleasesSource() throws Exception { public void testPendingSchemaOpPayloadReleasedOnFailure() { ByteBuf payload = Unpooled.buffer(8); ProducerImpl.OpSendMsg op = ProducerImpl.OpSendMsg.create( - Mockito.mock(LatencyHistogram.class), - Mockito.mock(MessageImpl.class), + mock(LatencyHistogram.class), + mock(MessageImpl.class), null, 1L, - Mockito.mock(SendCallback.class)); + 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 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(); + } } From 5031cf2804ffabdf533d46dc9504bdce5dc38fb7 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Thu, 10 Sep 2026 11:10:20 +0800 Subject: [PATCH 7/8] [fix][client] Address review: exercise send-path wiring and drop test reflection Motivation Review follow-up on the test coverage: - The failure regression tests called the release helpers directly, so a call site reverting to the bare method would not have failed any test. - The RawBatch encryption-failure test asserted only the source buffer: removing the partially-built-encrypted-output release still passed. - The SEND-frame parse test skipped 4 bytes too few, so CMD_SIZE was parsed as part of the command. - New tests reached private state through reflection, which CODING.md forbids in favor of @VisibleForTesting accessors or normal construction. Modifications - testSendPathFailureReleasesPayloadThroughTheStageHelpers drives the real sendAsync(): a failing compression stage must release the message payload, and a failing command serialization must release the compressed payload handed to it. Both call-site reversions (bare applyCompression / bare sendMessage) make it fail; verified by running them. - RawBatchMessageContainerImpl gains a @VisibleForTesting constructor taking the allocator; the encryption-failure test tracks every buffer the container allocates and asserts all are freed. Removing the encrypted-output release in the Throwable catch fails it; verified. - The SEND-frame parse skips both length fields (8 bytes) before parseFrom, so the command is parsed within its declared bounds. - The encryption and container tests build the producer through its real constructor (useConstructor + CALLS_REAL_METHODS): conf, log, client and msgCrypto hold real values, and the crypto mock is installed via conf.setMessageCrypto() through the same constructor branch production uses. No new reflection; producer stubs use the doAnswer form so registering a stub does not execute the real method. --- .../impl/RawBatchMessageContainerImpl.java | 9 ++ .../RawBatchMessageContainerImplTest.java | 19 ++- .../impl/BatchMessageContainerImplTest.java | 77 ++++++----- .../pulsar/client/impl/ProducerImplTest.java | 126 +++++++++++++++--- 4 files changed, 181 insertions(+), 50 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 3304a37bca1f0..f2104f05bcb9f 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 @@ -20,6 +20,7 @@ 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; @@ -58,6 +59,14 @@ public RawBatchMessageContainerImpl() { 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; 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 da12e8c223dd2..a902a32fd005a 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 @@ -24,17 +24,21 @@ 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; @@ -322,7 +326,16 @@ public void testToByteBufWithEncryptionWithInvalidEncryptKeys() { @Test public void testToByteBufReleasesPayloadWhenEncryptionFailsUnexpectedly() throws Exception { setEncryptionAndCompression(true, false); - RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(); + // 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 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); @@ -346,6 +359,10 @@ public void testToByteBufReleasesPayloadWhenEncryptionFailsUnexpectedly() throws // 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); } 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 10121be5222b3..b0ad454a350dd 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 @@ -22,11 +22,13 @@ 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.atLeastOnce; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -34,17 +36,23 @@ import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufAllocator; 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.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.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; @@ -271,19 +279,19 @@ public void testRecoveryAfterBatchBuildFailure(CompressionType compressionType) return compressed; }).when(producer).applyCompression(any()); AtomicBoolean fail = new AtomicBoolean(true); - when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> { + doAnswer(invocation -> { if (fail.get()) { throw new RuntimeException("mocked encryption failure"); } return invocation.getArgument(1); - }); - when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(invocation -> { + }).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); List> messages = addMessages(batchMessageContainer, 2); @@ -328,7 +336,7 @@ public void testRecoveryAfterEncryptionFailure() throws Exception { AtomicBoolean failSend = new AtomicBoolean(true); AtomicReference firstEncryptedRef = new AtomicReference<>(); - when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> { + 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()); @@ -338,8 +346,8 @@ public void testRecoveryAfterEncryptionFailure() throws Exception { firstEncryptedRef.set(encrypted); } return encrypted; - }); - when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(invocation -> { + }).when(producer).encryptMessage(any(), any()); + doAnswer(invocation -> { if (failSend.getAndSet(false)) { throw new RuntimeException("mocked send failure"); } @@ -348,7 +356,7 @@ public void testRecoveryAfterEncryptionFailure() throws Exception { 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> messages = addMessages(batchMessageContainer, 2); @@ -381,16 +389,16 @@ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean fai AtomicBoolean failOnce = new AtomicBoolean(true); if (failAtEncrypt) { - when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> { + doAnswer(invocation -> { if (failOnce.getAndSet(false)) { throw new RuntimeException("mocked encryption failure"); } return invocation.getArgument(1); - }); + }).when(producer).encryptMessage(any(), any()); } else { - when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> invocation.getArgument(1)); + doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any()); } - when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(invocation -> { + doAnswer(invocation -> { if (failAtSend && failOnce.getAndSet(false)) { throw new RuntimeException("mocked send failure"); } @@ -398,7 +406,7 @@ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean fai 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> messages = addMessages(batchMessageContainer, 3); @@ -420,10 +428,11 @@ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean fai // 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. + // 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(4); + header.skipBytes(8); parsed.parseFrom(header, cmdSize); assertEquals(parsed.getType(), BaseCommand.Type.SEND); header.resetReaderIndex(); @@ -447,13 +456,13 @@ private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean fai public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType compressionType) throws Exception { ProducerImpl producer = createTestProducer(compressionType); - when(producer.encryptMessage(any(), any())).thenAnswer(invocation -> invocation.getArgument(1)); + doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any()); List builtPairs = new ArrayList<>(); // The pair clears its component references when released, so track the buffers at build time. List builtHeaders = new ArrayList<>(); List builtPayloads = new ArrayList<>(); AtomicInteger sendCalls = new AtomicInteger(); - when(producer.sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any())).thenAnswer(invocation -> { + doAnswer(invocation -> { if (sendCalls.incrementAndGet() == 2) { throw new RuntimeException("mocked second sub-batch failure"); } @@ -466,7 +475,7 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType builtHeaders.add(header); builtPayloads.add(payload); return pair; - }); + }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any()); BatchMessageKeyBasedContainer container = new BatchMessageKeyBasedContainer(); container.setProducer(producer); @@ -516,34 +525,38 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType } private ProducerImpl createTestProducer(CompressionType compressionType) throws Exception { - ProducerImpl producer = mock(ProducerImpl.class); ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData(); producerConfigurationData.setCompressionType(compressionType); // Force the compression branch even for the tiny payloads used here, so the ZLIB case // actually compresses instead of silently taking the below-threshold no-compression path. producerConfigurationData.setCompressMinMsgBodySize(0); PulsarClientImpl pulsarClient = mock(PulsarClientImpl.class); + when(pulsarClient.newProducerId()).thenReturn(1L); + when(pulsarClient.getCnxPool()).thenReturn(mock(ConnectionPool.class)); when(pulsarClient.getMemoryLimitController()).thenReturn(mock(MemoryLimitController.class)); - try { - Field clientFiled = HandlerState.class.getDeclaredField("client"); - clientFiled.setAccessible(true); - clientFiled.set(producer, pulsarClient); - Field confFiled = ProducerBase.class.getDeclaredField("conf"); - confFiled.setAccessible(true); - confFiled.set(producer, producerConfigurationData); - } catch (Exception e) { - fail(e.getMessage()); - } - when(producer.getConfiguration()).thenReturn(producerConfigurationData); + 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)); // Mirror ProducerImpl.applyCompression semantics: encode into a new buffer and release // the source, so the container's ownership-transfer logic is exercised the real way. - when(producer.applyCompression(any())).thenAnswer(invocation -> { + // doAnswer-form stubbing: when-form would execute the real method while registering. + doAnswer(invocation -> { ByteBuf source = invocation.getArgument(0); ByteBuf compressed = PulsarByteBufAllocator.DEFAULT.buffer(source.readableBytes()); compressed.writeBytes(source); source.release(); return compressed; - }); + }).when(producer).applyCompression(any()); return producer; } 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 c1b9c90dd1ceb..3dc8724e5f832 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 @@ -20,7 +20,6 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.CALLS_REAL_METHODS; -import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.anyLong; @@ -28,6 +27,7 @@ import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; @@ -37,9 +37,16 @@ import static org.testng.Assert.assertTrue; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; +import io.netty.channel.EventLoopGroup; +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.Collections; +import java.util.Optional; import java.util.TreeSet; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.pulsar.client.api.CryptoKeyReader; @@ -47,7 +54,9 @@ 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.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; @@ -218,22 +227,55 @@ private ProducerConfigurationData encryptedProducerConf() { 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 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 { - ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); - FieldUtils.writeField(producer, "conf", encryptedProducerConf(), true); - - ByteBuf partial = Unpooled.buffer(64); - doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); 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()); - FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); + ProducerConfigurationData conf = encryptedProducerConf(); + conf.setMessageCrypto(msgCrypto); + ProducerImpl 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)) @@ -247,21 +289,17 @@ public void testEncryptMessageReleasesPartialBufferOnFailure() throws Exception /** The SEND crypto-failure action returns the unencrypted source; the partial buffer must not leak. */ @Test public void testEncryptMessageCryptoFailureActionSendReleasesPartialBuffer() throws Exception { - ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); - // Mock instances skip field initializers: provide the logger the SEND fallback branch uses. - FieldUtils.writeField(producer, "log", - mock(io.github.merlimat.slog.Logger.class, RETURNS_DEEP_STUBS), true); + 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); - FieldUtils.writeField(producer, "conf", conf, true); + ProducerImpl producer = constructProducer(mockedPulsarClient(), conf); ByteBuf partial = Unpooled.buffer(64); doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); - 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()); - FieldUtils.writeField(producer, "msgCrypto", msgCrypto, true); ByteBuf source = Unpooled.buffer(8); assertSame(producer.encryptMessage(new MessageMetadata(), source), source); @@ -378,4 +416,58 @@ public void testDeferredCommandConstructionFailureThenRecovery() throws Exceptio 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); + conf.setCompressMinMsgBodySize(0); + PulsarClientImpl client = mockedPulsarClient(); + when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024)); + ProducerImpl producer = constructProducer(client, conf); + producer.setState(ProducerImpl.State.Ready); + + // A failing compression stage must release the message payload. sendAsync() runs this stage on the + // caller thread and lets the runtime error propagate; the wiring under test is the buffer release. + doThrow(new RuntimeException("mocked compression failure")) + .when(producer).applyCompression(any()); + MessageImpl first = newMessage("first"); + SendCallback firstCallback = mock(SendCallback.class); + assertThatThrownBy(() -> producer.sendAsync(first, firstCallback)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked compression failure"); + verify(firstCallback, never()).sendComplete(any(), any()); + 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 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"); + } + + private static MessageImpl 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); + } } From e47b27f8cc82882fa823e31262dd152207079623 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Thu, 10 Sep 2026 12:41:15 +0800 Subject: [PATCH 8/8] [fix][client] Release the chunked payload base claim when a chunk build fails Motivation For chunked persistent messages, the non-last chunks retain their slices while the base payload's own ref-count claim is carried by the last chunk (its slice is never retained). When serialization failed on a non-last chunk, the send-path helper released that chunk's retained-slice claim, but the base's own claim was orphaned: every failed chunked send leaked the shared payload buffer. Modifications - The chunk loop in sendAsync() releases the base payload's own claim when a non-last chunk's serialization fails. The last chunk is excluded because its unretained slice already is the base claim, which the send-path helper releases; the persistence condition mirrors the slicing condition in serializeAndSendMessage(). - Wiring-level regression tests: chunked failures with and without compression assert the base is fully released (reverting the release fails them, verified); a deferred-schema test drives the real sendAsync() into the deferred branch and asserts the pendingPayload assignment, the failed-first-build retention, the resend rebuild, and the recycle release; the RawBatch PulsarClientException branch and the entry-bucket partial-build loop get the same tracked-buffer coverage. --- .../RawBatchMessageContainerImplTest.java | 40 +++++ .../pulsar/client/impl/ProducerImpl.java | 18 +- .../impl/BatchMessageContainerImplTest.java | 77 +++++++++ .../pulsar/client/impl/ProducerImplTest.java | 159 ++++++++++++++++++ 4 files changed, 291 insertions(+), 3 deletions(-) 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 a902a32fd005a..2fe8b8406693d 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 @@ -44,6 +44,7 @@ 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; @@ -366,4 +367,43 @@ public void testToByteBufReleasesPayloadWhenEncryptionFailsUnexpectedly() throws 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 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 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-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 4b5dacbb2e2a3..be27e2f94afe1 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 @@ -699,9 +699,21 @@ public void sendAsync(Message message, SendCallback callback) { 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); + } + throw t; + } readStartIndex = ((chunkId + 1) * payloadChunkSize); } } 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 b0ad454a350dd..9866e191046bf 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 @@ -42,6 +42,7 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Optional; import java.util.concurrent.CompletableFuture; @@ -524,6 +525,82 @@ public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType } } + /** + * The entry-bucket container builds its buckets through the same loop as the key-based container, so a later + * bucket failing to build must release the operations already built there as well: a regression in this loop + * would otherwise leak command buffers on every failed flush. + */ + @Test(dataProvider = "compressionTypes") + public void testEntryBucketPartialBuildFailureReleasesBuiltOps(CompressionType compressionType) + throws Exception { + ProducerImpl producer = createTestProducer(compressionType); + doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any()); + List builtPairs = new ArrayList<>(); + List builtHeaders = new ArrayList<>(); + List builtPayloads = new ArrayList<>(); + AtomicInteger sendCalls = new AtomicInteger(); + doAnswer(invocation -> { + if (sendCalls.incrementAndGet() == 2) { + throw new RuntimeException("mocked second bucket 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()); + + // "key-1" and "key-2" hash into different entry buckets with these splits. + EntryBucketBatchContainer container = + new EntryBucketBatchContainer(Arrays.asList(0x4000, 0x8000, 0xC000)); + container.setProducer(producer); + List> 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 ? "key-1" : "key-2"); + 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); + } + + // Bucket "key-1" builds its operation, bucket "key-2" fails: the built command must not leak. + assertThatThrownBy(container::createOpSendMsgs) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked second"); + assertEquals(builtPairs.size(), 1); + assertEquals(builtPairs.get(0).refCnt(), 0); + assertEquals(builtHeaders.get(0).refCnt(), 0); + if (compressionType != CompressionType.NONE) { + // Compression handed the payload over to the command, so it must be released as well. + assertEquals(builtPayloads.get(0).refCnt(), 0); + } else { + // Without compression the container still owns the payload buffer and reuses it on retry. + assertEquals(builtPayloads.get(0).refCnt(), 1); + } + + // All messages stay in their buckets and the retry produces a complete batch again. + assertEquals(container.getNumMessagesInBatch(), 4); + container.resetPayloadAfterFailedPublishing(); + List 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); + } + } + private ProducerImpl createTestProducer(CompressionType compressionType) throws Exception { ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData(); producerConfigurationData.setCompressionType(compressionType); 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 3dc8724e5f832..a49d8e591a583 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 @@ -32,6 +32,7 @@ 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; @@ -48,6 +49,7 @@ import java.util.TreeSet; import java.util.concurrent.CompletableFuture; 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; @@ -60,6 +62,7 @@ import org.apache.pulsar.client.impl.metrics.LatencyHistogram; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.protocol.ByteBufPair; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -464,10 +467,166 @@ public void testSendPathFailureReleasesPayloadThroughTheStageHelpers() throws Ex "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); + conf.setCompressMinMsgBodySize(0); + PulsarClientImpl client = mockedPulsarClient(); + when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024)); + ProducerImpl 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 message = newMessage(new byte[10 * 1024]); + ByteBuf payload = message.getDataBuffer(); + SendCallback callback = mock(SendCallback.class); + producer.sendAsync(message, callback); + + ArgumentCaptor 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); + conf.setCompressMinMsgBodySize(0); + PulsarClientImpl client = mockedPulsarClient(); + when(client.getMemoryLimitController()).thenReturn(new MemoryLimitController(1024 * 1024)); + ProducerImpl 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 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 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 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 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 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 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 newMessage(byte[] content) { + MessageMetadata metadata = new MessageMetadata(); + metadata.setPublishTime(System.currentTimeMillis()); + return MessageImpl.create(metadata, ByteBuffer.wrap(content), Schema.BYTES, null); + } }