From 1dbd1fa4db5fb239475c2341622bd5e28c7ee906 Mon Sep 17 00:00:00 2001 From: David Krauser Date: Thu, 20 Feb 2025 13:54:18 -0500 Subject: [PATCH] [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. --- ...post_body_additional_content.test.tsx.snap | 16 ++ webapp/platform/client/src/websocket.test.ts | 212 ++++++++++++++++++ webapp/platform/client/src/websocket.ts | 78 +++++-- 3 files changed, 287 insertions(+), 19 deletions(-) create mode 100644 webapp/platform/client/src/websocket.test.ts diff --git a/webapp/channels/src/components/post_view/post_body_additional_content/__snapshots__/post_body_additional_content.test.tsx.snap b/webapp/channels/src/components/post_view/post_body_additional_content/__snapshots__/post_body_additional_content.test.tsx.snap index 05e879c479..8de44734d2 100644 --- a/webapp/channels/src/components/post_view/post_body_additional_content/__snapshots__/post_body_additional_content.test.tsx.snap +++ b/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": "", diff --git a/webapp/platform/client/src/websocket.test.ts b/webapp/platform/client/src/websocket.test.ts new file mode 100644 index 0000000000..54e10b0478 --- /dev/null +++ b/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(); + }); +}); diff --git a/webapp/platform/client/src/websocket.ts b/webapp/platform/client/src/websocket.ts index bd478f291b..b4535b7353 100644 --- a/webapp/platform/client/src/websocket.ts +++ b/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 | null; + + constructor(config?: Partial) { 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) { - retryTime = MIN_WEBSOCKET_RETRY_TIME * this.connectFailCount * this.connectFailCount; - if (retryTime > MAX_WEBSOCKET_RETRY_TIME) { - retryTime = MAX_WEBSOCKET_RETRY_TIME; + let retryTime = this.config.minWebSocketRetryTime; + if (this.connectFailCount > this.config.maxWebSocketFails) { + retryTime = retryTime * this.connectFailCount * this.connectFailCount; + 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(); } }