wip: show how we could add headers per message

lib/instrumentation/kafkajs/producer.js

@@ -29,10 +29,21 @@ module.exports = function instrumentProducer({ shim, kafkajs }) {
           firstMessage.headers = firstMessage.headers ?? {}
         }
 
+        debugger
         return new MessageSpec({
           destinationName: data.topic,
           destinationType: shim.TOPIC,
-          headers: firstMessage.headers
+          messageHeaders: () => {
+            debugger
+            return data.messages.map((message) => {
+              if (message.headers) {
+                return message.headers
+              } else {
+                message.headers = {}
+                return message.headers
+              }
+            })
+          }
         })
       })
 

lib/shim/message-shim/index.js

@@ -246,6 +246,14 @@ function recordProduce(nodule, properties, recordNamer) {
       if (msgDesc.headers) {
         shim.insertCATRequestHeaders(msgDesc.headers, true)
       }
+
+      if (msgDesc.messageHeaders) {
+        const headers = msgDesc.messageHeaders()
+        for (const msgHeader of headers) {
+          debugger
+          shim.insertCATRequestHeaders(msgHeader, true)
+        }
+      }
     }
     msgDesc.recorder = genericRecorder
     return msgDesc

lib/shim/specs/message.js

@@ -45,6 +45,13 @@ class MessageSpec extends OperationSpec {
    * @type {Object<string, string>|null}
    */
   headers
+
+  /**
+   * Function to extract headers to insert DT/CAT headers 
+   *
+   * @type {Function}
+   */
+  messageHeaders
 
   /**
    * A function to handle the result of the instrumented message broker
@@ -81,6 +88,7 @@ class MessageSpec extends OperationSpec {
     this.destinationName = params.destinationName ?? null
     this.destinationType = params.destinationType ?? null
     this.headers = params.headers ?? null
+    this.messageHeaders = params.messageHeaders ?? null
     this.messageHandler = params.messageHandler ?? null
     this.queue = params.queue ?? null
     this.routingKey = params.routingKey ?? null

test/versioned/kafkajs/kafka.tap.js

@@ -14,7 +14,11 @@ const utils = require('./utils')
 const broker = `${params.kafka_host}:${params.kafka_port}`
 
 tap.beforeEach(async (t) => {
-  t.context.agent = helper.instrumentMockedAgent()
+  t.context.agent = helper.instrumentMockedAgent({
+    feature_flag: {
+      kafkajs_instrumentation: true
+    }
+  })
 
   const { Kafka, logLevel } = require('kafkajs')
   t.context.Kafka = Kafka
@@ -123,12 +127,13 @@ tap.test('send passes along DT headers', (t) => {
   agent.on('transactionFinished', (tx) => {
     t.equal(tx.isDistributedTrace, true)
 
-    const headers = {}
+    /*const headers = {}
     tx.traceContext.addTraceContextHeaders(headers)
     t.match(headers, {
       traceparent: /00-4bf92f3577b34da6a3ce929d0e0e4736-[a-z0-9]{16}-00/,
       tracestate: /42@nr=0-0-account_1-app_1-[a-z0-9]{16}-[a-z0-9]{16}-0-0.123000-1717426365982/
     })
+    */
 
     t.end()
   })
@@ -137,7 +142,7 @@ tap.test('send passes along DT headers', (t) => {
     // This acceptTraceContextPayload is how we are simulating that the agent
     // received a distributed trace context that has resulted in the Kafka
     // payload production.
-    tx.acceptTraceContextPayload(traceparent, tracestate)
+    //tx.acceptTraceContextPayload(traceparent, tracestate)
 
     await consumer.subscribe({ topic, fromBeginning: true })
 
@@ -146,6 +151,10 @@ tap.test('send passes along DT headers', (t) => {
       consumer.run({
         eachMessage: async ({ message: actualMessage }) => {
           t.equal(messages.includes(actualMessage.value.toString()), true)
+          const expectedTraceState = actualMessage.headers.tracestate.toString()
+          const expectedTraceParent = actualMessage.headers.traceparent.toString()
+          t.equal(expectedTraceState, tracestate)
+          t.equal(expectedTraceParent, traceparent)
           msgCount += 1
           if (msgCount === 3) {
             resolve()