dd-trace 5.80.0 → 5.82.0

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

@@ -1,10 +1,10 @@
 'use strict'
 
+const shimmer = require('../../datadog-shimmer')
 const {
   channel,
   addHook
 } = require('./helpers/instrument')
-const shimmer = require('../../datadog-shimmer')
 
 const requestStartCh = channel('apm:google-cloud-pubsub:request:start')
 const requestFinishCh = channel('apm:google-cloud-pubsub:request:finish')
@@ -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,26 +157,70 @@ 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)
   })
@@ -151,6 +228,51 @@ addHook({ name: '@google-cloud/pubsub', versions: ['>=1.2'], file: 'build/src/le
   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
 

package/packages/datadog-instrumentations/src/google-cloud-vertexai.js

@@ -1,9 +1,8 @@
 'use strict'
 
-const { addHook } = require('./helpers/instrument')
-const shimmer = require('../../datadog-shimmer')
-
 const vertexaiTracingChannel = require('dc-polyfill').tracingChannel('apm:vertexai:request')
+const shimmer = require('../../datadog-shimmer')
+const { addHook } = require('./helpers/instrument')
 
 function wrapGenerate (generate) {
   return function (request) {

package/packages/datadog-instrumentations/src/google-genai.js

@@ -0,0 +1,120 @@
+'use strict'
+
+const { channel, tracingChannel } = require('dc-polyfill')
+
+const shimmer = require('../../datadog-shimmer')
+const { addHook } = require('./helpers/instrument')
+
+const genaiTracingChannel = tracingChannel('apm:google:genai:request')
+const onStreamedChunkCh = channel('apm:google:genai:request:chunk')
+
+function wrapGenerateContent (method) {
+  return function wrappedGenerateContent (original) {
+    return function (...args) {
+      if (!genaiTracingChannel.start.hasSubscribers) {
+        return original.apply(this, args)
+      }
+
+      const normalizedName = normalizeMethodName(method)
+
+      const ctx = { args, methodName: normalizedName }
+
+      return genaiTracingChannel.start.runStores(ctx, () => {
+        let result
+        try {
+          result = original.apply(this, arguments)
+        } catch (error) {
+          finish(ctx, null, error)
+          throw error
+        } finally {
+          genaiTracingChannel.end.publish(ctx)
+        }
+        return result.then(response => {
+          if (response[Symbol.asyncIterator]) {
+            shimmer.wrap(response, Symbol.asyncIterator, iterator => wrapStreamIterator(iterator, ctx))
+          } else {
+            finish(ctx, response, null)
+          }
+          return response
+        }).catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+      })
+    }
+  }
+}
+
+function wrapStreamIterator (iterator, ctx) {
+  return function () {
+    const itr = iterator.apply(this, arguments)
+    shimmer.wrap(itr, 'next', next => function () {
+      return next.apply(this, arguments)
+        .then(res => {
+          const { done, value: chunk } = res
+          onStreamedChunkCh.publish({ ctx, chunk, done })
+
+          if (done) {
+            finish(ctx)
+          }
+
+          return res
+        })
+        .catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+    })
+
+    return itr
+  }
+}
+function finish (ctx, result, error) {
+  if (error) {
+    ctx.error = error
+    genaiTracingChannel.error.publish(ctx)
+  }
+
+  // streamed responses are handled and set separately
+  ctx.result ??= result
+
+  genaiTracingChannel.asyncEnd.publish(ctx)
+}
+// Hook the main package entry point
+addHook({
+  name: '@google/genai',
+  versions: ['>=1.19.0']
+}, exports => {
+  // Wrap GoogleGenAI to intercept when it creates Models instances
+  if (!exports.GoogleGenAI) return exports
+
+  shimmer.wrap(exports, 'GoogleGenAI', GoogleGenAI => {
+    return class extends GoogleGenAI {
+      constructor (...args) {
+        super(...args)
+
+        // We are patching the instance instead of the prototype because when it is compiled from
+        // typescript, the models property is not available on the prototype.
+        if (this.models) {
+          if (this.models.generateContent) {
+            shimmer.wrap(this.models, 'generateContent', wrapGenerateContent('generateContent'))
+          }
+          if (this.models.generateContentStream) {
+            shimmer.wrap(this.models, 'generateContentStream', wrapGenerateContent('generateContentStream'))
+          }
+          if (this.models.embedContent) {
+            shimmer.wrap(this.models, 'embedContent', wrapGenerateContent('embedContent'))
+          }
+        }
+      }
+    }
+  })
+  return exports
+})
+
+function normalizeMethodName (methodName) {
+  // Convert camelCase to snake_case and add Models prefix
+  return 'Models.' + methodName
+    .replaceAll(/([a-z0-9])([A-Z])/g, '$1_$2')
+    .toLowerCase()
+}

package/packages/datadog-instrumentations/src/graphql.js

@@ -1,10 +1,12 @@
 'use strict'
 
+const shimmer = require('../../datadog-shimmer')
 const {
   addHook,
   channel
 } = require('./helpers/instrument')
-const shimmer = require('../../datadog-shimmer')
+
+const ddGlobal = globalThis[Symbol.for('dd-trace')] ??= {}
 
 /** cached objects */
 
@@ -360,3 +362,21 @@ addHook({ name: 'graphql', file: 'validation/validate.js', versions: ['>=0.10']
 
   return validate
 })
+
+addHook({ name: 'graphql', file: 'language/printer.js', versions: ['>=0.10'] }, printer => {
+  // HACK: It's possible `graphql` is loaded before `@apollo/gateway` so we
+  //       can't use a channel as the latter plugin would load after the publish
+  //       happened. Not sure how to handle this so for now use a global.
+  ddGlobal.graphql_printer = printer
+  return printer
+})
+
+addHook({ name: 'graphql', file: 'language/visitor.js', versions: ['>=0.10'] }, visitor => {
+  ddGlobal.graphql_visitor = visitor
+  return visitor
+})
+
+addHook({ name: 'graphql', file: 'utilities/index.js', versions: ['>=0.10'] }, utilities => {
+  ddGlobal.graphql_utilities = utilities
+  return utilities
+})

package/packages/datadog-instrumentations/src/grpc/client.js

@@ -1,8 +1,8 @@
 'use strict'
 
-const types = require('./types')
 const { addHook, channel } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
+const types = require('./types')
 
 const patched = new WeakSet()
 const instances = new WeakMap()

package/packages/datadog-instrumentations/src/grpc/server.js

@@ -1,8 +1,8 @@
 'use strict'
 
-const types = require('./types')
 const { channel, addHook } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
+const types = require('./types')
 
 const startChannel = channel('apm:grpc:server:request:start')
 const asyncStartChannel = channel('apm:grpc:server:request:asyncStart')

package/packages/datadog-instrumentations/src/helpers/bundler-register.js

@@ -2,6 +2,7 @@
 
 const dc = require('dc-polyfill')
 
+const log = require('../../../dd-trace/src/log')
 const {
   filename,
   loadChannel,
@@ -9,7 +10,6 @@ const {
 } = require('./register.js')
 const hooks = require('./hooks')
 const instrumentations = require('./instrumentations')
-const log = require('../../../dd-trace/src/log')
 
 const CHANNEL = 'dd-trace:bundler:load'
 

package/packages/datadog-instrumentations/src/helpers/hook.js

@@ -1,6 +1,6 @@
 'use strict'
-const iitm = require('../../../dd-trace/src/iitm')
 const path = require('path')
+const iitm = require('../../../dd-trace/src/iitm')
 const ritm = require('../../../dd-trace/src/ritm')
 
 /**

package/packages/datadog-instrumentations/src/helpers/hooks.js

@@ -13,6 +13,7 @@ module.exports = {
   '@playwright/test': () => require('../playwright'),
   '@elastic/elasticsearch': () => require('../elasticsearch'),
   '@elastic/transport': () => require('../elasticsearch'),
+  '@google/genai': () => require('../google-genai'),
   '@google-cloud/pubsub': () => require('../google-cloud-pubsub'),
   '@google-cloud/vertexai': () => require('../google-cloud-vertexai'),
   '@graphql-tools/executor': () => require('../graphql'),
@@ -28,7 +29,8 @@ module.exports = {
   '@node-redis/client': () => require('../redis'),
   '@opensearch-project/opensearch': () => require('../opensearch'),
   '@opentelemetry/sdk-trace-node': () => require('../otel-sdk-trace'),
-  '@prisma/client': () => require('../prisma'),
+  '@prisma/client': { esmFirst: true, fn: () => require('../prisma') },
+  './runtime/library.js': () => require('../prisma'),
   '@redis/client': () => require('../redis'),
   '@smithy/smithy-client': () => require('../aws-sdk'),
   '@vitest/runner': { esmFirst: true, fn: () => require('../vitest') },
@@ -40,6 +42,7 @@ module.exports = {
   'aws-sdk': () => require('../aws-sdk'),
   bluebird: () => require('../bluebird'),
   'body-parser': () => require('../body-parser'),
+  bullmq: () => require('../bullmq'),
   bunyan: () => require('../bunyan'),
   'cassandra-driver': () => require('../cassandra-driver'),
   child_process: () => require('../child_process'),
@@ -82,6 +85,7 @@ module.exports = {
   kafkajs: () => require('../kafkajs'),
   langchain: () => require('../langchain'),
   ldapjs: () => require('../ldapjs'),
+  'light-my-request': () => require('../light-my-request'),
   'limitd-client': () => require('../limitd-client'),
   lodash: () => require('../lodash'),
   mariadb: () => require('../mariadb'),

package/packages/datadog-instrumentations/src/helpers/instrument.js

@@ -1,8 +1,9 @@
 'use strict'
 
+const { AsyncResource } = require('async_hooks')
 const dc = require('dc-polyfill')
 const instrumentations = require('./instrumentations')
-const { AsyncResource } = require('async_hooks')
+const rewriterInstrumentations = require('./rewriter/instrumentations')
 
 const channelMap = {}
 exports.channel = function (name) {
@@ -22,6 +23,15 @@ exports.tracingChannel = function (name) {
   return tc
 }
 
+exports.getHooks = function getHooks (names) {
+  names = [names].flat()
+
+  return rewriterInstrumentations
+    .map(inst => inst.module)
+    .filter(({ name }) => names.includes(name))
+    .map(({ name, versionRange, filePath }) => ({ name, versions: [versionRange], file: filePath }))
+}
+
 /**
  * @param {object} args
  * @param {string|string[]} args.name module name

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -1,15 +1,17 @@
 'use strict'
 
-const { channel } = require('dc-polyfill')
 const path = require('path')
-const satisfies = require('semifies')
-const Hook = require('./hook')
+const { channel } = require('dc-polyfill')
+const satisfies = require('../../../../vendor/dist/semifies')
 const requirePackageJson = require('../../../dd-trace/src/require-package-json')
 const log = require('../../../dd-trace/src/log')
-const checkRequireCache = require('./check-require-cache')
 const telemetry = require('../../../dd-trace/src/guardrails/telemetry')
 const { isInServerlessEnvironment } = require('../../../dd-trace/src/serverless')
 const { getEnvironmentVariables } = require('../../../dd-trace/src/config-helper')
+const checkRequireCache = require('./check-require-cache')
+const Hook = require('./hook')
+const { isRelativeRequire } = require('./shared-utils')
+const rewriter = require('./rewriter')
 
 const envs = getEnvironmentVariables()
 
@@ -48,6 +50,10 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
 const seenCombo = new Set()
 const allInstrumentations = {}
 
+for (const inst of disabledInstrumentations) {
+  rewriter.disable(inst)
+}
+
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue
@@ -99,6 +105,8 @@ for (const packageName of names) {
       hook[HOOK_SYMBOL] ??= new WeakSet()
       let matchesFile = moduleName === fullFilename
 
+      if (!matchesFile && isRelativeRequire(name)) matchesFile = true
+
       if (fullFilePattern) {
         // Some libraries include a hash in their filenames when installed,
         // so our instrumentation has to include a '.*' to match them for more than a single version.
@@ -214,7 +222,6 @@ function parseHookInstrumentationFileName (packageName) {
     hook = hook.fn
   }
   const hookString = hook.toString()
-
   const regex = /require\('([^']*)'\)/
   const match = hookString.match(regex)
 

package/packages/datadog-instrumentations/src/helpers/rewriter/compiler.js

@@ -0,0 +1,27 @@
+'use strict'
+
+let meriyah
+let astring
+let esquery
+
+module.exports = {
+  parse: (...args) => {
+    meriyah ??= require('../../../../../vendor/dist/meriyah')
+
+    return meriyah.parse(...args)
+  },
+
+  generate: (...args) => {
+    astring ??= require('../../../../../vendor/dist/astring')
+
+    return astring.generate(...args)
+  },
+
+  traverse: (ast, query, visitor) => {
+    esquery ??= require('../../../../../vendor/dist/esquery').default
+
+    const selector = esquery.parse(query)
+
+    return esquery.traverse(ast, selector, visitor)
+  },
+}