fix: memory leak in Agent (#4425)

Author: hexchain

lib/dispatcher/agent.js

@@ -45,27 +45,14 @@ class Agent extends DispatcherBase {
     }
 
     this[kOnConnect] = (origin, targets) => {
-      const result = this[kClients].get(origin)
-      if (result) {
-        result.count += 1
-      }
       this.emit('connect', origin, [this, ...targets])
     }
 
     this[kOnDisconnect] = (origin, targets, err) => {
-      const result = this[kClients].get(origin)
-      if (result) {
-        result.count -= 1
-        if (result.count <= 0) {
-          this[kClients].delete(origin)
-          result.dispatcher.destroy()
-        }
-      }
       this.emit('disconnect', origin, [this, ...targets], err)
     }
 
     this[kOnConnectionError] = (origin, targets, err) => {
-      // TODO: should this decrement result.count here?
       this.emit('connectionError', origin, [this, ...targets], err)
     }
   }
@@ -89,11 +76,33 @@ class Agent extends DispatcherBase {
     const result = this[kClients].get(key)
     let dispatcher = result && result.dispatcher
     if (!dispatcher) {
+      const closeClientIfUnused = (connected) => {
+        const result = this[kClients].get(key)
+        if (result) {
+          if (connected) result.count -= 1
+          if (result.count <= 0) {
+            this[kClients].delete(key)
+            result.dispatcher.close()
+          }
+        }
+      }
       dispatcher = this[kFactory](opts.origin, this[kOptions])
         .on('drain', this[kOnDrain])
-        .on('connect', this[kOnConnect])
-        .on('disconnect', this[kOnDisconnect])
-        .on('connectionError', this[kOnConnectionError])
+        .on('connect', (origin, targets) => {
+          const result = this[kClients].get(key)
+          if (result) {
+            result.count += 1
+          }
+          this[kOnConnect](origin, targets)
+        })
+        .on('disconnect', (origin, targets, err) => {
+          closeClientIfUnused(true)
+          this[kOnDisconnect](origin, targets, err)
+        })
+        .on('connectionError', (origin, targets, err) => {
+          closeClientIfUnused(false)
+          this[kOnConnectionError](origin, targets, err)
+        })
 
       this[kClients].set(key, { count: 0, dispatcher })
     }

test/issue-4244.js

@@ -0,0 +1,62 @@
+const { test, describe } = require('node:test')
+const assert = require('node:assert')
+const { createServer } = require('node:http')
+const { request, Agent, Pool } = require('..')
+
+// https://github.com/nodejs/undici/issues/4424
+describe('Agent should close inactive clients', () => {
+  test('without active connections', async (t) => {
+    const server = createServer({ keepAliveTimeout: 0 }, async (_req, res) => {
+      res.setHeader('connection', 'close')
+      res.writeHead(200)
+      res.end('ok')
+    }).listen(0)
+
+    t.after(() => server.close())
+
+    let p
+    const agent = new Agent({
+      factory: (origin, opts) => {
+        const pool = new Pool(origin, opts)
+        let _resolve, _reject
+        p = new Promise((resolve, reject) => {
+          _resolve = resolve
+          _reject = reject
+        })
+        pool.on('disconnect', () => {
+          setImmediate(() => pool.destroyed ? _resolve() : _reject(new Error('client not destroyed')))
+        })
+        return pool
+      }
+    })
+    const { statusCode } = await request(`http://localhost:${server.address().port}`, { dispatcher: agent })
+    assert.equal(statusCode, 200)
+
+    await p
+  })
+
+  test('in case of connection error', async (t) => {
+    let p
+    const agent = new Agent({
+      factory: (origin, opts) => {
+        const pool = new Pool(origin, opts)
+        let _resolve, _reject
+        p = new Promise((resolve, reject) => {
+          _resolve = resolve
+          _reject = reject
+        })
+        pool.on('connectionError', () => {
+          setImmediate(() => pool.destroyed ? _resolve() : _reject(new Error('client not destroyed')))
+        })
+        return pool
+      }
+    })
+    try {
+      await request('http://localhost:0', { dispatcher: agent })
+    } catch (_) {
+      // ignore
+    }
+
+    await p
+  })
+})