From 9ea64d7dbbfeec8520541817016b142a2ceb0136 Mon Sep 17 00:00:00 2001
From: Fedor Indutny <fedor@indutny.com>
Date: Fri, 12 May 2017 16:53:15 -0400
Subject: [PATCH] http: overridable keep-alive behavior of `Agent`

Introduce two overridable `Agent` methods:

* `keepSocketAlive(socket)`
* `reuseSocket(socket, req)`

These methods can be overridden by particular `Agent` class child to
make keep-alive behavior customizable.

Motivation: destroy persisted sockets after some configurable timeout.
It is very non-trivial to do it with available primitives. Such program
will most likely need to poke with undocumented events and methods of
`Agent`. With introduced API such behavior is easy to implement.
---
 doc/api/http.md                               | 36 ++++++++++
 lib/_http_agent.js                            | 22 ++++--
 test/parallel/test-http-keepalive-override.js | 67 +++++++++++++++++++
 3 files changed, 120 insertions(+), 5 deletions(-)
 create mode 100644 test/parallel/test-http-keepalive-override.js

diff --git a/doc/api/http.md b/doc/api/http.md
index 2ceeaf6dac04ea..493222d409d808 100644
--- a/doc/api/http.md
+++ b/doc/api/http.md
@@ -157,6 +157,42 @@ socket/stream from this function, or by passing the socket/stream to `callback`.
 
 `callback` has a signature of `(err, stream)`.
 
+### agent.keepSocketAlive(socket)
+<!-- YAML
+added: REPLACEME
+-->
+
+* `socket` {net.Socket}
+
+Called when `socket` is detached from a request and could be persisted by the
+Agent. Default behavior is to:
+
+```js
+socket.unref();
+socket.setKeepAlive(agent.keepAliveMsecs);
+```
+
+This method can be overridden by a particular `Agent` subclass. If this
+method returns a falsy value, the socket will be destroyed instead of persisting
+it for use with the next request.
+
+### agent.reuseSocket(socket, request)
+<!-- YAML
+added: REPLACEME
+-->
+
+* `socket` {net.Socket}
+* `request` {http.ClientRequest}
+
+Called when `socket` is attached to `request` after being persisted because of
+the keep-alive options. Default behavior is to:
+
+```js
+socket.ref();
+```
+
+This method can be overridden by a particular `Agent` subclass.
+
 ### agent.destroy()
 <!-- YAML
 added: v0.11.4
diff --git a/lib/_http_agent.js b/lib/_http_agent.js
index 60fff343dc66cf..9f0efe82d14313 100644
--- a/lib/_http_agent.js
+++ b/lib/_http_agent.js
@@ -90,15 +90,16 @@ function Agent(options) {
 
         if (count > self.maxSockets || freeLen >= self.maxFreeSockets) {
           socket.destroy();
-        } else {
+        } else if (self.keepSocketAlive(socket)) {
           freeSockets = freeSockets || [];
           self.freeSockets[name] = freeSockets;
-          socket.setKeepAlive(true, self.keepAliveMsecs);
-          socket.unref();
           socket[async_id_symbol] = -1;
           socket._httpMessage = null;
           self.removeSocket(socket, options);
           freeSockets.push(socket);
+        } else {
+          // Implementation doesn't want to keep socket alive
+          socket.destroy();
         }
       } else {
         socket.destroy();
@@ -169,13 +170,12 @@ Agent.prototype.addRequest = function addRequest(req, options, port/*legacy*/,
     // Assign the handle a new asyncId and run any init() hooks.
     socket._handle.asyncReset();
     socket[async_id_symbol] = socket._handle.getAsyncId();
-    debug('have free socket');
 
     // don't leak
     if (!this.freeSockets[name].length)
       delete this.freeSockets[name];
 
-    socket.ref();
+    this.reuseSocket(socket, req);
     req.onSocket(socket);
     this.sockets[name].push(socket);
   } else if (sockLen < this.maxSockets) {
@@ -306,6 +306,18 @@ Agent.prototype.removeSocket = function removeSocket(s, options) {
   }
 };
 
+Agent.prototype.keepSocketAlive = function keepSocketAlive(socket) {
+  socket.setKeepAlive(true, this.keepAliveMsecs);
+  socket.unref();
+
+  return true;
+};
+
+Agent.prototype.reuseSocket = function reuseSocket(socket, req) {
+  debug('have free socket');
+  socket.ref();
+};
+
 Agent.prototype.destroy = function destroy() {
   var sets = [this.freeSockets, this.sockets];
   for (var s = 0; s < sets.length; s++) {
diff --git a/test/parallel/test-http-keepalive-override.js b/test/parallel/test-http-keepalive-override.js
new file mode 100644
index 00000000000000..d25fc319747ad2
--- /dev/null
+++ b/test/parallel/test-http-keepalive-override.js
@@ -0,0 +1,67 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+
+const http = require('http');
+
+const server = http.createServer((req, res) => {
+  res.end('ok');
+}).listen(0, common.mustCall(() => {
+  const agent = http.Agent({
+    keepAlive: true,
+    maxSockets: 5,
+    maxFreeSockets: 2
+  });
+
+  const keepSocketAlive = agent.keepSocketAlive;
+  const reuseSocket = agent.reuseSocket;
+
+  let called = 0;
+  let expectedSocket;
+  agent.keepSocketAlive = common.mustCall((socket) => {
+    assert(socket);
+
+    called++;
+    if (called === 1) {
+      return false;
+    } else if (called === 2) {
+      expectedSocket = socket;
+      return keepSocketAlive.call(agent, socket);
+    }
+
+    assert.strictEqual(socket, expectedSocket);
+    return false;
+  }, 3);
+
+  agent.reuseSocket = common.mustCall((socket, req) => {
+    assert.strictEqual(socket, expectedSocket);
+    assert(req);
+
+    return reuseSocket.call(agent, socket, req);
+  }, 1);
+
+  function req(callback) {
+    http.request({
+      method: 'GET',
+      path: '/',
+      agent,
+      port: server.address().port
+    }, common.mustCall((res) => {
+      res.resume();
+      res.once('end', common.mustCall(() => {
+        setImmediate(callback);
+      }));
+    })).end();
+  }
+
+  // Should destroy socket instead of keeping it alive
+  req(common.mustCall(() => {
+    // Should keep socket alive
+    req(common.mustCall(() => {
+      // Should reuse the socket
+      req(common.mustCall(() => {
+        server.close();
+      }));
+    }));
+  }));
+}));