diff --git a/src/__tests__/fanout.test.ts b/src/__tests__/fanout.test.ts index bbc2ff02..9a2cde06 100644 --- a/src/__tests__/fanout.test.ts +++ b/src/__tests__/fanout.test.ts @@ -141,7 +141,7 @@ describe('FanOutManager', () => { )).toHaveLength(1) }) - it('destroy cleans up all state', async () => { + it('destroy cleans up all state idempotently', async () => { const mockSocket = { id: 'client-cleanup', send: vi.fn(), @@ -150,7 +150,33 @@ describe('FanOutManager', () => { fanOutManager.register('XLM/USDC', mockSocket) await fanOutManager.destroy() - await fanOutManager.destroy() + await fanOutManager.destroy() // Idempotent second call + expect(mockSocket.close).toHaveBeenCalledWith(1001, 'Server shutting down') + }) + + it('sends deliberate close code 1001 to all connected clients on shutdown (#347)', async () => { + const socketA = { id: 'client-a', send: vi.fn(), close: vi.fn() } + const socketB = { id: 'client-b', send: vi.fn(), close: vi.fn() } + + fanOutManager.register('XLM/USDC', socketA) + fanOutManager.register('BTC/USDC', socketB) + + await fanOutManager.destroy(1001, 'Server shutting down') + + expect(socketA.close).toHaveBeenCalledWith(1001, 'Server shutting down') + expect(socketB.close).toHaveBeenCalledWith(1001, 'Server shutting down') + }) + + it('respects flush window during shutdown before closing clients (#347)', async () => { + const socket = { id: 'flush-client', send: vi.fn(), close: vi.fn() } + fanOutManager.register('XLM/USDC', socket) + + const start = Date.now() + await fanOutManager.destroy(1001, 'Server shutting down', 50) + const elapsed = Date.now() - start + + expect(elapsed).toBeGreaterThanOrEqual(40) + expect(socket.close).toHaveBeenCalledWith(1001, 'Server shutting down') }) }) diff --git a/src/ws/fanout.ts b/src/ws/fanout.ts index d8430b49..64a0413d 100644 --- a/src/ws/fanout.ts +++ b/src/ws/fanout.ts @@ -117,6 +117,8 @@ interface InternalClient { class FanOutManager { private clients = new Map() // clientId -> client private pairClients = new Map>() // pairKey -> Set + private isDestroying = false + private isDestroyed = false // Redis pub/sub for cross-instance broadcasting private subscriber: typeof redis | null = null @@ -178,6 +180,7 @@ class FanOutManager { congestedSince: null, } + this.isDestroyed = false this.clients.set(id, client) let clients = this.pairClients.get(pairKey) if (!clients) { @@ -445,9 +448,14 @@ class FanOutManager { } /** - * Dispose all clients and clean up. + * Graceful Shutdown & Cleanup (#347) + * Sends close frames with code 1001 ('Server shutting down') and allows a flush window. + * Completely idempotent and safe against mid-fanout calls. */ - async destroy(): Promise { + async destroy(closeCode = 1001, reason = 'Server shutting down', flushWindowMs = 0): Promise { + if (this.isDestroyed || this.isDestroying) return + this.isDestroying = true + if (this.bufferSweep) { clearInterval(this.bufferSweep) this.bufferSweep = null @@ -459,10 +467,40 @@ class FanOutManager { // subscriber may not have connected — ignore } this.subscriber.disconnect() + this.subscriber = null + } + + // Allow in-flight sends / queued coalesced messages a brief window to flush (#347) + if (flushWindowMs > 0 && this.clients.size > 0) { + await new Promise(resolve => setTimeout(resolve, flushWindowMs)) } - for (const [id] of this.clients) { - this.unregister(id) + + const clientList = Array.from(this.clients.values()) + for (const client of clientList) { + try { + if (client.timer) { + clearTimeout(client.timer) + client.timer = null + } + client.destroyed = true + client.socket.close(closeCode, reason) + } catch { + // do not throw + } + this.unregister(client.id) } + + this.clients.clear() + this.pairClients.clear() + this.isDestroying = false + this.isDestroyed = true + } + + resetForTesting(): void { + this.clients.clear() + this.pairClients.clear() + this.isDestroyed = false + this.isDestroying = false } }