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

@@ -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<string, PooledConnection[]>();
   private waiters = new Map<string, ConnectionWaiter[]>();
+  private pendingConnections = new Map<string, number>();
+  private generations = new Map<string, number>();
+  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,7 +88,21 @@ class SSHConnectionPool {
     factory: () => Promise<Client>,
     existing: PooledConnection[],
   ): Promise<Client> {
+    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(),
@@ -91,6 +120,12 @@ class SSHConnectionPool {
       });
 
       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<Client>,
   ): Promise<Client> {
+    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();
   }

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<Client>((done) => {
+    resolve = done;
+  });
+  const client = new EventEmitter() as EventEmitter & {
+    end: ReturnType<typeof vi.fn>;
+    _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");
+  });
+});