dd-trace 2.3.1 → 2.4.2

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

@@ -0,0 +1,117 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const startICPCh = channel('apm:net:ipc:start')
+const asyncICPEndCh = channel('apm:net:ipc:async-end')
+const endICPCh = channel('apm:net:ipc:end')
+const errorICPCh = channel('apm:net:ipc:error')
+
+const startTCPCh = channel('apm:net:tcp:start')
+const asyncTCPEndCh = channel('apm:net:tcp:async-end')
+const endTCPCh = channel('apm:net:tcp:end')
+const errorTCPCh = channel('apm:net:tcp:error')
+
+const connectionCh = channel(`apm:net:tcp:connection`)
+
+addHook({ name: 'net' }, net => {
+  require('dns')
+
+  shimmer.wrap(net.Socket.prototype, 'connect', connect => function () {
+    if (!startICPCh.hasSubscribers || !startTCPCh.hasSubscribers) {
+      return connect.apply(this, arguments)
+    }
+
+    const options = getOptions(arguments)
+    const lastIndex = arguments.length - 1
+    const callback = arguments[lastIndex]
+
+    if (!options) return connect.apply(this, arguments)
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+
+    if (typeof callback === 'function') {
+      arguments[lastIndex] = asyncResource.bind(callback)
+    }
+
+    const protocol = options.path ? 'ipc' : 'tcp'
+
+    if (protocol === 'ipc') {
+      startICPCh.publish({ options })
+      setupListeners(this, 'ipc', asyncResource)
+    } else {
+      startTCPCh.publish({ options })
+      setupListeners(this, 'tcp', asyncResource)
+    }
+
+    try {
+      return connect.apply(this, arguments)
+    } catch (err) {
+      protocol === 'ipc' ? errorICPCh.publish(err) : errorTCPCh.publish(err)
+
+      throw err
+    } finally {
+      protocol === 'ipc' ? endICPCh.publish(undefined) : endTCPCh.publish(undefined)
+    }
+  })
+
+  return net
+})
+
+function getOptions (args) {
+  if (!args[0]) return
+
+  switch (typeof args[0]) {
+    case 'object':
+      if (Array.isArray(args[0])) return getOptions(args[0])
+      return args[0]
+    case 'string':
+      if (isNaN(parseFloat(args[0]))) {
+        return {
+          path: args[0]
+        }
+      }
+    case 'number': // eslint-disable-line no-fallthrough
+      return {
+        port: args[0],
+        host: typeof args[1] === 'string' ? args[1] : 'localhost'
+      }
+  }
+}
+
+function setupListeners (socket, protocol, ar) {
+  const events = ['connect', 'error', 'close', 'timeout']
+
+  const wrapListener = AsyncResource.bind(function (error) {
+    if (error) {
+      protocol === 'ipc' ? errorICPCh.publish(error) : errorTCPCh.publish(error)
+    }
+    protocol === 'ipc' ? asyncICPEndCh.publish(undefined) : asyncTCPEndCh.publish(undefined)
+  })
+
+  const localListener = AsyncResource.bind(function () {
+    connectionCh.publish({ socket })
+  })
+
+  const cleanupListener = function () {
+    socket.removeListener('connect', localListener)
+    events.forEach(event => {
+      socket.removeListener(event, wrapListener)
+      socket.removeListener(event, cleanupListener)
+    })
+  }
+
+  if (protocol === 'tcp') {
+    socket.once('connect', localListener)
+  }
+
+  events.forEach(event => {
+    socket.once(event, wrapListener)
+    socket.once(event, cleanupListener)
+  })
+}

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