async_hooks issues

Hi!

Thanks for this great library!

Unfortunately I have the same problem as mentioned here: #2404

I.e. the query-finish callback is called in a different execution context. Or something. I find the async hooks confusing but I need them to track which sql queries were initiated by which request in my api.

I have something reproducible!

const pg = require('pg')
const {AsyncLocalStorage} = require('async_hooks')

const asyncLocalStorage = new AsyncLocalStorage()

function getExecId() {
	return asyncLocalStorage.getStore()
}

async function sleep(ms) {
	return new Promise((resolve, reject) => {
		setTimeout(() => resolve(), ms)
	})
}

const pool = new pg.Pool({
	host: 'localhost',
	port: 5432,
	user: '...',
	password: '...',
	database: '...',
})

function doThings(id) {
	asyncLocalStorage.run(id, async () => {
		// everything in here should always see the same value from `getExecId()`
		console.log('query starts in', getExecId())
		pool.query('SELECT 1', r => {
			console.log('query finished in', getExecId())
		})
		console.log('finish', getExecId())
	})
}

;(async () => {
	doThings(1)
	await sleep(1000)
	console.log('')
	doThings(2)
})()

Output is

query starts in 1
finish 1
query finished in 1

query starts in 2
finish 2
query finished in 1     <--- should be 2 as well

This shows that the callback for pool.query is in the wrong context.

I think this can be solved in the library by having some class inherit AsyncResource and then calling runInAsyncScope for that callback, see here for an example.

[knilink]

one workaround is

import { promisify, callbackify } from 'node:util';
callbackify(promisify(pool.query)).call(pool, ...);

as mentioned in the doc “Context loss” sesstion:

If your code is callback-based, it is enough to promisify it with util.promisify() so it starts working with native promises.

so promisify seems to be able to recover the context.

here is an pure js example to demonstrate the issue.

import { promisify, callbackify } from 'node:util';
import { AsyncLocalStorage } from 'node:async_hooks';

const asyncLocalStorage = new AsyncLocalStorage();

class MyQueue {
  _queue = [];
  next() {
    setTimeout(() => {
      const cb = this._queue.shift();
      cb();
      if (this._queue.length) {
        this.next();
      }
    });
  }
  exec(cb) {
    this._queue.push(cb);
    if (this._queue.length === 1) {
      this.next();
    }
  }
  patchedExec = callbackify(promisify(this.exec));
}

const myQueue = new MyQueue();

function main() {
  for (let i = 0; i < 5; i++) {
    asyncLocalStorage.run(i, () => {
      myQueue.exec(() => {
        console.log(['exec', i, asyncLocalStorage.getStore()]);
      });
    });
  }
  for (let i = 0; i < 5; i++) {
    asyncLocalStorage.run(i, () => {
      myQueue.patchedExec(() => {
        console.log(['patchedExec', i, asyncLocalStorage.getStore()]);
      });
    });
  }
}

main();

for pg the callbacks were binded to the process which triggered client.connect()

[alxndrsn]

@Badestrand here’s a fix which should make sure all your callbacks execute in the correct context.

For me, this makes my integration tests marginally slower.

// Make sure that DB query callbacks are executed in the correct async context
// see: https://github.com/brianc/node-postgres/issues/2533
const { connect } = pool;

wrapQueryFn(pool);

pool.connect = async cb => {
  if(cb) {
    try {
      const client = await connect.call(pool);
      wrapQueryFn(client);
      return cb(undefined, client, client.release);
    } catch(err) {
      return cb(err);
    }
  } else {
    return connect.call(pool);
  }
};

function wrapQueryFn(db) {
  const { query } = db;

  if(query._async_context_protection) {
    return;
  }

  db.query = async (text, values, cb) => {
    if (typeof values === 'function') {
      cb = values;
      values = undefined;
    }

    if(cb) {
      try {
        return cb(undefined, await query.call(db, text, values));
      } catch(err) {
        return cb(err);
      }
    } else {
      return query.call(db, text, values);
    }
  };

  db.query._async_context_protection = true;
}