chore(tracing): remove async storage from mongo plugins

packages/datadog-instrumentations/src/mongodb-core.js

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -58,22 +57,6 @@ addHook({ name: 'mongodb-core', versions: ['~3.1.10'], file: 'lib/wireprotocol/2
   return WireProtocol
 })
 
-addHook({ name: 'mongodb', versions: ['>=3.5.4 <4.11.0'], file: 'lib/utils.js' }, util => {
-  shimmer.wrap(util, 'maybePromise', maybePromise => function (parent, callback, fn) {
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    const callbackIndex = arguments.length - 2
-
-    callback = arguments[callbackIndex]
-
-    if (typeof callback === 'function') {
-      arguments[callbackIndex] = asyncResource.bind(callback)
-    }
-
-    return maybePromise.apply(this, arguments)
-  })
-  return util
-})
-
 function wrapWp (wp) {
   shimmer.wrap(wp, 'command', command => wrapUnifiedCommand(command, 'command'))
   shimmer.wrap(wp, 'insert', insert => wrapUnifiedCommand(insert, 'insert', 'insert'))
@@ -148,61 +131,66 @@ function wrapCommand (command, operation, name) {
   return wrapped
 }
 
-function instrument (operation, command, ctx, args, server, ns, ops, options = {}) {
+function instrument (operation, command, mongoCtx, args, server, ns, ops, options = {}) {
   const name = options.name || (ops && Object.keys(ops)[0])
   const index = args.length - 1
-  let callback = args[index]
+  const callback = args[index]
 
-  if (typeof callback !== 'function') return command.apply(ctx, args)
+  if (typeof callback !== 'function') return command.apply(mongoCtx, args)
 
   const serverInfo = server && server.s && server.s.options
-  const callbackResource = new AsyncResource('bound-anonymous-fn')
-  const asyncResource = new AsyncResource('bound-anonymous-fn')
 
-  callback = callbackResource.bind(callback)
-
-  return asyncResource.runInAsyncScope(() => {
-    startCh.publish({ ns, ops, options: serverInfo, name })
-
-    args[index] = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (err, res) {
+  const ctx = {
+    ns,
+    ops,
+    options: serverInfo,
+    name
+  }
+  return startCh.runStores(ctx, () => {
+    args[index] = shimmer.wrapFunction(callback, callback => function (err, res) {
       if (err) {
-        errorCh.publish(err)
+        ctx.error = err
+        errorCh.publish(ctx)
       }
 
-      finishCh.publish()
-
-      if (callback) {
-        return callback.apply(this, arguments)
-      }
-    }))
+      return callback ? finishCh.runStores(ctx, callback, this, ...arguments) : finishCh.publish(ctx)
+    })
 
     try {
-      return command.apply(ctx, args)
+      return command.apply(mongoCtx, args)
     } catch (err) {
-      errorCh.publish(err)
+      ctx.error = err
+      errorCh.publish(ctx)
 
       throw err
     }
   })
 }
 
-function instrumentPromise (operation, command, ctx, args, server, ns, ops, options = {}) {
+function instrumentPromise (operation, command, mongoCtx, args, server, ns, ops, options = {}) {
   const name = options.name || (ops && Object.keys(ops)[0])
 
   const serverInfo = server && server.s && server.s.options
-  const asyncResource = new AsyncResource('bound-anonymous-fn')
 
-  return asyncResource.runInAsyncScope(() => {
-    startCh.publish({ ns, ops, options: serverInfo, name })
+  const ctx = {
+    ns,
+    ops,
+    options: serverInfo,
+    name
+  }
 
-    const promise = command.apply(ctx, args)
+  return startCh.runStores(ctx, () => {
+    const promise = command.apply(mongoCtx, args)
 
     return promise.then(function (res) {
-      finishCh.publish()
-      return res
+      ctx.result = res
+      return finishCh.runStores(ctx, () => {
+        return res
+      })
     }, function (err) {
-      errorCh.publish(err)
-      finishCh.publish()
+      ctx.error = err
+      errorCh.publish(ctx)
+      finishCh.publish(ctx)
 
       throw err
     })

packages/datadog-instrumentations/src/mongodb.js

@@ -4,8 +4,7 @@ require('./mongodb-core')
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -41,19 +40,16 @@ addHook({ name: 'mongodb', versions: ['>=3.3 <5', '5', '>=6'] }, mongodb => {
           return method.apply(this, arguments)
         }
 
-        const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-        return asyncResource.runInAsyncScope(() => {
-          const filters = [arguments[0]]
-          if (useTwoArguments) {
-            filters.push(arguments[1])
-          }
+        const ctx = {
+          filters: [arguments[0]],
+          methodName
+        }
 
-          startCh.publish({
-            filters,
-            methodName
-          })
+        if (useTwoArguments) {
+          ctx.filters.push(arguments[1])
+        }
 
+        return startCh.runStores(ctx, () => {
           return method.apply(this, arguments)
         })
       }

packages/datadog-plugin-mongodb-core/src/index.js

@@ -24,7 +24,9 @@ class MongodbCorePlugin extends DatabasePlugin {
     )
   }
 
-  start ({ ns, ops, options = {}, name }) {
+  bindStart (ctx) {
+    const { ns, ops, options = {}, name } = ctx
+
     // heartbeat commands can be disabled if this.config.heartbeatEnabled is false
     if (!this.config.heartbeatEnabled && isHeartbeat(ops, this.config)) {
       return
@@ -44,11 +46,13 @@ class MongodbCorePlugin extends DatabasePlugin {
         'out.host': options.host,
         'out.port': options.port
       }
-    })
+    }, ctx)
     const comment = this.injectDbmComment(span, ops.comment, service)
     if (comment) {
       ops.comment = comment
     }
+
+    return ctx.currentStore
   }
 
   getPeerService (tags) {