dd-trace 5.101.0 → 5.103.0

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

@@ -7,6 +7,10 @@ const {
   channel,
   addHook,
 } = require('./helpers/instrument')
+const {
+  clientToCluster,
+  cloneMessages,
+} = require('./helpers/kafka')
 
 const producerStartCh = channel('apm:kafkajs:produce:start')
 const producerCommitCh = channel('apm:kafkajs:produce:commit')
@@ -24,6 +28,26 @@ const batchConsumerErrorCh = channel('apm:kafkajs:consume-batch:error')
 
 const disabledHeaderWeakSet = new WeakSet()
 
+addHook({ name: 'kafkajs', file: 'src/producer/index.js', versions: ['>=1.4'] }, (createProducer) =>
+  shimmer.wrapFunction(createProducer, original => function wrappedCreateProducer (params) {
+    const producer = original(params)
+    if (params?.cluster) {
+      clientToCluster.set(producer, params.cluster)
+    }
+    return producer
+  })
+)
+
+addHook({ name: 'kafkajs', file: 'src/consumer/index.js', versions: ['>=1.4'] }, (createConsumer) =>
+  shimmer.wrapFunction(createConsumer, original => function wrappedCreateConsumer (params) {
+    const consumer = original(params)
+    if (params?.cluster) {
+      clientToCluster.set(consumer, params.cluster)
+    }
+    return consumer
+  })
+)
+
 addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKafka) => {
   class Kafka extends BaseKafka {
     constructor (options) {
@@ -36,103 +60,121 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
 
   shimmer.wrap(Kafka.prototype, 'producer', createProducer => function () {
     const producer = createProducer.apply(this, arguments)
-    const send = producer.send
+    const originalSend = producer.send
     const bootstrapServers = this._brokers
+    const cluster = clientToCluster.get(producer)
 
-    const kafkaClusterIdPromise = getKafkaClusterId(this)
-
-    producer.send = function () {
-      const wrappedSend = (clusterId) => {
-        const { topic, messages = [] } = arguments[0]
-
-        const ctx = {
-          bootstrapServers,
-          clusterId,
-          disableHeaderInjection: disabledHeaderWeakSet.has(producer),
-          messages,
-          topic,
-        }
+    producer.send = function (...args) {
+      if (!producerStartCh.hasSubscribers) {
+        return originalSend.apply(this, args)
+      }
 
-        for (const message of messages) {
-          if (message !== null && typeof message === 'object' && !ctx.disableHeaderInjection) {
-            message.headers = message.headers || {}
-          }
-        }
+      // Fast path: kafkajs has fetched metadata, so clusterId is already on
+      // the broker pool.
+      const metadata = cluster?.brokerPool?.metadata
+      if (metadata) {
+        return runSend.call(this, args, metadata.clusterId)
+      }
 
-        return producerStartCh.runStores(ctx, () => {
-          try {
-            const result = send.apply(this, arguments)
-            result.then(
-              (res) => {
-                ctx.result = res
-                producerFinishCh.publish(ctx)
-                producerCommitCh.publish(ctx)
-              },
-              (err) => {
-                ctx.error = err
-                if (err) {
-                  // Fixes bug where we would inject message headers for kafka brokers that don't support headers
-                  // (version <0.11). On the error, we disable header injection.
-                  // Unfortunately the error name / type is not more specific.
-                  // This approach is implemented by other tracers as well.
-                  if (err.name === 'KafkaJSProtocolError' && err.type === 'UNKNOWN') {
-                    disabledHeaderWeakSet.add(producer)
-                    log.error(
-                      // eslint-disable-next-line @stylistic/max-len
-                      'Kafka Broker responded with UNKNOWN_SERVER_ERROR (-1). Please look at broker logs for more information. Tracer message header injection for Kafka is disabled.'
-                    )
-                  }
-                  producerErrorCh.publish(err)
-                }
-                producerFinishCh.publish(ctx)
-              })
+      // Slow path, taken at most once per producer connect cycle. Prime the
+      // metadata fetch kafkajs's send would do internally a few stack frames
+      // later. `sharedPromiseTo` collapses our call and kafkajs's call into a
+      // single round trip, so total latency is unchanged.
+      if (typeof cluster?.refreshMetadataIfNecessary !== 'function') {
+        return runSend.call(this, args)
+      }
+      return cluster.refreshMetadataIfNecessary().then(
+        () => runSend.call(this, args, cluster.brokerPool?.metadata?.clusterId),
+        () => runSend.call(this, args)
+      )
+    }
 
-            return result
-          } catch (e) {
-            ctx.error = e
-            producerErrorCh.publish(ctx)
-            producerFinishCh.publish(ctx)
-            throw e
-          }
-        })
+    function runSend (args, clusterId) {
+      const arg0 = args[0]
+      const topic = arg0?.topic
+      const inputMessages = Array.isArray(arg0?.messages) ? arg0.messages : []
+      const disableHeaderInjection = disabledHeaderWeakSet.has(producer)
+
+      // Hand kafkajs and the plugin a shallow clone so injection writes to
+      // tracer-owned objects instead of the caller's. With injection
+      // disabled the clone must not seed `headers: {}` either: brokers that
+      // reject any header field cannot recover otherwise.
+      let messages = inputMessages
+      if (inputMessages.length > 0) {
+        messages = cloneMessages(inputMessages, !disableHeaderInjection)
+        args[0] = { ...arg0, messages }
       }
 
-      if (isPromise(kafkaClusterIdPromise)) {
-        // promise is not resolved
-        return kafkaClusterIdPromise.then((clusterId) => {
-          return wrappedSend(clusterId)
-        })
+      const ctx = {
+        bootstrapServers,
+        clusterId,
+        disableHeaderInjection,
+        messages,
+        topic,
       }
 
-      // promise is already resolved
-      return wrappedSend(kafkaClusterIdPromise)
+      return producerStartCh.runStores(ctx, () => {
+        try {
+          const result = originalSend.apply(this, args)
+          result.then(
+            (res) => {
+              ctx.result = res
+              producerFinishCh.publish(ctx)
+              producerCommitCh.publish(ctx)
+            },
+            (error) => {
+              ctx.error = error
+              if (error) {
+                if (error.name === 'KafkaJSProtocolError' && error.type === 'UNKNOWN') {
+                  disabledHeaderWeakSet.add(producer)
+                  log.error(
+                    // eslint-disable-next-line @stylistic/max-len
+                    'Kafka Broker responded with UNKNOWN_SERVER_ERROR (-1). Please look at broker logs for more information. Tracer message header injection for Kafka is disabled.'
+                  )
+                }
+                producerErrorCh.publish(error)
+              }
+              producerFinishCh.publish(ctx)
+            }
+          )
+          return result
+        } catch (error) {
+          ctx.error = error
+          producerErrorCh.publish(ctx)
+          producerFinishCh.publish(ctx)
+          throw error
+        }
+      })
     }
+
     return producer
   })
 
-  shimmer.wrap(Kafka.prototype, 'consumer', createConsumer => function () {
+  shimmer.wrap(Kafka.prototype, 'consumer', createConsumer => function (...args) {
     if (!consumerStartCh.hasSubscribers) {
-      return createConsumer.apply(this, arguments)
+      return createConsumer.apply(this, args)
     }
 
-    const kafkaClusterIdPromise = getKafkaClusterId(this)
-    let resolvedClusterId = null
+    const consumer = createConsumer.apply(this, arguments)
+    const cluster = clientToCluster.get(consumer)
+    const groupId = arguments[0].groupId
+
+    const readClusterId = () => cluster?.brokerPool?.metadata?.clusterId
 
-    const eachMessageExtractor = (args, clusterId) => {
+    const eachMessageExtractor = (args) => {
       const { topic, partition, message } = args[0]
-      return { topic, partition, message, groupId, clusterId }
+      return { topic, partition, message, groupId, clusterId: readClusterId() }
     }
 
-    const eachBatchExtractor = (args, clusterId) => {
+    const eachBatchExtractor = (args) => {
       const { batch } = args[0]
       const { topic, partition, messages } = batch
-      return { topic, partition, messages, groupId, clusterId }
+      return { topic, partition, messages, groupId, clusterId: readClusterId() }
     }
 
-    const consumer = createConsumer.apply(this, arguments)
-
     consumer.on(consumer.events.COMMIT_OFFSETS, (event) => {
       const { payload: { groupId: commitGroupId, topics } } = event
+      const clusterId = readClusterId()
       const commitList = []
       for (const { topic, partitions } of topics) {
         for (const { partition, offset } of partitions) {
@@ -141,7 +183,7 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
             partition,
             offset,
             topic,
-            clusterId: resolvedClusterId,
+            clusterId,
           })
         }
       }
@@ -149,118 +191,67 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     })
 
     const run = consumer.run
-    const groupId = arguments[0].groupId
 
     consumer.run = function ({ eachMessage, eachBatch, ...runArgs }) {
-      const wrapConsume = (clusterId) => {
-        // In kafkajs COMMIT_OFFSETS always happens in the context of one synchronous run
-        // So this will always reference a correct cluster id
-        resolvedClusterId = 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 not resolved
-        return kafkaClusterIdPromise.then((clusterId) => {
-          return wrapConsume(clusterId)
-        })
-      }
-
-      // promise is already resolved
-      return wrapConsume(kafkaClusterIdPromise)
+      return run({
+        eachMessage: wrappedCallback(
+          eachMessage,
+          consumerStartCh,
+          consumerFinishCh,
+          consumerErrorCh,
+          eachMessageExtractor
+        ),
+        eachBatch: wrappedCallback(
+          eachBatch,
+          batchConsumerStartCh,
+          batchConsumerFinishCh,
+          batchConsumerErrorCh,
+          eachBatchExtractor
+        ),
+        ...runArgs,
+      })
     }
+
     return consumer
   })
+
   return Kafka
 })
 
-const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs, clusterId) => {
-  return typeof fn === 'function'
-    ? function (...args) {
-      const extractedArgs = extractArgs(args, clusterId)
-      const ctx = {
-        extractedArgs,
-      }
+const wrappedCallback = (fn, startCh, finishCh, errorCh, extractArgs) => {
+  if (typeof fn !== 'function') return fn
+  return function (...args) {
+    const ctx = {
+      extractedArgs: extractArgs(args),
+    }
 
-      return startCh.runStores(ctx, () => {
-        try {
-          const result = fn.apply(this, args)
-          if (result && typeof result.then === 'function') {
-            result.then(
-              (res) => {
-                ctx.result = res
-                finishCh.publish(ctx)
-              },
-              (err) => {
-                ctx.error = err
-                if (err) {
-                  errorCh.publish(ctx)
-                }
-                finishCh.publish(ctx)
-              })
-          } else {
-            finishCh.publish(ctx)
-          }
-          return result
-        } catch (e) {
-          ctx.error = e
-          errorCh.publish(ctx)
+    return startCh.runStores(ctx, () => {
+      try {
+        const result = fn.apply(this, args)
+        if (result && typeof result.then === 'function') {
+          result.then(
+            (res) => {
+              ctx.result = res
+              finishCh.publish(ctx)
+            },
+            (error) => {
+              ctx.error = error
+              if (error) {
+                errorCh.publish(ctx)
+              }
+              finishCh.publish(ctx)
+            }
+          )
+        } else {
           finishCh.publish(ctx)
-          throw e
         }
-      })
-    }
-    : 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
+        return result
+      } catch (error) {
+        ctx.error = error
+        errorCh.publish(ctx)
+        finishCh.publish(ctx)
+        throw error
+      }
     })
-}
-
-function isPromise (obj) {
-  return !!obj && (typeof obj === 'object' || typeof obj === 'function') && typeof obj.then === 'function'
+  }
 }

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

@@ -32,41 +32,41 @@ addHook({
   versions: ['>=2'],
   file: 'lib/knex-builder/Knex.js',
 }, Knex => {
-  shimmer.wrap(Knex.Client.prototype, 'raw', raw => function () {
+  shimmer.wrap(Knex.Client.prototype, 'raw', raw => function (...args) {
     if (!startRawQueryCh.hasSubscribers) {
-      return raw.apply(this, arguments)
+      return raw.apply(this, args)
     }
 
-    const sql = arguments[0]
+    const sql = args[0]
 
     // Skip query done by Knex to get the value used for undefined
     if (sql === 'DEFAULT') {
-      return raw.apply(this, arguments)
+      return raw.apply(this, args)
     }
 
     const context = { sql, dialect: this.dialect }
     return startRawQueryCh.runStores(context, () => {
-      const rawResult = raw.apply(this, arguments)
-      shimmer.wrap(rawResult, 'then', originalThen => function () {
+      const rawResult = raw.apply(this, args)
+      shimmer.wrap(rawResult, 'then', originalThen => function (...args) {
         return rawQuerySubscribes.runStores(context, () => {
-          arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
-          if (arguments[1]) arguments[1] = wrapCallbackWithFinish(arguments[1], finish, context)
+          args[0] = wrapCallbackWithFinish(args[0], finish, context)
+          if (args[1]) args[1] = wrapCallbackWithFinish(args[1], finish, context)
 
-          const originalThenResult = originalThen.apply(this, arguments)
+          const originalThenResult = originalThen.apply(this, args)
 
-          shimmer.wrap(originalThenResult, 'catch', originalCatch => function () {
-            arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
-            return originalCatch.apply(this, arguments)
+          shimmer.wrap(originalThenResult, 'catch', originalCatch => function (...args) {
+            args[0] = wrapCallbackWithFinish(args[0], finish, context)
+            return originalCatch.apply(this, args)
           })
 
           return originalThenResult
         })
       })
 
-      shimmer.wrap(rawResult, 'asCallback', originalAsCallback => function () {
+      shimmer.wrap(rawResult, 'asCallback', originalAsCallback => function (...args) {
         return rawQuerySubscribes.runStores(context, () => {
-          arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
-          return originalAsCallback.apply(this, arguments)
+          args[0] = wrapCallbackWithFinish(args[0], finish, context)
+          return originalAsCallback.apply(this, args)
         })
       })
 
@@ -80,7 +80,7 @@ addHook({
 function wrapCallbackWithFinish (callback, finish, context) {
   if (typeof callback !== 'function') return callback
 
-  return shimmer.wrapFunction(callback, callback => function () {
-    finish(context, () => callback.apply(this, arguments))
+  return shimmer.wrapFunction(callback, callback => function (...args) {
+    finish(context, () => callback.apply(this, args))
   })
 }

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

@@ -14,8 +14,8 @@ const routeChannel = channel('apm:koa:request:route')
 const originals = new WeakMap()
 
 function wrapCallback (callback) {
-  return function callbackWithTrace () {
-    const handleRequest = callback.apply(this, arguments)
+  return function callbackWithTrace (...args) {
+    const handleRequest = callback.apply(this, args)
 
     if (typeof handleRequest !== 'function') return handleRequest
 
@@ -28,14 +28,14 @@ function wrapCallback (callback) {
 }
 
 function wrapUse (use) {
-  return function useWithTrace () {
-    const result = use.apply(this, arguments)
+  return function useWithTrace (...args) {
+    const result = use.apply(this, args)
 
-    if (!Array.isArray(this.middleware)) return result
+    if (Array.isArray(this.middleware)) {
+      const fn = this.middleware.pop()
 
-    const fn = this.middleware.pop()
-
-    this.middleware.push(wrapMiddleware(fn))
+      this.middleware.push(wrapMiddleware(fn))
+    }
 
     return result
   }
@@ -54,8 +54,8 @@ function wrapRegister (register) {
 }
 
 function wrapRouterUse (use) {
-  return function useWithTrace () {
-    const router = use.apply(this, arguments)
+  return function useWithTrace (...args) {
+    const router = use.apply(this, args)
 
     for (const layer of router.stack) {
       wrapStack(layer)
@@ -144,10 +144,10 @@ function fulfill (ctx, error) {
 }
 
 function wrapNext (req, next) {
-  return shimmer.wrapFunction(next, next => function () {
+  return shimmer.wrapFunction(next, next => function (...args) {
     nextChannel.publish({ req })
 
-    return next.apply(this, arguments)
+    return next.apply(this, args)
   })
 }
 

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

@@ -71,12 +71,12 @@ addHook({ name: 'ldapjs', versions: ['>=2'] }, ldapjs => {
     return search.apply(this, arguments)
   })
 
-  shimmer.wrap(ldapjs.Client.prototype, '_send', _send => function () {
-    const callbackIndex = getCallbackArgIndex(arguments)
+  shimmer.wrap(ldapjs.Client.prototype, '_send', _send => function (...args) {
+    const callbackIndex = getCallbackArgIndex(args)
     if (callbackIndex > -1) {
-      const callback = arguments[callbackIndex]
+      const callback = args[callbackIndex]
       // eslint-disable-next-line n/handle-callback-err
-      arguments[callbackIndex] = shimmer.wrapFunction(callback, callback => function (err, corkedEmitter) {
+      args[callbackIndex] = shimmer.wrapFunction(callback, callback => function (err, corkedEmitter) {
         if (corkedEmitter !== null && typeof corkedEmitter === 'object' && typeof corkedEmitter.on === 'function') {
           wrapEmitter(corkedEmitter)
         }
@@ -84,7 +84,7 @@ addHook({ name: 'ldapjs', versions: ['>=2'] }, ldapjs => {
       })
     }
 
-    return _send.apply(this, arguments)
+    return _send.apply(this, args)
   })
 
   shimmer.wrap(ldapjs.Client.prototype, 'bind', bind => function (dn, password, controls, callback) {

package/packages/datadog-instrumentations/src/light-my-request.js

@@ -69,8 +69,8 @@ function wrapDispatchFunc (dispatchFunc) {
     // Also wrap end() as fallback
     const originalEnd = res.end
     if (originalEnd) {
-      res.end = function wrappedEnd () {
-        const result = originalEnd.apply(this, arguments)
+      res.end = function wrappedEnd (...args) {
+        const result = originalEnd.apply(this, args)
         // Trigger finish if events don't fire
         setImmediate(onFinish)
         return result

package/packages/datadog-instrumentations/src/limitd-client.js

@@ -4,10 +4,10 @@ const shimmer = require('../../datadog-shimmer')
 const { addHook, AsyncResource } = require('./helpers/instrument')
 
 function wrapRequest (original) {
-  return function () {
-    const id = arguments.length - 1
-    arguments[id] = AsyncResource.bind(arguments[id])
-    return original.apply(this, arguments)
+  return function (...args) {
+    const id = args.length - 1
+    args[id] = AsyncResource.bind(args[id])
+    return original.apply(this, args)
   }
 }
 

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

@@ -12,13 +12,13 @@ addHook({ name: 'lodash', versions: ['>=4'] }, lodash => {
     lodash,
     instrumentedLodashFn,
     lodashFn => {
-      return function () {
+      return function (...args) {
         if (!lodashOperationCh.hasSubscribers) {
-          return lodashFn.apply(this, arguments)
+          return lodashFn.apply(this, args)
         }
 
-        const result = lodashFn.apply(this, arguments)
-        const message = { operation: lodashFn.name, arguments, result }
+        const result = lodashFn.apply(this, args)
+        const message = { operation: lodashFn.name, arguments: args, result }
         lodashOperationCh.publish(message)
 
         return message.result

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

@@ -12,13 +12,13 @@ const errorCh = channel('apm:mariadb:query:error')
 const skipCh = channel('apm:mariadb:pool:skip')
 
 function wrapCommandStart (start, ctx) {
-  return shimmer.wrapFunction(start, start => function () {
-    if (!startCh.hasSubscribers) return start.apply(this, arguments)
+  return shimmer.wrapFunction(start, start => function (...args) {
+    if (!startCh.hasSubscribers) return start.apply(this, args)
 
     const { reject, resolve } = this
     shimmer.wrap(this, 'resolve', function wrapResolve () {
-      return function () {
-        return finishCh.runStores(ctx, resolve, this, ...arguments)
+      return function (...args) {
+        return finishCh.runStores(ctx, resolve, this, ...args)
       }
     })
 
@@ -32,7 +32,7 @@ function wrapCommandStart (start, ctx) {
       }
     })
 
-    return startCh.runStores(ctx, start, this, ...arguments)
+    return startCh.runStores(ctx, start, this, ...args)
   })
 }
 
@@ -128,25 +128,25 @@ function wrapPoolBase (PoolBase) {
 // so instead we just skip instrumentation completely to avoid memory leaks
 // and/or orphan spans.
 function wrapPoolMethod (createConnection) {
-  return function () {
-    return skipCh.runStores({}, createConnection, this, ...arguments)
+  return function (...args) {
+    return skipCh.runStores({}, createConnection, this, ...args)
   }
 }
 
 function wrapPoolGetConnectionMethod (getConnection) {
-  return function wrappedGetConnection () {
-    const cb = arguments[arguments.length - 1]
-    if (typeof cb !== 'function') return getConnection.apply(this, arguments)
+  return function wrappedGetConnection (...args) {
+    const cb = args[args.length - 1]
+    if (typeof cb !== 'function') return getConnection.apply(this, args)
 
     const ctx = {}
 
-    arguments[arguments.length - 1] = function () {
-      return connectionFinishCh.runStores(ctx, cb, this, ...arguments)
+    args[args.length - 1] = function (...args) {
+      return connectionFinishCh.runStores(ctx, cb, this, ...args)
     }
 
     connectionStartCh.publish(ctx)
 
-    return getConnection.apply(this, arguments)
+    return getConnection.apply(this, args)
   }
 }