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/QueryChannel.java

@@ -23,6 +23,7 @@
 import io.axoniq.axonserver.grpc.query.QueryResponse;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 
 /**
  * Communication channel with AxonServer for Query related interactions.
@@ -71,4 +72,22 @@ SubscriptionQueryResult subscriptionQuery(QueryRequest query,
      * @return a {@link CompletableFuture} of {@link Void} to react when all query handlers have been unsubscribed
      */
     CompletableFuture<Void> prepareDisconnect();
+
+    /**
+     * Registers listener that is going to be invoked when query with given {@code queryId} is completed. Once the query
+     * is completed, listener will be removed.
+     *
+     * @param queryId  identifier of the query we are interested when is completed
+     * @param listener to be invoked when the query is completed
+     */
+    Registration registerQueryCompleteListener(String queryId, Runnable listener);
+
+    /**
+     * Registers consumer to accepts number of permits to be invoked when the query {@code queryId} consumer requests more updates.
+     * Once the query is completed, listener will be removed.
+     *
+     * @param queryId  identifier of the query we are interested when is completed
+     * @param consumer that accepts number of permits to be invoked when the query consumer requests more updates
+     */
+    Registration registerQueryFlowControlListener(String queryId, Consumer<Long> consumer);
 }

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

@@ -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_CONTROLL, this::flowControlQuery);
         queryServiceStub = QueryServiceGrpc.newStub(channel);
     }
 
@@ -191,6 +193,22 @@ public void complete() {
         result.complete();
     }
 
+    private void completeStreamingQuery(QueryProviderInbound complete, ReplyChannel<QueryProviderOutbound> result) {
+        queryQueryFlowControlListener.remove(complete.getQueryComplete().getRequestId());
+        Runnable listener = queryCompleteListeners.remove(complete.getQueryComplete().getRequestId());
+        if (listener != null) {
+            listener.run();
+        }
+        result.complete();
+    }
+
+    private void flowControlQuery(QueryProviderInbound flowControl, ReplyChannel<QueryProviderOutbound> queryProviderOutboundReplyChannel) {
+        queryQueryFlowControlListener.getOrDefault(flowControl.getQueryFlowControll().getRequestId(),
+                        l -> logger.debug("Received flow control message for unknown request id {}", flowControl.getQueryFlowControll().getRequestId()))
+                .accept(flowControl.getQueryFlowControll().getPermits());
+        queryProviderOutboundReplyChannel.complete();
+    }
+
     @Override
     public void connect() {
         if (!queryHandlers.isEmpty()) {
@@ -339,7 +357,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())) {
@@ -437,6 +455,30 @@ public CompletableFuture<Void> prepareDisconnect() {
         return future;
     }
 
+    private final Map<String, Runnable> queryCompleteListeners = new ConcurrentHashMap<>();
+    private final Map<String, Consumer<Long>> queryQueryFlowControlListener = new ConcurrentHashMap<>();
+
+
+    @Override
+    public Registration registerQueryCompleteListener(String queryId, Runnable listener) {
+        queryCompleteListeners.put(queryId, listener);
+
+        return () -> {
+            queryCompleteListeners.remove(queryId);
+            return CompletableFuture.completedFuture(null);
+        };
+    }
+
+    @Override
+    public Registration registerQueryFlowControlListener(String queryId, Consumer<Long> consumer) {
+        queryQueryFlowControlListener.put(queryId, consumer);
+
+        return () -> {
+            queryQueryFlowControlListener.remove(queryId);
+            return CompletableFuture.completedFuture(null);
+        };
+    }
+
     private void cancelAllSubscriptionQueries() {
         subscriptionQueries.forEach((k, v) -> subscriptionQueries.remove(k).forEach(Registration::cancel));
     }