dd-trace 4.47.1 → 4.49.0

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

@@ -52,45 +52,59 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     const send = producer.send
     const bootstrapServers = this._brokers
 
-    producer.send = function () {
-      const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
+    const kafkaClusterIdPromise = getKafkaClusterId(this)
 
-      return innerAsyncResource.runInAsyncScope(() => {
-        if (!producerStartCh.hasSubscribers) {
-          return send.apply(this, arguments)
-        }
+    producer.send = function () {
+      const wrappedSend = (clusterId) => {
+        const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
 
-        try {
-          const { topic, messages = [] } = arguments[0]
-          for (const message of messages) {
-            if (message !== null && typeof message === 'object') {
-              message.headers = message.headers || {}
-            }
+        return innerAsyncResource.runInAsyncScope(() => {
+          if (!producerStartCh.hasSubscribers) {
+            return send.apply(this, arguments)
           }
-          producerStartCh.publish({ topic, messages, bootstrapServers })
-
-          const result = send.apply(this, arguments)
-
-          result.then(
-            innerAsyncResource.bind(res => {
-              producerFinishCh.publish(undefined)
-              producerCommitCh.publish(res)
-            }),
-            innerAsyncResource.bind(err => {
-              if (err) {
-                producerErrorCh.publish(err)
+
+          try {
+            const { topic, messages = [] } = arguments[0]
+            for (const message of messages) {
+              if (message !== null && typeof message === 'object') {
+                message.headers = message.headers || {}
               }
-              producerFinishCh.publish(undefined)
-            })
-          )
+            }
+            producerStartCh.publish({ topic, messages, bootstrapServers, clusterId })
 
-          return result
-        } catch (e) {
-          producerErrorCh.publish(e)
-          producerFinishCh.publish(undefined)
-          throw e
-        }
-      })
+            const result = send.apply(this, arguments)
+
+            result.then(
+              innerAsyncResource.bind(res => {
+                producerFinishCh.publish(undefined)
+                producerCommitCh.publish(res)
+              }),
+              innerAsyncResource.bind(err => {
+                if (err) {
+                  producerErrorCh.publish(err)
+                }
+                producerFinishCh.publish(undefined)
+              })
+            )
+
+            return result
+          } catch (e) {
+            producerErrorCh.publish(e)
+            producerFinishCh.publish(undefined)
+            throw e
+          }
+        })
+      }
+
+      if (!isPromise(kafkaClusterIdPromise)) {
+        // promise is already resolved
+        return wrappedSend(kafkaClusterIdPromise)
+      } else {
+        // promise is not resolved
+        return kafkaClusterIdPromise.then((clusterId) => {
+          return wrappedSend(clusterId)
+        })
+      }
     }
     return producer
   })
@@ -100,15 +114,17 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
       return createConsumer.apply(this, arguments)
     }
 
-    const eachMessageExtractor = (args) => {
+    const kafkaClusterIdPromise = getKafkaClusterId(this)
+
+    const eachMessageExtractor = (args, clusterId) => {
       const { topic, partition, message } = args[0]
-      return { topic, partition, message, groupId }
+      return { topic, partition, message, groupId, clusterId }
     }
 
-    const eachBatchExtractor = (args) => {
+    const eachBatchExtractor = (args, clusterId) => {
       const { batch } = args[0]
       const { topic, partition, messages } = batch
-      return { topic, partition, messages, groupId }
+      return { topic, partition, messages, groupId, clusterId }
     }
 
     const consumer = createConsumer.apply(this, arguments)
@@ -116,43 +132,53 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     consumer.on(consumer.events.COMMIT_OFFSETS, commitsFromEvent)
 
     const run = consumer.run
-
     const groupId = arguments[0].groupId
+
     consumer.run = function ({ eachMessage, eachBatch, ...runArgs }) {
-      eachMessage = wrapFunction(
-        eachMessage,
-        consumerStartCh,
-        consumerFinishCh,
-        consumerErrorCh,
-        eachMessageExtractor
-      )
-
-      eachBatch = wrapFunction(
-        eachBatch,
-        batchConsumerStartCh,
-        batchConsumerFinishCh,
-        batchConsumerErrorCh,
-        eachBatchExtractor
-      )
-
-      return run({
-        eachMessage,
-        eachBatch,
-        ...runArgs
-      })
+      const wrapConsume = (clusterId) => {
+        return run({
+          eachMessage: wrappedCallback(
+            eachMessage,
+            consumerStartCh,
+            consumerFinishCh,
+            consumerErrorCh,
+            eachMessageExtractor,
+            clusterId
+          ),
+          eachBatch: wrappedCallback(
+            eachBatch,
+            batchConsumerStartCh,
+            batchConsumerFinishCh,
+            batchConsumerErrorCh,
+            eachBatchExtractor,
+            clusterId
+          ),
+          ...runArgs
+        })
+      }
+
+      if (!isPromise(kafkaClusterIdPromise)) {
+        // promise is already resolved
+        return wrapConsume(kafkaClusterIdPromise)
+      } else {
+        // promise is not resolved
+        return kafkaClusterIdPromise.then((clusterId) => {
+          return wrapConsume(clusterId)
+        })
+      }
     }
-
     return consumer
   })
   return Kafka
 })
 
