remove async storage from http2 server instrumentation

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

@@ -5,8 +5,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 
@@ -30,14 +29,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 })
-      }
+    if (eventName !== 'close') return emit.apply(this, arguments)
 
+    return finishServerCh.runStores({ req: this.req, ...ctx }, () => {
       return emit.apply(this, arguments)
     })
   }
@@ -51,18 +47,17 @@ 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 => 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-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

@@ -12,11 +12,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 }) {
+  bindStart (ctx) {
+    const { req, res } = ctx
+
     const store = storage('legacy').getStore()
     const span = web.startSpan(
       this.tracer,
@@ -39,14 +41,20 @@ class Http2ServerPlugin extends ServerPlugin {
       context.res.writeHead = web.wrapWriteHead(context)
       context.instrumented = true
     }
+
+    return ctx.currentStore
   }
 
-  finish ({ req }) {
+  bindFinish (ctx) {
+    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.parentStore
   }
 
   error (error) {

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

@@ -20,7 +20,7 @@ const supportedProxies = {
   }
 }
 
-function createInferredProxySpan (headers, childOf, tracer, context) {
+function createInferredProxySpan (headers, childOf, tracer, reqCtx, store, traceCtx) {
   if (!headers) {
     return null
   }
@@ -57,13 +57,18 @@ function createInferredProxySpan (headers, childOf, tracer, context) {
   )
 
   tracer.scope().activate(span)
-  context.inferredProxySpan = span
+  reqCtx.inferredProxySpan = span
   childOf = span
 
   log.debug('Successfully created inferred proxy span.')
 
   setInferredProxySpanTags(span, proxyContext)
 
+  if (traceCtx) {
+    traceCtx.parentStore = store
+    traceCtx.currentStore = { ...store, span }
+  }
+
   return childOf
 }
 

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

@@ -10,6 +10,7 @@ const kinds = require('../../../../../ext/kinds')
 const urlFilter = require('./urlfilter')
 const { ERROR_MESSAGE, ERROR_TYPE, ERROR_STACK } = require('../../constants')
 const { createInferredProxySpan, finishInferredProxySpan } = require('./inferred_proxy')
+const { storage } = require('../../../../datadog-core')
 
 let extractIp
 
@@ -92,7 +93,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
@@ -101,7 +102,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
@@ -257,21 +258,28 @@ 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)
 
+    const store = storage('legacy').getStore()
+
     // 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, store, traceCtx)
       if (proxySpan) {
         childOf = proxySpan
       }
     }
 
     const span = tracer.startSpan(name, { childOf, links: childOf?._links })
 
+    if (traceCtx) {
+      traceCtx.parentStore = store
+      traceCtx.currentStore = { ...store, span }
+    }
+
     return span
   },