Automatically reconnect to the stream if a heartbeat is not received in time

Author: jkodumal

build.gradle

@@ -32,7 +32,7 @@ dependencies {
     compile "com.google.guava:guava:19.0"
     compile "joda-time:joda-time:2.9.3"
     compile "org.slf4j:slf4j-api:1.7.21"
-    compile group: "com.launchdarkly", name: "okhttp-eventsource", version: "0.2.3", changing: true
+    compile group: "com.launchdarkly", name: "okhttp-eventsource", version: "1.0.0-SNAPSHOT", changing: true
     compile "redis.clients:jedis:2.8.1"
     testCompile "org.easymock:easymock:3.4"
     testCompile 'junit:junit:4.12'

src/main/java/com/launchdarkly/client/StreamProcessor.java

@@ -4,13 +4,18 @@
 import com.launchdarkly.eventsource.EventHandler;
 import com.launchdarkly.eventsource.EventSource;
 import com.launchdarkly.eventsource.MessageEvent;
+import com.launchdarkly.eventsource.ReadyState;
 import okhttp3.Headers;
+import org.joda.time.DateTime;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 class StreamProcessor implements UpdateProcessor {
@@ -20,12 +25,15 @@ class StreamProcessor implements UpdateProcessor {
   private static final String INDIRECT_PUT = "indirect/put";
   private static final String INDIRECT_PATCH = "indirect/patch";
   private static final Logger logger = LoggerFactory.getLogger(StreamProcessor.class);
+  private static final int DEAD_CONNECTION_INTERVAL_SECONDS = 300;
 
   private final FeatureStore store;
   private final LDConfig config;
   private final String sdkKey;
   private final FeatureRequestor requestor;
-  private EventSource es;
+  private final ScheduledExecutorService heartbeatDetectorService;
+  private DateTime lastHeartbeat;
+  private volatile EventSource es;
   private AtomicBoolean initialized = new AtomicBoolean(false);
 
 
@@ -34,6 +42,8 @@ class StreamProcessor implements UpdateProcessor {
     this.config = config;
     this.sdkKey = sdkKey;
     this.requestor = requestor;
+    this.heartbeatDetectorService = Executors.newScheduledThreadPool(1);
+    heartbeatDetectorService.scheduleAtFixedRate(new HeartbeatDetector(), 1, 1, TimeUnit.MINUTES);
   }
 
   @Override
@@ -50,7 +60,7 @@ public Future<Void> start() {
 
       @Override
       public void onOpen() throws Exception {
-
+        lastHeartbeat = DateTime.now();
       }
 
       @Override
@@ -100,6 +110,12 @@ public void onMessage(String name, MessageEvent event) throws Exception {
         }
       }
 
+      @Override
+      public void onComment(String comment) {
+        logger.debug("Received a heartbeat");
+        lastHeartbeat = DateTime.now();
+      }
+
       @Override
       public void onError(Throwable throwable) {
         logger.error("Encountered EventSource error: " + throwable.getMessage());
@@ -125,6 +141,14 @@ public void close() throws IOException {
     if (store != null) {
       store.close();
     }
+    if (heartbeatDetectorService != null) {
+      heartbeatDetectorService.shutdownNow();
+      try {
+        heartbeatDetectorService.awaitTermination(100, TimeUnit.MILLISECONDS);
+      } catch (InterruptedException e) {
+        logger.error("Encountered an exception terminating heartbeat detector: " + e.getMessage());
+      }
+    }
   }
 
   @Override
@@ -171,4 +195,21 @@ int version() {
     }
 
   }
+
+  private final class HeartbeatDetector implements Runnable {
+
+    @Override
+    public void run() {
+      DateTime fiveMinutesAgo = DateTime.now().minusSeconds(DEAD_CONNECTION_INTERVAL_SECONDS);
+      if (lastHeartbeat.isBefore(fiveMinutesAgo) && es.getState() == ReadyState.OPEN) {
+        try {
+          logger.info("Stream stopped receiving heartbeats- reconnecting.");
+          es.close();
+          start();
+        } catch (IOException e) {
+          logger.warn("Encountered exception closing stream connection: " + e.getMessage());
+        }
+      }
+    }
+  }
 }