diff --git a/packages/datadog-plugin-google-cloud-pubsub/src/consumer.js b/packages/datadog-plugin-google-cloud-pubsub/src/consumer.js index e833e4b219..cef852e087 100644 --- a/packages/datadog-plugin-google-cloud-pubsub/src/consumer.js +++ b/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 } diff --git a/packages/datadog-plugin-google-cloud-pubsub/src/producer.js b/packages/datadog-plugin-google-cloud-pubsub/src/producer.js index 86e5cdf534..82183e7090 100644 --- a/packages/datadog-plugin-google-cloud-pubsub/src/producer.js +++ b/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 diff --git a/packages/datadog-plugin-google-cloud-pubsub/test/dsm.spec.js b/packages/datadog-plugin-google-cloud-pubsub/test/dsm.spec.js index 7ff87aeb94..4b68844efe 100644 --- a/packages/datadog-plugin-google-cloud-pubsub/test/dsm.spec.js +++ b/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