-const wrapFunction = (fn, startCh, finishCh, errorCh, extractArgs) => {
+const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs, clusterId) => {
   return typeof fn === 'function'
     ? function (...args) {
       const innerAsyncResource = new AsyncResource('bound-anonymous-fn')
       return innerAsyncResource.runInAsyncScope(() => {
-        const extractedArgs = extractArgs(args)
+        const extractedArgs = extractArgs(args, clusterId)
+
         startCh.publish(extractedArgs)
         try {
           const result = fn.apply(this, args)
@@ -179,3 +205,37 @@ const wrapFunction = (fn, startCh, finishCh, errorCh, extractArgs) => {
     }
     : fn
 }
+
+const getKafkaClusterId = (kafka) => {
+  if (kafka._ddKafkaClusterId) {
+    return kafka._ddKafkaClusterId
+  }
+
+  if (!kafka.admin) {
+    return null
+  }
+
+  const admin = kafka.admin()
+
+  if (!admin.describeCluster) {
+    return null
+  }
+
+  return admin.connect()
+    .then(() => {
+      return admin.describeCluster()
+    })
+    .then((clusterInfo) => {
+      const clusterId = clusterInfo?.clusterId
+      kafka._ddKafkaClusterId = clusterId
+      admin.disconnect()
+      return clusterId
+    })
+    .catch((error) => {
+      throw error
+    })
+}
+
+function isPromise (obj) {
+  return !!obj && (typeof obj === 'object' || typeof obj === 'function') && typeof obj.then === 'function'
+}

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

@@ -1,6 +1,6 @@
 const { addHook, channel } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
-const { getCallSites } = require('../../../dd-trace/src/plugins/util/test')
+const { getCallSites } = require('../../../dd-trace/src/plugins/util/stacktrace')
 const { testToStartLine } = require('./utils')
 
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')

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

@@ -280,12 +280,12 @@ function getOnFailHandler (isMain) {
 }
 
 function getOnTestRetryHandler () {
-  return function (test) {
+  return function (test, err) {
     const asyncResource = getTestAsyncResource(test)
     if (asyncResource) {
       const isFirstAttempt = test._currentRetry === 0
       asyncResource.runInAsyncScope(() => {
-        testRetryCh.publish(isFirstAttempt)
+        testRetryCh.publish({ isFirstAttempt, err })
       })
     }
     const key = getTestToArKey(test)

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

@@ -0,0 +1,37 @@
+'use strict'
+
+const shimmer = require('../../datadog-shimmer')
+const { channel, addHook, AsyncResource } = require('./helpers/instrument')
+
+const multerReadCh = channel('datadog:multer:read:finish')
+
+function publishRequestBodyAndNext (req, res, next) {
+  return shimmer.wrapFunction(next, next => function () {
+    if (multerReadCh.hasSubscribers && req) {
+      const abortController = new AbortController()
+      const body = req.body
+
+      multerReadCh.publish({ req, res, body, abortController })
+
+      if (abortController.signal.aborted) return
+    }
+
+    return next.apply(this, arguments)
+  })
+}
+
+addHook({
+  name: 'multer',
+  file: 'lib/make-middleware.js',
+  versions: ['^1.4.4-lts.1']
+}, makeMiddleware => {
+  return shimmer.wrapFunction(makeMiddleware, makeMiddleware => function () {
+    const middleware = makeMiddleware.apply(this, arguments)
+
+    return shimmer.wrapFunction(middleware, middleware => function wrapMulterMiddleware (req, res, next) {
+      const nextResource = new AsyncResource('bound-anonymous-fn')
+      arguments[2] = nextResource.bind(publishRequestBodyAndNext(req, res, next))
+      return middleware.apply(this, arguments)
+    })
+  })
+})

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

@@ -6,11 +6,14 @@ const {
   AsyncResource
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
+const semver = require('semver')
 
-addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connection => {
+addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Connection, version) => {
   const startCh = channel('apm:mysql2:query:start')
   const finishCh = channel('apm:mysql2:query:finish')
   const errorCh = channel('apm:mysql2:query:error')
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3')
 
   shimmer.wrap(Connection.prototype, 'addCommand', addCommand => function (cmd) {
     if (!startCh.hasSubscribers) return addCommand.apply(this, arguments)
@@ -28,6 +31,76 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
     return asyncResource.bind(addCommand, this).apply(this, arguments)
   })
 
+  shimmer.wrap(Connection.prototype, 'query', query => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return query.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const addCommand = this.addCommand
+      this.addCommand = function (cmd) { return cmd }
+
+      let queryCommand
+      try {
+        queryCommand = query.apply(this, arguments)
+      } finally {
+        this.addCommand = addCommand
+      }
+
+      cb = queryCommand.onResult
+
+      process.nextTick(() => {
+        if (cb) {
+          cb(abortController.signal.reason)
+        } else {
+          queryCommand.emit('error', abortController.signal.reason)
+        }
+
+        if (shouldEmitEndAfterQueryAbort) {
+          queryCommand.emit('end')
+        }
+      })
+
+      return queryCommand
+    }
+
+    return query.apply(this, arguments)
+  })
+
+  shimmer.wrap(Connection.prototype, 'execute', execute => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return execute.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const addCommand = this.addCommand
+      this.addCommand = function (cmd) { return cmd }
+
+      let result
+      try {
+        result = execute.apply(this, arguments)
+      } finally {
+        this.addCommand = addCommand
+      }
+
+      result?.onResult(abortController.signal.reason)
+
+      return result
+    }
+
+    return execute.apply(this, arguments)
+  })
+
   return Connection
 
   function bindExecute (cmd, execute, asyncResource) {
@@ -79,3 +152,149 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, Connec
     }, cmd))
   }
 })
