fix: mongodb async context fix

.tav.yml

@@ -143,22 +143,22 @@ mongodb-core:
 mongodb-3:
   name: mongodb
   versions: '>=3.3 <4'
-  commands: node test/instrumentation/modules/mongodb.test.js
+  commands: node test/instrumentation/modules/mongodb/mongodb.test.js
 mongodb-4:
   name: mongodb
   versions: '>=4 <5'
   node: '>=12.9.0'
-  commands: node test/instrumentation/modules/mongodb.test.js
+  commands: node test/instrumentation/modules/mongodb/mongodb.test.js
 mongodb-5:
   name: mongodb
   versions: '>=5 <6'
   node: '>=14.20.1'
-  commands: node test/instrumentation/modules/mongodb.test.js
+  commands: node test/instrumentation/modules/mongodb/mongodb.test.js
 mongodb:
   name: mongodb
   versions: '>=6 <7'
   node: '>=15.0.0'
-  commands: node test/instrumentation/modules/mongodb.test.js
+  commands: node test/instrumentation/modules/mongodb/mongodb.test.js
 
 # Bluebird is effectively deprecated (https://github.com/petkaantonov/bluebird#%EF%B8%8Fnote%EF%B8%8F).
 # Testing the full set of supported bluebird releases (`>=2 <4`) is currently

lib/instrumentation/index.js

