fix: ensure onCancelled callback is properly invoked (#376)

The onCancelled callback was not being called when a subscription was
stopped by the user or when the onEvent handler threw an exception.
This fixes both cases by propagating cancellation through the consumer
and catching handler exceptions in ReadResponseObserver.

Also reorders onSubscriptionConfirmation to complete the future before
calling the listener, preventing race conditions.

Co-authored-by: William Chong <william-chong@outlook.com>

Author: github-actions[bot]

src/main/java/io/kurrent/dbclient/ReadResponseObserver.java

@@ -52,6 +52,8 @@ void cancel(String reason, Throwable cause) {
         this.requestStream.cancel(reason, cause);
         if (cause instanceof StreamNotFoundException)
             this.consumer.onStreamNotFound(((StreamNotFoundException) cause).getStreamName());
+        else
+            this.consumer.onCancelled(cause);
     }
 
     void manageFlowControl() {
@@ -105,6 +107,18 @@ public void onNext(StreamsOuterClass.ReadResp value) {
             return;
         }
 
+        try {
+            handleMessage(value);
+        } catch (Exception ex) {
+            logger.error("Exception thrown by subscription handler", ex);
+            cancel("subscription handler threw an exception", ex);
+            return;
+        }
+
+        manageFlowControl();
+    }
+
+    private void handleMessage(StreamsOuterClass.ReadResp value) {
         if (value.hasEvent())
             consumer.onEvent(ResolvedEvent.fromWire(value.getEvent()));
         else if (value.hasConfirmation())
@@ -136,8 +150,6 @@ else if (value.hasLastAllStreamPosition()) {
         } else {
             logger.warn("received unknown message variant");
         }
-
-        manageFlowControl();
     }
 
     @Override
@@ -149,6 +161,7 @@ public void onError(Throwable t) {
             StatusRuntimeException e = (StatusRuntimeException) t;
 
             if (e.getStatus().getCode() == Status.Code.CANCELLED) {
+                this.consumer.onCancelled(null);
                 return;
             }
 

src/main/java/io/kurrent/dbclient/SubscriptionStreamConsumer.java

@@ -32,8 +32,8 @@ public void onEvent(ResolvedEvent event) {
     @Override
     public void onSubscriptionConfirmation(String subscriptionId) {
         this.subscription = new Subscription(this.internal, subscriptionId, this.checkpointer);
-        this.listener.onConfirmation(this.subscription);
         this.future.complete(this.subscription);
+        this.listener.onConfirmation(this.subscription);
     }
 
     @Override

src/test/java/io/kurrent/dbclient/MiscTests.java

@@ -1,8 +1,10 @@
 package io.kurrent.dbclient;
 
+import org.junit.platform.suite.api.SelectClasses;
 import org.junit.platform.suite.api.SelectPackages;
 import org.junit.platform.suite.api.Suite;
 
 @Suite
 @SelectPackages("io.kurrent.dbclient.misc")
+@SelectClasses(SubscriptionStreamConsumerTests.class)
 public class MiscTests {}

src/test/java/io/kurrent/dbclient/SubscriptionStreamConsumerTests.java

@@ -0,0 +1,79 @@
+package io.kurrent.dbclient;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class SubscriptionStreamConsumerTests {
+    private static final SubscriptionTracingCallback NO_OP_TRACING =
+        (subscriptionId, event, action) -> action.run();
+
+    @Test
+    public void testOnCancelledCompletesFutureExceptionallyBeforeConfirmation() {
+        CompletableFuture<Subscription> future = new CompletableFuture<>();
+        AtomicBoolean listenerCalled = new AtomicBoolean(false);
+        AtomicReference<Throwable> listenerException = new AtomicReference<>();
+
+        SubscriptionStreamConsumer consumer = new SubscriptionStreamConsumer(
+            new SubscriptionListener() {
+                @Override
+                public void onCancelled(Subscription subscription, Throwable exception) {
+                    listenerCalled.set(true);
+                    listenerException.set(exception);
+                }
+            },
+            null,
+            future,
+            NO_OP_TRACING
+        );
+
+        RuntimeException error = new RuntimeException("server went away");
+        consumer.onCancelled(error);
+
+        Assertions.assertTrue(future.isCompletedExceptionally());
+        ExecutionException ex = Assertions.assertThrows(ExecutionException.class, future::get);
+        Assertions.assertSame(error, ex.getCause());
+        Assertions.assertTrue(listenerCalled.get());
+        Assertions.assertSame(error, listenerException.get());
+    }
+
+    @Test
+    public void testOnCancelledAfterConfirmationDoesNotAffectFuture() {
+        CompletableFuture<Subscription> future = new CompletableFuture<>();
+        AtomicBoolean listenerCalled = new AtomicBoolean(false);
+        AtomicReference<Subscription> listenerSubscription = new AtomicReference<>();
+
+        SubscriptionStreamConsumer consumer = new SubscriptionStreamConsumer(
+            new SubscriptionListener() {
+                @Override
+                public void onCancelled(Subscription subscription, Throwable exception) {
+                    listenerCalled.set(true);
+                    listenerSubscription.set(subscription);
+                }
+            },
+            null,
+            future,
+            NO_OP_TRACING
+        );
+
+        consumer.onSubscribe(new org.reactivestreams.Subscription() {
+            @Override public void request(long n) {}
+            @Override public void cancel() {}
+        });
+
+        consumer.onSubscriptionConfirmation("test-sub-id");
+        Assertions.assertTrue(future.isDone());
+        Assertions.assertFalse(future.isCompletedExceptionally());
+
+        consumer.onCancelled(null);
+
+        Assertions.assertFalse(future.isCompletedExceptionally(), "future should remain successfully completed");
+        Assertions.assertTrue(listenerCalled.get());
+        Assertions.assertNotNull(listenerSubscription.get());
+        Assertions.assertEquals("test-sub-id", listenerSubscription.get().getSubscriptionId());
+    }
+}

src/test/java/io/kurrent/dbclient/streams/SubscriptionTests.java

@@ -12,7 +12,9 @@
 import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 public interface SubscriptionTests extends ConnectionAware {
     @Test
@@ -185,4 +187,80 @@ public void onCaughtUp(Subscription subscription, Instant timestamp, Long stream
         caughtUp.await();
         subscription.stop();
     }
+
+    @Test
+    @Timeout(value = 2, unit = TimeUnit.MINUTES)
+    default void testOnCancelledIsInvokedWhenSubscriptionIsStopped() throws Throwable {
+        KurrentDBClient client = getDefaultClient();
+        String streamName = generateName();
+
+        client.appendToStream(
+            streamName,
+            EventData.builderAsJson(generateName(), "{}".getBytes()).build()
+        ).get();
+
+        final CountDownLatch onEventReceived = new CountDownLatch(1);
+        final CountDownLatch onCancelledReceived = new CountDownLatch(1);
+        final AtomicReference<Throwable> cancelException = new AtomicReference<>();
+        final AtomicBoolean onCancelledInvoked = new AtomicBoolean(false);
+
+        Subscription subscription = client.subscribeToStream(streamName, new SubscriptionListener() {
+            @Override
+            public void onEvent(Subscription subscription, ResolvedEvent event) {
+                onEventReceived.countDown();
+            }
+
+            @Override
+            public void onCancelled(Subscription subscription, Throwable exception) {
+                onCancelledInvoked.set(true);
+                cancelException.set(exception);
+                onCancelledReceived.countDown();
+            }
+        }).get();
+
+        onEventReceived.await();
+        subscription.stop();
+        onCancelledReceived.await();
+
+        Assertions.assertTrue(onCancelledInvoked.get());
+        Assertions.assertNull(cancelException.get());
+    }
+
+    @Test
+    @Timeout(value = 2, unit = TimeUnit.MINUTES)
+    default void testOnCancelledIsInvokedWhenOnEventThrows() throws Throwable {
+        KurrentDBClient client = getDefaultClient();
+        String streamName = generateName();
+
+        final CountDownLatch onCancelledReceived = new CountDownLatch(1);
+        final AtomicReference<Throwable> cancelException = new AtomicReference<>();
+        final AtomicBoolean onCancelledInvoked = new AtomicBoolean(false);
+
+        client.subscribeToStream(streamName, new SubscriptionListener() {
+            @Override
+            public void onEvent(Subscription subscription, ResolvedEvent event) {
+                throw new RuntimeException("failure");
+            }
+
+            @Override
+            public void onCancelled(Subscription subscription, Throwable exception) {
+                onCancelledInvoked.set(true);
+                cancelException.set(exception);
+                onCancelledReceived.countDown();
+            }
+        }).get();
+
+        client.appendToStream(
+            streamName,
+            EventData.builderAsJson(generateName(), "{}".getBytes()).build()
+        ).get();
+
+        onCancelledReceived.await();
+
+        Assertions.assertTrue(onCancelledInvoked.get());
+        Assertions.assertNotNull(cancelException.get());
+        Assertions.assertInstanceOf(RuntimeException.class, cancelException.get());
+        Assertions.assertEquals("failure", cancelException.get().getMessage());
+    }
+
 }