From 8859620afb64934d82251d372f01bb5741f8eeab Mon Sep 17 00:00:00 2001
From: Robert Nagy <ronagy@icloud.com>
Date: Tue, 11 Apr 2023 09:55:38 +0200
Subject: [PATCH] fix: only ref socket when in use

Refs: https://github.com/nodejs/undici/issues/2026#issuecomment-1502713911
---
 lib/client.js       | 17 ++++-------------
 lib/core/request.js | 19 ++++++++++++++++++-
 lib/core/symbols.js |  1 -
 3 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/lib/client.js b/lib/client.js
index b230c368dab..49af3ecdd8d 100644
--- 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
@@ -1107,6 +1105,9 @@ async function connect (client) {
       .on('close', onSocketClose)
 
     client[kSocket] = socket
+    if (typeof client[kSocket].unref === 'function') {
+      client[kSocket].unref()
+    }
 
     if (channels.connected.hasSubscribers) {
       channels.connected.publish({
@@ -1194,16 +1195,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 +1367,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
index c82159e5ba0..bc02f739899 100644
--- 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 = {}
 
@@ -214,10 +215,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)
   }
 
@@ -253,6 +260,11 @@ class Request {
     if (channels.trailers.hasSubscribers) {
       channels.trailers.publish({ request: this, trailers })
     }
+
+    if (typeof this[kSocket].unref === 'function') {
+      this[kSocket].unref()
+    }
+
     return this[kHandler].onComplete(trailers)
   }
 
@@ -261,10 +273,15 @@ class Request {
       channels.error.publish({ request: this, error })
     }
 
+    if (typeof this[kSocket].unref === 'function') {
+      this[kSocket].unref()
+    }
+
     if (this.aborted) {
       return
     }
     this.aborted = true
+
     return this[kHandler].onError(error)
   }
 
diff --git a/lib/core/symbols.js b/lib/core/symbols.js
index c852107a72a..06f346a5b67 100644
--- 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'),