From a972359cad44e76745fb0c5555ddd669bfbdd88c Mon Sep 17 00:00:00 2001
From: Robert Nagy <ronagy@icloud.com>
Date: Thu, 8 Dec 2022 17:15:08 +0100
Subject: [PATCH] perf: Keep weak reference to session for re-use

---
 lib/core/connect.js | 83 +++++++++++++++++++++++++++++++++------------
 1 file changed, 62 insertions(+), 21 deletions(-)

diff --git a/lib/core/connect.js b/lib/core/connect.js
index a73f3c32b9a..0e09040f876 100644
--- a/lib/core/connect.js
+++ b/lib/core/connect.js
@@ -4,6 +4,7 @@ const net = require('net')
 const assert = require('assert')
 const util = require('./util')
 const { InvalidArgumentError, ConnectTimeoutError } = require('./errors')
+
 let tls // include tls conditionally since it is not always available
 
 // TODO: session re-use does not wait for the first
@@ -11,15 +12,71 @@ let tls // include tls conditionally since it is not always available
 // resolve the same servername multiple times even when
 // re-use is enabled.
 
+let SessionCache
+if (global.FinalizationRegistry) {
+  SessionCache = class WeakSessionCache {
+    constructor (maxCachedSessions) {
+      this._maxCachedSessions = maxCachedSessions
+      this._sessionCache = new Map()
+      this._sessionRegistry = new global.FinalizationRegistry((key) => {
+        if (this._sessionCache.size < this._maxCachedSessions) {
+          return
+        }
+
+        const ref = this._sessionCache.get(key)
+        if (ref !== undefined && ref.deref() === undefined) {
+          this._sessionCache.delete(key)
+        }
+      })
+    }
+
+    get (sessionKey) {
+      const ref = this._sessionCache.get(sessionKey)
+      return ref ? ref.deref() : null
+    }
+
+    set (sessionKey, session) {
+      if (this._maxCachedSessions === 0) {
+        return
+      }
+
+      this._sessionCache.set(sessionKey, new WeakRef(session))
+      this._sessionRegistry.register(session, sessionKey)
+    }
+  }
+} else {
+  SessionCache = class SimpleSessionCache {
+    constructor (maxCachedSessions) {
+      this._maxCachedSessions = maxCachedSessions
+      this._sessionCache = new Map()
+    }
+
+    get (sessionKey) {
+      return this._sessionCache.get(sessionKey)
+    }
+
+    set (sessionKey, session) {
+      if (this._sessionCache.size >= this._maxCachedSessions) {
+        // remove the oldest session
+        const { value: oldestKey } = this._sessionCache.keys().next()
+        this._sessionCache.delete(oldestKey)
+      }
+
+      this._sessionCache.set(sessionKey, session)
+    }
+  }
+}
+
 function buildConnector ({ maxCachedSessions, socketPath, timeout, ...opts }) {
   if (maxCachedSessions != null && (!Number.isInteger(maxCachedSessions) || maxCachedSessions < 0)) {
     throw new InvalidArgumentError('maxCachedSessions must be a positive integer or zero')
   }
 
   const options = { path: socketPath, ...opts }
-  const sessionCache = new Map()
+  const sessionCache = maxCachedSessions === 0
+    ? null
+    : new SessionCache(maxCachedSessions == null ? 100 : maxCachedSessions)
   timeout = timeout == null ? 10e3 : timeout
-  maxCachedSessions = maxCachedSessions == null ? 100 : maxCachedSessions
 
   return function connect ({ hostname, host, protocol, port, servername, localAddress, httpSocket }, callback) {
     let socket
@@ -30,7 +87,7 @@ function buildConnector ({ maxCachedSessions, socketPath, timeout, ...opts }) {
       servername = servername || options.servername || util.getServerName(host) || null
 
       const sessionKey = servername || hostname
-      const session = sessionCache.get(sessionKey) || null
+      const session = sessionCache?.get(sessionKey) || null
 
       assert(sessionKey)
 
@@ -47,24 +104,8 @@ function buildConnector ({ maxCachedSessions, socketPath, timeout, ...opts }) {
 
       socket
         .on('session', function (session) {
-          // cache is disabled
-          if (maxCachedSessions === 0) {
-            return
-          }
-
-          if (sessionCache.size >= maxCachedSessions) {
-            // remove the oldest session
-            const { value: oldestKey } = sessionCache.keys().next()
-            sessionCache.delete(oldestKey)
-          }
-
-          sessionCache.set(sessionKey, session)
-        })
-        .on('error', function (err) {
-          if (sessionKey && err.code !== 'UND_ERR_INFO') {
-            // TODO (fix): Only delete for session related errors.
-            sessionCache.delete(sessionKey)
-          }
+          // TODO (fix): Can a session become invalid once established? Don't think so?
+          sessionCache?.set(sessionKey, session)
         })
     } else {
       assert(!httpSocket, 'httpSocket can only be sent on TLS update')