fix: ensure correct run context for 'mongodb' instrumentation

.ci/.jenkins_tav.yml

@@ -23,6 +23,7 @@ TAV:
   - koa-router
   - memcached
   - mimic-response
+  - mongodb
   - mongodb-core
   - mysql
   - mysql2

.ci/scripts/test.sh

@@ -195,7 +195,7 @@ elif [[ -n "${TAV_MODULE}" ]]; then
     tedious)
       DOCKER_COMPOSE_FILE=docker-compose-mssql.yml
       ;;
-    mongodb-core)
+    mongodb|mongodb-core)
       DOCKER_COMPOSE_FILE=docker-compose-mongodb.yml
       ;;
     pg|knex)

.tav.yml

@@ -98,9 +98,17 @@ pg-new-node:
 mongodb-core:
   versions: '>=1.2.19 <4'
   commands: node test/instrumentation/modules/mongodb-core.test.js
-mongodb:
-  versions: '>=3.3'
+
+mongodb-3:
+  name: mongodb
+  versions: '>=3.3 <4'
+  commands: node test/instrumentation/modules/mongodb.test.js
+mongodb-4:
+  name: mongodb
+  versions: '>=4 <5'
+  node: '>=12'
   commands: node test/instrumentation/modules/mongodb.test.js
+
 bluebird:
   versions: '>=2 <4'
   commands:

CHANGELOG.asciidoc

