From cbca9056315f2f687f7ec83f7bf3729da07d51d5 Mon Sep 17 00:00:00 2001 From: Allard Buijze Date: Tue, 28 Oct 2025 12:08:10 +0100 Subject: [PATCH 1/2] Added support for streaming initial results of subscription queries Instead of relying on the INITIAL_RESULT request in the subscription query, the initial result is now considered a regular query that is executed on its own channel. The subscription query grpc request is now only used for streaming the updates. --- .../query/SubscriptionQueryResult.java | 11 +- .../query/impl/QueryChannelImpl.java | 33 ++--- .../query/impl/SubscriptionQueryStream.java | 117 ++++++++++++++---- .../ConnectorRunner_SubscriptionQuery.java | 14 ++- .../query/QueryChannelIntegrationTest.java | 52 ++++---- .../query/StreamingQueryIntegrationTest.java | 8 +- 6 files changed, 149 insertions(+), 86 deletions(-) diff --git a/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java b/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java index db392cd6..2cebc7c1 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java @@ -30,15 +30,14 @@ public interface SubscriptionQueryResult { /** - * Returns a CompletableFuture that completes when the initial result of the query is available. If an error - * occurred while querying, the CompletableFuture completes exceptionally. + * Returns a ResultStream that represent the initial result of the subscription query. *

- * Invoking this method will send a request for the initial result, if that hasn't been requested before. Subsequent - * invocations of this method will return the same CompletableFuture instance. + * Invoking this method will send a request for the initial result if that hasn't been requested before. Further + * invocations of this method will return the same ResultStream instance. * - * @return a CompletableFuture that completes with the initial result of the subscription query + * @return a ResultStream that provides the initial result of the subscription query */ - CompletableFuture initialResult(); + ResultStream initialResult(); /** * Returns the stream of updates to the queried projection. The stream can be read in a blocking and non-blocking diff --git a/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java b/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java index ecd3987c..12575fa2 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java @@ -367,7 +367,7 @@ public ResultStream query(QueryRequest query) { if (query.getMessageIdentifier().isEmpty()) { query = query.toBuilder().setMessageIdentifier(UUID.randomUUID().toString()).build(); } - AbstractBufferedStream results = new AbstractBufferedStream( + AbstractBufferedStream results = new AbstractBufferedStream<>( clientIdentification.getClientId(), 32, 8 ) { @@ -386,7 +386,7 @@ public void close() { outboundStream().cancel("Client cancelled the stream.", null); } }; - if ("".equals(query.getMessageIdentifier())) { + if (query.getMessageIdentifier().isEmpty()) { logger.debug("No message identifier has been set on the query. Adding a random identifier now."); query = query.toBuilder().setMessageIdentifier(UUID.randomUUID().toString()).build(); } @@ -407,10 +407,11 @@ public SubscriptionQueryResult subscriptionQuery(QueryRequest query, finalQuery = query; } String subscriptionId = finalQuery.getMessageIdentifier(); - CompletableFuture initialResultFuture = new CompletableFuture<>(); SubscriptionQueryStream subscriptionStream = new SubscriptionQueryStream( - subscriptionId, initialResultFuture, QueryChannelImpl.this.clientIdentification.getClientId(), - bufferSize, fetchSize + subscriptionId, + clientIdentification.getClientId(), () -> QueryChannelImpl.this.query(finalQuery), + bufferSize, + fetchSize ); StreamObserver upstream = queryServiceStub.subscription(subscriptionStream); subscriptionStream.enableFlowControl(); @@ -422,25 +423,14 @@ public SubscriptionQueryResult subscriptionQuery(QueryRequest query, upstream.onNext(SubscriptionQueryRequest.newBuilder().setSubscribe(subscriptionQuery).build()); return new SubscriptionQueryResult() { - private final AtomicBoolean initialResultRequested = new AtomicBoolean(); - @Override - public CompletableFuture initialResult() { - if (!initialResultFuture.isDone() && !initialResultRequested.getAndSet(true)) { - SubscriptionQuery.Builder initialResultRequest = - SubscriptionQuery.newBuilder() - .setQueryRequest(finalQuery) - .setSubscriptionIdentifier(subscriptionId); - upstream.onNext(SubscriptionQueryRequest.newBuilder() - .setGetInitialResult(initialResultRequest) - .build()); - } - return initialResultFuture; + public ResultStream initialResult() { + return subscriptionStream.initialStream(); } @Override public ResultStream updates() { - return subscriptionStream.buffer(); + return subscriptionStream.updatesBuffer(); } }; } @@ -588,7 +578,7 @@ private boolean booleanProcessingValue(List instructions, } private void handleQuery(QueryProviderInbound inbound, ReplyChannel result) { - doHandleQuery(inbound, new ReplyChannel() { + doHandleQuery(inbound, new ReplyChannel<>() { @Override public void send(QueryResponse response) { result.send(QueryProviderOutbound.newBuilder().setQueryResponse(response).build()); @@ -615,7 +605,6 @@ public void completeWithError(ErrorMessage errorMessage) { public void completeWithError(ErrorCategory errorCategory, String message) { result.completeWithError(errorCategory, message); } - }); } @@ -623,7 +612,7 @@ private void getInitialResult(QueryProviderInbound query, ReplyChannel() { + new ReplyChannel<>() { @Override public void send(QueryResponse response) { diff --git a/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java b/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java index c89a72d6..0bbdfda9 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java @@ -32,7 +32,10 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.concurrent.CompletableFuture; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; /** * A {@link FlowControlledStream} implementation to return the results of a {@link SubscriptionQuery}. @@ -42,28 +45,29 @@ public class SubscriptionQueryStream extends FlowControlledStream initialResultFuture; private final AbstractBufferedStream updateBuffer; - private final String clientId; + private final Supplier> initialResultSupplier; + + private final AtomicReference> initialResult = new AtomicReference<>(); /** * Instantiates a {@link SubscriptionQueryStream} to stream {@link SubscriptionQuery} results. * - * @param subscriptionQueryId the identifier of this subscription query - * @param initialResultFuture the initial result of the subscription query - * @param clientId the identifier of the client initiating the subscription query - * @param bufferSize the size of the update {@link #buffer()} - * @param fetchSize the number of updates to be consumed prior to refilling the update {@link #buffer()} + * @param subscriptionQueryId the identifier of this subscription query + * @param clientId the identifier of the client initiating the subscription query + * @param initialResultSupplier a {@link Supplier} providing the initial {@link ResultStream} of the subscription query + * @param bufferSize the size of the update {@link #updatesBuffer()} + * @param fetchSize the number of updates to be consumed prior to refilling the update + * {@link #updatesBuffer()} */ public SubscriptionQueryStream(String subscriptionQueryId, - CompletableFuture initialResultFuture, String clientId, + Supplier> initialResultSupplier, int bufferSize, int fetchSize) { super(clientId, bufferSize, fetchSize); - this.clientId = clientId; + this.initialResultSupplier = initialResultSupplier; this.subscriptionQueryId = subscriptionQueryId; - this.initialResultFuture = initialResultFuture; this.updateBuffer = new SubscriptionQueryUpdateBuffer(clientId, subscriptionQueryId, bufferSize, fetchSize); } @@ -72,10 +76,78 @@ public SubscriptionQueryStream(String subscriptionQueryId, * * @return the {@link ResultStream} buffering the {@link QueryUpdate}s to this subscription query */ - public ResultStream buffer() { + public ResultStream updatesBuffer() { return updateBuffer; } + /** + * + * Returns the {@link ResultStream} buffering the initial results to this subscription query. + * + * @return the {@link ResultStream} buffering the initial results to this subscription query + */ + public ResultStream initialStream() { + return initialResult.updateAndGet(i -> i == null ? buildInitialResultStream() : i); + } + + private ResultStream buildInitialResultStream() { + ResultStream delegate = initialResultSupplier.get(); + delegate.onAvailable(() -> { + if (delegate.getError().isPresent()) { + // close the update stream, since initial results completed with an error + updatesBuffer().close(); + } + }); + return new ResultStream<>() { + @Override + public QueryResponse peek() { + return delegate.peek(); + } + + @Override + public QueryResponse nextIfAvailable() { + return delegate.nextIfAvailable(); + } + + @Override + public QueryResponse nextIfAvailable(long timeout, TimeUnit unit) throws InterruptedException { + return delegate.nextIfAvailable(timeout, unit); + } + + @Override + public QueryResponse next() throws InterruptedException { + return delegate.next(); + } + + @Override + public void onAvailable(Runnable callback) { + delegate.onAvailable(() -> { + if (delegate.getError().isPresent()) { + // close the update stream, since initial results completed with an error + updatesBuffer().close(); + } + callback.run(); + }); + } + + @Override + public void close() { + delegate.close(); + SubscriptionQueryStream.this.updatesBuffer().close(); + } + + @Override + public boolean isClosed() { + return delegate.isClosed(); + } + + @Override + public Optional getError() { + return delegate.getError(); + } + }; + } + @Override public void onNext(SubscriptionQueryResponse value) { switch (value.getResponseCase()) { @@ -99,26 +171,26 @@ public void onNext(SubscriptionQueryResponse value) { value.getCompleteExceptionally().getClientId() ); updateBuffer.onError(exception); - if (!initialResultFuture.isDone()) { - initialResultFuture.completeExceptionally(exception); + ResultStream initial = initialResult.get(); + if (initial != null && !initial.isClosed()) { + initial.close(); } break; case INITIAL_RESULT: - logger.debug("Received subscription query initial result. Subscription Id: {}. Message Id: {}.", - value.getSubscriptionIdentifier(), - value.getMessageIdentifier()); - initialResultFuture.complete(value.getInitialResult()); - break; default: logger.info("Received unsupported message from SubscriptionQuery. " - + "It doesn't declare one of the expected types"); + + "It doesn't declare one of the expected types: {}", + value.getResponseCase()); break; } } @Override public void onError(Throwable t) { - initialResultFuture.completeExceptionally(t); + ResultStream initialResult = this.initialResult.get(); + if (initialResult != null && !initialResult.isClosed()) { + initialResult.close(); + } updateBuffer.onError(t); try { @@ -136,9 +208,6 @@ public void onError(Throwable t) { @Override public void onCompleted() { updateBuffer.onCompleted(); - if (!initialResultFuture.isDone()) { - initialResultFuture.completeExceptionally(new AxonServerException(ErrorCategory.QUERY_DISPATCH_ERROR, "Subscription query has already been completed", clientId)); - } } @Override diff --git a/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java b/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java index a3124f70..b72af194 100644 --- a/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java +++ b/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java @@ -61,7 +61,17 @@ public static void main(String[] args) throws Exception { .setQuery("java.lang.String") .setPayload(SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello world")).build()).build(), SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello")).build(), 1000, 100); - result.initialResult().whenComplete((r, e) -> System.out.println("Initial result: " + r.getPayload().getData().toStringUtf8())); + ResultStream initialResult = result.initialResult(); + initialResult.onAvailable(() -> { + QueryResponse r = initialResult.nextIfAvailable(); + if (r != null) { + System.out.println("Initial result: " + r.getPayload().getData().toStringUtf8()); + } else if (initialResult.getError().isPresent()) { + System.out.println("Initial result error: " + initialResult.getError().get().getMessage()); + } else if (initialResult.isClosed()) { + System.out.println("Initial result stream closed"); + } + }); ResultStream updates = result.updates(); long deadline = System.currentTimeMillis() + 15000; while (!updates.isClosed()) { @@ -73,7 +83,6 @@ public static void main(String[] args) throws Exception { updates.close(); } } - result.initialResult().get(5, TimeUnit.SECONDS); } finally { contextConnection.disconnect(); testSubject.shutdown(); @@ -95,6 +104,7 @@ public static void main(String[] args) { @Override public void handle(QueryRequest q, ReplyChannel r) { System.out.println("Handled query"); + r.send(QueryResponse.newBuilder().setRequestIdentifier(q.getMessageIdentifier()).setPayload(q.getPayload()).build()); r.sendLast(QueryResponse.newBuilder().setRequestIdentifier(q.getMessageIdentifier()).setPayload(q.getPayload()).build()); } diff --git a/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java index 7d82ba61..4fc08ffa 100644 --- a/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java +++ b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java @@ -32,11 +32,7 @@ import io.axoniq.axonserver.grpc.query.QueryResponse; import io.axoniq.axonserver.grpc.query.QueryUpdate; import io.axoniq.axonserver.grpc.query.SubscriptionQuery; -import io.grpc.StatusRuntimeException; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.RepeatedTest; -import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.*; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -54,17 +50,8 @@ import java.util.stream.Collectors; import static io.axoniq.axonserver.connector.impl.ObjectUtils.doIfNotNull; -import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertFalseWithin; -import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertTrueWithin; -import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertWithin; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.api.Assertions.fail; +import static io.axoniq.axonserver.connector.testutils.AssertUtils.*; +import static org.junit.jupiter.api.Assertions.*; class QueryChannelIntegrationTest extends AbstractAxonServerIntegrationTest { @@ -119,7 +106,12 @@ void testPermitsAreRenewedForSubscriptionQuery() throws Exception { for (int i = 0; i < 1000; i++) { SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); // the initial result may be requested when the subscription query was already closed. Therefore we accept exceptionally completed results. - results.add(result.initialResult().exceptionally(e -> null).whenComplete((r, e) -> result.updates().close())); + CompletableFuture completed = new CompletableFuture<>(); + results.add(completed); + result.initialResult().onAvailable(() -> { + completed.complete(null); + result.updates().close(); + }); } for (CompletableFuture result : results) { @@ -127,7 +119,8 @@ void testPermitsAreRenewedForSubscriptionQuery() throws Exception { } SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); - assertDoesNotThrow(() -> result.initialResult().get(5, TimeUnit.SECONDS)); + result.initialResult().nextIfAvailable(5, TimeUnit.SECONDS); + assertFalse(result.initialResult().getError().isPresent()); } @Test @@ -254,8 +247,14 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan .build(), 100, 10); - assertTrueWithin(1, TimeUnit.SECONDS, () -> subscriptionQuery.initialResult().isDone()); - assertFalse(subscriptionQuery.initialResult().isCompletedExceptionally()); + assertTrueWithin(1, TimeUnit.SECONDS, () -> { + // we need to consume staged messages to see the "close" signal + do { + // just consume + } while (subscriptionQuery.initialResult().nextIfAvailable() != null); + return subscriptionQuery.initialResult().isClosed(); + }); + assertFalse(subscriptionQuery.initialResult().getError().isPresent()); assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); updateHandlerRef.get().sendUpdate(QueryUpdate.newBuilder().setPayload(SerializedObject.newBuilder().setType(String.class.getName()).setData(ByteString.copyFromUtf8("Hello")).build()).build()); @@ -314,7 +313,7 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan SerializedObject.newBuilder().setType("update").build(), 100, 10); - assertEquals(subscriptionId, subscriptionQuery.initialResult().get(1, TimeUnit.SECONDS) + assertEquals(subscriptionId, subscriptionQuery.initialResult().nextIfAvailable(1, TimeUnit.SECONDS) .getRequestIdentifier()); assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); logger.info("Sending update"); @@ -470,13 +469,8 @@ void testUnsupportedSubscriptionQueryReturnsNoHandler() throws InterruptedExcept SerializedObject.newBuilder().setType("update").build(), 100, 10); assertWithin(1, TimeUnit.SECONDS, () -> assertTrue(result.updates().isClosed())); - try { - result.initialResult().get(1, TimeUnit.SECONDS); - fail("Expected an exception to be reported"); - } catch (ExecutionException e) { - assertTrue(e.getCause() instanceof StatusRuntimeException); - assertEquals(ErrorCategory.NO_HANDLER_FOR_QUERY.errorCode(), ((StatusRuntimeException) e.getCause()).getStatus().getDescription()); - } + QueryResponse firstMessage = result.initialResult().nextIfAvailable(1, TimeUnit.SECONDS); + assertTrue(firstMessage.hasErrorMessage()); } @Test @@ -491,7 +485,7 @@ void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws Int SerializedObject.newBuilder().setType("update").build(), 100, 10); - result.initialResult().get(1, TimeUnit.SECONDS); + result.initialResult().nextIfAvailable(1, TimeUnit.SECONDS); assertNull(result.updates().nextIfAvailable()); assertFalse(result.updates().isClosed()); diff --git a/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java b/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java index 38151ceb..c4517507 100644 --- a/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java +++ b/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java @@ -32,6 +32,7 @@ import io.axoniq.axonserver.grpc.query.QueryResponse; import org.jetbrains.annotations.NotNull; import org.junit.jupiter.api.*; +import org.testcontainers.shaded.org.awaitility.Awaitility; import java.util.ArrayList; import java.util.List; @@ -253,9 +254,10 @@ void testStreamingQueryResultsInError() throws InterruptedException, TimeoutExce ResultStream resultStream = connection2.queryChannel() .query(queryRequest); - QueryResponse response = resultStream.next(); - assertTrue(response.hasErrorMessage()); + QueryResponse response = resultStream.nextIfAvailable(1, TimeUnit.SECONDS); + assertNull(response); + Awaitility.await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> assertTrue(resultStream.isClosed())); } @Test @@ -327,7 +329,7 @@ public void request(long requested) { .setErrorCode("errorCode") .setErrorMessage(ErrorMessage.getDefaultInstance()) .build(); - responseHandler.sendLast(response); + responseHandler.completeWithError(ErrorMessage.newBuilder().setErrorCode("12345").build()); } @Override From bed0ba5f94f805ec0eb1e8410d4a34cf96d83a8b Mon Sep 17 00:00:00 2001 From: Allard Buijze Date: Tue, 28 Oct 2025 21:39:39 +0100 Subject: [PATCH 2/2] Reinstated the original method for backwards compatibility The method is deprecated but will still accept the first result from a stream. Marked old methods and methods using unnecessary parameters as deprecated --- .../connector/query/QueryChannel.java | 21 +- .../query/SubscriptionQueryResult.java | 18 +- .../query/impl/QueryChannelImpl.java | 24 +- .../query/impl/SubscriptionQueryStream.java | 20 +- .../ConnectorRunner_SubscriptionQuery.java | 4 +- .../query/QueryChannelIntegrationTest.java | 47 +- .../QueryChannelLegacyApiIntegrationTest.java | 516 ++++++++++++++++++ .../query/StreamingQueryIntegrationTest.java | 13 +- 8 files changed, 612 insertions(+), 51 deletions(-) create mode 100644 src/test/java/io/axoniq/axonserver/connector/query/QueryChannelLegacyApiIntegrationTest.java diff --git a/src/main/java/io/axoniq/axonserver/connector/query/QueryChannel.java b/src/main/java/io/axoniq/axonserver/connector/query/QueryChannel.java index 186c9fa8..4ad0f94a 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/QueryChannel.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/QueryChannel.java @@ -52,16 +52,33 @@ public interface QueryChannel { * an initial query result followed by a stream of updates. * * @param query the subscription {@link QueryRequest} to send - * @param updateType the type of updates expected from this subscription query * @param bufferSize the number of updates to be buffered by the update result stream * @param fetchSize the number of updates to be consumed from the update stream prior to refilling it * @return the {@link SubscriptionQueryResult} containing the initial result and update stream */ SubscriptionQueryResult subscriptionQuery(QueryRequest query, - SerializedObject updateType, int bufferSize, int fetchSize); + /** + * Sends out a subscription {@code query} to AxonServer for routing to the appropriate handler. Allows for receiving + * an initial query result followed by a stream of updates. + * + * @param query the subscription {@link QueryRequest} to send + * @param updateType the type of updates expected from this subscription query + * @param bufferSize the number of updates to be buffered by the update result stream + * @param fetchSize the number of updates to be consumed from the update stream prior to refilling it + * @return the {@link SubscriptionQueryResult} containing the initial result and update stream + * @deprecated in favor of {@link #subscriptionQuery(QueryRequest, int, int)}. UpdateType is not used. + */ + @Deprecated(since = "2025.2.0") + default SubscriptionQueryResult subscriptionQuery(QueryRequest query, + @SuppressWarnings("unused") SerializedObject updateType, + int bufferSize, + int fetchSize) { + return subscriptionQuery(query, bufferSize, fetchSize); + } + /** * Prepares this {@link QueryChannel} to disconnect, by unsubscribing all registered query handlers. Will wait with * a certain cut off until all acknowledgments of unsubscribing have been received. diff --git a/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java b/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java index 2cebc7c1..41aec004 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java @@ -29,6 +29,22 @@ */ public interface SubscriptionQueryResult { + /** + * Returns a CompletableFuture that completes when the initial result of the query is available. If an error + * occurred while querying, the CompletableFuture completes exceptionally. + *

+ * Invoking this method will send a request for the initial result, if that hasn't been requested before. Subsequent + * invocations of this method will return the same CompletableFuture instance. + *

+ * Note that calling this method will not prevent the {@link #initialResults()} call from retrieving a new set + * of initial results, or vice versa. + * + * @return a CompletableFuture that completes with the initial result of the subscription query + * @deprecated in favor of {@link #initialResults()}, which returns a stream of results. + */ + @Deprecated(since = "2025.2.0", forRemoval = true) + CompletableFuture initialResult(); + /** * Returns a ResultStream that represent the initial result of the subscription query. *

@@ -37,7 +53,7 @@ public interface SubscriptionQueryResult { * * @return a ResultStream that provides the initial result of the subscription query */ - ResultStream initialResult(); + ResultStream initialResults(); /** * Returns the stream of updates to the queried projection. The stream can be read in a blocking and non-blocking diff --git a/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java b/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java index 12575fa2..87acca6b 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java @@ -43,7 +43,6 @@ import io.axoniq.axonserver.grpc.InstructionAck; import io.axoniq.axonserver.grpc.ProcessingInstruction; import io.axoniq.axonserver.grpc.ProcessingKey; -import io.axoniq.axonserver.grpc.SerializedObject; import io.axoniq.axonserver.grpc.control.ClientIdentification; import io.axoniq.axonserver.grpc.query.QueryComplete; import io.axoniq.axonserver.grpc.query.QueryProviderInbound; @@ -396,7 +395,6 @@ public void close() { @Override public SubscriptionQueryResult subscriptionQuery(QueryRequest query, - SerializedObject updateResponseType, int bufferSize, int fetchSize) { logger.trace("Sending subscription query over QueryChannel for context '{}'.", context); @@ -407,8 +405,9 @@ public SubscriptionQueryResult subscriptionQuery(QueryRequest query, finalQuery = query; } String subscriptionId = finalQuery.getMessageIdentifier(); + CompletableFuture initialResultFuture = new CompletableFuture<>(); SubscriptionQueryStream subscriptionStream = new SubscriptionQueryStream( - subscriptionId, + subscriptionId, initialResultFuture, clientIdentification.getClientId(), () -> QueryChannelImpl.this.query(finalQuery), bufferSize, fetchSize @@ -418,13 +417,28 @@ public SubscriptionQueryResult subscriptionQuery(QueryRequest query, SubscriptionQuery subscriptionQuery = SubscriptionQuery.newBuilder() .setQueryRequest(finalQuery) .setSubscriptionIdentifier(subscriptionId) - .setUpdateResponseType(updateResponseType) .build(); upstream.onNext(SubscriptionQueryRequest.newBuilder().setSubscribe(subscriptionQuery).build()); return new SubscriptionQueryResult() { + private final AtomicBoolean initialResultRequested = new AtomicBoolean(); + + @Override + public CompletableFuture initialResult() { + if (!initialResultFuture.isDone() && !initialResultRequested.getAndSet(true)) { + SubscriptionQuery.Builder initialResultRequest = + SubscriptionQuery.newBuilder() + .setQueryRequest(finalQuery) + .setSubscriptionIdentifier(subscriptionId); + upstream.onNext(SubscriptionQueryRequest.newBuilder() + .setGetInitialResult(initialResultRequest) + .build()); + } + return initialResultFuture; + } + @Override - public ResultStream initialResult() { + public ResultStream initialResults() { return subscriptionStream.initialStream(); } diff --git a/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java b/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java index 0bbdfda9..d93e8649 100644 --- a/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java +++ b/src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java @@ -32,6 +32,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.concurrent.CompletableFuture; import java.util.Optional; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; @@ -45,7 +46,9 @@ public class SubscriptionQueryStream extends FlowControlledStream initialResultFuture; private final AbstractBufferedStream updateBuffer; + private final String clientId; private final Supplier> initialResultSupplier; private final AtomicReference> initialResult = new AtomicReference<>(); @@ -61,13 +64,16 @@ public class SubscriptionQueryStream extends FlowControlledStream initialResultFuture, String clientId, Supplier> initialResultSupplier, int bufferSize, int fetchSize) { super(clientId, bufferSize, fetchSize); + this.clientId = clientId; this.initialResultSupplier = initialResultSupplier; this.subscriptionQueryId = subscriptionQueryId; + this.initialResultFuture = initialResultFuture; this.updateBuffer = new SubscriptionQueryUpdateBuffer(clientId, subscriptionQueryId, bufferSize, fetchSize); } @@ -171,12 +177,16 @@ public void onNext(SubscriptionQueryResponse value) { value.getCompleteExceptionally().getClientId() ); updateBuffer.onError(exception); - ResultStream initial = initialResult.get(); - if (initial != null && !initial.isClosed()) { - initial.close(); + if (!initialResultFuture.isDone()) { + initialResultFuture.completeExceptionally(exception); } break; case INITIAL_RESULT: + logger.debug("Received subscription query initial result. Subscription Id: {}. Message Id: {}.", + value.getSubscriptionIdentifier(), + value.getMessageIdentifier()); + initialResultFuture.complete(value.getInitialResult()); + break; default: logger.info("Received unsupported message from SubscriptionQuery. " + "It doesn't declare one of the expected types: {}", @@ -187,6 +197,7 @@ public void onNext(SubscriptionQueryResponse value) { @Override public void onError(Throwable t) { + initialResultFuture.completeExceptionally(t); ResultStream initialResult = this.initialResult.get(); if (initialResult != null && !initialResult.isClosed()) { initialResult.close(); @@ -208,6 +219,9 @@ public void onError(Throwable t) { @Override public void onCompleted() { updateBuffer.onCompleted(); + if (!initialResultFuture.isDone()) { + initialResultFuture.completeExceptionally(new AxonServerException(ErrorCategory.QUERY_DISPATCH_ERROR, "Subscription query has already been completed", clientId)); + } } @Override diff --git a/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java b/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java index b72af194..04f65d21 100644 --- a/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java +++ b/src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java @@ -60,8 +60,8 @@ public static void main(String[] args) throws Exception { .addProcessingInstructions(ProcessingInstruction.newBuilder().setKey(ProcessingKey.NR_OF_RESULTS).setValue(MetaDataValue.newBuilder().setNumberValue(20).build()).build()) .setQuery("java.lang.String") .setPayload(SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello world")).build()).build(), - SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello")).build(), 1000, 100); - ResultStream initialResult = result.initialResult(); + 1000, 100); + ResultStream initialResult = result.initialResults(); initialResult.onAvailable(() -> { QueryResponse r = initialResult.nextIfAvailable(); if (r != null) { diff --git a/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java index 4fc08ffa..6dcf0f16 100644 --- a/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java +++ b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelIntegrationTest.java @@ -104,11 +104,11 @@ void testPermitsAreRenewedForSubscriptionQuery() throws Exception { // Ensure more messages are produced than available permits. for (int i = 0; i < 1000; i++) { - SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); + SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), 100, 10); // the initial result may be requested when the subscription query was already closed. Therefore we accept exceptionally completed results. CompletableFuture completed = new CompletableFuture<>(); results.add(completed); - result.initialResult().onAvailable(() -> { + result.initialResults().onAvailable(() -> { completed.complete(null); result.updates().close(); }); @@ -118,9 +118,9 @@ void testPermitsAreRenewedForSubscriptionQuery() throws Exception { assertDoesNotThrow(() -> result.get(10, TimeUnit.SECONDS)); } - SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); - result.initialResult().nextIfAvailable(5, TimeUnit.SECONDS); - assertFalse(result.initialResult().getError().isPresent()); + SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), 100, 10); + result.initialResults().nextIfAvailable(5, TimeUnit.SECONDS); + assertFalse(result.initialResults().getError().isPresent()); } @Test @@ -182,9 +182,8 @@ void testSubscribedHandlersReconnectAfterConnectionFailure() throws Exception { void testSubscriptionQueryAllowsEmptyMessageId() { QueryChannel queryChannel = connection1.queryChannel(); QueryRequest queryRequest = QueryRequest.newBuilder().build(); - SerializedObject serializedObject = SerializedObject.newBuilder().build(); assertDoesNotThrow( - () -> queryChannel.subscriptionQuery(queryRequest, serializedObject, 5, 1)); + () -> queryChannel.subscriptionQuery(queryRequest, 5, 1)); } @RepeatedTest(10) @@ -242,19 +241,16 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan .setMessageIdentifier("id") .setQuery("testQuery") .build(), - SerializedObject.newBuilder() - .setType("update") - .build(), 100, 10); assertTrueWithin(1, TimeUnit.SECONDS, () -> { // we need to consume staged messages to see the "close" signal do { // just consume - } while (subscriptionQuery.initialResult().nextIfAvailable() != null); - return subscriptionQuery.initialResult().isClosed(); + } while (subscriptionQuery.initialResults().nextIfAvailable() != null); + return subscriptionQuery.initialResults().isClosed(); }); - assertFalse(subscriptionQuery.initialResult().getError().isPresent()); + assertFalse(subscriptionQuery.initialResults().getError().isPresent()); assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); updateHandlerRef.get().sendUpdate(QueryUpdate.newBuilder().setPayload(SerializedObject.newBuilder().setType(String.class.getName()).setData(ByteString.copyFromUtf8("Hello")).build()).build()); @@ -310,10 +306,9 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan .setMessageIdentifier(subscriptionId) .setQuery("testQuery") .build(), - SerializedObject.newBuilder().setType("update").build(), 100, 10); - assertEquals(subscriptionId, subscriptionQuery.initialResult().nextIfAvailable(1, TimeUnit.SECONDS) + assertEquals(subscriptionId, subscriptionQuery.initialResults().nextIfAvailable(1, TimeUnit.SECONDS) .getRequestIdentifier()); assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); logger.info("Sending update"); @@ -358,7 +353,6 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan SubscriptionQueryResult subscriptionQuery = connection2.queryChannel().subscriptionQuery(QueryRequest.newBuilder() .setMessageIdentifier(subscriptionId) .setQuery("testQuery").build(), - SerializedObject.newBuilder().setType("update").build(), 100, 10); assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); @@ -413,20 +407,15 @@ void testReconnectFinishesQueriesInTransit() throws InterruptedException { return null; } }).collect(Collectors.toList()); - actualMessages.forEach(r -> { - assertFalse(r.hasErrorMessage()); - }); - + actualMessages.forEach(r -> assertFalse(r.hasErrorMessage())); } @Test void testDisconnectFinishesQueriesInTransit() throws InterruptedException { Queue> queriesInProgress = new ConcurrentLinkedQueue<>(); QueryChannel queryChannel = connection1.queryChannel(); - queryChannel.registerQueryHandler((command, reply) -> { - CompletableFuture result = new CompletableFuture<>(); - queriesInProgress.add(reply); - }, new QueryDefinition("testQuery", String.class)); + queryChannel.registerQueryHandler((command, reply) -> + queriesInProgress.add(reply), new QueryDefinition("testQuery", String.class)); QueryChannel queryChannel2 = connection2.queryChannel(); @@ -461,20 +450,19 @@ void testDisconnectFinishesQueriesInTransit() throws InterruptedException { } @Test - void testUnsupportedSubscriptionQueryReturnsNoHandler() throws InterruptedException, TimeoutException { + void testUnsupportedSubscriptionQueryReturnsNoHandler() throws InterruptedException { String subscriptionId = UUID.randomUUID().toString(); SubscriptionQueryResult result = connection1.queryChannel().subscriptionQuery(QueryRequest.newBuilder() .setMessageIdentifier(subscriptionId) .setQuery("testQuery").build(), - SerializedObject.newBuilder().setType("update").build(), 100, 10); assertWithin(1, TimeUnit.SECONDS, () -> assertTrue(result.updates().isClosed())); - QueryResponse firstMessage = result.initialResult().nextIfAvailable(1, TimeUnit.SECONDS); + QueryResponse firstMessage = result.initialResults().nextIfAvailable(1, TimeUnit.SECONDS); assertTrue(firstMessage.hasErrorMessage()); } @Test - void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws InterruptedException, ExecutionException, TimeoutException { + void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws InterruptedException, TimeoutException { String subscriptionId = UUID.randomUUID().toString(); connection2.queryChannel().registerQueryHandler(this::mockHandler, new QueryDefinition("testQuery", String.class)) .awaitAck(1, TimeUnit.SECONDS); @@ -482,10 +470,9 @@ void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws Int SubscriptionQueryResult result = connection1.queryChannel().subscriptionQuery(QueryRequest.newBuilder() .setMessageIdentifier(subscriptionId) .setQuery("testQuery").build(), - SerializedObject.newBuilder().setType("update").build(), 100, 10); - result.initialResult().nextIfAvailable(1, TimeUnit.SECONDS); + result.initialResults().nextIfAvailable(1, TimeUnit.SECONDS); assertNull(result.updates().nextIfAvailable()); assertFalse(result.updates().isClosed()); diff --git a/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelLegacyApiIntegrationTest.java b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelLegacyApiIntegrationTest.java new file mode 100644 index 00000000..529c173f --- /dev/null +++ b/src/test/java/io/axoniq/axonserver/connector/query/QueryChannelLegacyApiIntegrationTest.java @@ -0,0 +1,516 @@ +/* + * Copyright (c) 2020-2021. AxonIQ + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.axoniq.axonserver.connector.query; + +import com.google.protobuf.ByteString; +import io.axoniq.axonserver.connector.AbstractAxonServerIntegrationTest; +import io.axoniq.axonserver.connector.AxonServerConnection; +import io.axoniq.axonserver.connector.AxonServerConnectionFactory; +import io.axoniq.axonserver.connector.AxonServerException; +import io.axoniq.axonserver.connector.ErrorCategory; +import io.axoniq.axonserver.connector.Registration; +import io.axoniq.axonserver.connector.ReplyChannel; +import io.axoniq.axonserver.connector.ResultStream; +import io.axoniq.axonserver.connector.impl.ContextConnection; +import io.axoniq.axonserver.connector.query.impl.QueryChannelImpl; +import io.axoniq.axonserver.grpc.SerializedObject; +import io.axoniq.axonserver.grpc.query.QueryRequest; +import io.axoniq.axonserver.grpc.query.QueryResponse; +import io.axoniq.axonserver.grpc.query.QueryUpdate; +import io.axoniq.axonserver.grpc.query.SubscriptionQuery; +import io.grpc.StatusRuntimeException; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.RepeatedTest; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +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.AtomicReference; +import java.util.stream.Collectors; + +import static io.axoniq.axonserver.connector.impl.ObjectUtils.doIfNotNull; +import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertFalseWithin; +import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertTrueWithin; +import static io.axoniq.axonserver.connector.testutils.AssertUtils.assertWithin; +import static org.junit.jupiter.api.Assertions.*; + +/** + * Tests the legacy api for subscription queries, which uses the single initial value request via the subscription query + * call. + */ +@SuppressWarnings("removal") +class QueryChannelLegacyApiIntegrationTest extends AbstractAxonServerIntegrationTest { + + private static final CompletableFuture COMPLETED_FUTURE = CompletableFuture.completedFuture(null); + private static final Logger logger = LoggerFactory.getLogger(QueryChannelIntegrationTest.class); + private AxonServerConnectionFactory connectionFactory1; + private AxonServerConnection connection1; + private AxonServerConnectionFactory connectionFactory2; + private AxonServerConnection connection2; + + @BeforeEach + void setUp() { + connectionFactory1 = AxonServerConnectionFactory.forClient(getClass().getSimpleName() + "_Handler") + .connectTimeout(1500, TimeUnit.MILLISECONDS) + .reconnectInterval(500, TimeUnit.MILLISECONDS) + .routingServers(axonServerAddress) + .queryPermits(100) + .build(); + + connection1 = connectionFactory1.connect("default"); + + connectionFactory2 = AxonServerConnectionFactory.forClient(getClass().getSimpleName() + "_Sender") + .connectTimeout(1500, TimeUnit.MILLISECONDS) + .reconnectInterval(500, TimeUnit.MILLISECONDS) + .routingServers(axonServerAddress) + .queryPermits(100) + .build(); + connection2 = connectionFactory2.connect("default"); + } + + @AfterEach + void tearDown() { + connectionFactory1.shutdown(); + connectionFactory2.shutdown(); + } + + @Test + void testPermitsAreRenewedForSubscriptionQuery() throws Exception { + QueryChannel queryChannel = connection1.queryChannel(); + queryChannel.registerQueryHandler(new SubscriptionQueryHandler(), new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + queryChannel.registerQueryHandler(new SubscriptionQueryHandler(), new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + QueryChannel queryChannel2 = connection2.queryChannel(); + queryChannel2.registerQueryHandler(new SubscriptionQueryHandler(), new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + List> results = new ArrayList<>(); + + // Ensure more messages are produced than available permits. + for (int i = 0; i < 1000; i++) { + SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); + // the initial result may be requested when the subscription query was already closed. Therefore we accept exceptionally completed results. + results.add(result.initialResult().exceptionally(e -> null).whenComplete((r, e) -> result.updates().close())); + } + + for (CompletableFuture result : results) { + assertDoesNotThrow(() -> result.get(10, TimeUnit.SECONDS)); + } + + SubscriptionQueryResult result = queryChannel2.subscriptionQuery(QueryRequest.newBuilder().setMessageIdentifier(UUID.randomUUID().toString()).setQuery("testQuery").build(), SerializedObject.newBuilder().build(), 100, 10); + assertDoesNotThrow(() -> result.initialResult().get(5, TimeUnit.SECONDS)); + } + + @Test + void testUnsubscribedHandlersDoesNotReceiveQueries() throws Exception { + QueryChannel queryChannel = connection1.queryChannel(); + Registration registration = queryChannel.registerQueryHandler(this::mockHandler, new QueryDefinition("testQuery", "testResult")); + + registration.cancel().join(); + + Thread.sleep(100); // An ACK will only confirm receipt of the instruction, not the processing of it. + + ResultStream result = connection2.queryChannel().query(QueryRequest.newBuilder().setQuery("testQuery").build()); + + QueryResponse queryResponse = result.nextIfAvailable(2, TimeUnit.SECONDS); + + assertNotNull(queryResponse); + + assertTrue(queryResponse.hasErrorMessage()); + + axonServerProxy.disable(); + + assertFalseWithin(100, TimeUnit.MILLISECONDS, connection1::isReady); + + axonServerProxy.enable(); + + assertTrueWithin(5, TimeUnit.SECONDS, connection1::isReady); + assertTrueWithin(5, TimeUnit.SECONDS, connection2::isReady); + + ResultStream result2 = connection2.queryChannel().query(QueryRequest.newBuilder().setQuery("testQuery").build()); + QueryResponse actual = result2.nextIfAvailable(1, TimeUnit.SECONDS); + assertNotNull(actual); + assertTrue(actual.hasErrorMessage()); + } + + @Test + void testSubscribedHandlersReconnectAfterConnectionFailure() throws Exception { + QueryChannel queryChannel = connection1.queryChannel(); + queryChannel.registerQueryHandler(this::mockHandler, new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + axonServerProxy.disable(); + + assertTrueWithin(1, TimeUnit.SECONDS, connection1::isConnectionFailed); + + axonServerProxy.enable(); + + assertTrueWithin(3, TimeUnit.SECONDS, connection1::isReady); + + // on some environments, the subscription action may be delayed a little + assertWithin(1500, TimeUnit.MILLISECONDS, () -> { + ResultStream result = connection2.queryChannel().query(QueryRequest.newBuilder().setQuery("testQuery").build()); + QueryResponse queryResponse = result.nextIfAvailable(500, TimeUnit.MILLISECONDS); + assertNotNull(queryResponse); + assertEquals("", queryResponse.getErrorMessage().getMessage()); + }); + } + + @Test + void testSubscriptionQueryAllowsEmptyMessageId() { + QueryChannel queryChannel = connection1.queryChannel(); + QueryRequest queryRequest = QueryRequest.newBuilder().build(); + SerializedObject serializedObject = SerializedObject.newBuilder().build(); + assertDoesNotThrow( + () -> queryChannel.subscriptionQuery(queryRequest, serializedObject, 5, 1)); + } + + @RepeatedTest(10) + void testQueryChannelConsideredConnectedWhenNoHandlersSubscribed() throws IOException { + QueryChannelImpl queryChannel = (QueryChannelImpl) connection1.queryChannel(); + // just to make sure that no attempt was made to connect, since there are no handlers + assertTrue(queryChannel.isReady()); + + // make sure no real connection can be set up + axonServerProxy.disable(); + assertWithin(1, TimeUnit.SECONDS, () -> assertFalse(connection1.isConnected())); + + assertTrue(queryChannel.isReady()); + assertFalse(connection1.isConnected()); + + Registration registration = queryChannel.registerQueryHandler((q, r) -> r.complete(), + new QueryDefinition("test", String.class)); + AxonServerException exception = assertThrows(AxonServerException.class, () -> registration.awaitAck(1, TimeUnit.SECONDS)); + assertEquals(ErrorCategory.INSTRUCTION_ACK_ERROR, exception.getErrorCategory()); + + // because of the attempt to set up a connection, it may need a few milliseconds to discover that's not possible + assertWithin(1, TimeUnit.SECONDS, () -> assertFalse(queryChannel.isReady())); + + axonServerProxy.enable(); + + // verify connection is established + assertWithin(2, TimeUnit.SECONDS, () -> assertTrue(queryChannel.isReady())); + } + + @Test + void testSubscriptionQueryCancelledOnDisconnect() throws Exception { + connection2.controlChannel().enableHeartbeat(100, 100, TimeUnit.MILLISECONDS); + QueryChannel queryChannel = connection1.queryChannel(); + AtomicReference updateHandlerRef = new AtomicReference<>(); + queryChannel.registerQueryHandler( + new QueryHandler() { + @Override + public void handle(QueryRequest query, ReplyChannel responseHandler) { + mockHandler(query, responseHandler); + } + + @Override + public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHandler updateHandler) { + updateHandlerRef.set(updateHandler); + return () -> { + updateHandlerRef.set(null); + return COMPLETED_FUTURE; + }; + } + }, new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + SubscriptionQueryResult subscriptionQuery = connection2.queryChannel() + .subscriptionQuery(QueryRequest.newBuilder() + .setMessageIdentifier("id") + .setQuery("testQuery") + .build(), + SerializedObject.newBuilder() + .setType("update") + .build(), + 100, 10); + + assertTrueWithin(1, TimeUnit.SECONDS, () -> subscriptionQuery.initialResult().isDone()); + assertFalse(subscriptionQuery.initialResult().isCompletedExceptionally()); + assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); + + updateHandlerRef.get().sendUpdate(QueryUpdate.newBuilder().setPayload(SerializedObject.newBuilder().setType(String.class.getName()).setData(ByteString.copyFromUtf8("Hello")).build()).build()); + + assertWithin(2, TimeUnit.SECONDS, () -> assertNotNull(subscriptionQuery.updates().nextIfAvailable())); + + axonServerProxy.disable(); + + assertFalseWithin(1, TimeUnit.SECONDS, connection1::isConnected); + assertFalseWithin(1, TimeUnit.SECONDS, connection2::isConnected); + + assertWithin(1, TimeUnit.SECONDS, () -> { + assertTrue(subscriptionQuery.updates().isClosed()); + assertNull(updateHandlerRef.get()); + }); + axonServerProxy.enable(); + + assertTrueWithin(2, TimeUnit.SECONDS, connection1::isReady); + } + + @RepeatedTest(20) + void testClosingSubscriptionQueryFromSenderStopsUpdateStream() throws InterruptedException, TimeoutException, ExecutionException { + QueryChannel queryChannel = connection1.queryChannel(); + AtomicReference updateHandlerRef = new AtomicReference<>(); + String subscriptionId = UUID.randomUUID().toString(); + queryChannel.registerQueryHandler(new QueryHandler() { + @Override + public void handle(QueryRequest query, ReplyChannel responseHandler) { + logger.info("Handling query"); + mockHandler(query, responseHandler); + } + + @Override + public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHandler updateHandler) { + logger.info("Registering update handler for subscription query"); + if (!subscriptionId.equals(query.getQueryRequest().getMessageIdentifier())) { + logger.warn("Received old subscription query. Ignoring"); + return null; + } + updateHandlerRef.set(updateHandler); + return () -> { + logger.info("Clearing update handler"); + updateHandlerRef.set(null); + return COMPLETED_FUTURE; + }; + } + }, new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + Thread.sleep(100); + + SubscriptionQueryResult subscriptionQuery = connection2.queryChannel().subscriptionQuery(QueryRequest.newBuilder() + .setMessageIdentifier(subscriptionId) + .setQuery("testQuery") + .build(), + SerializedObject.newBuilder().setType("update").build(), + 100, 10); + + assertEquals(subscriptionId, subscriptionQuery.initialResult().get(1, TimeUnit.SECONDS) + .getRequestIdentifier()); + assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); + logger.info("Sending update"); + updateHandlerRef.get().sendUpdate(QueryUpdate.newBuilder().build()); + + assertWithin(2, TimeUnit.SECONDS, () -> assertNotNull(subscriptionQuery.updates().nextIfAvailable())); + + subscriptionQuery.updates().close(); + + assertWithin(1, TimeUnit.SECONDS, () -> { + assertNull(subscriptionQuery.updates().nextIfAvailable()); + assertTrue(subscriptionQuery.updates().isClosed(), "Client side update stream should have been closed"); + assertNull(updateHandlerRef.get(), "Expected updateHandler to have been deregistered"); + }); + } + + @Test + void testClosingSubscriptionQueryFromProviderStopsUpdateStream() throws InterruptedException, TimeoutException { + QueryChannel queryChannel = connection1.queryChannel(); + AtomicReference updateHandlerRef = new AtomicReference<>(); + String subscriptionId = UUID.randomUUID().toString(); + queryChannel.registerQueryHandler(new QueryHandler() { + @Override + public void handle(QueryRequest query, ReplyChannel responseHandler) { + mockHandler(query, responseHandler); + } + + @Override + public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHandler updateHandler) { + if (!subscriptionId.equals(query.getQueryRequest().getMessageIdentifier())) { + return null; + } + updateHandlerRef.set(updateHandler); + return () -> { + updateHandlerRef.set(null); + return COMPLETED_FUTURE; + }; + } + }, new QueryDefinition("testQuery", "testResult")) + .awaitAck(1, TimeUnit.SECONDS); + + SubscriptionQueryResult subscriptionQuery = connection2.queryChannel().subscriptionQuery(QueryRequest.newBuilder() + .setMessageIdentifier(subscriptionId) + .setQuery("testQuery").build(), + SerializedObject.newBuilder().setType("update").build(), + 100, 10); + + assertWithin(1, TimeUnit.SECONDS, () -> assertNotNull(updateHandlerRef.get())); + + updateHandlerRef.get().sendUpdate(QueryUpdate.newBuilder().build()); + updateHandlerRef.get().complete(); + + ResultStream updates = subscriptionQuery.updates(); + assertNotNull(updates.nextIfAvailable(1, TimeUnit.SECONDS)); + + assertNull(updates.nextIfAvailable()); + + assertWithin(1, TimeUnit.SECONDS, () -> { + assertTrue(updates.isClosed(), "Expected client side to be deregistered"); + assertNull(updateHandlerRef.get(), "Expected UpdateHandler to be deregistered"); + }); + } + + @Test + void testReconnectFinishesQueriesInTransit() throws InterruptedException { + Queue> queriesInProgress = new ConcurrentLinkedQueue<>(); + QueryChannel queryChannel = connection1.queryChannel(); + queryChannel.registerQueryHandler((command, reply) -> queriesInProgress.add(reply), + new QueryDefinition("testQuery", String.class)); + + QueryChannel queryChannel2 = connection2.queryChannel(); + + assertWithin(1, TimeUnit.SECONDS, () -> connection1.isReady()); + + Thread.sleep(100); // this is because #5 (https://github.com/AxonIQ/axonserver-connector-java/issues/5) isn't implemented yet + + List> actual = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + actual.add(queryChannel2.query(QueryRequest.newBuilder().setQuery("testQuery").build())); + } + + assertWithin(1, TimeUnit.SECONDS, () -> assertEquals(10, queriesInProgress.size())); + + ((ContextConnection) connection1).getManagedChannel().requestReconnect(); + ((QueryChannelImpl) connection1.queryChannel()).reconnect(); + + assertWithin(1, TimeUnit.SECONDS, () -> assertTrue(connection1.isReady())); + + while (!queriesInProgress.isEmpty()) { + doIfNotNull(queriesInProgress.poll(), r -> r.sendLast(QueryResponse.newBuilder().setPayload(SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Works"))).build())); + } + + List actualMessages = actual.stream().map(r -> { + try { + return r.nextIfAvailable(1, TimeUnit.SECONDS); + } catch (InterruptedException e) { + return null; + } + }).collect(Collectors.toList()); + actualMessages.forEach(r -> { + assertFalse(r.hasErrorMessage()); + }); + + } + + @Test + void testDisconnectFinishesQueriesInTransit() throws InterruptedException { + Queue> queriesInProgress = new ConcurrentLinkedQueue<>(); + QueryChannel queryChannel = connection1.queryChannel(); + queryChannel.registerQueryHandler((command, reply) -> + queriesInProgress.add(reply), new QueryDefinition("testQuery", String.class)); + + QueryChannel queryChannel2 = connection2.queryChannel(); + + assertWithin(1, TimeUnit.SECONDS, () -> connection1.isReady()); + + Thread.sleep(100); // this is because #5 (https://github.com/AxonIQ/axonserver-connector-java/issues/5) isn't implemented yet + + List> actual = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + actual.add(queryChannel2.query(QueryRequest.newBuilder().setQuery("testQuery").build())); + } + + assertWithin(1, TimeUnit.SECONDS, () -> assertEquals(10, queriesInProgress.size())); + + // a call to disconnect waits for disconnection to be completed. That isn't compatible with our test here. + CompletableFuture.runAsync(() -> connection1.disconnect()); + + while (!queriesInProgress.isEmpty()) { + doIfNotNull(queriesInProgress.poll(), r -> r.sendLast(QueryResponse.newBuilder().setPayload(SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Works"))).build())); + } + + List actualMessages = actual.stream().map(r -> { + try { + return r.nextIfAvailable(1, TimeUnit.SECONDS); + } catch (InterruptedException e) { + return fail(e); + } + }).collect(Collectors.toList()); + actualMessages.forEach(r -> assertFalse(r.hasErrorMessage())); + + assertWithin(1, TimeUnit.SECONDS, () -> assertTrue(((ContextConnection) connection1).getManagedChannel().isTerminated())); + } + + @Test + void testUnsupportedSubscriptionQueryReturnsNoHandler() throws InterruptedException, TimeoutException { + String subscriptionId = UUID.randomUUID().toString(); + SubscriptionQueryResult result = connection1.queryChannel().subscriptionQuery(QueryRequest.newBuilder() + .setMessageIdentifier(subscriptionId) + .setQuery("testQuery").build(), + SerializedObject.newBuilder().setType("update").build(), + 100, 10); + assertWithin(1, TimeUnit.SECONDS, () -> assertTrue(result.updates().isClosed())); + try { + result.initialResult().get(1, TimeUnit.SECONDS); + fail("Expected an exception to be reported"); + } catch (ExecutionException e) { + assertInstanceOf(StatusRuntimeException.class, e.getCause()); + assertEquals(ErrorCategory.NO_HANDLER_FOR_QUERY.errorCode(), ((StatusRuntimeException) e.getCause()).getStatus().getDescription()); + } + } + + @Test + void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws InterruptedException, ExecutionException, TimeoutException { + String subscriptionId = UUID.randomUUID().toString(); + connection2.queryChannel().registerQueryHandler(this::mockHandler, new QueryDefinition("testQuery", String.class)) + .awaitAck(1, TimeUnit.SECONDS); + + SubscriptionQueryResult result = connection1.queryChannel().subscriptionQuery(QueryRequest.newBuilder() + .setMessageIdentifier(subscriptionId) + .setQuery("testQuery").build(), + SerializedObject.newBuilder().setType("update").build(), + 100, 10); + + result.initialResult().get(1, TimeUnit.SECONDS); + assertNull(result.updates().nextIfAvailable()); + assertFalse(result.updates().isClosed()); + + result.updates().close(); + + assertTrue(result.updates().isClosed()); + } + + private void mockHandler(QueryRequest query, ReplyChannel responseHandler) { + responseHandler.sendLast(QueryResponse.newBuilder().setRequestIdentifier(query.getMessageIdentifier()).setPayload(query.getPayload()).build()); + } + + private class SubscriptionQueryHandler implements QueryHandler { + + @Override + public void handle(QueryRequest query, ReplyChannel responseHandler) { + mockHandler(query, responseHandler); + } + + @Override + public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHandler updateHandler) { + updateHandler.sendUpdate(QueryUpdate.newBuilder().build()); + return () -> CompletableFuture.completedFuture(null); + } + } +} \ No newline at end of file diff --git a/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java b/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java index c4517507..b2887f3b 100644 --- a/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java +++ b/src/test/java/io/axoniq/axonserver/connector/query/StreamingQueryIntegrationTest.java @@ -32,7 +32,6 @@ import io.axoniq.axonserver.grpc.query.QueryResponse; import org.jetbrains.annotations.NotNull; import org.junit.jupiter.api.*; -import org.testcontainers.shaded.org.awaitility.Awaitility; import java.util.ArrayList; import java.util.List; @@ -46,6 +45,7 @@ import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.*; +import static org.testcontainers.shaded.org.awaitility.Awaitility.await; class StreamingQueryIntegrationTest extends AbstractAxonServerIntegrationTest { @@ -256,8 +256,9 @@ void testStreamingQueryResultsInError() throws InterruptedException, TimeoutExce QueryResponse response = resultStream.nextIfAvailable(1, TimeUnit.SECONDS); - assertNull(response); - Awaitility.await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> assertTrue(resultStream.isClosed())); + assertNotNull(response); + assertTrue(response.hasErrorMessage()); + await().until(resultStream::isClosed); } @Test @@ -329,7 +330,7 @@ public void request(long requested) { .setErrorCode("errorCode") .setErrorMessage(ErrorMessage.getDefaultInstance()) .build(); - responseHandler.completeWithError(ErrorMessage.newBuilder().setErrorCode("12345").build()); + responseHandler.sendLast(response); } @Override @@ -346,10 +347,6 @@ private static class SequencingQueryHandler implements QueryHandler { private final String prefix; private final AtomicLong maxResults; - public SequencingQueryHandler() { - this(Long.MAX_VALUE); - } - public SequencingQueryHandler(long maxResults) { this(maxResults, ""); }