Add WebSocket client ping implementation (#30293)

This commit introduces new functionality on the client side to send PING messages over the websocket. If the server doesn't respond within PING_INTERVAL (currently 30 seconds), the connection is closed and re-created. This will allow us to find broken connections more quickly.

webapp/channels/src/components/post_view/post_body_additional_content/__snapshots__/post_body_additional_content.test.tsx.snap

@@ -100,6 +100,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "closeCallback": null,
         "closeListeners": Set {},
         "config": Object {
+          "clientPingInterval": 30000,
           "maxWebSocketFails": 7,
           "maxWebSocketRetryTime": 300000,
           "minWebSocketRetryTime": 3000,
@@ -118,6 +119,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "messageListeners": Set {},
         "missedEventCallback": null,
         "missedMessageListeners": Set {},
+        "pingInterval": null,
         "postedAck": false,
         "reconnectCallback": null,
         "reconnectListeners": Set {},
@@ -156,6 +158,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "closeCallback": null,
         "closeListeners": Set {},
         "config": Object {
+          "clientPingInterval": 30000,
           "maxWebSocketFails": 7,
           "maxWebSocketRetryTime": 300000,
           "minWebSocketRetryTime": 3000,
@@ -174,6 +177,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "messageListeners": Set {},
         "missedEventCallback": null,
         "missedMessageListeners": Set {},
+        "pingInterval": null,
         "postedAck": false,
         "reconnectCallback": null,
         "reconnectListeners": Set {},

webapp/platform/client/src/websocket.test.ts

@@ -92,7 +92,7 @@ describe('websocketclient', () => {
 
         mockWebSocket.close();
 
-        jest.advanceTimersByTime(40);
+        jest.advanceTimersByTime(100);
 
         client.close();
         expect(openSpy).toHaveBeenCalledTimes(2);
@@ -209,4 +209,198 @@ describe('websocketclient', () => {
 
         jest.useRealTimers();
     });
+
+    test('should stay connected after ping response', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 1,
+            reconnectJitterRange: 1,
+            clientPingInterval: 1,
+        });
+
+        let numPings = 0;
+        let numPongs = 0;
+        mockWebSocket.send = (evt) => {
+            const msg = JSON.parse(evt);
+
+            if (msg.action !== 'ping') {
+                return;
+            }
+            numPings++;
+
+            const rsp = {
+                text: 'pong',
+                seq_reply: msg.seq,
+            };
+
+            if (mockWebSocket.onmessage) {
+                mockWebSocket.onmessage({data: JSON.stringify(rsp)});
+                numPongs++;
+            }
+        };
+
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+        const closeSpy = jest.spyOn(mockWebSocket, 'close');
+
+        client.initialize('mock.url');
+        mockWebSocket.open();
+
+        jest.advanceTimersByTime(30);
+
+        client.close();
+
+        expect(openSpy).toBeCalledTimes(1);
+        expect(closeSpy).toBeCalledTimes(1);
+        expect(numPings).toBeGreaterThan(10);
+        expect(numPongs).toBeGreaterThan(10);
+
+        jest.useRealTimers();
+    });
+
+    test('should reconnect after no ping response', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 1,
+            reconnectJitterRange: 1,
+            clientPingInterval: 1,
+        });
+
+        let numPings = 0;
+        let numPongs = 0;
+        mockWebSocket.send = (evt) => {
+            const msg = JSON.parse(evt);
+
+            if (msg.action !== 'ping') {
+                return;
+            }
+            numPings++;
+
+            // stop responding after three pings
+            if (numPings > 3) {
+                return;
+            }
+
+            const rsp = {
+                text: 'pong',
+                seq_reply: msg.seq,
+            };
+
+            if (mockWebSocket.onmessage) {
+                mockWebSocket.onmessage({data: JSON.stringify(rsp)});
+                numPongs++;
+            }
+        };
+
+        mockWebSocket.open = jest.fn(mockWebSocket.open);
+        mockWebSocket.close = jest.fn(() => {
+            mockWebSocket.readyState = WebSocket.CLOSED;
+            if (mockWebSocket.onclose) {
+                mockWebSocket.onclose();
+            }
+            if ((mockWebSocket.close as jest.Mock).mock.calls.length > 2) {
+                client.close();
+            }
+        });
+
+        client.initialize('mock.url');
+        mockWebSocket.open();
+
+        jest.advanceTimersByTime(30);
+
+        client.close();
+
+        expect(mockWebSocket.open).toBeCalledTimes(3);
+        expect(mockWebSocket.close).toBeCalledTimes(3);
+        expect(numPings).toBe(6);
+        expect(numPongs).toBe(3);
+
+        jest.useRealTimers();
+    });
+
+    test('should reset ping interval state when reconnecting during pending ping', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 1,
+            reconnectJitterRange: 1,
+            clientPingInterval: 15,
+        });
+
+        let numPings = 0;
+        let numPongs = 0;
+        mockWebSocket.send = (evt) => {
+            const msg = JSON.parse(evt);
+
+            if (msg.action !== 'ping') {
+                return;
+            }
+            numPings++;
+
+            // don't respond to first ping
+            if (numPings === 1) {
+                return;
+            }
+
+            const rsp = {
+                text: 'pong',
+                seq_reply: msg.seq,
+            };
+
+            if (mockWebSocket.onmessage) {
+                mockWebSocket.onmessage({data: JSON.stringify(rsp)});
+                numPongs++;
+            }
+        };
+
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+        const closeSpy = jest.spyOn(mockWebSocket, 'close');
+
+        client.initialize('mock.url');
+        mockWebSocket.open();
+
+        // Let first ping happen
+        jest.advanceTimersByTime(25);
+        expect(numPings).toBe(1);
+        expect(numPongs).toBe(0);
+
+        // Close and reopen connection before ping timeout
+        mockWebSocket.close();
+
+        // Let new connection run for a while to ensure no immediate reconnect
+        jest.advanceTimersByTime(100);
+        client.close();
+
+        expect(numPings).toBe(7);
+        expect(numPongs).toBe(numPings - 1); // Ensure we only skipped the first response
+        expect(openSpy).toHaveBeenCalledTimes(2); // Initial open and one reconnect
+        expect(closeSpy).toHaveBeenCalledTimes(2); // Manual close and final close
+
+        jest.useRealTimers();
+    });
 });

