Stop dispatching as soon as Consumer#close() is called

References #142

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java

@@ -236,7 +236,6 @@ private static class SubscriptionTracker {
     private volatile boolean hasReceivedSomething = false;
     private volatile byte subscriptionIdInClient;
     private volatile ClientSubscriptionsManager manager;
-    private volatile boolean closing = false;
 
     private SubscriptionTracker(
         StreamConsumer consumer,
@@ -254,7 +253,6 @@ private SubscriptionTracker(
     }
 
     synchronized void cancel() {
-      this.closing = true;
       if (this.manager != null) {
         LOGGER.debug("Removing consumer from manager " + this.consumer);
         this.manager.remove(this);
@@ -263,10 +261,6 @@ synchronized void cancel() {
       }
     }
 
-    boolean isClosing() {
-      return this.closing;
-    }
-
     synchronized void assign(byte subscriptionIdInClient, ClientSubscriptionsManager manager) {
       this.subscriptionIdInClient = subscriptionIdInClient;
       this.manager = manager;
@@ -420,7 +414,7 @@ private ClientSubscriptionsManager(
           (client, subscriptionId, offset, messageCount, dataSize) -> {
             SubscriptionTracker subscriptionTracker =
                 subscriptionTrackers.get(subscriptionId & 0xFF);
-            if (subscriptionTracker != null && !subscriptionTracker.isClosing()) {
+            if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) {
               client.credit(subscriptionId, 1);
             } else {
               LOGGER.debug(

src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java

@@ -95,6 +95,13 @@ class StreamConsumer implements Consumer {
         messageHandlerWithOrWithoutTracking = messageHandler;
       }
 
+      MessageHandler closedAwareMessageHandler =
+          (context, message) -> {
+            if (!closed.get()) {
+              messageHandlerWithOrWithoutTracking.handle(context, message);
+            }
+          };
+
       Runnable init =
           () -> {
             this.closingCallback =
@@ -104,7 +111,7 @@ class StreamConsumer implements Consumer {
                     offsetSpecification,
                     this.name,
                     subscriptionListener,
-                    messageHandlerWithOrWithoutTracking);
+                    closedAwareMessageHandler);
 
             this.status = Status.RUNNING;
           };

src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java

@@ -185,6 +185,51 @@ void consume() throws Exception {
     consumer.close();
   }
 
+  @Test
+  void closeOnCondition() throws Exception {
+    int messageCount = 50_000;
+    CountDownLatch publishLatch = new CountDownLatch(messageCount);
+    Client client =
+        cf.get(
+            new Client.ClientParameters()
+                .publishConfirmListener((publisherId, publishingId) -> publishLatch.countDown()));
+
+    client.declarePublisher(b(1), null, stream);
+    IntStream.range(0, messageCount)
+        .forEach(
+            i ->
+                client.publish(
+                    b(1),
+                    Collections.singletonList(
+                        client.messageBuilder().addData("".getBytes()).build())));
+
+    assertThat(publishLatch.await(10, TimeUnit.SECONDS)).isTrue();
+
+    int messagesToProcess = 20_000;
+
+    CountDownLatch consumeLatch = new CountDownLatch(1);
+    AtomicInteger receivedMessages = new AtomicInteger();
+    AtomicInteger processedMessages = new AtomicInteger();
+
+    Consumer consumer =
+        environment.consumerBuilder().stream(stream)
+            .offset(OffsetSpecification.first())
+            .messageHandler(
+                (context, message) -> {
+                  if (receivedMessages.incrementAndGet() <= messagesToProcess) {
+                    processedMessages.incrementAndGet();
+                  }
+                  if (receivedMessages.get() == messagesToProcess) {
+                    consumeLatch.countDown();
+                  }
+                })
+            .build();
+
+    assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue();
+    consumer.close();
+    assertThat(processedMessages).hasValue(messagesToProcess);
+  }
+
   @Test
   void creatingConsumerOnNonExistingStreamShouldThrowException() {
     String nonExistingStream = UUID.randomUUID().toString();