Fix MaxQueueSize semaphore release leak in createOpSendMsg

apache · Aug 8, 2022 · 2fd6ac9 · 2fd6ac9
1 parent 752225a
commit 2fd6ac9
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 43 deletions.
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/ProducerSemaphoreTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/ProducerSemaphoreTest.java
@@ -18,18 +18,15 @@
  */
 package org.apache.pulsar.client.impl;
 
-import static org.mockito.ArgumentMatchers.any;
 import java.lang.reflect.Field;
 import java.nio.charset.StandardCharsets;
-import java.util.concurrent.TimeUnit;
 import lombok.Cleanup;
 import org.apache.pulsar.client.api.MessageId;
 import org.apache.pulsar.client.api.ProducerConsumerBase;
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.common.api.proto.MessageMetadata;
 import org.apache.pulsar.common.util.FutureUtil;
-import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
@@ -58,6 +55,39 @@ public void cleanup() throws Exception {
         super.internalCleanup();
     }
 
+    @Test(timeOut = 10_000)
+    public void testProducerSemaphoreInvalidMessage() throws Exception {
+        final int pendingQueueSize = 100;
+
+        @Cleanup
+        ProducerImpl<byte[]> producer = (ProducerImpl<byte[]>) pulsarClient.newProducer()
+                .topic("testProducerSemaphoreAcquire")
+                .maxPendingMessages(pendingQueueSize)
+                .enableBatching(true)
+                .create();
+
+        this.stopBroker();
+
+        Field maxMessageSizeFiled = ClientCnx.class.getDeclaredField("maxMessageSize");
+        maxMessageSizeFiled.setAccessible(true);
+        maxMessageSizeFiled.set(null, 2);
+
+        try {
+            producer.send("semaphore-test".getBytes(StandardCharsets.UTF_8));
+            Assert.fail("can not reach here");
+        } catch (PulsarClientException.InvalidMessageException ex) {
+            Assert.assertEquals(producer.getSemaphore().get().availablePermits(), pendingQueueSize);
+        }
+
+        producer.conf.setBatchingEnabled(false);
+        try {
+            producer.send("semaphore-test".getBytes(StandardCharsets.UTF_8));
+            Assert.fail("can not reach here");
+        } catch (PulsarClientException.InvalidMessageException ex) {
+            Assert.assertEquals(producer.getSemaphore().get().availablePermits(), pendingQueueSize);
+        }
+    }
+
     @Test(timeOut = 30000)
     public void testProducerSemaphoreAcquireAndRelease() throws PulsarClientException, ExecutionException, InterruptedException {
 
@@ -209,44 +239,4 @@ public void testEnsureNotBlockOnThePendingQueue() throws Exception {
         Assert.assertEquals(producer.getSemaphore().get().availablePermits(), pendingQueueSize);
         Assert.assertFalse(producer.isErrorStat());
     }
-
-    @Test(timeOut = 10_000)
-    public void testBatchMessageSendTimeoutProducerSemaphoreRelease() throws Exception {
-        final int pendingQueueSize = 10;
-        @Cleanup
-        ProducerImpl<byte[]> producer =
-                (ProducerImpl<byte[]>) pulsarClient.newProducer()
-                        .topic("testProducerSemaphoreRelease")
-                        .sendTimeout(2, TimeUnit.SECONDS)
-                        .maxPendingMessages(pendingQueueSize)
-                        .enableBatching(true)
-                        .batchingMaxPublishDelay(100, TimeUnit.MILLISECONDS)
-                        .batchingMaxBytes(15)
-                        .create();
-        this.stopBroker();
-        try {
-            ProducerImpl<byte[]> spyProducer = Mockito.spy(producer);
-            // Make the pendingMessages not empty
-            spyProducer.newMessage().value("semaphore-test".getBytes(StandardCharsets.UTF_8)).sendAsync();
-            spyProducer.newMessage().value("semaphore-test".getBytes(StandardCharsets.UTF_8)).sendAsync();
-
-            Field batchMessageContainerField = ProducerImpl.class.getDeclaredField("batchMessageContainer");
-            batchMessageContainerField.setAccessible(true);
-            BatchMessageContainerImpl batchMessageContainer =
-                    (BatchMessageContainerImpl) batchMessageContainerField.get(spyProducer);
-            batchMessageContainer.setProducer(spyProducer);
-            Mockito.doThrow(new PulsarClientException.CryptoException("crypto error")).when(spyProducer)
-                    .encryptMessage(any(), any());
-
-            try {
-                spyProducer.newMessage().value("memory-test".getBytes(StandardCharsets.UTF_8)).sendAsync().get();
-            } catch (Exception e) {
-                throw PulsarClientException.unwrap(e);
-            }
-
-            throw new IllegalStateException("can not reach here");
-        } catch (PulsarClientException.TimeoutException ex) {
-            Assert.assertEquals(producer.getSemaphore().get().availablePermits(), pendingQueueSize);
-        }
-    }
 }
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
@@ -199,6 +199,7 @@ public boolean isMultiBatches() {
     public OpSendMsg createOpSendMsg() throws IOException {
         ByteBuf encryptedPayload = producer.encryptMessage(messageMetadata, getCompressedBatchMetadataAndPayload());
         if (encryptedPayload.readableBytes() > ClientCnx.getMaxMessageSize()) {
+            producer.semaphoreRelease(messages.size());
             discard(new PulsarClientException.InvalidMessageException(
                     "Message size is bigger than " + ClientCnx.getMaxMessageSize() + " bytes"));
             return null;