Add agent socket reuse

lib/agent.js

@@ -1,4 +1,5 @@
 const tcp = require('bare-tcp')
+const HTTPClientConnection = require('./client-connection')
 
 module.exports = class HTTPAgent {
   constructor (opts = {}) {
@@ -11,7 +12,12 @@ module.exports = class HTTPAgent {
       timeout = -1
     } = opts
 
+    this._sockets = new Map()
+    this._freeSockets = new Map()
+
     this._keepAlive = typeof keepAlive === 'number' ? keepAlive : keepAlive ? keepAliveMsecs : -1
+
+    this._opts = { ...opts }
     this._maxSockets = maxSockets
     this._maxTotalSockets = maxTotalSockets
     this._maxFreeSockets = maxFreeSockets
@@ -27,7 +33,87 @@ module.exports = class HTTPAgent {
   }
 
   keepSocketAlive (socket) {
-    return false
+    if (this._keepAlive === -1) return false
+
+    socket.setKeepAlive(true, this._keepAlive)
+    socket.unref()
+
+    return true
+  }
+
+  getName (opts) {
+    return `${opts.host}:${opts.port}`
+  }
+
+  addRequest (req, opts) {
+    opts = { ...opts, ...this._opts }
+
+    const name = this.getName(opts)
+
+    let socket
+
+    if (this._freeSockets.has(name)) {
+      const sockets = this._freeSockets.get(name)
+      socket = sockets.values().next().value
+      sockets.delete(socket)
+      if (sockets.size === 0) this._freeSockets.delete(name)
+
+      this.reuseSocket(socket, req)
+    } else {
+      socket = this.createConnection(opts)
+
+      socket
+        .on('free', () => this._onfree(socket, name))
+        .on('close', () => this._onremove(socket, name))
+    }
+
+    let sockets = this._sockets.get(name)
+    if (sockets === undefined) {
+      sockets = new Set()
+      this._sockets.set(name, sockets)
+    }
+
+    sockets.add(socket)
+
+    req.socket = socket
+
+    const connection = HTTPClientConnection.from(socket, opts)
+
+    connection.req = req
+  }
+
+  destroy () {
+    for (const set of [this._sockets, this._freeSockets]) {
+      for (const [, sockets] of set) {
+        for (const socket of sockets) socket.destroy()
+      }
+    }
+  }
+
+  _onfree (socket, name) {
+    if (this.keepSocketAlive(socket)) {
+      this._onremove(socket, name, false)
+
+      let sockets = this._freeSockets.get(name)
+      if (sockets === undefined) {
+        sockets = new Set()
+        this._freeSockets.set(name, sockets)
+      }
+
+      sockets.add(socket)
+    } else {
+      socket.end()
+    }
+  }
+
+  _onremove (socket, name, all = true) {
+    for (const set of all ? [this._sockets, this._freeSockets] : [this._sockets]) {
+      const sockets = set.get(name)
+      if (sockets === undefined) continue
+
+      sockets.delete(socket)
+      if (sockets.size === 0) set.delete(name)
+    }
   }
 
   static global = new this({ keepAlive: 1000, timeout: 5000 })

lib/client-connection.js

@@ -5,6 +5,16 @@ const errors = require('./errors')
 const empty = Buffer.alloc(0)
 
 module.exports = class HTTPClientConnection {
+  static _connections = new WeakMap()
+
+  static for (socket) {
+    return this._connections.get(socket) || null
+  }
+
+  static from (socket, opts) {
+    return this.for(socket) || new this(socket, opts)
+  }
+
   constructor (socket, opts = {}) {
     const {
       IncomingMessage = HTTPIncomingMessage
@@ -21,6 +31,7 @@ module.exports = class HTTPClientConnection {
     this._length = -1
     this._read = 0
     this._buffer = null
+    this._idle = true
 
     this._onerror = this._onerror.bind(this)
     this._onclose = this._onclose.bind(this)
@@ -36,6 +47,12 @@ module.exports = class HTTPClientConnection {
       .on('data', this._ondata)
       .on('drain', this._ondrain)
       .on('timeout', this._ontimeout)
+
+    HTTPClientConnection._connections.set(socket, this)
+  }
+
+  get idle () {
+    return this._idle
   }
 
   _onerror (err) {
@@ -51,6 +68,8 @@ module.exports = class HTTPClientConnection {
   }
 
   _ondata (data) {
+    this._idle = false
+
     if (this._state === constants.state.IN_BODY) return this._onbody(data)
 
     if (this._buffer !== null) {
@@ -181,13 +200,7 @@ module.exports = class HTTPClientConnection {
   }
 
   _onupgrade (head) {
-    this.socket
-      .off('error', this._onerror)
-      .off('close', this._onclose)
-      .off('end', this._onend)
-      .off('data', this._ondata)
-      .off('drain', this._ondrain)
-      .off('timeout', this._ontimeout)
+    this._ondetach()
 
     const req = this.req
 
@@ -206,18 +219,31 @@ module.exports = class HTTPClientConnection {
   _onfinished () {
     if (this.res) this.res.push(null)
     if (this.req) this.req._continueFinal()
-
-    this.socket.end()
   }
 
   _onreset () {
     this._state = constants.state.BEFORE_HEAD
     this._length = -1
     this._read = 0
     this._buffer = null
+    this._idle = true
+
+    this.socket.emit('free')
   }
 
   _ondrain () {
     if (this.req) this.req._continueWrite()
   }
+
+  _ondetach () {
+    this.socket
+      .off('error', this._onerror)
+      .off('close', this._onclose)
+      .off('end', this._onend)
+      .off('data', this._ondata)
+      .off('drain', this._ondrain)
+      .off('timeout', this._ontimeout)
+
+    HTTPClientConnection._connections.delete(this.socket)
+  }
 }

lib/client-request.js

@@ -1,6 +1,5 @@
 const HTTPAgent = require('./agent')
 const HTTPOutgoingMessage = require('./outgoing-message')
-const HTTPClientConnection = require('./client-connection')
 
 module.exports = class HTTPClientRequest extends HTTPOutgoingMessage {
   constructor (opts = {}, onresponse = null) {
@@ -17,19 +16,14 @@ module.exports = class HTTPClientRequest extends HTTPOutgoingMessage {
     const host = opts.host = opts.host || 'localhost'
     const port = opts.port = opts.port || 80
 
-    const {
-      connection = new HTTPClientConnection(agent.createConnection(opts), opts)
-    } = opts
+    super()
 
-    super(connection.socket)
-
-    connection.req = this
+    agent.addRequest(this, opts)
 
     this.method = method
     this.path = path
     this.headers = { host: host + ':' + port, ...opts.headers }
 
-    this._connection = connection
     this._chunked = method !== 'GET' && method !== 'HEAD'
 
     this._pendingFinal = null

lib/server-connection.js

@@ -1,10 +1,17 @@
+const tcp = require('bare-tcp')
 const HTTPIncomingMessage = require('./incoming-message')
 const HTTPServerResponse = require('./server-response')
 const constants = require('./constants')
 
 const empty = Buffer.alloc(0)
 
 module.exports = class HTTPServerConnection {
+  static _connections = new WeakMap()
+
+  static for (socket) {
+    return this._connections.get(socket) || null
+  }
+
   constructor (server, socket, opts = {}) {
     const {
       IncomingMessage = HTTPIncomingMessage,
@@ -24,6 +31,7 @@ module.exports = class HTTPServerConnection {
     this._length = -1
     this._read = 0
     this._buffer = null
+    this._idle = true
 
     this._onerror = this._onerror.bind(this)
     this._ondata = this._ondata.bind(this)
@@ -36,14 +44,22 @@ module.exports = class HTTPServerConnection {
       .on('drain', this._ondrain)
       .on('timeout', this._ontimeout)
 
+    HTTPServerConnection._connections.set(socket, this)
+
     if (this.server.timeout) socket.setTimeout(this.server.timeout)
   }
 
+  get idle () {
+    return this._idle
+  }
+
   _onerror (err) {
     this.socket.destroy(err)
   }
 
   _ondata (data) {
+    this._idle = false
+
     if (this._state === constants.state.IN_BODY) return this._onbody(data)
 
     if (this._buffer !== null) {
@@ -176,11 +192,7 @@ module.exports = class HTTPServerConnection {
   }
 
   _onupgrade (head) {
-    this.socket
-      .off('error', this._onerror)
-      .off('data', this._ondata)
-      .off('drain', this._ondrain)
-      .off('timeout', this._ontimeout)
+    this._ondetach()
 
     const req = this.req
 
@@ -207,9 +219,24 @@ module.exports = class HTTPServerConnection {
     this._length = -1
     this._read = 0
     this._buffer = null
+    this._idle = true
+
+    if (this.server._state & tcp.constants.state.CLOSING) {
+      this.socket.destroy()
+    }
   }
 
   _ondrain () {
     if (this.res) this.res._continueWrite()
   }
+
+  _ondetach () {
+    this.socket
+      .off('error', this._onerror)
+      .off('data', this._ondata)
+      .off('drain', this._ondrain)
+      .off('timeout', this._ontimeout)
+
+    HTTPServerConnection._connections.delete(this.socket)
+  }
 }

lib/server.js

@@ -28,4 +28,16 @@ module.exports = class HTTPServer extends TCPServer {
 
     return this
   }
+
+  close (onclose) {
+    super.close(onclose)
+
+    for (const socket of this._connections) {
+      const connection = HTTPServerConnection.for(socket)
+
+      if (connection && connection.idle) {
+        socket.destroy()
+      }
+    }
+  }
 }