[fix][client] Fix buffer ownership on the send failure paths

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<MessageMetadata, MessageMetadata> 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

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<Boolean> 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;
         }

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<ByteBuf> allocated = new ArrayList<>();
+        ByteBufAllocator trackingAllocator = mock(ByteBufAllocator.class);
+        doAnswer(invocation -> {
+            ByteBuf buffer = Unpooled.buffer(invocation.getArgument(0));
+            allocated.add(buffer);
+            return buffer;
+        }).when(trackingAllocator).buffer(anyInt());
+        RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(trackingAllocator);
+        container.setCryptoKeyReader(cryptoKeyReader);
+        container.add(createMessage("my-topic", "hi-1", 0), null);
+
+        // Replace the real crypto with one whose encrypt() throws an unexpected RuntimeException, so the batch
+        // payload is built (getCompressedBatchMetadataAndPayload) and then encryption fails outside the
+        // PulsarClientException contract.
+        MessageCrypto<MessageMetadata, MessageMetadata> crypto = mock(MessageCrypto.class);
+        when(crypto.getMaxOutputSize(anyInt())).thenReturn(128);
+        doThrow(new RuntimeException("mocked crypto failure"))
+                .when(crypto).encrypt(anySet(), any(), any(), any(), any());
+        container.setMsgCryptoForTesting(crypto);
+
+        Throwable e = null;
+        try {
+            container.toByteBuf();
+        } catch (Throwable ex) {
+            e = ex;
+        }
+        Assert.assertEquals(e.getClass(), RuntimeException.class);
+        Assert.assertTrue(e.getMessage().contains("mocked crypto failure"));
+        // The compressed batch payload must have been released instead of leaked; the container keeps its
+        // (now released) buffer reference until the caller recovers, mirroring the producer path.
+        Assert.assertEquals(container.batchedMessageMetadataAndPayload.refCnt(), 0);
+        // The partially built encrypted output buffer must have been released as well, not only the source.
+        for (ByteBuf buffer : allocated) {
+            Assert.assertEquals(buffer.refCnt(), 0);
+        }
+
+        container.discard(null);
+    }
+
+    /**
+     * A crypto provider failing with a {@link PulsarClientException} must release the compressed batch payload
+     * and the partially built encrypted buffer and discard the batch, so the compactor can reuse the container.
+     */
+    @Test
+    public void testToByteBufReleasesPayloadAndDiscardsWhenEncryptionFailsWithClientException() throws Exception {
+        setEncryptionAndCompression(true, false);
+        List<ByteBuf> allocated = new ArrayList<>();
+        ByteBufAllocator trackingAllocator = mock(ByteBufAllocator.class);
+        doAnswer(invocation -> {
+            ByteBuf buffer = Unpooled.buffer(invocation.getArgument(0));
+            allocated.add(buffer);
+            return buffer;
+        }).when(trackingAllocator).buffer(anyInt());
+        RawBatchMessageContainerImpl container = new RawBatchMessageContainerImpl(trackingAllocator);
+        container.setCryptoKeyReader(cryptoKeyReader);
+        container.add(createMessage("my-topic", "hi-1", 0), null);
+
+        MessageCrypto<MessageMetadata, MessageMetadata> crypto = mock(MessageCrypto.class);
+        when(crypto.getMaxOutputSize(anyInt())).thenReturn(128);
+        doThrow(new PulsarClientException("mocked crypto failure"))
+                .when(crypto).encrypt(anySet(), any(), any(), any(), any());
+        container.setMsgCryptoForTesting(crypto);
+
+        try {
+            container.toByteBuf();
+            Assert.fail("expected the encryption failure to propagate");
+        } catch (RuntimeException e) {
+            Assert.assertTrue(e.getMessage().contains("Failed to encrypt payload"));
+            Assert.assertTrue(e.getCause() instanceof PulsarClientException);
+        }
+        // Unlike the unexpected-Throwable branch, the PulsarClientException branch discards the batch so the
+        // container is empty and reusable for the next flush.
+        Assert.assertEquals(container.getNumMessagesInBatch(), 0);
+        for (ByteBuf buffer : allocated) {
+            Assert.assertEquals(buffer.refCnt(), 0);
+        }
+    }
 }