dd-trace 5.51.0 → 5.53.0

package/packages/datadog-instrumentations/src/confluentinc-kafka-javascript.js

@@ -0,0 +1,406 @@
+'use strict'
+
+const {
+  addHook,
+  channel,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const log = require('../../dd-trace/src/log')
+
+// Create channels for Confluent Kafka JavaScript
+const channels = {
+  producerStart: channel('apm:@confluentinc/kafka-javascript:produce:start'),
+  producerFinish: channel('apm:@confluentinc/kafka-javascript:produce:finish'),
+  producerError: channel('apm:@confluentinc/kafka-javascript:produce:error'),
+  producerCommit: channel('apm:@confluentinc/kafka-javascript:produce:commit'),
+  consumerStart: channel('apm:@confluentinc/kafka-javascript:consume:start'),
+  consumerFinish: channel('apm:@confluentinc/kafka-javascript:consume:finish'),
+  consumerError: channel('apm:@confluentinc/kafka-javascript:consume:error'),
+  consumerCommit: channel('apm:@confluentinc/kafka-javascript:consume:commit'),
+
+  // batch operations
+  batchConsumerStart: channel('apm:@confluentinc/kafka-javascript:consume-batch:start'),
+  batchConsumerFinish: channel('apm:@confluentinc/kafka-javascript:consume-batch:finish'),
+  batchConsumerError: channel('apm:@confluentinc/kafka-javascript:consume-batch:error'),
+  batchConsumerCommit: channel('apm:@confluentinc/kafka-javascript:consume-batch:commit')
+}
+
+const disabledHeaderWeakSet = new WeakSet()
+
+// we need to store the offset per partition per topic for the consumer to track offsets for DSM
+const latestConsumerOffsets = new Map()
+
+// Customize the instrumentation for Confluent Kafka JavaScript
+addHook({ name: '@confluentinc/kafka-javascript', versions: ['>=1.0.0'] }, (module) => {
+  // Hook native module classes first
+  instrumentBaseModule(module)
+
+  // Then hook KafkaJS if it exists
+  if (module.KafkaJS) {
+    instrumentKafkaJS(module.KafkaJS)
+  }
+
+  return module
+})
+
+function instrumentBaseModule (module) {
+  // Helper function to wrap producer classes
+  function wrapProducerClass (ProducerClass, className) {
+    return shimmer.wrap(module, className, function wrapProducer (Original) {
+      return function wrappedProducer () {
+        const producer = new Original(...arguments)
+
+        // Hook the produce method
+        if (typeof producer?.produce === 'function') {
+          shimmer.wrap(producer, 'produce', function wrapProduce (produce) {
+            return function wrappedProduce (topic, partition, message, key, timestamp, opaque) {
+              if (!channels.producerStart.hasSubscribers) {
+                return produce.apply(this, arguments)
+              }
+
+              const brokers = this.globalConfig?.['bootstrap.servers']
+
+              const asyncResource = new AsyncResource('bound-anonymous-fn')
+              return asyncResource.runInAsyncScope(() => {
+                try {
+                  channels.producerStart.publish({
+                    topic,
+                    messages: [{ key, value: message }],
+                    bootstrapServers: brokers
+                  })
+
+                  const result = produce.apply(this, arguments)
+
+                  channels.producerCommit.publish(undefined)
+                  channels.producerFinish.publish(undefined)
+                  return result
+                } catch (error) {
+                  channels.producerError.publish(error)
+                  channels.producerFinish.publish(undefined)
+                  throw error
+                }
+              })
+            }
+          })
+        }
+
+        return producer
+      }
+    })
+  }
+
+  // Helper function to wrap consumer classes
+  function wrapConsumerClass (ConsumerClass, className) {
+    return shimmer.wrap(module, className, function wrapConsumer (Original) {
+      return function wrappedConsumer () {
+        const consumer = new Original(...arguments)
+        const groupId = this.groupId || (arguments[0]?.['group.id'])
+
+        // Wrap the consume method
+        if (typeof consumer?.consume === 'function') {
+          shimmer.wrap(consumer, 'consume', function wrapConsume (consume) {
+            return function wrappedConsume (numMessages, callback) {
+              if (!channels.consumerStart.hasSubscribers) {
+                return consume.apply(this, arguments)
+              }
+
+              if (!callback && typeof numMessages === 'function') {
+                callback = numMessages
+              }
+
+              // Handle callback-based consumption
+              if (typeof callback === 'function') {
+                return consume.call(this, numMessages, function wrappedCallback (err, messages) {
+                  if (messages && messages.length > 0) {
+                    messages.forEach(message => {
+                      channels.consumerStart.publish({
+                        topic: message?.topic,
+                        partition: message?.partition,
+                        message,
+                        groupId
+                      })
+                      updateLatestOffset(message?.topic, message?.partition, message?.offset, groupId)
+                    })
+                  }
+
+                  if (err) {
+                    channels.consumerError.publish(err)
+                  }
+
+                  try {
+                    const result = callback.apply(this, arguments)
+                    channels.consumerFinish.publish(undefined)
+                    return result
+                  } catch (error) {
+                    channels.consumerError.publish(error)
+                    channels.consumerFinish.publish(undefined)
+                    throw error
+                  }
+                })
+              }
+
+              // If no callback is provided, just pass through
+              return consume.apply(this, arguments)
+            }
+          })
+
+          // Wrap the commit method for handling offset commits
+          if (consumer && typeof consumer.commit === 'function') {
+            shimmer.wrap(consumer, 'commit', wrapCommit)
+          }
+        }
+
+        return consumer
+      }
+    })
+  }
+
+  // Wrap Producer and KafkaProducer classes if they exist
+  if (typeof module.Producer === 'function') {
+    wrapProducerClass(module.Producer, 'Producer')
+  }
+  if (typeof module.KafkaProducer === 'function') {
+    wrapProducerClass(module.KafkaProducer, 'KafkaProducer')
+  }
+
+  // Wrap Consumer and KafkaConsumer classes if they exist
+  if (typeof module.Consumer === 'function') {
+    wrapConsumerClass(module.Consumer, 'Consumer')
+  }
+  if (typeof module.KafkaConsumer === 'function') {
+    wrapConsumerClass(module.KafkaConsumer, 'KafkaConsumer')
+  }
+}
+
+function instrumentKafkaJS (kafkaJS) {
+  // Hook the Kafka class if it exists
+  if (typeof kafkaJS?.Kafka === 'function') {
+    shimmer.wrap(kafkaJS, 'Kafka', function wrapKafka (OriginalKafka) {
+      return function KafkaWrapper (options) {
+        const kafka = new OriginalKafka(options)
+        const kafkaJSOptions = options?.kafkaJS || options
+        const brokers = kafkaJSOptions?.brokers ? kafkaJSOptions.brokers.join(',') : ''
+
+        // Store brokers for later use
+        kafka._ddBrokers = brokers
+
+        // Wrap the producer method if it exists
+        if (typeof kafka?.producer === 'function') {
+          shimmer.wrap(kafka, 'producer', function wrapProducerMethod (producerMethod) {
+            return function wrappedProducerMethod () {
+              const producer = producerMethod.apply(this, arguments)
+
+              if (!brokers && arguments?.[0]?.['bootstrap.servers']) {
+                kafka._ddBrokers = arguments[0]['bootstrap.servers']
+              }
+
+              // Wrap the send method of the producer
+              if (producer && typeof producer.send === 'function') {
+                shimmer.wrap(producer, 'send', function wrapSend (send) {
+                  return function wrappedSend (payload) {
+                    if (!channels.producerStart.hasSubscribers) {
+                      return send.apply(this, arguments)
+                    }
+
+                    const asyncResource = new AsyncResource('bound-anonymous-fn')
+                    return asyncResource.runInAsyncScope(() => {
+                      try {
+                        channels.producerStart.publish({
+                          topic: payload?.topic,
+                          messages: payload?.messages || [],
+                          bootstrapServers: kafka._ddBrokers,
+                          disableHeaderInjection: disabledHeaderWeakSet.has(producer)
+                        })
+
+                        const result = send.apply(this, arguments)
+
+                        result.then(
+                          asyncResource.bind(res => {
+                            channels.producerCommit.publish(res)
+                            channels.producerFinish.publish(undefined)
+                          }),
+                          asyncResource.bind(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. Tnfortunately the error name / type is not more specific.
+                              // This approach is implemented by other tracers as well.
+                              if (err.name === 'KafkaJSError' && err.type === 'ERR_UNKNOWN') {
+                                disabledHeaderWeakSet.add(producer)
+                                log.error('Kafka Broker responded with UNKNOWN_SERVER_ERROR (-1). ' +
+                                  'Please look at broker logs for more information. ' +
+                                  'Tracer message header injection for Kafka is disabled.')
+                              }
+                              channels.producerError.publish(err)
+                            }
+                            channels.producerFinish.publish(undefined)
+                          })
+                        )
+
+                        return result
+                      } catch (e) {
+                        channels.producerError.publish(e)
+                        channels.producerFinish.publish(undefined)
+                        throw e
+                      }
+                    })
+                  }
+                })
+              }
+
+              return producer
+            }
+          })
+        }
+
+        // Wrap the consumer method if it exists
+        if (typeof kafka?.consumer === 'function') {
+          shimmer.wrap(kafka, 'consumer', function wrapConsumerMethod (consumerMethod) {
+            return function wrappedConsumerMethod (config) {
+              const consumer = consumerMethod.apply(this, arguments)
+              const groupId = getGroupId(config)
+
+              // Wrap the run method for handling message consumption
+              if (typeof consumer?.run === 'function') {
+                shimmer.wrap(consumer, 'run', function wrapRun (run) {
+                  return function wrappedRun (options) {
+                    if (!channels.consumerStart.hasSubscribers) {
+                      return run.apply(this, arguments)
+                    }
+
+                    const eachMessage = options?.eachMessage
+                    const eachBatch = options?.eachBatch
+                    if (eachMessage) {
+                      options.eachMessage = wrapKafkaCallback(
+                        eachMessage,
+                        {
+                          startCh: channels.consumerStart,
+                          commitCh: channels.consumerCommit,
+                          finishCh: channels.consumerFinish,
+                          errorCh: channels.consumerError
+                        },
+                        (payload) => {
+                          return {
+                            topic: payload?.topic,
+                            partition: payload?.partition,
+                            offset: payload?.message?.offset,
+                            message: payload?.message,
+                            groupId
+                          }
+                        })
+                    } else if (eachBatch) {
+                      options.eachBatch = wrapKafkaCallback(
+                        eachBatch,
+                        {
+                          startCh: channels.batchConsumerStart,
+                          commitCh: channels.batchConsumerCommit,
+                          finishCh: channels.batchConsumerFinish,
+                          errorCh: channels.batchConsumerError
+                        },
+                        (payload) => {
+                          const { batch } = payload
+                          return {
+                            topic: batch?.topic,
+                            partition: batch?.partition,
+                            offset: batch?.messages[batch?.messages?.length - 1]?.offset,
+                            messages: batch?.messages,
+                            groupId
+                          }
+                        }
+                      )
+                    }
+
+                    return run.apply(this, arguments)
+                  }
+                })
+              }
+
+              // Wrap the commit method for handling offset commits
+              if (typeof consumer?.commitOffsets === 'function') {
+                shimmer.wrap(consumer, 'commitOffsets', wrapCommit)
+              }
+
+              return consumer
+            }
+          })
+        }
+
+        return kafka
+      }
+    })
+  }
+}
+
+function wrapCommit (commit) {
+  return function wrappedCommit (options) {
+    if (!channels.consumerCommit.hasSubscribers) {
+      return commit.apply(this, arguments)
+    }
+
+    const result = commit.apply(this, arguments)
+    channels.consumerCommit.publish(getLatestOffsets())
+    latestConsumerOffsets.clear()
+    return result
+  }
+}
+
+function wrapKafkaCallback (callback, { startCh, commitCh, finishCh, errorCh }, getPayload) {
+  return function wrappedKafkaCallback (payload) {
+    const commitPayload = getPayload(payload)
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    return asyncResource.runInAsyncScope(() => {
+      startCh.publish(commitPayload)
+
+      updateLatestOffset(commitPayload?.topic, commitPayload?.partition, commitPayload?.offset, commitPayload?.groupId)
+
+      try {
+        const result = callback.apply(this, arguments)
+
+        if (result && typeof result.then === 'function') {
+          return result
+            .then(asyncResource.bind(res => {
+              finishCh.publish(undefined)
+              return res
+            }))
+            .catch(asyncResource.bind(err => {
+              errorCh.publish(err)
+              finishCh.publish(undefined)
+              throw err
+            }))
+        } else {
+          finishCh.publish(undefined)
+          return result
+        }
+      } catch (error) {
+        errorCh.publish(error)
+        finishCh.publish(undefined)
+        throw error
+      }
+    })
+  }
+}
+
+function getGroupId (config) {
+  if (!config) return ''
+  if (config.kafkaJS?.groupId) return config.kafkaJS.groupId
+  if (config?.groupId) return config.groupId
+  if (config['group.id']) return config['group.id']
+  return ''
+}
+
+function updateLatestOffset (topic, partition, offset, groupId) {
+  const key = `${topic}:${partition}`
+  latestConsumerOffsets.set(key, {
+    topic,
+    partition,
+    offset,
+    groupId
+  })
+}
+
+function getLatestOffsets () {
+  return Array.from(latestConsumerOffsets.values())
+}

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

