diff --git a/packages/pg-pool/index.js b/packages/pg-pool/index.js index ab514fa88..e37c3c0cc 100644 --- a/packages/pg-pool/index.js +++ b/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 diff --git a/packages/pg-pool/test/async-context-reset.js b/packages/pg-pool/test/async-context-reset.js new file mode 100644 index 000000000..c4e667964 --- /dev/null +++ b/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)) + }) + }) + }) + }) + }) +})