dd-trace 2.3.1 → 2.4.0

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/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'
+    }
+  }
+}

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

@@ -0,0 +1,66 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+addHook({ name: 'tedious', versions: [ '>=1.0.0' ] }, tedious => {
+  const startCh = channel('apm:tedious:request:start')
+  const asyncEndCh = channel('apm:tedious:request:async-end')
+  const endCh = channel('apm:tedious:request:end')
+  const errorCh = channel('apm:tedious:request:error')
+  shimmer.wrap(tedious.Connection.prototype, 'makeRequest', makeRequest => function (request) {
+    if (!startCh.hasSubscribers) {
+      return request.apply(this, arguments)
+    }
+
+    const queryOrProcedure = getQueryOrProcedure(request)
+
+    if (!queryOrProcedure) {
+      return makeRequest.apply(this, arguments)
+    }
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+
+    const connectionConfig = this.config
+
+    startCh.publish({ queryOrProcedure, connectionConfig })
+
+    const cb = asyncResource.bind(request.callback, request)
+    request.callback = AsyncResource.bind(function (error) {
+      if (error) {
+        errorCh.publish(error)
+      }
+      asyncEndCh.publish(undefined)
+
+      return cb.apply(this, arguments)
+    }, null, request)
+
+    try {
+      return makeRequest.apply(this, arguments)
+    } catch (error) {
+      errorCh.publish(error)
+
+      throw error
+    } finally {
+      endCh.publish(undefined)
+    }
+  })
+
+  return tedious
+})
+
+function getQueryOrProcedure (request) {
+  if (!request.parameters) return
+
+  const statement = request.parametersByName.statement || request.parametersByName.stmt
+
+  if (!statement) {
+    return request.sqlTextOrProcedure
+  }
+
+  return statement.value
+}