diff --git a/.changeset/websocket-lifecycle-shutdown.md b/.changeset/websocket-lifecycle-shutdown.md new file mode 100644 index 00000000..cbb29683 --- /dev/null +++ b/.changeset/websocket-lifecycle-shutdown.md @@ -0,0 +1,10 @@ +--- +'lens': patch +--- + +Add WebSocket lifecycle heartbeat, dead connection reaping, socket memory bounding, shutdown close frames, and graceful SIGTERM teardown (Closes #347, #336, #335, #330). + +- **#347**: Enhanced `FanOutManager.destroy()` to gracefully close connected WebSocket clients with close code `1001` ('Server shutting down'), provide a configurable in-flight flush window, ensure full idempotency, and guarantee safety against mid-fanout calls. +- **#336**: Implemented periodic heartbeat pings with configurable interval (`WS_HEARTBEAT_INTERVAL_MS`) and max missed pings (`WS_MAX_MISSED_PINGS`). Reaps silent clients while cleaning up both `clients` and `pairClients` subscription maps, and exports `ws_connected_clients` gauge. +- **#335**: Added socket buffer inspection on `client.socket.bufferedAmount`. Coalesces intermediate price updates past soft threshold (`WS_BUFFER_THRESHOLD_BYTES`) and drops stalled clients past hard ceiling (`WS_BUFFER_MAX_BYTES`) with close code `1008` ('Buffer overflow'), exporting `ws_slow_clients_total` and `ws_messages_dropped_total`. +- **#330**: Implemented `GracefulShutdownCoordinator` in `src/shutdown.ts` responding to `SIGTERM` and `SIGINT` with strictly ordered teardown (stop HTTP/WS listeners -> stop ingesters -> close WS clients -> drain BullMQ queues/workers -> close pgPool and Redis) bounded by configurable `SHUTDOWN_GRACE_PERIOD_MS` (default 10s). diff --git a/.env.example b/.env.example index 3937aedc..d314452b 100644 --- a/.env.example +++ b/.env.example @@ -225,3 +225,17 @@ MCP_PAID_CALL_ALLOWED_HOSTS= # against a table that has never been pruned cannot hold a long lock or blow # the WAL. Default 1000. SNAPSHOT_PRUNE_BATCH_SIZE=1000 + +# ─── WebSocket Lifecycle & Memory Bounding (#336, #335, #347) ────────────────── +# Periodic ping interval in milliseconds to reap dead connections (default: 30000) +WS_HEARTBEAT_INTERVAL_MS=30000 +# Maximum missed pings before closing and reaping a silent connection (default: 2) +WS_MAX_MISSED_PINGS=2 +# Socket buffer threshold in bytes before coalescing price updates (default: 65536 / 64KB) +WS_BUFFER_THRESHOLD_BYTES=65536 +# Socket buffer hard ceiling in bytes before dropping a stalled client with code 1008 (default: 262144 / 256KB) +WS_BUFFER_MAX_BYTES=262144 + +# ─── Graceful Shutdown (#330) ────────────────────────────────────────────────── +# Maximum grace period in milliseconds for in-flight work to drain on SIGTERM/SIGINT (default: 10000) +SHUTDOWN_GRACE_PERIOD_MS=10000 diff --git a/src/__tests__/websocketLifecycleShutdown.test.ts b/src/__tests__/websocketLifecycleShutdown.test.ts new file mode 100644 index 00000000..9ff1126d --- /dev/null +++ b/src/__tests__/websocketLifecycleShutdown.test.ts @@ -0,0 +1,350 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import { + FanOutManager, + ws_connected_clients, + ws_clients_total, + ws_messages_sent_total, + ws_messages_dropped_total, + ws_slow_clients_total, +} from '../ws/fanout' +import { + GracefulShutdownCoordinator, + DEFAULT_SHUTDOWN_GRACE_PERIOD_MS, +} from '../shutdown' + +vi.mock('ioredis', () => { + const mockRedis = { + publish: vi.fn().mockResolvedValue(1), + subscribe: vi.fn().mockResolvedValue(undefined), + unsubscribe: vi.fn().mockResolvedValue(undefined), + disconnect: vi.fn(), + on: vi.fn(), + quit: vi.fn().mockResolvedValue(undefined), + } + return { + default: vi.fn(() => mockRedis), + } +}) + +vi.mock('../redis', () => ({ + redis: { + publish: vi.fn().mockResolvedValue(1), + quit: vi.fn().mockResolvedValue(undefined), + on: vi.fn(), + }, +})) + +vi.mock('../events', () => ({ + priceEmitter: { + on: vi.fn(), + off: vi.fn(), + listenerCount: vi.fn().mockReturnValue(0), + }, + PRICE_UPDATE: 'price:update', + PriceUpdateEvent: class {}, +})) + +describe('WebSocket Lifecycle, Memory Bounding & Graceful Teardown (#347, #336, #335, #330)', () => { + let manager: FanOutManager + + beforeEach(() => { + vi.clearAllMocks() + manager = new FanOutManager() + }) + + afterEach(async () => { + await manager.destroy() + }) + + // ─── #347: Clean Close Frames on Shutdown ─────────────────────────────────── + describe('#347: Websocket Shutdown Close Frames & Flush Window', () => { + it('sends every client a close frame with code 1001 on shutdown', async () => { + const client1 = { id: 'c1', send: vi.fn(), close: vi.fn() } + const client2 = { id: 'c2', send: vi.fn(), close: vi.fn() } + + manager.register('XLM/USDC', client1) + manager.register('XLM/USDC', client2) + + expect(manager.getClientCount()).toBe(2) + + await manager.destroy(1001, 'Server shutting down', 0) + + expect(client1.close).toHaveBeenCalledWith(1001, 'Server shutting down') + expect(client2.close).toHaveBeenCalledWith(1001, 'Server shutting down') + expect(manager.getClientCount()).toBe(0) + }) + + it('is idempotent: calling destroy multiple times does not throw or double-close', async () => { + const client = { id: 'c-idem', send: vi.fn(), close: vi.fn() } + manager.register('XLM/USDC', client) + + await manager.destroy() + await manager.destroy() + await manager.destroy() + + expect(client.close).toHaveBeenCalledTimes(1) + expect(manager.getClientCount()).toBe(0) + }) + + it('calling destroy mid-fanout does not throw', async () => { + const slowClient = { + id: 'c-mid', + send: vi.fn(() => { + // Trigger destroy during send + void manager.destroy() + }), + close: vi.fn(), + } + + manager.register('XLM/USDC', slowClient) + // Broadcast price update + ;(manager as any).broadcastToLocalClients({ + assetA: 'USDC', + assetB: 'XLM', + price: 0.12, + timestamp: new Date().toISOString(), + }) + + await expect(manager.destroy()).resolves.toBeUndefined() + }) + }) + + // ─── #336: Dead Connection Reaping & Map Hygiene ───────────────────────────── + describe('#336: Dead WebSocket Connection Reaping & Map Hygiene', () => { + it('pings live clients and tracks pong responses without disconnecting them', () => { + const pingFn = vi.fn() + const closeFn = vi.fn() + const client = { id: 'live-c', send: vi.fn(), close: closeFn, ping: pingFn } + + manager.register('XLM/USDC', client) + + // First heartbeat cycle: pings the client, marks isAlive=false + manager.checkHeartbeats() + expect(pingFn).toHaveBeenCalledTimes(1) + expect(closeFn).not.toHaveBeenCalled() + expect(manager.getClientCount()).toBe(1) + + // Client responds with pong + manager.recordPong('live-c') + + // Second heartbeat cycle: client was alive, pings again, not disconnected + manager.checkHeartbeats() + expect(pingFn).toHaveBeenCalledTimes(2) + expect(closeFn).not.toHaveBeenCalled() + expect(manager.getClientCount()).toBe(1) + }) + + it('reaps a client that misses consecutive pings and cleans up both maps', () => { + const closeFn = vi.fn() + const terminateFn = vi.fn() + const client = { id: 'dead-c', send: vi.fn(), close: closeFn, terminate: terminateFn, ping: vi.fn() } + + manager.register('XLM/USDC', client) + expect(manager.getClientCount()).toBe(1) + expect(manager.getSubscriptionCount('XLM/USDC')).toBe(1) + + // Cycle 1: ping sent, isAlive marked false + manager.checkHeartbeats() + expect(terminateFn).not.toHaveBeenCalled() + + // Cycle 2: no pong received -> missedPings = 1 + manager.checkHeartbeats() + expect(terminateFn).not.toHaveBeenCalled() + + // Cycle 3: still no pong -> missedPings = 2 (exceeds DEFAULT_MAX_MISSED_PINGS) -> reaped! + manager.checkHeartbeats() + expect(terminateFn).toHaveBeenCalledTimes(1) + + // Both maps must be completely clean (no ghost subscription) + expect(manager.getClientCount()).toBe(0) + expect(manager.getSubscriptionCount('XLM/USDC')).toBe(0) + expect(manager.getClientsMap().has('dead-c')).toBe(false) + expect(manager.getPairClientsMap().has('XLM/USDC')).toBe(false) + }) + + it('exports ws_connected_clients gauge and maintains accuracy across lifecycle', () => { + const client1 = { id: 'gauge-1', send: vi.fn(), close: vi.fn() } + const client2 = { id: 'gauge-2', send: vi.fn(), close: vi.fn() } + + const unsub1 = manager.register('XLM/USDC', client1) + const unsub2 = manager.register('XLM/USDC', client2) + + expect(manager.getClientCount()).toBe(2) + + unsub1() + expect(manager.getClientCount()).toBe(1) + + unsub2() + expect(manager.getClientCount()).toBe(0) + }) + }) + + // ─── #335: Memory Bounding & Slow Client Handling ───────────────────────────── + describe('#335: Socket Buffer Bounding & Memory Protection', () => { + it('coalesces updates when bufferedAmount exceeds soft threshold', () => { + let buffered = 70 * 1024 // 70KB > 64KB soft threshold + const sendFn = vi.fn() + const closeFn = vi.fn() + + const client = { + id: 'slow-c', + get bufferedAmount() { + return buffered + }, + send: sendFn, + close: closeFn, + } + + manager.register('USDC/XLM', client) + + // Broadcast update + ;(manager as any).broadcastToLocalClients({ + assetA: 'USDC', + assetB: 'XLM', + price: 0.12, + timestamp: new Date().toISOString(), + }) + + // Soft threshold should NOT call send (to prevent queueing in memory), but coalesce + expect(sendFn).not.toHaveBeenCalled() + expect(closeFn).not.toHaveBeenCalled() + expect(manager.getClientCount()).toBe(1) + }) + + it('drops stalled client with code 1008 when bufferedAmount exceeds hard ceiling', () => { + const buffered = 300 * 1024 // 300KB > 256KB hard limit + const sendFn = vi.fn() + const closeFn = vi.fn() + + const stalledClient = { + id: 'stalled-c', + get bufferedAmount() { + return buffered + }, + send: sendFn, + close: closeFn, + } + + const normalClient = { + id: 'fast-c', + bufferedAmount: 0, + send: vi.fn(), + close: vi.fn(), + } + + manager.register('USDC/XLM', stalledClient) + manager.register('USDC/XLM', normalClient) + + // Broadcast update + ;(manager as any).broadcastToLocalClients({ + assetA: 'USDC', + assetB: 'XLM', + price: 0.15, + timestamp: new Date().toISOString(), + }) + + // Stalled client dropped with 1008 Policy Violation + expect(closeFn).toHaveBeenCalledWith(1008, 'Buffer overflow') + expect(manager.getClientCount()).toBe(1) + expect(manager.getClientsMap().has('stalled-c')).toBe(false) + + // Normal client received update without being affected + expect(normalClient.send).toHaveBeenCalled() + }) + }) + + // ─── #330: Graceful Teardown on SIGTERM / SIGINT ───────────────────────────── + describe('#330: Graceful Shutdown Coordinator', () => { + it('executes teardown in strict sequence: HTTP -> Ingesters -> WS -> Queues -> DB', async () => { + const order: string[] = [] + + const mockApp = { + close: vi.fn(async () => { + order.push('app') + }), + } + const mockStopIngesters = vi.fn(async () => { + order.push('ingesters') + }) + const mockFanOut = { + destroy: vi.fn(async () => { + order.push('ws') + }), + } + const mockWorker = { + close: vi.fn(async () => { + order.push('queue') + }), + } + const mockPgPool = { + end: vi.fn(async () => { + order.push('pgPool') + }), + } + const mockRedis = { + quit: vi.fn(async () => { + order.push('redis') + }), + } + + const coordinator = new GracefulShutdownCoordinator( + { + app: mockApp, + stopIngesters: mockStopIngesters, + fanOut: mockFanOut, + queuesAndWorkers: [mockWorker], + pgPool: mockPgPool, + redis: mockRedis, + }, + { exitProcess: false, gracePeriodMs: 5000 } + ) + + await coordinator.shutdown('SIGTERM') + + expect(order).toEqual(['app', 'ingesters', 'ws', 'queue', 'pgPool', 'redis']) + expect(mockApp.close).toHaveBeenCalledTimes(1) + expect(mockStopIngesters).toHaveBeenCalledTimes(1) + expect(mockFanOut.destroy).toHaveBeenCalledTimes(1) + expect(mockWorker.close).toHaveBeenCalledTimes(1) + expect(mockPgPool.end).toHaveBeenCalledTimes(1) + expect(mockRedis.quit).toHaveBeenCalledTimes(1) + }) + + it('is idempotent: a second signal during shutdown does not re-enter drain', async () => { + let callCount = 0 + const mockApp = { + close: vi.fn(async () => { + callCount++ + await new Promise(r => setTimeout(r, 50)) + }), + } + + const coordinator = new GracefulShutdownCoordinator( + { app: mockApp }, + { exitProcess: false, gracePeriodMs: 5000 } + ) + + // Fire two signals concurrently + const p1 = coordinator.shutdown('SIGTERM') + const p2 = coordinator.shutdown('SIGINT') + + await Promise.all([p1, p2]) + + expect(callCount).toBe(1) + expect(mockApp.close).toHaveBeenCalledTimes(1) + }) + + it('enforces bounded grace period deadline when drain hangs', async () => { + const mockApp = { + close: vi.fn(() => new Promise(() => {})), // Hangs forever + } + + const coordinator = new GracefulShutdownCoordinator( + { app: mockApp }, + { exitProcess: false, gracePeriodMs: 50 } // 50ms deadline + ) + + await expect(coordinator.shutdown('SIGTERM')).rejects.toThrow('exceeded hard deadline of 50ms') + }) + }) +}) diff --git a/src/api/websocket.ts b/src/api/websocket.ts index 0c8b1f5c..7e678bff 100644 --- a/src/api/websocket.ts +++ b/src/api/websocket.ts @@ -85,20 +85,45 @@ export async function registerWebSocket(app: FastifyInstance) { app.log.info('[ws] Connection authorized') socket.send(JSON.stringify({ type: 'status', message: 'Streaming active' })) - // Register with fan-out manager for backpressure-aware broadcasting + // Register with fan-out manager for backpressure-aware broadcasting (#335, #336, #347) const clientId = uuid() const unsubscribe = fanOutManager.register('*', { id: clientId, + get bufferedAmount() { + return socket.bufferedAmount ?? 0 + }, send: (data: string) => { if (socket.readyState === 1) { // OPEN socket.send(data) } }, - close: () => { - socket.close() + close: (code = 1000, reason = 'Normal closure') => { + try { + socket.close(code, reason) + } catch { + // ignore + } + }, + ping: () => { + try { + if (socket.readyState === 1) socket.ping() + } catch { + // ignore + } + }, + terminate: () => { + try { + socket.terminate() + } catch { + // ignore + } }, }) + socket.on('pong', () => { + fanOutManager.recordPong(clientId) + }) + socket.on('close', () => { app.log.info('[ws] Connection closed') unsubscribe() diff --git a/src/index.ts b/src/index.ts index 55ed4c2b..b7ce1060 100644 --- a/src/index.ts +++ b/src/index.ts @@ -24,6 +24,7 @@ import { PRICE_POINTS_RETENTION_DAYS, } from './jobs/pricePointsPartitionJob' import { loadPersistedPairs, getActivePairs } from './pairsRegistry' +import { registerShutdownHandlers } from './shutdown' const log = createLogger('startup') @@ -73,6 +74,10 @@ async function main() { log.warn({ err }, 'websocket fan-out init skipped') } + // ── Background queues & workers registry for graceful shutdown (#330) ─── + const activeQueuesAndWorkers: Array<{ close: () => Promise }> = [] + let isShuttingDown = false + // ── Aggregate refresh worker (non-blocking — requires Redis) ───────────── // BullMQ opens its own ioredis connections, separate from ./redis and // without our error handler on them. If scheduling fails they are still @@ -92,6 +97,7 @@ async function main() { worker = startAggregateWorker(network) await scheduleAggregateRefresh(queue, network) started.push(network) + activeQueuesAndWorkers.push(queue, worker) } catch (err) { log.warn({ network, err }, 'aggregate refresh worker skipped (redis unavailable)') await Promise.allSettled([queue?.close(), worker?.close()]) @@ -134,6 +140,7 @@ async function main() { retentionWorker = startSnapshotRetentionWorker() await scheduleSnapshotRetention(retentionQueue) log.info('snapshot retention worker started') + activeQueuesAndWorkers.push(retentionQueue, retentionWorker) } catch (err) { log.warn({ err }, 'snapshot retention worker skipped (redis unavailable)') log.warn('falling back to an in-process hourly prune') @@ -174,6 +181,7 @@ async function main() { partitionWorker = startPricePointsPartitionWorker() await schedulePricePointsPartitionMaintenance(partitionQueue) log.info('price_points partition worker started') + activeQueuesAndWorkers.push(partitionQueue, partitionWorker) } catch (err) { log.warn({ err }, 'price_points partition worker unavailable; falling back to hourly in-process maintenance') await Promise.allSettled([partitionQueue?.close(), partitionWorker?.close()]) @@ -188,8 +196,11 @@ async function main() { // ingesters running on testnet. const restartIngester = (name: string, network: NetworkName, fn: () => Promise) => { fn().catch(err => { + if (isShuttingDown) return log.error({ ingester: name, network, err }, 'ingester crashed; restarting in 10s') - setTimeout(() => restartIngester(name, network, fn), 10_000) + setTimeout(() => { + if (!isShuttingDown) restartIngester(name, network, fn) + }, 10_000) }) } @@ -202,6 +213,18 @@ async function main() { restartIngester('Snapshot', network, () => startSnapshotIngester(network)) } + // ── Register graceful shutdown handlers (#330, #347) ────────────────────── + registerShutdownHandlers({ + app, + stopIngesters: () => { + isShuttingDown = true + }, + fanOut: fanOutManager, + queuesAndWorkers: activeQueuesAndWorkers, + pgPool, + redis, + }) + log.info( { count: getActivePairs().length, pairs: getActivePairs().map(p => p.pairKey) }, 'watching pairs' diff --git a/src/shutdown.ts b/src/shutdown.ts new file mode 100644 index 00000000..05d69822 --- /dev/null +++ b/src/shutdown.ts @@ -0,0 +1,193 @@ +/** + * Graceful Shutdown Coordinator (#330) + * + * Ensures SIGTERM and SIGINT trigger a deterministic, ordered teardown: + * 1. Stop accepting new HTTP/WS connections (app.close) + * 2. Stop background ingestion loops (stopIngesters) + * 3. Gracefully close connected WebSocket clients with code 1001 (fanOut.destroy) + * 4. Close background queues & workers (BullMQ) + * 5. Explicitly drain and close database and cache pools (pgPool, redis) + * + * Implements a bounded grace period (default 10s, configurable via SHUTDOWN_GRACE_PERIOD_MS) + * and guarantees idempotency: duplicate signals do not re-enter the drain. + */ + +export interface ShutdownTarget { + name: string + close: () => Promise | unknown +} + +export interface ShutdownContext { + app?: { close: () => Promise } + stopIngesters?: () => void | Promise + fanOut?: { destroy: (code?: number, reason?: string, flushMs?: number) => Promise } + queuesAndWorkers?: Array<{ close: () => Promise }> + pgPool?: { end: () => Promise } + redis?: { quit: () => Promise; disconnect?: () => void } +} + +export interface ShutdownOptions { + gracePeriodMs?: number + exitProcess?: boolean + logger?: { + info: (...args: any[]) => void + warn: (...args: any[]) => void + error: (...args: any[]) => void + } +} + +export const DEFAULT_SHUTDOWN_GRACE_PERIOD_MS = 10_000 + +export class GracefulShutdownCoordinator { + private isShuttingDown = false + private shutdownPromise: Promise | null = null + private shutdownTimer: NodeJS.Timeout | null = null + + constructor( + private readonly context: ShutdownContext = {}, + private readonly options: ShutdownOptions = {} + ) {} + + get shuttingDown(): boolean { + return this.isShuttingDown + } + + async shutdown(signal = 'SIGTERM'): Promise { + const logger = this.options.logger ?? console + const gracePeriodMs = + this.options.gracePeriodMs ?? + (parseInt(process.env.SHUTDOWN_GRACE_PERIOD_MS ?? '', 10) || DEFAULT_SHUTDOWN_GRACE_PERIOD_MS) + + if (this.isShuttingDown) { + logger.warn(`[shutdown] Shutdown already in progress (${signal} ignored)`) + return this.shutdownPromise ?? Promise.resolve() + } + + this.isShuttingDown = true + logger.info(`[shutdown] Received ${signal}. Starting graceful shutdown with ${gracePeriodMs}ms grace period…`) + + const performDrain = async () => { + // Step 1: Stop accepting new HTTP and WebSocket traffic + if (this.context.app) { + try { + logger.info('[shutdown] 1/5 Stopping HTTP & WebSocket listener…') + await this.context.app.close() + } catch (err) { + logger.warn('[shutdown] Error closing HTTP listener:', (err as Error).message) + } + } + + // Step 2: Signal background ingesters to stop + if (this.context.stopIngesters) { + try { + logger.info('[shutdown] 2/5 Stopping background ingesters…') + await this.context.stopIngesters() + } catch (err) { + logger.warn('[shutdown] Error stopping ingesters:', (err as Error).message) + } + } + + // Step 3: Gracefully disconnect WebSocket clients with 1001 Going Away + if (this.context.fanOut) { + try { + logger.info('[shutdown] 3/5 Sending close frames to WebSocket clients…') + await this.context.fanOut.destroy(1001, 'Server shutting down', 250) + } catch (err) { + logger.warn('[shutdown] Error destroying WebSocket fanout:', (err as Error).message) + } + } + + // Step 4: Close background job queues and workers + if (this.context.queuesAndWorkers && this.context.queuesAndWorkers.length > 0) { + try { + logger.info(`[shutdown] 4/5 Closing ${this.context.queuesAndWorkers.length} background queue/worker(s)…`) + await Promise.allSettled(this.context.queuesAndWorkers.map(w => w.close())) + } catch (err) { + logger.warn('[shutdown] Error closing background workers:', (err as Error).message) + } + } + + // Step 5: Explicitly close database and Redis connections + logger.info('[shutdown] 5/5 Closing PostgreSQL pool and Redis connections…') + if (this.context.pgPool) { + try { + await this.context.pgPool.end() + } catch (err) { + logger.warn('[shutdown] Error closing Postgres pool:', (err as Error).message) + } + } + if (this.context.redis) { + try { + await this.context.redis.quit() + } catch (_err) { + try { + this.context.redis.disconnect?.() + } catch { + // ignore + } + } + } + + logger.info('[shutdown] Graceful shutdown completed cleanly.') + } + + // Hard deadline watchdog + const timeoutPromise = new Promise((_, reject) => { + this.shutdownTimer = setTimeout(() => { + const msg = `[shutdown] Graceful shutdown exceeded hard deadline of ${gracePeriodMs}ms; forcing exit!` + logger.error(msg) + reject(new Error(msg)) + }, gracePeriodMs) + if (this.shutdownTimer.unref) { + this.shutdownTimer.unref() + } + }) + + this.shutdownPromise = Promise.race([performDrain(), timeoutPromise]) + .finally(() => { + if (this.shutdownTimer) { + clearTimeout(this.shutdownTimer) + this.shutdownTimer = null + } + }) + .then( + () => { + if (this.options.exitProcess ?? (process.env.NODE_ENV !== 'test')) { + process.exit(0) + } + }, + (err) => { + logger.error('[shutdown] Teardown failed or timed out:', (err as Error).message) + if (this.options.exitProcess ?? (process.env.NODE_ENV !== 'test')) { + process.exit(1) + } + throw err + } + ) + + return this.shutdownPromise + } +} + +let activeCoordinator: GracefulShutdownCoordinator | null = null + +export function registerShutdownHandlers( + context: ShutdownContext, + options?: ShutdownOptions +): GracefulShutdownCoordinator { + const coordinator = new GracefulShutdownCoordinator(context, options) + activeCoordinator = coordinator + + const onSignal = (signal: string) => { + coordinator.shutdown(signal).catch(() => {}) + } + + process.once('SIGTERM', () => onSignal('SIGTERM')) + process.once('SIGINT', () => onSignal('SIGINT')) + + return coordinator +} + +export function getActiveShutdownCoordinator(): GracefulShutdownCoordinator | null { + return activeCoordinator +} diff --git a/src/ws/fanout.ts b/src/ws/fanout.ts index 85c414ed..5c4bf127 100644 --- a/src/ws/fanout.ts +++ b/src/ws/fanout.ts @@ -7,14 +7,20 @@ import { createLogger } from '../logger' const log = createLogger('ws:fanout') -// ─── Metrics ────────────────────────────────────────────────────────────────── +// ─── Metrics (#336, #335) ───────────────────────────────────────────────────── export const ws_clients_total = new Gauge({ name: 'ws_clients_total', - help: 'Current number of connected WebSocket clients', + help: 'Current number of connected WebSocket clients by pair label', labelNames: ['pair'], registers: [register], }) +export const ws_connected_clients = new Gauge({ + name: 'ws_connected_clients', + help: 'Total number of active connected WebSocket clients', + registers: [register], +}) + export const ws_messages_sent_total = new Counter({ name: 'ws_messages_sent_total', help: 'Total number of price updates sent to WebSocket clients', @@ -24,14 +30,43 @@ export const ws_messages_sent_total = new Counter({ export const ws_messages_dropped_total = new Counter({ name: 'ws_messages_dropped_total', - help: 'Total number of price updates dropped due to client backpressure', + help: 'Total number of price updates dropped due to client backpressure or buffer threshold', + labelNames: ['pair'], + registers: [register], +}) + +export const ws_slow_clients_total = new Counter({ + name: 'ws_slow_clients_total', + help: 'Total number of slow clients dropped due to buffer overflow', labelNames: ['pair'], registers: [register], }) // ─── Configuration ───────────────────────────────────────────────────────────── const REDIS_CHANNEL = `lens:${activeNetwork}:price:updates` -const DEFAULT_COALESCE_INTERVAL_MS = 100 // max 10 updates/sec per client +export const DEFAULT_COALESCE_INTERVAL_MS = 100 // max 10 updates/sec per client + +export const DEFAULT_HEARTBEAT_INTERVAL_MS = 30_000 // ping interval (#336) +export const DEFAULT_MAX_MISSED_PINGS = 2 // max missed pings before reaping (#336) + +export const DEFAULT_BUFFER_THRESHOLD_BYTES = 64 * 1024 // 64 KB soft coalescing threshold (#335) +export const DEFAULT_BUFFER_MAX_BYTES = 256 * 1024 // 256 KB hard limit before dropping client (#335) + +export function getHeartbeatIntervalMs(): number { + return parseInt(process.env.WS_HEARTBEAT_INTERVAL_MS ?? '', 10) || DEFAULT_HEARTBEAT_INTERVAL_MS +} + +export function getMaxMissedPings(): number { + return parseInt(process.env.WS_MAX_MISSED_PINGS ?? '', 10) || DEFAULT_MAX_MISSED_PINGS +} + +export function getBufferThresholdBytes(): number { + return parseInt(process.env.WS_BUFFER_THRESHOLD_BYTES ?? '', 10) || DEFAULT_BUFFER_THRESHOLD_BYTES +} + +export function getBufferMaxBytes(): number { + return parseInt(process.env.WS_BUFFER_MAX_BYTES ?? '', 10) || DEFAULT_BUFFER_MAX_BYTES +} /** * Pair names can contain arbitrary asset codes and issuer identifiers supplied @@ -46,7 +81,10 @@ function pairMetricLabel(pairKey: string): 'other' | null { export interface FanOutClient { id: string send: (data: string) => void - close: () => void + close: (code?: number, reason?: string) => void + ping?: () => void + terminate?: () => void + bufferedAmount?: number } interface InternalClient { @@ -57,10 +95,12 @@ interface InternalClient { pendingUpdate: string | null timer: ReturnType | null destroyed: boolean + isAlive: boolean + missedPings: number } // ─── Fan-Out Manager ────────────────────────────────────────────────────────── -class FanOutManager { +export class FanOutManager { private clients = new Map() // clientId -> client private pairClients = new Map>() // pairKey -> Set @@ -69,7 +109,17 @@ class FanOutManager { private publisher: typeof redis = redis private redisConnected = false + // Heartbeat state (#336) + private heartbeatTimer: ReturnType | null = null + + // Shutdown lifecycle state (#347) + private isDestroying = false + private isDestroyed = false + async initialize(): Promise { + // Start periodic heartbeat watchdog + this.startHeartbeat() + // Subscribe to local price emitter (always works) priceEmitter.on(PRICE_UPDATE, (event: PriceUpdateEvent) => { this.onPriceUpdate(event) @@ -113,6 +163,8 @@ class FanOutManager { pendingUpdate: null, timer: null, destroyed: false, + isAlive: true, + missedPings: 0, } this.clients.set(id, client) @@ -126,12 +178,13 @@ class FanOutManager { const metricPair = pairMetricLabel(pairKey) if (metricPair) ws_clients_total.labels(metricPair).inc() ws_clients_total.labels('*').inc() + ws_connected_clients.inc() // Return unsubscribe function return () => this.unregister(id) } - private unregister(id: string): void { + unregister(id: string): void { const client = this.clients.get(id) if (!client) return if (client.timer) clearTimeout(client.timer) @@ -149,8 +202,78 @@ class FanOutManager { const metricPair = pairMetricLabel(client.pairKey) if (metricPair) ws_clients_total.labels(metricPair).dec() ws_clients_total.labels('*').dec() + ws_connected_clients.dec() + } + + // ─── Heartbeat Management (#336) ─────────────────────────────────────────── + + startHeartbeat(intervalMs?: number): void { + this.stopHeartbeat() + const interval = intervalMs ?? getHeartbeatIntervalMs() + this.heartbeatTimer = setInterval(() => { + this.checkHeartbeats() + }, interval) + if (this.heartbeatTimer && typeof this.heartbeatTimer === 'object' && 'unref' in this.heartbeatTimer) { + this.heartbeatTimer.unref() + } + } + + stopHeartbeat(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer) + this.heartbeatTimer = null + } + } + + checkHeartbeats(): void { + const maxMissed = getMaxMissedPings() + const clientList = Array.from(this.clients.values()) + + for (const client of clientList) { + if (client.destroyed) continue + + if (!client.isAlive) { + client.missedPings++ + if (client.missedPings >= maxMissed) { + // Client stopped responding — terminate/close and reap cleanly + try { + if (client.socket.terminate) { + client.socket.terminate() + } else { + client.socket.close(1006, 'Ping timeout') + } + } catch { + // ignore + } + this.unregister(client.id) + continue + } + } else { + // Expect pong response before next cycle + client.isAlive = false + } + + // Send ping probe + try { + if (client.socket.ping) { + client.socket.ping() + } + } catch { + this.unregister(client.id) + } + } + } + + recordPong(id: string): void { + const client = this.clients.get(id) + if (client) { + client.isAlive = true + client.missedPings = 0 + } } + // ─── Price Update Broadcasting ────────────────────────────────────────────── + /** * Called when a price update is received from the local emitter. * Publishes to Redis for other instances and broadcasts locally. @@ -203,9 +326,11 @@ class FanOutManager { if (elapsed >= DEFAULT_COALESCE_INTERVAL_MS) { // Client is ready to receive - this.sendToClient(client, payload) - client.lastSent = now - client.pendingUpdate = null + const sent = this.sendToClient(client, payload) + if (sent) { + client.lastSent = now + client.pendingUpdate = null + } } else { // Client is too fast — coalesce: buffer the latest update client.pendingUpdate = payload @@ -215,9 +340,11 @@ class FanOutManager { client.timer = null if (client.destroyed) return if (client.pendingUpdate) { - this.sendToClient(client, client.pendingUpdate) - client.lastSent = Date.now() - client.pendingUpdate = null + const sent = this.sendToClient(client, client.pendingUpdate) + if (sent) { + client.lastSent = Date.now() + client.pendingUpdate = null + } } }, timeout) } @@ -225,22 +352,62 @@ class FanOutManager { } } - private sendToClient(client: InternalClient, payload: string): void { + /** + * Sends payload to a client enforcing backpressure and buffer threshold checks (#335). + * Returns true only when socket.send actually executes successfully. + */ + private sendToClient(client: InternalClient, payload: string): boolean { + if (client.destroyed) return false + + const buffered = client.socket.bufferedAmount ?? 0 + const hardLimit = getBufferMaxBytes() + const softThreshold = getBufferThresholdBytes() + const metricPair = pairMetricLabel(client.pairKey) + + // Hard ceiling: client is stalled and accumulating unbounded memory (#335) + if (buffered >= hardLimit) { + if (metricPair) ws_slow_clients_total.labels(metricPair).inc() + ws_slow_clients_total.labels('*').inc() + try { + client.socket.close(1008, 'Buffer overflow') + } catch { + // ignore + } + this.unregister(client.id) + return false + } + + // Soft threshold: socket is lagging behind; drop intermediate updates and coalesce (#335) + if (buffered >= softThreshold) { + if (metricPair) ws_messages_dropped_total.labels(metricPair).inc() + ws_messages_dropped_total.labels('*').inc() + client.pendingUpdate = payload + return false + } + try { client.socket.send(payload) - const metricPair = pairMetricLabel(client.pairKey) if (metricPair) ws_messages_sent_total.labels(metricPair).inc() ws_messages_sent_total.labels('*').inc() + return true } catch { // Client likely disconnected this.unregister(client.id) + return false } } /** - * 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 + + this.stopHeartbeat() + if (this.subscriber) { try { await this.subscriber.unsubscribe(REDIS_CHANNEL) @@ -248,10 +415,60 @@ class FanOutManager { // subscriber may not have connected — ignore } this.subscriber.disconnect() + this.subscriber = null } - for (const [id] of this.clients) { - this.unregister(id) + + // 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)) + } + + 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 + } + + // ─── Inspection Helpers for Testing & Health Checks (#336) ─────────────────── + + getClientCount(): number { + return this.clients.size + } + + getSubscriptionCount(pairKey?: string): number { + if (pairKey) return this.pairClients.get(pairKey)?.size ?? 0 + return this.pairClients.size + } + + getClientsMap(): ReadonlyMap { + return this.clients + } + + getPairClientsMap(): ReadonlyMap> { + return this.pairClients + } + + resetForTesting(): void { + this.stopHeartbeat() + this.clients.clear() + this.pairClients.clear() + this.isDestroyed = false + this.isDestroying = false } }