fix(pg-pool): avoid stale async context on client reuse.

packages/pg-pool/index.js

@@ -1,5 +1,6 @@
 'use strict'
 const EventEmitter = require('events').EventEmitter
+const { AsyncResource } = require('async_hooks')
 
 const NOOP = function () {}
 
@@ -18,8 +19,65 @@ class IdleItem {
 }
 
 class PendingItem {
-  constructor(callback) {
+  constructor(callback, asyncResource) {
     this.callback = callback
+    this.asyncResource = asyncResource
+  }
+}
+
+// Passed to `handle.asyncReset()` purely as diagnostic info describing what
+// kind of handle is being reset; mirrors Node's own `http.Agent`, which
+// resets a pooled socket's handle the same way when reusing it for a new
+// request (see `ReusedHandle` in `lib/_http_agent.js`).
+class ReusedHandle {
+  constructor(type, handle) {
+    this.type = type
+    this.handle = handle
+  }
+}
+
+// A client's underlying socket is a native resource: whatever async context
+// (and anything an `AsyncLocalStorage` consumer, e.g. APM instrumentation,
+// stored in it -- such as a whole in-flight trace) was active the moment it
+// connected stays reachable from it for as long as the socket itself is
+// alive, because that's what later events on it get dispatched with. That's
+// fine for a one-shot connection, but a pooled client is deliberately reused
+// across many unrelated logical operations over its lifetime, so without
+// this, whichever one happened to be active when the connection was first
+// opened would stay pinned in memory for the life of the pool.
+//
+// `asyncReset` is the API Node provides for exactly this situation -- handing
+// a long-lived native resource to a new, unrelated logical operation -- and
+// it's what `http.Agent` uses when handing a pooled, already-connected
+// socket to a new request. Call it here, at the equivalent moment for a
+// pooled `pg` connection, so the socket forgets the previous operation's
+// context instead of carrying it forward indefinitely.
+//
+// `asyncReset` alone is not quite enough on every Node version still in
+// support, though: on Node <24, `AsyncLocalStorage` additionally stamps a
+// resource with an internal `kResourceStore` symbol at creation time (used
+// to drive `executionAsyncResource()`-style APIs), and `asyncReset` does
+// not itself refresh that symbol -- confirmed via heap snapshot, even
+// though it does correctly fix which context *later events* dispatch into.
+// The stale value just sits there, unrelated to dispatch, reachable for as
+// long as the socket exists. Node 24+'s newer context-tracking does not
+// have this property at all. There is no public API for clearing it, so
+// this reaches for the symbol directly -- undesirable in general, but the
+// only way to avoid this version-dependent retention until Node exposes
+// one.
+function resetClientAsyncContext(client) {
+  const handle = client && client.connection && client.connection.stream && client.connection.stream._handle
+  if (!handle || typeof handle.asyncReset !== 'function') {
+    return
+  }
+
+  handle.asyncReset(new ReusedHandle(handle.getProviderType(), handle))
+
+  const resourceStoreSymbol = Object.getOwnPropertySymbols(handle).find((symbol) =>
+    symbol.toString().includes('kResourceStore')
+  )
+  if (resourceStoreSymbol) {
+    handle[resourceStoreSymbol] = undefined
   }
 }
 
@@ -196,6 +254,15 @@ class Pool extends EventEmitter {
     const response = promisify(this.Promise, cb)
     const result = response.result
 
+    // Captures the context that's active right now, at the moment this
+    // call was made. A request serviced immediately (or via a same-tick
+    // `process.nextTick`) naturally keeps this context without any help,
+    // but one that has to wait in `_pendingQueue` until a connection is
+    // freed by a later, unrelated `release()` would otherwise be serviced
+    // in *that* release's context instead of its own caller's -- see
+    // `_acquireClient`.
+    const asyncResource = new AsyncResource('pg-pool.connect')
+
     // if we don't have to connect a new client, don't do so
     if (this._isFull() || this._idle.length) {
       // if we have idle clients schedule a pulse immediately
@@ -204,7 +271,7 @@ class Pool extends EventEmitter {
       }
 
       if (!this.options.connectionTimeoutMillis) {
-        this._pendingQueue.push(new PendingItem(response.callback))
+        this._pendingQueue.push(new PendingItem(response.callback, asyncResource))
         return result
       }
 
@@ -213,7 +280,7 @@ class Pool extends EventEmitter {
         response.callback(err, res, done)
       }
 
-      const pendingItem = new PendingItem(queueCallback)
+      const pendingItem = new PendingItem(queueCallback, asyncResource)
 
       // set connection timeout on checking out an existing client
       const tid = setTimeout(() => {
@@ -232,7 +299,7 @@ class Pool extends EventEmitter {
       return result
     }
 
-    this.newClient(new PendingItem(response.callback))
+    this.newClient(new PendingItem(response.callback, asyncResource))
 
     return result
   }
@@ -333,36 +400,69 @@ class Pool extends EventEmitter {
 
   // acquire a client for a pending work item
   _acquireClient(client, pendingItem, idleListener, isNew) {
-    if (isNew) {
-      this.emit('connect', client)
-    }
+    // Drop pg-pool's own reference to the captured resource as soon as
+    // it's no longer needed, rather than leaving it reachable from
+    // `pendingItem` for however long something else (e.g. a callback
+    // closure kept around by the client's own connection) happens to keep
+    // that object alive -- it's only needed for the synchronous dispatch
+    // below.
+    const asyncResource = pendingItem.asyncResource
+    pendingItem.asyncResource = null
+
+    const dispatch = () => {
+      // Handing an already-connected client to a (possibly different)
+      // logical operation than the one that last used it -- reset its
+      // socket's native handle so it stops propagating whatever async
+      // context was active the last time it was used.
+      if (!isNew) {
+        resetClientAsyncContext(client)
+      }
+
+      if (isNew) {
+        this.emit('connect', client)
+      }
 
-    this.emit('acquire', client)
+      this.emit('acquire', client)
 
-    client.release = this._releaseOnce(client, idleListener)
+      client.release = this._releaseOnce(client, idleListener)
 
-    client.removeListener('error', idleListener)
+      client.removeListener('error', idleListener)
 
-    if (!pendingItem.timedOut) {
-      if (isNew && this.options.verify) {
-        this.options.verify(client, (err) => {
-          if (err) {
-            client.release(err)
-            return pendingItem.callback(err, undefined, NOOP)
-          }
+      if (!pendingItem.timedOut) {
+        if (isNew && this.options.verify) {
+          this.options.verify(client, (err) => {
+            if (err) {
+              client.release(err)
+              return pendingItem.callback(err, undefined, NOOP)
+            }
 
+            pendingItem.callback(undefined, client, client.release)
+          })
+        } else {
           pendingItem.callback(undefined, client, client.release)
-        })
-      } else {
-        pendingItem.callback(undefined, client, client.release)
-      }
-    } else {
-      if (isNew && this.options.verify) {
-        this.options.verify(client, client.release)
+        }
       } else {
-        client.release()
+        if (isNew && this.options.verify) {
+          this.options.verify(client, client.release)
+        } else {
+          client.release()
+        }
       }
     }
+
+    if (isNew || !asyncResource) {
+      return dispatch()
+    }
+
+    // `_acquireClient` can run inside a completely unrelated call stack --
+    // e.g. triggered by a different client's `release()` -- rather than the
+    // call stack of whichever caller is actually receiving this client.
+    // Dispatch inside the context captured when that caller originally
+    // called `connect()`/`query()`, so `resetClientAsyncContext` above (and
+    // the `pendingItem.callback` it leads to) see the right context, the
+    // same way `http.Agent` replays a queued request's original async
+    // context when handing it a freed socket.
+    return asyncResource.runInAsyncScope(dispatch)
   }
 
   // returns a function that wraps _release and throws if called more than once

packages/pg-pool/test/async-context-reset.js

@@ -0,0 +1,104 @@
+'use strict'
+const Pool = require('../')
+const expect = require('expect.js')
+const { AsyncLocalStorage } = require('async_hooks')
+
+describe('async context on client reuse', () => {
+  it('resets the socket handle async context when handing out an already-connected client', async () => {
+    // Force every connect() after the first to reuse the SAME client.
+    const pool = new Pool({ max: 1 })
+
+    try {
+      const first = await pool.connect()
+      const handle = first.connection.stream._handle
+      expect(handle).to.be.ok()
+      expect(typeof handle.asyncReset).to.equal('function')
+
+      let resetCalls = 0
+      const originalAsyncReset = handle.asyncReset
+      handle.asyncReset = function (...args) {
+        resetCalls++
+        return originalAsyncReset.apply(this, args)
+      }
+
+      first.release()
+
+      const second = await pool.connect()
+      try {
+        expect(second).to.equal(first)
+        expect(resetCalls).to.equal(1)
+      } finally {
+        second.release()
+      }
+    } finally {
+      await pool.end()
+    }
+  })
+
+  it('clears the legacy kResourceStore context when handing out a reused client', async () => {
+    // On Node <24, AsyncLocalStorage stamps a resource with an internal
+    // kResourceStore symbol at creation time; this only happens while an
+    // ALS context is actually active, so one has to be running here for
+    // the symbol to exist at all.
+    const als = new AsyncLocalStorage()
+    const pool = new Pool({ max: 1 })
+
+    try {
+      const first = await als.run('first-connect', () => pool.connect())
+      const handle = first.connection.stream._handle
+      const resourceStoreSymbol = Object.getOwnPropertySymbols(handle).find((symbol) =>
+        symbol.toString().includes('kResourceStore')
+      )
+
+      if (!resourceStoreSymbol) {
+        // Node 24+ tracks async context differently and never stamps this
+        // symbol in the first place, so there's nothing to clear.
+        first.release()
+        return
+      }
+
+      expect(handle[resourceStoreSymbol]).to.equal('first-connect')
+
+      first.release()
+
+      const second = await pool.connect()
+      try {
+        expect(second).to.equal(first)
+        expect(handle[resourceStoreSymbol]).to.equal(undefined)
+      } finally {
+        second.release()
+      }
+    } finally {
+      await pool.end()
+    }
+  })
+
+  it('does not leak the connecting context into a later query on a reused client', (done) => {
+    const als = new AsyncLocalStorage()
+    const pool = new Pool({ max: 1 })
+
+    als.run('first-context', () => {
+      pool.query('SELECT 1', (err) => {
+        if (err) {
+          return pool.end(() => done(err))
+        }
+
+        als.run('second-context', () => {
+          pool.query('SELECT 1', (err2) => {
+            if (err2) {
+              return pool.end(() => done(err2))
+            }
+
+            let assertionError = null
+            try {
+              expect(als.getStore()).to.equal('second-context')
+            } catch (e) {
+              assertionError = e
+            }
+            pool.end(() => done(assertionError))
+          })
+        })
+      })
+    })
+  })
+})