fix: only ref socket when in use

Refs: #2026 (comment)
nodejs · Apr 11, 2023 · 0dea450 · 0dea450
1 parent b20405e
commit 0dea450
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 18 deletions.
diff --git a/lib/client.js b/lib/client.js
@@ -42,7 +42,6 @@ const {
   kConnected,
   kConnecting,
   kNeedDrain,
-  kNoRef,
   kKeepAliveDefaultTimeout,
   kHostHeader,
   kPendingIdx,
@@ -1091,7 +1090,6 @@ async function connect (client) {
 
     assert(socket)
 
-    socket[kNoRef] = false
     socket[kWriting] = false
     socket[kReset] = false
     socket[kBlocking] = false
@@ -1108,6 +1106,10 @@ async function connect (client) {
 
     client[kSocket] = socket
 
+    if (typeof client[kSocket].unref === 'function') {
+      client[kSocket].unref()
+    }
+
     if (channels.connected.hasSubscribers) {
       channels.connected.publish({
         connectParams: {
@@ -1194,16 +1196,6 @@ function _resume (client, sync) {
     const socket = client[kSocket]
 
     if (socket && !socket.destroyed) {
-      if (client[kSize] === 0) {
-        if (!socket[kNoRef] && socket.unref) {
-          socket.unref()
-          socket[kNoRef] = true
-        }
-      } else if (socket[kNoRef] && socket.ref) {
-        socket.ref()
-        socket[kNoRef] = false
-      }
-
       if (client[kSize] === 0) {
         if (socket[kParser].timeoutType !== TIMEOUT_IDLE) {
           socket[kParser].setTimeout(client[kKeepAliveTimeoutValue], TIMEOUT_IDLE)
@@ -1376,7 +1368,7 @@ function write (client, request) {
       errorRequest(client, request, err || new RequestAbortedError())
 
       util.destroy(socket, new InformationalError('aborted'))
-    })
+    }, socket)
   } catch (err) {
     errorRequest(client, request, err)
   }

diff --git a/lib/core/request.js b/lib/core/request.js
@@ -29,6 +29,7 @@ const headerCharRegex = /[^\t\x20-\x7e\x80-\xff]/
 const invalidPathRegex = /[^\u0021-\u00ff]/
 
 const kHandler = Symbol('handler')
+const kSocket = Symbol('socket')
 
 const channels = {}
 
@@ -191,6 +192,7 @@ class Request {
 
     this.servername = util.getServerName(this.host)
 
+    this[kSocket] = null
     this[kHandler] = handler
 
     if (channels.create.hasSubscribers) {
@@ -214,10 +216,16 @@ class Request {
     }
   }
 
-  onConnect (abort) {
+  onConnect (abort, socket) {
     assert(!this.aborted)
     assert(!this.completed)
 
+    this[kSocket] = socket
+
+    if (typeof this[kSocket].ref === 'function') {
+      this[kSocket].ref()
+    }
+
     return this[kHandler].onConnect(abort)
   }
 
@@ -249,22 +257,31 @@ class Request {
   onComplete (trailers) {
     assert(!this.aborted)
 
+    if (typeof this[kSocket].unref === 'function') {
+      this[kSocket].unref()
+    }
+
     this.completed = true
     if (channels.trailers.hasSubscribers) {
       channels.trailers.publish({ request: this, trailers })
     }
+
     return this[kHandler].onComplete(trailers)
   }
 
   onError (error) {
+    assert(!this.aborted)
+
+    if (typeof this[kSocket].unref === 'function') {
+      this[kSocket].unref()
+    }
+
     if (channels.error.hasSubscribers) {
       channels.error.publish({ request: this, error })
     }
 
-    if (this.aborted) {
-      return
-    }
     this.aborted = true
+
     return this[kHandler].onError(error)
   }
 

diff --git a/lib/core/symbols.js b/lib/core/symbols.js
@@ -19,7 +19,6 @@ module.exports = {
   kServerName: Symbol('server name'),
   kLocalAddress: Symbol('local address'),
   kHost: Symbol('host'),
-  kNoRef: Symbol('no ref'),
   kBodyUsed: Symbol('used'),
   kRunning: Symbol('running'),
   kBlocking: Symbol('blocking'),