fix: Ensures response end handler is invoked just once

confluentinc · Mar 4, 2022 · 931fb07 · 931fb07
1 parent ea5d6d7
commit 931fb07
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 0 deletions.
diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/api/server/QueryStreamHandler.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/api/server/QueryStreamHandler.java
@@ -37,6 +37,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.kafka.common.utils.Time;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -169,6 +170,11 @@ private void handleQueryPublisher(
     final QueryResponseMetadata metadata;
     Optional<String> completionMessage = Optional.empty();
     Optional<String> limitMessage = Optional.of("Limit Reached");
+    // The end handler can be called twice if the connection is closed by the client.  The
+    // call to response.end() resulting from queryPublisher.close() may result in a second
+    // call to the end handler, which will mess up metrics, so we ensure that this called just
+    // once by keeping track of the calls.
+    final AtomicBoolean endedResponse = new AtomicBoolean(false);
 
     if (queryPublisher.isPullQuery()) {
       metadata = new QueryResponseMetadata(
@@ -180,6 +186,10 @@ private void handleQueryPublisher(
 
       // When response is complete, publisher should be closed
       routingContext.response().endHandler(v -> {
+        if (endedResponse.getAndSet(true)) {
+          log.warn("Connection already closed so just returning");
+          return;
+        }
         queryPublisher.close();
         metricsCallbackHolder.reportMetrics(
             routingContext.response().getStatusCode(),
@@ -195,6 +205,10 @@ private void handleQueryPublisher(
           preparePushProjectionSchema(queryPublisher.geLogicalSchema()));
 
       routingContext.response().endHandler(v -> {
+        if (endedResponse.getAndSet(true)) {
+          log.warn("Connection already closed so just returning");
+          return;
+        }
         queryPublisher.close();
         metricsCallbackHolder.reportMetrics(
             routingContext.response().getStatusCode(),
@@ -215,6 +229,10 @@ private void handleQueryPublisher(
 
       // When response is complete, publisher should be closed and query unregistered
       routingContext.response().endHandler(v -> {
+        if (endedResponse.getAndSet(true)) {
+          log.warn("Connection already closed so just returning");
+          return;
+        }
         query.close();
         metricsCallbackHolder.reportMetrics(
             routingContext.response().getStatusCode(),

diff --git a/ksqldb-rest-app/src/test/java/io/confluent/ksql/api/server/QueryStreamHandlerTest.java b/ksqldb-rest-app/src/test/java/io/confluent/ksql/api/server/QueryStreamHandlerTest.java
@@ -19,6 +19,7 @@
 import static org.hamcrest.Matchers.notNullValue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -155,4 +156,41 @@ public void shouldSucceed_pushQuery() {
     assertThat(subscriber.getValue(), notNullValue());
     verify(pushQueryHolder).close();
   }
+
+  @Test
+  public void shouldSucceed_scalablePushQuery() {
+    // Given:
+    when(queryPublisher.isPullQuery()).thenReturn(false);
+    when(queryPublisher.isScalablePushQuery()).thenReturn(true);
+    final QueryStreamArgs req = new QueryStreamArgs("select * from foo emit changes;",
+        Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
+    givenRequest(req);
+
+    // When:
+    handler.handle(routingContext);
+    endHandler.getValue().handle(null);
+
+    // Then:
+    assertThat(subscriber.getValue(), notNullValue());
+    verify(queryPublisher).close();
+  }
+
+  @Test
+  public void verifyEndHandlerNotCalledTwice_scalablePushQuery() {
+    // Given:
+    when(queryPublisher.isPullQuery()).thenReturn(false);
+    when(queryPublisher.isScalablePushQuery()).thenReturn(true);
+    final QueryStreamArgs req = new QueryStreamArgs("select * from foo emit changes;",
+        Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
+    givenRequest(req);
+
+    // When:
+    handler.handle(routingContext);
+    endHandler.getValue().handle(null);
+    endHandler.getValue().handle(null);
+
+    // Then:
+    assertThat(subscriber.getValue(), notNullValue());
+    verify(queryPublisher, times(1)).close();
+  }
 }