remove async storage from moleculer (#5956)

* remove async storage from moleculer

Author: wconti27

packages/datadog-instrumentations/src/moleculer/client.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { channel, addHook, AsyncResource } = require('../helpers/instrument')
+const { channel, addHook } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 
 const startChannel = channel('apm:moleculer:call:start')
@@ -9,29 +9,28 @@ const errorChannel = channel('apm:moleculer:call:error')
 
 function wrapCall (call) {
   return function (actionName, params, opts) {
-    const callResource = new AsyncResource('bound-anonymous-fn')
-
     opts = arguments[2] = opts || {}
     opts.meta = opts.meta || {}
 
     arguments.length = Math.max(3, arguments.length)
 
-    return callResource.runInAsyncScope(() => {
-      startChannel.publish({ actionName, params, opts })
-
+    const ctx = { actionName, params, opts }
+    return startChannel.runStores(ctx, () => {
       const promise = call.apply(this, arguments)
       const broker = this
-      const ctx = promise.ctx
+      ctx.promiseCtx = promise.ctx
+      ctx.broker = broker
 
       return promise
         .then(
           result => {
-            finishChannel.publish({ broker, ctx })
+            finishChannel.publish(ctx)
             return result
           },
           error => {
-            errorChannel.publish(error)
-            finishChannel.publish({ broker, ctx })
+            ctx.error = error
+            errorChannel.publish(ctx)
+            finishChannel.publish(ctx)
             throw error
           }
         )

packages/datadog-instrumentations/src/moleculer/server.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { channel, addHook, AsyncResource } = require('../helpers/instrument')
+const { channel, addHook } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 
 const startChannel = channel('apm:moleculer:action:start')
@@ -24,27 +24,26 @@ function createMiddleware () {
     localAction (next, action) {
       const broker = this
 
-      return shimmer.wrapFunction(next, next => function datadogMiddleware (ctx) {
-        const actionResource = new AsyncResource('bound-anonymous-fn')
-
-        return actionResource.runInAsyncScope(() => {
-          startChannel.publish({ action, ctx, broker })
-
+      return shimmer.wrapFunction(next, next => function datadogMiddleware (middlewareCtx) {
+        const ctx = { action, middlewareCtx, broker }
+        return startChannel.runStores(ctx, () => {
           try {
-            return next(ctx).then(
+            return next(middlewareCtx).then(
               result => {
-                finishChannel.publish()
+                finishChannel.publish(ctx)
                 return result
               },
               error => {
-                errorChannel.publish(error)
-                finishChannel.publish()
+                ctx.error = error
+                errorChannel.publish(ctx)
+                finishChannel.publish(ctx)
                 throw error
               }
             )
           } catch (e) {
-            errorChannel.publish(e)
-            finishChannel.publish()
+            ctx.error = e
+            errorChannel.publish(ctx)
+            finishChannel.publish(ctx)
           }
         })
       })

packages/datadog-plugin-moleculer/src/client.js

@@ -7,29 +7,35 @@ class MoleculerClientPlugin extends ClientPlugin {
   static get id () { return 'moleculer' }
   static get operation () { return 'call' }
 
-  start ({ actionName, opts }) {
+  bindStart (ctx) {
+    const { actionName, opts } = ctx
+
     const span = this.startSpan(this.operationName(), {
       service: this.config.service || this.serviceName(),
       resource: actionName,
       kind: 'client'
-    })
+    }, ctx)
 
     this.tracer.inject(span, 'text_map', opts.meta)
+
+    return ctx.currentStore
   }
 
-  finish ({ broker, ctx }) {
-    const span = this.activeSpan
+  finish (ctx) {
+    const { promiseCtx, broker } = ctx
+
+    const span = ctx.currentStore.span || this.activeSpan
 
-    if (ctx) {
-      const endpoint = ctx.endpoint || {}
+    if (promiseCtx) {
+      const endpoint = promiseCtx.endpoint || {}
       const node = endpoint.node || {}
 
       this.addHost(node.hostname, node.port)
 
-      span.addTags(moleculerTags(broker, ctx, this.config))
+      span.addTags(moleculerTags(broker, promiseCtx, this.config))
     }
 
-    super.finish()
+    super.finish(ctx)
   }
 }
 

packages/datadog-plugin-moleculer/src/server.js

@@ -7,19 +7,23 @@ class MoleculerServerPlugin extends ServerPlugin {
   static get id () { return 'moleculer' }
   static get operation () { return 'action' }
 
-  start ({ action, ctx, broker }) {
-    const followsFrom = this.tracer.extract('text_map', ctx.meta)
+  bindStart (ctx) {
+    const { action, middlewareCtx, broker } = ctx
+
+    const followsFrom = this.tracer.extract('text_map', middlewareCtx.meta)
     this.startSpan(this.operationName(), {
-      childOf: followsFrom || this.activeSpan,
+      childOf: followsFrom || ctx?.currentStore?.span || this.activeSpan,
       service: this.config.service || this.serviceName(),
       resource: action.name,
       kind: 'server',
       type: 'web',
       meta: {
         'resource.name': action.name,
-        ...moleculerTags(broker, ctx, this.config)
+        ...moleculerTags(broker, middlewareCtx, this.config)
       }
-    })
+    }, ctx)
+
+    return ctx.currentStore
   }
 }