@@ -1,5 +1,6 @@
 'use strict'
 
+const { errorMonitor } = require('events')
 const {
   channel,
   addHook,
@@ -186,7 +187,7 @@ addHook({ name: 'couchbase', file: 'lib/bucket.js', versions: ['^2.6.12'] }, Buc
         finishCh.publish(undefined)
       }))
 
-      emitter.once('error', asyncResource.bind((error) => {
+      emitter.once(errorMonitor, asyncResource.bind((error) => {
         errorCh.publish(error)
         finishCh.publish(undefined)
       }))

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

@@ -8,6 +8,7 @@ const log = require('../../dd-trace/src/log')
 const testStartCh = channel('ci:cucumber:test:start')
 const testRetryCh = channel('ci:cucumber:test:retry')
 const testFinishCh = channel('ci:cucumber:test:finish') // used for test steps too
+const testFnCh = channel('ci:cucumber:test:fn')
 
 const testStepStartCh = channel('ci:cucumber:test-step:start')
 
@@ -52,7 +53,7 @@ const patched = new WeakSet()
 
 const lastStatusByPickleId = new Map()
 const numRetriesByPickleId = new Map()
-const numAttemptToAsyncResource = new Map()
+const numAttemptToCtx = new Map()
 const newTestsByTestFullname = new Map()
 
 let eventDataCollector = null
@@ -227,16 +228,12 @@ function wrapRun (pl, isLatestVersion) {
   patched.add(pl)
 
   shimmer.wrap(pl.prototype, 'run', run => function () {
-    if (!testStartCh.hasSubscribers) {
+    if (!testFinishCh.hasSubscribers) {
       return run.apply(this, arguments)
     }
 
     let numAttempt = 0
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    numAttemptToAsyncResource.set(numAttempt, asyncResource)
-
     const testFileAbsolutePath = this.pickle.uri
 
     const testSourceLine = this.gherkinDocument?.feature?.location?.line
@@ -247,9 +244,9 @@ function wrapRun (pl, isLatestVersion) {
       testSourceLine,
       isParallel: !!process.env.CUCUMBER_WORKER_ID
     }
-    asyncResource.runInAsyncScope(() => {
-      testStartCh.publish(testStartPayload)
-    })
+    const ctx = testStartPayload
+    numAttemptToCtx.set(numAttempt, ctx)
+    testStartCh.runStores(ctx, () => { })
     const promises = {}
     try {
       this.eventBroadcaster.on('envelope', shimmer.wrapFunction(null, () => async (testCase) => {
@@ -265,7 +262,7 @@ function wrapRun (pl, isLatestVersion) {
               // ignore error
             }
 
-            const failedAttemptAsyncResource = numAttemptToAsyncResource.get(numAttempt)
+            const failedAttemptCtx = numAttemptToCtx.get(numAttempt)
             const isFirstAttempt = numAttempt++ === 0
             const isAtrRetry = !isFirstAttempt && isFlakyTestRetriesEnabled
 
@@ -273,23 +270,19 @@ function wrapRun (pl, isLatestVersion) {
               await promises.hitBreakpointPromise
             }
 
-            failedAttemptAsyncResource.runInAsyncScope(() => {
-              // the current span will be finished and a new one will be created
-              testRetryCh.publish({ isFirstAttempt, error, isAtrRetry })
-            })
+            // the current span will be finished and a new one will be created
+            testRetryCh.publish({ isFirstAttempt, error, isAtrRetry, ...failedAttemptCtx.currentStore })
 
-            const newAsyncResource = new AsyncResource('bound-anonymous-fn')
-            numAttemptToAsyncResource.set(numAttempt, newAsyncResource)
+            const newCtx = { ...testStartPayload, promises }
+            numAttemptToCtx.set(numAttempt, newCtx)
 
-            newAsyncResource.runInAsyncScope(() => {
-              testStartCh.publish({ ...testStartPayload, promises }) // a new span will be created
-            })
+            testStartCh.runStores(newCtx, () => { })
           }
         }
       }))
       let promise
 
-      asyncResource.runInAsyncScope(() => {
+      testFnCh.runStores(ctx, () => {
         promise = run.apply(this, arguments)
       })
       promise.finally(async () => {
@@ -309,6 +302,7 @@ function wrapRun (pl, isLatestVersion) {
         let isAttemptToFixRetry = false
         let hasFailedAllRetries = false
         let hasPassedAllRetries = false
+        let hasFailedAttemptToFix = false
         let isDisabled = false
         let isQuarantined = false
 
@@ -330,6 +324,7 @@ function wrapRun (pl, isLatestVersion) {
               }, { pass: 0, fail: 0 })
               hasFailedAllRetries = fail === testManagementAttemptToFixRetries + 1
               hasPassedAllRetries = pass === testManagementAttemptToFixRetries + 1
+              hasFailedAttemptToFix = fail > 0
             }
           }
         }
@@ -341,38 +336,40 @@ function wrapRun (pl, isLatestVersion) {
           isEfdRetry = numRetries > 0
         }
 
-        const attemptAsyncResource = numAttemptToAsyncResource.get(numAttempt)
+        const attemptCtx = numAttemptToCtx.get(numAttempt)
 
         const error = getErrorFromCucumberResult(result)
 
         if (promises.hitBreakpointPromise) {
           await promises.hitBreakpointPromise
         }
-        attemptAsyncResource.runInAsyncScope(() => {
-          testFinishCh.publish({
-            status,
-            skipReason,
-            error,
-            isNew,
-            isEfdRetry,
-            isFlakyRetry: numAttempt > 0,
-            isAttemptToFix,
-            isAttemptToFixRetry,
-            hasFailedAllRetries,
-            hasPassedAllRetries,
-            isDisabled,
-            isQuarantined
-          })
+        testFinishCh.publish({
+          status,
+          skipReason,
+          error,
+          isNew,
+          isEfdRetry,
+          isFlakyRetry: numAttempt > 0,
+          isAttemptToFix,
+          isAttemptToFixRetry,
+          hasFailedAllRetries,
+          hasPassedAllRetries,
+          hasFailedAttemptToFix,
+          isDisabled,
+          isQuarantined,
+          ...attemptCtx.currentStore
         })
       })
       return promise
     } catch (err) {
-      errorCh.publish(err)
-      throw err
+      ctx.err = err
+      errorCh.runStores(ctx, () => {
+        throw err
+      })
     }
   })
   shimmer.wrap(pl.prototype, 'runStep', runStep => function () {
-    if (!testStepStartCh.hasSubscribers) {
+    if (!testFinishCh.hasSubscribers) {
       return runStep.apply(this, arguments)
     }
     const testStep = arguments[0]
@@ -384,9 +381,8 @@ function wrapRun (pl, isLatestVersion) {
       resource = testStep.isHook ? 'hook' : testStep.pickleStep.text
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    return asyncResource.runInAsyncScope(() => {
-      testStepStartCh.publish({ resource })
+    const ctx = { resource }
+    return testStepStartCh.runStores(ctx, () => {
       try {
         const promise = runStep.apply(this, arguments)
 
@@ -395,12 +391,14 @@ function wrapRun (pl, isLatestVersion) {
             ? getStatusFromResultLatest(result)
             : getStatusFromResult(result)
 
-          testFinishCh.publish({ isStep: true, status, skipReason, errorMessage })
+          testFinishCh.publish({ isStep: true, status, skipReason, errorMessage, ...ctx.currentStore })
         })
         return promise
       } catch (err) {
-        errorCh.publish(err)
-        throw err
+        ctx.err = err
+        errorCh.runStores(ctx, () => {
+          throw err
+        })
       }
     })
   })

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

@@ -1,6 +1,6 @@
 'use strict'
 
-const { channel, addHook, AsyncResource } = require('./helpers/instrument')
+const { channel, addHook } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
 const rrtypes = {
@@ -53,7 +53,7 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
   const errorCh = channel(prefix + ':error')
 
   const wrapped = function () {
-    const cb = AsyncResource.bind(arguments[arguments.length - 1])
+    const cb = arguments[arguments.length - 1]
     if (
       !startCh.hasSubscribers ||
       arguments.length < expectedArgs ||
@@ -62,30 +62,32 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
       return fn.apply(this, arguments)
     }
 
-    const startArgs = Array.from(arguments)
-    startArgs.pop() // gets rid of the callback
+    const args = Array.from(arguments)
+    args.pop() // gets rid of the callback
     if (rrtype) {
-      startArgs.push(rrtype)
+      args.push(rrtype)
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish(startArgs)
+    const ctx = { args }
 
-      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error, result) {
+    return startCh.runStores(ctx, () => {
+      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function (error, result, ...args) {
         if (error) {
-          errorCh.publish(error)
+          ctx.error = error
+          errorCh.publish(ctx)
         }
-        finishCh.publish(result)
-        cb.apply(this, arguments)
-      }))
+
+        ctx.result = result
+        finishCh.runStores(ctx, cb, this, error, result, ...args)
+      })
 
       try {
         return fn.apply(this, arguments)
       // TODO deal with promise versions when we support `dns/promises`
       } catch (error) {
         error.stack // trigger getting the stack at the original throwing point
-        errorCh.publish(error)
+        ctx.error = error
+        errorCh.publish(ctx)
 
         throw error
       }