dd-trace 2.4.2 → 2.7.0

package/packages/datadog-instrumentations/src/google-cloud-pubsub.js

@@ -0,0 +1,100 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const requestStartCh = channel('apm:google-cloud-pubsub:request:start')
+const requestFinishCh = channel('apm:google-cloud-pubsub:request:finish')
+const requestErrorCh = channel('apm:google-cloud-pubsub:request:error')
+
+const receiveStartCh = channel(`apm:google-cloud-pubsub:receive:start`)
+const receiveFinishCh = channel('apm:google-cloud-pubsub:receive:finish')
+const receiveErrorCh = channel('apm:google-cloud-pubsub:receive:error')
+
+addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
+  const PubSub = obj.PubSub
+  const Subscription = obj.Subscription
+
+  shimmer.wrap(PubSub.prototype, 'request', request => function (cfg = { reqOpts: {} }, cb) {
+    if (!requestStartCh.hasSubscribers) {
+      return request.apply(this, arguments)
+    }
+
+    const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
+    const outerAsyncResource = new AsyncResource('bound-anonymous-fn')
+
+    return innerAsyncResource.runInAsyncScope(() => {
+      let messages = []
+      if (cfg.reqOpts && cfg.method === 'publish') {
+        messages = cfg.reqOpts.messages
+      }
+
+      requestStartCh.publish({ cfg, projectId: this.projectId, messages })
+      cb = outerAsyncResource.bind(cb)
+
+      const fn = () => {
+        arguments[1] = innerAsyncResource.bind(function (error) {
+          if (error) {
+            requestErrorCh.publish(error)
+          }
+          requestFinishCh.publish(undefined)
+          return cb.apply(this, arguments)
+        })
+        return request.apply(this, arguments)
+      }
+
+      try {
+        return fn.apply(this, arguments)
+      } catch (e) {
+        requestErrorCh.publish(e)
+        throw e
+      }
+    })
+  })
+
+  shimmer.wrap(Subscription.prototype, 'emit', emit => function (eventName, message) {
+    if (eventName !== 'message' || !message) return emit.apply(this, arguments)
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+
+    return asyncResource.runInAsyncScope(() => {
+      try {
+        return emit.apply(this, arguments)
+      } catch (err) {
+        receiveErrorCh.publish({ err, message })
+        throw err
+      }
+    })
+  })
+
+  return obj
+})
+
+addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/lease-manager.js' }, (obj) => {
+  const LeaseManager = obj.LeaseManager
+
+  shimmer.wrap(LeaseManager.prototype, '_dispense', dispense => function (message) {
+    if (receiveStartCh.hasSubscribers) {
+      receiveStartCh.publish({ message })
+    }
+    return dispense.apply(this, arguments)
+  })
+
+  shimmer.wrap(LeaseManager.prototype, 'remove', remove => function (message) {
+    receiveFinishCh.publish({ message })
+    return remove.apply(this, arguments)
+  })
+
+  shimmer.wrap(LeaseManager.prototype, 'clear', clear => function () {
+    for (const message of this._messages) {
+      receiveFinishCh.publish({ message })
+    }
+    return clear.apply(this, arguments)
+  })
+
+  return obj
+})

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

@@ -9,7 +9,7 @@ const shimmer = require('../../../datadog-shimmer')
 const startServerCh = channel('apm:http:server:request:start')
 const endServerCh = channel('apm:http:server:request:end')
 const errorServerCh = channel('apm:http:server:request:error')
