Added support for asyncModeEnabled=true

Author: wleroux

src/main/java/graphql/servlet/AbstractGraphQLHttpServlet.java

@@ -30,6 +30,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -291,7 +292,7 @@ public String executeQuery(String query) {
 
     private void doRequestAsync(HttpServletRequest request, HttpServletResponse response, HttpRequestHandler handler) {
         if (configuration.isAsyncServletModeEnabled()) {
-            AsyncContext asyncContext = request.startAsync();
+            AsyncContext asyncContext = request.startAsync(request, response);
             HttpServletRequest asyncRequest = (HttpServletRequest) asyncContext.getRequest();
             HttpServletResponse asyncResponse = (HttpServletResponse) asyncContext.getResponse();
             new Thread(() -> doRequest(asyncRequest, asyncResponse, handler, asyncContext)).start();
@@ -344,12 +345,22 @@ private void query(GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQL
             resp.setContentType(APPLICATION_EVENT_STREAM_UTF8);
             resp.setStatus(STATUS_OK);
 
-            HttpServletRequest req = invocationInput.getContext().getHttpServletRequest().get();
-            AsyncContext asyncContext = req.startAsync(req, resp);
-            asyncContext.setTimeout(60 * 1000);
+            HttpServletRequest req = invocationInput.getContext().getHttpServletRequest().orElseThrow(IllegalStateException::new);
+            boolean isInAsyncThread = req.isAsyncStarted();
+            AsyncContext asyncContext = isInAsyncThread ? req.getAsyncContext() : req.startAsync(req, resp);
+            asyncContext.setTimeout(configuration.getSubscriptionTimeout());
             AtomicReference<Subscription> subscriptionRef = new AtomicReference<>();
             asyncContext.addListener(new SubscriptionAsyncListener(subscriptionRef));
-            ((Publisher<ExecutionResult>) result.getData()).subscribe(new ExecutionResultSubscriber(subscriptionRef, asyncContext, graphQLObjectMapper));
+            ExecutionResultSubscriber subscriber = new ExecutionResultSubscriber(subscriptionRef, asyncContext, graphQLObjectMapper);
+            ((Publisher<ExecutionResult>) result.getData()).subscribe(subscriber);
+            if (isInAsyncThread) {
+                // We need to delay the completion of async context until after the subscription has terminated, otherwise the AsyncContext is prematurely closed.
+                try {
+                    subscriber.await();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                }
+            }
         }
     }
 
@@ -480,6 +491,7 @@ private static class ExecutionResultSubscriber implements Subscriber<ExecutionRe
         private final AtomicReference<Subscription> subscriptionRef;
         private final AsyncContext asyncContext;
         private final GraphQLObjectMapper graphQLObjectMapper;
+        private final CountDownLatch completedLatch = new CountDownLatch(1);
 
         public ExecutionResultSubscriber(AtomicReference<Subscription> subscriptionRef, AsyncContext asyncContext, GraphQLObjectMapper graphQLObjectMapper) {
             this.subscriptionRef = subscriptionRef;
@@ -507,11 +519,17 @@ public void onNext(ExecutionResult executionResult) {
         @Override
         public void onError(Throwable t) {
             asyncContext.complete();
+            completedLatch.countDown();
         }
 
         @Override
         public void onComplete() {
             asyncContext.complete();
+            completedLatch.countDown();
+        }
+
+        public void await() throws InterruptedException {
+            completedLatch.await();
         }
     }
 }

src/main/java/graphql/servlet/GraphQLConfiguration.java

@@ -13,6 +13,7 @@ public class GraphQLConfiguration {
     private GraphQLObjectMapper objectMapper;
     private List<GraphQLServletListener> listeners;
     private boolean asyncServletModeEnabled;
+    private long subscriptionTimeout;
 
     public static GraphQLConfiguration.Builder with(GraphQLSchema schema) {
         return with(new DefaultGraphQLSchemaProvider(schema));
@@ -26,12 +27,13 @@ public static GraphQLConfiguration.Builder with(GraphQLInvocationInputFactory in
         return new Builder(invocationInputFactory);
     }
 
-    private GraphQLConfiguration(GraphQLInvocationInputFactory invocationInputFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper objectMapper, List<GraphQLServletListener> listeners, boolean asyncServletModeEnabled) {
+    private GraphQLConfiguration(GraphQLInvocationInputFactory invocationInputFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper objectMapper, List<GraphQLServletListener> listeners, boolean asyncServletModeEnabled, long subscriptionTimeout) {
         this.invocationInputFactory = invocationInputFactory;
         this.queryInvoker = queryInvoker;
         this.objectMapper = objectMapper;
         this.listeners = listeners;
         this.asyncServletModeEnabled = asyncServletModeEnabled;
+        this.subscriptionTimeout = subscriptionTimeout;
     }
 
     public GraphQLInvocationInputFactory getInvocationInputFactory() {
@@ -62,6 +64,10 @@ public boolean remove(GraphQLServletListener listener) {
         return listeners.remove(listener);
     }
 
+    public long getSubscriptionTimeout() {
+        return subscriptionTimeout;
+    }
+
     public static class Builder {
 
         private GraphQLInvocationInputFactory.Builder invocationInputFactoryBuilder;
@@ -70,6 +76,7 @@ public static class Builder {
         private GraphQLObjectMapper objectMapper = GraphQLObjectMapper.newBuilder().build();
         private List<GraphQLServletListener> listeners = new ArrayList<>();
         private boolean asyncServletModeEnabled = false;
+        private long subscriptionTimeout = 0;
 
         private Builder(GraphQLInvocationInputFactory.Builder invocationInputFactoryBuilder) {
             this.invocationInputFactoryBuilder = invocationInputFactoryBuilder;
@@ -115,13 +122,19 @@ public Builder with(GraphQLRootObjectBuilder rootObjectBuilder) {
             return this;
         }
 
+        public Builder with(long subscriptionTimeout) {
+            this.subscriptionTimeout = subscriptionTimeout;
+            return this;
+        }
+
         public GraphQLConfiguration build() {
             return new GraphQLConfiguration(
                     this.invocationInputFactory != null ? this.invocationInputFactory : invocationInputFactoryBuilder.build(),
                     queryInvoker,
                     objectMapper,
                     listeners,
-                    asyncServletModeEnabled
+                    asyncServletModeEnabled,
+                    subscriptionTimeout
             );
         }
 

src/main/java/graphql/servlet/SimpleGraphQLHttpServlet.java

@@ -30,6 +30,21 @@ public SimpleGraphQLHttpServlet(GraphQLInvocationInputFactory invocationInputFac
                 .build();
     }
 
+    /**
+     * @deprecated use {@link GraphQLHttpServlet} instead
+     */
+    @Deprecated
+    public SimpleGraphQLHttpServlet(GraphQLInvocationInputFactory invocationInputFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQLObjectMapper, List<GraphQLServletListener> listeners, boolean asyncServletMode, long subscriptionTimeout) {
+        super(listeners);
+        this.configuration = GraphQLConfiguration.with(invocationInputFactory)
+                .with(queryInvoker)
+                .with(graphQLObjectMapper)
+                .with(listeners != null ? listeners : new ArrayList<>())
+                .with(asyncServletMode)
+                .with(subscriptionTimeout)
+                .build();
+    }
+
     private SimpleGraphQLHttpServlet(GraphQLConfiguration configuration) {
         this.configuration = Objects.requireNonNull(configuration, "configuration is required");
     }
@@ -77,6 +92,7 @@ public static class Builder {
         private GraphQLObjectMapper graphQLObjectMapper = GraphQLObjectMapper.newBuilder().build();
         private List<GraphQLServletListener> listeners;
         private boolean asyncServletMode;
+        private long subscriptionTimeout;
 
         Builder(GraphQLInvocationInputFactory invocationInputFactory) {
             this.invocationInputFactory = invocationInputFactory;
@@ -102,13 +118,19 @@ public Builder withListeners(List<GraphQLServletListener> listeners) {
             return this;
         }
 
+        public Builder withSubscriptionTimeout(long subscriptionTimeout) {
+            this.subscriptionTimeout = subscriptionTimeout;
+            return this;
+        }
+
         @Deprecated
         public SimpleGraphQLHttpServlet build() {
             GraphQLConfiguration configuration = GraphQLConfiguration.with(invocationInputFactory)
                     .with(queryInvoker)
                     .with(graphQLObjectMapper)
                     .with(listeners != null ? listeners : new ArrayList<>())
                     .with(asyncServletMode)
+                    .with(subscriptionTimeout)
                     .build();
             return new SimpleGraphQLHttpServlet(configuration);
         }