dd-trace 5.103.0 → 5.105.0

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

@@ -8,6 +8,7 @@ const {
   addHook,
 } = require('./helpers/instrument')
 const {
+  brokerSupportsMessageHeaders,
   clientToCluster,
   cloneMessages,
 } = require('./helpers/kafka')
@@ -26,7 +27,7 @@ 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) {
@@ -61,39 +62,64 @@ 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 originalSend = producer.send
+    const originalSendBatch = producer.sendBatch
     const bootstrapServers = this._brokers
     const cluster = clientToCluster.get(producer)
 
-    producer.send = function (...args) {
-      if (!producerStartCh.hasSubscribers) {
-        return originalSend.apply(this, args)
+    let disableHeaderInjection = false
+
+    let refreshHeaderSupport = () => {
+      if (!brokerSupportsMessageHeaders(cluster?.brokerPool)) {
+        disableHeaderInjection = true
+        refreshHeaderSupport = noop
+        log.info('kafkajs broker negotiated Produce <v3; tracer header injection disabled.')
       }
+    }
 
-      // Fast path: kafkajs has fetched metadata, so clusterId is already on
-      // the broker pool.
+    /**
+     * Resolve the negotiated clusterId once and hand it to `call`. Fast path reads
+     * `cluster.brokerPool.metadata` synchronously when kafkajs already fetched it.
+     * Slow path primes `refreshMetadataIfNecessary`, which `sharedPromiseTo`
+     * deduplicates with kafkajs's own internal fetch so total latency is unchanged.
+     *
+     * @param {(clusterId: string | undefined) => Promise<unknown>} call
+     */
+    const withClusterId = (call) => {
       const metadata = cluster?.brokerPool?.metadata
       if (metadata) {
-        return runSend.call(this, args, metadata.clusterId)
+        refreshHeaderSupport()
+        return call(metadata.clusterId)
       }
-
-      // 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 call()
       }
       return cluster.refreshMetadataIfNecessary().then(
-        () => runSend.call(this, args, cluster.brokerPool?.metadata?.clusterId),
-        () => runSend.call(this, args)
+        () => {
+          refreshHeaderSupport()
+          return call(cluster.brokerPool?.metadata?.clusterId)
+        },
+        () => call()
       )
     }
 
+    producer.send = function (...args) {
+      if (!producerStartCh.hasSubscribers) {
+        return originalSend.apply(this, args)
+      }
+      return withClusterId((clusterId) => runSend.call(this, args, clusterId))
+    }
+
+    producer.sendBatch = function (...args) {
+      if (!producerStartCh.hasSubscribers) {
+        return originalSendBatch.apply(this, args)
+      }
+      return withClusterId((clusterId) => runSendBatch.call(this, args, clusterId))
+    }
+
     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
@@ -125,8 +151,13 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
             (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') {
-                  disabledHeaderWeakSet.add(producer)
+                  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.'
@@ -147,6 +178,98 @@ addHook({ name: 'kafkajs', file: 'src/index.js', versions: ['>=1.4'] }, (BaseKaf
       })
     }
 
+    function runSendBatch (args, clusterId) {
+      const arg0 = args[0]
+      const inputTopicMessages = Array.isArray(arg0?.topicMessages) ? arg0.topicMessages : []
+      if (inputTopicMessages.length === 0) {
+        return originalSendBatch.apply(this, args)
+      }
+
+      // One ctx per topicMessages entry — kafkajs implements `send` as a single-entry
+      // `sendBatch` (`producer/messageProducer.js`), so one span per entry is the same
+      // unit `send` already produces. Cloning only happens for valid arrays so kafkajs
+      // still sees and rejects a caller's malformed `messages` field.
+      const outputEntries = new Array(inputTopicMessages.length)
+      const ctxList = []
+      let cloned = false
+      for (let i = 0; i < inputTopicMessages.length; i++) {
+        const entry = inputTopicMessages[i]
+        const topic = entry?.topic
+        const rawMessages = entry?.messages
+        let entryMessages = rawMessages
+        if (Array.isArray(rawMessages) && rawMessages.length > 0) {
+          entryMessages = cloneMessages(rawMessages, !disableHeaderInjection)
+          outputEntries[i] = { ...entry, messages: entryMessages }
+          cloned = true
+        } else {
+          outputEntries[i] = entry
+        }
+        ctxList.push({
+          bootstrapServers,
+          clusterId,
+          disableHeaderInjection,
+          messages: Array.isArray(entryMessages) ? entryMessages : [],
+          topic,
+        })
+      }
+      if (cloned) {
+        args[0] = { ...arg0, topicMessages: outputEntries }
+      }
+
+      for (const ctx of ctxList) {
+        producerStartCh.runStores(ctx, noop)
+      }
+
+      let result
+      try {
+        result = originalSendBatch.apply(this, args)
+      } catch (error) {
+        failProduceBatch(ctxList, error)
+        throw error
+      }
+
+      result.then(
+        (res) => {
+          for (const ctx of ctxList) {
+            ctx.result = res
+            producerFinishCh.publish(ctx)
+          }
+          // kafkajs returns a single aggregated response covering every topic;
+          // commit fires once so the plugin's `setOffset` loop runs once per
+          // entry of the response, not once per span.
+          producerCommitCh.publish(ctxList[0])
+        },
+        (error) => failProduceBatch(ctxList, error)
+      )
+
+      return result
+    }
+
+    /**
+     * Tag every open ctx with the shared error, then publish error + finish so the
+     * plugin closes each span. The mixed-version safety net (broker advertised
+     * Produce v3+ but the leader rejected the headers) fires at most once per
+     * failed batch and short-circuits subsequent sends to the disabled path.
+     *
+     * @param {Array<object>} ctxList
+     * @param {Error} error
+     */
+    function failProduceBatch (ctxList, error) {
+      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.'
+        )
+      }
+      for (const ctx of ctxList) {
+        ctx.error = error
+        producerErrorCh.publish(ctx)
+        producerFinishCh.publish(ctx)
+      }
+    }
+
     return producer
   })
 

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)