Merge pull request #112 from graphql-java-kickstart/bugfix/91

oliemansm · web-flow · commit 841bf58557f2 · 2018-10-13T08:12:55.000+01:00
Fix Apollo keep alive implementation
diff --git a/src/main/java/graphql/servlet/ApolloSubscriptionConnectionListener.java b/src/main/java/graphql/servlet/ApolloSubscriptionConnectionListener.java
@@ -1,17 +1,42 @@
 package graphql.servlet;
 
+import java.time.Duration;
 import java.util.Optional;
 
 public interface ApolloSubscriptionConnectionListener extends SubscriptionConnectionListener {
 
+    long KEEP_ALIVE_INTERVAL_SEC = 15;
+
     String CONNECT_RESULT_KEY = "CONNECT_RESULT";
 
     default boolean isKeepAliveEnabled() {
-        return false;
+        return true;
     }
 
-    default Optional<Object> onConnect(Object payload) {
+    default Optional<Object> onConnect(Object payload) throws SubscriptionException {
         return Optional.empty();
     }
 
+    default Duration getKeepAliveInterval() {
+        return Duration.ofSeconds(KEEP_ALIVE_INTERVAL_SEC);
+    }
+
+    static ApolloSubscriptionConnectionListener createWithKeepAliveDisabled() {
+        return new ApolloSubscriptionConnectionListener() {
+            @Override
+            public boolean isKeepAliveEnabled() {
+                return false;
+            }
+        };
+    }
+
+    static ApolloSubscriptionConnectionListener createWithKeepAliveInterval(Duration interval) {
+        return new ApolloSubscriptionConnectionListener() {
+            @Override
+            public Duration getKeepAliveInterval() {
+                return interval;
+            }
+        };
+    }
+
 }
diff --git a/src/main/java/graphql/servlet/GraphQLWebsocketServlet.java b/src/main/java/graphql/servlet/GraphQLWebsocketServlet.java
@@ -30,15 +30,9 @@ public class GraphQLWebsocketServlet extends Endpoint {
     private static final CloseReason ERROR_CLOSE_REASON = new CloseReason(CloseReason.CloseCodes.UNEXPECTED_CONDITION, "Internal Server Error");
     private static final CloseReason SHUTDOWN_CLOSE_REASON = new CloseReason(CloseReason.CloseCodes.UNEXPECTED_CONDITION, "Server Shut Down");
 
-    private static final List<SubscriptionProtocolFactory> subscriptionProtocolFactories = Collections.singletonList(new ApolloSubscriptionProtocolFactory());
-    private static final SubscriptionProtocolFactory fallbackSubscriptionProtocolFactory = new FallbackSubscriptionProtocolFactory();
-    private static final List<String> allSubscriptionProtocols;
-
-    static {
-        allSubscriptionProtocols = Stream.concat(subscriptionProtocolFactories.stream(), Stream.of(fallbackSubscriptionProtocolFactory))
-                .map(SubscriptionProtocolFactory::getProtocol)
-                .collect(Collectors.toList());
-    }
+    private final List<SubscriptionProtocolFactory> subscriptionProtocolFactories;
+    private final SubscriptionProtocolFactory fallbackSubscriptionProtocolFactory;
+    private final List<String> allSubscriptionProtocols;
 
     private final Map<Session, WsSessionSubscriptions> sessionSubscriptionCache = new HashMap<>();
     private final SubscriptionHandlerInput subscriptionHandlerInput;
@@ -52,6 +46,12 @@ public GraphQLWebsocketServlet(GraphQLQueryInvoker queryInvoker, GraphQLInvocati
 
     public GraphQLWebsocketServlet(GraphQLQueryInvoker queryInvoker, GraphQLInvocationInputFactory invocationInputFactory, GraphQLObjectMapper graphQLObjectMapper, SubscriptionConnectionListener subscriptionConnectionListener) {
         this.subscriptionHandlerInput = new SubscriptionHandlerInput(invocationInputFactory, queryInvoker, graphQLObjectMapper, subscriptionConnectionListener);
+
+        subscriptionProtocolFactories = Collections.singletonList(new ApolloSubscriptionProtocolFactory(subscriptionHandlerInput));
+        fallbackSubscriptionProtocolFactory = new FallbackSubscriptionProtocolFactory(subscriptionHandlerInput);
+        allSubscriptionProtocols = Stream.concat(subscriptionProtocolFactories.stream(), Stream.of(fallbackSubscriptionProtocolFactory))
+                                         .map(SubscriptionProtocolFactory::getProtocol)
+                                         .collect(Collectors.toList());
     }
 
     @Override
@@ -119,7 +119,7 @@ public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request,
         }
 
         SubscriptionProtocolFactory subscriptionProtocolFactory = getSubscriptionProtocolFactory(protocol);
-        sec.getUserProperties().put(PROTOCOL_HANDLER_REQUEST_KEY, subscriptionProtocolFactory.createHandler(subscriptionHandlerInput));
+        sec.getUserProperties().put(PROTOCOL_HANDLER_REQUEST_KEY, subscriptionProtocolFactory.createHandler());
 
         if (request.getHeaders().get(HandshakeResponse.SEC_WEBSOCKET_ACCEPT) != null) {
             response.getHeaders().put(HandshakeResponse.SEC_WEBSOCKET_ACCEPT, allSubscriptionProtocols);
@@ -165,7 +165,7 @@ public boolean isShutDown() {
         return isShutDown.get();
     }
 
-    private static SubscriptionProtocolFactory getSubscriptionProtocolFactory(List<String> accept) {
+    private SubscriptionProtocolFactory getSubscriptionProtocolFactory(List<String> accept) {
         for (String protocol : accept) {
             for (SubscriptionProtocolFactory subscriptionProtocolFactory : subscriptionProtocolFactories) {
                 if (subscriptionProtocolFactory.getProtocol().equals(protocol)) {
diff --git a/src/main/java/graphql/servlet/SubscriptionException.java b/src/main/java/graphql/servlet/SubscriptionException.java
@@ -0,0 +1,18 @@
+package graphql.servlet;
+
+public class SubscriptionException extends Exception {
+
+    private Object payload;
+
+    public SubscriptionException() {
+    }
+
+    public SubscriptionException(Object payload) {
+        this.payload = payload;
+    }
+
+    public Object getPayload() {
+        return payload;
+    }
+
+}
diff --git a/src/main/java/graphql/servlet/internal/ApolloSubscriptionKeepAliveRunner.java b/src/main/java/graphql/servlet/internal/ApolloSubscriptionKeepAliveRunner.java
@@ -0,0 +1,64 @@
+package graphql.servlet.internal;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.websocket.Session;
+import java.time.Duration;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+class ApolloSubscriptionKeepAliveRunner {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ApolloSubscriptionKeepAliveRunner.class);
+
+    private static final int EXECUTOR_POOL_SIZE = 10;
+
+    private final ScheduledExecutorService executor;
+    private final SubscriptionSender sender;
+    private final ApolloSubscriptionProtocolHandler.OperationMessage keepAliveMessage;
+    private final Map<Session, Future<?>> futures;
+    private final long keepAliveIntervalSeconds;
+
+    ApolloSubscriptionKeepAliveRunner(SubscriptionSender sender, Duration keepAliveInterval) {
+        this.sender = Objects.requireNonNull(sender);
+        this.keepAliveMessage = ApolloSubscriptionProtocolHandler.OperationMessage.newKeepAliveMessage();
+        this.executor = Executors.newScheduledThreadPool(EXECUTOR_POOL_SIZE);
+        this.futures = new ConcurrentHashMap<>();
+        this.keepAliveIntervalSeconds = keepAliveInterval.getSeconds();
+    }
+
+    void keepAlive(Session session) {
+        futures.computeIfAbsent(session, this::startKeepAlive);
+    }
+
+    private ScheduledFuture<?> startKeepAlive(Session session) {
+        return executor.scheduleAtFixedRate(() -> {
+            try {
+                if (session.isOpen()) {
+                    sender.send(session, keepAliveMessage);
+                } else {
+                    LOG.warn("Session appears to be closed. Aborting keep alive");
+                    abort(session);
+                }
+            } catch (Throwable t) {
+                LOG.error("Cannot send keep alive message. Aborting keep alive", t);
+                abort(session);
+            }
+        }, 0, keepAliveIntervalSeconds, TimeUnit.SECONDS);
+    }
+
+    void abort(Session session) {
+        Future<?> future = futures.remove(session);
+        if (future != null) {
+            future.cancel(true);
+        }
+    }
+
+}
diff --git a/src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolFactory.java b/src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolFactory.java
@@ -1,15 +1,30 @@
 package graphql.servlet.internal;
 
+import graphql.servlet.ApolloSubscriptionConnectionListener;
+
 /**
  * @author Andrew Potter
  */
 public class ApolloSubscriptionProtocolFactory extends SubscriptionProtocolFactory {
-    public ApolloSubscriptionProtocolFactory() {
+    private final SubscriptionHandlerInput subscriptionHandlerInput;
+    private final SubscriptionSender subscriptionSender;
+    private final ApolloSubscriptionKeepAliveRunner keepAliveRunner;
+    private final ApolloSubscriptionConnectionListener connectionListener;
+
+    public ApolloSubscriptionProtocolFactory(SubscriptionHandlerInput subscriptionHandlerInput) {
         super("graphql-ws");
+        this.subscriptionHandlerInput = subscriptionHandlerInput;
+        this.connectionListener = subscriptionHandlerInput.getSubscriptionConnectionListener()
+                .filter(ApolloSubscriptionConnectionListener.class::isInstance)
+                .map(ApolloSubscriptionConnectionListener.class::cast)
+                .orElse(new ApolloSubscriptionConnectionListener() {});
+        subscriptionSender =
+            new SubscriptionSender(subscriptionHandlerInput.getGraphQLObjectMapper().getJacksonMapper());
+        keepAliveRunner = new ApolloSubscriptionKeepAliveRunner(subscriptionSender, connectionListener.getKeepAliveInterval());
     }
 
     @Override
-    public SubscriptionProtocolHandler createHandler(SubscriptionHandlerInput subscriptionHandlerInput) {
-        return new ApolloSubscriptionProtocolHandler(subscriptionHandlerInput);
+    public SubscriptionProtocolHandler createHandler() {
+        return new ApolloSubscriptionProtocolHandler(subscriptionHandlerInput, connectionListener, subscriptionSender, keepAliveRunner);
     }
 }
diff --git a/src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolHandler.java b/src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolHandler.java
@@ -6,6 +6,7 @@
 import graphql.ExecutionResult;
 import graphql.servlet.ApolloSubscriptionConnectionListener;
 import graphql.servlet.GraphQLSingleInvocationInput;
+import graphql.servlet.SubscriptionException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,14 +34,18 @@ public class ApolloSubscriptionProtocolHandler extends SubscriptionProtocolHandl
     private static final CloseReason TERMINATE_CLOSE_REASON = new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "client requested " + GQL_CONNECTION_TERMINATE.getType());
 
     private final SubscriptionHandlerInput input;
+    private final SubscriptionSender sender;
+    private final ApolloSubscriptionKeepAliveRunner keepAliveRunner;
     private final ApolloSubscriptionConnectionListener connectionListener;
 
-    public ApolloSubscriptionProtocolHandler(SubscriptionHandlerInput subscriptionHandlerInput) {
+    public ApolloSubscriptionProtocolHandler(SubscriptionHandlerInput subscriptionHandlerInput,
+                                             ApolloSubscriptionConnectionListener connectionListener,
+                                             SubscriptionSender subscriptionSender,
+                                             ApolloSubscriptionKeepAliveRunner keepAliveRunner) {
         this.input = subscriptionHandlerInput;
-        this.connectionListener = subscriptionHandlerInput.getSubscriptionConnectionListener()
-                .filter(ApolloSubscriptionConnectionListener.class::isInstance)
-                .map(ApolloSubscriptionConnectionListener.class::cast)
-                .orElse(new ApolloSubscriptionConnectionListener() {});
+        this.connectionListener = connectionListener;
+        this.sender = subscriptionSender;
+        this.keepAliveRunner = keepAliveRunner;
     }
 
     @Override
@@ -54,20 +59,20 @@ public void onMessage(HandshakeRequest request, Session session, WsSessionSubscr
             return;
         }
 
-        switch(message.getType()) {
+        switch (message.getType()) {
             case GQL_CONNECTION_INIT:
                 try {
                     Optional<Object> connectionResponse = connectionListener.onConnect(message.getPayload());
                     connectionResponse.ifPresent(it -> session.getUserProperties().put(ApolloSubscriptionConnectionListener.CONNECT_RESULT_KEY, it));
-                } catch (Throwable t) {
-                    sendMessage(session, OperationMessage.Type.GQL_CONNECTION_ERROR, t.getMessage());
+                } catch (SubscriptionException e) {
+                    sendMessage(session, OperationMessage.Type.GQL_CONNECTION_ERROR, message.getId(), e.getPayload());
                     return;
                 }
 
                 sendMessage(session, OperationMessage.Type.GQL_CONNECTION_ACK, message.getId());
 
                 if (connectionListener.isKeepAliveEnabled()) {
-                    sendMessage(session, OperationMessage.Type.GQL_CONNECTION_KEEP_ALIVE, message.getId());
+                    keepAliveRunner.keepAlive(session);
                 }
                 break;
 
@@ -86,6 +91,7 @@ public void onMessage(HandshakeRequest request, Session session, WsSessionSubscr
                 break;
 
             case GQL_CONNECTION_TERMINATE:
+                keepAliveRunner.abort(session);
                 try {
                     session.close(TERMINATE_CLOSE_REASON);
                 } catch (IOException e) {
@@ -112,7 +118,7 @@ private GraphQLSingleInvocationInput createInvocationInput(Session session, Oper
     private void handleSubscriptionStart(Session session, WsSessionSubscriptions subscriptions, String id, ExecutionResult executionResult) {
         executionResult = input.getGraphQLObjectMapper().sanitizeErrors(executionResult);
 
-        if(input.getGraphQLObjectMapper().areErrorsPresent(executionResult)) {
+        if (input.getGraphQLObjectMapper().areErrorsPresent(executionResult)) {
             sendMessage(session, OperationMessage.Type.GQL_ERROR, id, input.getGraphQLObjectMapper().convertSanitizedExecutionResult(executionResult, false));
             return;
         }
@@ -127,11 +133,13 @@ protected void sendDataMessage(Session session, String id, Object payload) {
 
     @Override
     protected void sendErrorMessage(Session session, String id) {
+        keepAliveRunner.abort(session);
         sendMessage(session, GQL_ERROR, id);
     }
 
     @Override
     protected void sendCompleteMessage(Session session, String id) {
+        keepAliveRunner.abort(session);
         sendMessage(session, GQL_COMPLETE, id);
     }
 
@@ -140,13 +148,7 @@ private void sendMessage(Session session, OperationMessage.Type type, String id)
     }
 
     private void sendMessage(Session session, OperationMessage.Type type, String id, Object payload) {
-        try {
-            session.getBasicRemote().sendText(input.getGraphQLObjectMapper().getJacksonMapper().writeValueAsString(
-                new OperationMessage(type, id, payload)
-            ));
-        } catch (IOException e) {
-            throw new RuntimeException("Error sending subscription response", e);
-        }
+        sender.send(session, new OperationMessage(type, id, payload));
     }
 
     @JsonInclude(JsonInclude.Include.NON_NULL)
@@ -164,6 +166,10 @@ public OperationMessage(Type type, String id, Object payload) {
             this.payload = payload;
         }
 
+        static OperationMessage newKeepAliveMessage() {
+            return new OperationMessage(Type.GQL_CONNECTION_KEEP_ALIVE, null, null);
+        }
+
         public Type getType() {
             return type;
         }
diff --git a/src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolFactory.java b/src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolFactory.java
@@ -4,12 +4,15 @@
  * @author Andrew Potter
  */
 public class FallbackSubscriptionProtocolFactory extends SubscriptionProtocolFactory {
-    public FallbackSubscriptionProtocolFactory() {
+    private final SubscriptionHandlerInput subscriptionHandlerInput;
+
+    public FallbackSubscriptionProtocolFactory(SubscriptionHandlerInput subscriptionHandlerInput) {
         super("");
+        this.subscriptionHandlerInput = subscriptionHandlerInput;
     }
 
     @Override
-    public SubscriptionProtocolHandler createHandler(SubscriptionHandlerInput subscriptionHandlerInput) {
+    public SubscriptionProtocolHandler createHandler() {
         return new FallbackSubscriptionProtocolHandler(subscriptionHandlerInput);
     }
 }
diff --git a/src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolHandler.java b/src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolHandler.java
@@ -11,9 +11,11 @@
 public class FallbackSubscriptionProtocolHandler extends SubscriptionProtocolHandler {
 
     private final SubscriptionHandlerInput input;
+    private final SubscriptionSender sender;
 
     public FallbackSubscriptionProtocolHandler(SubscriptionHandlerInput subscriptionHandlerInput) {
         this.input = subscriptionHandlerInput;
+        sender = new SubscriptionSender(subscriptionHandlerInput.getGraphQLObjectMapper().getJacksonMapper());
     }
 
     @Override
@@ -32,11 +34,7 @@ public void onMessage(HandshakeRequest request, Session session, WsSessionSubscr
 
     @Override
     protected void sendDataMessage(Session session, String id, Object payload) {
-        try {
-            session.getBasicRemote().sendText(input.getGraphQLObjectMapper().getJacksonMapper().writeValueAsString(payload));
-        } catch (IOException e) {
-            throw new RuntimeException("Error sending subscription response", e);
-        }
+        sender.send(session, payload);
     }
 
     @Override
diff --git a/src/main/java/graphql/servlet/internal/SubscriptionProtocolFactory.java b/src/main/java/graphql/servlet/internal/SubscriptionProtocolFactory.java
@@ -14,5 +14,5 @@ public String getProtocol() {
         return protocol;
     }
 
-    public abstract SubscriptionProtocolHandler createHandler(SubscriptionHandlerInput subscriptionHandlerInput);
+    public abstract SubscriptionProtocolHandler createHandler();
 }
diff --git a/src/main/java/graphql/servlet/internal/SubscriptionProtocolHandler.java b/src/main/java/graphql/servlet/internal/SubscriptionProtocolHandler.java
@@ -55,13 +55,13 @@ public void onNext(ExecutionResult executionResult) {
                 @Override
                 public void onError(Throwable throwable) {
                     log.error("Subscription error", throwable);
-                    subscriptions.cancel(id);
+                    unsubscribe(subscriptions, id);
                     sendErrorMessage(session, id);
                 }
 
                 @Override
                 public void onComplete() {
-                    subscriptions.cancel(id);
+                    unsubscribe(subscriptions, id);
                     sendCompleteMessage(session, id);
                 }
             });
diff --git a/src/main/java/graphql/servlet/internal/SubscriptionSender.java b/src/main/java/graphql/servlet/internal/SubscriptionSender.java
diff --git a/src/main/java/graphql/servlet/internal/WsSessionSubscriptions.java b/src/main/java/graphql/servlet/internal/WsSessionSubscriptions.java

Original file line number	Diff line number	Diff line change
`@@ -4,12 +4,15 @@`
`4`	`4`	`* @author Andrew Potter`
`5`	`5`	`*/`
`6`	`6`	`public class FallbackSubscriptionProtocolFactory extends SubscriptionProtocolFactory {`
`7`		`- public FallbackSubscriptionProtocolFactory() {`
	`7`	`+ private final SubscriptionHandlerInput subscriptionHandlerInput;`
	`8`	`+`
	`9`	`+ public FallbackSubscriptionProtocolFactory(SubscriptionHandlerInput subscriptionHandlerInput) {`
`8`	`10`	`super("");`
	`11`	`+ this.subscriptionHandlerInput = subscriptionHandlerInput;`
`9`	`12`	`}`
`10`	`13`
`11`	`14`	`@Override`
`12`		`- public SubscriptionProtocolHandler createHandler(SubscriptionHandlerInput subscriptionHandlerInput) {`
	`15`	`+ public SubscriptionProtocolHandler createHandler() {`
`13`	`16`	`return new FallbackSubscriptionProtocolHandler(subscriptionHandlerInput);`
`14`	`17`	`}`
`15`	`18`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,5 +14,5 @@ public String getProtocol() {`
`14`	`14`	`return protocol;`
`15`	`15`	`}`
`16`	`16`
`17`		`- public abstract SubscriptionProtocolHandler createHandler(SubscriptionHandlerInput subscriptionHandlerInput);`
	`17`	`+ public abstract SubscriptionProtocolHandler createHandler();`
`18`	`18`	`}`