[MM-62924] WebSocketClient reconnection unit tests (#30135)

This commit adds some reconnection-specific unit tests to the WebSocketClient class, and fixes a few minor bugs:

- If `close()` is called during a reconnection delay, the socket won't close.
- If `send()` is called during a reconnection delay, we immediately try to reconnect (and don't respect any configured delays).
- If `initialize()` is called during a reconnection delay, we immediately try to reconnect (and don't respect any configured delays).
- If we receive two disconnection events, we'll attempt to reconnect with two new connections.

To allow for testing, I also needed to make the WebSocketClient more configurable. Specifically, the config allows for mocking the underlying websockets, and to control timeout delays.

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

@@ -99,6 +99,13 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
       WebSocketClient {
         "closeCallback": null,
         "closeListeners": Set {},
+        "config": Object {
+          "maxWebSocketFails": 7,
+          "maxWebSocketRetryTime": 300000,
+          "minWebSocketRetryTime": 3000,
+          "newWebSocketFn": [Function],
+          "reconnectJitterRange": 2000,
+        },
         "conn": null,
         "connectFailCount": 0,
         "connectionId": "",
@@ -114,6 +121,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "postedAck": false,
         "reconnectCallback": null,
         "reconnectListeners": Set {},
+        "reconnectTimeout": null,
         "responseCallbacks": Object {},
         "responseSequence": 1,
         "serverHostname": "",
@@ -147,6 +155,13 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
       WebSocketClient {
         "closeCallback": null,
         "closeListeners": Set {},
+        "config": Object {
+          "maxWebSocketFails": 7,
+          "maxWebSocketRetryTime": 300000,
+          "minWebSocketRetryTime": 3000,
+          "newWebSocketFn": [Function],
+          "reconnectJitterRange": 2000,
+        },
         "conn": null,
         "connectFailCount": 0,
         "connectionId": "",
@@ -162,6 +177,7 @@ exports[`PostBodyAdditionalContent with a normal link Should render the plugin c
         "postedAck": false,
         "reconnectCallback": null,
         "reconnectListeners": Set {},
+        "reconnectTimeout": null,
         "responseCallbacks": Object {},
         "responseSequence": 1,
         "serverHostname": "",

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

@@ -0,0 +1,212 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+import WebSocketClient from './websocket';
+
+// Define some WebSocket globals that aren't defined in node
+if (typeof WebSocket === 'undefined') {
+    (global as any).WebSocket = {
+        CONNECTING: 0, OPEN: 1, CLOSING: 2, CLOSED: 3,
+    };
+}
+
+class MockWebSocket {
+    readonly binaryType: BinaryType = 'blob';
+    readonly bufferedAmount: number = 0;
+    readonly extensions: string = '';
+
+    readonly CONNECTING = WebSocket.CONNECTING;
+    readonly OPEN = WebSocket.OPEN;
+    readonly CLOSING = WebSocket.CLOSING;
+    readonly CLOSED = WebSocket.CLOSED;
+
+    public url: string = '';
+    readonly protocol: string = '';
+    public readyState: number = WebSocket.CONNECTING;
+
+    public onopen: (() => void) | null = null;
+    public onclose: (() => void) | null = null;
+    public onerror: (() => void) | null = null;
+    public onmessage: ((evt: any) => void) | null = null;
+
+    open() {
+        this.readyState = WebSocket.OPEN;
+        if (this.onopen) {
+            this.onopen();
+        }
+    }
+
+    close() {
+        this.readyState = WebSocket.CLOSED;
+        if (this.onclose) {
+            this.onclose();
+        }
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    send(msg: any) { }
+    addEventListener() { }
+    removeEventListener() { }
+    dispatchEvent(): boolean {
+        return false;
+    }
+}
+
+describe('websocketclient', () => {
+    test('initialize should register connection callbacks', () => {
+        const mockWebSocket = new MockWebSocket();
+
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                return mockWebSocket;
+            },
+        });
+
+        client.initialize('mock.url');
+
+        expect(mockWebSocket.onopen).toBeTruthy();
+        expect(mockWebSocket.onclose).toBeTruthy();
+
+        client.close();
+    });
+
+    test('should reconnect on websocket close', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                mockWebSocket.open();
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 10,
+            reconnectJitterRange: 10,
+        });
+
+        client.initialize('mock.url');
+        expect(openSpy).toHaveBeenCalledTimes(1);
+
+        mockWebSocket.close();
+
+        jest.advanceTimersByTime(40);
+
+        client.close();
+        expect(openSpy).toHaveBeenCalledTimes(2);
+
+        jest.useRealTimers();
+    });
+
+    test('should close during reconnection delay', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 50,
+            reconnectJitterRange: 1,
+        });
+
+        const initializeSpy = jest.spyOn(client, 'initialize');
+        client.initialize('mock.url');
+        mockWebSocket.open();
+        mockWebSocket.close();
+
+        jest.advanceTimersByTime(10);
+
+        client.close();
+
+        jest.advanceTimersByTime(80);
+
+        client.close();
+        expect(initializeSpy).toBeCalledTimes(1);
+        expect(openSpy).toBeCalledTimes(1);
+
+        jest.useRealTimers();
+    });
+
+    test('should not re-open if initialize called during reconnection delay', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 50,
+            reconnectJitterRange: 1,
+        });
+
+        const initializeSpy = jest.spyOn(client, 'initialize');
+        client.initialize('mock.url');
+        mockWebSocket.open();
+        mockWebSocket.close();
+
+        jest.advanceTimersByTime(10);
+
+        client.initialize('mock.url');
+        expect(initializeSpy).toBeCalledTimes(2);
+        expect(openSpy).toBeCalledTimes(1);
+
+        jest.advanceTimersByTime(80);
+
+        client.close();
+        expect(initializeSpy).toBeCalledTimes(3);
+        expect(openSpy).toBeCalledTimes(2);
+
+        jest.useRealTimers();
+    });
+
+    test('should not register second reconnection timeout if onclose called twice', () => {
+        jest.useFakeTimers();
+
+        const mockWebSocket = new MockWebSocket();
+        const openSpy = jest.spyOn(mockWebSocket, 'open');
+
+        const client = new WebSocketClient({
+            newWebSocketFn: (url: string) => {
+                mockWebSocket.url = url;
+                if (mockWebSocket.onopen) {
+                    mockWebSocket.open();
+                }
+                return mockWebSocket;
+            },
+            minWebSocketRetryTime: 50,
+            reconnectJitterRange: 1,
+        });
+
+        const initializeSpy = jest.spyOn(client, 'initialize');
+        client.initialize('mock.url');
+        mockWebSocket.open();
+        mockWebSocket.close();
+
+        jest.advanceTimersByTime(10);
+
+        mockWebSocket.close();
+
+        jest.advanceTimersByTime(80);
+
+        client.close();
+        expect(initializeSpy).toBeCalledTimes(2);
+        expect(openSpy).toBeCalledTimes(2);
+
+        jest.useRealTimers();
+    });
+});

webapp/platform/client/src/websocket.ts

@@ -1,11 +1,6 @@
 // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
 // See LICENSE.txt for license information.
 
-const MAX_WEBSOCKET_FAILS = 7;
-const MIN_WEBSOCKET_RETRY_TIME = 3000; // 3 sec
-const MAX_WEBSOCKET_RETRY_TIME = 300000; // 5 mins
-const JITTER_RANGE = 2000; // 2 sec
-
 const WEBSOCKET_HELLO = 'hello';
 
 export type MessageListener = (msg: WebSocketMessage) => void;
@@ -15,7 +10,27 @@ export type MissedMessageListener = () => void;
 export type ErrorListener = (event: Event) => void;
 export type CloseListener = (connectFailCount: number) => void;
 
+export type WebSocketClientConfig = {
+    maxWebSocketFails: number;
+    minWebSocketRetryTime: number;
+    maxWebSocketRetryTime: number;
+    reconnectJitterRange: number;
+    newWebSocketFn: (url: string) => WebSocket;
+}
+
+const defaultWebSocketClientConfig: WebSocketClientConfig = {
+    maxWebSocketFails: 7,
+    minWebSocketRetryTime: 3000, // 3 seconds
+    maxWebSocketRetryTime: 300000, // 5 minutes
+    reconnectJitterRange: 2000, // 2 seconds
+    newWebSocketFn: (url: string) => {
+        return new WebSocket(url);
+    },
+};
+
 export default class WebSocketClient {
+    private config: WebSocketClientConfig;
+
     private conn: WebSocket | null;
     private connectionUrl: string | null;
 
@@ -71,7 +86,9 @@ export default class WebSocketClient {
     private serverHostname: string | null;
     private postedAck: boolean;
 
-    constructor() {
+    private reconnectTimeout: ReturnType<typeof setTimeout> | null;
+
+    constructor(config?: Partial<WebSocketClientConfig>) {
         this.conn = null;
         this.connectionUrl = null;
         this.responseSequence = 1;
@@ -81,6 +98,8 @@ export default class WebSocketClient {
         this.connectionId = '';
         this.serverHostname = '';
         this.postedAck = false;
+        this.reconnectTimeout = null;
+        this.config = {...defaultWebSocketClientConfig, ...config};
     }
 
     // on connect, only send auth cookie and blank state.
@@ -91,6 +110,13 @@ export default class WebSocketClient {
             return;
         }
 
+        // We have a timeout waiting to re-initialize the websocket.
+        // We should wait until that fires before initializing,
+        // otherwise we may not respect the configured backoff.
+        if (this.reconnectTimeout) {
+            return;
+        }
+
         if (connectionUrl == null) {
             console.log('websocket must have connection url'); //eslint-disable-line no-console
             return;
@@ -107,7 +133,12 @@ export default class WebSocketClient {
         // Add connection id, and last_sequence_number to the query param.
         // We cannot use a cookie because it will bleed across tabs.
         // We cannot also send it as part of the auth_challenge, because the session cookie is already sent with the request.
-        this.conn = new WebSocket(`${connectionUrl}?connection_id=${this.connectionId}&sequence_number=${this.serverSequence}${this.postedAck ? '&posted_ack=true' : ''}`);
+        const websocketUrl = `${connectionUrl}?connection_id=${this.connectionId}&sequence_number=${this.serverSequence}${this.postedAck ? '&posted_ack=true' : ''}`;
+        if (this.config.newWebSocketFn) {
+            this.conn = this.config.newWebSocketFn(websocketUrl);
+        } else {
+            this.conn = new WebSocket(websocketUrl);
+        }
         this.connectionUrl = connectionUrl;
 
         this.conn.onopen = () => {
@@ -141,22 +172,28 @@ export default class WebSocketClient {
             this.closeCallback?.(this.connectFailCount);
             this.closeListeners.forEach((listener) => listener(this.connectFailCount));
 
-            let retryTime = MIN_WEBSOCKET_RETRY_TIME;
-
             // If we've failed a bunch of connections then start backing off
-            if (this.connectFailCount > MAX_WEBSOCKET_FAILS) {
+            let retryTime = this.config.minWebSocketRetryTime;
-                retryTime = MIN_WEBSOCKET_RETRY_TIME * this.connectFailCount * this.connectFailCount;
+            if (this.connectFailCount > this.config.maxWebSocketFails) {
-                if (retryTime > MAX_WEBSOCKET_RETRY_TIME) {
+                retryTime = retryTime * this.connectFailCount * this.connectFailCount;
-                    retryTime = MAX_WEBSOCKET_RETRY_TIME;
+                if (retryTime > this.config.maxWebSocketRetryTime) {
+                    retryTime = this.config.maxWebSocketRetryTime;
                 }
             }
 
             // Applying jitter to avoid thundering herd problems.
-            retryTime += Math.random() * JITTER_RANGE;
+            retryTime += Math.random() * this.config.reconnectJitterRange;
 
-            setTimeout(
+            // If we already have a reconnect timeout waiting,
+            // we should let that handle the next connection.
+            if (this.reconnectTimeout) {
+                return;
+            }
+
+            this.reconnectTimeout = setTimeout(
                 () => {
-                    this.initialize(connectionUrl, token, postedAck);
+                    this.reconnectTimeout = null;
+                    this.initialize(this.connectionUrl, token, this.postedAck);
                 },
                 retryTime,
             );
@@ -358,6 +395,10 @@ export default class WebSocketClient {
     close() {
         this.connectFailCount = 0;
         this.responseSequence = 1;
+        if (this.reconnectTimeout) {
+            clearTimeout(this.reconnectTimeout);
+            this.reconnectTimeout = null;
+        }
         if (this.conn && this.conn.readyState === WebSocket.OPEN) {
             this.conn.onclose = () => {};
             this.conn.close();
@@ -377,11 +418,10 @@ export default class WebSocketClient {
             this.responseCallbacks[msg.seq] = responseCallback;
         }
 
+        // Only try to send the message if the websocket is open.
+        // If the websocket is closed here, we will drop the message.
         if (this.conn && this.conn.readyState === WebSocket.OPEN) {
             this.conn.send(JSON.stringify(msg));
-        } else if (!this.conn || this.conn.readyState === WebSocket.CLOSED) {
-            this.conn = null;
-            this.initialize();
         }
     }