DNS with less AsyncResources

packages/datadog-instrumentations/src/dns.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 rrtypes = {
@@ -48,11 +48,12 @@ function patchResolveShorthands (prototype) {
 
 function wrap (prefix, fn, expectedArgs, rrtype) {
   const startCh = channel(prefix + ':start')
-  const finishCh = channel(prefix + ':finish')
+  const asyncEndCh = channel(prefix + ':async_end')
+  const endCh = channel(prefix + ':end')
   const errorCh = channel(prefix + ':error')
 
   const wrapped = function () {
-    const cb = AsyncResource.bind(arguments[arguments.length - 1])
+    const cb = arguments[arguments.length - 1]
     if (
       !startCh.hasSubscribers ||
       arguments.length < expectedArgs ||
@@ -67,28 +68,31 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
       startArgs.push(rrtype)
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish(startArgs)
+    const context = { args: startArgs }
+    startCh.publish(context)
 
-      arguments[arguments.length - 1] = asyncResource.bind(function (error, result) {
-        if (error) {
-          errorCh.publish(error)
-        }
-        finishCh.publish(result)
-        cb.apply(this, arguments)
-      })
+    arguments[arguments.length - 1] = function (error, result) {
+      if (error) {
+        context.error = error
+        errorCh.publish(context)
+      }
+      context.result = result
+      asyncEndCh.publish(context)
+      cb.apply(this, arguments)
+    }
 
-      try {
-        return fn.apply(this, arguments)
+    try {
+      return fn.apply(this, arguments)
       // TODO deal with promise versions when we support `dns/promises`
-      } catch (error) {
-        error.stack // trigger getting the stack at the original throwing point
-        errorCh.publish(error)
-
-        throw error
-      }
-    })
+    } catch (error) {
+      error.stack // trigger getting the stack at the original throwing point
+      context.error = error
+      errorCh.publish(context)
+
+      throw error
+    } finally {
+      endCh.publish(context)
+    }
   }
 
   return shimmer.wrap(fn, wrapped)

packages/datadog-plugin-dns/src/lookup.js

@@ -6,7 +6,7 @@ class DNSLookupPlugin extends ClientPlugin {
   static get name () { return 'dns' }
   static get operation () { return 'lookup' }
 
-  start ([hostname]) {
+  start ({ args: [hostname] }) {
     this.startSpan('dns.lookup', {
       service: this.config.service,
       resource: hostname,
@@ -19,7 +19,7 @@ class DNSLookupPlugin extends ClientPlugin {
     })
   }
 
-  finish (result) {
+  finish ({ result }) {
     const span = this.activeSpan
 
     if (Array.isArray(result)) {

packages/datadog-plugin-dns/src/lookup_service.js

@@ -6,7 +6,7 @@ class DNSLookupServicePlugin extends ClientPlugin {
   static get name () { return 'dns' }
   static get operation () { return 'lookup_service' }
 
-  start ([address, port]) {
+  start ({ args: [address, port] }) {
     this.startSpan('dns.lookup_service', {
       service: this.config.service,
       resource: `${address}:${port}`,

packages/datadog-plugin-dns/src/resolve.js

@@ -6,7 +6,7 @@ class DNSResolvePlugin extends ClientPlugin {
   static get name () { return 'dns' }
   static get operation () { return 'resolve' }
 
-  start ([hostname, maybeType]) {
+  start ({ args: [hostname, maybeType] }) {
     const rrtype = typeof maybeType === 'string' ? maybeType : 'A'
 
     this.startSpan('dns.resolve', {

packages/datadog-plugin-dns/src/reverse.js

@@ -6,7 +6,7 @@ class DNSReversePlugin extends ClientPlugin {
   static get name () { return 'dns' }
   static get operation () { return 'reverse' }
 
-  start ([ip]) {
+  start ({ args: [ip] }) {
     this.startSpan('dns.reverse', {
       service: this.config.service,
       resource: ip,

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

@@ -16,6 +16,11 @@ class OutgoingPlugin extends TracingPlugin {
     this.addHost(url.hostname, url.port)
   }
 
+  asyncEnd (...args) {
+    super.asyncEnd(...args)
+    this.exit(...args)
+  }
+
   addHost (hostname, port) {
     const span = this.activeSpan
 

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

@@ -30,6 +30,7 @@ module.exports = class Plugin {
     this._subscriptions = []
     this._enabled = false
     this._tracer = tracer
+    this._storesByContext = new WeakMap()
   }
 
   get tracer () {
@@ -41,6 +42,10 @@ module.exports = class Plugin {
     storage.enterWith({ ...store, span })
   }
 
+  exit (ctx) {
+    storage.enterWith(this.getStoreByContext(ctx))
+  }
+
   // TODO: Implement filters on resource name for all plugins.
   /** Prevents creation of spans here and for all async descendants. */
   skip () {
@@ -61,6 +66,14 @@ module.exports = class Plugin {
     }
   }
 
+  setStoreByContext (context, store) {
+    this._storesByContext.set(context, store)
+  }
+
+  getStoreByContext (context) {
+    return this._storesByContext.get(context)
+  }
+
   configure (config) {
     if (typeof config === 'boolean') {
       config = { enabled: config }

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

@@ -12,16 +12,28 @@ class TracingPlugin extends Plugin {
     this.operation = this.constructor.operation
 
     this.addTraceSub('start', message => {
+      if (message && typeof message === 'object') {
+        this.setStoreByContext(message, storage.getStore())
+      }
       this.start(message)
     })
 
     this.addTraceSub('error', err => {
       this.error(err)
     })
 
+    // TODO remove the finish event everywhere
     this.addTraceSub('finish', message => {
       this.finish(message)
     })
+
+    this.addTraceSub('async_end', message => {
+      this.asyncEnd(message)
+    })
+
+    this.addTraceSub('end', message => {
+      this.end(message)
+    })
   }
 
   get activeSpan () {
@@ -46,7 +58,19 @@ class TracingPlugin extends Plugin {
     this.activeSpan.finish()
   }
 
+  asyncEnd (...args) {
+    this.finish(...args)
+    // any non-IncomingPlugin plugins need to call exit here.
+  }
+
+  end (...args) {
+    this.exit(...args)
+  }
+
   error (error) {
+    if (error && typeof error === 'object' && error.error) {
+      error = error.error
+    }
     this.addError(error)
   }
 
@@ -83,7 +107,7 @@ class TracingPlugin extends Plugin {
 
     analyticsSampler.sample(span, this.config.measured)
 
-    storage.enterWith({ ...store, span })
+    this.enter(span, store)
 
     return span
   }