Make HttpRequestHandlerImpl.handle() async when possible (fixes #259)

Author: olim7t

graphql-java-kickstart/src/main/java/graphql/kickstart/execution/GraphQLInvoker.java

@@ -8,6 +8,7 @@
 import graphql.kickstart.execution.input.GraphQLBatchedInvocationInput;
 import graphql.kickstart.execution.input.GraphQLInvocationInput;
 import graphql.kickstart.execution.input.GraphQLSingleInvocationInput;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import lombok.AllArgsConstructor;
@@ -27,25 +28,36 @@ public CompletableFuture<ExecutionResult> executeAsync(GraphQLSingleInvocationIn
   }
 
   public GraphQLQueryResult query(GraphQLInvocationInput invocationInput) {
+    return queryAsync(invocationInput).join();
+  }
+
+  public CompletableFuture<GraphQLQueryResult> queryAsync(GraphQLInvocationInput invocationInput) {
     if (invocationInput instanceof GraphQLSingleInvocationInput) {
-      return GraphQLQueryResult.create(query((GraphQLSingleInvocationInput) invocationInput));
+      return executeAsync((GraphQLSingleInvocationInput)invocationInput).thenApply(GraphQLQueryResult::create);
     }
     GraphQLBatchedInvocationInput batchedInvocationInput = (GraphQLBatchedInvocationInput) invocationInput;
-    return GraphQLQueryResult.create(query(batchedInvocationInput));
+    return executeAsync(batchedInvocationInput).thenApply(GraphQLQueryResult::create);
   }
 
-  private ExecutionResult query(GraphQLSingleInvocationInput singleInvocationInput) {
-    return executeAsync(singleInvocationInput).join();
+  private CompletableFuture<List<ExecutionResult>> executeAsync(GraphQLBatchedInvocationInput batchedInvocationInput) {
+    GraphQL graphQL = batchedDataLoaderGraphQLBuilder.newGraphQL(batchedInvocationInput, graphQLBuilder);
+    return sequence(
+        batchedInvocationInput.getExecutionInputs().stream()
+            .map(executionInput -> proxy.executeAsync(graphQL, executionInput))
+            .collect(toList()));
   }
 
-  private List<ExecutionResult> query(GraphQLBatchedInvocationInput batchedInvocationInput) {
-    GraphQL graphQL = batchedDataLoaderGraphQLBuilder.newGraphQL(batchedInvocationInput, graphQLBuilder);
-    return batchedInvocationInput.getExecutionInputs().stream()
-        .map(executionInput -> proxy.executeAsync(graphQL, executionInput))
-        .collect(toList())
-        .stream()
-        .map(CompletableFuture::join)
-        .collect(toList());
+  @SuppressWarnings({"unchecked", "rawtypes"})
+  private <T> CompletableFuture<List<T>> sequence(List<CompletableFuture<T>> futures) {
+    CompletableFuture[] futuresArray = futures.toArray(new CompletableFuture[0]);
+    return CompletableFuture.allOf(futuresArray).thenApply(aVoid -> {
+      List<T> result =  new ArrayList<>();
+      for (CompletableFuture future : futuresArray) {
+        assert future.isDone(); // per the API contract of allOf()
+        result.add((T) future.join());
+      }
+      return result;
+    });
   }
 }
 

graphql-java-servlet/src/main/java/graphql/kickstart/servlet/HttpRequestHandlerImpl.java

@@ -9,6 +9,9 @@
 import graphql.kickstart.execution.input.GraphQLInvocationInput;
 import graphql.kickstart.execution.input.GraphQLSingleInvocationInput;
 import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.concurrent.CompletableFuture;
+import javax.servlet.AsyncContext;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
@@ -48,37 +51,60 @@ public void handle(HttpServletRequest request, HttpServletResponse response) thr
 
   private void execute(GraphQLInvocationInput invocationInput, HttpServletRequest request,
       HttpServletResponse response) {
-    try {
-      GraphQLQueryResult queryResult = invoke(invocationInput, request, response);
+    if (request.isAsyncSupported()) {
+      AsyncContext asyncContext = request.startAsync(request, response);
+      asyncContext.setTimeout(configuration.getAsyncTimeout());
+      invoke(invocationInput, request, response)
+          .thenAccept(result -> writeResultResponse(result, request, response))
+          .exceptionally(t -> writeErrorResponse(t, response))
+          .thenAccept(aVoid -> asyncContext.complete());
+    } else {
+      try {
+        GraphQLQueryResult result = invoke(invocationInput, request, response).join();
+        writeResultResponse(result, request, response);
+      } catch (Throwable t) {
+        writeErrorResponse(t, response);
+      }
+    }
+  }
 
-      QueryResponseWriter queryResponseWriter = QueryResponseWriter.createWriter(queryResult, configuration.getObjectMapper(),
-          configuration.getSubscriptionTimeout());
+  private void writeResultResponse(GraphQLQueryResult queryResult, HttpServletRequest request,
+      HttpServletResponse response) {
+    QueryResponseWriter queryResponseWriter = QueryResponseWriter.createWriter(queryResult, configuration.getObjectMapper(),
+        configuration.getSubscriptionTimeout());
+    try {
       queryResponseWriter.write(request, response);
-    } catch (Throwable t) {
-      response.setStatus(STATUS_BAD_REQUEST);
-      log.info("Bad GET request: path was not \"/schema.json\" or no query variable named \"query\" given");
-      log.debug("Possibly due to exception: ", t);
+    } catch (IOException e) {
+      // will be processed by the exceptionally() call below
+      throw new UncheckedIOException(e);
     }
   }
 
-  private GraphQLQueryResult invoke(GraphQLInvocationInput invocationInput, HttpServletRequest request,
+  private Void writeErrorResponse(Throwable t, HttpServletResponse response) {
+    response.setStatus(STATUS_BAD_REQUEST);
+    log.info("Bad GET request: path was not \"/schema.json\" or no query variable named \"query\" given");
+    log.debug("Possibly due to exception: ", t);
+    return null;
+  }
+
+  private CompletableFuture<GraphQLQueryResult> invoke(GraphQLInvocationInput invocationInput, HttpServletRequest request,
       HttpServletResponse response) {
     if (invocationInput instanceof GraphQLSingleInvocationInput) {
-      return graphQLInvoker.query(invocationInput);
+      return graphQLInvoker.queryAsync(invocationInput);
     }
     return invokeBatched((GraphQLBatchedInvocationInput) invocationInput, request, response);
   }
 
-  private GraphQLQueryResult invokeBatched(GraphQLBatchedInvocationInput batchedInvocationInput,
+  private CompletableFuture<GraphQLQueryResult> invokeBatched(GraphQLBatchedInvocationInput batchedInvocationInput,
       HttpServletRequest request,
       HttpServletResponse response) {
     BatchInputPreProcessor preprocessor = configuration.getBatchInputPreProcessor();
     BatchInputPreProcessResult result = preprocessor.preProcessBatch(batchedInvocationInput, request, response);
     if (result.isExecutable()) {
-      return graphQLInvoker.query(result.getBatchedInvocationInput());
+      return graphQLInvoker.queryAsync(result.getBatchedInvocationInput());
     }
 
-    return GraphQLQueryResult.createError(result.getStatusCode(), result.getStatusMessage());
+    return CompletableFuture.completedFuture(GraphQLQueryResult.createError(result.getStatusCode(), result.getStatusMessage()));
   }
 
 }