Skip to content

Commit e93f79d

Browse files
fix(dsm): move DSM logic from bindStart to start for google-cloud-pubsub
Move DSM checkpoint logic from bindStart to start in both consumer and producer plugins so DSM context is properly bound in the async store. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 7990d7d commit e93f79d

2 files changed

Lines changed: 31 additions & 16 deletions

File tree

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,20 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
9696
})
9797
}
9898

99+
start (ctx) {
100+
if (!this.config.dsmEnabled) return
101+
const { message } = ctx
102+
if (!message?.attributes) return
103+
104+
const { span } = ctx.currentStore
105+
const subscription = message._subscriber._subscription
106+
const topic = subscription?.metadata?.topic || message.attributes?.['pubsub.topic']
107+
const payloadSize = getMessageSize(message)
108+
this.tracer.decodeDataStreamsContext(message.attributes)
109+
this.tracer
110+
.setCheckpoint(['direction:in', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
111+
}
112+
99113
bindStart (ctx) {
100114
const { message } = ctx
101115
const subscription = message._subscriber._subscription
@@ -184,13 +198,6 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
184198
}
185199
}
186200

187-
if (this.config.dsmEnabled && message?.attributes) {
188-
const payloadSize = getMessageSize(message)
189-
this.tracer.decodeDataStreamsContext(message.attributes)
190-
this.tracer
191-
.setCheckpoint(['direction:in', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
192-
}
193-
194201
return ctx.currentStore
195202
}
196203

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

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,23 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
3636
if (topicName) attributes['pubsub.topic'] = topicName
3737
}
3838

39+
start (ctx) {
40+
if (!this.config.dsmEnabled) return
41+
const { request } = ctx
42+
const messages = request.messages || []
43+
const topic = request.topic
44+
const { span } = ctx.currentStore
45+
46+
for (const msg of messages) {
47+
const dataStreamsContext = this.tracer.setCheckpoint(
48+
['direction:out', `topic:${topic}`, 'type:google-pubsub'],
49+
span,
50+
getHeadersSize(msg)
51+
)
52+
DsmPathwayCodec.encode(dataStreamsContext, msg.attributes)
53+
}
54+
}
55+
3956
bindStart (ctx) {
4057
const { request, api, projectId } = ctx
4158
if (api !== 'publish') return
@@ -128,15 +145,6 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
128145
if (batchTraceIdUpper) {
129146
msg.attributes['_dd.pubsub_request.p.tid'] = batchTraceIdUpper
130147
}
131-
132-
if (this.config.dsmEnabled) {
133-
const dataStreamsContext = this.tracer.setCheckpoint(
134-
['direction:out', `topic:${topic}`, 'type:google-pubsub'],
135-
batchSpan,
136-
getHeadersSize(msg)
137-
)
138-
DsmPathwayCodec.encode(dataStreamsContext, msg.attributes)
139-
}
140148
}
141149

142150
ctx.batchSpan = batchSpan

0 commit comments

Comments
 (0)