From c27f0eade200daecf69aeb1b503a0c9328fbd96a Mon Sep 17 00:00:00 2001 From: David Krauser Date: Tue, 11 Mar 2025 11:48:09 -0400 Subject: [PATCH] [MM-62925] introduce websocket client-side ping (#8633) 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. --- app/client/websocket/index.test.ts | 244 +++++++++++++++++++++++++++-- app/client/websocket/index.ts | 44 +++++- app/constants/websocket.ts | 2 + test/timer_helpers.ts | 20 +++ 4 files changed, 290 insertions(+), 20 deletions(-) create mode 100644 test/timer_helpers.ts diff --git a/app/client/websocket/index.test.ts b/app/client/websocket/index.test.ts index 4a8e0f74b..6100388ea 100644 --- a/app/client/websocket/index.test.ts +++ b/app/client/websocket/index.test.ts @@ -6,6 +6,7 @@ import {WebSocketReadyState, getOrCreateWebSocketClient} from '@mattermost/react import {WebsocketEvents} from '@constants'; import DatabaseManager from '@database/manager'; import {getConfigValue} from '@queries/servers/system'; +import {enableFakeTimers, disableFakeTimers, advanceTimers} from '@test/timer_helpers'; import {hasReliableWebsocket} from '@utils/config'; import {logDebug, logInfo, logError, logWarning} from '@utils/log'; @@ -44,27 +45,41 @@ const mockedHasReliableWebsocket = jest.mocked(hasReliableWebsocket); const mockedGetOrCreateWebSocketClient = jest.mocked(getOrCreateWebSocketClient); describe('WebSocketClient', () => { - let client: WebSocketClient; const serverUrl = 'https://example.com'; const token = 'test-token'; - const mockConn = { + const createMockConn = () => ({ onOpen: jest.fn(), onClose: jest.fn(), onError: jest.fn(), onMessage: jest.fn(), - open: jest.fn(), - close: jest.fn(), - invalidate: jest.fn(), send: jest.fn(), - readyState: WebSocketReadyState.OPEN, - }; - const mockClient = {client: mockConn}; - mockedGetOrCreateWebSocketClient.mockResolvedValue(mockClient as any); - mockedHasReliableWebsocket.mockReturnValue(false); + readyState: WebSocketReadyState.CLOSED, + open() { + this.readyState = WebSocketReadyState.OPEN; + this.onOpen.mock.calls[0][0]({}); + }, + close() { + this.readyState = WebSocketReadyState.CLOSED; + this.onClose.mock.calls[0][0]({}); + }, + }); + + let mockConn: ReturnType; + let client: WebSocketClient; beforeEach(() => { + mockConn = createMockConn(); + const mockClient = {client: mockConn}; + mockedGetOrCreateWebSocketClient.mockResolvedValue(mockClient as any); + mockedHasReliableWebsocket.mockReturnValue(false); client = new WebSocketClient(serverUrl, token); + enableFakeTimers(); + }); + + afterEach(() => { + client.close(true); + disableFakeTimers(); }); it('should initialize the WebSocketClient', async () => { @@ -85,7 +100,7 @@ describe('WebSocketClient', () => { await client.initialize({}, true); - mockConn.onOpen.mock.calls[0][0](); + expect(mockConn.readyState).toBe(WebSocketReadyState.OPEN); expect(logInfo).toHaveBeenCalledWith('websocket connected to', 'wss://example.com/api/v4/websocket'); expect(firstConnectCallback).toHaveBeenCalled(); @@ -97,8 +112,6 @@ describe('WebSocketClient', () => { await client.initialize(); - mockConn.onOpen.mock.calls[0][0](); - expect(logInfo).toHaveBeenCalledWith('websocket re-established connection to', 'wss://example.com/api/v4/websocket'); expect(reconnectCallback).toHaveBeenCalled(); }); @@ -130,12 +143,35 @@ describe('WebSocketClient', () => { await client.initialize(); - mockConn.onClose.mock.calls[0][0]({}); + mockConn.close(); expect(logInfo).toHaveBeenCalledWith('websocket closed', 'wss://example.com/api/v4/websocket'); expect(closeCallback).toHaveBeenCalled(); }); + it('should handle WebSocket close event - reconnect', async () => { + enableFakeTimers(); + + const closeCallback = jest.fn(); + client.setCloseCallback(closeCallback); + + const connectingCallback = jest.fn(); + client.setConnectingCallback(connectingCallback); + + await client.initialize(); + + expect(connectingCallback).toHaveBeenCalledTimes(1); + expect(closeCallback).toHaveBeenCalledTimes(0); + + mockConn.close(); + + await advanceTimers(6000); + + expect(connectingCallback).toHaveBeenCalledTimes(2); + expect(closeCallback).toHaveBeenCalledTimes(1); + expect(mockConn.readyState).toBe(WebSocketReadyState.OPEN); + }); + it('should handle WebSocket close event - tls handshake error', async () => { await client.initialize(); const message = {code: 1015, reason: 'tls handshake error'}; @@ -242,9 +278,16 @@ describe('WebSocketClient', () => { client.sendUserTypingEvent('channel1', 'parent1'); - expect(mockConn.send).toHaveBeenCalledWith(JSON.stringify({ - action: 'user_typing', + expect(mockConn.send).toHaveBeenNthCalledWith(1, JSON.stringify({ + action: 'authentication_challenge', seq: 1, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(2, JSON.stringify({ + action: 'user_typing', + seq: 2, data: { channel_id: 'channel1', parent_id: 'parent1', @@ -253,7 +296,7 @@ describe('WebSocketClient', () => { }); it('should fail to send user typing event', async () => { - client.close(); + client.close(true); client.sendUserTypingEvent('channel1', 'parent1'); expect(mockConn.send).not.toHaveBeenCalled(); @@ -264,4 +307,171 @@ describe('WebSocketClient', () => { expect(client.isConnected()).toBe(true); }); + + it('should send ping messages on interval and handle pong responses', async () => { + await client.initialize(); + + // Wait until we get a client PING + await advanceTimers(30100); + expect(mockConn.send).toHaveBeenNthCalledWith(1, JSON.stringify({ + action: 'authentication_challenge', + seq: 1, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(2, JSON.stringify({ + action: 'ping', + seq: 2, + })); + + // Second ping should be sent if we got a pong response + const pongMessage = {data: {text: WebsocketEvents.PONG}, seq_reply: 2, status: 'OK'}; + mockConn.onMessage.mock.calls[0][0]({message: pongMessage}); + + // Second ping should be sent if we got a pong response + await advanceTimers(30100); + expect(mockConn.send).toHaveBeenNthCalledWith(3, JSON.stringify({ + action: 'ping', + seq: 3, + })); + + // Verify ping sequence increments + const pongMessage2 = {data: {text: WebsocketEvents.PONG}, seq_reply: 3, status: 'OK'}; + mockConn.onMessage.mock.calls[0][0]({message: pongMessage2}); + + await advanceTimers(30100); + expect(mockConn.send).toHaveBeenNthCalledWith(4, JSON.stringify({ + action: 'ping', + seq: 4, + })); + }); + + it('should handle ping timeouts and reconnect', async () => { + await client.initialize(); + + // Send first ping + await advanceTimers(30100); + expect(mockConn.send).toHaveBeenNthCalledWith(1, JSON.stringify({ + action: 'authentication_challenge', + seq: 1, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(2, JSON.stringify({ + action: 'ping', + seq: 2, + })); + + // No pong received, next interval should trigger close + await advanceTimers(30100); + expect(mockConn.onClose).toHaveBeenCalled(); + + // Reset mock and verify reconnect behavior + mockConn.onClose.mockClear(); + mockConn.send.mockClear(); + + // Should attempt to reconnect after timeout + await advanceTimers(3000); + + // Should start pinging again after reconnect + mockConn.onOpen.mock.calls[0][0](); + await advanceTimers(30000); + + expect(mockConn.send).toHaveBeenNthCalledWith(2, JSON.stringify({ + action: 'authentication_challenge', + seq: 2, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(3, JSON.stringify({ + action: 'ping', + seq: 3, + })); + }); + + it('should clear ping interval on close', async () => { + enableFakeTimers(); + + await client.initialize(); + mockConn.onOpen.mock.calls[0][0](); // Complete the connection + mockConn.send.mockClear(); // Clear the initial authentication call + + // Advance timer - no ping should be sent + await advanceTimers(20000); + + client.close(true); + + // Advance timer - no ping should be sent + await advanceTimers(20000); + + expect(mockConn.send).not.toHaveBeenCalled(); + }); + + it('should reset ping interval state when reconnecting during pending ping', async () => { + enableFakeTimers(); + + const closeCallback = jest.fn(); + client.setCloseCallback(closeCallback); + + const connectingCallback = jest.fn(); + client.setConnectingCallback(connectingCallback); + + await client.initialize(); + + // Wait until we get a client PING + await advanceTimers(30100); + + expect(mockConn.send).toHaveBeenNthCalledWith(1, JSON.stringify({ + action: 'authentication_challenge', + seq: 1, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(2, JSON.stringify({ + action: 'ping', + seq: 2, + })); + + mockConn.close(); + + // Let connection reconnect + await advanceTimers(5100); + + // And client for PINGs to start again + await advanceTimers(30100); + + expect(connectingCallback).toHaveBeenCalledTimes(2); + expect(closeCallback).toHaveBeenCalledTimes(1); + expect(mockConn.send).toHaveBeenNthCalledWith(3, JSON.stringify({ + action: 'authentication_challenge', + seq: 1, + data: { + token: 'test-token', + }, + })); + expect(mockConn.send).toHaveBeenNthCalledWith(4, JSON.stringify({ + action: 'ping', + seq: 2, + })); + + // Second ping should be sent if we got a pong response + const pongMessage = {data: {text: WebsocketEvents.PONG}, seq_reply: 2, status: 'OK'}; + mockConn.onMessage.mock.calls[0][0]({message: pongMessage}); + + // Ensure we continue to get client PINGs + await advanceTimers(30100); + expect(mockConn.send).toHaveBeenNthCalledWith(5, JSON.stringify({ + action: 'ping', + seq: 3, + })); + + expect(connectingCallback).toHaveBeenCalledTimes(2); + expect(closeCallback).toHaveBeenCalledTimes(1); + + disableFakeTimers(); + }); }); diff --git a/app/client/websocket/index.ts b/app/client/websocket/index.ts index fd56a7320..79bf3c6be 100644 --- a/app/client/websocket/index.ts +++ b/app/client/websocket/index.ts @@ -15,6 +15,7 @@ const MAX_WEBSOCKET_FAILS = 7; const WEBSOCKET_TIMEOUT = toMilliseconds({seconds: 30}); const MIN_WEBSOCKET_RETRY_TIME = toMilliseconds({seconds: 3}); const MAX_WEBSOCKET_RETRY_TIME = toMilliseconds({minutes: 5}); +const PING_INTERVAL = toMilliseconds({seconds: 30}); const DEFAULT_OPTIONS = { forceConnection: true, }; @@ -30,6 +31,9 @@ export default class WebSocketClient { private serverUrl: string; private connectFailCount = 0; + private pingInterval: NodeJS.Timeout | undefined; + private waitingForPong: boolean = false; + // The first time we connect to a server (on init or login) // we do the sync out of the websocket lifecycle. // This is used to avoid calling twice to the sync logic. @@ -138,6 +142,7 @@ export default class WebSocketClient { // the websocket will call onClose then onError then initialize again with readyState CLOSED, we need to open it again if (this.conn.readyState === WebSocketReadyState.CLOSED) { clearTimeout(this.connectionTimeout); + clearInterval(this.pingInterval); this.conn.open(); } return; @@ -149,6 +154,7 @@ export default class WebSocketClient { this.conn!.onOpen(() => { clearTimeout(this.connectionTimeout); + clearInterval(this.pingInterval); // No need to reset sequence number here. if (!reliableWebSockets) { @@ -177,11 +183,30 @@ export default class WebSocketClient { } } + this.waitingForPong = false; + this.pingInterval = setInterval( + () => { + if (!this.waitingForPong) { + this.waitingForPong = true; + this.ping(); + return; + } + + // We are not calling this.close() because we need to auto-restart. + this.responseSequence = 1; + clearInterval(this.pingInterval); + this.conn?.close(); + }, + PING_INTERVAL, + ); + this.connectFailCount = 0; }); this.conn!.onClose((ev) => { clearTimeout(this.connectionTimeout); + clearInterval(this.pingInterval); + this.conn = undefined; this.responseSequence = 1; @@ -253,6 +278,9 @@ export default class WebSocketClient { if (msg.error) { logWarning(msg); } + if (msg.data?.text === WebsocketEvents.PONG) { + this.waitingForPong = false; + } } else if (this.eventCallback) { if (reliableWebSockets) { // We check the hello packet, which is always the first packet in a stream. @@ -339,15 +367,28 @@ export default class WebSocketClient { this.connectFailCount = 0; this.responseSequence = 1; clearTimeout(this.connectionTimeout); + clearInterval(this.pingInterval); this.conn?.close(); } public invalidate() { clearTimeout(this.connectionTimeout); + clearInterval(this.pingInterval); this.conn?.invalidate(); this.conn = undefined; } + private ping() { + const msg = { + action: WebsocketEvents.PING, + seq: this.responseSequence++, + }; + + if (this.conn && this.conn.readyState === WebSocketReadyState.OPEN) { + this.conn.send(JSON.stringify(msg)); + } + } + private sendMessage(action: string, data: any) { const msg = { action, @@ -357,9 +398,6 @@ export default class WebSocketClient { if (this.conn && this.conn.readyState === WebSocketReadyState.OPEN) { this.conn.send(JSON.stringify(msg)); - } else if (!this.conn || this.conn.readyState === WebSocketReadyState.CLOSED) { - this.conn = undefined; - this.initialize(this.token); } } diff --git a/app/constants/websocket.ts b/app/constants/websocket.ts index d4034e213..9eebe9bf0 100644 --- a/app/constants/websocket.ts +++ b/app/constants/websocket.ts @@ -40,6 +40,8 @@ const WebsocketEvents = { PREFERENCES_DELETED: 'preferences_deleted', EPHEMERAL_MESSAGE: 'ephemeral_message', STATUS_CHANGED: 'status_change', + PING: 'ping', + PONG: 'pong', HELLO: 'hello', WEBRTC: 'webrtc', REACTION_ADDED: 'reaction_added', diff --git a/test/timer_helpers.ts b/test/timer_helpers.ts new file mode 100644 index 000000000..82c0e5e83 --- /dev/null +++ b/test/timer_helpers.ts @@ -0,0 +1,20 @@ +// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved. +// See LICENSE.txt for license information. + +// Helper functions to handle Jest timer issues with async code +// These are needed because Jest's fake timers don't play well with Promise-based code. +// The combination of fake timers for time advancement + real timers for nextTick +// allows us to properly test async timing behavior. + +export const enableFakeTimers = () => { + jest.useFakeTimers({doNotFake: ['nextTick']}); +}; + +export const disableFakeTimers = () => { + jest.useRealTimers(); +}; + +export const advanceTimers = async (ms: number) => { + jest.advanceTimersByTime(ms); + await new Promise(process.nextTick); +};