remove async storage from elasticsearch instrumentation (#5935)

* remove async storage from elasticsearch

Author: wconti27

packages/datadog-instrumentations/src/elasticsearch.js

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -73,33 +72,30 @@ function createWrapRequest (name) {
 
       if (!params) return request.apply(this, arguments)
 
-      const parentResource = new AsyncResource('bound-anonymous-fn')
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-      return asyncResource.runInAsyncScope(() => {
-        startCh.publish({ params })
-
+      const ctx = { params }
+      return startCh.runStores(ctx, () => {
         try {
           const lastIndex = arguments.length - 1
           cb = arguments[lastIndex]
 
           if (typeof cb === 'function') {
-            cb = parentResource.bind(cb)
-
-            arguments[lastIndex] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error) {
-              finish(params, error)
-              return cb.apply(null, arguments)
-            }))
+            arguments[lastIndex] = shimmer.wrapFunction(cb, cb => function (error) {
+              if (error) {
+                ctx.error = error
+                errorCh.publish(ctx)
+              }
+              return finishCh.runStores(ctx, cb, null, ...arguments)
+            })
             return request.apply(this, arguments)
           }
           const promise = request.apply(this, arguments)
           if (promise && typeof promise.then === 'function') {
-            const onResolve = asyncResource.bind(() => finish(params))
-            const onReject = asyncResource.bind(e => finish(params, e))
+            const onResolve = () => finish(ctx)
+            const onReject = e => finish(ctx, e)
 
             promise.then(onResolve, onReject)
           } else {
-            finish(params)
+            finish(ctx)
           }
           return promise
         } catch (err) {
@@ -112,11 +108,12 @@ function createWrapRequest (name) {
     }
   }
 
-  function finish (params, error) {
+  function finish (ctx, error) {
     if (error) {
+      ctx.error = error
       errorCh.publish(error)
     }
-    finishCh.publish({ params })
+    return finishCh.runStores(ctx, () => {})
   }
 }
 

packages/datadog-instrumentations/src/sharedb.js

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -37,30 +36,25 @@ addHook({ name: 'sharedb', versions: ['>=1'], file: 'lib/agent.js' }, Agent => {
   const errorCh = channel('apm:sharedb:request:error')
 
   shimmer.wrap(Agent.prototype, '_handleMessage', origHandleMessageFn => function (request, callback) {
-    const callbackResource = new AsyncResource('bound-anonymous-fn')
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
     const action = request.a
     const actionName = getReadableActionName(action)
 
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish({ actionName, request })
-
-      callback = callbackResource.bind(callback)
-
-      arguments[1] = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (error, res) {
+    const ctx = { actionName, request }
+    return startCh.runStores(ctx, () => {
+      arguments[1] = shimmer.wrapFunction(callback, callback => function (error, res) {
         if (error) {
+          ctx.error = error
           errorCh.publish(error)
         }
-        finishCh.publish({ request, res })
-
-        return callback.apply(this, arguments)
-      }))
+        ctx.res = res
+        return finishCh.runStores(ctx, callback, this, ...arguments)
+      })
 
       try {
         return origHandleMessageFn.apply(this, arguments)
       } catch (error) {
-        errorCh.publish(error)
+        ctx.error = error
+        errorCh.publish(ctx)
 
         throw error
       }

packages/datadog-plugin-elasticsearch/src/index.js

@@ -5,7 +5,9 @@ const DatabasePlugin = require('../../dd-trace/src/plugins/database')
 class ElasticsearchPlugin extends DatabasePlugin {
   static get id () { return 'elasticsearch' }
 
-  start ({ params }) {
+  bindStart (ctx) {
+    const { params } = ctx
+
     const body = getBody(params.body || params.bulkBody)
 
     this.startSpan(this.operationName(), {
@@ -20,13 +22,19 @@ class ElasticsearchPlugin extends DatabasePlugin {
         [`${this.system}.body`]: body,
         [`${this.system}.params`]: JSON.stringify(params.querystring || params.query)
       }
-    })
+    }, ctx)
+
+    return ctx.currentStore
   }
 
-  finish ({ params }) {
+  bindFinish (ctx) {
+    const { params } = ctx
+
     const span = this.activeSpan
     this.config.hooks.query(span, params)
-    super.finish({ params })
+    super.finish(ctx)
+
+    return ctx.parentStore
   }
 }
 

packages/datadog-plugin-sharedb/src/index.js

@@ -5,27 +5,35 @@ const ServerPlugin = require('../../dd-trace/src/plugins/server')
 class SharedbPlugin extends ServerPlugin {
   static get id () { return 'sharedb' }
 
-  start ({ actionName, request }) {
+  bindStart (ctx) {
+    const { actionName, request } = ctx
+
     const span = this.startSpan('sharedb.request', {
       service: this.config.service,
       resource: getReadableResourceName(actionName, request.c, request.q),
       kind: 'server',
       meta: {
         'sharedb.action': actionName
       }
-    })
+    }, ctx)
 
     if (this.config.hooks && this.config.hooks.receive) {
       this.config.hooks.receive(span, request)
     }
+
+    return ctx.currentStore
   }
 
-  finish ({ request, res }) {
-    const span = this.activeSpan
+  bindFinish (ctx) {
+    const { request, res } = ctx
+
+    const span = ctx.currentStore.span
     if (this.config.hooks && this.config.hooks.reply) {
       this.config.hooks.reply(span, request, res)
     }
-    super.finish()
+    super.finish(ctx)
+
+    return ctx.parentStore
   }
 }