Streaming query support

pom.xml

@@ -49,7 +49,7 @@
         <netty.tcnative.version>2.0.34.Final</netty.tcnative.version>
         <junit.jupiter.version>5.8.1</junit.jupiter.version>
 
-        <axonserver.api.version>4.5</axonserver.api.version>
+        <axonserver.api.version>4.6.0-SNAPSHOT</axonserver.api.version>
 
         <jacoco-maven-plugin.version>0.8.7</jacoco-maven-plugin.version>
     </properties>

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

@@ -155,7 +155,7 @@ protected T take() throws InterruptedException {
      * @return he first entry of this buffer without removing it, or {@code null} if it is empty
      */
     protected T peek() {
-        return validate(buffer.peek(), false);
+        return validate(buffer.peek(), true);
     }
 
     private T validate(T peek, boolean nullOnTerminal) {

src/main/java/io/axoniq/axonserver/connector/query/QueryHandler.java

@@ -18,6 +18,7 @@
 
 import io.axoniq.axonserver.connector.Registration;
 import io.axoniq.axonserver.connector.ReplyChannel;
+import io.axoniq.axonserver.connector.query.impl.NoopFlowControl;
 import io.axoniq.axonserver.grpc.query.QueryRequest;
 import io.axoniq.axonserver.grpc.query.QueryResponse;
 import io.axoniq.axonserver.grpc.query.QueryUpdate;
@@ -40,6 +41,22 @@ public interface QueryHandler {
      */
     void handle(QueryRequest query, ReplyChannel<QueryResponse> responseHandler);
 
+    /**
+     * Handle the given {@code query}, using given {@code responseHandler} to send the response(s). It is flow control
+     * aware - messages should be sent via {@code responseHandler} when requested.
+     * <p>
+     * Note that the query <em>must</em> be completed using {@link ReplyChannel#complete()} or {@link
+     * ReplyChannel#sendLast(Object)}.
+     *
+     * @param query the message representing the query request
+     * @param responseHandler the handler to send responses with
+     * @return a {@link FlowControl} to request more responses and also to cancel sending responses
+     */
+    default FlowControl stream(QueryRequest query, ReplyChannel<QueryResponse> responseHandler) {
+        handle(query, responseHandler);
+        return NoopFlowControl.INSTANCE;
+    }
+
     /**
      * Registers an incoming subscription query request, represented by given {@code query}, using given {@code
      * updateHandler} to send updates when the projection for this query changes.
@@ -56,6 +73,24 @@ default Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHa
         return null;
     }
 
+    /**
+     * Flow control for queries.
+     */
+    interface FlowControl {
+
+        /**
+         * Requests {@code requested} amount of responses to be sent.
+         *
+         * @param requested number of responses to be sent
+         */
+        void request(long requested);
+
+        /**
+         * Completes response sending.
+         */
+        void complete();
+    }
+
     /**
      * Interface describing a stream of updates to a subscription query.
      */

src/main/java/io/axoniq/axonserver/connector/query/impl/MultiFlowControl.java

@@ -0,0 +1,42 @@
+package io.axoniq.axonserver.connector.query.impl;
+
+import io.axoniq.axonserver.connector.Registration;
+import io.axoniq.axonserver.connector.query.QueryHandler;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * An implementation of {@link QueryHandler.FlowControl} that delegates operations to several registered {@link
+ * QueryHandler.FlowControl} implementations.
+ */
+public class MultiFlowControl implements QueryHandler.FlowControl {
+
+    private final List<QueryHandler.FlowControl> delegates = new CopyOnWriteArrayList<>();
+
+    /**
+     * Adds given {@code flowControl} to the list of delegates to be invoked once the {@link QueryHandler.FlowControl}
+     * operations is triggered.
+     *
+     * @param flowControl the {@link QueryHandler.FlowControl} implementation to be registered as a delegate
+     * @return a registration which can be used to remove this {@link QueryHandler.FlowControl} from the list of
+     * delegates
+     */
+    public Registration add(QueryHandler.FlowControl flowControl) {
+        delegates.add(flowControl);
+        return () -> CompletableFuture.runAsync(() -> delegates.remove(flowControl));
+    }
+
+    @Override
+    public void request(long requested) {
+        delegates.parallelStream()
+                 .forEach(flowControl -> flowControl.request(requested));
+    }
+
+    @Override
+    public void complete() {
+        delegates.parallelStream()
+                 .forEach(QueryHandler.FlowControl::complete);
+    }
+}

src/main/java/io/axoniq/axonserver/connector/query/impl/NoopFlowControl.java

@@ -0,0 +1,24 @@
+package io.axoniq.axonserver.connector.query.impl;
+
+import io.axoniq.axonserver.connector.query.QueryHandler;
+
+/**
+ * NOOP implementation of {@link QueryHandler.FlowControl}.
+ */
+public enum NoopFlowControl implements QueryHandler.FlowControl {
+
+    /**
+     * Singleton instance of {@link NoopFlowControl}.
+     */
+    INSTANCE;
+
+    @Override
+    public void request(long requested) {
+        // noop
+    }
+
+    @Override
+    public void complete() {
+        // noop
+    }
+}

src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java

@@ -92,7 +92,7 @@ public class QueryChannelImpl extends AbstractAxonServerChannel<QueryProviderOut
     private final Object queryHandlerMonitor = new Object();
     private final Map<String, Set<Registration>> subscriptionQueries = new ConcurrentHashMap<>();
     private final QueryServiceGrpc.QueryServiceStub queryServiceStub;
-    private final Set<CompletableFuture<?>> queriesInProgress = ConcurrentHashMap.newKeySet();
+    private final Map<String, QueryInProgress> queriesInProgress = new ConcurrentHashMap<>();
     private final AtomicBoolean subscriptionsCompleted = new AtomicBoolean(false);
 
     /**
@@ -122,6 +122,8 @@ public QueryChannelImpl(ClientIdentification clientIdentification,
         instructionHandlers.put(SubscriptionQueryRequest.RequestCase.GET_INITIAL_RESULT, this::getInitialResult);
         instructionHandlers.put(SubscriptionQueryRequest.RequestCase.SUBSCRIBE, this::subscribeToQueryUpdates);
         instructionHandlers.put(SubscriptionQueryRequest.RequestCase.UNSUBSCRIBE, this::unsubscribeToQueryUpdates);
+        instructionHandlers.put(QueryProviderInbound.RequestCase.QUERY_COMPLETE, this::completeStreamingQuery);
+        instructionHandlers.put(QueryProviderInbound.RequestCase.QUERY_FLOW_CONTROL, this::flowControlQuery);
         queryServiceStub = QueryServiceGrpc.newStub(channel);
     }
 
@@ -191,6 +193,18 @@ public void complete() {
         result.complete();
     }
 
+    private void completeStreamingQuery(QueryProviderInbound complete, ReplyChannel<QueryProviderOutbound> result) {
+        queriesInProgress.getOrDefault(complete.getQueryComplete().getRequestId(), QueryInProgress.noop())
+                         .complete();
+        result.complete();
+    }
+
+    private void flowControlQuery(QueryProviderInbound flowControl, ReplyChannel<QueryProviderOutbound> result) {
+        queriesInProgress.getOrDefault(flowControl.getQueryFlowControl().getRequestId(), QueryInProgress.noop())
+                         .request(flowControl.getQueryFlowControl().getPermits());
+        result.complete();
+    }
+
     @Override
     public void connect() {
         if (!queryHandlers.isEmpty()) {
@@ -339,7 +353,7 @@ protected QueryResponse terminalMessage() {
 
             @Override
             public void close() {
-                // this is a one-way stream. No need to close it.
+                outboundStream().cancel("Client cancelled the stream.", null);
             }
         };
         if ("".equals(query.getMessageIdentifier())) {
@@ -410,7 +424,9 @@ public synchronized void disconnect() {
             if (!queriesInProgress.isEmpty()) {
                 logger.info("Disconnect requested. Waiting for {} queries to be completed", queriesInProgress.size());
             }
-            return CompletableFuture.allOf(queriesInProgress.stream()
+            return CompletableFuture.allOf(queriesInProgress.values()
+                                                            .stream()
+                                                            .map(QueryInProgress::completeAsync)
                                                             .reduce(CompletableFuture::allOf)
                                                             .orElseGet(() -> CompletableFuture.completedFuture(null)));
         }).thenAccept(previousStream -> doIfNotNull(previousOutbound, StreamObserver::onCompleted));
@@ -451,10 +467,12 @@ private void doHandleQuery(QueryProviderInbound query, ReplyChannel<QueryRespons
     }
 
     private void doHandleQuery(QueryRequest query, ReplyChannel<QueryResponse> responseChannel) {
-        CompletableFuture<?> completionHandle = new CompletableFuture<>();
-        queriesInProgress.add(completionHandle);
-        completionHandle.whenComplete((r, e) -> queriesInProgress.remove(completionHandle));
-        ReplyChannel<QueryResponse> responseHandler = new CloseAwareReplyChannelAdapter(responseChannel, () -> completionHandle.complete(null));
+        MultiFlowControl multiFlowControl = new MultiFlowControl();
+        Runnable removeQuery = () -> queriesInProgress.remove(query.getMessageIdentifier());
+        QueryInProgress queryInProgress = new QueryInProgress(removeQuery, multiFlowControl);
+        queriesInProgress.putIfAbsent(query.getMessageIdentifier(), queryInProgress);
+        ReplyChannel<QueryResponse> responseHandler = new CloseAwareReplyChannelAdapter(responseChannel,
+                                                                                        queryInProgress::complete);
         Set<QueryHandler> handlers = queryHandlers.getOrDefault(query.getQuery(), Collections.emptySet());
         if (handlers.isEmpty()) {
             responseHandler.sendNack();
@@ -470,47 +488,50 @@ private void doHandleQuery(QueryRequest query, ReplyChannel<QueryResponse> respo
         responseHandler.sendAck();
 
         AtomicInteger completeCounter = new AtomicInteger(handlers.size());
-        handlers.forEach(queryHandler -> queryHandler.handle(query, new ReplyChannel<QueryResponse>() {
-            @Override
-            public void send(QueryResponse response) {
-                if (!query.getMessageIdentifier().equals(response.getRequestIdentifier())) {
-                    logger.debug("RequestIdentifier not properly set, modifying message");
-                    QueryResponse newResponse = response.toBuilder()
-                                                        .setRequestIdentifier(query.getMessageIdentifier())
-                                                        .build();
-                    responseHandler.send(newResponse);
-                } else {
-                    responseHandler.send(response);
+        handlers.forEach(queryHandler -> {
+            ReplyChannel<QueryResponse> replyChannel = new ReplyChannel<QueryResponse>() {
+                @Override
+                public void send(QueryResponse response) {
+                    if (!query.getMessageIdentifier().equals(response.getRequestIdentifier())) {
+                        logger.debug("RequestIdentifier not properly set, modifying message");
+                        QueryResponse newResponse = response.toBuilder()
+                                                            .setRequestIdentifier(query.getMessageIdentifier())
+                                                            .build();
+                        responseHandler.send(newResponse);
+                    } else {
+                        responseHandler.send(response);
+                    }
                 }
-            }
 
-            @Override
-            public void complete() {
-                if (completeCounter.decrementAndGet() == 0) {
-                    responseHandler.complete();
+                @Override
+                public void complete() {
+                    if (completeCounter.decrementAndGet() == 0) {
+                        responseHandler.complete();
+                    }
                 }
-            }
 
-            @Override
-            public void completeWithError(ErrorMessage errorMessage) {
-                responseHandler.completeWithError(errorMessage);
-            }
+                @Override
+                public void completeWithError(ErrorMessage errorMessage) {
+                    responseHandler.completeWithError(errorMessage);
+                }
 
-            @Override
-            public void completeWithError(ErrorCategory errorCategory, String message) {
-                responseHandler.completeWithError(errorCategory, message);
-            }
+                @Override
+                public void completeWithError(ErrorCategory errorCategory, String message) {
+                    responseHandler.completeWithError(errorCategory, message);
+                }
 
-            @Override
-            public void sendNack(ErrorMessage errorMessage) {
-                responseHandler.sendNack(errorMessage);
-            }
+                @Override
+                public void sendNack(ErrorMessage errorMessage) {
+                    responseHandler.sendNack(errorMessage);
+                }
 
-            @Override
-            public void sendAck() {
-                responseHandler.sendAck();
-            }
-        }));
+                @Override
+                public void sendAck() {
+                    responseHandler.sendAck();
+                }
+            };
+            multiFlowControl.add(queryHandler.stream(query, replyChannel));
+        });
     }
 
     private void handleQuery(QueryProviderInbound inbound, ReplyChannel<QueryProviderOutbound> result) {
@@ -601,6 +622,37 @@ public void sendAck() {
         );
     }
 
+    private static class QueryInProgress {
+
+        private final CompletableFuture<?> closeHandler;
+        private final QueryHandler.FlowControl flowControl;
+
+        public static QueryInProgress noop() {
+            return new QueryInProgress(() -> { }, NoopFlowControl.INSTANCE);
+        }
+
+        public QueryInProgress(Runnable closeHandler, QueryHandler.FlowControl flowControl) {
+            this.closeHandler = new CompletableFuture<>();
+            this.flowControl = flowControl;
+            this.closeHandler.whenComplete((r, e) -> {
+                flowControl.complete();
+                closeHandler.run();
+            });
+        }
+
+        public CompletableFuture<?> completeAsync() {
+            return closeHandler;
+        }
+
+        public void complete() {
+            completeAsync().complete(null);
+        }
+
+        public void request(long requested) {
+            flowControl.request(requested);
+        }
+    }
+
     private static class CloseAwareReplyChannelAdapter implements ReplyChannel<QueryResponse> {
         private final ReplyChannel<QueryResponse> delegate;
         private final Runnable onClose;

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

@@ -382,10 +382,8 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan
     void testReconnectFinishesQueriesInTransit() throws InterruptedException {
         Queue<ReplyChannel<QueryResponse>> queriesInProgress = new ConcurrentLinkedQueue<>();
         QueryChannel queryChannel = connection1.queryChannel();
-        queryChannel.registerQueryHandler((command, reply) -> {
-            CompletableFuture<QueryResponse> 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();