remove async storage from http2 server instrumentation

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

@@ -5,14 +5,15 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 
 const startServerCh = channel('apm:http2:server:request:start')
 const errorServerCh = channel('apm:http2:server:request:error')
 const finishServerCh = channel('apm:http2:server:request:finish')
+// this channel is for wrapping the response emit method and handling store context, it doesn't have any subscribers
+const responseCh = channel('apm:http2:server:response:emit')
 
 const names = ['http2', 'node:http2']
 
@@ -30,16 +31,11 @@ function wrapCreateServer (createServer) {
   }
 }
 
-function wrapResponseEmit (emit) {
-  const asyncResource = new AsyncResource('bound-anonymous-fn')
+function wrapResponseEmit (emit, ctx) {
   return function (eventName, event) {
-    return asyncResource.runInAsyncScope(() => {
-      if (eventName === 'close' && finishServerCh.hasSubscribers) {
-        finishServerCh.publish({ req: this.req })
-      }
-
-      return emit.apply(this, arguments)
-    })
+    ctx.req = this.req
+    ctx.eventName = eventName
+    return finishServerCh.runStores(ctx, emit, this, ...arguments)
   }
 }
 function wrapEmit (emit) {
@@ -51,18 +47,19 @@ function wrapEmit (emit) {
     if (eventName === 'request') {
       res.req = req
 
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-      return asyncResource.runInAsyncScope(() => {
-        startServerCh.publish({ req, res })
-
-        shimmer.wrap(res, 'emit', wrapResponseEmit)
+      const ctx = { req, res }
+      return startServerCh.runStores(ctx, () => {
+        shimmer.wrap(res, 'emit', emit => responseCh.runStores(ctx, () => {
+          return wrapResponseEmit(emit, ctx)
+        }))
 
         try {
           return emit.apply(this, arguments)
-        } catch (err) {
-          errorServerCh.publish(err)
+        } catch (error) {
+          ctx.error = error
+          errorServerCh.publish(ctx)
 
-          throw err
+          throw error
         }
       })
     }

packages/datadog-plugin-child_process/test/index.spec.js

@@ -52,6 +52,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',
@@ -75,6 +76,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',
@@ -100,6 +102,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',
@@ -126,6 +129,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',
@@ -153,6 +157,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',
@@ -180,6 +185,7 @@ describe('Child process plugin', () => {
         expect(tracerStub.startSpan).to.have.been.calledOnceWithExactly(
           'command_execution',
           {
+            startTime: undefined,
             childOf: undefined,
             tags: {
               component: 'subprocess',

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

@@ -11,16 +11,16 @@ class HttpServerPlugin extends ServerPlugin {
     return 'http'
   }
 
+  static get prefix () {
+    return 'apm:http:server:request'
+  }
+
   constructor (...args) {
     super(...args)
     this._parentStore = undefined
     this.addTraceSub('exit', message => this.exit(message))
   }
 
-  addTraceSub (eventName, handler) {
-    this.addSub(`apm:${this.constructor.id}:server:${this.operation}:${eventName}`, handler)
-  }
-
   start ({ req, res, abortController }) {
     const store = storage('legacy').getStore()
     const span = web.startSpan(

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

@@ -3,7 +3,6 @@
 // Plugin temporarily disabled. See https://github.com/DataDog/dd-trace-js/issues/312
 
 const ServerPlugin = require('../../dd-trace/src/plugins/server')
-const { storage } = require('../../datadog-core')
 const web = require('../../dd-trace/src/plugins/util/web')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 
@@ -12,12 +11,13 @@ class Http2ServerPlugin extends ServerPlugin {
     return 'http2'
   }
 
-  addTraceSub (eventName, handler) {
-    this.addSub(`apm:${this.constructor.id}:server:${this.operation}:${eventName}`, handler)
+  static get prefix () {
+    return 'apm:http2:server:request'
   }
 
-  start ({ req, res }) {
-    const store = storage('legacy').getStore()
+  bindStart (ctx) {
+    const { req, res } = ctx
+
     const span = web.startSpan(
       this.tracer,
       {
@@ -26,28 +26,43 @@ class Http2ServerPlugin extends ServerPlugin {
       },
       req,
       res,
-      this.operationName()
+      this.operationName(),
+      ctx
     )
 
     span.setTag(COMPONENT, this.constructor.id)
     span._integrationName = this.constructor.id
 
-    this.enter(span, { ...store, req, res })
+    ctx.currentStore.req = req
+    ctx.currentStore.res = res
 
     const context = web.getContext(req)
 
     if (!context.instrumented) {
       context.res.writeHead = web.wrapWriteHead(context)
       context.instrumented = true
     }
+
+    return ctx.currentStore
   }
 
-  finish ({ req }) {
+  bindFinish (ctx) {
+    if (ctx.eventName !== 'close') return ctx.currentStore
+
+    const { req } = ctx
+
     const context = web.getContext(req)
 
     if (!context || !context.res) return // Not created by a http.Server instance.
 
     web.finishAll(context)
+
+    return ctx.currentStore
+  }
+
+  finish (ctx) {
+    // we let bindFinish handle the finish, but keep this method because we don't want to finish the span
+    // early for a response event that is not a 'close' event, which prevents tags from being set during web.finishAll
   }
 
   error (error) {

packages/dd-trace/src/plugins/plugin.js

@@ -60,7 +60,7 @@ module.exports = class Plugin {
   }
 
   get tracer () {
-    return this._tracer._tracer
+    return this._tracer?._tracer || this._tracer
   }
 
   enter (span, store) {

packages/dd-trace/src/plugins/tracing.js

@@ -102,25 +102,28 @@ class TracingPlugin extends Plugin {
     }
   }
 
-  startSpan (name, { childOf, kind, meta, metrics, service, resource, type } = {}, enterOrCtx = true) {
+  startSpan (name, options = {}, enterOrCtx = true) {
     const store = storage('legacy').getStore()
-    if (store && childOf === undefined) {
-      childOf = store.span
+    if (store && options.childOf === undefined) {
+      options.childOf = store.span
     }
 
-    const span = this.tracer.startSpan(name, {
-      childOf,
+    const tracer = options.tracer || this.tracer
+
+    const span = tracer.startSpan(name, {
+      startTime: options.startTime,
+      childOf: options.childOf,
       tags: {
-        [COMPONENT]: this.component,
-        'service.name': service || this.tracer._service,
-        'resource.name': resource,
-        'span.kind': kind,
-        'span.type': type,
-        ...meta,
-        ...metrics
+        [COMPONENT]: options.component || this.component,
+        'service.name': options.service || options.meta?.service || tracer._service,
+        'resource.name': options.resource,
+        'span.kind': options.kind,
+        'span.type': options.type,
+        ...options.meta,
+        ...options.metrics
       },
-      integrationName: this.component,
-      links: childOf?._links
+      integrationName: options.integrationName || options.component || this.component,
+      links: options.childOf?._links
     })
 
     analyticsSampler.sample(span, this.config.measured)

packages/dd-trace/src/plugins/util/inferred_proxy.js

@@ -22,7 +22,7 @@ const supportedProxies = {
   }
 }
 
-function createInferredProxySpan (headers, childOf, tracer, context) {
+function createInferredProxySpan (headers, childOf, tracer, reqCtx, traceCtx, startSpanHelper) {
   if (!headers) {
     return null
   }
@@ -41,26 +41,22 @@ function createInferredProxySpan (headers, childOf, tracer, context) {
 
   log.debug('Successfully extracted inferred span info %s for proxy:', proxyContext, proxyContext.proxySystemName)
 
-  const span = tracer.startSpan(
-    proxySpanInfo.spanName,
-    {
-      childOf,
-      type: 'web',
-      startTime: proxyContext.requestTime,
-      integrationName: proxySpanInfo.component,
-      tags: {
-        service: proxyContext.domainName || tracer._config.service,
-        component: proxySpanInfo.component,
-        [SPAN_TYPE]: 'web',
-        [HTTP_METHOD]: proxyContext.method,
-        [HTTP_URL]: proxyContext.domainName + proxyContext.path,
-        stage: proxyContext.stage
-      }
+  const span = startSpanHelper(tracer, proxySpanInfo.spanName, {
+    childOf,
+    type: 'web',
+    startTime: proxyContext.requestTime,
+    integrationName: proxySpanInfo.component,
+    meta: {
+      service: proxyContext.domainName || tracer._config.service,
+      component: proxySpanInfo.component,
+      [SPAN_TYPE]: 'web',
+      [HTTP_METHOD]: proxyContext.method,
+      [HTTP_URL]: proxyContext.domainName + proxyContext.path,
+      stage: proxyContext.stage
     }
-  )
+  }, traceCtx)
 
-  tracer.scope().activate(span)
-  context.inferredProxySpan = span
+  reqCtx.inferredProxySpan = span
   childOf = span
 
   log.debug('Successfully created inferred proxy span.')

packages/dd-trace/src/plugins/util/web.js

@@ -37,6 +37,17 @@ const HTTP2_HEADER_PATH = ':path'
 const contexts = new WeakMap()
 const ends = new WeakMap()
 
+const TracingPlugin = require('../tracing')
+
+function startSpanHelper (tracer, name, options, traceCtx, config = {}) {
+  return TracingPlugin.prototype.startSpan.call(
+    { component: 'web', config },
+    name,
+    { ...options, tracer },
+    traceCtx
+  )
+}
+
 const web = {
   TYPE: WEB,
 
@@ -93,7 +104,7 @@ const web = {
     analyticsSampler.sample(span, config.measured, true)
   },
 
-  startSpan (tracer, config, req, res, name) {
+  startSpan (tracer, config, req, res, name, traceCtx) {
     const context = this.patch(req)
 
     let span
@@ -102,7 +113,7 @@ const web = {
       context.span.context()._name = name
       span = context.span
     } else {
-      span = web.startChildSpan(tracer, name, req)
+      span = web.startChildSpan(tracer, name, req, traceCtx)
     }
 
     context.tracer = tracer
@@ -163,10 +174,11 @@ const web = {
     const tracer = context.tracer
     const childOf = this.active(req)
     const config = context.config
+    const traceCtx = context.traceCtx
 
     if (config.middleware === false) return this.bindAndWrapMiddlewareErrors(fn, req, tracer, childOf)
 
-    const span = tracer.startSpan(name, { childOf })
+    const span = startSpanHelper(tracer, name, { childOf }, traceCtx, config)
 
     analyticsSampler.sample(span, config.measured)
 
@@ -258,20 +270,20 @@ const web = {
   },
 
   // Extract the parent span from the headers and start a new span as its child
-  startChildSpan (tracer, name, req) {
+  startChildSpan (tracer, name, req, traceCtx) {
     const headers = req.headers
-    const context = contexts.get(req)
+    const reqCtx = contexts.get(req)
     let childOf = tracer.extract(FORMAT_HTTP_HEADERS, headers)
 
     // we may have headers signaling a router proxy span should be created (such as for AWS API Gateway)
     if (tracer._config?.inferredProxyServicesEnabled) {
-      const proxySpan = createInferredProxySpan(headers, childOf, tracer, context)
+      const proxySpan = createInferredProxySpan(headers, childOf, tracer, reqCtx, traceCtx, startSpanHelper)
       if (proxySpan) {
         childOf = proxySpan
       }
     }
 
-    const span = tracer.startSpan(name, { childOf, links: childOf?._links })
+    const span = startSpanHelper(tracer, name, { childOf }, traceCtx)
 
     return span
   },