tracing: remove asyncresource from kafka instrumentations

.github/workflows/appsec.yml

@@ -302,3 +302,35 @@ jobs:
       - uses: ./.github/actions/node/active-lts
       - run: yarn test:appsec:plugins:ci
       - uses: codecov/codecov-action@ad3126e916f78f00edff4ed0317cf185271ccc2d # v5.4.2
+
+  kafka:
+    runs-on: ubuntu-latest
+    services:
+      kafka:
+        image: apache/kafka-native:3.8.0-rc2
+        env:
+          KAFKA_PROCESS_ROLES: broker,controller
+          KAFKA_NODE_ID: '1'
+          KAFKA_LISTENERS: PLAINTEXT://:9092,CONTROLLER://:9093
+          KAFKA_CONTROLLER_QUORUM_VOTERS: [email protected]:9093
+          KAFKA_CONTROLLER_LISTENER_NAMES: CONTROLLER
+          KAFKA_CLUSTER_ID: r4zt_wrqTRuT7W2NJsB_GA
+          KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://127.0.0.1:9092
+          KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT
+          KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT
+          KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: '1'
+          KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS: '0'
+        ports:
+          - 9092:9092
+          - 9093:9093
+    env:
+      PLUGINS: kafkajs
+      SERVICES: kafka
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: ./.github/actions/node/oldest-maintenance-lts
+      - uses: ./.github/actions/install
+      - run: yarn test:appsec:plugins:ci
+      - uses: ./.github/actions/node/active-lts
+      - run: yarn test:appsec:plugins:ci
+      - uses: codecov/codecov-action@ad3126e916f78f00edff4ed0317cf185271ccc2d # v5.4.2

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

@@ -2,8 +2,7 @@
 
 const {
   addHook,
-  channel,
-  AsyncResource
+  channel
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -32,7 +31,6 @@ 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)
@@ -55,30 +53,31 @@ function instrumentBaseModule (module) {
         // 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) {
+            return function wrappedProduce (topic, partition, message, key, timestamp, opaque, headers) {
               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 ctx = {}
+              ctx.topic = topic
+              ctx.messages = [{ key, value: message }]
+              ctx.bootstrapServers = brokers
 
-                  const result = produce.apply(this, arguments)
+              return channels.producerStart.runStores(ctx, () => {
+                try {
+                  headers = convertHeaders(ctx.messages[0].headers)
+                  const result = produce.apply(this, [topic, partition, message, key, timestamp, opaque, headers])
 
-                  channels.producerCommit.publish(undefined)
-                  channels.producerFinish.publish(undefined)
+                  ctx.res = result
+                  channels.producerCommit.publish(ctx)
+                  channels.producerFinish.runStores(ctx, () => {})
                   return result
                 } catch (error) {
-                  channels.producerError.publish(error)
-                  channels.producerFinish.publish(undefined)
+                  ctx.err = error
+                  channels.producerError.publish(ctx)
+                  channels.producerFinish.runStores(ctx, () => {})
                   throw error
                 }
               })
@@ -110,32 +109,37 @@ function instrumentBaseModule (module) {
                 callback = numMessages
               }
 
+              const ctx = {}
               // 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
-                      })
+                      ctx.topic = message?.topic
+                      ctx.partition = message?.partition
+                      ctx.message = message
+                      ctx.groupId = groupId
+
+                      channels.consumerStart.runStores(ctx, () => {})
                       updateLatestOffset(message?.topic, message?.partition, message?.offset, groupId)
                     })
                   }
 
                   if (err) {
-                    channels.consumerError.publish(err)
+                    ctx.err = err
+                    channels.consumerError.publish(ctx)
                   }
 
                   try {
                     const result = callback.apply(this, arguments)
-                    channels.consumerFinish.publish(undefined)
+                    if (messages && messages.length > 0) {
+                      channels.consumerFinish.runStores(ctx, () => {})
+                    }
                     return result
                   } catch (error) {
-                    channels.consumerError.publish(error)
-                    channels.consumerFinish.publish(undefined)
+                    ctx.err = error
+                    channels.consumerError.publish(ctx)
+                    channels.consumerFinish.runStores(ctx, () => {})
                     throw error
                   }
                 })
@@ -204,44 +208,42 @@ function instrumentKafkaJS (kafkaJS) {
                       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 ctx = {}
+                    ctx.topic = payload?.topic
+                    ctx.messages = payload?.messages || []
+                    ctx.bootstrapServers = kafka._ddBrokers
+                    ctx.disableHeaderInjection = disabledHeaderWeakSet.has(producer)
 
+                    return channels.producerStart.runStores(ctx, () => {
+                      try {
                         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)
+                        result.then((res) => {
+                          ctx.res = res
+                          channels.producerCommit.publish(ctx)
+                          channels.producerFinish.publish(undefined)
+                        }, (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.producerFinish.publish(undefined)
-                          })
-                        )
+                            ctx.err = err
+                            channels.producerError.publish(ctx)
+                          }
+                          channels.producerFinish.publish(undefined)
+                        })
 
                         return result
                       } catch (e) {
-                        channels.producerError.publish(e)
+                        ctx.err = e
+                        channels.producerError.publish(ctx)
                         channels.producerFinish.publish(undefined)
                         throw e
                       }
@@ -350,33 +352,36 @@ function wrapKafkaCallback (callback, { startCh, commitCh, finishCh, errorCh },
   return function wrappedKafkaCallback (payload) {
     const commitPayload = getPayload(payload)
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-    return asyncResource.runInAsyncScope(() => {
-      startCh.publish(commitPayload)
+    const ctx = {}
+    ctx.extractedArgs = commitPayload
 
+    return startCh.runStores(ctx, () => {
       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)
+            .then((res) => {
+              ctx.res = res
+              finishCh.runStores(ctx, () => {})
               return res
-            }))
-            .catch(asyncResource.bind(err => {
-              errorCh.publish(err)
-              finishCh.publish(undefined)
+            })
+            .catch((err) => {
+              ctx.err = err
+              errorCh.publish(ctx)
+              finishCh.runStores(ctx, () => {})
               throw err
-            }))
+            })
         } else {
-          finishCh.publish(undefined)
+          finishCh.runStores(ctx, () => {})
           return result
         }
       } catch (error) {
-        errorCh.publish(error)
-        finishCh.publish(undefined)
+        ctx.err = error
+        errorCh.publish(ctx)
+        finishCh.runStores(ctx, () => {})
         throw error
       }
     })
@@ -404,3 +409,8 @@ function updateLatestOffset (topic, partition, offset, groupId) {
 function getLatestOffsets () {
   return Array.from(latestConsumerOffsets.values())
 }
+
+function convertHeaders (headers) {
+  // convert headers from object to array of objects with 1 key and value per array entry
+  return Object.entries(headers).map(([key, value]) => ({ [key.toString()]: value.toString() }))
+}