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..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 @@ -18,7 +18,10 @@ */ package org.apache.pulsar.client.impl; +import com.google.common.annotations.VisibleForTesting; import io.netty.buffer.ByteBuf; +import io.netty.buffer.ByteBufAllocator; +import io.netty.util.ReferenceCountUtil; import java.nio.ByteBuffer; import java.util.Set; import org.apache.pulsar.client.api.CryptoKeyReader; @@ -56,26 +59,44 @@ 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; } - 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 +112,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 +194,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..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 @@ -21,19 +21,30 @@ import static org.apache.pulsar.common.api.proto.CompressionType.NONE; import static org.apache.pulsar.common.api.proto.CompressionType.ZSTD; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.AssertJUnit.assertFalse; import static org.testng.AssertJUnit.assertTrue; import io.netty.buffer.ByteBuf; +import io.netty.buffer.ByteBufAllocator; import io.netty.buffer.Unpooled; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.charset.Charset; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Optional; import org.apache.pulsar.client.api.CryptoKeyReader; import org.apache.pulsar.client.api.EncryptionKeyInfo; import org.apache.pulsar.client.api.MessageCrypto; +import org.apache.pulsar.client.api.PulsarClientException; import org.apache.pulsar.client.api.Schema; import org.apache.pulsar.client.impl.crypto.MessageCryptoBc; import org.apache.pulsar.common.api.EncryptionContext; @@ -308,4 +319,91 @@ 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); + // 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); + + // 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); + // The partially built encrypted output buffer must have been released as well, not only the source. + for (ByteBuf buffer : allocated) { + Assert.assertEquals(buffer.refCnt(), 0); + } + + container.discard(null); + } + + /** + * A crypto provider failing with a {@link PulsarClientException} must release the compressed batch payload + * and the partially built encrypted buffer and discard the batch, so the compactor can reuse the container. + */ + @Test + public void testToByteBufReleasesPayloadAndDiscardsWhenEncryptionFailsWithClientException() throws Exception { + setEncryptionAndCompression(true, false); + List 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-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..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 @@ -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,57 @@ 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(); + } + } + + /** + * Releases the command of an operation that was built from this container but never reached the send queue: + * in multi-batch mode a later sub-batch can fail to build after this one already produced its operation. + * The payload claim is shared with the container when its ownership never left (no compression or + * encryption): take it out of the pair before releasing, so the container keeps the buffer and + * {@link #resetPayloadAfterFailedPublishing()} reuses it when the messages are retried. + */ + void releaseOrphanedOpCmd(ProducerImpl.OpSendMsg op) { + if (batchPayloadOwned) { + // 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) { if (uncompressedSize > maxBatchSize) { maxBatchSize = uncompressedSize; @@ -255,6 +319,7 @@ public void clear() { minEntryBucketHash = Integer.MAX_VALUE; maxEntryBucketHash = Integer.MIN_VALUE; batchedMessageMetadataAndPayload = null; + batchPayloadOwned = false; currentTxnidMostBits = -1L; currentTxnidLeastBits = -1L; batchAllocatedSizeBytes = 0; @@ -272,10 +337,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 +362,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 +400,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 +422,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 +451,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/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/main/java/org/apache/pulsar/client/impl/ProducerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java index 052055b329a62..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 @@ -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) @@ -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); } } @@ -829,9 +841,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 @@ -840,22 +852,17 @@ 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 the deferred command is built; if the op is failed before + // that happens, recycle() releases it instead of orphaning the buffer. + op.pendingPayload = encryptedPayload; final MessageMetadata finalMsgMetadata = msgMetadata; - op.rePopulate = () -> { - if (msgMetadata.hasChunkId()) { - // The message metadata is shared between all chunks in a large message - // We need to reset the chunk id for each call of this method - // It's safe to do that because there is only 1 thread to manipulate this message metadata - finalMsgMetadata.setChunkId(chunkId); - } - op.cmd = sendMessage(producerId, sequenceId, numMessages, messageId, finalMsgMetadata, - encryptedPayload); - }; + op.rePopulate = () -> buildDeferredCommand(op, finalMsgMetadata, producerId, sequenceId, + numMessages, messageId, chunkId); } op.setNumMessagesInBatch(numMessages); op.setBatchSizeByte(encryptedPayload.readableBytes()); @@ -990,9 +997,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 +1010,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,7 +1022,76 @@ 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; + } + } + + /** + * Applies compression and, when the codec fails, releases the source payload instead of orphaning it: + * applyCompression() releases its input only after the codec succeeds. For chunked messages this also + * returns the slice's claim on the shared payload buffer. + */ + ByteBuf applyCompressionOrReleaseSource(ByteBuf source) { + try { + return applyCompression(source); + } catch (Throwable t) { + ReferenceCountUtil.safeRelease(source); + throw t; + } + } + + /** + * Applies encryption and, when the crypto fails, releases the source payload instead of orphaning it: + * encryptMessage() leaves the source with the caller on failure (its internal partial output buffer is + * released inside). For chunked messages this also returns the slice's claim on the shared payload + * buffer. + */ + ByteBuf encryptMessageOrReleaseSource(MessageMetadata msgMetadata, ByteBuf source) + throws PulsarClientException { + try { + return encryptMessage(msgMetadata, source); + } catch (Throwable t) { + ReferenceCountUtil.safeRelease(source); + throw t; + } + } + + /** + * Builds the command of an op whose serialization was deferred until the schema registration completed. + * The deferred payload stays owned by the op ({@code pendingPayload}) until the command is built: a + * failed construction must not release it, because the op remains pending and the next resend rebuilds + * from the same buffer; {@code recycle()} releases it when the op is failed instead. + */ + void buildDeferredCommand(OpSendMsg op, MessageMetadata msgMetadata, long producerId, long sequenceId, + int numMessages, MessageId messageId, int chunkId) { + if (msgMetadata.hasChunkId()) { + // The message metadata is shared between all chunks in a large message. We need to reset the + // chunk id for each call of this method. It's safe to do that because there is only 1 thread + // to manipulate this message metadata. + msgMetadata.setChunkId(chunkId); } + op.cmd = sendMessage(producerId, sequenceId, numMessages, messageId, msgMetadata, op.pendingPayload); + // The payload's ownership moved into the command. + op.pendingPayload = null; } protected ByteBufPair sendMessage(long producerId, long sequenceId, int numMessages, @@ -1613,6 +1694,7 @@ protected static final class OpSendMsg { MessageImpl msg; List> msgs; ByteBufPair cmd; + ByteBuf pendingPayload; SendCallback callback; Runnable rePopulate; ChunkedMessageCtx chunkedMessageCtx; @@ -1633,6 +1715,7 @@ void initialize() { msg = null; msgs = null; cmd = null; + pendingPayload = null; callback = null; rePopulate = null; sequenceId = -1L; @@ -1759,6 +1842,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/BatchMessageContainerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java index cc40cd27f8e96..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 @@ -18,25 +18,48 @@ */ package org.apache.pulsar.client.impl; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.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; import static org.testng.Assert.fail; +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.Arrays; +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; +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 +253,402 @@ 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); + doAnswer(invocation -> { + if (fail.get()) { + throw new RuntimeException("mocked encryption failure"); + } + return invocation.getArgument(1); + }).when(producer).encryptMessage(any(), any()); + doAnswer(invocation -> { + ByteBuf payload = invocation.getArgument(5); + ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer(); + header.writeInt(4 + 4 + payload.readableBytes()); + header.writeInt(0); + return ByteBufPair.get(header, payload); + }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any()); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + 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<>(); + doAnswer(invocation -> { + // Real encryption allocates a new buffer and releases the source payload. + ByteBuf source = invocation.getArgument(1); + ByteBuf encrypted = PulsarByteBufAllocator.DEFAULT.buffer(source.readableBytes()); + encrypted.writeBytes(source); + source.release(); + if (failSend.get()) { + firstEncryptedRef.set(encrypted); + } + return encrypted; + }).when(producer).encryptMessage(any(), any()); + doAnswer(invocation -> { + if (failSend.getAndSet(false)) { + throw new RuntimeException("mocked send failure"); + } + ByteBuf payload = invocation.getArgument(5); + ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer(); + header.writeInt(4 + 4 + payload.readableBytes()); + header.writeInt(0); + return ByteBufPair.get(header, payload); + }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any()); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + List> messages = addMessages(batchMessageContainer, 2); + + assertThatThrownBy(batchMessageContainer::createOpSendMsg).isInstanceOf(RuntimeException.class); + // The encrypted payload is orphaned once sendMessage fails; it must be released, not leaked. + assertEquals(firstEncryptedRef.get().refCnt(), 0); + batchMessageContainer.resetPayloadAfterFailedPublishing(); + + ProducerImpl.OpSendMsg op = batchMessageContainer.createOpSendMsg(); + assertNotNull(op); + assertNotNull(op.cmd); + op.cmd.release(); + batchMessageContainer.clear(); + messages.forEach(ReferenceCountUtil::safeRelease); + } + + /** + * Without compression or encryption, a failed build must still leave a retry that produces a well-formed, + * parseable SEND frame with balanced ref-counts. + */ + @Test + public void testNoCompressionBuildFailureProducesValidFrame() throws Exception { + assertValidSendFrameAfterFailure(true, false); // failure in encryptMessage + assertValidSendFrameAfterFailure(false, true); // failure in sendMessage + } + + private void assertValidSendFrameAfterFailure(boolean failAtEncrypt, boolean failAtSend) throws Exception { + ProducerImpl producer = createTestProducer(CompressionType.NONE); + + AtomicBoolean failOnce = new AtomicBoolean(true); + if (failAtEncrypt) { + doAnswer(invocation -> { + if (failOnce.getAndSet(false)) { + throw new RuntimeException("mocked encryption failure"); + } + return invocation.getArgument(1); + }).when(producer).encryptMessage(any(), any()); + } else { + doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any()); + } + doAnswer(invocation -> { + if (failAtSend && failOnce.getAndSet(false)) { + throw new RuntimeException("mocked send failure"); + } + MessageMetadata metadata = invocation.getArgument(4); + ByteBuf payload = invocation.getArgument(5); + return Commands.newSend(0L, metadata.hasSequenceId() ? metadata.getSequenceId() : 0L, 1, + Commands.ChecksumType.Crc32c, metadata, payload); + }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any()); + + BatchMessageContainerImpl batchMessageContainer = new BatchMessageContainerImpl(producer); + List> messages = addMessages(batchMessageContainer, 3); + + // First build fails after the batch payload was produced; ProducerImpl.batchMessageAndSend() then resets. + assertThatThrownBy(batchMessageContainer::createOpSendMsg).isInstanceOf(RuntimeException.class); + batchMessageContainer.resetPayloadAfterFailedPublishing(); + + // The retry must succeed and produce a well-formed SEND frame. + ProducerImpl.OpSendMsg op = batchMessageContainer.createOpSendMsg(); + assertNotNull(op); + assertNotNull(op.cmd); + + ByteBufPair cmd = op.cmd; + ByteBuf header = cmd.getFirst(); + ByteBuf payloadBuf = cmd.getSecond(); + int totalSize = header.getInt(0); + int cmdSize = header.getInt(4); + // The total-size field must equal the number of bytes that follow it. + assertEquals(totalSize, cmd.readableBytes() - 4, + "TOTAL_SIZE must equal the number of bytes following the total-size field"); + // The command must parse cleanly as a SEND command. Skip both length fields: TOTAL_SIZE and + // CMD_SIZE (getInt above reads absolutely and does not move the reader index). + BaseCommand parsed = new BaseCommand(); + header.markReaderIndex(); + header.skipBytes(8); + parsed.parseFrom(header, cmdSize); + assertEquals(parsed.getType(), BaseCommand.Type.SEND); + header.resetReaderIndex(); + + // Ref-counts must be balanced: the op owns the batch buffer exactly once, and it is freed once. + assertEquals(payloadBuf.refCnt(), 1); + cmd.release(); + assertEquals(payloadBuf.refCnt(), 0); + + batchMessageContainer.clear(); + messages.forEach(ReferenceCountUtil::safeRelease); + } + + /** + * In multi-batch mode, a later sub-batch can fail to build after an earlier one already produced its + * operation. The already-built operations never reach the send queue, so their commands must be released + * by the container — otherwise every failed flush leaks command buffers, and repeated retries grow the + * direct memory usage. + */ + @Test(dataProvider = "compressionTypes") + public void testMultiBatchesPartialBuildFailureReleasesBuiltOps(CompressionType compressionType) + throws Exception { + ProducerImpl producer = createTestProducer(compressionType); + doAnswer(invocation -> invocation.getArgument(1)).when(producer).encryptMessage(any(), any()); + List 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(); + doAnswer(invocation -> { + if (sendCalls.incrementAndGet() == 2) { + throw new RuntimeException("mocked second sub-batch failure"); + } + ByteBuf payload = invocation.getArgument(5); + ByteBuf header = PulsarByteBufAllocator.DEFAULT.buffer(); + header.writeInt(4 + 4 + payload.readableBytes()); + header.writeInt(0); + ByteBufPair pair = ByteBufPair.get(header, payload); + builtPairs.add(pair); + builtHeaders.add(header); + builtPayloads.add(payload); + return pair; + }).when(producer).sendMessage(anyLong(), anyLong(), anyLong(), anyInt(), any(), any()); + + BatchMessageKeyBasedContainer container = new BatchMessageKeyBasedContainer(); + container.setProducer(producer); + List> messages = new ArrayList<>(); + try { + for (int i = 0; i < 4; i++) { + MessageMetadata messageMetadata = new MessageMetadata(); + messageMetadata.setSequenceId(i); + messageMetadata.setProducerName("producer"); + messageMetadata.setPublishTime(System.currentTimeMillis()); + messageMetadata.setPartitionKey(i < 2 ? "a" : "b"); + ByteBuffer payload = ByteBuffer.wrap(("payload-" + i).getBytes(StandardCharsets.UTF_8)); + MessageImpl message = MessageImpl.create(messageMetadata, payload, Schema.BYTES, null); + messages.add(message); + container.add(message, null); + } + + // Sub-batch "a" builds its operation, sub-batch "b" fails: the built command must not leak. + assertThatThrownBy(container::createOpSendMsgs) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked second"); + assertEquals(builtPairs.size(), 1); + // The pair itself is released (returned to its recycler) and the op recycled, not just the + // components freed. + assertEquals(builtPairs.get(0).refCnt(), 0); + assertEquals(builtHeaders.get(0).refCnt(), 0); + 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 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); + } + } + + /** + * 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); + // 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)); + 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. + // 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; + } + + 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; + } } 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..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 @@ -18,22 +18,51 @@ */ package org.apache.pulsar.client.impl; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyInt; +import static org.mockito.Mockito.anyLong; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; +import 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 java.util.concurrent.atomic.AtomicReference; import org.apache.commons.lang3.reflect.FieldUtils; +import org.apache.pulsar.client.api.CryptoKeyReader; +import org.apache.pulsar.client.api.MessageCrypto; +import org.apache.pulsar.client.api.ProducerCryptoFailureAction; import org.apache.pulsar.client.api.PulsarClientException; import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.impl.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; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -193,4 +222,411 @@ 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; + } + + /** 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 { + MessageCrypto msgCrypto = mock(MessageCrypto.class); + when(msgCrypto.getMaxOutputSize(anyInt())).thenReturn(64); + doThrow(new RuntimeException("mocked encryption failure")) + .when(msgCrypto).encrypt(any(), any(), any(), any(), any()); + ProducerConfigurationData conf = encryptedProducerConf(); + conf.setMessageCrypto(msgCrypto); + ProducerImpl producer = constructProducer(mockedPulsarClient(), conf); + + ByteBuf partial = Unpooled.buffer(64); + doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); + + ByteBuf source = Unpooled.buffer(8); + assertThatThrownBy(() -> producer.encryptMessage(new MessageMetadata(), source)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked encryption failure"); + assertEquals(partial.refCnt(), 0, "the partially built encrypted buffer must not leak"); + assertEquals(source.refCnt(), 1, "the source payload stays with the caller"); + source.release(); + } + + /** The SEND crypto-failure action returns the unencrypted source; the partial buffer must not leak. */ + @Test + public void testEncryptMessageCryptoFailureActionSendReleasesPartialBuffer() throws Exception { + MessageCrypto msgCrypto = mock(MessageCrypto.class); + when(msgCrypto.getMaxOutputSize(anyInt())).thenReturn(64); + doThrow(new PulsarClientException("mocked encryption failure")) + .when(msgCrypto).encrypt(any(), any(), any(), any(), any()); + ProducerConfigurationData conf = encryptedProducerConf(); + conf.setMessageCrypto(msgCrypto); + conf.setCryptoFailureAction(ProducerCryptoFailureAction.SEND); + ProducerImpl producer = constructProducer(mockedPulsarClient(), conf); + + ByteBuf partial = Unpooled.buffer(64); + doReturn(partial).when(producer).allocateEncryptedBuffer(anyInt()); + + ByteBuf source = Unpooled.buffer(8); + assertSame(producer.encryptMessage(new MessageMetadata(), source), source); + assertEquals(partial.refCnt(), 0, "the partially built encrypted buffer must not leak"); + source.release(); + } + + /** A failed command serialization must release the payload instead of orphaning it. */ + @Test + public void testSendMessageFailureReleasesPayload() throws Exception { + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new RuntimeException("mocked serialization failure")) + .when(producer) + .sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any()); + + ByteBuf payload = Unpooled.buffer(8); + assertThatThrownBy(() -> producer.sendMessageOrReleasePayload( + 1, 1, 1, null, new MessageMetadata(), payload)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked serialization failure"); + assertEquals(payload.refCnt(), 0, "the payload must be released on a failed serialization"); + } + + /** A failing compression stage must release the source payload the codec left with the caller. */ + @Test + public void testApplyCompressionFailureReleasesSource() throws Exception { + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new RuntimeException("mocked compression failure")) + .when(producer) + .applyCompression(any()); + + ByteBuf source = Unpooled.buffer(8); + assertThatThrownBy(() -> producer.applyCompressionOrReleaseSource(source)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked compression failure"); + assertEquals(source.refCnt(), 0, "the source payload must be released on a failed compression"); + } + + /** A failing encryption stage must release the source payload encryptMessage() left with the caller. */ + @Test + public void testEncryptMessageFailureReleasesSource() throws Exception { + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + doThrow(new PulsarClientException("mocked encryption failure")) + .when(producer) + .encryptMessage(any(), any()); + + ByteBuf source = Unpooled.buffer(8); + assertThatThrownBy(() -> producer.encryptMessageOrReleaseSource(new MessageMetadata(), source)) + .isInstanceOf(PulsarClientException.class) + .hasMessageContaining("mocked encryption failure"); + assertEquals(source.refCnt(), 0, "the source payload must be released on a failed encryption"); + } + + /** + * An op whose command is deferred until the schema is registered holds its payload; when the op is + * failed before the command was built (send timeout, producer close), recycle() must release it. + */ + @Test + public void testPendingSchemaOpPayloadReleasedOnFailure() { + ByteBuf payload = Unpooled.buffer(8); + ProducerImpl.OpSendMsg op = ProducerImpl.OpSendMsg.create( + mock(LatencyHistogram.class), + mock(MessageImpl.class), + null, + 1L, + mock(SendCallback.class)); + op.pendingPayload = payload; + + op.recycle(); + + assertEquals(payload.refCnt(), 0, "the deferred payload must be released when the op is recycled"); + } + + /** + * A deferred command whose first construction fails must keep the payload alive: the op stays pending + * and the next resend (reconnect) rebuilds the command from the same buffer instead of touching a + * released one. + */ + @Test + public void testDeferredCommandConstructionFailureThenRecovery() throws Exception { + ProducerImpl producer = mock(ProducerImpl.class, CALLS_REAL_METHODS); + AtomicInteger sendCalls = new AtomicInteger(); + ByteBufPair builtCmd = mock(ByteBufPair.class); + doAnswer(invocation -> { + if (sendCalls.incrementAndGet() == 1) { + throw new RuntimeException("mocked header allocation failure"); + } + return builtCmd; + }).when(producer).sendMessage(anyLong(), anyLong(), anyInt(), any(), any(), any()); + + ByteBuf payload = Unpooled.buffer(8); + ProducerImpl.OpSendMsg op = ProducerImpl.OpSendMsg.create( + mock(LatencyHistogram.class), + mock(MessageImpl.class), + null, + 1L, + mock(SendCallback.class)); + op.pendingPayload = payload; + + // First construction fails: the payload must stay with the op for the retry. + assertThatThrownBy(() -> producer.buildDeferredCommand(op, new MessageMetadata(), 1, 1, 1, null, -1)) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("mocked header allocation failure"); + assertEquals(payload.refCnt(), 1, "the payload must stay alive for the next resend"); + assertEquals(op.pendingPayload, payload, "the op keeps owning the deferred payload"); + assertNull(op.cmd); + + // The resend after reconnect rebuilds from the same buffer and hands it to the command. + producer.buildDeferredCommand(op, new MessageMetadata(), 1, 1, 1, null, -1); + assertEquals(op.cmd, builtCmd, "the retry must rebuild the command"); + assertNull(op.pendingPayload, "the payload's ownership moved into the command"); + // Recycling the op now must not release the payload again. + op.recycle(); + assertEquals(payload.refCnt(), 1, "the payload belongs to the command now, not to the op"); + payload.release(); + } + + /** + * Exercises the send-path wiring: {@code sendAsync()} itself must route through the stage helpers, so a + * failing compression or serialization stage releases the payload buffers instead of orphaning them. The + * helper-level tests above cannot detect a call site reverting to the bare method. + */ + @Test + public void testSendPathFailureReleasesPayloadThroughTheStageHelpers() throws Exception { + ProducerConfigurationData conf = new ProducerConfigurationData(); + conf.setBatchingEnabled(false); + 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"); + } + + /** + * 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); + } }