diff --git a/src/routes/health.ts b/src/routes/health.ts index 92050f3..d962971 100644 --- a/src/routes/health.ts +++ b/src/routes/health.ts @@ -54,10 +54,14 @@ export function healthRoutes(): Hono { checks.db = false; } - // Check WebSocket subsystem — saturated WS means new clients can't connect + // Check WebSocket subsystem — saturated WS means new clients can't connect. + // Use liveConnections (the `ws` library's own live socket count) rather + // than totalConnections, which lags behind during the async auth/rate-limit + // chain and would under-report true saturation under connection bursts. try { const wsMetrics = getWebSocketMetrics(); - const utilization = wsMetrics.totalConnections / wsMetrics.limits.maxGlobalConnections; + const liveConnections = wsMetrics.liveConnections ?? wsMetrics.totalConnections; + const utilization = liveConnections / wsMetrics.limits.maxGlobalConnections; checks.ws = utilization < 0.95; // degraded if >95% of connection slots used } catch { checks.ws = false; diff --git a/src/routes/ws.ts b/src/routes/ws.ts index b2ee32e..7cd81b7 100644 --- a/src/routes/ws.ts +++ b/src/routes/ws.ts @@ -201,6 +201,15 @@ const metrics: Metrics = { lastResetTime: Date.now(), }; +// metrics.totalConnections only reflects sockets that have finished the +// async IP-blocklist/auth-ban/rate-limit chain in the "connection" handler +// below — a burst of concurrent handshakes can sit open at the `ws` library +// level for that entire await chain before metrics.totalConnections is ever +// incremented. Holding the live wss instance lets callers (health checks) +// read wss.clients.size, the library's own real-time count, instead of +// trusting our slower-to-update bookkeeping. +let liveWss: WebSocketServer | null = null; + /** * Safely send a JSON-serializable payload to a WebSocket client. * @@ -458,6 +467,11 @@ export function getWebSocketMetrics(): any { return { totalConnections: metrics.totalConnections, + // The live count from the `ws` library itself — includes sockets still + // mid-handshake in the async auth/rate-limit chain that totalConnections + // hasn't counted yet. Falls back to totalConnections if the WS server + // hasn't been set up yet (e.g. in tests that don't call setupWebSocket). + liveConnections: liveWss ? liveWss.clients.size : metrics.totalConnections, connectionsPerSlab: Object.fromEntries(metrics.connectionsPerSlab), messagesPerSec: parseFloat((metrics.messagesReceived / elapsedSec).toFixed(2)), bytesPerSec: parseInt((metrics.bytesSent / elapsedSec).toFixed(0), 10), @@ -472,6 +486,7 @@ export function getWebSocketMetrics(): any { export function setupWebSocket(server: Server): WebSocketServer { const wss = new WebSocketServer({ server, maxPayload: 1024 }); + liveWss = wss; // Idempotent re-entry: drop any listeners left over from a previous call // (tests, hot reload, restart) before re-registering. The eventBus is a diff --git a/tests/routes/health.test.ts b/tests/routes/health.test.ts index 6e08651..14ca7fe 100644 --- a/tests/routes/health.test.ts +++ b/tests/routes/health.test.ts @@ -54,6 +54,41 @@ describe("health routes", () => { vi.mocked(getSupabase).mockReturnValue(mockSupabase); }); + describe("WS liveness check (BUG-110)", () => { + it("flags ws as degraded when liveConnections (not totalConnections) is near the cap", async () => { + mockConnection.getSlot.mockResolvedValue(100); + mockSupabase.select.mockResolvedValue({ count: 5, error: null }); + // totalConnections (our slower bookkeeping) looks fine, but + // liveConnections (the real wss.clients.size) is saturated — the + // check must use the latter. + vi.mocked(getWebSocketMetrics).mockReturnValue({ + totalConnections: 0, + liveConnections: 999, + limits: { maxGlobalConnections: 1000 }, + }); + + const app = healthRoutes(); + const res = await app.request("/health"); + const data = await res.json(); + expect(data.checks.ws).toBe(false); + expect(data.status).not.toBe("ok"); + }); + + it("falls back to totalConnections when liveConnections is absent", async () => { + mockConnection.getSlot.mockResolvedValue(100); + mockSupabase.select.mockResolvedValue({ count: 5, error: null }); + vi.mocked(getWebSocketMetrics).mockReturnValue({ + totalConnections: 5, + limits: { maxGlobalConnections: 1000 }, + }); + + const app = healthRoutes(); + const res = await app.request("/health"); + const data = await res.json(); + expect(data.checks.ws).toBe(true); + }); + }); + it("should return 200 with ok status when RPC and DB work", async () => { mockConnection.getSlot.mockResolvedValue(123456789); mockSupabase.select.mockResolvedValue({ count: 5, error: null }); diff --git a/tests/routes/ws-live-metrics.test.ts b/tests/routes/ws-live-metrics.test.ts new file mode 100644 index 0000000..f592232 --- /dev/null +++ b/tests/routes/ws-live-metrics.test.ts @@ -0,0 +1,120 @@ +/** + * WS live connection metrics (BUG-110). + * + * metrics.totalConnections is only incremented once a connection clears the + * async IP-blocklist/auth-ban/rate-limit chain in the "connection" handler — + * a socket that's open at the `ws` library level but still mid-handshake is + * invisible to it. getWebSocketMetrics().liveConnections should instead + * reflect wss.clients.size, the library's own real-time count, so health + * checks built on it don't under-report true connection-slot pressure. + * + * This test stalls a connection inside the isAuthBanned() await (by mocking + * the shared store to never resolve it) to deterministically reproduce the + * gap between the two counters. + */ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import http from "node:http"; +import WebSocket from "ws"; + +let resolveAuthBanned: (() => void) | null = null; + +vi.mock("../../src/middleware/shared-store.js", () => ({ + getSharedStore: () => ({ + isAuthBanned: () => + new Promise((resolve) => { + resolveAuthBanned = () => resolve(false); + }), + recordAuthFailure: vi.fn().mockResolvedValue({ count: 0, bannedUntil: 0 }), + getConnectionCount: vi.fn().mockResolvedValue(0), + incrementConnectionCount: vi.fn().mockResolvedValue(undefined), + decrementConnectionCount: vi.fn().mockResolvedValue(undefined), + evictExpiredAuthFailures: vi.fn().mockResolvedValue(undefined), + }), +})); + +vi.mock("@percolator/shared", () => ({ + createLogger: vi.fn(() => ({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + })), + eventBus: { on: vi.fn(), off: vi.fn() }, + getSupabase: vi.fn(), + sanitizeSlabAddress: vi.fn((s: string) => s), + sendInfoAlert: vi.fn(), +})); + +function waitForOpen(ws: WebSocket): Promise { + return new Promise((resolve, reject) => { + ws.once("open", resolve); + ws.once("error", reject); + }); +} + +function waitForClose(ws: WebSocket): Promise { + return new Promise((resolve) => { + if (ws.readyState === WebSocket.CLOSED) return resolve(); + ws.once("close", () => resolve()); + }); +} + +describe("WS live connection metrics (BUG-110)", () => { + let server: http.Server; + + beforeEach(() => { + vi.resetModules(); + resolveAuthBanned = null; + process.env.NODE_ENV = "test"; + process.env.WS_AUTH_REQUIRED = "false"; + }); + + afterEach(async () => { + delete process.env.WS_AUTH_REQUIRED; + await new Promise((resolve) => server.close(() => resolve())); + }); + + it("reports liveConnections from the real wss instance while a connection is mid-handshake", async () => { + const { setupWebSocket, getWebSocketMetrics } = await import( + "../../src/routes/ws.js" + ); + + server = http.createServer(); + setupWebSocket(server as unknown as import("node:http").Server); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as { port: number }; + + const ws = new WebSocket(`ws://127.0.0.1:${port}/`); + try { + await waitForOpen(ws); + + // The socket is open at the ws-library level, but the server-side + // handler is stuck awaiting isAuthBanned() — it never reached + // clients.add(). + const metrics = getWebSocketMetrics(); + expect(metrics.totalConnections).toBe(0); + expect(metrics.liveConnections).toBe(1); + } finally { + // Unblock the stalled handler (even on assertion failure) so the + // connection can close and server.close() in afterEach doesn't hang. + resolveAuthBanned?.(); + await new Promise((r) => setTimeout(r, 20)); + ws.close(); + await waitForClose(ws); + } + }); + + it("falls back to totalConnections when no connection is mid-handshake", async () => { + const { setupWebSocket, getWebSocketMetrics } = await import( + "../../src/routes/ws.js" + ); + + server = http.createServer(); + setupWebSocket(server as unknown as import("node:http").Server); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + + const metrics = getWebSocketMetrics(); + expect(metrics.totalConnections).toBe(0); + expect(metrics.liveConnections).toBe(0); + }); +});