feat(NODE-5639): keep executing context on queued commands

src/cmap/connection_pool.ts

@@ -1,3 +1,4 @@
+import { AsyncResource } from 'async_hooks';
 import { clearTimeout, setTimeout } from 'timers';
 
 import type { ObjectId } from '../bson';
@@ -356,7 +357,9 @@ export class ConnectionPool extends TypedEventEmitter<ConnectionPoolEvents> {
       new ConnectionCheckOutStartedEvent(this)
     );
 
-    const waitQueueMember: WaitQueueMember = { callback };
+    const callbackWithContext = AsyncResource.bind(callback);
+
+    const waitQueueMember: WaitQueueMember = { callback: callbackWithContext };
     const waitQueueTimeoutMS = this.options.waitQueueTimeoutMS;
     if (waitQueueTimeoutMS) {
       waitQueueMember.timer = setTimeout(() => {

test/unit/cmap/connection_pool.test.js

@@ -1,5 +1,6 @@
 'use strict';
 
+const { AsyncLocalStorage } = require('async_hooks');
 const { ConnectionPool } = require('../../mongodb');
 const { WaitQueueTimeoutError } = require('../../mongodb');
 const mock = require('../../tools/mongodb-mock/index');
@@ -148,6 +149,34 @@ describe('Connection Pool', function () {
     });
   });
 
+  it('should propagate context on checkout callbacks', function (done) {
+    mockMongod.setMessageHandler(request => {
+      const doc = request.document;
+      if (isHello(doc)) {
+        request.reply(mock.HELLO);
+      }
+    });
+
+    const storage = new AsyncLocalStorage();
+    const context = { id: Math.random() };
+    const pool = new ConnectionPool(stubServer, {
+      maxPoolSize: 1,
+      waitQueueTimeoutMS: 200,
+      hostAddress: mockMongod.hostAddress()
+    });
+
+    pool.ready();
+    storage.enterWith(context);
+
+    pool.checkOut((err, conn) => {
+      expect(err).to.not.exist;
+      expect(conn).to.exist;
+      expect(storage.getStore()).to.eq(context);
+      storage.exit();
+      done();
+    });
+  });
+
   describe('minPoolSize population', function () {
     let clock, timerSandbox;
     beforeEach(() => {