Pass in keep alive duration

Author: oliemansm

src/main/java/graphql/servlet/ApolloSubscriptionConnectionListener.java

@@ -1,9 +1,12 @@
 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() {
@@ -14,4 +17,8 @@ default Optional<Object> onConnect(Object payload) throws SubscriptionException
         return Optional.empty();
     }
 
+    default Duration getKeepAliveInterval() {
+        return Duration.ofSeconds(KEEP_ALIVE_INTERVAL_SEC);
+    }
+
 }

src/main/java/graphql/servlet/internal/ApolloSubscriptionKeepAliveRunner.java

@@ -4,7 +4,9 @@
 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;
@@ -16,19 +18,20 @@ class ApolloSubscriptionKeepAliveRunner {
 
     private static final Logger LOG = LoggerFactory.getLogger(ApolloSubscriptionKeepAliveRunner.class);
 
-    private static final long KEEP_ALIVE_INTERVAL_SEC = 15;
     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) {
-        this.sender = sender;
+    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) {
@@ -48,7 +51,7 @@ private ScheduledFuture<?> startKeepAlive(Session session) {
                 LOG.error("Cannot send keep alive message. Aborting keep alive", t);
                 abort(session);
             }
-        }, 0, KEEP_ALIVE_INTERVAL_SEC, TimeUnit.SECONDS);
+        }, 0, keepAliveIntervalSeconds, TimeUnit.SECONDS);
     }
 
     void abort(Session session) {

src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolFactory.java

@@ -1,23 +1,30 @@
 package graphql.servlet.internal;
 
+import graphql.servlet.ApolloSubscriptionConnectionListener;
+
 /**
  * @author Andrew Potter
  */
 public class ApolloSubscriptionProtocolFactory extends SubscriptionProtocolFactory {
     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);
+        keepAliveRunner = new ApolloSubscriptionKeepAliveRunner(subscriptionSender, connectionListener.getKeepAliveInterval());
     }
 
     @Override
     public SubscriptionProtocolHandler createHandler() {
-        return new ApolloSubscriptionProtocolHandler(subscriptionHandlerInput, subscriptionSender, keepAliveRunner);
+        return new ApolloSubscriptionProtocolHandler(subscriptionHandlerInput, connectionListener, subscriptionSender, keepAliveRunner);
     }
 }

src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolHandler.java

@@ -39,13 +39,11 @@ public class ApolloSubscriptionProtocolHandler extends SubscriptionProtocolHandl
     private final ApolloSubscriptionConnectionListener connectionListener;
 
     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;
     }