Merge pull request #118 from pusher/fix-pong-timeout-disconnect

Immediately set connection state to disconnected due to a pong timeout

Author: Luís Fonseca

src/main/java/com/pusher/client/connection/websocket/WebSocketClientWrapper.java

@@ -22,15 +22,14 @@
 public class WebSocketClientWrapper extends WebSocketClient {
 
     private static final String WSS_SCHEME = "wss";
-    private final WebSocketListener webSocketListener;
+    private WebSocketListener webSocketListener;
 
     public WebSocketClientWrapper(final URI uri, final Proxy proxy, final WebSocketListener webSocketListener) throws SSLException {
         super(uri);
 
         if (uri.getScheme().equals(WSS_SCHEME)) {
             try {
-                SSLContext sslContext = null;
-                sslContext = SSLContext.getInstance("TLS");
+                SSLContext sslContext = SSLContext.getInstance("TLS");
                 sslContext.init(null, null, null); // will use java's default
                                                    // key and trust store which
                                                    // is sufficient unless you
@@ -58,21 +57,36 @@ public WebSocketClientWrapper(final URI uri, final Proxy proxy, final WebSocketL
 
     @Override
     public void onOpen(final ServerHandshake handshakedata) {
-        webSocketListener.onOpen(handshakedata);
+        if (webSocketListener != null) {
+            webSocketListener.onOpen(handshakedata);
+        }
     }
 
     @Override
     public void onMessage(final String message) {
-        webSocketListener.onMessage(message);
+        if (webSocketListener != null) {
+            webSocketListener.onMessage(message);
+        }
     }
 
     @Override
     public void onClose(final int code, final String reason, final boolean remote) {
-        webSocketListener.onClose(code, reason, remote);
+        if (webSocketListener != null) {
+            webSocketListener.onClose(code, reason, remote);
+        }
     }
 
     @Override
     public void onError(final Exception ex) {
-        webSocketListener.onError(ex);
+        if (webSocketListener != null) {
+            webSocketListener.onError(ex);
+        }
+    }
+
+    /**
+     * Removes the WebSocketListener so that the underlying WebSocketClient doesn't expose any listener events.
+     */
+    public void removeWebSocketListener() {
+        webSocketListener = null;
     }
 }

src/main/java/com/pusher/client/connection/websocket/WebSocketConnection.java

@@ -40,7 +40,7 @@ public class WebSocketConnection implements InternalConnection, WebSocketListene
     private final Proxy proxy;
 
     private volatile ConnectionState state = ConnectionState.DISCONNECTED;
-    private WebSocketClient underlyingConnection;
+    private WebSocketClientWrapper underlyingConnection;
     private String socketId;
 
     public WebSocketConnection(
@@ -251,18 +251,18 @@ public void run() {
 
     @Override
     public void onClose(final int code, final String reason, final boolean remote) {
+        if (state == ConnectionState.DISCONNECTED) {
+            log.error("Received close from underlying socket when already disconnected. " + "Close code ["
+                    + code + "], Reason [" + reason + "], Remote [" + remote + "]");
+            return;
+        }
+
         activityTimer.cancelTimeouts();
 
         factory.queueOnEventThread(new Runnable() {
             @Override
             public void run() {
-                if (state != ConnectionState.DISCONNECTED) {
-                    updateState(ConnectionState.DISCONNECTED);
-                }
-                else {
-                    log.error("Received close from underlying socket when already disconnected. " + "Close code ["
-                            + code + "], Reason [" + reason + "], Remote [" + remote + "]");
-                }
+                updateState(ConnectionState.DISCONNECTED);
                 factory.shutdownThreads();
             }
         });
@@ -342,7 +342,15 @@ private synchronized void schedulePongCheck() {
                 @Override
                 public void run() {
                     log.debug("Timed out awaiting pong from server - disconnecting");
+
+                    underlyingConnection.removeWebSocketListener();
+
                     disconnect();
+
+                    // Proceed immediately to handle the close
+                    // The WebSocketClient will attempt a graceful WebSocket shutdown by exchanging the close frames
+                    // but may not succeed if this disconnect was called due to pong timeout...
+                    onClose(-1, "Pong timeout", false);
                 }
             }, pongTimeout, TimeUnit.MILLISECONDS);
         }

src/main/java/com/pusher/client/util/Factory.java

@@ -62,7 +62,7 @@ public synchronized InternalConnection getConnection(final String apiKey, final
         return connection;
     }
 
-    public WebSocketClient newWebSocketClientWrapper(final URI uri, final Proxy proxy, final WebSocketListener webSocketListener) throws SSLException {
+    public WebSocketClientWrapper newWebSocketClientWrapper(final URI uri, final Proxy proxy, final WebSocketListener webSocketListener) throws SSLException {
         return new WebSocketClientWrapper(uri, proxy, webSocketListener);
     }
 

src/test/java/com/pusher/client/connection/websocket/WebSocketClientWrapperTest.java

@@ -13,6 +13,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.internal.verification.NoMoreInteractions;
 import org.mockito.runners.MockitoJUnitRunner;
 
 @RunWith(MockitoJUnitRunner.class)
@@ -52,4 +53,15 @@ public void testOnErrorIsDelegatedToTheListener() {
         wrapper.onError(e);
         verify(mockListener).onError(e);
     }
+
+    @Test
+    public void testRemoveWebSocketListener() {
+        wrapper.onClose(1, "reason", true);
+        verify(mockListener).onClose(1, "reason", true);
+
+        wrapper.removeWebSocketListener();
+
+        wrapper.onClose(1, "reason", true);
+        verify(mockListener, new NoMoreInteractions()).onClose(1, "reason", true);
+    }
 }

src/test/java/com/pusher/client/connection/websocket/WebSocketConnectionTest.java

@@ -8,6 +8,7 @@
 import java.net.Proxy;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 
 import javax.net.ssl.SSLException;
 
@@ -29,8 +30,8 @@
 @RunWith(MockitoJUnitRunner.class)
 public class WebSocketConnectionTest {
 
-    private static final long ACTIVITY_TIMEOUT = 120000;
-    private static final long PONG_TIMEOUT = 30000;
+    private static final long ACTIVITY_TIMEOUT = 500;
+    private static final long PONG_TIMEOUT = 500;
     private static final String URL = "ws://ws.example.com/";
     private static final String EVENT_NAME = "my-event";
     private static final String CONN_ESTABLISHED_EVENT = "{\"event\":\"pusher:connection_established\",\"data\":\"{\\\"socket_id\\\":\\\"21112.816204\\\"}\"}";
@@ -278,6 +279,24 @@ public void testDisconnectInDisconnectingStateIsIgnored() {
         verify(mockEventListener, times(3)).onConnectionStateChange(any(ConnectionStateChange.class));
     }
 
+    @Test
+    public void testPongTimeoutResultsInDisconnect() throws InterruptedException {
+        when(factory.getTimers()).thenReturn(new ScheduledThreadPoolExecutor(2));
+
+        connection.connect();
+        connection.onMessage(CONN_ESTABLISHED_EVENT);
+
+        verify(mockUnderlyingConnection, timeout((int) (ACTIVITY_TIMEOUT + PONG_TIMEOUT))).close();
+
+        verify(mockEventListener).onConnectionStateChange(
+                new ConnectionStateChange(ConnectionState.CONNECTED, ConnectionState.DISCONNECTING));
+
+        verify(mockEventListener).onConnectionStateChange(
+                new ConnectionStateChange(ConnectionState.DISCONNECTING, ConnectionState.DISCONNECTED));
+
+        assertEquals(ConnectionState.DISCONNECTED, connection.getState());
+    }
+
     /* end of tests */
 
     private void connect() {