dd-trace 5.80.0 → 5.82.0

package/packages/datadog-plugin-express/src/index.js

@@ -1,8 +1,8 @@
 'use strict'
 
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 const ExpressTracingPlugin = require('./tracing')
 const ExpressCodeOriginForSpansPlugin = require('./code_origin')
-const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 
 class ExpressPlugin extends CompositePlugin {
   static id = 'express'

package/packages/datadog-plugin-fastify/src/index.js

@@ -1,8 +1,8 @@
 'use strict'
 
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 const FastifyTracingPlugin = require('./tracing')
 const FastifyCodeOriginForSpansPlugin = require('./code_origin')
-const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 
 class FastifyPlugin extends CompositePlugin {
   static id = 'fastify'

package/packages/datadog-plugin-google-cloud-pubsub/src/client.js

@@ -8,11 +8,11 @@ class GoogleCloudPubsubClientPlugin extends ClientPlugin {
   static operation = 'request'
 
   start (ctx) {
-    const { request, api, projectId } = ctx
+    const { request, api, projectId, storedContext } = ctx
 
     if (api === 'publish') return
 
-    this.startSpan(this.operationName(), {
+    const spanOptions = {
       service: this.config.service || this.serviceName(),
       resource: [api, request.name].filter(Boolean).join(' '),
       kind: this.constructor.kind,
@@ -20,7 +20,17 @@ class GoogleCloudPubsubClientPlugin extends ClientPlugin {
         'pubsub.method': api,
         'gcloud.project_id': projectId
       }
-    }, ctx)
+    }
+
+    /**
+     * Use stored context from consumer plugin to link acknowledge span to message processing span.
+     * Without this, the acknowledge span would be orphaned (no async context available).
+     */
+    if (storedContext?.span) {
+      spanOptions.childOf = storedContext.span.context()
+    }
+
+    this.startSpan(this.operationName(), spanOptions, ctx)
 
     return ctx.currentStore
   }

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

@@ -2,30 +2,188 @@
 
 const { getMessageSize } = require('../../dd-trace/src/datastreams')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
+const SpanContext = require('../../dd-trace/src/opentracing/span_context')
+const id = require('../../dd-trace/src/id')
+const { storage } = require('../../datadog-core')
+
+/**
+ * PULL SUBSCRIPTION: Service explicitly pulls messages from Pub/Sub and processes them.
+ * Async context storage for linking acknowledge() API calls back to original message spans.
+ *
+ * Problem: message.ack() is often called asynchronously outside the original handler,
+ * losing async context. The acknowledge() API only receives ackIds (strings), not Messages.
+ *
+ * Solution: Store context by Message (WeakMap), lookup by ackId (Map<ackId, WeakRef<Message>>).
+ * WeakRef allows Messages to be GC'd even if never acknowledged (network failures, crashes).
+ */
+const messageToContext = new WeakMap() // Message -> context (auto-cleanup on GC)
+const ackIdToMessage = new Map() // ackId -> WeakRef<Message> (needs cleanup)
+
+const ackMapCleanup = new FinalizationRegistry((ackId) => {
+  ackIdToMessage.delete(ackId) // Remove orphaned ackId when Message is GC'd
+})
 
 class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
   static id = 'google-cloud-pubsub'
   static operation = 'receive'
 
+  constructor (...args) {
+    super(...args)
+
+    /**
+     * Channel: message:ack-store
+     * When message.ack() is called, capture and store the current async context.
+     * Use WeakRef to allow Message GC even if never acknowledged.
+     */
+    this.addSub('apm:google-cloud-pubsub:message:ack-store', (ctx) => {
+      const { message, ackId } = ctx
+      const currentStore = storage('legacy').getStore()
+
+      if (currentStore) {
+        messageToContext.set(message, currentStore)
+        ackIdToMessage.set(ackId, new WeakRef(message))
+        ackMapCleanup.register(message, ackId, message)
+      }
+    })
+
+    /**
+     * Channel: message:ack-retrieve
+     * When acknowledge() API is called, retrieve stored context for the ackIds.
+     * Map: ackId -> WeakRef<Message> -> stored context.
+     * Clean up immediately after acknowledge (happy path), or via FinalizationRegistry (GC).
+     */
+    this.addSub('apm:google-cloud-pubsub:message:ack-retrieve', ({ ackIds, api, ctx }) => {
+      for (const ackId of ackIds) {
+        const weakRef = ackIdToMessage.get(ackId)
+        if (weakRef) {
+          const message = weakRef.deref()
+          if (message) {
+            const storedContext = messageToContext.get(message)
+            if (storedContext) {
+              ctx.storedContext = storedContext
+              break
+            }
+          }
+        }
+      }
+
+      if (api === 'acknowledge') {
+        ackIds.forEach(ackId => ackIdToMessage.delete(ackId))
+      }
+    })
+  }
+
+  #reconstructPubSubRequestContext (attrs) {
+    const traceIdLower = attrs['_dd.pubsub_request.trace_id']
+    const spanId = attrs['_dd.pubsub_request.span_id']
+    const traceIdUpper = attrs['_dd.pubsub_request.p.tid']
+
+    if (!traceIdLower || !spanId) return null
+
+    const traceId128 = traceIdUpper ? traceIdUpper + traceIdLower : traceIdLower.padStart(32, '0')
+    const traceId = id(traceId128, 16)
+    const parentId = id(spanId, 16)
+
+    const tags = {}
+    if (traceIdUpper) tags['_dd.p.tid'] = traceIdUpper
+
+    return new SpanContext({
+      traceId,
+      spanId: parentId,
+      tags
+    })
+  }
+
   bindStart (ctx) {
     const { message } = ctx
     const subscription = message._subscriber._subscription
-    const topic = subscription.metadata && subscription.metadata.topic
-    const childOf = this.tracer.extract('text_map', message.attributes) || null
+    const topic = subscription?.metadata?.topic || message.attributes?.['pubsub.topic']
+
+    const batchRequestTraceId = message.attributes?.['_dd.pubsub_request.trace_id']
+    const batchRequestSpanId = message.attributes?.['_dd.pubsub_request.span_id']
+
+    // Parse batch metadata once upfront for reuse
+    const batchSizeStr = message.attributes?.['_dd.batch.size']
+    const batchIndexStr = message.attributes?.['_dd.batch.index']
+    const batchSize = batchSizeStr ? Number(batchSizeStr) : undefined
+    const batchIndex = batchIndexStr ? Number(batchIndexStr) : undefined
+
+    let childOf = this.tracer.extract('text_map', message.attributes)
+
+    const isFirstMessage = batchIndex === 0
+    if (isFirstMessage && batchRequestSpanId) {
+      const pubsubRequestContext = this.#reconstructPubSubRequestContext(message.attributes)
+      if (pubsubRequestContext) {
+        childOf = pubsubRequestContext
+      }
+    }
+
+    const topicName = topic?.slice(topic.lastIndexOf('/') + 1) ??
+      subscription.name.slice(subscription.name.lastIndexOf('/') + 1)
+    const baseService = this.tracer._service || 'unknown'
+    const serviceName = this.config.service || `${baseService}-pubsub`
+    const meta = {
+      'gcloud.project_id': subscription.pubsub.projectId,
+      'pubsub.topic': topic,
+      'span.kind': 'consumer',
+      'pubsub.subscription_type': 'pull',
+      'pubsub.span_type': 'message_processing',
+      'messaging.operation': 'receive',
+      base_service: baseService,
+      service_override_type: 'custom'
+    }
+
+    if (batchRequestTraceId && batchRequestSpanId) {
+      meta['pubsub.batch.request_trace_id'] = batchRequestTraceId
+      // Use JSON format like producer for proper span link parsing
+      meta['_dd.span_links'] = JSON.stringify([{
+        trace_id: batchRequestTraceId,
+        span_id: batchRequestSpanId,
+        flags: 0
+      }])
+    }
+
+    const metrics = {
+      'pubsub.ack': 0
+    }
+
+    if (batchSize) {
+      metrics['pubsub.batch.message_count'] = batchSize
+      metrics['pubsub.batch.size'] = batchSize
+    }
+    if (batchIndex !== undefined) {
+      metrics['pubsub.batch.message_index'] = batchIndex
+      metrics['pubsub.batch.index'] = batchIndex
+    }
+
+    if (batchSize && batchIndex !== undefined) {
+      meta['pubsub.batch.description'] = `Message ${batchIndex + 1} of ${batchSize}`
+    }
 
     const span = this.startSpan({
       childOf,
-      resource: topic,
+      resource: `Message from ${topicName}`,
       type: 'worker',
-      meta: {
-        'gcloud.project_id': subscription.pubsub.projectId,
-        'pubsub.topic': topic
-      },
-      metrics: {
-        'pubsub.ack': 0
-      }
+      service: serviceName,
+      meta,
+      metrics
     }, ctx)
 
+    if (message.id) {
+      span.setTag('pubsub.message_id', message.id)
+    }
+    if (message.publishTime) {
+      span.setTag('pubsub.publish_time', message.publishTime.toISOString())
+    }
+
+    if (message.attributes) {
+      const publishStartTime = message.attributes['x-dd-publish-start-time']
+      if (publishStartTime) {
+        const deliveryDuration = Date.now() - Number(publishStartTime)
+        span.setTag('pubsub.delivery_duration_ms', deliveryDuration)
+      }
+    }
+
     if (this.config.dsmEnabled && message?.attributes) {
       const payloadSize = getMessageSize(message)
       this.tracer.decodeDataStreamsContext(message.attributes)
@@ -38,14 +196,15 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
 
   bindFinish (ctx) {
     const { message } = ctx
-    const span = ctx.currentStore.span
+    const span = ctx.currentStore?.span
+
+    if (!span) return ctx.parentStore
 
     if (message?._handled) {
       span.setTag('pubsub.ack', 1)
     }
 
     super.finish()
-
     return ctx.parentStore
   }
 }

package/packages/datadog-plugin-google-cloud-pubsub/src/index.js

@@ -1,11 +1,10 @@
 'use strict'
 
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 const ProducerPlugin = require('./producer')
 const ConsumerPlugin = require('./consumer')
 const ClientPlugin = require('./client')
-const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 
-// TODO: Consider splitting channels for publish/receive in the instrumentation.
 class GoogleCloudPubsubPlugin extends CompositePlugin {
   static id = 'google-cloud-pubsub'
   static get plugins () {

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

@@ -2,42 +2,189 @@
 
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
 const { DsmPathwayCodec, getHeadersSize } = require('../../dd-trace/src/datastreams')
+const id = require('../../dd-trace/src/id')
 
 class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
   static id = 'google-cloud-pubsub'
   static operation = 'request'
 
+  constructor (...args) {
+    super(...args)
+
+    /**
+     * Inject trace context into individual messages published via Topic.publish() or
+     * Topic.publishMessage(). This happens before the batch publish API call, allowing
+     * each message to propagate trace context to consumers.
+     */
+    this.addSub('apm:google-cloud-pubsub:message:publish', this.handleMessagePublish.bind(this))
+  }
+
+  handleMessagePublish ({ attributes, pubsub, topicName }) {
+    // Skip if message already has trace context from upstream
+    if (attributes['x-datadog-trace-id'] || attributes.traceparent) return
+
+    const activeSpan = this.tracer.scope().active()
+    if (!activeSpan) return
+
+    // Inject current span's trace context into message attributes
+    this.tracer.inject(activeSpan, 'text_map', attributes)
+
+    const traceIdUpperBits = activeSpan.context()._trace.tags['_dd.p.tid']
+    if (traceIdUpperBits) attributes['_dd.p.tid'] = traceIdUpperBits
+
+    if (pubsub) attributes['gcloud.project_id'] = pubsub.projectId
+    if (topicName) attributes['pubsub.topic'] = topicName
+  }
+
   bindStart (ctx) {
     const { request, api, projectId } = ctx
-
     if (api !== 'publish') return
 
     const messages = request.messages || []
     const topic = request.topic
-    const span = this.startSpan({ // TODO: rename
-      resource: `${api} ${topic}`,
+    const messageCount = messages.length
+    const hasTraceContext = messages[0]?.attributes?.['x-datadog-trace-id']
+
+    /**
+     * Batch Publishing Strategy:
+     * - Create one "batch span" representing the entire publish operation
+     * - If messages already have trace context (from upstream), use the first message's
+     *   context as the parent, and create span links for messages 2-N
+     * - Otherwise, use the current active span as parent
+     * - Inject batch span context + metadata into all message attributes for downstream
+     *   consumers to reconstruct the trace and understand batch relationships
+     */
+    const spanLinkData = hasTraceContext
+      ? messages.slice(1).map(msg => this.#extractSpanLink(msg.attributes)).filter(Boolean)
+      : []
+
+    const firstAttrs = messages[0]?.attributes
+    const parentData = firstAttrs?.['x-datadog-trace-id'] && firstAttrs['x-datadog-parent-id']
+      ? {
+          traceId: firstAttrs['x-datadog-trace-id'],
+          spanId: firstAttrs['x-datadog-parent-id'],
+          traceIdUpper: firstAttrs['_dd.p.tid'],
+          samplingPriority: firstAttrs['x-datadog-sampling-priority']
+        }
+      : null
+
+    const lastSlash = topic.lastIndexOf('/')
+    const topicName = lastSlash === -1 ? topic : topic.slice(lastSlash + 1)
+    const batchSpan = this.startSpan({
+      childOf: parentData ? this.#extractParentContext(parentData) : undefined,
+      resource: `${api} to Topic ${topicName}`,
       meta: {
         'gcloud.project_id': projectId,
-        'pubsub.method': api, // TODO: remove
-        'pubsub.topic': topic
+        'pubsub.method': api,
+        'pubsub.topic': topic,
+        'span.kind': 'producer',
+        '_dd.base_service': this.tracer._service,
+        '_dd.serviceoverride.type': 'integration',
+        'pubsub.linked_message_count': spanLinkData.length || undefined,
+        operation: messageCount > 1 ? 'batched.pubsub.request' : 'pubsub.request'
+      },
+      metrics: {
+        'pubsub.batch.message_count': messageCount,
+        'pubsub.batch': messageCount > 1 ? true : undefined
       }
     }, ctx)
 
-    for (const msg of messages) {
-      if (!msg.attributes) {
-        msg.attributes = {}
+    const spanCtx = batchSpan.context()
+    // Get 128-bit trace ID and span ID as hex strings
+    const fullTraceIdHex = spanCtx.toTraceId(true)
+    const batchSpanIdHex = spanCtx.toSpanId(true)
+    // Extract lower 64 bits (last 16 hex chars) for trace ID
+    const batchTraceIdHex = fullTraceIdHex.slice(-16)
+    const batchTraceIdUpper = spanCtx._trace.tags['_dd.p.tid']
+
+    if (spanLinkData.length) {
+      batchSpan.setTag('_dd.span_links', JSON.stringify(
+        spanLinkData.map(link => ({
+          trace_id: link.traceId,
+          span_id: link.spanId,
+          flags: link.samplingPriority || 0
+        }))
+      ))
+    }
+
+    messages.forEach((msg, i) => {
+      msg.attributes ??= {}
+
+      if (!hasTraceContext) {
+        this.tracer.inject(batchSpan, 'text_map', msg.attributes)
+      }
+
+      Object.assign(msg.attributes, {
+        '_dd.pubsub_request.trace_id': batchTraceIdHex,
+        '_dd.pubsub_request.span_id': batchSpanIdHex,
+        '_dd.batch.size': String(messageCount),
+        '_dd.batch.index': String(i),
+        'gcloud.project_id': projectId,
+        'pubsub.topic': topic,
+        'x-dd-publish-start-time': String(Math.floor(batchSpan._startTime))
+      })
+
+      if (batchTraceIdUpper) {
+        msg.attributes['_dd.pubsub_request.p.tid'] = batchTraceIdUpper
       }
-      this.tracer.inject(span, 'text_map', msg.attributes)
+
       if (this.config.dsmEnabled) {
-        const payloadSize = getHeadersSize(msg)
-        const dataStreamsContext = this.tracer
-          .setCheckpoint(['direction:out', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
+        const dataStreamsContext = this.tracer.setCheckpoint(
+          ['direction:out', `topic:${topic}`, 'type:google-pubsub'],
+          batchSpan,
+          getHeadersSize(msg)
+        )
         DsmPathwayCodec.encode(dataStreamsContext, msg.attributes)
       }
-    }
+    })
 
+    ctx.batchSpan = batchSpan
     return ctx.currentStore
   }
+
+  bindFinish (ctx) {
+    if (ctx.batchSpan && !ctx.batchSpan._duration) ctx.batchSpan.finish()
+    return super.bindFinish(ctx)
+  }
+
+  bindError (ctx) {
+    if (ctx.error && ctx.batchSpan) {
+      ctx.batchSpan.setTag('error', ctx.error)
+    }
+    return ctx.parentStore
+  }
+
+  #extractSpanLink (attrs) {
+    if (!attrs?.['x-datadog-trace-id'] || !attrs['x-datadog-parent-id']) return null
+
+    // Convert to hex strings
+    const lowerHex = id(attrs['x-datadog-trace-id']).toString(16)
+    const spanIdHex = id(attrs['x-datadog-parent-id']).toString(16)
+
+    // Build full 128-bit trace ID
+    const traceIdHex = attrs['_dd.p.tid']
+      ? attrs['_dd.p.tid'] + lowerHex
+      : lowerHex.padStart(32, '0')
+
+    return {
+      traceId: traceIdHex,
+      spanId: spanIdHex,
+      samplingPriority: attrs['x-datadog-sampling-priority']
+        ? Number.parseInt(attrs['x-datadog-sampling-priority'], 10)
+        : undefined
+    }
+  }
+
+  #extractParentContext (data) {
+    const carrier = {
+      'x-datadog-trace-id': data.traceId,
+      'x-datadog-parent-id': data.spanId
+    }
+    if (data.traceIdUpper) carrier['_dd.p.tid'] = data.traceIdUpper
+    if (data.samplingPriority) carrier['x-datadog-sampling-priority'] = String(data.samplingPriority)
+
+    return this.tracer.extract('text_map', carrier)
+  }
 }
 
 module.exports = GoogleCloudPubsubProducerPlugin