webapp/platform/client/src/websocket.ts

@@ -16,6 +16,7 @@ export type WebSocketClientConfig = {
     maxWebSocketRetryTime: number;
     reconnectJitterRange: number;
     newWebSocketFn: (url: string) => WebSocket;
+    clientPingInterval: number;
 }
 
 const defaultWebSocketClientConfig: WebSocketClientConfig = {
@@ -26,6 +27,7 @@ const defaultWebSocketClientConfig: WebSocketClientConfig = {
     newWebSocketFn: (url: string) => {
         return new WebSocket(url);
     },
+    clientPingInterval: 30000, // 30 seconds
 };
 
 export default class WebSocketClient {
@@ -86,6 +88,8 @@ export default class WebSocketClient {
     private serverHostname: string | null;
     private postedAck: boolean;
 
+    private pingInterval: ReturnType<typeof setInterval> | null;
+
     private reconnectTimeout: ReturnType<typeof setTimeout> | null;
 
     constructor(config?: Partial<WebSocketClientConfig>) {
@@ -100,6 +104,7 @@ export default class WebSocketClient {
         this.postedAck = false;
         this.reconnectTimeout = null;
         this.config = {...defaultWebSocketClientConfig, ...config};
+        this.pingInterval = null;
     }
 
     // on connect, only send auth cookie and blank state.
@@ -156,6 +161,28 @@ export default class WebSocketClient {
                 this.firstConnectListeners.forEach((listener) => listener());
             }
 
+            this.stopPingInterval();
+            var waitingForPong = false;
+            this.pingInterval = setInterval(
+                () => {
+                    if (!waitingForPong) {
+                        waitingForPong = true;
+                        this.ping(() => {
+                            waitingForPong = false;
+                        });
+                        return;
+                    }
+
+                    console.log('ping received no response within time limit: re-establishing websocket'); //eslint-disable-line no-console
+
+                    // We are not calling this.close() because we need to auto-restart.
+                    this.connectFailCount = 0;
+                    this.responseSequence = 1;
+                    this.stopPingInterval();
+                    this.conn?.close(); // Will auto-reconnect after configured retry time
+                },
+                this.config.clientPingInterval);
+
             this.connectFailCount = 0;
         };
 
@@ -172,6 +199,9 @@ export default class WebSocketClient {
             this.closeCallback?.(this.connectFailCount);
             this.closeListeners.forEach((listener) => listener(this.connectFailCount));
 
+            // Make sure we stop pinging if the connection is closed
+            this.stopPingInterval();
+
             // If we've failed a bunch of connections then start backing off
             let retryTime = this.config.minWebSocketRetryTime;
             if (this.connectFailCount > this.config.maxWebSocketFails) {
@@ -399,6 +429,7 @@ export default class WebSocketClient {
             clearTimeout(this.reconnectTimeout);
             this.reconnectTimeout = null;
         }
+        this.stopPingInterval();
         if (this.conn && this.conn.readyState === WebSocket.OPEN) {
             this.conn.onclose = () => {};
             this.conn.close();
@@ -407,6 +438,28 @@ export default class WebSocketClient {
         }
     }
 
+    stopPingInterval() {
+        if (this.pingInterval) {
+            clearInterval(this.pingInterval);
+            this.pingInterval = null;
+        }
+    }
+
+    ping(responseCallback?: (msg: any) => void) {
+        const msg = {
+            action: 'ping',
+            seq: this.responseSequence++,
+        };
+
+        if (responseCallback) {
+            this.responseCallbacks[msg.seq] = responseCallback;
+        }
+
+        if (this.conn && this.conn.readyState === WebSocket.OPEN) {
+            this.conn.send(JSON.stringify(msg));
+        }
+    }
+
     sendMessage(action: string, data: any, responseCallback?: (msg: any) => void) {
         const msg = {
             action,