Connector must invoke callback on registration error

When registering a command or event handler, the returned registration should invoke any registered acknowledgement callbacks, whether that registration was successful or not. An additional callback method is added to allow distinguishing between these situations.

This issue is the root cause of AxonIQ/AxonFramework#3938

Author: abuijze

src/main/java/io/axoniq/axonserver/connector/Registration.java

@@ -16,19 +16,21 @@
 
 package io.axoniq.axonserver.connector;
 
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Consumer;
 
 /**
- * Interface describing an instruction to perform a registration, which can be cancelled.
+ * Interface describing an instruction to perform a registration, which can be canceled.
  */
 @FunctionalInterface
 public interface Registration {
 
     /**
      * Cancel the registration from which this instance was returned. Does nothing if the registration has already been
-     * cancelled, or when the registration was undone by another mechanism (such as a new registration overriding this
+     * canceled, or when the registration was undone by another mechanism (such as a new registration overriding this
      * one).
      *
      * @return a {@link CompletableFuture} of {@link Void} to react when {@code this} {@link Registration} has been
@@ -44,7 +46,7 @@ public interface Registration {
      * @param timeout the duration to wait until the operation has been acknowledged
      * @param unit    the {@link TimeUnit} for the given {@code timeout} to wait until the operation has been
      *                acknowledged
-     * @return {@code this} {@link Registration} to support a fluent API
+     * @return this Registration to support a fluent API
      * @throws TimeoutException     is thrown when the given {@code timeout} and {@code unit} is surpassed
      * @throws InterruptedException is thrown when the thread waiting for the acknowledgement is interrupted
      */
@@ -56,13 +58,31 @@ default Registration awaitAck(long timeout, TimeUnit unit) throws TimeoutExcepti
      * Registers the given {@code runnable} to {@code this} {@link Registration} to be executed when the acknowledgement
      * of {@code this} {@link Registration} is received. Allows for the addition of further logic to {@code this
      * Registration}, like invoking {@link #awaitAck(long, TimeUnit)} for example.
+     * <p/>
+     * The given {@code runnable} is invoked, regardless of successful or failed acknowledgement. Use
+     * {@link #onAck(Consumer)} to register an action that needs to distinguish between successful and failed
+     * registration.
      *
      * @param runnable the {@link Runnable} to execute when the acknowledgement of {@code this} {@link Registration} is
      *                 received
-     * @return {@code this} {@link Registration} to support a fluent API
+     * @return this Registration to support a fluent API
      */
     default Registration onAck(Runnable runnable) {
         runnable.run();
         return this;
     }
+
+    /**
+     * Registers the given {@code runnable} to {@code this} {@link Registration} to be executed when the acknowledgement
+     * of {@code this} {@link Registration} is received. Allows for the addition of further logic to
+     * {@code this Registration}, like invoking {@link #awaitAck(long, TimeUnit)} for example.
+     *
+     * @param action the action to execute when the acknowledgement of {@code this} {@link Registration} is received,
+     *               either normally or exceptionally.
+     * @return this Registration to support a fluent API
+     */
+    default Registration onAck(Consumer<Optional<Throwable>> action) {
+        action.accept(Optional.empty());
+        return this;
+    }
 }

src/main/java/io/axoniq/axonserver/connector/impl/AsyncRegistration.java

@@ -20,10 +20,12 @@
 import io.axoniq.axonserver.connector.ErrorCategory;
 import io.axoniq.axonserver.connector.Registration;
 
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 /**
@@ -69,7 +71,15 @@ public Registration awaitAck(long timeout, TimeUnit unit) throws TimeoutExceptio
 
     @Override
     public Registration onAck(Runnable runnable) {
-        requestAck.thenRun(runnable);
+        requestAck.exceptionally(e -> null).thenRun(runnable);
+        return this;
+    }
+
+    @Override
+    public Registration onAck(Consumer<Optional<Throwable>> action) {
+        requestAck.thenApply(ignored -> Optional.<Throwable>empty())
+                  .exceptionally(Optional::of)
+                  .thenAccept(action);
         return this;
     }
 }

src/test/java/io/axoniq/axonserver/connector/command/CommandChannelIntegrationTest.java

@@ -16,6 +16,7 @@
 
 package io.axoniq.axonserver.connector.command;
 
+import eu.rekawek.toxiproxy.model.Toxic;
 import io.axoniq.axonserver.connector.AbstractAxonServerIntegrationTest;
 import io.axoniq.axonserver.connector.AxonServerConnection;
 import io.axoniq.axonserver.connector.AxonServerConnectionFactory;
@@ -33,17 +34,21 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static io.axoniq.axonserver.connector.impl.ObjectUtils.doIfNotNull;
 import static io.axoniq.axonserver.connector.impl.ObjectUtils.silently;
 import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertWithin;
 import static org.junit.jupiter.api.Assertions.*;
+import static org.testcontainers.shaded.org.awaitility.Awaitility.await;
 
 class CommandChannelIntegrationTest extends AbstractAxonServerIntegrationTest {
 
@@ -59,13 +64,15 @@ void setUp() {
                                                                    "client1")
                                                         .routingServers(axonServerAddress)
                                                         .forceReconnectViaRoutingServers(false)
+                                                        .connectTimeout(1500, TimeUnit.MILLISECONDS)
                                                         .reconnectInterval(500, TimeUnit.MILLISECONDS)
                                                         .build();
         connection1 = connectionFactory1.connect("default");
 
         connectionFactory2 = AxonServerConnectionFactory.forClient(getClass().getSimpleName(),
                                                                    "client2")
                                                         .routingServers(axonServerAddress)
+                                                        .connectTimeout(1500, TimeUnit.MILLISECONDS)
                                                         .reconnectInterval(500, TimeUnit.MILLISECONDS)
                                                         .forceReconnectViaRoutingServers(false)
                                                         .build();
@@ -74,9 +81,13 @@ void setUp() {
     }
 
     @AfterEach
-    void tearDown() {
+    void tearDown() throws IOException {
         silently(connectionFactory1, AxonServerConnectionFactory::shutdown);
         silently(connectionFactory2, AxonServerConnectionFactory::shutdown);
+        axonServerProxy.enable();
+        for (Toxic toxic : axonServerProxy.toxics().getAll()) {
+            toxic.remove();
+        }
     }
 
     @Test
@@ -293,6 +304,22 @@ void unsubscribingHandlerReturnsUnknownHandlerForCommand() throws TimeoutExcepti
 
     }
 
+    @Test
+    void subscribingHandlersWithAxonServerUnavailableAcknowledgesSubscription() throws IOException {
+        axonServerProxy.disable();
+
+        AtomicBoolean acked = new AtomicBoolean(false);
+        AtomicReference<Optional<Throwable>> ackError = new AtomicReference<>();
+        Registration registration = connection1.commandChannel().registerCommandHandler(this::mockHandler,
+                                                                                        100,
+                                                                                        "testCommand");
+
+        registration.onAck(() -> acked.set(true))
+                    .onAck(ackError::set);
+        await().until(ackError::get, v -> v != null && v.isPresent());
+        await().untilTrue(acked);
+    }
+
     private CompletableFuture<CommandResponse> mockHandler(Command command) {
         return CompletableFuture.completedFuture(CommandResponse.getDefaultInstance());
     }

src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java

@@ -17,6 +17,7 @@
 package io.axoniq.axonserver.connector.query;
 
 import com.google.protobuf.ByteString;
+import eu.rekawek.toxiproxy.model.Toxic;
 import io.axoniq.axonserver.connector.AbstractAxonServerIntegrationTest;
 import io.axoniq.axonserver.connector.AxonServerConnection;
 import io.axoniq.axonserver.connector.AxonServerConnectionFactory;
@@ -39,19 +40,22 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.Queue;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static io.axoniq.axonserver.connector.impl.ObjectUtils.doIfNotNull;
 import static io.axoniq.axonserver.connector.testutils.AssertUtils.*;
 import static org.junit.jupiter.api.Assertions.*;
+import static org.testcontainers.shaded.org.awaitility.Awaitility.await;
 
 class QueryChannelIntegrationTest extends AbstractAxonServerIntegrationTest {
 
@@ -83,9 +87,13 @@ void setUp() {
     }
 
     @AfterEach
-    void tearDown() {
+    void tearDown() throws Exception {
         connectionFactory1.shutdown();
         connectionFactory2.shutdown();
+        axonServerProxy.enable();
+        for (Toxic toxic : axonServerProxy.toxics().getAll()) {
+            toxic.remove();
+        }
     }
 
     @Test
@@ -481,6 +489,22 @@ void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws Int
         assertTrue(result.updates().isClosed());
     }
 
+    @Test
+    void subscribingHandlersWithAxonServerUnavailableAcknowledgesSubscription() throws IOException {
+        axonServerProxy.disable();
+
+        AtomicBoolean acked = new AtomicBoolean(false);
+        AtomicReference<Optional<Throwable>> ackError = new AtomicReference<>();
+        Registration registration = connection1.queryChannel().registerQueryHandler(this::mockHandler,
+                                                                                    new QueryDefinition("testQuery",
+                                                                                                        String.class));
+
+        registration.onAck(() -> acked.set(true))
+                    .onAck(ackError::set);
+        await().until(ackError::get, v -> v != null && v.isPresent());
+        await().untilTrue(acked);
+    }
+
     private void mockHandler(QueryRequest query, ReplyChannel<QueryResponse> responseHandler) {
         responseHandler.sendLast(QueryResponse.newBuilder().setRequestIdentifier(query.getMessageIdentifier()).setPayload(query.getPayload()).build());
     }