[improve][client] Terminate consumer.receive() when consumer is closed (#24550)

Author: 3pacccccc

pulsar-broker/src/test/java/org/apache/pulsar/client/api/SimpleProducerConsumerTest.java

@@ -21,6 +21,7 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.pulsar.common.naming.TopicName.PARTITIONED_TOPIC_SUFFIX;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
@@ -4139,6 +4140,73 @@ public void testReceiveAsyncCompletedWhenClosing() throws Exception {
         countDownLatch3.await();
     }
 
+    @Test(timeOut = 100000)
+    public void consumerReceiveThrowExceptionWhenConsumerClose() throws Exception {
+        Consumer<byte[]> consumer = pulsarClient
+                .newConsumer()
+                .topic("persistent://my-property/my-ns/my-topic2")
+                .receiverQueueSize(10)
+                .subscriptionType(SubscriptionType.Shared)
+                .subscriptionName("my-sub")
+                .subscribe();
+
+        Thread thread = new Thread(() -> {
+            try {
+                // sleep 0.1 second to close consumer to ensure consumer.receive() is triggerd
+                Thread.sleep(1000);
+                consumer.close();
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        thread.start();
+
+        assertThatThrownBy(
+                () -> consumer.receive()
+        )
+                .isInstanceOf(PulsarClientException.class)
+                .hasMessage("java.lang.InterruptedException: Queue is terminated")
+                .hasCauseInstanceOf(InterruptedException.class);
+    }
+
+
+    @Test(timeOut = 100000)
+    public void multiThreadConsumerReceiveThrowExceptionWhenConsumerClose() throws Exception {
+        Consumer<byte[]> consumer = pulsarClient
+                .newConsumer()
+                .topic("persistent://my-property/my-ns/my-topic2")
+                .receiverQueueSize(10)
+                .subscriptionType(SubscriptionType.Shared)
+                .subscriptionName("my-sub")
+                .subscribe();
+        int threadCount = 10;
+
+        CountDownLatch terminateCompletedLatch = new CountDownLatch(threadCount);
+        CountDownLatch allThreadReadyLatch = new CountDownLatch(threadCount);
+        AtomicInteger interruptedThreadCount = new AtomicInteger(0);
+        for (int i = 0; i < threadCount; i++) {
+            new Thread(() -> {
+                allThreadReadyLatch.countDown();
+                try {
+                    consumer.receive();
+                    fail("thread should have been interrupted");
+                } catch (PulsarClientException e) {
+                    terminateCompletedLatch.countDown();
+                    interruptedThreadCount.incrementAndGet();
+                }
+            }).start();
+        }
+        // all threads should be ready in at most 3 seconds
+        assertTrue(allThreadReadyLatch.await(3, TimeUnit.SECONDS));
+        // close consumer, and all threads should be interrupted by thrown PulsarClientException
+        consumer.close();
+        // all threads should be terminated in at most 3 seconds
+        assertTrue(terminateCompletedLatch.await(3, TimeUnit.SECONDS));
+        // Verify all threads were properly terminated
+        assertEquals(interruptedThreadCount.get(), threadCount);
+    }
+
+
     @Test(timeOut = 20000)
     public void testResetPosition() throws Exception {
         final String topicName = "persistent://my-property/my-ns/testResetPosition";

pulsar-broker/src/test/java/org/apache/pulsar/client/impl/ZeroQueueSizeTest.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pulsar.client.impl;
 
+import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;
@@ -422,6 +423,38 @@ public void testZeroQueueSizeMessageRedeliveryForAsyncReceive()
         producer.close();
     }
 
+    @Test(timeOut = 30000)
+    public void testZeroQueueGetExceptionWhenReceiveBatchMessage() throws PulsarClientException {
+
+        int batchMessageDelayMs = 100;
+        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic("persistent://prop-xyz/use/ns-abc/topic1")
+                .subscriptionName("my-subscriber-name").subscriptionType(SubscriptionType.Shared).receiverQueueSize(0)
+                .subscribe();
+
+        ProducerBuilder<byte[]> producerBuilder = pulsarClient.newProducer()
+                .topic("persistent://prop-xyz/use/ns-abc/topic1")
+                .messageRoutingMode(MessageRoutingMode.SinglePartition);
+
+        producerBuilder.enableBatching(true).batchingMaxPublishDelay(batchMessageDelayMs, TimeUnit.MILLISECONDS)
+                .batchingMaxMessages(5);
+
+        Producer<byte[]> producer = producerBuilder.create();
+
+        // send a batch message to trigger zeroQueueConsumer to close
+        for (int i = 0; i < 10; i++) {
+            String message = "my-message-" + i;
+            producer.sendAsync(message.getBytes());
+        }
+
+        // when zeroQueueConsumer receive a batch message, it will close and receive method will throw exception
+        assertThatThrownBy(
+                consumer::receive
+        )
+                .isInstanceOf(PulsarClientException.class)
+                .hasMessage("java.lang.InterruptedException: Queue is terminated")
+                .hasCauseInstanceOf(InterruptedException.class);
+    }
+
     @Test(timeOut = 30000)
     public void testPauseAndResume() throws Exception {
         final String topicName = "persistent://prop/ns-abc/zero-queue-pause-and-resume";

pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java

@@ -1243,17 +1243,8 @@ private void closeConsumerTasks() {
         }
         negativeAcksTracker.close();
         stats.getStatTimeout().ifPresent(Timeout::cancel);
-        if (poolMessages) {
-            releasePooledMessagesAndStopAcceptNew();
-        }
-    }
-
-    /**
-     * If enabled pooled messages, we should release the messages after closing consumer and stop accept the new
-     * messages.
-     */
-    private void releasePooledMessagesAndStopAcceptNew() {
-        incomingMessages.terminate(message -> message.release());
+        //terminate incomingMessages queue, stop accept the new messages and waking up blocked thread.
+        incomingMessages.terminate(Message::release);
         clearIncomingMessages();
     }
 

pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/GrowableArrayBlockingQueue.java

@@ -194,6 +194,9 @@ public T take() throws InterruptedException {
 
         try {
             while (SIZE_UPDATER.get(this) == 0) {
+                if (terminated) {
+                    throw new InterruptedException("Queue is terminated");
+                }
                 isNotEmpty.await();
             }
 
@@ -220,6 +223,9 @@ public T poll(long timeout, TimeUnit unit) throws InterruptedException {
                 if (timeoutNanos <= 0) {
                     return null;
                 }
+                if (terminated) {
+                    return null;
+                }
 
                 timeoutNanos = isNotEmpty.awaitNanos(timeoutNanos);
             }
@@ -422,7 +428,8 @@ public String toString() {
     }
 
     /**
-     * Make the queue not accept new items. if there are still new data trying to enter the queue, it will be handed
+     * Make the queue not accept new items and waking up blocked consume.
+     * if there are still new data trying to enter the queue, it will be handed
      * by {@param itemAfterTerminatedHandler}.
      */
     public void terminate(@Nullable Consumer<T> itemAfterTerminatedHandler) {
@@ -436,6 +443,14 @@ public void terminate(@Nullable Consumer<T> itemAfterTerminatedHandler) {
         } finally {
             tailLock.unlockWrite(stamp);
         }
+
+        // Signal waiting consumer threads to prevent indefinite blocking after termination
+        headLock.lock();
+        try {
+            isNotEmpty.signalAll();
+        } finally {
+            headLock.unlock();
+        }
     }
 
     public boolean isTerminated() {

pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/GrowableArrayBlockingQueueTest.java

@@ -208,6 +208,117 @@ public void pollTimeout2() throws Exception {
         latch.await();
     }
 
+    /**
+     * test that multi-threads calling take(), and then terminate the queue, these threads will be properly interrupted.
+     */
+    @Test(timeOut = 10000)
+    public void testTakeBlockingThreadsTermination() throws InterruptedException {
+        GrowableArrayBlockingQueue<Integer> queue = new GrowableArrayBlockingQueue<>(2);
+        int threadCount = 10;
+        CountDownLatch terminateCompletedLatch = new CountDownLatch(threadCount);
+        CountDownLatch allThreadReadyLatch = new CountDownLatch(threadCount);
+        AtomicInteger interruptedThreadCount = new AtomicInteger(0);
+        for (int i = 0; i < threadCount; i++) {
+            new Thread(() -> {
+                try {
+                    allThreadReadyLatch.countDown();
+                    queue.take();
+                    fail("thread should have been interrupted");
+                } catch (InterruptedException e) {
+                    // Expected interruption, record and notify
+                    terminateCompletedLatch.countDown();
+                    interruptedThreadCount.incrementAndGet();
+                }
+            }).start();
+        }
+        // all threads should be ready in at most 1 second
+        assertTrue(allThreadReadyLatch.await(1, TimeUnit.SECONDS));
+        // Terminate the queue - this should interrupt all threads blocked on take()
+        queue.terminate(null);
+        // all threads should be interrupted in at most 1 second
+        assertTrue(terminateCompletedLatch.await(1, TimeUnit.SECONDS));
+        // Verify all threads were properly interrupted
+        assertEquals(interruptedThreadCount.get(), threadCount);
+    }
+
+    /**
+     * test that multi-threads calling poll(), and then terminate the queue,
+     * these threads will be terminated by return null value.
+     */
+    @Test(timeOut = 10000)
+    public void testPollBlockingThreadsTermination() throws InterruptedException {
+        GrowableArrayBlockingQueue<Integer> queue = new GrowableArrayBlockingQueue<>(2);
+        int threadCount = 10;
+        CountDownLatch terminateCompletedLatch = new CountDownLatch(threadCount);
+        CountDownLatch allThreadReadyLatch = new CountDownLatch(threadCount);
+        AtomicInteger terminateThreadCount = new AtomicInteger(0);
+        for (int i = 0; i < threadCount; i++) {
+            new Thread(() -> {
+                try {
+                    allThreadReadyLatch.countDown();
+                    // poll message at a very long timeout, so it will keep pending
+                    Integer poll = queue.poll(1, TimeUnit.HOURS);
+                    // should return a null value if queue is terminated
+                    assertNull(poll);
+                    terminateCompletedLatch.countDown();
+                    terminateThreadCount.incrementAndGet();
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            }).start();
+        }
+        // all threads should be ready in at most 1 second
+        assertTrue(allThreadReadyLatch.await(1, TimeUnit.SECONDS));
+        // Terminate the queue - this should make all threads return null value
+        queue.terminate(null);
+        // all threads should be terminated in at most 1 second
+        assertTrue(terminateCompletedLatch.await(1, TimeUnit.SECONDS));
+        assertEquals(terminateThreadCount.get(), threadCount);
+    }
+
+    /**
+     * test that multi-threads calling take() and poll() mix, and then terminate the queue,
+     * the queue will signal all these threads properly.
+     */
+    @Test(timeOut = 10000)
+    public void testPollTakeMixBlockingThreadsTermination() throws InterruptedException {
+        GrowableArrayBlockingQueue<Integer> queue = new GrowableArrayBlockingQueue<>(2);
+        int threadCount = 10;
+        CountDownLatch terminateCompletedLatch = new CountDownLatch(threadCount);
+        CountDownLatch allThreadReadyLatch = new CountDownLatch(threadCount);
+        AtomicInteger terminateThreadCount = new AtomicInteger(0);
+        for (int i = 0; i < threadCount; i++) {
+            int finalI = i;
+            new Thread(() -> {
+                allThreadReadyLatch.countDown();
+                if (finalI % 2 == 0) {
+                    try {
+                        queue.take();
+                    } catch (InterruptedException e) {
+                        terminateCompletedLatch.countDown();
+                        terminateThreadCount.incrementAndGet();
+                    }
+                } else {
+                    try {
+                        Integer poll = queue.poll(1, TimeUnit.HOURS);
+                        assertNull(poll);
+                        terminateCompletedLatch.countDown();
+                        terminateThreadCount.incrementAndGet();
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+            }).start();
+        }
+        // all threads should be ready in at most 1 second
+        assertTrue(allThreadReadyLatch.await(1, TimeUnit.SECONDS));
+        // Terminate the queue - this should make all threads return null value
+        queue.terminate(null);
+        // all threads should be terminated in at most 1 second
+        assertTrue(terminateCompletedLatch.await(1, TimeUnit.SECONDS));
+        assertEquals(terminateThreadCount.get(), threadCount);
+    }
+
     @Test
     public void removeTest() {
         BlockingQueue<Integer> queue = new GrowableArrayBlockingQueue<>(4);