@@ -43,10 +43,11 @@ var MODULES = [
   ['@elastic/elasticsearch', '@elastic/elasticsearch-canary'],
   '@opentelemetry/api',
   '@opentelemetry/sdk-metrics',
-  '@node-redis/client/dist/lib/client',
-  '@node-redis/client/dist/lib/client/commands-queue',
-  '@redis/client/dist/lib/client',
-  '@redis/client/dist/lib/client/commands-queue',
+  // XXX: check what is the path for theses ones
+  '@node-redis/client/dist/lib/client/index.js',
+  '@node-redis/client/dist/lib/client/commands-queue.js',
+  '@redis/client/dist/lib/client/index.js',
+  '@redis/client/dist/lib/client/commands-queue.js',
   'apollo-server-core',
   'aws-sdk',
   'bluebird',
@@ -73,6 +74,7 @@ var MODULES = [
   'mimic-response',
   'mongodb-core',
   'mongodb',
+  'mongodb/lib/cmap/connection_pool.js',
   'mysql',
   'mysql2',
   'next/dist/server/api-utils/node',
@@ -134,6 +136,12 @@ function Instrumentation(agent) {
   // because the array is given to the hook(...) call.
   this._patches = new NamedArray();
 
+  // NOTE: since RITM is using internals it will ask for
+  // the version for each sub-module. As a consequence agent will
+  // do many unnecesary reads to get the package version.
+  // A cache will avoid these reads
+  this._versionsCache = new Map();
+
   for (let modName of MODULES) {
     if (!Array.isArray(modName)) modName = [modName];
     const pathName = modName[0];
@@ -142,7 +150,9 @@ function Instrumentation(agent) {
       // Lazy require so that we don't have to use `require.resolve` which
       // would fail in combination with Webpack. For more info see:
       // https://github.com/elastic/apm-agent-nodejs/pull/957
-      return require(`./modules/${pathName}.js`)(...args);
+      const modExt = pathName.endsWith('.js') ? '' : '.js';
+      const modPath = `./modules/${pathName}${modExt}`;
+      return require(modPath)(...args);
     });
   }
 
@@ -358,35 +368,37 @@ Instrumentation.prototype._startHook = function () {
 
   this._agent.logger.debug('adding hooks to Node.js module loader');
 
-  this._ritmHook = new RitmHook(this._patches.keys, function (
-    exports,
-    name,
-    basedir,
-  ) {
-    const enabled = self._isModuleEnabled(name);
-    var pkg, version;
-
-    const isHandlingLambda =
-      self._lambdaHandlerInfo && self._lambdaHandlerInfo.module === name;
-
-    if (!isHandlingLambda && basedir) {
-      pkg = path.join(basedir, 'package.json');
-      try {
-        version = JSON.parse(fs.readFileSync(pkg)).version;
-      } catch (e) {
-        self._agent.logger.debug(
-          'could not shim %s module: %s',
-          name,
-          e.message,
-        );
-        return exports;
+  this._ritmHook = new RitmHook(
+    this._patches.keys,
+    // Need to enable internals so we hook into sub-modules even
+    // if the main module is hooked
+    // see: https://github.com/elastic/apm-agent-nodejs/pull/3657/files#r1344193251
+    { internals: true },
+    function (exports, name, basedir) {
+      const enabled = self._isModuleEnabled(name);
+      var version;
+
+      const isHandlingLambda =
+        self._lambdaHandlerInfo && self._lambdaHandlerInfo.module === name;
+
+      if (!isHandlingLambda && basedir) {
+        try {
+          version = self._getPackageVersion(name, basedir);
+        } catch (e) {
+          self._agent.logger.debug(
+            'could not shim %s module: %s',
+            name,
+            e.message,
+          );
+          return exports;
+        }
+      } else {
+        version = process.versions.node;
       }
-    } else {
-      version = process.versions.node;
-    }
 
-    return self._patchModule(exports, name, version, enabled, false);
-  });
+      return self._patchModule(exports, name, version, enabled, false);
+    },
+  );
 
   this._iitmHook = IitmHook(
     Object.keys(IITM_MODULES),
@@ -412,22 +424,26 @@ Instrumentation.prototype._startHook = function () {
 };
 
 Instrumentation.prototype._getPackageVersion = function (modName, modBaseDir) {
+  const modPath = path.join(modBaseDir, 'package.json');
+  if (this._versionsCache.has(modPath)) {
+    return this._versionsCache.get(modPath);
+  }
+
   try {
-    const version = JSON.parse(
-      fs.readFileSync(path.join(modBaseDir, 'package.json')),
-    ).version;
+    const version = JSON.parse(fs.readFileSync(modPath)).version;
     if (typeof version === 'string') {
+      this._versionsCache.set(modPath, version);
       return version;
-    } else {
-      return undefined;
     }
   } catch (err) {
     this._agent.logger.debug(
       { modName, modBaseDir, err },
       'could not load package version',
     );
-    return undefined;
   }
+
+  this._versionsCache.set(modPath, undefined);
+  return undefined;
 };
 
 /**

lib/instrumentation/modules/mongodb/lib/cmap/connection_pool.js

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and other contributors where applicable.
+ * Licensed under the BSD 2-Clause License; you may not use this file except in
+ * compliance with the BSD 2-Clause License.
+ */
+
+'use strict';
+
+const { AsyncResource } = require('async_hooks');
+
+const semver = require('semver');
+
+module.exports = (mod, agent, { version, enabled }) => {
+  if (!enabled) return mod;
+  if (!semver.satisfies(version, '>=3.3 <7.0')) {
+    agent.logger.debug(
+      'mongodb version %s not instrumented (mongodb <3.3 is instrumented via mongodb-core)',
+      version,
+    );
+    return mod;
+  }
+
+  if (mod.ConnectionPool) {
+    class ConnectionPoolTraced extends mod.ConnectionPool {
+      checkOut(callback) {
+        return super.checkOut(AsyncResource.bind(callback));
+      }
+    }
+
+    Object.defineProperty(mod, 'ConnectionPool', {
+      enumerable: true,
+      get: function () {
+        return ConnectionPoolTraced;
+      },
+    });
+
+    return mod;
+  }
+};

test/_is_mongodb_incompat.js

@@ -7,6 +7,7 @@
 'use strict';
 
 var semver = require('semver');
+const { safeGetPackageVersion } = require('./_utils');
 
 /**
  * Return whether the current 'mongodb' version is incompatible with the
@@ -26,7 +27,7 @@ var semver = require('semver');
  */
 function isMongodbIncompat() {
   const nodeVer = process.version;
-  const mongodbVer = require('mongodb/package.json').version;
+  const mongodbVer = safeGetPackageVersion('mongodb');
   const msg = `mongodb@${mongodbVer} is incompatible with node@${nodeVer}`;
 
   if (semver.satisfies(mongodbVer, '4.x')) {