Ensure refresh of permits on subscription query

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

@@ -54,7 +54,7 @@ default void sendLast(T outboundMessage) {
 
     /**
      * Sends an receipt acknowledgement if one hasn't been sent yet for this instruction. If not explicitly sent, it
-     * will be send once the {@link #complete()} method is invoked.
+     * will be sent once the {@link #complete()} method is invoked.
      * <p>
      * If the incoming instruction has no instruction ID, this method does nothing.
      */

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

@@ -134,6 +134,7 @@ private void unsubscribeToQueryUpdates(QueryProviderInbound query, ReplyChannel<
         if (registration != null) {
             registration.forEach(Registration::cancel);
         }
+        result.complete();
     }
 
     private void subscribeToQueryUpdates(QueryProviderInbound query, ReplyChannel<QueryProviderOutbound> result) {
@@ -185,6 +186,7 @@ public void complete() {
                 ).add(registration);
             }
         });
+        result.complete();
     }
 
     @Override
@@ -533,7 +535,7 @@ public void send(QueryResponse response) {
 
                     @Override
                     public void complete() {
-                        result.sendAck();
+                        result.complete();
                     }
 
                     @Override

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020. AxonIQ
+ * 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.
@@ -134,6 +134,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

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021. AxonIQ
+ * 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.
@@ -40,6 +40,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -50,6 +52,7 @@
 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;
@@ -61,7 +64,7 @@
 class QueryChannelIntegrationTest extends AbstractAxonServerIntegrationTest {
 
     private static final CompletableFuture<Void> COMPLETED_FUTURE = CompletableFuture.completedFuture(null);
-
+    private static final Logger logger = LoggerFactory.getLogger(QueryChannelIntegrationTest.class);
     private AxonServerConnectionFactory connectionFactory1;
     private AxonServerConnection connection1;
     private AxonServerConnectionFactory connectionFactory2;
@@ -92,6 +95,35 @@ void tearDown() {
         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<CompletableFuture<?>> 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();
@@ -317,7 +349,7 @@ public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHan
                 };
             }
         }, new QueryDefinition("testQuery", "testResult"))
-            .awaitAck(1, TimeUnit.SECONDS);
+                    .awaitAck(1, TimeUnit.SECONDS);
 
         SubscriptionQueryResult subscriptionQuery = connection2.queryChannel().subscriptionQuery(QueryRequest.newBuilder()
                                                                                                              .setMessageIdentifier(subscriptionId)
@@ -385,4 +417,18 @@ void testSubscriptionQueryReturnsNoUpdatesOnUnsupportedSubscription() throws Int
     private void mockHandler(QueryRequest query, ReplyChannel<QueryResponse> responseHandler) {
         responseHandler.sendLast(QueryResponse.newBuilder().setRequestIdentifier(query.getMessageIdentifier()).setPayload(query.getPayload()).build());
     }
+
+    private class SubscriptionQueryHandler implements QueryHandler {
+
+        @Override
+        public void handle(QueryRequest query, ReplyChannel<QueryResponse> responseHandler) {
+            mockHandler(query, responseHandler);
+        }
+
+        @Override
+        public Registration registerSubscriptionQuery(SubscriptionQuery query, UpdateHandler updateHandler) {
+            updateHandler.sendUpdate(QueryUpdate.newBuilder().build());
+            return () -> CompletableFuture.completedFuture(null);
+        }
+    }
 }