dd-trace 2.2.1 → 2.4.0

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

@@ -0,0 +1,122 @@
+const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+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 hookErrorCh = channel('ci:mocha:hook:error')
+const parameterizedTestCh = channel('ci:mocha:test:parameterize')
+const testRunEndCh = channel('ci:mocha:run:end')
+
+function isRetry (test) {
+  return test._currentRetry !== undefined && test._currentRetry !== 0
+}
+
+function getAllTestsInSuite (root) {
+  const tests = []
+  function getTests (suiteOrTest) {
+    suiteOrTest.tests.forEach(test => {
+      tests.push(test)
+    })
+    suiteOrTest.suites.forEach(suite => {
+      getTests(suite)
+    })
+  }
+  getTests(root)
+  return tests
+}
+
+addHook({
+  name: 'mocha',
+  versions: ['>=5.2.0'],
+  file: 'lib/runner.js'
+}, (Runner) => {
+  shimmer.wrap(Runner.prototype, 'runTest', runTest => function () {
+    if (!testStartCh.hasSubscribers) {
+      return runTest.apply(this, arguments)
+    }
+
+    if (!isRetry(this.test)) {
+      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'
+      }
+
+      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) {
+      return runTests.apply(this, arguments)
+    }
+    this.once('end', AsyncResource.bind(() => {
+      testRunEndCh.publish(undefined)
+    }))
+    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)
+  })
+
+  shimmer.wrap(Runner.prototype, 'fail', fail => function (hook, error) {
+    if (!hookErrorCh.hasSubscribers) {
+      return fail.apply(this, arguments)
+    }
+    if (error && hook.ctx && hook.ctx.currentTest) {
+      error.message = `${hook.title}: ${error.message}`
+      hookErrorCh.publish({ test: hook.ctx.currentTest, error })
+    }
+    return fail.apply(this, arguments)
+  })
+
+  return Runner
+})
+
+addHook({
+  name: 'mocha-each',
+  versions: ['>=2.0.1']
+}, (mochaEach) => {
+  return shimmer.wrap(mochaEach, function () {
+    const [params] = arguments
+    const { it, ...rest } = mochaEach.apply(this, arguments)
+    return {
+      it: function (name) {
+        parameterizedTestCh.publish({ name, params })
+        it.apply(this, arguments)
+      },
+      ...rest
+    }
+  })
+})

package/packages/datadog-instrumentations/src/mongodb-core.js

