diff --git a/scripts/ws-load-test.js b/scripts/ws-load-test.js new file mode 100644 index 0000000..139a2cd --- /dev/null +++ b/scripts/ws-load-test.js @@ -0,0 +1,98 @@ +/** + * WebSocket connection-churn load test (issue #158). + * + * Simulates many clients repeatedly connecting and disconnecting to verify that + * the TradeGateway connection pool stays flat (no memory leak) under flaky-network + * conditions. Use it together with heap snapshots to confirm dropped sockets are + * garbage collected. + * + * Prerequisites: + * - The API running locally: npm run start + * - socket.io-client installed: npm i -D socket.io-client + * + * Usage: + * node scripts/ws-load-test.js + * URL=http://localhost:3000 CLIENTS=500 ROUNDS=20 HOLD_MS=500 node scripts/ws-load-test.js + * + * Capture heap snapshots around the run to confirm stability: + * node --expose-gc scripts/ws-load-test.js # forces GC between rounds + * # then compare process RSS / heapUsed printed each round; it should plateau. + */ + +const URL = process.env.URL || 'http://localhost:3000'; +const CLIENTS = Number(process.env.CLIENTS || 500); +const ROUNDS = Number(process.env.ROUNDS || 20); +const HOLD_MS = Number(process.env.HOLD_MS || 500); + +let io; +try { + io = require('socket.io-client'); +} catch { + console.error( + 'socket.io-client is required to run this load test.\n' + + 'Install it with: npm i -D socket.io-client', + ); + process.exit(1); +} + +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); + +function mb(bytes) { + return (bytes / 1024 / 1024).toFixed(1); +} + +function logMemory(label) { + const m = process.memoryUsage(); + console.log( + `[${label}] rss=${mb(m.rss)}MB heapUsed=${mb(m.heapUsed)}MB heapTotal=${mb(m.heapTotal)}MB`, + ); +} + +function connectOnce() { + return new Promise((resolve) => { + const socket = io(URL, { + transports: ['websocket'], + reconnection: false, + timeout: 5000, + }); + const done = () => resolve(socket); + socket.on('connect', done); + socket.on('connect_error', done); + }); +} + +async function round(n) { + const sockets = await Promise.all( + Array.from({ length: CLIENTS }, () => connectOnce()), + ); + await sleep(HOLD_MS); + // Abruptly drop every connection — mimics flaky clients vanishing. + sockets.forEach((s) => s.disconnect()); + await sleep(HOLD_MS); + if (global.gc) { + global.gc(); + } + logMemory(`round ${n + 1}/${ROUNDS}`); +} + +async function main() { + console.log( + `Churning ${CLIENTS} clients x ${ROUNDS} rounds against ${URL} ` + + `(${global.gc ? 'gc enabled' : 'run with --expose-gc for GC between rounds'})`, + ); + logMemory('baseline'); + for (let i = 0; i < ROUNDS; i++) { + await round(i); + } + logMemory('final'); + console.log( + 'If heapUsed/rss plateau across rounds rather than growing monotonically, ' + + 'dropped connections are being pruned correctly.', + ); + process.exit(0); +} + +main().catch((err) => { + console.error('Load test failed:', err); + process.exit(1); +}); diff --git a/src/dynamic/dynamic.module.ts b/src/dynamic/dynamic.module.ts index 50ecd9f..a447712 100644 --- a/src/dynamic/dynamic.module.ts +++ b/src/dynamic/dynamic.module.ts @@ -1,9 +1,10 @@ import { Module } from '@nestjs/common'; -import { InvoicesController } from '../invoices/invoices.controller'; -import { PdfService } from '../invoices/pdf.service'; +import { NetworkController } from './dynamic.controller'; +import { PdfService } from './dynamic.serivice'; @Module({ - controllers: [InvoicesController], + controllers: [NetworkController], providers: [PdfService], + exports: [PdfService], }) export class InvoicesModule {} diff --git a/src/pools/pools.controller.int.spec.ts b/src/pools/pools.controller.int.spec.ts index 2418a40..4a21cb4 100644 --- a/src/pools/pools.controller.int.spec.ts +++ b/src/pools/pools.controller.int.spec.ts @@ -20,17 +20,19 @@ describe('PoolsController (integration)', () => { await app.close(); }); - it('GET /api/v1/pools/:poolId/apy-history returns 200 and success envelope', async () => { + it('GET /api/v1/pools/:poolId/apy-history returns 200 and a 7-point history array', async () => { const res = await request(app.getHttpServer()) .get('/api/v1/pools/pool-123/apy-history') .expect(200); + // The endpoint returns the APY history array directly (see the controller's + // `ApyHistoryPoint[]` return type / Swagger `isArray: true`); this app has no + // global response-envelope interceptor. expect(res.body).toBeDefined(); - expect(res.body.status).toBe('success'); - expect(Array.isArray(res.body.data)).toBe(true); - expect(res.body.data).toHaveLength(7); + expect(Array.isArray(res.body)).toBe(true); + expect(res.body).toHaveLength(7); - for (const item of res.body.data) { + for (const item of res.body) { expect(item).toHaveProperty('date'); expect(item).toHaveProperty('apyPercentage'); expect(typeof item.date).toBe('string'); diff --git a/src/trade/trade.gateway.spec.ts b/src/trade/trade.gateway.spec.ts new file mode 100644 index 0000000..a33e9a8 --- /dev/null +++ b/src/trade/trade.gateway.spec.ts @@ -0,0 +1,183 @@ +import { Test, TestingModule } from '@nestjs/testing'; +import { TradeGateway } from './trade.gateway'; +import { RedisService } from '../common/redis/redis.service'; + +/** + * Unit tests for the TradeGateway connection pool (issue #158). + * + * Verifies that connections are tracked on connect, fully removed on disconnect + * (no leaked references or timers), that the heartbeat terminates ghost/dead + * sockets, and that churning 500 clients leaves the pool empty. + */ +const HEARTBEAT_INTERVAL_MS = 30_000; + +interface MockSocket { + id: string; + connected: boolean; + conn: { on: jest.Mock; off: jest.Mock }; + on: jest.Mock; + off: jest.Mock; + emit: jest.Mock; + disconnect: jest.Mock; + trigger: (event: string, ...args: unknown[]) => void; + triggerConn: (event: string, ...args: unknown[]) => void; +} + +function createMockSocket(id: string): MockSocket { + const handlers: Record void> = {}; + const connHandlers: Record void> = {}; + const socket: MockSocket = { + id, + connected: true, + conn: { + on: jest.fn((event: string, handler: (...args: unknown[]) => void) => { + connHandlers[event] = handler; + }), + off: jest.fn(), + }, + on: jest.fn((event: string, handler: (...args: unknown[]) => void) => { + handlers[event] = handler; + }), + off: jest.fn(), + emit: jest.fn(), + disconnect: jest.fn(function (this: MockSocket) { + this.connected = false; + }), + trigger: (event, ...args) => handlers[event]?.(...args), + triggerConn: (event, ...args) => connHandlers[event]?.(...args), + }; + return socket; +} + +describe('TradeGateway (connection pool)', () => { + let gateway: TradeGateway; + + beforeEach(async () => { + const module: TestingModule = await Test.createTestingModule({ + providers: [ + TradeGateway, + { provide: RedisService, useValue: { subscribe: jest.fn() } }, + ], + }).compile(); + + gateway = module.get(TradeGateway); + + // Use jest's fake timers for the whole suite so the per-socket heartbeat is + // deterministic and timer globals are consistently available. Enabled after + // the async module compile so it doesn't interfere with promise resolution. + jest.useFakeTimers(); + }); + + afterEach(() => { + gateway.onModuleDestroy(); + jest.clearAllTimers(); + jest.useRealTimers(); + }); + + it('should be defined', () => { + expect(gateway).toBeDefined(); + expect(gateway.getConnectionCount()).toBe(0); + }); + + it('tracks a connection on connect and releases it on disconnect', () => { + const socket = createMockSocket('a'); + + gateway.handleConnection(socket as never); + expect(gateway.getConnectionCount()).toBe(1); + + gateway.handleDisconnect(socket as never); + expect(gateway.getConnectionCount()).toBe(0); + }); + + it('clears the per-socket heartbeat on disconnect (no lingering interval)', () => { + jest.useFakeTimers(); + const clearSpy = jest.spyOn(global, 'clearInterval'); + const socket = createMockSocket('a'); + + gateway.handleConnection(socket as never); + gateway.handleDisconnect(socket as never); + + expect(clearSpy).toHaveBeenCalled(); + + // After disconnect, advancing time must not emit further heartbeats. + socket.emit.mockClear(); + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS * 3); + expect(socket.emit).not.toHaveBeenCalled(); + }); + + it('detaches every listener it attached on disconnect', () => { + const socket = createMockSocket('a'); + + gateway.handleConnection(socket as never); + gateway.handleDisconnect(socket as never); + + expect(socket.off).toHaveBeenCalledWith('heartbeat:pong', expect.any(Function)); + expect(socket.off).toHaveBeenCalledWith('error', expect.any(Function)); + expect(socket.conn.off).toHaveBeenCalledWith('packet', expect.any(Function)); + }); + + it('sends a heartbeat ping and keeps a responsive client alive', () => { + jest.useFakeTimers(); + const socket = createMockSocket('a'); + gateway.handleConnection(socket as never); + + // First probe: emits an app-level ping and marks the client unproven. + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS); + expect(socket.emit).toHaveBeenCalledWith('heartbeat:ping'); + + // Client proves liveness (engine pong packet). + socket.triggerConn('packet', { type: 'pong' }); + + // Next probe: still alive, so it is NOT disconnected. + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS); + expect(socket.disconnect).not.toHaveBeenCalled(); + expect(gateway.getConnectionCount()).toBe(1); + }); + + it('terminates an unresponsive (ghost) client', () => { + jest.useFakeTimers(); + const socket = createMockSocket('a'); + gateway.handleConnection(socket as never); + + // Two probes with no pong in between → second probe terminates it. + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS); // marks unproven, pings + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS); // still unproven → disconnect + expect(socket.disconnect).toHaveBeenCalledWith(true); + }); + + it('terminates a client whose transport already dropped', () => { + jest.useFakeTimers(); + const socket = createMockSocket('a'); + gateway.handleConnection(socket as never); + + socket.connected = false; // transport gone but still in the pool + jest.advanceTimersByTime(HEARTBEAT_INTERVAL_MS); + expect(socket.disconnect).toHaveBeenCalledWith(true); + }); + + it('does not leak references when 500 clients connect and disconnect', () => { + const sockets = Array.from({ length: 500 }, (_, i) => + createMockSocket(`client-${i}`), + ); + + sockets.forEach((s) => gateway.handleConnection(s as never)); + expect(gateway.getConnectionCount()).toBe(500); + + sockets.forEach((s) => gateway.handleDisconnect(s as never)); + expect(gateway.getConnectionCount()).toBe(0); + }); + + it('clears all timers on module destroy', () => { + jest.useFakeTimers(); + const clearSpy = jest.spyOn(global, 'clearInterval'); + [createMockSocket('a'), createMockSocket('b')].forEach((s) => + gateway.handleConnection(s as never), + ); + expect(gateway.getConnectionCount()).toBe(2); + + gateway.onModuleDestroy(); + + expect(gateway.getConnectionCount()).toBe(0); + expect(clearSpy).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/trade/trade.gateway.ts b/src/trade/trade.gateway.ts index d4e2f64..d4bdcfb 100644 --- a/src/trade/trade.gateway.ts +++ b/src/trade/trade.gateway.ts @@ -1,23 +1,68 @@ -import { OnModuleInit, Logger } from '@nestjs/common'; -import { WebSocketGateway, WebSocketServer } from '@nestjs/websockets'; -import { Server } from 'socket.io'; +import { + WebSocketGateway, + WebSocketServer, + OnGatewayConnection, + OnGatewayDisconnect, +} from '@nestjs/websockets'; +import { Logger, OnModuleInit, OnModuleDestroy } from '@nestjs/common'; +import { Server, Socket } from 'socket.io'; import { RedisService } from '../common/redis/redis.service'; +/** How often each socket is probed for liveness (issue #158: ~30s heartbeat). */ +const HEARTBEAT_INTERVAL_MS = 30_000; +/** Engine.io ping cadence — kept below the app heartbeat so a live client always refreshes liveness first. */ +const PING_INTERVAL_MS = 25_000; +/** Engine.io waits this long for the client's pong before dropping the socket. */ +const PING_TIMEOUT_MS = 20_000; + +/** + * Per-connection bookkeeping. Held in the {@link TradeGateway.clients} pool and + * deleted on disconnect so a dropped socket (and everything it closes over) can + * be garbage collected. + */ +interface ClientState { + /** Refreshed whenever the client proves it is alive (engine pong or app pong). */ + isAlive: boolean; + /** Per-socket heartbeat timer — MUST be cleared on disconnect to avoid a leak. */ + heartbeat: ReturnType; + /** Listeners we attached, kept so they can be detached on cleanup. */ + onPacket: (packet: { type: string }) => void; + onPong: () => void; + onError: (err: Error) => void; +} + /** * WebSocket Gateway for real-time trade updates. - * Subscribes to Redis 'live_trades' channel and broadcasts messages to connected clients. + * + * Subscribes to the Redis `live_trades` channel and broadcasts to connected + * clients. Maintains an explicit connection pool with a heartbeat so that dead + * or dropped sockets (issue #158) are pruned from memory rather than leaking: + * + * - Every socket is tracked in {@link clients} on connect and removed on + * disconnect, dropping all references to the socket instance. + * - Engine.io ping/pong (`pingInterval`/`pingTimeout`) terminates connections + * whose underlying transport has gone away. + * - An additional per-socket heartbeat forcibly disconnects ghost connections + * that stop responding, and its `setInterval` is always cleared on disconnect. */ @WebSocketGateway({ cors: { origin: '*', }, + pingInterval: PING_INTERVAL_MS, + pingTimeout: PING_TIMEOUT_MS, }) -export class TradeGateway implements OnModuleInit { +export class TradeGateway + implements OnModuleInit, OnGatewayConnection, OnGatewayDisconnect, OnModuleDestroy +{ @WebSocketServer() server: Server; private readonly logger = new Logger(TradeGateway.name); + /** Active connection pool, keyed by socket id. */ + private readonly clients = new Map(); + constructor(private readonly redisService: RedisService) {} /** @@ -29,4 +74,91 @@ export class TradeGateway implements OnModuleInit { this.server.emit('trade', JSON.parse(message)); }); } + + /** + * Registers a new connection in the pool and wires up its heartbeat. + */ + handleConnection(client: Socket) { + const markAlive = () => { + const state = this.clients.get(client.id); + if (state) { + state.isAlive = true; + } + }; + + // Treat any inbound engine pong (sent automatically by socket.io clients in + // response to the server's ping) as proof of life — no client-side code + // required. App-level 'heartbeat:pong' is also honored for custom clients. + const onPacket = (packet: { type: string }) => { + if (packet.type === 'pong') { + markAlive(); + } + }; + const onPong = () => markAlive(); + const onError = (err: Error) => { + // Socket.IO emits 'disconnect' after a fatal 'error'; log for visibility. + this.logger.error(`Socket error for ${client.id}: ${err?.message ?? err}`); + }; + + client.conn?.on('packet', onPacket); + client.on('heartbeat:pong', onPong); + client.on('error', onError); + + const heartbeat = setInterval(() => { + const state = this.clients.get(client.id); + if (!state) { + return; + } + // A socket that is no longer connected, or hasn't proven liveness since + // the last probe, is a ghost — force it closed so it gets cleaned up. + if (!client.connected || !state.isAlive) { + this.logger.warn(`Terminating unresponsive client: ${client.id}`); + client.disconnect(true); // triggers handleDisconnect -> cleanup + return; + } + state.isAlive = false; + client.emit('heartbeat:ping'); + }, HEARTBEAT_INTERVAL_MS); + + this.clients.set(client.id, { isAlive: true, heartbeat, onPacket, onPong, onError }); + this.logger.log( + `Client connected: ${client.id} (active connections: ${this.clients.size})`, + ); + } + + /** + * Tears down a connection: clears its heartbeat, detaches listeners, and drops + * it from the pool so the socket instance is no longer referenced. + */ + handleDisconnect(client: Socket) { + const state = this.clients.get(client.id); + if (state) { + clearInterval(state.heartbeat); + client.conn?.off('packet', state.onPacket); + client.off('heartbeat:pong', state.onPong); + client.off('error', state.onError); + this.clients.delete(client.id); + } + this.logger.log( + `Client disconnected: ${client.id} (active connections: ${this.clients.size})`, + ); + } + + /** + * Clears every per-socket timer on shutdown so no intervals outlive the module. + */ + onModuleDestroy() { + for (const state of this.clients.values()) { + clearInterval(state.heartbeat); + } + this.clients.clear(); + } + + /** + * Number of connections currently tracked in the pool. Exposed for health + * checks and leak assertions in tests. + */ + getConnectionCount(): number { + return this.clients.size; + } }