dd-trace 5.102.0 → 5.104.0

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

@@ -7,6 +7,11 @@ const {
   channel,
   addHook,
 } = require('./helpers/instrument')
+const {
+  brokerSupportsMessageHeaders,
+  clientToCluster,
+  cloneMessages,
+} = require('./helpers/kafka')
 
 const producerStartCh = channel('apm:kafkajs:produce:start')
 const producerCommitCh = channel('apm:kafkajs:produce:commit')
@@ -22,7 +27,27 @@ const batchConsumerStartCh = channel('apm:kafkajs:consume-batch:start')
 const batchConsumerFinishCh = channel('apm:kafkajs:consume-batch:finish')
 const batchConsumerErrorCh = channel('apm:kafkajs:consume-batch:error')
 
-const disabledHeaderWeakSet = new WeakSet()
+const noop = () => {}
+
+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 {
@@ -34,79 +59,113 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     }
   }
 
-  shimmer.wrap(Kafka.prototype, 'producer', createProducer => function (...args) {
-    const producer = createProducer.apply(this, args)
-    const send = producer.send
+  shimmer.wrap(Kafka.prototype, 'producer', createProducer => function () {
+    const producer = createProducer.apply(this, arguments)
+    const originalSend = producer.send
     const bootstrapServers = this._brokers
+    const cluster = clientToCluster.get(producer)
+
+    let disableHeaderInjection = false
 
-    const kafkaClusterIdPromise = getKafkaClusterId(this)
+    let refreshHeaderSupport = () => {
+      if (!brokerSupportsMessageHeaders(cluster?.brokerPool)) {
+        disableHeaderInjection = true
+        refreshHeaderSupport = noop
+        log.info('kafkajs broker negotiated Produce <v3; tracer header injection disabled.')
+      }
+    }
 
     producer.send = function (...args) {
-      const wrappedSend = (clusterId) => {
-        const { topic, messages = [] } = args[0]
-
-        const ctx = {
-          bootstrapServers,
-          clusterId,
-          disableHeaderInjection: disabledHeaderWeakSet.has(producer),
-          messages,
-          topic,
-        }
+      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 versions and clusterId
+      // are already on the broker pool.
+      const metadata = cluster?.brokerPool?.metadata
+      if (metadata) {
+        refreshHeaderSupport()
+        return runSend.call(this, args, metadata.clusterId)
+      }
 
-        return producerStartCh.runStores(ctx, () => {
-          try {
-            const result = send.apply(this, args)
-            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)
-              })
-
-            return result
-          } catch (e) {
-            ctx.error = e
-            producerErrorCh.publish(ctx)
-            producerFinishCh.publish(ctx)
-            throw e
-          }
-        })
+      // 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(
+        () => {
+          refreshHeaderSupport()
+          return runSend.call(this, args, cluster.brokerPool?.metadata?.clusterId)
+        },
+        () => runSend.call(this, args)
+      )
+    }
+
+    function runSend (args, clusterId) {
+      const arg0 = args[0]
+      const topic = arg0?.topic
+      const inputMessages = Array.isArray(arg0?.messages) ? arg0.messages : []
+
+      // 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) {
+                // Safety net for mixed-version clusters where the seed
+                // broker advertised Produce v3+ but the leader we shipped to
+                // could not parse the headers, surfacing as
+                // KafkaJSProtocolError UNKNOWN (server error code -1).
+                if (error.name === 'KafkaJSProtocolError' && error.type === 'UNKNOWN') {
+                  disableHeaderInjection = true
+                  refreshHeaderSupport = noop
+                  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
   })
 
@@ -115,24 +174,26 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
       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 eachMessageExtractor = (args, clusterId) => {
+    const readClusterId = () => cluster?.brokerPool?.metadata?.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, args)
-
     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 +202,7 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
             partition,
             offset,
             topic,
-            clusterId: resolvedClusterId,
+            clusterId,
           })
         }
       }
@@ -149,118 +210,67 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
     })
 
     const run = consumer.run
-    const groupId = args[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/mariadb.js

@@ -93,7 +93,7 @@ function createWrapQueryCallback (options) {
       }
 
       if (typeof cb === 'function') {
-        arguments[arguments.length - 1] = shimmer.wrapFunction(cb, wrapper)
+        arguments[arguments.length - 1] = shimmer.wrapCallback(cb, wrapper)
       } else {
         arguments.length += 1
         arguments[arguments.length - 1] = wrapper()

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

@@ -23,7 +23,8 @@ addHook({ name: 'memcached', versions: ['>=2.2'] }, Memcached => {
 
       const ctx = { client, server, query }
       startCh.runStores(ctx, () => {
-        query.callback = shimmer.wrapFunction(query.callback, callback => function (err) {
+        if (typeof query.callback !== 'function') return
+        query.callback = shimmer.wrapCallback(query.callback, callback => function (err) {
           if (err) {
             ctx.error = err
             errorCh.publish(ctx)