From 2d907537b42c93fd053d114abeee036153963249 Mon Sep 17 00:00:00 2001
From: Luke Karrys <luke@lukekarrys.com>
Date: Fri, 29 Mar 2024 15:22:32 -0700
Subject: [PATCH] Synchronously update internal sockets length so http.Agent
 pooling is used

Fixes #299
---
 packages/agent-base/src/index.ts | 71 ++++++++++++++++++++++++++++----
 packages/agent-base/test/test.ts | 49 ++++++++++++++++++++++
 2 files changed, 112 insertions(+), 8 deletions(-)

diff --git a/packages/agent-base/src/index.ts b/packages/agent-base/src/index.ts
index f3758fb9..da58f528 100644
--- a/packages/agent-base/src/index.ts
+++ b/packages/agent-base/src/index.ts
@@ -77,26 +77,81 @@ export abstract class Agent extends http.Agent {
 			);
 	}
 
+	// In order to support async signatures in `connect()` and Node's native
+	// connection pooling in http.Agent, the array of sockets for each origin has
+	// to be updated syncronously. This is so the length of the array is accurate
+	// when `addRequest()` is next called. We achieve this by creating a fake socket
+	// and adding it to this.sockets and incrementing totalSocketCount.
+	private incrementSockets(name: string) {
+		// If maxSockets and maxTotalSockets are both Infinity then there is no need
+		// to create a fake socket because Node.js native connection pooling will
+		// never be invoked.
+		if (this.maxSockets === Infinity && this.maxTotalSockets === Infinity) {
+			return null;
+		}
+		if (!this.sockets[name]) {
+			// All instances of `sockets` are expected TypeScript errors. The alternative is to
+			// add it as a private property of this class but that will break TypeScript subclassing.
+			// @ts-expect-error `sockets` is readonly in `@types/node` but we need to write to it
+			this.sockets[name] = [];
+		}
+		const fakeSocket = new net.Socket({ writable: false });
+		// @ts-expect-error
+		this.sockets[name].push(fakeSocket);
+		// @ts-expect-error `totalSocketCount` isn't defined in `@types/node`
+		this.totalSocketCount++;
+		return fakeSocket;
+	}
+
+	private decrementSockets(name: string, socket: null | net.Socket) {
+		if (!this.sockets[name] || socket === null) {
+			return;
+		}
+		// @ts-expect-error
+		const index = this.sockets[name].indexOf(socket);
+		if (index !== -1) {
+			// @ts-expect-error
+			this.sockets[name].splice(index, 1);
+			// @ts-expect-error
+			this.totalSocketCount--;
+			// @ts-expect-error
+			if (this.sockets[name].length === 0) {
+				// @ts-expect-error
+				delete this.sockets[name];
+			}
+		}
+	}
+
 	createSocket(
 		req: http.ClientRequest,
 		options: AgentConnectOpts,
 		cb: (err: Error | null, s?: Duplex) => void
 	) {
+		// @ts-expect-error `getName()` isn't defined in `@types/node`
+		const name = this.getName(options);
+		const fakeSocket = this.incrementSockets(name);
 		const connectOpts = {
 			...options,
 			secureEndpoint: this.isSecureEndpoint(options),
 		};
 		Promise.resolve()
 			.then(() => this.connect(req, connectOpts))
-			.then((socket) => {
-				if (socket instanceof http.Agent) {
-					// @ts-expect-error `addRequest()` isn't defined in `@types/node`
-					return socket.addRequest(req, connectOpts);
+			.then(
+				(socket) => {
+					this.decrementSockets(name, fakeSocket);
+					if (socket instanceof http.Agent) {
+						// @ts-expect-error `addRequest()` isn't defined in `@types/node`
+						return socket.addRequest(req, connectOpts);
+					}
+					this[INTERNAL].currentSocket = socket;
+					// @ts-expect-error `createSocket()` isn't defined in `@types/node`
+					super.createSocket(req, options, cb);
+				},
+				(err) => {
+					this.decrementSockets(name, fakeSocket);
+					cb(err);
 				}
-				this[INTERNAL].currentSocket = socket;
-				// @ts-expect-error `createSocket()` isn't defined in `@types/node`
-				super.createSocket(req, options, cb);
-			}, cb);
+			);
 	}
 
 	createConnection(): Duplex {
diff --git a/packages/agent-base/test/test.ts b/packages/agent-base/test/test.ts
index 6f49a84d..d4a631a9 100644
--- a/packages/agent-base/test/test.ts
+++ b/packages/agent-base/test/test.ts
@@ -310,6 +310,55 @@ describe('Agent (TypeScript)', () => {
 		});
 	});
 
+	it('should support `keepAlive: true` with `maxSockets`', async () => {
+		let reqCount = 0;
+		let connectCount = 0;
+
+		class MyAgent extends Agent {
+			async connect(_req: http.ClientRequest, opts: AgentConnectOpts) {
+				connectCount++;
+				assert(opts.secureEndpoint === false);
+				await sleep(10);
+				return net.connect(opts);
+			}
+		}
+		const agent = new MyAgent({ keepAlive: true, maxSockets: 1 });
+
+		const server = http.createServer(async (req, res) => {
+			expect(req.headers.connection).toEqual('keep-alive');
+			reqCount++;
+			await sleep(10);
+			res.end();
+		});
+		const addr = await listen(server);
+
+		try {
+			const resPromise = req(new URL('/foo', addr), { agent });
+			const res2Promise = req(new URL('/another', addr), { agent });
+
+			const res = await resPromise;
+			expect(reqCount).toEqual(1);
+			expect(connectCount).toEqual(1);
+			expect(res.headers.connection).toEqual('keep-alive');
+
+			res.resume();
+			const s1 = res.socket;
+			await once(s1, 'free');
+
+			const res2 = await res2Promise;
+			expect(reqCount).toEqual(2);
+			expect(connectCount).toEqual(1);
+			expect(res2.headers.connection).toEqual('keep-alive');
+			assert(res2.socket === s1);
+
+			res2.resume();
+			await once(res2.socket, 'free');
+		} finally {
+			agent.destroy();
+			server.close();
+		}
+	});
+
 	describe('"https" module', () => {
 		it('should work for basic HTTPS requests', async () => {
 			let gotReq = false;