[SVLS-7168] Create GCP PubSub Push Subscriptions Plugin (#6260)

* Add plugin for GCP PubSub Push Subscriptions
* feat: add producer-side batch message handling with span linking
* feat: add span linking from delivery span to pubsub.request
* feat: add span linking and batch metadata to pull-based consumer

- Collect span links from messages 2-N (first becomes parent)
- Extract parent context from first message trace context
- Create pubsub.request span with span links metadata
- Inject batch metadata into all messages (_dd.pubsub_request.*, _dd.batch.*)
- Add 128-bit trace ID support (_dd.p.tid)
- Add operation tag for batched vs single requests
- HTTP plugin now checks if a delivery span is active in storage
- If found, uses delivery span as parent for http.request
- Ensures proper span hierarchy for push subscriptions
- Extract pubsub.request span ID from message attributes
- Add span link correlation tags
- Calculate delivery duration from publish start time
- Add batch size and index tags for batched messages

Author: nina9753

package.json

@@ -48,7 +48,7 @@
     "test:llmobs:plugins:ci": "yarn services && nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs:plugins",
     "test:openfeature": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/openfeature/*.spec.js\"",
     "test:openfeature:ci": "nyc --no-clean --include \"packages/dd-trace/src/openfeature/**/*.js\" -- npm run test:openfeature",
-    "test:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/@($(echo ${SPEC:-'*'}))*.spec.js\"",
+    "test:plugins": "mocha --expose-gc -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/@($(echo ${SPEC:-'*'}))*.spec.js\"",
     "test:plugins:ci": "yarn services && nyc --no-clean --include \"packages/datadog-plugin-@($(echo $PLUGINS))/src/**/*.js\" -- npm run test:plugins",
     "test:plugins:ci:flaky": "yarn services && nyc --no-clean --include \"packages/datadog-plugin-@($(echo $PLUGINS))/src/**/*.js\" -- npm run test:plugins -- --bail --retries 2",
     "test:plugins:upstream": "node ./packages/dd-trace/test/plugins/suite.js",

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

@@ -14,6 +14,21 @@ 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')
 
+/**
+ * Message-level channels for trace context injection and async context propagation:
+ * - messagePublishCh: Inject trace context when Topic.publish() is called
+ * - messageAckStoreCh: Store async context when message.ack() is called (PULL)
+ * - messageAckRetrieveCh: Retrieve stored context when acknowledge() API is called (PULL)
+ * - messageStoreCh: Propagate context when Subscription emits 'message' event (PULL)
+ */
+const messagePublishCh = channel('apm:google-cloud-pubsub:message:publish')
+const messageAckStoreCh = channel('apm:google-cloud-pubsub:message:ack-store')
+const messageAckRetrieveCh = channel('apm:google-cloud-pubsub:message:ack-retrieve')
+const messageStoreCh = channel('apm:google-cloud-pubsub:message:store')
+
+// WeakMap for passing context between LeaseManager._dispense and LeaseManager.remove
+const messageContexts = new WeakMap()
+
 const publisherMethods = [
   'createTopic',
   'updateTopic',
@@ -65,6 +80,21 @@ function wrapMethod (method) {
     if (!requestStartCh.hasSubscribers) return method.apply(this, arguments)
 
     const ctx = { request, api, projectId: this.auth._cachedProjectId }
+
+    /**
+     * For acknowledge/modifyAckDeadline: retrieve stored context from consumer plugin.
+     * These APIs only have ackIds (no Message objects), so async context is lost.
+     * Plugin sets ctx.storedContext, which client.js uses to link the acknowledge span.
+     */
+    const isAckOperation = api === 'acknowledge' || api === 'modifyAckDeadline'
+    if (isAckOperation && request?.ackIds?.length > 0) {
+      messageAckRetrieveCh.publish({
+        ackIds: request.ackIds,
+        api,
+        ctx
+      })
+    }
+
     return requestStartCh.runStores(ctx, () => {
       const cb = arguments[arguments.length - 1]
 
@@ -74,12 +104,11 @@ function wrapMethod (method) {
             ctx.error = error
             requestErrorCh.publish(ctx)
           }
-
           return requestFinishCh.runStores(ctx, cb, this, ...arguments)
         })
-
         return method.apply(this, arguments)
       }
+
       return method.apply(this, arguments)
         .then(
           response => {
@@ -108,12 +137,16 @@ function massWrap (obj, methods, wrapper) {
 addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
   const Subscription = obj.Subscription
 
+  /**
+   * PULL: Intercept 'message' events and propagate async context via runStores.
+   * This ensures the consumer plugin creates spans in the correct context.
+   */
   shimmer.wrap(Subscription.prototype, 'emit', emit => function (eventName, message) {
     if (eventName !== 'message' || !message) return emit.apply(this, arguments)
 
-    const ctx = {}
+    const ctx = { message }
     try {
-      return emit.apply(this, arguments)
+      return messageStoreCh.runStores(ctx, emit, this, ...arguments)
     } catch (err) {
       ctx.error = err
       receiveErrorCh.publish(ctx)
@@ -124,33 +157,122 @@ addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
   return obj
 })
 
+addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/subscriber.js' }, (obj) => {
+  const Message = obj.Message
+
+  /**
+   * PULL: Capture async context when message.ack() is called.
+   * This is our last chance before context is lost. The acknowledge() API call happens
+   * later (often batched), and we'll retrieve this stored context to link spans.
+   * Flow: message.ack() -> store context -> acknowledge() API -> retrieve context
+   */
+  if (Message?.prototype?.ack) {
+    shimmer.wrap(Message.prototype, 'ack', originalAck => function () {
+      if (this.ackId) {
+        const ctx = {
+          message: this,
+          ackId: this.ackId
+        }
+
+        return messageAckStoreCh.runStores(ctx, originalAck, this, ...arguments)
+      }
+
+      return originalAck.apply(this, arguments)
+    })
+  }
+
+  return obj
+})
+
+/**
+ * PULL: Hook LeaseManager to track message lifecycle (dispense/remove/clear).
+ * _dispense: Message given to handler -> create span
+ * remove: Message removed from lease (ack/nack/timeout) -> finish span
+ * clear: Subscription closed -> finish all remaining spans
+ */
 addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/lease-manager.js' }, (obj) => {
   const LeaseManager = obj.LeaseManager
-  const ctx = {}
+  if (!LeaseManager) {
+    return obj
+  }
 
   shimmer.wrap(LeaseManager.prototype, '_dispense', dispense => function (message) {
-    if (receiveStartCh.hasSubscribers) {
-      ctx.message = message
-      return receiveStartCh.runStores(ctx, dispense, this, ...arguments)
-    }
-    return dispense.apply(this, arguments)
+    const ctx = { message }
+    messageContexts.set(message, ctx)
+
+    return receiveStartCh.runStores(ctx, dispense, this, ...arguments)
   })
 
   shimmer.wrap(LeaseManager.prototype, 'remove', remove => function (message) {
-    return receiveFinishCh.runStores(ctx, remove, this, ...arguments)
+    const ctx = messageContexts.get(message)
+    if (ctx) {
+      messageContexts.delete(message)
+    }
+
+    return receiveFinishCh.runStores(ctx || { message }, remove, this, ...arguments)
   })
 
   shimmer.wrap(LeaseManager.prototype, 'clear', clear => function () {
-    for (const message of this._messages) {
-      ctx.message = message
-      receiveFinishCh.publish(ctx)
+    if (this._messages) {
+      for (const message of this._messages.values()) {
+        const ctx = messageContexts.get(message)
+        if (ctx) {
+          receiveFinishCh.publish(ctx)
+          messageContexts.delete(message)
+        }
+      }
     }
     return clear.apply(this, arguments)
   })
 
   return obj
 })
 
+/**
+ * Inject trace context into individual messages via Topic.publish()/publishMessage().
+ * Flow: User calls topic.publish() -> inject context (here) -> SDK batches messages ->
+ * publish() API called -> producer plugin creates batch span + metadata
+ */
+addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
+  if (!obj.Topic?.prototype) return obj
+
+  if (typeof obj.Topic.prototype.publishMessage === 'function') {
+    shimmer.wrap(obj.Topic.prototype, 'publishMessage', publishMessage => {
+      return function (data, attributesOrCallback, callback) {
+        if (data && typeof data === 'object') {
+          if (!data.attributes) data.attributes = {}
+          messagePublishCh.publish({
+            attributes: data.attributes,
+            pubsub: this.pubsub,
+            topicName: this.name
+          })
+        }
+        return publishMessage.apply(this, arguments)
+      }
+    })
+  }
+
+  if (typeof obj.Topic.prototype.publish === 'function') {
+    shimmer.wrap(obj.Topic.prototype, 'publish', publish => function (buffer, attributesOrCallback, callback) {
+      if (typeof attributesOrCallback === 'function' || !attributesOrCallback) {
+        arguments[1] = {}
+        arguments[2] = attributesOrCallback
+      }
+
+      messagePublishCh.publish({
+        attributes: arguments[1],
+        pubsub: this.pubsub,
+        topicName: this.name,
+        buffer
+      })
+
+      return publish.apply(this, arguments)
+    })
+  }
+
+  return obj
+})
+
 addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'] }, (obj) => {
   const { PublisherClient, SchemaServiceClient, SubscriberClient } = obj.v1
 

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

@@ -8,19 +8,29 @@ 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,
       meta: {
         '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
   }