From 703e8cd037a5e29ce846e5206988ce3732412f5c Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Fri, 28 Aug 2026 10:36:31 +0800 Subject: [PATCH] fix: enforce SSH pool connection limits (#1353) * fix: enforce SSH pool connection limits * fix: discard stale pooled connections --- src/backend/hosts/ssh-connection-pool.ts | 84 ++++++++++++++----- .../tests/hosts/ssh-connection-pool.test.ts | 68 +++++++++++++++ 2 files changed, 133 insertions(+), 19 deletions(-) create mode 100644 src/backend/tests/hosts/ssh-connection-pool.test.ts diff --git a/src/backend/hosts/ssh-connection-pool.ts b/src/backend/hosts/ssh-connection-pool.ts index 0f0691b7..b4cf0b91 100644 --- a/src/backend/hosts/ssh-connection-pool.ts +++ b/src/backend/hosts/ssh-connection-pool.ts @@ -20,9 +20,12 @@ const DEFAULT_MAX_WAIT_MS = 30_000; const IDLE_MAX_AGE_MS = 10 * 60 * 1000; const CLEANUP_INTERVAL_MS = 2 * 60 * 1000; -class SSHConnectionPool { +export class SSHConnectionPool { private connections = new Map(); private waiters = new Map(); + private pendingConnections = new Map(); + private generations = new Map(); + private destroyed = false; private maxConnectionsPerHost = DEFAULT_MAX_CONNECTIONS_PER_HOST; private maxWaitMs = DEFAULT_MAX_WAIT_MS; private cleanupInterval: NodeJS.Timeout; @@ -31,6 +34,18 @@ class SSHConnectionPool { this.cleanupInterval = setInterval(() => { this.cleanup(); }, CLEANUP_INTERVAL_MS); + this.cleanupInterval.unref(); + } + + private hasCapacity(key: string, connections: PooledConnection[]): boolean { + return ( + connections.length + (this.pendingConnections.get(key) || 0) < + this.maxConnectionsPerHost + ); + } + + private invalidatePendingConnections(key: string): void { + this.generations.set(key, (this.generations.get(key) || 0) + 1); } private isConnectionHealthy(client: Client): boolean { @@ -73,24 +88,44 @@ class SSHConnectionPool { factory: () => Promise, existing: PooledConnection[], ): Promise { - const client = await factory(); - const pooled: PooledConnection = { - client, - lastUsed: Date.now(), - inUse: true, - hostKey: key, - }; - existing.push(pooled); - this.connections.set(key, existing); + const generation = this.generations.get(key) || 0; + this.pendingConnections.set( + key, + (this.pendingConnections.get(key) || 0) + 1, + ); + try { + const client = await factory(); + if (this.destroyed || (this.generations.get(key) || 0) !== generation) { + try { + client.end(); + } catch { + // expected + } + throw new Error(`SSH connection pool cleared for ${key}`); + } + const pooled: PooledConnection = { + client, + lastUsed: Date.now(), + inUse: true, + hostKey: key, + }; + existing.push(pooled); + this.connections.set(key, existing); - client.on("end", () => { - this.removeConnection(key, client); - }); - client.on("close", () => { - this.removeConnection(key, client); - }); + client.on("end", () => { + this.removeConnection(key, client); + }); + client.on("close", () => { + this.removeConnection(key, client); + }); - return client; + return client; + } finally { + const pending = (this.pendingConnections.get(key) || 1) - 1; + if (pending === 0) this.pendingConnections.delete(key); + else this.pendingConnections.set(key, pending); + this.wakeWaiter(key); + } } private enqueueWaiter( @@ -147,7 +182,7 @@ class SSHConnectionPool { } } - if (connections.length < this.maxConnectionsPerHost) { + if (this.hasCapacity(key, connections)) { const waiter = queue.shift()!; if (queue.length === 0) this.waiters.delete(key); else this.waiters.set(key, queue); @@ -173,6 +208,9 @@ class SSHConnectionPool { key: string, factory: () => Promise, ): Promise { + if (this.destroyed) { + throw new Error("SSH connection pool destroyed"); + } let connections = this.connections.get(key) || []; const available = connections.find((conn) => !conn.inUse); @@ -186,7 +224,7 @@ class SSHConnectionPool { } } - if (connections.length < this.maxConnectionsPerHost) { + if (this.hasCapacity(key, connections)) { return this.createPooledClient(key, factory, connections); } @@ -217,6 +255,7 @@ class SSHConnectionPool { } clearKeyConnections(key: string): void { + this.invalidatePendingConnections(key); const connections = this.connections.get(key) || []; for (const conn of connections) { try { @@ -263,6 +302,12 @@ class SSHConnectionPool { } clearAllConnections(): void { + const keys = new Set([ + ...this.connections.keys(), + ...this.pendingConnections.keys(), + ...this.waiters.keys(), + ]); + for (const key of keys) this.invalidatePendingConnections(key); for (const key of [...this.waiters.keys()]) { this.rejectWaiters(key, "SSH connection pool destroyed"); } @@ -279,6 +324,7 @@ class SSHConnectionPool { } destroy(): void { + this.destroyed = true; clearInterval(this.cleanupInterval); this.clearAllConnections(); } diff --git a/src/backend/tests/hosts/ssh-connection-pool.test.ts b/src/backend/tests/hosts/ssh-connection-pool.test.ts new file mode 100644 index 00000000..ba3f4700 --- /dev/null +++ b/src/backend/tests/hosts/ssh-connection-pool.test.ts @@ -0,0 +1,68 @@ +import { EventEmitter } from "node:events"; +import { describe, expect, it, vi } from "vitest"; +import type { Client } from "ssh2"; +import { SSHConnectionPool } from "../../hosts/ssh-connection-pool.js"; + +function deferredClient() { + let resolve!: (client: Client) => void; + const promise = new Promise((done) => { + resolve = done; + }); + const client = new EventEmitter() as EventEmitter & { + end: ReturnType; + _sock: { destroyed: boolean; writable: boolean }; + }; + client.end = vi.fn(); + client._sock = { destroyed: false, writable: true }; + return { client: client as unknown as Client, promise, resolve }; +} + +describe("SSHConnectionPool", () => { + it("counts pending factories against the per-host connection limit", async () => { + const pool = new SSHConnectionPool(); + const clients = Array.from({ length: 4 }, deferredClient); + let factoryCalls = 0; + const factory = vi.fn(() => clients[factoryCalls++].promise); + + const requests = Array.from({ length: 4 }, () => + pool.getConnection("same-host", factory), + ); + await vi.waitFor(() => expect(factory).toHaveBeenCalledTimes(3)); + + clients[0].resolve(clients[0].client); + const first = await requests[0]; + pool.releaseConnection("same-host", first); + + await expect(requests[3]).resolves.toBe(first); + expect(factory).toHaveBeenCalledTimes(3); + + clients[1].resolve(clients[1].client); + clients[2].resolve(clients[2].client); + await Promise.all([requests[1], requests[2]]); + pool.destroy(); + }); + + it("discards a pending connection when its host pool is cleared", async () => { + const pool = new SSHConnectionPool(); + const pending = deferredClient(); + const request = pool.getConnection("cleared-host", () => pending.promise); + + pool.clearKeyConnections("cleared-host"); + pending.resolve(pending.client); + + await expect(request).rejects.toThrow( + "SSH connection pool cleared for cleared-host", + ); + expect(pending.client.end).toHaveBeenCalledOnce(); + pool.destroy(); + }); + + it("rejects new connections after the pool is destroyed", async () => { + const pool = new SSHConnectionPool(); + pool.destroy(); + + await expect( + pool.getConnection("host", () => deferredClient().promise), + ).rejects.toThrow("SSH connection pool destroyed"); + }); +});