fix(google-cloud-pubsub): fix DSM context propagation for google-cloud-pubsub

packages/datadog-plugin-google-cloud-pubsub/src/consumer.js

@@ -96,6 +96,20 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
     })
   }
 
+  start (ctx) {
+    if (!this.config.dsmEnabled) return
+    const { message } = ctx
+    if (!message?.attributes) return
+
+    const { span } = ctx.currentStore
+    const subscription = message._subscriber._subscription
+    const topic = subscription?.metadata?.topic || message.attributes?.['pubsub.topic']
+    const payloadSize = getMessageSize(message)
+    this.tracer.decodeDataStreamsContext(message.attributes)
+    this.tracer
+      .setCheckpoint(['direction:in', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
+  }
+
   bindStart (ctx) {
     const { message } = ctx
     const subscription = message._subscriber._subscription
@@ -181,13 +195,6 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
       }
     }
 
-    if (this.config.dsmEnabled && message?.attributes) {
-      const payloadSize = getMessageSize(message)
-      this.tracer.decodeDataStreamsContext(message.attributes)
-      this.tracer
-        .setCheckpoint(['direction:in', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
-    }
-
     return ctx.currentStore
   }
 

packages/datadog-plugin-google-cloud-pubsub/src/producer.js

@@ -36,6 +36,25 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
     if (topicName) attributes['pubsub.topic'] = topicName
   }
 
+  start (ctx) {
+    if (!this.config.dsmEnabled) return
+    if (ctx.api !== 'publish' || !ctx.currentStore) return
+    const { request } = ctx
+    const messages = request.messages || []
+    const topic = request.topic
+    const { span } = ctx.currentStore
+
+    for (const msg of messages) {
+      const attributes = msg.attributes ??= {}
+      const dataStreamsContext = this.tracer.setCheckpoint(
+        ['direction:out', `topic:${topic}`, 'type:google-pubsub'],
+        span,
+        getHeadersSize(msg)
+      )
+      DsmPathwayCodec.encode(dataStreamsContext, attributes)
+    }
+  }
+
   bindStart (ctx) {
     const { request, api, projectId } = ctx
     if (api !== 'publish') return
@@ -113,7 +132,6 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
 
     const messageCountStr = String(messageCount)
     const startTimeStr = String(Math.floor(batchSpan._startTime))
-    const dsmEnabled = this.config.dsmEnabled
 
     for (let i = 0; i < messageCount; i++) {
       const msg = messages[i]
@@ -137,15 +155,6 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
       if (batchTraceIdUpper) {
         attributes['_dd.pubsub_request.p.tid'] = batchTraceIdUpper
       }
-
-      if (dsmEnabled) {
-        const dataStreamsContext = this.tracer.setCheckpoint(
-          ['direction:out', `topic:${topic}`, 'type:google-pubsub'],
-          batchSpan,
-          getHeadersSize(msg)
-        )
-        DsmPathwayCodec.encode(dataStreamsContext, attributes)
-      }
     }
 
     ctx.batchSpan = batchSpan

packages/datadog-plugin-google-cloud-pubsub/test/dsm.spec.js

@@ -128,6 +128,72 @@ describe('Plugin', () => {
           })
         })
 
+        describe('concurrent context isolation', () => {
+          it('Should maintain separate DSM context for interleaved consume-produce flows', async () => {
+            const setCheckpointSpy = sinon.spy(DataStreamsProcessor.prototype, 'setCheckpoint')
+
+            try {
+              const topicAIn = (await pubsub.createTopic(`dsm-iso-a-in-${id()}`))[0]
+              const topicBIn = (await pubsub.createTopic(`dsm-iso-b-in-${id()}`))[0]
+              const topicAOut = (await pubsub.createTopic(`dsm-iso-a-out-${id()}`))[0]
+              const topicBOut = (await pubsub.createTopic(`dsm-iso-b-out-${id()}`))[0]
+
+              const subA = (await topicAIn.createSubscription(`sub-a-${id()}`))[0]
+              const subB = (await topicBIn.createSubscription(`sub-b-${id()}`))[0]
+
+              const fullTopicAIn = topicAIn.metadata?.name || topicAIn.name
+              const fullTopicBIn = topicBIn.metadata?.name || topicBIn.name
+              const fullTopicAOut = topicAOut.metadata?.name || topicAOut.name
+              const fullTopicBOut = topicBOut.metadata?.name || topicBOut.name
+
+              // Synchronization: both consumers must receive before either produces
+              let resolveAEntered, resolveBEntered
+              const aEntered = new Promise(resolve => { resolveAEntered = resolve })
+              const bEntered = new Promise(resolve => { resolveBEntered = resolve })
+              let doneCount = 0
+              const allDone = new Promise(resolve => {
+                const check = () => { if (++doneCount === 2) resolve() }
+                subA.on('message', async (msg) => {
+                  msg.ack()
+                  resolveAEntered()
+                  await bEntered
+                  await publish(topicAOut, { data: Buffer.from('from-a') })
+                  check()
+                })
+                subB.on('message', async (msg) => {
+                  msg.ack()
+                  resolveBEntered()
+                  await aEntered
+                  await publish(topicBOut, { data: Buffer.from('from-b') })
+                  check()
+                })
+              })
+
+              await publish(topicAIn, { data: Buffer.from('msg-a') })
+              await publish(topicBIn, { data: Buffer.from('msg-b') })
+
+              await allDone
+
+              const calls = setCheckpointSpy.getCalls()
+              const checkpoint = (dir, topic) => calls.find(c =>
+                c.args[0].includes(`direction:${dir}`) && c.args[0].includes(`topic:${topic}`)
+              )
+
+              const consumeA = checkpoint('in', fullTopicAIn)
+              const consumeB = checkpoint('in', fullTopicBIn)
+              const produceA = checkpoint('out', fullTopicAOut)
+              const produceB = checkpoint('out', fullTopicBOut)
+
+              assert.ok(produceA?.args[2], 'Process A produce should have a parent DSM context')
+              assert.ok(produceB?.args[2], 'Process B produce should have a parent DSM context')
+              assert.deepStrictEqual(produceA.args[2].hash, consumeA.returnValue.hash)
+              assert.deepStrictEqual(produceB.args[2].hash, consumeB.returnValue.hash)
+            } finally {
+              setCheckpointSpy.restore()
+            }
+          })
+        })
+
         describe('it should set a message payload size', () => {
           let recordCheckpointSpy