@@ -0,0 +1,179 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const startCh = channel(`apm:mongodb:query:start`)
+const endCh = channel(`apm:mongodb:query:end`)
+const asyncEndCh = channel(`apm:mongodb:query:async-end`)
+const errorCh = channel(`apm:mongodb:query:error`)
+
+addHook({ name: 'mongodb-core', versions: ['2 - 3.1.9'] }, Server => {
+  const serverProto = Server.Server.prototype
+  shimmer.wrap(serverProto, 'command', command => wrapCommand(command, 'command'))
+  shimmer.wrap(serverProto, 'insert', insert => wrapCommand(insert, 'insert', 'insert'))
+  shimmer.wrap(serverProto, 'update', update => wrapCommand(update, 'update', 'update'))
+  shimmer.wrap(serverProto, 'remove', remove => wrapCommand(remove, 'remove', 'remove'))
+
+  const cursorProto = Server.Cursor.prototype
+  shimmer.wrap(cursorProto, '_getmore', _getmore => wrapCursor(_getmore, 'getMore', 'getMore'))
+  shimmer.wrap(cursorProto, '_find', _find => wrapQuery(_find, '_find'))
+  shimmer.wrap(cursorProto, 'kill', kill => wrapCursor(kill, 'killCursors', 'killCursors'))
+  return Server
+})
+
+addHook({ name: 'mongodb', versions: ['>=4'], file: 'lib/cmap/connection.js' }, Connection => {
+  const proto = Connection.Connection.prototype
+  shimmer.wrap(proto, 'command', command => wrapConnectionCommand(command, 'command'))
+  shimmer.wrap(proto, 'query', query => wrapConnectionCommand(query, 'query'))
+  return Connection
+})
+
+addHook({ name: 'mongodb', versions: ['>=3.3 <4'], file: 'lib/core/wireprotocol/index.js' }, wp => wrapWp(wp))
+
+addHook({ name: 'mongodb-core', versions: ['>=3.2'], file: 'lib/wireprotocol/index.js' }, wp => wrapWp(wp))
+
+addHook({ name: 'mongodb-core', versions: ['~3.1.10'], file: 'lib/wireprotocol/3_2_support.js' }, WireProtocol => {
+  shimmer.wrap(WireProtocol.prototype, 'command', command => wrapUnifiedCommand(command, 'command'))
+  return WireProtocol
+})
+
+addHook({ name: 'mongodb-core', versions: ['~3.1.10'], file: 'lib/wireprotocol/2_6_support.js' }, WireProtocol => {
+  shimmer.wrap(WireProtocol.prototype, 'command', command => wrapUnifiedCommand(command, 'command'))
+  return WireProtocol
+})
+
+addHook({ name: 'mongodb', versions: ['>=3.5.4'], file: 'lib/utils.js' }, util => {
+  shimmer.wrap(util, 'maybePromise', maybePromise => function (parent, callback, fn) {
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const callbackIndex = arguments.length - 2
+
+    callback = arguments[callbackIndex]
+
+    if (typeof callback === 'function') {
+      arguments[callbackIndex] = asyncResource.bind(callback)
+    }
+
+    return maybePromise.apply(this, arguments)
+  })
+  return util
+})
+
+function wrapWp (wp) {
+  shimmer.wrap(wp, 'command', command => wrapUnifiedCommand(command, 'command'))
+  shimmer.wrap(wp, 'insert', insert => wrapUnifiedCommand(insert, 'insert', 'insert'))
+  shimmer.wrap(wp, 'update', update => wrapUnifiedCommand(update, 'update', 'update'))
+  shimmer.wrap(wp, 'remove', remove => wrapUnifiedCommand(remove, 'remove', 'remove'))
+  shimmer.wrap(wp, 'query', query => wrapUnifiedCommand(query, 'query'))
+  shimmer.wrap(wp, 'getMore', getMore => wrapUnifiedCommand(getMore, 'getMore', 'getMore'))
+  shimmer.wrap(wp, 'killCursors', killCursors => wrapUnifiedCommand(killCursors, 'killCursors', 'killCursors'))
+  return wp
+}
+
+function wrapUnifiedCommand (command, operation, name) {
+  const wrapped = function (server, ns, ops) {
+    if (!startCh.hasSubscribers) {
+      return command.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    return instrument(asyncResource, operation, command, this, arguments, server, ns, ops, { name })
+  }
+  return shimmer.wrap(command, wrapped)
+}
+
+function wrapConnectionCommand (command, operation, name) {
+  const wrapped = function (ns, ops) {
+    if (!startCh.hasSubscribers) {
+      return command.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const hostParts = typeof this.address === 'string' ? this.address.split(':') : ''
+    const options = hostParts.length === 2
+      ? { host: hostParts[0], port: hostParts[1] }
+      : {} // no port means the address is a random UUID so no host either
+    const topology = { s: { options } }
+
+    ns = `${ns.db}.${ns.collection}`
+    return instrument(asyncResource, operation, command, this, arguments, topology, ns, ops, { name })
+  }
+  return shimmer.wrap(command, wrapped)
+}
+
+function wrapQuery (query, operation, name) {
+  const wrapped = function () {
+    if (!startCh.hasSubscribers) {
+      return query.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const pool = this.server.s.pool
+    const ns = this.ns
+    const ops = this.cmd
+    return instrument(asyncResource, operation, query, this, arguments, pool, ns, ops)
+  }
+
+  return shimmer.wrap(query, wrapped)
+}
+
+function wrapCursor (cursor, operation, name) {
+  const wrapped = function () {
+    if (!startCh.hasSubscribers) {
+      return cursor.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const pool = this.server.s.pool
+    const ns = this.ns
+    return instrument(asyncResource, operation, cursor, this, arguments, pool, ns, {}, { name })
+  }
+  return shimmer.wrap(cursor, wrapped)
+}
+
+function wrapCommand (command, operation, name) {
+  const wrapped = function (ns, ops) {
+    if (!startCh.hasSubscribers) {
+      return command.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    return instrument(asyncResource, operation, command, this, arguments, this, ns, ops, { name })
+  }
+  return shimmer.wrap(command, wrapped)
+}
+
+function instrument (ar, operation, command, ctx, args, server, ns, ops, options = {}) {
+  const name = options.name || (ops && Object.keys(ops)[0])
+  const index = args.length - 1
+  let callback = args[index]
+
+  if (typeof callback !== 'function') return command.apply(ctx, args)
+
+  callback = ar.bind(callback)
+
+  const serverInfo = server && server.s && server.s.options
+
+  startCh.publish({ ns, ops, options: serverInfo, name })
+
+  args[index] = AsyncResource.bind(function (err, res) {
+    if (err) {
+      errorCh.publish(err)
+    }
+
+    asyncEndCh.publish(undefined)
+
+    if (callback) {
+      return callback.apply(this, arguments)
+    }
+  })
+
+  try {
+    return command.apply(ctx, args)
+  } catch (err) {
+    errorCh.publish(err)
+
+    throw err
+  } finally {
+    endCh.publish(undefined)
+  }
+}

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

@@ -0,0 +1,75 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const startCh = channel('apm:pg:query:start')
+const asyncEndCh = channel('apm:pg:query:async-end')
+const endCh = channel('apm:pg:query:end')
+const errorCh = channel('apm:pg:query:error')
+
+addHook({ name: 'pg', versions: ['>=4'] }, pg => {
+  shimmer.wrap(pg.Client.prototype, 'query', query => wrapQuery(query))
+  return pg
+})
+
+addHook({ name: 'pg', file: 'lib/native/index.js', versions: ['>=4'] }, Client => {
+  shimmer.wrap(Client.prototype, 'query', query => wrapQuery(query))
+  return Client
+})
+
+function wrapQuery (query) {
+  return function () {
+    if (!startCh.hasSubscribers) {
+      return query.apply(this, arguments)
+    }
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+
+    const retval = query.apply(this, arguments)
+
+    const queryQueue = this.queryQueue || this._queryQueue
+    const activeQuery = this.activeQuery || this._activeQuery
+    const pgQuery = queryQueue[queryQueue.length - 1] || activeQuery
+
+    if (!pgQuery) {
+      return retval
+    }
+    const statement = pgQuery.text
+
+    startCh.publish({ params: this.connectionParameters, statement })
+
+    const finish = AsyncResource.bind(function (error) {
+      if (error) {
+        errorCh.publish(error)
+      }
+      asyncEndCh.publish(undefined)
+    })
+
+    if (pgQuery.callback) {
+      const originalCallback = asyncResource.bind(pgQuery.callback)
+      pgQuery.callback = function (err, res) {
+        finish(err)
+        return originalCallback.apply(this, arguments)
+      }
+    } else if (pgQuery.once) {
+      pgQuery
+        .once('error', finish)
+        .once('end', () => finish())
+    } else {
+      pgQuery.then(() => finish(), finish)
+    }
+
+    try {
+      return retval
+    } catch (err) {
+      errorCh.publish(err)
+    } finally {
+      endCh.publish(undefined)
+    }
+  }
+}

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

@@ -21,14 +21,14 @@ addHook({ name: '@node-redis/client', file: 'dist/lib/client/commands-queue.js',
     const name = command[0]
     const args = command.slice(1)
 
-    startSpan(this, name, args)
+    start(this, name, args)
 
     const res = addCommand.apply(this, arguments)
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const onResolve = asyncResource.bind(() => finish(asyncEndCh, errorCh))
+    const onReject = asyncResource.bind(err => finish(asyncEndCh, errorCh, err))
 
-    res.then(
-      () => finish(asyncEndCh, errorCh),
-      err => finish(asyncEndCh, errorCh, err)
-    )
+    res.then(onResolve, onReject)
     endCh.publish(undefined)
     return res
   })
@@ -45,7 +45,7 @@ addHook({ name: 'redis', versions: ['>=2.6 <4'] }, redis => {
 
     const cb = asyncResource.bind(options.callback)
 
-    startSpan(this, options.command, options.args)
+    start(this, options.command, options.args)
 
     options.callback = AsyncResource.bind(wrapCallback(asyncEndCh, errorCh, cb))
 
@@ -70,7 +70,7 @@ addHook({ name: 'redis', versions: ['>=0.12 <2.6'] }, redis => {
 
     const asyncResource = new AsyncResource('bound-anonymous-fn')
 
-    startSpan(this, command, args)
+    start(this, command, args)
 
     if (typeof callback === 'function') {
       const cb = asyncResource.bind(callback)
@@ -95,7 +95,7 @@ addHook({ name: 'redis', versions: ['>=0.12 <2.6'] }, redis => {
   return redis
 })
 
-function startSpan (client, command, args) {
+function start (client, command, args) {
   const db = client.selected_db
   const connectionOptions = client.connection_options || client.connection_option || client.connectionOption || {}
   startCh.publish({ db, command, args, connectionOptions })

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

@@ -0,0 +1,224 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const circularBufferConstructor = Symbol('circularBufferConstructor')
+const inFlightDeliveries = Symbol('inFlightDeliveries')
+
+const patched = new WeakSet()
+const dispatchCh = channel('apm:rhea:dispatch')
+const errorCh = channel('apm:rhea:error')
+const asyncEndCh = channel('apm:rhea:async-end')
+const endCh = channel('apm:rhea:end')
+
+const contexts = new WeakMap()
+
+addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/link.js' }, obj => {
+  const startSendCh = channel('apm:rhea:send:start')
+  const startReceiveCh = channel('apm:rhea:receive:start')
+
+  const Sender = obj.Sender
+  const Receiver = obj.Receiver
+  shimmer.wrap(Sender.prototype, 'send', send => function (msg, tag, format) {
+    if (!canTrace(this)) {
+      // we can't handle disconnects or ending spans, so we can't safely instrument
+      return send.apply(this, arguments)
+    }
+
+    const { host, port } = getHostAndPort(this.connection)
+
+    const targetAddress = this.options && this.options.target &&
+      this.options.target.address ? this.options.target.address : undefined
+
+    startSendCh.publish({ targetAddress, host, port, msg })
+    const delivery = send.apply(this, arguments)
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const context = {
+      asyncResource
+    }
+    contexts.set(delivery, context)
+
+    addToInFlightDeliveries(this.connection, delivery)
+    try {
+      return delivery
+    } catch (err) {
+      errorCh.publish(err)
+
+      throw err
+    } finally {
+      endCh.publish(undefined)
+    }
+  })
+
+  shimmer.wrap(Receiver.prototype, 'dispatch', dispatch => function (eventName, msgObj) {
+    if (!canTrace(this)) {
+      // we can't handle disconnects or ending spans, so we can't safely instrument
+      return dispatch.apply(this, arguments)
+    }
+
+    if (eventName === 'message' && msgObj) {
+      startReceiveCh.publish({ msgObj, connection: this.connection })
+
+      if (msgObj.delivery) {
+        const asyncResource = new AsyncResource('bound-anonymous-fn')
+        const context = {
+          asyncResource
+        }
+        contexts.set(msgObj.delivery, context)
+        msgObj.delivery.update = wrapDeliveryUpdate(msgObj.delivery, msgObj.delivery.update)
+        addToInFlightDeliveries(this.connection, msgObj.delivery)
+      }
+      try {
+        return dispatch.apply(this, arguments)
+      } catch (err) {
+        errorCh.publish(err)
+
+        throw err
+      } finally {
+        endCh.publish(undefined)
+      }
+    }
+
+    return dispatch.apply(this, arguments)
+  })
+  return obj
+})
+
+addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/connection.js' }, Connection => {
+  shimmer.wrap(Connection.prototype, 'dispatch', dispatch => function (eventName, obj) {
+    if (eventName === 'disconnected') {
+      const error = obj.error || this.saved_error
+      if (this[inFlightDeliveries]) {
+        this[inFlightDeliveries].forEach(delivery => {
+          const context = contexts.get(delivery)
+          const asyncResource = context.asyncResource
+          asyncResource.runInAsyncScope(() => {
+            errorCh.publish(error)
+            finish(delivery, null)
+          })
+        })
+      }
+    }
+    return dispatch.apply(this, arguments)
+  })
+  return Connection
+})
+
+addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/session.js' }, (Session) => {
+  patchCircularBuffer(Session.prototype, Session)
+  return Session
+})
+
+function canTrace (link) {
+  return link.connection && link.session && link.session.outgoing
+}
+
+function getHostAndPort (connection) {
+  let host
+  let port
+  if (connection && connection.options) {
+    host = connection.options.host
+    port = connection.options.port
+  }
+  return { host, port }
+}
+
+function wrapDeliveryUpdate (obj, update) {
+  const context = contexts.get(obj)
+  const asyncResource = context.asyncResource
+  if (obj && asyncResource) {
+    const cb = asyncResource.bind(update)
+    return AsyncResource.bind(function wrappedUpdate (settled, stateData) {
+      const state = getStateFromData(stateData)
+      dispatchCh.publish({ state })
+      return cb.apply(this, arguments)
+    })
+  }
+  return function wrappedUpdate (settled, stateData) {
+    return update.apply(this, arguments)
+  }
+}
+
+function patchCircularBuffer (proto, Session) {
+  Object.defineProperty(proto, 'outgoing', {
+    configurable: true,
+    get () {},
+    set (outgoing) {
+      delete proto.outgoing // removes the setter on the prototype
+      this.outgoing = outgoing // assigns on the instance, like normal
+      if (outgoing) {
+        let CircularBuffer
+        if (outgoing.deliveries) {
+          CircularBuffer = outgoing.deliveries.constructor
+        }
+        if (CircularBuffer && !patched.has(CircularBuffer.prototype)) {
+          shimmer.wrap(CircularBuffer.prototype, 'pop_if', popIf => function (fn) {
+            arguments[0] = AsyncResource.bind(function (entry) {
+              const context = contexts.get(entry)
+              const asyncResource = context.asyncResource
+              let shouldPop
+              if (asyncResource) {
+                fn = asyncResource.bind(fn, this)
+                shouldPop = fn(entry)
+                if (shouldPop && asyncResource) {
+                  const remoteState = entry.remote_state
+                  const state = remoteState && remoteState.constructor
+                    ? entry.remote_state.constructor.composite_type : undefined
+                  asyncResource.runInAsyncScope(() => {
+                    return finish(entry, state)
+                  })
+                }
+              }
+
+              return shouldPop
+            })
+            return popIf.apply(this, arguments)
+          })
+          patched.add(CircularBuffer.prototype)
+          const Session = proto.constructor
+          if (Session) {
+            Session[circularBufferConstructor] = CircularBuffer
+          }
+        }
+      }
+    }
+  })
+}
+
+function addToInFlightDeliveries (connection, delivery) {
+  let deliveries = connection[inFlightDeliveries]
+  if (!deliveries) {
+    deliveries = new Set()
+    connection[inFlightDeliveries] = deliveries
+  }
+  deliveries.add(delivery)
+}
+
+function finish (delivery, state) {
+  const obj = contexts.get(delivery)
+  if (obj) {
+    if (state) {
+      dispatchCh.publish({ state })
+    }
+    asyncEndCh.publish(undefined)
+    if (obj.connection && obj.connection[inFlightDeliveries]) {
+      obj.connection[inFlightDeliveries].delete(delivery)
+    }
+  }
+}
+
+function getStateFromData (stateData) {
+  if (stateData && stateData.descriptor && stateData.descriptor) {
+    switch (stateData.descriptor.value) {
+      case 0x24: return 'accepted'
+      case 0x25: return 'rejected'
+      case 0x26: return 'released'
+      case 0x27: return 'modified'
+    }
+  }
+}