From 3d1bdc254c5a8edd081e2186f3ae2fb1211b7c2e Mon Sep 17 00:00:00 2001
From: rickyes <ives199511@gmail.com>
Date: Thu, 28 May 2020 23:15:55 +0800
Subject: [PATCH] http: add maxTotalSockets to agent class

Add maxTotalSockets to determine how many sockets an agent can open.
Unlike maxSockets, The maxTotalSockets does not count by per origin.

PR-URL: https://github.com/nodejs/node/pull/33617
Backport-PR-URL: https://github.com/nodejs/node/pull/35396
Fixes: https://github.com/nodejs/node/issues/31942
Reviewed-By: Robert Nagy <ronagy@icloud.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
---
 doc/api/http.md                               |  10 ++
 lib/_http_agent.js                            |  58 ++++++++-
 .../test-http-agent-maxtotalsockets.js        | 113 ++++++++++++++++++
 3 files changed, 177 insertions(+), 4 deletions(-)
 create mode 100644 test/parallel/test-http-agent-maxtotalsockets.js

diff --git a/doc/api/http.md b/doc/api/http.md
index d7e8c4caa852f1..88949e697f7b83 100644
--- a/doc/api/http.md
+++ b/doc/api/http.md
@@ -285,6 +285,16 @@ added: v0.3.6
 By default set to `Infinity`. Determines how many concurrent sockets the agent
 can have open per origin. Origin is the returned value of [`agent.getName()`][].
 
+### `agent.maxTotalSockets`
+<!-- YAML
+added: REPLACEME
+-->
+
+* {number}
+
+By default set to `Infinity`. Determines how many concurrent sockets the agent
+can have open. Unlike `maxSockets`, this parameter applies across all origins.
+
 ### `agent.requests`
 <!-- YAML
 added: v0.5.9
