dd-trace 2.6.0 → 2.8.0

package/packages/datadog-instrumentations/src/kafkajs.js

@@ -0,0 +1,112 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const producerStartCh = channel('apm:kafkajs:produce:start')
+const producerFinishCh = channel('apm:kafkajs:produce:finish')
+const producerErrorCh = channel('apm:kafkajs:produce:error')
+
+const consumerStartCh = channel('apm:kafkajs:consume:start')
+const consumerFinishCh = channel('apm:kafkajs:consume:finish')
+const consumerErrorCh = channel('apm:kafkajs:consume:error')
+
+addHook({ name: 'kafkajs', versions: ['>=1.4'] }, (obj) => {
+  const Kafka = obj.Kafka
+  shimmer.wrap(Kafka.prototype, 'producer', createProducer => function () {
+    const producer = createProducer.apply(this, arguments)
+    const send = producer.send
+
+    producer.send = function () {
+      const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
+
+      return innerAsyncResource.runInAsyncScope(() => {
+        if (!producerStartCh.hasSubscribers) {
+          return send.apply(this, arguments)
+        }
+
+        try {
+          const { topic, messages = [] } = arguments[0]
+          for (const message of messages) {
+            if (typeof message === 'object') {
+              message.headers = message.headers || {}
+            }
+          }
+          producerStartCh.publish({ topic, messages })
+
+          const result = send.apply(this, arguments)
+
+          result.then(
+            innerAsyncResource.bind(() => producerFinishCh.publish(undefined)),
+            innerAsyncResource.bind(err => {
+              if (err) {
+                producerErrorCh.publish(err)
+              }
+              producerFinishCh.publish(undefined)
+            })
+          )
+
+          return result
+        } catch (e) {
+          producerErrorCh.publish(e)
+          producerFinishCh.publish(undefined)
+          throw e
+        }
+      })
+    }
+    return producer
+  })
+
+  shimmer.wrap(Kafka.prototype, 'consumer', createConsumer => function () {
+    if (!consumerStartCh.hasSubscribers) {
+      return createConsumer.apply(this, arguments)
+    }
+
+    const consumer = createConsumer.apply(this, arguments)
+    const run = consumer.run
+
+    consumer.run = function ({ eachMessage, ...runArgs }) {
+      if (typeof eachMessage !== 'function') return run({ eachMessage, ...runArgs })
+
+      return run({
+        eachMessage: function (...eachMessageArgs) {
+          const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
+          return innerAsyncResource.runInAsyncScope(() => {
+            const { topic, partition, message } = eachMessageArgs[0]
+            consumerStartCh.publish({ topic, partition, message })
+            try {
+              const result = eachMessage.apply(this, eachMessageArgs)
+
+              if (result && typeof result.then === 'function') {
+                result.then(
+                  innerAsyncResource.bind(() => consumerFinishCh.publish(undefined)),
+                  innerAsyncResource.bind(err => {
+                    if (err) {
+                      consumerErrorCh.publish(err)
+                    }
+                    consumerFinishCh.publish(undefined)
+                  })
+                )
+              } else {
+                consumerFinishCh.publish(undefined)
+              }
+
+              return result
+            } catch (e) {
+              consumerErrorCh.publish(e)
+              consumerFinishCh.publish(undefined)
+              throw e
+            }
+          })
+        },
+        ...runArgs
+      })
+    }
+    return consumer
+  })
+  return obj
+})

package/packages/datadog-instrumentations/src/knex.js

@@ -0,0 +1,20 @@
+'use strict'
+
+const { addHook } = require('./helpers/instrument')
+const { wrapThen } = require('./helpers/promise')
+const shimmer = require('../../datadog-shimmer')
+
+patch('lib/query/builder.js')
+patch('lib/raw.js')
+patch('lib/schema/builder.js')
+
+function patch (file) {
+  addHook({
+    name: 'knex',
+    versions: ['>=0.8.0'],
+    file
+  }, Builder => {
+    shimmer.wrap(Builder.prototype, 'then', wrapThen)
+    return Builder
+  })
+}

package/packages/datadog-instrumentations/src/koa.js