-const asyncEndServerCh = channel('apm:http:server:request:async-end')
+const asyncEndServerCh = channel('apm:http:server:request:finish')
 
 addHook({ name: 'https' }, http => {
   // http.ServerResponse not present on https

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

@@ -0,0 +1,175 @@
+'use strict'
+
+const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const testStartCh = channel('ci:jest:test:start')
+const testSkippedCh = channel('ci:jest:test:skip')
+const testRunEndCh = channel('ci:jest:test:end')
+const testErrCh = channel('ci:jest:test:err')
+const testSuiteEnd = channel('ci:jest:test-suite:end')
+
+const {
+  getTestSuitePath,
+  getTestParametersString
+} = require('../../dd-trace/src/plugins/util/test')
+
+const { getFormattedJestTestParameters } = require('../../datadog-plugin-jest/src/util')
+
+const specStatusToTestStatus = {
+  'pending': 'skip',
+  'disabled': 'skip',
+  'todo': 'skip',
+  'passed': 'pass',
+  'failed': 'fail'
+}
+
+const asyncResources = new WeakMap()
+const originalTestFns = new WeakMap()
+
+// based on https://github.com/facebook/jest/blob/main/packages/jest-circus/src/formatNodeAssertErrors.ts#L41
+function formatJestError (errors) {
+  let error
+  if (Array.isArray(errors)) {
+    const [originalError, asyncError] = errors
+    if (originalError === null || !originalError.stack) {
+      error = asyncError
+      error.message = originalError
+    } else {
+      error = originalError
+    }
+  } else {
+    error = errors
+  }
+  return error
+}
+
+function getWrappedEnvironment (BaseEnvironment) {
+  return class DatadogEnvironment extends BaseEnvironment {
+    constructor (config, context) {
+      super(config, context)
+      const rootDir = config.globalConfig ? config.globalConfig.rootDir : config.rootDir
+      this.testSuite = getTestSuitePath(context.testPath, rootDir)
+      this.nameToParams = {}
+      this.global._ddtrace = global._ddtrace
+    }
+    async teardown () {
+      super.teardown().finally(() => {
+        testSuiteEnd.publish()
+      })
+    }
+
+    async handleTestEvent (event, state) {
+      if (super.handleTestEvent) {
+        await super.handleTestEvent(event, state)
+      }
+
+      const globalExpect = this.global.expect
+      const setNameToParams = (name, params) => { this.nameToParams[name] = params }
+
+      if (event.name === 'setup') {
+        shimmer.wrap(this.global.test, 'each', each => function () {
+          const testParameters = getFormattedJestTestParameters(arguments)
+          const eachBind = each.apply(this, arguments)
+          return function () {
+            const [testName] = arguments
+            setNameToParams(testName, testParameters)
+            return eachBind.apply(this, arguments)
+          }
+        })
+      }
+      if (event.name === 'test_start') {
+        const testParameters = getTestParametersString(this.nameToParams, event.test.name)
+
+        // Async resource for this test is created here
+        // It is used later on by the test_done handler
+        const asyncResource = new AsyncResource('bound-anonymous-fn')
+        asyncResources.set(event.test, asyncResource)
+        asyncResource.runInAsyncScope(() => {
+          testStartCh.publish({
+            name: globalExpect.getState().currentTestName,
+            suite: this.testSuite,
+            runner: 'jest-circus',
+            testParameters
+          })
+          originalTestFns.set(event.test, event.test.fn)
+          event.test.fn = asyncResource.bind(event.test.fn)
+        })
+      }
+      if (event.name === 'test_done') {
+        const asyncResource = asyncResources.get(event.test)
+        asyncResource.runInAsyncScope(() => {
+          let status = 'pass'
+          if (event.test.errors && event.test.errors.length) {
+            status = 'fail'
+            const formattedError = formatJestError(event.test.errors[0])
+            testErrCh.publish(formattedError)
+          }
+          testRunEndCh.publish(status)
+          // restore in case it is retried
+          event.test.fn = originalTestFns.get(event.test)
+        })
+      }
+      if (event.name === 'test_skip' || event.name === 'test_todo') {
+        testSkippedCh.publish({
+          name: globalExpect.getState().currentTestName,
+          suite: this.testSuite,
+          runner: 'jest-circus'
+        })
+      }
+    }
+  }
+}
+
+function getTestEnvironment (pkg) {
+  if (pkg.default) {
+    const wrappedTestEnvironment = getWrappedEnvironment(pkg.default)
+    pkg.default = wrappedTestEnvironment
+    pkg.TestEnvironment = wrappedTestEnvironment
+    return pkg
+  }
+  return getWrappedEnvironment(pkg)
+}
+
+addHook({
+  name: 'jest-environment-node',
+  versions: ['>=24.8.0']
+}, getTestEnvironment)
+
+addHook({
+  name: 'jest-environment-jsdom',
+  versions: ['>=24.8.0']
+}, getTestEnvironment)
+
+addHook({
+  name: 'jest-jasmine2',
+  versions: ['>=24.8.0'],
+  file: 'build/jasmineAsyncInstall.js'
+}, (jasmineAsyncInstallExport) => {
+  return function (globalConfig, globalInput) {
+    globalInput._ddtrace = global._ddtrace
+    shimmer.wrap(globalInput.jasmine.Spec.prototype, 'execute', execute => function (onComplete) {
+      const asyncResource = new AsyncResource('bound-anonymous-fn')
+      asyncResource.runInAsyncScope(() => {
+        const testSuite = getTestSuitePath(this.result.testPath, globalConfig.rootDir)
+        testStartCh.publish({
+          name: this.getFullName(),
+          suite: testSuite,
+          runner: 'jest-jasmine2'
+        })
+        const spec = this
+        const callback = asyncResource.bind(function () {
+          if (spec.result.failedExpectations && spec.result.failedExpectations.length) {
+            const formattedError = formatJestError(spec.result.failedExpectations[0].error)
+            testErrCh.publish(formattedError)
+          }
+          testRunEndCh.publish(specStatusToTestStatus[spec.result.status])
+          onComplete.apply(this, arguments)
+        })
+        arguments[0] = callback
+        execute.apply(this, arguments)
+      })
+    })
+    return jasmineAsyncInstallExport.default(globalConfig, globalInput)
+  }
+})

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 exitChannel = channel('apm:koa:middleware:exit')
+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] = AsyncResource.bind(next)
+      }
+
+      try {
+        const result = fn.apply(this, arguments)
+
+        if (result && typeof result.then === 'function') {
+          return result.then(
+            result => {
+              exit(ctx, layer)
+              return result
+            },
+            err => {
+              exit(ctx, layer, err)
+              throw err
+            }
+          )
+        } else {
+          exit(ctx, layer)
+          return result
+        }
+      } catch (e) {
+        exit(ctx, layer, e)
+        throw e
+      }
+    })
+  }
+}
+
+function exit (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 })
+  }
+
+  exitChannel.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
+})