@@ -52,6 +52,8 @@ Notes:
 
 * Fixes for run context handling for 'pg' instrumentation. ({issues}2430[#2430])
 
+* Fixes for run context handling for 'mongodb' instrumentation. ({issues}2512[#2512])
+
 
 [[release-notes-3.26.0]]
 ==== 3.26.0 2021/12/07

examples/trace-mongodb.js

@@ -0,0 +1,80 @@
+// A small example showing Elastic APM tracing the 'mongodb' package.
+//
+// This assumes a MongoDB server running on localhost. You can use:
+//    npm run docker:start mongodb
+// to start a MongoDB docker container. Then `npm run docker:stop` to stop it.
+
+const apm = require('../').start({ // elastic-apm-node
+  serviceName: 'example-trace-mongodb',
+  logUncaughtExceptions: true
+})
+
+const MongoClient = require('mongodb').MongoClient
+
+const DB_NAME = 'example-trace-mongodb'
+const url = 'mongodb://localhost:27017'
+
+async function usingPromises () {
+  // For tracing spans to be created, there must be an active transaction.
+  // Typically, a transaction is automatically started for incoming HTTP
+  // requests to a Node.js server. However, because this script is not running
+  // an HTTP server, we manually start a transaction. More details at:
+  // https://www.elastic.co/guide/en/apm/agent/nodejs/current/custom-transactions.html
+  const t1 = apm.startTransaction('t1')
+
+  const client = new MongoClient(url)
+  try {
+    await client.connect()
+
+    const database = client.db(DB_NAME)
+    const coll = database.collection('breakfast')
+
+    let res = await coll.insertMany([
+      { item: 'spam', n: 0 },
+      { item: 'ham', n: 1 },
+      { item: 'eggs', n: 2 }
+    ])
+    console.log('insertMany:', res)
+
+    res = await coll.findOne({ item: 'eggs' })
+    console.log('findOne eggs:', res)
+
+    coll.findOne({ item: 'ham' }, function (err, res) {
+      console.log('findOne ham: err=%s res=%s', err && err.message, res)
+    })
+
+    await coll.deleteMany({})
+
+    res = await coll.findOne({ item: 'eggs' })
+    console.log('findOne eggs:', res)
+  } finally {
+    await client.close()
+    t1.end()
+  }
+}
+
+function usingCallbacks () {
+  const t2 = apm.startTransaction('t2-callback-style')
+
+  MongoClient.connect(url, function (err, client) {
+    console.log('connect: err=%s', err && err.message)
+    if (err) {
+      throw err
+    }
+
+    const db = client.db(DB_NAME)
+    const coll = db.collection('breakfast')
+    coll.insertMany([
+      { item: 'spam', n: 0 },
+      { item: 'ham', n: 1 },
+      { item: 'eggs', n: 2 }
+    ], { w: 1 }, function (err, res) {
+      console.log('insertMany: err=%s res=%s', err && err.message, res)
+      t2.end()
+      client.close()
+    })
+  })
+}
+
+usingPromises()
+  .finally(usingCallbacks)

lib/instrumentation/modules/mongodb.js

@@ -1,7 +1,9 @@
 'use strict'
 
 const semver = require('semver')
+
 const { getDBDestination } = require('../context')
+const shimmer = require('../shimmer')
 
 // Match expected `<hostname>:<port>`, e.g. "mongo:27017", "::1:27017",
 // "127.0.0.1:27017".
@@ -14,25 +16,30 @@ module.exports = (mongodb, agent, { version, enabled }) => {
     return mongodb
   }
 
+  const ins = agent._instrumentation
+
   const activeSpans = new Map()
   if (mongodb.instrument) {
     const listener = mongodb.instrument()
-
     listener.on('started', onStart)
     listener.on('succeeded', onEnd)
     listener.on('failed', onEnd)
   } else if (mongodb.MongoClient) {
     // mongodb 4.0+ removed the instrument() method in favor of
-    // listeners on the instantiated client objects.
-    class MongoClient extends mongodb.MongoClient {
+    // listeners on the instantiated client objects. There are two mechanisms
+    // to get a client:
+    // 1. const client = new mongodb.MongoClient(...)
+    // 2. const client = await MongoClient.connect(...)
+    class MongoClientTraced extends mongodb.MongoClient {
       constructor () {
-        // The `command*` events are only sent if `monitorCommands: true`.
-        if (!arguments[1]) {
-          arguments[1] = { monitorCommands: true }
-        } else if (arguments[1].monitorCommands !== true) {
-          arguments[1] = Object.assign({}, arguments[1], { monitorCommands: true })
+        // The `command*` events are only emitted if `options.monitorCommands: true`.
+        const args = Array.prototype.slice.call(arguments)
+        if (!args[1]) {
+          args[1] = { monitorCommands: true }
+        } else if (args[1].monitorCommands !== true) {
+          args[1] = Object.assign({}, args[1], { monitorCommands: true })
         }
-        super(...arguments)
+        super(...args)
         this.on('commandStarted', onStart)
         this.on('commandSucceeded', onEnd)
         this.on('commandFailed', onEnd)
@@ -44,13 +51,53 @@ module.exports = (mongodb, agent, { version, enabled }) => {
       {
         enumerable: true,
         get: function () {
-          return MongoClient
+          return MongoClientTraced
         }
       }
     )
+    shimmer.wrap(mongodb.MongoClient, 'connect', wrapConnect)
+  } else {
+    agent.logger.warn('could not instrument mongodb@%s', version)
   }
   return mongodb
 
+  // Wrap the MongoClient.connect(url, options?, callback?) static method.
+  // It calls back with `function (err, client)` or returns a Promise that
+  // resolves to the client.
+  // https://github.com/mongodb/node-mongodb-native/blob/v4.2.1/src/mongo_client.ts#L503-L511
+  function wrapConnect (origConnect) {
+    return function wrappedConnect (url, options, callback) {
+      if (typeof options === 'function') {
+        callback = options
+        options = {}
+      }
+      options = options || {}
+      if (!options.monitorCommands) {
+        options.monitorCommands = true
+      }
+      if (typeof callback === 'function') {
+        return origConnect.call(this, url, options, function wrappedCallback (err, client) {
+          if (err) {
+            callback(err)
+          } else {
+            client.on('commandStarted', onStart)
+            client.on('commandSucceeded', onEnd)
+            client.on('commandFailed', onEnd)
+            callback(err, client)
+          }
+        })
+      } else {
+        const p = origConnect.call(this, url, options, callback)
+        p.then(client => {
+          client.on('commandStarted', onStart)
+          client.on('commandSucceeded', onEnd)
+          client.on('commandFailed', onEnd)
+        })
+        return p
+      }
+    }
+  }
+
   function onStart (event) {
     // `event` is a `CommandStartedEvent`
     // https://github.com/mongodb/specifications/blob/master/source/command-monitoring/command-monitoring.rst#api
@@ -70,7 +117,7 @@ module.exports = (mongodb, agent, { version, enabled }) => {
       event.commandName
     ].join('.')
 
-    const span = agent.startSpan(name, 'db', 'mongodb', 'query')
+    const span = ins.createSpan(name, 'db', 'mongodb', 'query')
     if (span) {
       activeSpans.set(event.requestId, span)
 

package.json

@@ -163,7 +163,7 @@
     "mimic-response": "^2.1.0",
     "mkdirp": "^0.5.1",
     "module-details-from-path": "^1.0.3",
-    "mongodb": "^4.1.0",
+    "mongodb": "^4.2.1",
     "mongodb-core": "^3.2.7",
     "mysql": "^2.18.1",
     "mysql2": "^2.1.0",