feat: add ack context map and producer improvements for batching

- Add ack context map to preserve trace context across batched acknowledges
- Update producer to use batchSpan._startTime for accurate publish time
- Add explicit parent span support in client plugin
- Wrap Message.ack() to store context before batched gRPC acknowledge
- Update Subscription.emit to properly handle storage context
- Sync auto-load improvements from Branch 1

Author: nina9753

packages/datadog-instrumentations/src/google-cloud-pubsub.js

@@ -5,16 +5,17 @@ const {
   addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
+const { storage } = require('../../datadog-core')
 
 // Auto-load push subscription plugin for push subscribers
 // This ensures pubsub.delivery spans work even if app doesn't import @google-cloud/pubsub
+// The plugin's constructor subscribes to apm:http:server:request:start immediately (no tracer needed)
+// The tracer is only used at request time when delivery spans are created
 try {
-  const tracer = require('../../dd-trace')
   const PushSubscriptionPlugin = require('../../datadog-plugin-google-cloud-pubsub/src/pubsub-push-subscription')
-  if (tracer._tracer) {
-    const handler = new PushSubscriptionPlugin(tracer._tracer, {})
-    handler.configure({})
-  }
+  const pushPlugin = new PushSubscriptionPlugin(null, {})
+  // Plugin instance kept alive to maintain channel subscription
+  pushPlugin.configure({})
 } catch {
   // Silent - push subscription plugin is optional
 }
@@ -27,6 +28,10 @@ const receiveStartCh = channel('apm:google-cloud-pubsub:receive:start')
 const receiveFinishCh = channel('apm:google-cloud-pubsub:receive:finish')
 const receiveErrorCh = channel('apm:google-cloud-pubsub:receive:error')
 
+// Global map to store ackId -> span context for batched acknowledges
+// This allows us to restore the correct trace context when the batched gRPC call happens
+const ackContextMap = new Map()
+
 const publisherMethods = [
   'createTopic',
   'updateTopic',
@@ -77,7 +82,75 @@ function wrapMethod (method) {
   return function (request) {
     if (!requestStartCh.hasSubscribers) return method.apply(this, arguments)
 
+    // For acknowledge/modifyAckDeadline, try to restore span context from stored map
+    let restoredStore = null
+    if (api === 'acknowledge' || api === 'modifyAckDeadline') {
+      if (request && request.ackIds && request.ackIds.length > 0) {
+        // Try to find a stored context for any of these ack IDs
+        for (const ackId of request.ackIds) {
+          const storedContext = ackContextMap.get(ackId)
+          if (storedContext) {
+            restoredStore = storedContext
+            break
+          }
+        }
+        
+        // Clean up ackIds from the map ONLY for acknowledge, not modifyAckDeadline
+        // ModifyAckDeadline happens first (lease extension), then acknowledge happens later
+        if (api === 'acknowledge') {
+          request.ackIds.forEach(ackId => {
+            if (ackContextMap.has(ackId)) {
+              ackContextMap.delete(ackId)
+            }
+          })
+        }
+      }
+    }
+
     const ctx = { request, api, projectId: this.auth._cachedProjectId }
+    
+    // If we have a restored context, run in that context
+    if (restoredStore) {
+      // CRITICAL: Add the parent span to ctx so the plugin uses it as parent
+      const parentSpan = restoredStore.span
+      if (parentSpan) {
+        ctx.parentSpan = parentSpan
+      }
+      const self = this
+      const args = arguments
+      return storage('legacy').run(restoredStore, () => {
+        return requestStartCh.runStores(ctx, () => {
+          const cb = args[args.length - 1]
+
+          if (typeof cb === 'function') {
+            args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (error) {
+              if (error) {
+                ctx.error = error
+                requestErrorCh.publish(ctx)
+              }
+              return requestFinishCh.runStores(ctx, cb, this, ...arguments)
+            })
+            return method.apply(self, args)
+          }
+
+          return method.apply(self, args)
+            .then(
+              response => {
+                requestFinishCh.publish(ctx)
+                return response
+              },
+              error => {
+                ctx.error = error
+                requestErrorCh.publish(ctx)
+                requestFinishCh.publish(ctx)
+                throw error
+              }
+            )
+        })
+      })
+    }
+
+    // Otherwise run normally
     return requestStartCh.runStores(ctx, () => {
       const cb = arguments[arguments.length - 1]
 
@@ -123,7 +196,12 @@ addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
   shimmer.wrap(Subscription.prototype, 'emit', emit => function (eventName, message) {
     if (eventName !== 'message' || !message) return emit.apply(this, arguments)
 
-    const ctx = {}
+    // Get the current async context store (should contain the pubsub.delivery span)
+    const store = storage('legacy').getStore()
+    
+    // If we have a span in the store, the context is properly set up
+    // The user's message handler will now run in this context and see the active span
+    const ctx = { message, store }
     try {
       return emit.apply(this, arguments)
     } catch (err) {
@@ -136,6 +214,40 @@ addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
   return obj
 })
 
+// Wrap message.ack() - must hook the subscriber-message.js file directly since Message is not exported from main module
+addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/subscriber.js' }, (obj) => {
+  const Message = obj.Message
+  
+  
+  if (Message && Message.prototype && Message.prototype.ack) {
+    shimmer.wrap(Message.prototype, 'ack', originalAck => function () {
+      // Capture the current active span and create a store with its context
+      const currentStore = storage('legacy').getStore()
+      const activeSpan = currentStore && currentStore.span
+      
+      if (activeSpan) {
+        
+        // CRITICAL: We must store a context that reflects the span's actual trace
+        // The span might have been created with a custom parent (reparented to pubsub.request)
+        // but the async storage might still contain the original context.
+        // So we create a fresh store with the span to ensure the correct trace is preserved.
+        const storeWithSpanContext = { ...currentStore, span: activeSpan }
+        
+        if (this.ackId) {
+          ackContextMap.set(this.ackId, storeWithSpanContext)
+        }
+      } else {
+      }
+      
+      return originalAck.apply(this, arguments)
+    })
+    
+  } else {
+  }
+
+  return obj
+})
+
 addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/lease-manager.js' }, (obj) => {
   const LeaseManager = obj.LeaseManager
   const ctx = {}

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

@@ -12,15 +12,23 @@ class GoogleCloudPubsubClientPlugin extends ClientPlugin {
 
     if (api === 'publish') return
 
-    this.startSpan(this.operationName(), {
+    const explicitParent = ctx.parentSpan // From restored context in wrapMethod
+    const spanOptions = {
       service: this.config.service || this.serviceName(),
       resource: [api, request.name].filter(Boolean).join(' '),
       kind: this.constructor.kind,
       meta: {
         'pubsub.method': api,
         'gcloud.project_id': projectId
       }
-    }, ctx)
+    }
+    
+    // If we have an explicit parent span (from restored context), use it
+    if (explicitParent) {
+      spanOptions.childOf = explicitParent.context()
+    }
+
+    this.startSpan(this.operationName(), spanOptions, ctx)
 
     return ctx.currentStore
   }

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

@@ -87,15 +87,14 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
         '_dd.batch.size': String(messageCount),
         '_dd.batch.index': String(i),
         'gcloud.project_id': projectId,
-        'pubsub.topic': topic
+        'pubsub.topic': topic,
+        'x-dd-publish-start-time': String(Math.floor(batchSpan._startTime))
       })
 
       if (batchTraceIdUpper) {
         msg.attributes['_dd.pubsub_request.p.tid'] = batchTraceIdUpper
       }
 
-      msg.attributes['x-dd-publish-start-time'] ??= String(Date.now())
-
       if (this.config.dsmEnabled) {
         const dataStreamsContext = this.tracer.setCheckpoint(
           ['direction:out', `topic:${topic}`, 'type:google-pubsub'],