diff --git a/lib/_http_agent.js b/lib/_http_agent.js
index d75d9f341347c6..73b9ed2c32896e 100644
--- a/lib/_http_agent.js
+++ b/lib/_http_agent.js
@@ -22,6 +22,7 @@
 'use strict';
 
 const {
+  NumberIsNaN,
   ObjectKeys,
   ObjectSetPrototypeOf,
   ObjectValues,
@@ -34,7 +35,15 @@ let debug = require('internal/util/debuglog').debuglog('http', (fn) => {
   debug = fn;
 });
 const { async_id_symbol } = require('internal/async_hooks').symbols;
+const {
+  codes: {
+    ERR_OUT_OF_RANGE,
+  },
+} = require('internal/errors');
+const { validateNumber } = require('internal/validators');
+
 const kOnKeylog = Symbol('onkeylog');
+const kRequestOptions = Symbol('requestOptions');
 // New Agent code.
 
 // The largest departure from the previous implementation is that
@@ -81,6 +90,17 @@ function Agent(options) {
   this.keepAlive = this.options.keepAlive || false;
   this.maxSockets = this.options.maxSockets || Agent.defaultMaxSockets;
   this.maxFreeSockets = this.options.maxFreeSockets || 256;
+  this.maxTotalSockets = this.options.maxTotalSockets;
+  this.totalSocketCount = 0;
+
+  if (this.maxTotalSockets !== undefined) {
+    validateNumber(this.maxTotalSockets, 'maxTotalSockets');
+    if (this.maxTotalSockets <= 0 || NumberIsNaN(this.maxTotalSockets))
+      throw new ERR_OUT_OF_RANGE('maxTotalSockets', '> 0',
+                                 this.maxTotalSockets);
+  } else {
+    this.maxTotalSockets = Infinity;
+  }
 
   this.on('free', (socket, options) => {
     const name = this.getName(options);
@@ -113,7 +133,9 @@ function Agent(options) {
         if (this.sockets[name])
           count += this.sockets[name].length;
 
-        if (count > this.maxSockets || freeLen >= this.maxFreeSockets) {
+        if (this.totalSocketCount > this.maxTotalSockets ||
+            count > this.maxSockets ||
+            freeLen >= this.maxFreeSockets) {
           socket.destroy();
         } else if (this.keepSocketAlive(socket)) {
           freeSockets = freeSockets || [];
@@ -236,7 +258,9 @@ Agent.prototype.addRequest = function addRequest(req, options, port/* legacy */,
     this.reuseSocket(socket, req);
     setRequestSocket(this, req, socket);
     this.sockets[name].push(socket);
-  } else if (sockLen < this.maxSockets) {
+    this.totalSocketCount++;
+  } else if (sockLen < this.maxSockets &&
+             this.totalSocketCount < this.maxTotalSockets) {
     debug('call onSocket', sockLen, freeLen);
     // If we are under maxSockets create a new one.
     this.createSocket(req, options, handleSocketCreation(this, req, true));
@@ -246,6 +270,10 @@ Agent.prototype.addRequest = function addRequest(req, options, port/* legacy */,
     if (!this.requests[name]) {
       this.requests[name] = [];
     }
+
+    // Used to create sockets for pending requests from different origin
+    req[kRequestOptions] = options;
+
     this.requests[name].push(req);
   }
 };
@@ -275,7 +303,8 @@ Agent.prototype.createSocket = function createSocket(req, options, cb) {
       this.sockets[name] = [];
     }
     this.sockets[name].push(s);
-    debug('sockets', name, this.sockets[name].length);
+    this.totalSocketCount++;
+    debug('sockets', name, this.sockets[name].length, this.totalSocketCount);
     installListeners(this, s, options);
     cb(null, s);
   };
@@ -376,17 +405,38 @@ Agent.prototype.removeSocket = function removeSocket(s, options) {
         // Don't leak
         if (sockets[name].length === 0)
           delete sockets[name];
+        this.totalSocketCount--;
       }
     }
   }
 
+  let req;
   if (this.requests[name] && this.requests[name].length) {
     debug('removeSocket, have a request, make a socket');
-    const req = this.requests[name][0];
+    req = this.requests[name][0];
+  } else {
+    // TODO(rickyes): this logic will not be FIFO across origins.
+    // There might be older requests in a different origin, but
+    // if the origin which releases the socket has pending requests
+    // that will be prioritized.
+    for (const prop in this.requests) {
+      // Check whether this specific origin is already at maxSockets
+      if (this.sockets[prop] && this.sockets[prop].length) break;
+      debug('removeSocket, have a request with different origin,' +
+        ' make a socket');
+      req = this.requests[prop][0];
+      options = req[kRequestOptions];
+      break;
+    }
+  }
+
+  if (req && options) {
+    req[kRequestOptions] = undefined;
     // If we have pending requests and a socket gets closed make a new one
     const socketCreationHandler = handleSocketCreation(this, req, false);
     this.createSocket(req, options, socketCreationHandler);
   }
+
 };
 
 Agent.prototype.keepSocketAlive = function keepSocketAlive(socket) {
diff --git a/test/parallel/test-http-agent-maxtotalsockets.js b/test/parallel/test-http-agent-maxtotalsockets.js
new file mode 100644
index 00000000000000..c44c8db627d330
--- /dev/null
+++ b/test/parallel/test-http-agent-maxtotalsockets.js
@@ -0,0 +1,113 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const http = require('http');
+const Countdown = require('../common/countdown');
+
+assert.throws(() => new http.Agent({
+  maxTotalSockets: 'test',
+}), {
+  code: 'ERR_INVALID_ARG_TYPE',
+  name: 'TypeError',
+  message: 'The "maxTotalSockets" argument must be of type number. ' +
+    "Received type string ('test')",
+});
+
+[-1, 0, NaN].forEach((item) => {
+  assert.throws(() => new http.Agent({
+    maxTotalSockets: item,
+  }), {
+    code: 'ERR_OUT_OF_RANGE',
+    name: 'RangeError',
+    message: 'The value of "maxTotalSockets" is out of range. ' +
+      `It must be > 0. Received ${item}`,
+  });
+});
+
+assert.ok(new http.Agent({
+  maxTotalSockets: Infinity,
+}));
+
+function start(param = {}) {
+  const { maxTotalSockets, maxSockets } = param;
+
+  const agent = new http.Agent({
+    keepAlive: true,
+    keepAliveMsecs: 1000,
+    maxTotalSockets,
+    maxSockets,
+    maxFreeSockets: 3
+  });
+
+  const server = http.createServer(common.mustCall((req, res) => {
+    res.end('hello world');
+  }, 6));
+  const server2 = http.createServer(common.mustCall((req, res) => {
+    res.end('hello world');
+  }, 6));
+
+  server.keepAliveTimeout = 0;
+  server2.keepAliveTimeout = 0;
+
+  const countdown = new Countdown(12, () => {
+    assert.strictEqual(getRequestCount(), 0);
+    agent.destroy();
+    server.close();
+    server2.close();
+  });
+
+  function handler(s) {
+    for (let i = 0; i < 6; i++) {
+      http.get({
+        host: 'localhost',
+        port: s.address().port,
+        agent,
+        path: `/${i}`,
+      }, common.mustCall((res) => {
+        assert.strictEqual(res.statusCode, 200);
+        res.resume();
+        res.on('end', common.mustCall(() => {
+          for (const key of Object.keys(agent.sockets)) {
+            assert(agent.sockets[key].length <= maxSockets);
+          }
+          assert(getTotalSocketsCount() <= maxTotalSockets);
+          countdown.dec();
+        }));
+      }));
+    }
+  }
+
+  function getTotalSocketsCount() {
+    let num = 0;
+    for (const key of Object.keys(agent.sockets)) {
+      num += agent.sockets[key].length;
+    }
+    return num;
+  }
+
+  function getRequestCount() {
+    let num = 0;
+    for (const key of Object.keys(agent.requests)) {
+      num += agent.requests[key].length;
+    }
+    return num;
+  }
+
+  server.listen(0, common.mustCall(() => handler(server)));
+  server2.listen(0, common.mustCall(() => handler(server2)));
+}
+
+// If maxTotalSockets is larger than maxSockets,
+// then the origin check will be skipped
+// when the socket is removed.
+[{
+  maxTotalSockets: 2,
+  maxSockets: 3,
+}, {
+  maxTotalSockets: 3,
+  maxSockets: 2,
+}, {
+  maxTotalSockets: 2,
+  maxSockets: 2,
+}].forEach(start);