@@ -0,0 +1,159 @@
+'use strict'
+
+const shimmer = require('../../datadog-shimmer')
+const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+
+const enterChannel = channel('apm:koa:middleware:enter')
+const errorChannel = channel('apm:koa:middleware:error')
+const nextChannel = channel('apm:koa:middleware:next')
+const handleChannel = channel('apm:koa:request:handle')
+const routeChannel = channel('apm:koa:request:route')
+
+const originals = new WeakMap()
+
+function wrapCallback (callback) {
+  return function callbackWithTrace () {
+    const handleRequest = callback.apply(this, arguments)
+
+    if (typeof handleRequest !== 'function') return handleRequest
+
+    return function handleRequestWithTrace (req, res) {
+      handleChannel.publish({ req, res })
+
+      return handleRequest.apply(this, arguments)
+    }
+  }
+}
+
+function wrapUse (use) {
+  return function useWithTrace () {
+    const result = use.apply(this, arguments)
+
+    if (!Array.isArray(this.middleware)) return result
+
+    const fn = this.middleware.pop()
+
+    this.middleware.push(wrapMiddleware(fn))
+
+    return result
+  }
+}
+
+function wrapRegister (register) {
+  return function registerWithTrace (path, methods, middleware, opts) {
+    const route = register.apply(this, arguments)
+
+    if (!Array.isArray(path) && route && Array.isArray(route.stack)) {
+      wrapStack(route)
+    }
+
+    return route
+  }
+}
+
+function wrapRouterUse (use) {
+  return function useWithTrace () {
+    const router = use.apply(this, arguments)
+
+    router.stack.forEach(wrapStack)
+
+    return router
+  }
+}
+
+function wrapStack (layer) {
+  layer.stack = layer.stack.map(middleware => {
+    if (typeof middleware !== 'function') return middleware
+
+    const original = originals.get(middleware)
+
+    middleware = original || middleware
+
+    const handler = shimmer.wrap(middleware, wrapMiddleware(middleware, layer))
+
+    originals.set(handler, middleware)
+
+    return handler
+  })
+}
+
+function wrapMiddleware (fn, layer) {
+  if (typeof fn !== 'function') return fn
+
+  const name = fn.name
+
+  return function (ctx, next) {
+    if (!ctx || !enterChannel.hasSubscribers) return fn.apply(this, arguments)
+
+    const middlewareResource = new AsyncResource('bound-anonymous-fn')
+    const req = ctx.req
+
+    return middlewareResource.runInAsyncScope(() => {
+      enterChannel.publish({ req, name })
+
+      if (typeof next === 'function') {
+        arguments[1] = next
+      }
+
+      try {
+        const result = fn.apply(this, arguments)
+
+        if (result && typeof result.then === 'function') {
+          return result.then(
+            result => {
+              fulfill(ctx, layer)
+              return result
+            },
+            err => {
+              fulfill(ctx, layer, err)
+              throw err
+            }
+          )
+        } else {
+          fulfill(ctx, layer)
+          return result
+        }
+      } catch (e) {
+        fulfill(ctx, layer, e)
+        throw e
+      }
+    })
+  }
+}
+
+function fulfill (ctx, layer, error) {
+  if (error) {
+    errorChannel.publish(error)
+  }
+
+  const req = ctx.req
+  const route = ctx.routePath || (layer && layer.path)
+
+  // TODO: make sure that the parent class cannot override this in `enter`
+  if (route) {
+    routeChannel.publish({ req, route })
+  }
+
+  nextChannel.publish(ctx)
+}
+
+addHook({ name: 'koa', versions: ['>=2'] }, Koa => {
+  shimmer.wrap(Koa.prototype, 'callback', wrapCallback)
+  shimmer.wrap(Koa.prototype, 'use', wrapUse)
+
+  return Koa
+})
+
+addHook({ name: '@koa/router', versions: ['>=8'] }, Router => {
+  shimmer.wrap(Router.prototype, 'register', wrapRegister)
+  shimmer.wrap(Router.prototype, 'use', wrapRouterUse)
+
+  return Router
+})
+
+addHook({ name: 'koa-router', versions: ['>=7'] }, Router => {
+  shimmer.wrap(Router.prototype, 'register', wrapRegister)
+  shimmer.wrap(Router.prototype, 'use', wrapRouterUse)
+
+  return Router
+})

package/packages/datadog-instrumentations/src/limitd-client.js

@@ -0,0 +1,21 @@
+'use strict'
+
+const { addHook, AsyncResource } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+function wrapRequest (original) {
+  return function () {
+    const id = arguments.length - 1
+    arguments[id] = AsyncResource.bind(arguments[id])
+    return original.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: 'limitd-client',
+  versions: ['>=2.8']
+}, LimitdClient => {
+  shimmer.wrap(LimitdClient.prototype, '_directRequest', wrapRequest)
+  shimmer.wrap(LimitdClient.prototype, '_retriedRequest', wrapRequest)
+  return LimitdClient
+})

