diff --git a/README.md b/README.md index 5e1c154..c679ef0 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,58 @@ stock.on('message', (message) => { }); ``` +### Health Check + +The WebSocket client can monitor connection liveness using an app-level +JSON ping/pong (`{ event: 'ping' }` / `{ event: 'pong' }`). It is disabled by +default. When enabled, on each interval tick the client checks whether **any** +inbound message has arrived since the last ping it sent (freshness check): + +- If nothing arrived since the last ping, that counts as a *miss* and the + consecutive-miss counter is incremented. +- If any inbound message (a `pong`, market data, or anything else) arrived, + the counter is reset to `0`. +- Once the consecutive-miss counter reaches `maxMissedPongs`, the client + disconnects and stops the timer. + +| Option | Type | Default | Description | +| ---------------- | --------- | ------- | ----------------------------------------------------------------- | +| `enabled` | `boolean` | `false` | Enables the health-check ping/pong. | +| `pingInterval` | `number` | `30000` | Interval in milliseconds between health-check pings. | +| `maxMissedPongs` | `number` | `2` | Consecutive misses (no inbound messages) before disconnecting. | + +```js +const client = new WebSocketClient({ + apiKey: 'YOUR_API_KEY', + healthCheck: { + enabled: true, + pingInterval: 30000, + maxMissedPongs: 2, + }, +}); +``` + +#### Disconnect reason + +When the client disconnects because of a health-check timeout, the +`disconnect` event carries a second argument `{ reason: 'health-check-timeout' }`. +Normal or manual disconnects emit `disconnect` **without** a second argument +(it is `undefined`). Listeners that only read the first argument (the close +event) continue to work unchanged. + +```js +const stock = client.stock; + +stock.on('disconnect', (event, info) => { + if (info?.reason === 'health-check-timeout') { + console.log('Health check timed out, reconnecting...'); + stock.connect().then(() => { + stock.subscribe({ channel: 'trades', symbol: '2330' }); + }); + } +}); +``` + ## License [MIT](LICENSE) diff --git a/src/websocket/client.ts b/src/websocket/client.ts index d3f8f0b..36e714d 100644 --- a/src/websocket/client.ts +++ b/src/websocket/client.ts @@ -16,10 +16,17 @@ export interface WebSocketClientOptions { healthCheck?: HealthCheckConfig; } +export interface DisconnectReason { + reason: 'health-check-timeout'; +} + export class WebSocketClient extends events.EventEmitter { private socket!: WebSocket; - private missedPongs = 0; + private consecutiveMisses = 0; + private lastMessageAt = 0; + private lastPingAt = 0; private pingTimerId: ReturnType | undefined; + private pendingDisconnectReason: DisconnectReason | undefined; constructor(protected readonly options: WebSocketClientOptions) { super(); @@ -28,9 +35,17 @@ export class WebSocketClient extends events.EventEmitter { public connect() { this.socket = new WebSocket(this.options.url); this.socket.onopen = () => this.emit(CONNECT_EVENT); - this.socket.onmessage = event => this.emit(MESSAGE_EVENT, event.data); + this.socket.onmessage = event => { + // Any inbound message counts as freshness. + this.lastMessageAt = Date.now(); + this.emit(MESSAGE_EVENT, event.data); + }; this.socket.onerror = event => this.emit(ERROR_EVENT, event.error); - this.socket.onclose = event => this.emit(DISCONNECT_EVENT, event); + this.socket.onclose = event => { + const reason = this.pendingDisconnectReason; + this.pendingDisconnectReason = undefined; + this.emit(DISCONNECT_EVENT, event, reason); + }; this.on(CONNECT_EVENT, () => this.authenticate()); this.on(MESSAGE_EVENT, message => this.handleMessage(message)); @@ -44,22 +59,36 @@ export class WebSocketClient extends events.EventEmitter { this.socket.close(); if (this.pingTimerId) { clearInterval(this.pingTimerId); + this.pingTimerId = undefined; } } - private detectConnectionStatus(state?: string) { + private detectConnectionStatus() { if (!this.options.healthCheck?.enabled) return; + // Clamp to >= 1: maxMissedPongs of 0 would disconnect a healthy connection + // on the first tick (consecutiveMisses starts at 0, and 0 >= 0). + const maxMissed = Math.max(1, this.options.healthCheck.maxMissedPongs ?? 2); + + // Freshness check: did anything arrive since our last ping? + if (this.lastMessageAt < this.lastPingAt) { + this.consecutiveMisses += 1; + } else { + this.consecutiveMisses = 0; + } + + if (this.consecutiveMisses >= maxMissed) { + this.pendingDisconnectReason = { reason: 'health-check-timeout' }; + this.disconnect(); + return; + } + try { - this.ping({ state }); - this.missedPongs += 1; - const maxMissed = this.options.healthCheck.maxMissedPongs ?? 2; - if (this.missedPongs > maxMissed) { - this.disconnect(); - return; - } + this.ping({}); + this.lastPingAt = Date.now(); } catch (error) { console.error(`Failed to send ping: ${error}`); + this.pendingDisconnectReason = { reason: 'health-check-timeout' }; this.disconnect(); return; } @@ -107,6 +136,11 @@ export class WebSocketClient extends events.EventEmitter { // Start health check if enabled if (this.options.healthCheck?.enabled) { const interval = this.options.healthCheck.pingInterval ?? 30000; + this.consecutiveMisses = 0; + // Seed timestamps so the first tick treats the connection as fresh. + const now = Date.now(); + this.lastMessageAt = now; + this.lastPingAt = now; this.pingTimerId = setInterval(() => { this.detectConnectionStatus(); }, interval); @@ -117,9 +151,6 @@ export class WebSocketClient extends events.EventEmitter { this.emit(UNAUTHENTICATED_EVENT, data); } } - if (event === 'pong') { - this.missedPongs = 0; - } } catch (err) {} } } diff --git a/test/websocket-client.spec.ts b/test/websocket-client.spec.ts index e408144..d19e7d2 100644 --- a/test/websocket-client.spec.ts +++ b/test/websocket-client.spec.ts @@ -541,7 +541,7 @@ describe('WebSocketClient', () => { stock.disconnect(); }); - it('should handle detectConnectionStatus when health check is disabled', async () => { + it('should not send any ping when health check is disabled', async () => { const client = new WebSocketClient({ apiKey: 'api-key' }); @@ -568,10 +568,10 @@ describe('WebSocketClient', () => { stock.disconnect(); }); - it('should handle pong message and reset missed pongs counter', async () => { + it('should clear ping timer when disconnected', async () => { const client = new WebSocketClient({ apiKey: 'api-key', - healthCheck: { enabled: true, pingInterval: 50 } + healthCheck: { enabled: true, pingInterval: 100 } }); const stock = client.stock; @@ -582,24 +582,32 @@ describe('WebSocketClient', () => { server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); await promise; - // First ping + // Wait for first ping await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); - // Server sends pong (should reset missedPongs to 0) - server.send(JSON.stringify({ event: 'pong' })); - - // Second ping - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + const messagesBefore = server.messages.length; + // Disconnect stock.disconnect(); + await server.closed; + + // Wait to ensure no more pings are sent + await new Promise(resolve => setTimeout(resolve, 200)); + + // No new messages should have been received after disconnect + expect(server.messages.length).toBe(messagesBefore); }); + }); - it('should clear ping timer when disconnected', async () => { + describe('Freshness-based detection', () => { + it('should NOT disconnect while any inbound message keeps the connection fresh', async () => { const client = new WebSocketClient({ apiKey: 'api-key', - healthCheck: { enabled: true, pingInterval: 100 } + healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 2 } }); const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); const promise = stock.connect(); await server.connected; @@ -608,25 +616,71 @@ describe('WebSocketClient', () => { server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); await promise; - // Wait for first ping - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + // For several cycles, send some inbound data (NOT a pong) after each ping. + // Any inbound message resets freshness, so no disconnect should happen. + for (let i = 0; i < 4; i++) { + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + server.send(JSON.stringify({ event: 'data', data: { foo: 'bar' } })); + } - const messagesBefore = server.messages.length; + expect(disconnectCb).not.toHaveBeenCalled(); + stock.disconnect(); + }); - // Disconnect + it('should treat a pong like any other inbound message and stay fresh', async () => { + const client = new WebSocketClient({ + apiKey: 'api-key', + healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 2 } + }); + const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); + + const promise = stock.connect(); + await server.connected; + await expect(server).toReceiveMessage(JSON.stringify({ event: 'auth', data: { apikey: 'api-key' } })); + + server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); + await promise; + + for (let i = 0; i < 4; i++) { + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + server.send(JSON.stringify({ event: 'pong' })); + } + + expect(disconnectCb).not.toHaveBeenCalled(); stock.disconnect(); - await server.closed; + }); - // Wait to ensure no more pings are sent - await new Promise(resolve => setTimeout(resolve, 200)); + it('should clamp maxMissedPongs to >= 1 (0 must not disconnect a healthy connection)', async () => { + const client = new WebSocketClient({ + apiKey: 'api-key', + healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 0 } + }); + const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); - // No new messages should have been received after disconnect - expect(server.messages.length).toBe(messagesBefore); + const promise = stock.connect(); + await server.connected; + await expect(server).toReceiveMessage(JSON.stringify({ event: 'auth', data: { apikey: 'api-key' } })); + + server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); + await promise; + + // Keep the connection fresh; with a correct clamp it should never disconnect. + for (let i = 0; i < 3; i++) { + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + server.send(JSON.stringify({ event: 'data', data: {} })); + } + + expect(disconnectCb).not.toHaveBeenCalled(); + stock.disconnect(); }); }); describe('Error handling and auto-disconnect', () => { - it('should disconnect when maxMissedPongs is exceeded', async () => { + it('should disconnect after maxMissedPongs consecutive misses (no inbound messages)', async () => { const client = new WebSocketClient({ apiKey: 'api-key', healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 2 } @@ -642,17 +696,65 @@ describe('WebSocketClient', () => { server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); await promise; - // Receive 3 pings without sending pong (missedPongs will exceed maxMissedPongs=2) - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + // Receive pings without sending anything back -> consecutive misses accumulate await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); - // Should disconnect after third ping + // Should disconnect once misses reach maxMissedPongs + await server.closed; + expect(disconnectCb).toHaveBeenCalled(); + }); + + it('should carry { reason: "health-check-timeout" } in the disconnect event on timeout', async () => { + const client = new WebSocketClient({ + apiKey: 'api-key', + healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 2 } + }); + const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); + + const promise = stock.connect(); + await server.connected; + await expect(server).toReceiveMessage(JSON.stringify({ event: 'auth', data: { apikey: 'api-key' } })); + + server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); + await promise; + await server.closed; expect(disconnectCb).toHaveBeenCalled(); + // Second arg carries the health-check timeout reason. + expect(disconnectCb.mock.calls[0][1]).toEqual({ reason: 'health-check-timeout' }); }); - it('should handle ping send failure and disconnect', async () => { + it('should accumulate consecutive misses and reset them when a message arrives', async () => { + const client = new WebSocketClient({ + apiKey: 'api-key', + healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 3 } + }); + const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); + + const promise = stock.connect(); + await server.connected; + await expect(server).toReceiveMessage(JSON.stringify({ event: 'auth', data: { apikey: 'api-key' } })); + + server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); + await promise; + + // Two misses, then a message resets the counter, then more pings. + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + // Reset freshness before reaching 3 misses. + server.send(JSON.stringify({ event: 'data', data: {} })); + await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); + + expect(disconnectCb).not.toHaveBeenCalled(); + stock.disconnect(); + }); + + it('should handle ping send failure and disconnect with the timeout reason', async () => { const client = new WebSocketClient({ apiKey: 'api-key', healthCheck: { enabled: true, pingInterval: 50 } @@ -670,8 +772,6 @@ describe('WebSocketClient', () => { await promise; // Mock socket.send to throw an error on the next ping - // @ts-ignore - accessing private property for testing - const originalSend = stock.socket.send; // @ts-ignore stock.socket.send = jest.fn(() => { throw new Error('Network error'); @@ -682,18 +782,21 @@ describe('WebSocketClient', () => { expect(consoleErrorSpy).toHaveBeenCalledWith(expect.stringContaining('Failed to send ping')); expect(disconnectCb).toHaveBeenCalled(); + expect(disconnectCb.mock.calls[0][1]).toEqual({ reason: 'health-check-timeout' }); consoleErrorSpy.mockRestore(); }); }); - describe('Integration scenarios', () => { - it('should continue health check across multiple ping-pong cycles', async () => { + describe('Normal disconnect', () => { + it('should emit disconnect with undefined second arg on manual disconnect', async () => { const client = new WebSocketClient({ apiKey: 'api-key', - healthCheck: { enabled: true, pingInterval: 50, maxMissedPongs: 2 } + healthCheck: { enabled: true, pingInterval: 100 } }); const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); const promise = stock.connect(); await server.connected; @@ -702,22 +805,31 @@ describe('WebSocketClient', () => { server.send(JSON.stringify({ event: 'authenticated', data: { message: 'Authenticated successfully' } })); await promise; - // Cycle 1: ping -> pong - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); - server.send(JSON.stringify({ event: 'pong' })); + stock.disconnect(); + await server.closed; - // Cycle 2: ping -> pong - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); - server.send(JSON.stringify({ event: 'pong' })); + expect(disconnectCb).toHaveBeenCalled(); + // No reason payload on a normal/manual disconnect. + expect(disconnectCb.mock.calls[0][1]).toBeUndefined(); + }); - // Cycle 3: ping -> pong - await expect(server).toReceiveMessage(JSON.stringify({ event: 'ping', data: {} })); - server.send(JSON.stringify({ event: 'pong' })); + it('should emit disconnect with undefined second arg when the server closes', async () => { + const client = new WebSocketClient({ apiKey: 'api-key' }); + const stock = client.stock; + const disconnectCb = jest.fn(); + stock.once('disconnect', disconnectCb); - // Connection should still be healthy - stock.disconnect(); + stock.connect(); + await server.connected; + server.close(); + await server.closed; + + expect(disconnectCb).toHaveBeenCalled(); + expect(disconnectCb.mock.calls[0][1]).toBeUndefined(); }); + }); + describe('Integration scenarios', () => { it('should work alongside normal subscribe/unsubscribe operations', async () => { const client = new WebSocketClient({ apiKey: 'api-key',