+
+addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, version) => {
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3')
+
+  shimmer.wrap(Pool.prototype, 'query', query => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return query.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      const getConnection = this.getConnection
+      this.getConnection = function () {}
+
+      let queryCommand
+      try {
+        queryCommand = query.apply(this, arguments)
+      } finally {
+        this.getConnection = getConnection
+      }
+
+      process.nextTick(() => {
+        if (queryCommand.onResult) {
+          queryCommand.onResult(abortController.signal.reason)
+        } else {
+          queryCommand.emit('error', abortController.signal.reason)
+        }
+
+        if (shouldEmitEndAfterQueryAbort) {
+          queryCommand.emit('end')
+        }
+      })
+
+      return queryCommand
+    }
+
+    return query.apply(this, arguments)
+  })
+
+  shimmer.wrap(Pool.prototype, 'execute', execute => function (sql, values, cb) {
+    if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
+
+    if (typeof sql === 'object') sql = sql?.sql
+
+    if (!sql) return execute.apply(this, arguments)
+
+    const abortController = new AbortController()
+    startOuterQueryCh.publish({ sql, abortController })
+
+    if (abortController.signal.aborted) {
+      if (typeof values === 'function') {
+        cb = values
+      }
+
+      process.nextTick(() => {
+        cb(abortController.signal.reason)
+      })
+      return
+    }
+
+    return execute.apply(this, arguments)
+  })
+
+  return Pool
+})
+
+// PoolNamespace.prototype.query does not exist in mysql2<2.3.0
+addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['>=2.3.0'] }, PoolCluster => {
+  const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
+  const wrappedPoolNamespaces = new WeakSet()
+
+  shimmer.wrap(PoolCluster.prototype, 'of', of => function () {
+    const poolNamespace = of.apply(this, arguments)
+
+    if (startOuterQueryCh.hasSubscribers && !wrappedPoolNamespaces.has(poolNamespace)) {
+      shimmer.wrap(poolNamespace, 'query', query => function (sql, values, cb) {
+        if (typeof sql === 'object') sql = sql?.sql
+
+        if (!sql) return query.apply(this, arguments)
+
+        const abortController = new AbortController()
+        startOuterQueryCh.publish({ sql, abortController })
+
+        if (abortController.signal.aborted) {
+          const getConnection = this.getConnection
+          this.getConnection = function () {}
+
+          let queryCommand
+          try {
+            queryCommand = query.apply(this, arguments)
+          } finally {
+            this.getConnection = getConnection
+          }
+
+          process.nextTick(() => {
+            if (queryCommand.onResult) {
+              queryCommand.onResult(abortController.signal.reason)
+            } else {
+              queryCommand.emit('error', abortController.signal.reason)
+            }
+
+            queryCommand.emit('end')
+          })
+
+          return queryCommand
+        }
+
+        return query.apply(this, arguments)
+      })
+
+      shimmer.wrap(poolNamespace, 'execute', execute => function (sql, values, cb) {
+        if (typeof sql === 'object') sql = sql?.sql
+
+        if (!sql) return execute.apply(this, arguments)
+
+        const abortController = new AbortController()
+        startOuterQueryCh.publish({ sql, abortController })
+
+        if (abortController.signal.aborted) {
+          if (typeof values === 'function') {
+            cb = values
+          }
+
+          process.nextTick(() => {
+            cb(abortController.signal.reason)
+          })
+
+          return
+        }
+
+        return execute.apply(this, arguments)
+      })
+
+      wrappedPoolNamespaces.add(poolNamespace)
+    }
+
+    return poolNamespace
+  })
+
+  return PoolCluster
+})

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

