fix(#469): allow graphql-java future to be canceled

Author: oliemansm

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

@@ -0,0 +1,26 @@
+package graphql.kickstart.execution;
+
+import graphql.ExecutionResult;
+import graphql.kickstart.execution.input.GraphQLInvocationInput;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+class FutureBatchedExecutionResult implements FutureExecutionResult {
+
+  @Getter
+  private final GraphQLInvocationInput invocationInput;
+  private final CompletableFuture<List<ExecutionResult>> batched;
+
+  @Override
+  public CompletableFuture<GraphQLQueryResult> thenApplyQueryResult() {
+    return batched.thenApply(GraphQLQueryResult::create);
+  }
+
+  @Override
+  public void cancel() {
+    batched.cancel(true);
+  }
+}

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

@@ -0,0 +1,26 @@
+package graphql.kickstart.execution;
+
+import graphql.kickstart.execution.input.GraphQLInvocationInput;
+import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+class FutureErrorExecutionResult implements FutureExecutionResult {
+
+  private final GraphQLErrorQueryResult errorQueryResult;
+
+  @Override
+  public CompletableFuture<GraphQLQueryResult> thenApplyQueryResult() {
+    return CompletableFuture.completedFuture(errorQueryResult);
+  }
+
+  @Override
+  public GraphQLInvocationInput getInvocationInput() {
+    return null;
+  }
+
+  @Override
+  public void cancel() {
+    // nothing to do
+  }
+}

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

@@ -0,0 +1,29 @@
+package graphql.kickstart.execution;
+
+import graphql.ExecutionResult;
+import graphql.kickstart.execution.input.GraphQLInvocationInput;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+public interface FutureExecutionResult {
+
+  static FutureExecutionResult single(
+      GraphQLInvocationInput invocationInput, CompletableFuture<ExecutionResult> single) {
+    return new FutureSingleExecutionResult(invocationInput, single);
+  }
+
+  static FutureExecutionResult batched(
+      GraphQLInvocationInput invocationInput, CompletableFuture<List<ExecutionResult>> batched) {
+    return new FutureBatchedExecutionResult(invocationInput, batched);
+  }
+
+  static FutureExecutionResult error(GraphQLErrorQueryResult result) {
+    return new FutureErrorExecutionResult(result);
+  }
+
+  CompletableFuture<GraphQLQueryResult> thenApplyQueryResult();
+
+  GraphQLInvocationInput getInvocationInput();
+
+  void cancel();
+}

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

@@ -0,0 +1,25 @@
+package graphql.kickstart.execution;
+
+import graphql.ExecutionResult;
+import graphql.kickstart.execution.input.GraphQLInvocationInput;
+import java.util.concurrent.CompletableFuture;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+class FutureSingleExecutionResult implements FutureExecutionResult {
+
+  @Getter
+  private final GraphQLInvocationInput invocationInput;
+  private final CompletableFuture<ExecutionResult> single;
+
+  @Override
+  public CompletableFuture<GraphQLQueryResult> thenApplyQueryResult() {
+    return single.thenApply(GraphQLQueryResult::create);
+  }
+
+  @Override
+  public void cancel() {
+    single.cancel(true);
+  }
+}

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

@@ -22,6 +22,15 @@ public class GraphQLInvoker {
   private final BatchedDataLoaderGraphQLBuilder batchedDataLoaderGraphQLBuilder;
   private final GraphQLInvokerProxy proxy = GraphQL::executeAsync;
 
+  public FutureExecutionResult execute(GraphQLInvocationInput invocationInput) {
+    if (invocationInput instanceof GraphQLSingleInvocationInput) {
+      return FutureExecutionResult.single(
+          invocationInput, executeAsync((GraphQLSingleInvocationInput) invocationInput));
+    }
+    return FutureExecutionResult.batched(
+        invocationInput, executeAsync((GraphQLBatchedInvocationInput) invocationInput));
+  }
+
   public CompletableFuture<ExecutionResult> executeAsync(
       GraphQLSingleInvocationInput invocationInput) {
     GraphQL graphQL = graphQLBuilder.build(invocationInput.getSchema());

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

@@ -2,6 +2,7 @@
 
 import graphql.ExecutionResult;
 import graphql.ExecutionResultImpl;
+import graphql.kickstart.execution.FutureExecutionResult;
 import graphql.kickstart.execution.GraphQLInvoker;
 import graphql.kickstart.execution.GraphQLQueryResult;
 import graphql.kickstart.execution.error.GenericGraphQLError;
@@ -13,6 +14,7 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.Optional;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicReference;
@@ -40,7 +42,7 @@ public void execute(
     if (request.isAsyncSupported()) {
       invokeAndHandleAsync(invocationInput, request, response, listenerHandler);
     } else {
-      invokeAndHandle(invocationInput, request, response, listenerHandler).join();
+      handle(invoke(invocationInput, request, response), request, response, listenerHandler).join();
     }
   }
 
@@ -54,10 +56,10 @@ private void invokeAndHandleAsync(
             ? request.getAsyncContext()
             : request.startAsync(request, response);
     asyncContext.setTimeout(configuration.getAsyncTimeout());
-    AtomicReference<CompletableFuture<Void>> futureHolder = new AtomicReference<>();
+    AtomicReference<FutureExecutionResult> futureHolder = new AtomicReference<>();
     AsyncTimeoutListener timeoutListener =
         event -> {
-          Optional.ofNullable(futureHolder.get()).ifPresent(it -> it.cancel(true));
+          Optional.ofNullable(futureHolder.get()).ifPresent(FutureExecutionResult::cancel);
           writeResultResponse(
               invocationInput,
               GraphQLQueryResult.create(
@@ -68,22 +70,27 @@ private void invokeAndHandleAsync(
         };
     asyncContext.addListener(timeoutListener);
     asyncContext.start(
-        () ->
-            futureHolder.set(
-                invokeAndHandle(invocationInput, request, response, listenerHandler)
-                    .thenAccept(aVoid -> asyncContext.complete())));
+        () -> {
+          FutureExecutionResult futureResult = invoke(invocationInput, request, response);
+          futureHolder.set(futureResult);
+          handle(futureResult, request, response, listenerHandler);
+        });
   }
 
-  private CompletableFuture<Void> invokeAndHandle(
-      GraphQLInvocationInput invocationInput,
+  private CompletableFuture<Void> handle(
+      FutureExecutionResult futureResult,
       HttpServletRequest request,
       HttpServletResponse response,
       ListenerHandler listenerHandler) {
-    return invoke(invocationInput, request, response)
-        .thenAccept(it -> writeResultResponse(invocationInput, it, request, response))
+    return futureResult
+        .thenApplyQueryResult()
+        .thenAccept(
+            it -> writeResultResponse(futureResult.getInvocationInput(), it, request, response))
         .thenAccept(it -> listenerHandler.onSuccess())
         .exceptionally(
-            t -> writeErrorResponse(invocationInput, request, response, listenerHandler, t))
+            t ->
+                writeErrorResponse(
+                    futureResult.getInvocationInput(), request, response, listenerHandler, t))
         .thenAccept(it -> listenerHandler.onFinally());
   }
 
@@ -106,22 +113,32 @@ private Void writeErrorResponse(
       HttpServletResponse response,
       ListenerHandler listenerHandler,
       Throwable t) {
+    Throwable cause = getCause(t);
     if (!response.isCommitted()) {
       writeResultResponse(
-          invocationInput, GraphQLQueryResult.create(toErrorResult(t)), request, response);
-      listenerHandler.onError(t);
+          invocationInput, GraphQLQueryResult.create(toErrorResult(cause)), request, response);
+      listenerHandler.onError(cause);
     } else {
       log.warn(
           "Cannot write GraphQL response, because the HTTP response is already committed. It most likely timed out.");
     }
     return null;
   }
 
+  private Throwable getCause(Throwable t) {
+    return t instanceof CompletionException && t.getCause() != null ? t.getCause() : t;
+  }
+
   private ExecutionResult toErrorResult(Throwable t) {
     String message =
-        t instanceof CompletionException && t.getCause() != null
-            ? t.getCause().getMessage()
+        t instanceof CancellationException
+            ? "Execution canceled because timeout of "
+                + configuration.getAsyncTimeout()
+                + " millis was reached"
             : t.getMessage();
+    if (message == null) {
+      message = "Unexpected error occurred";
+    }
     return new ExecutionResultImpl(new GenericGraphQLError(message));
   }
 
@@ -130,28 +147,28 @@ protected QueryResponseWriter createWriter(
     return queryResponseWriterFactory.createWriter(invocationInput, queryResult, configuration);
   }
 
-  private CompletableFuture<GraphQLQueryResult> invoke(
+  private FutureExecutionResult invoke(
       GraphQLInvocationInput invocationInput,
       HttpServletRequest request,
       HttpServletResponse response) {
     if (invocationInput instanceof GraphQLSingleInvocationInput) {
-      return graphQLInvoker.queryAsync(invocationInput);
+      return graphQLInvoker.execute(invocationInput);
     }
     return invokeBatched((GraphQLBatchedInvocationInput) invocationInput, request, response);
   }
 
-  private CompletableFuture<GraphQLQueryResult> invokeBatched(
+  private FutureExecutionResult invokeBatched(
       GraphQLBatchedInvocationInput batchedInvocationInput,
       HttpServletRequest request,
       HttpServletResponse response) {
     BatchInputPreProcessor preprocessor = configuration.getBatchInputPreProcessor();
     BatchInputPreProcessResult result =
         preprocessor.preProcessBatch(batchedInvocationInput, request, response);
     if (result.isExecutable()) {
-      return graphQLInvoker.queryAsync(result.getBatchedInvocationInput());
+      return graphQLInvoker.execute(result.getBatchedInvocationInput());
     }
 
-    return CompletableFuture.completedFuture(
+    return FutureExecutionResult.error(
         GraphQLQueryResult.createError(result.getStatusCode(), result.getStatusMessage()));
   }
 }

graphql-java-servlet/src/test/groovy/graphql/kickstart/servlet/cache/CachingHttpRequestInvokerTest.groovy

@@ -1,6 +1,7 @@
 package graphql.kickstart.servlet.cache
 
 import graphql.ExecutionResult
+import graphql.kickstart.execution.FutureExecutionResult
 import graphql.kickstart.execution.GraphQLInvoker
 import graphql.kickstart.execution.GraphQLObjectMapper
 import graphql.kickstart.execution.GraphQLQueryResult
@@ -39,6 +40,7 @@ class CachingHttpRequestInvokerTest extends Specification {
     graphqlInvoker = Mock(GraphQLInvoker)
     graphqlObjectMapper = Mock(GraphQLObjectMapper)
     outputStreamMock = Mock(ServletOutputStream)
+    graphqlInvoker.execute(invocationInputMock) >> FutureExecutionResult.single(invocationInputMock, CompletableFuture.completedFuture(Mock(GraphQLQueryResult)))
     cachingInvoker = new CachingHttpRequestInvoker(configuration, httpRequestInvokerMock, cacheReaderMock)
 
     configuration.getResponseCacheManager() >> responseCacheManagerMock