package/packages/datadog-instrumentations/src/memcached.js

@@ -10,8 +10,7 @@ const shimmer = require('../../datadog-shimmer')
 addHook({ name: 'memcached', versions: ['>=2.2'] }, Memcached => {
   const startCh = channel('apm:memcached:command:start')
   const startWithArgsCh = channel('apm:memcached:command:start:with-args')
-  const asyncEndCh = channel('apm:memcached:command:async-end')
-  const endCh = channel('apm:memcached:command:end')
+  const finishCh = channel('apm:memcached:command:finish')
   const errorCh = channel('apm:memcached:command:error')
 
   shimmer.wrap(Memcached.prototype, 'command', command => function (queryCompiler, server) {
@@ -19,34 +18,37 @@ addHook({ name: 'memcached', versions: ['>=2.2'] }, Memcached => {
       return command.apply(this, arguments)
     }
 
+    const callbackResource = new AsyncResource('bound-anonymous-fn')
     const asyncResource = new AsyncResource('bound-anonymous-fn')
 
     const client = this
 
-    const wrappedQueryCompiler = function () {
+    const wrappedQueryCompiler = asyncResource.bind(function () {
       const query = queryCompiler.apply(this, arguments)
-      const callback = asyncResource.bind(query.callback)
+      const callback = callbackResource.bind(query.callback)
 
-      query.callback = AsyncResource.bind(function (err) {
+      query.callback = asyncResource.bind(function (err) {
         if (err) {
           errorCh.publish(err)
         }
-        asyncEndCh.publish(undefined)
+        finishCh.publish()
 
         return callback.apply(this, arguments)
       })
       startWithArgsCh.publish({ client, server, query })
 
       return query
-    }
+    })
+
+    return asyncResource.runInAsyncScope(() => {
+      startCh.publish()
 
-    startCh.publish(undefined)
+      arguments[0] = wrappedQueryCompiler
 
-    arguments[0] = wrappedQueryCompiler
+      const result = command.apply(this, arguments)
 
-    const result = command.apply(this, arguments)
-    endCh.publish(undefined)
-    return result
+      return result
+    })
   })
 
   return Memcached

package/packages/datadog-instrumentations/src/mocha.js

@@ -4,12 +4,11 @@ const shimmer = require('../../datadog-shimmer')
 const testStartCh = channel('ci:mocha:test:start')
 const errorCh = channel('ci:mocha:test:error')
 const skipCh = channel('ci:mocha:test:skip')
-const testEndCh = channel('ci:mocha:test:end')
-const testAsyncEndCh = channel('ci:mocha:test:async-end')
-const suiteEndCh = channel('ci:mocha:suite:end')
+const testFinishCh = channel('ci:mocha:test:finish')
+const suiteFinishCh = channel('ci:mocha:suite:finish')
 const hookErrorCh = channel('ci:mocha:hook:error')
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')
-const testRunEndCh = channel('ci:mocha:run:end')
+const testRunFinishCh = channel('ci:mocha:run:finish')
 
 // TODO: remove when root hooks and fixtures are implemented
 const patched = new WeakSet()
@@ -38,59 +37,58 @@ function mochaHook (Runner) {
   patched.add(Runner)
 
   shimmer.wrap(Runner.prototype, 'runTest', runTest => function () {
-    if (!testStartCh.hasSubscribers) {
+    if (!testStartCh.hasSubscribers || isRetry(this.test)) {
       return runTest.apply(this, arguments)
     }
 
-    if (!isRetry(this.test)) {
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    return asyncResource.runInAsyncScope(() => {
       testStartCh.publish(this.test)
-    }
-
-    this.once('test end', AsyncResource.bind(() => {
-      let status
 
-      if (this.test.pending) {
-        status = 'skipped'
-      } else if (this.test.state !== 'failed' && !this.test.timedOut) {
-        status = 'pass'
-      } else {
-        status = 'fail'
+      this.once('test end', AsyncResource.bind(() => {
+        let status
+
+        if (this.test.pending) {
+          status = 'skipped'
+        } else if (this.test.state !== 'failed' && !this.test.timedOut) {
+          status = 'pass'
+        } else {
+          status = 'fail'
+        }
+
+        testFinishCh.publish(status)
+      }))
+
+      this.once('fail', AsyncResource.bind((test, err) => {
+        errorCh.publish(err)
+      }))
+
+      this.once('pending', AsyncResource.bind((test) => {
+        skipCh.publish(test)
+      }))
+
+      try {
+        return runTest.apply(this, arguments)
+      } catch (err) {
+        errorCh.publish(err)
+        throw err
       }
-
-      testAsyncEndCh.publish(status)
-    }))
-
-    this.once('fail', AsyncResource.bind((test, err) => {
-      errorCh.publish(err)
-    }))
-
-    this.once('pending', AsyncResource.bind((test) => {
-      skipCh.publish(test)
-    }))
-
-    try {
-      return runTest.apply(this, arguments)
-    } catch (err) {
-      errorCh.publish(err)
-      throw err
-    } finally {
-      testEndCh.publish(undefined)
-    }
+    })
   })
 
   shimmer.wrap(Runner.prototype, 'runTests', runTests => function () {
-    if (!suiteEndCh.hasSubscribers) {
+    if (!suiteFinishCh.hasSubscribers) {
       return runTests.apply(this, arguments)
     }
     this.once('end', AsyncResource.bind(() => {
-      testRunEndCh.publish(undefined)
+      testRunFinishCh.publish()
     }))
     runTests.apply(this, arguments)
     const suite = arguments[0]
     // We call `getAllTestsInSuite` with the root suite so every skipped test
     // should already have an associated test span.
     const tests = getAllTestsInSuite(suite)
-    suiteEndCh.publish(tests)
+    suiteFinishCh.publish(tests)
   })
 
   shimmer.wrap(Runner.prototype, 'fail', fail => function (hook, error) {

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

@@ -0,0 +1,46 @@
+'use strict'
+
+const { channel, addHook, AsyncResource } = require('../helpers/instrument')
+const shimmer = require('../../../datadog-shimmer')
+
+const startChannel = channel('apm:moleculer:call:start')
+const finishChannel = channel('apm:moleculer:call:finish')
+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 promise = call.apply(this, arguments)
+      const broker = this
+      const ctx = promise.ctx
+
+      return promise
+        .then(
+          result => {
+            finishChannel.publish({ broker, ctx })
+            return result
+          },
+          error => {
+            errorChannel.publish(error)
+            finishChannel.publish({ broker, ctx })
+            throw error
+          }
+        )
+    })
+  }
+}
+
+addHook({ name: 'moleculer', versions: ['>=0.14'] }, moleculer => {
+  shimmer.wrap(moleculer.ServiceBroker.prototype, 'call', wrapCall)
+
+  return moleculer
+})

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

@@ -0,0 +1,59 @@
+'use strict'
+
+const { channel, addHook, AsyncResource } = require('../helpers/instrument')
+const shimmer = require('../../../datadog-shimmer')
+
+const startChannel = channel('apm:moleculer:action:start')
+const finishChannel = channel('apm:moleculer:action:finish')
+const errorChannel = channel('apm:moleculer:action:error')
+
+function wrapRegisterMiddlewares (registerMiddlewares) {
+  return function (userMiddlewares) {
+    if (this.middlewares && this.middlewares.add) {
+      this.middlewares.add(createMiddleware())
+    }
+
+    return registerMiddlewares.apply(this, arguments)
+  }
+}
+
+function createMiddleware () {
+  return {
+    name: 'Datadog',
+
+    localAction (next, action) {
+      const broker = this
+
+      return function datadogMiddleware (ctx) {
+        const actionResource = new AsyncResource('bound-anonymous-fn')
+
+        return actionResource.runInAsyncScope(() => {
+          startChannel.publish({ action, ctx, broker })
+
+          try {
+            return next(ctx).then(
+              result => {
+                finishChannel.publish()
+                return result
+              },
+              error => {
+                errorChannel.publish(error)
+                finishChannel.publish()
+                throw error
+              }
+            )
+          } catch (e) {
+            errorChannel.publish(e)
+            finishChannel.publish()
+          }
+        })
+      }
+    }
+  }
+}
+
+addHook({ name: 'moleculer', versions: ['>=0.14'] }, moleculer => {
+  shimmer.wrap(moleculer.ServiceBroker.prototype, 'registerMiddlewares', wrapRegisterMiddlewares)
+
+  return moleculer
+})

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

@@ -0,0 +1,4 @@
+'use strict'
+
+require('./moleculer/client')
+require('./moleculer/server')