@@ -3,8 +3,8 @@
 const { addHook } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
-const tracingChannel = require('dc-polyfill').tracingChannel
-const ch = tracingChannel('apm:openai:request')
+const dc = require('dc-polyfill')
+const ch = dc.tracingChannel('apm:openai:request')
 
 const V4_PACKAGE_SHIMS = [
   {

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

@@ -0,0 +1,127 @@
+const shimmer = require('../../datadog-shimmer')
+const { addHook } = require('./helpers/instrument')
+
+const dc = require('dc-polyfill')
+const serializeChannel = dc.channel('apm:protobufjs:serialize-start')
+const deserializeChannel = dc.channel('apm:protobufjs:deserialize-end')
+
+function wrapSerialization (messageClass) {
+  if (messageClass?.encode) {
+    shimmer.wrap(messageClass, 'encode', original => function () {
+      if (!serializeChannel.hasSubscribers) {
+        return original.apply(this, arguments)
+      }
+      serializeChannel.publish({ messageClass: this })
+      return original.apply(this, arguments)
+    })
+  }
+}
+
+function wrapDeserialization (messageClass) {
+  if (messageClass?.decode) {
+    shimmer.wrap(messageClass, 'decode', original => function () {
+      if (!deserializeChannel.hasSubscribers) {
+        return original.apply(this, arguments)
+      }
+      const result = original.apply(this, arguments)
+      deserializeChannel.publish({ messageClass: result })
+      return result
+    })
+  }
+}
+
+function wrapSetup (messageClass) {
+  if (messageClass?.setup) {
+    shimmer.wrap(messageClass, 'setup', original => function () {
+      const result = original.apply(this, arguments)
+
+      wrapSerialization(messageClass)
+      wrapDeserialization(messageClass)
+
+      return result
+    })
+  }
+}
+
+function wrapProtobufClasses (root) {
+  if (!root) {
+    return
+  }
+
+  if (root.decode) {
+    wrapSetup(root)
+  }
+
+  if (root.nestedArray) {
+    for (const subRoot of root.nestedArray) {
+      wrapProtobufClasses(subRoot)
+    }
+  }
+}
+
+function wrapReflection (protobuf) {
+  const reflectionMethods = [
+    {
+      target: protobuf.Root,
+      name: 'fromJSON'
+    },
+    {
+      target: protobuf.Type.prototype,
+      name: 'fromObject'
+    }
+  ]
+
+  reflectionMethods.forEach(method => {
+    shimmer.wrap(method.target, method.name, original => function () {
+      const result = original.apply(this, arguments)
+      if (result.nested) {
+        for (const type in result.nested) {
+          wrapSetup(result.nested[type])
+        }
+      }
+      if (result.$type) {
+        wrapSetup(result.$type)
+      }
+      return result
+    })
+  })
+}
+
+function isPromise (obj) {
+  return !!obj && (typeof obj === 'object' || typeof obj === 'function') && typeof obj.then === 'function'
+}
+
+addHook({
+  name: 'protobufjs',
+  versions: ['>=6.8.0']
+}, protobuf => {
+  shimmer.wrap(protobuf.Root.prototype, 'load', original => function () {
+    const result = original.apply(this, arguments)
+    if (isPromise(result)) {
+      return result.then(root => {
+        wrapProtobufClasses(root)
+        return root
+      })
+    } else {
+      // If result is not a promise, directly wrap the protobuf classes
+      wrapProtobufClasses(result)
+      return result
+    }
+  })
+
+  shimmer.wrap(protobuf.Root.prototype, 'loadSync', original => function () {
+    const root = original.apply(this, arguments)
+    wrapProtobufClasses(root)
+    return root
+  })
+
+  shimmer.wrap(protobuf, 'Type', Original => function () {
+    const typeInstance = new Original(...arguments)
+    wrapSetup(typeInstance)
+    return typeInstance
+  })
+
+  wrapReflection(protobuf)
+
+  return protobuf
+})