dd-trace 6.13.0 → 6.15.0

package/packages/datadog-instrumentations/src/helpers/pool-acquire.js

@@ -19,6 +19,9 @@ const { performance } = require('node:perf_hooks')
  * @typedef {{ deferred: boolean }} PoolWaitTransfer
  * @typedef {{ transfer: PoolWaitTransfer, waitTime: number }} DeferredPoolWait
  * @typedef {{ pending: boolean, pool: object }} DeferredPoolQueryAcquire
+ * @typedef {'mariadb'|'mysql'|'mysql2'|'pg'} DatabaseDriver
+ * @typedef {{ connection?: object, owner: object, pending: boolean }} PromisePoolQueryAcquire
+ * @typedef {{ waitTime?: number }} PoolAcquireCapture
  * @typedef {{
  *   connectionFinishCh: Channel,
  *   acquireStartCh: Channel,
@@ -41,6 +44,12 @@ let currentPoolQueryAcquire
 /** @type {import('node:async_hooks').AsyncLocalStorage<DeferredPoolQueryAcquire>|undefined} */
 let deferredPoolQueryAcquireStorage
 
+/** @type {import('node:async_hooks').AsyncLocalStorage<PromisePoolQueryAcquire>|undefined} */
+let promisePoolQueryAcquireStorage
+
+/** @type {PoolAcquireCapture|undefined} */
+let currentPoolAcquireCapture
+
 /** @type {WeakMap<object, ClusterAcquire>} */
 const deferredClusterAcquires = new WeakMap()
 
@@ -52,6 +61,110 @@ const poolWaitTimes = []
 const deferredPoolWaitTimes = new WeakMap()
 let deferredPoolWaitCount = 0
 
+/** @type {Record<string, { acquireStartCh: Channel, acquireFinishCh: Channel }>} */
+const promisePoolAcquireChannels = {}
+
+/** @type {WeakMap<object, Record<string, unknown>>} */
+const promisePoolConnectionConfigs = new WeakMap()
+
+const deferredPromisePoolWaitTransfer = { deferred: true }
+
+/**
+ * Mark a promise-returning ORM operation as the owner of its next pool acquisition.
+ *
+ * @param {Function} method
+ * @param {(receiver: object, args: ArgumentsLike) => object|undefined} resolveOwner
+ * @param {(receiver: object, args: ArgumentsLike) => DatabaseDriver|undefined} resolveDriver
+ * @returns {Function}
+ */
+function wrapPromisePoolQueryMethod (method, resolveOwner, resolveDriver) {
+  return function () {
+    const driver = resolveDriver(this, arguments)
+    if (driver === undefined || !getPromisePoolAcquireChannels(driver).acquireStartCh.hasSubscribers) {
+      return method.apply(this, arguments)
+    }
+
+    const owner = resolveOwner(this, arguments)
+    if (owner === undefined) return method.apply(this, arguments)
+
+    const storage = getPromisePoolQueryAcquireStorage()
+    return storage.run({ owner, pending: true }, () => method.apply(this, arguments))
+  }
+}
+
+/**
+ * Instrument a promise-returning ORM pool acquisition using the underlying driver contract.
+ *
+ * @param {Function} method
+ * @param {(receiver: object, args: ArgumentsLike) => DatabaseDriver|undefined} resolveDriver
+ * @param {(receiver: object, args: ArgumentsLike) => Record<string, unknown>} resolveConfig
+ * @param {(receiver: object, args: ArgumentsLike) => boolean} hasIdleConnection
+ * @returns {Function}
+ */
+function wrapPromisePoolAcquire (method, resolveDriver, resolveConfig, hasIdleConnection) {
+  return function () {
+    const driver = resolveDriver(this, arguments)
+    if (driver === undefined) return method.apply(this, arguments)
+
+    const channels = getPromisePoolAcquireChannels(driver)
+    if (!channels.acquireStartCh.hasSubscribers) return method.apply(this, arguments)
+
+    const queryAcquire = promisePoolQueryAcquireStorage?.getStore()
+    const internal = queryAcquire?.owner === this && queryAcquire.pending
+    if (internal) queryAcquire.pending = false
+    const start = hasIdleConnection(this, arguments) ? undefined : performance.now()
+    const capture = {}
+    let ctx
+
+    if (!internal) {
+      ctx = createPromisePoolAcquireContext(driver, resolveConfig(this, arguments))
+      channels.acquireStartCh.publish(ctx)
+    }
+
+    let result
+    try {
+      result = runWithPoolAcquireCapture(capture, method, this, arguments)
+    } catch (error) {
+      if (internal) queryAcquire.pending = true
+      finishPromisePoolAcquire(driver, channels, ctx, start, capture, error, resolveConfig(this, arguments))
+      throw error
+    }
+
+    return result.then(connection => {
+      const waitTime = capture.waitTime ?? acquireWait(start)
+      if (internal) {
+        queryAcquire.connection = connection
+        setPoolWaitTime(connection, waitTime)
+      } else {
+        finishPromisePoolAcquireContext(driver, channels, ctx, connection, waitTime)
+      }
+      return connection
+    }, error => {
+      if (internal) queryAcquire.pending = true
+      finishPromisePoolAcquire(driver, channels, ctx, start, capture, error, resolveConfig(this, arguments))
+      throw error
+    })
+  }
+}
+
+/**
+ * Clear an acquisition that was never consumed by a driver query.
+ *
+ * @param {Function} method
+ * @returns {Function}
+ */
+function wrapPoolRelease (method) {
+  return function (connection) {
+    const queryAcquire = promisePoolQueryAcquireStorage?.getStore()
+    if (queryAcquire?.owner === this && queryAcquire.connection === connection) {
+      queryAcquire.connection = undefined
+      queryAcquire.pending = true
+    }
+    if (connection !== undefined) clearPoolWaitTime(connection)
+    return method.apply(this, arguments)
+  }
+}
+
 /**
  * @param {Function} method
  * @param {Channel} connectionStartCh
@@ -194,6 +307,25 @@ function wrapSynchronousPoolGetConnection (getConnection, channels) {
 
     const pool = this
     const ctx = {}
+    const capture = currentPoolAcquireCapture
+    if (capture !== undefined) {
+      const start = acquireStart(pool)
+
+      arguments[0] = function () {
+        capture.waitTime = acquireWait(start)
+        return connectionFinishCh.runStores(ctx, callback, this, ...arguments)
+      }
+      if (queuedCallbacks !== undefined) queuedCallbacks.add(arguments[0])
+
+      connectionStartCh.publish(ctx)
+      try {
+        return getConnection.apply(pool, arguments)
+      } catch (error) {
+        capture.waitTime = acquireWait(start)
+        throw error
+      }
+    }
+
     const acquire = currentPoolQueryAcquire
     const acquireCtx = acquire === undefined && acquireStartCh.hasSubscribers
       ? { conf: pool.config.connectionConfig }
@@ -325,6 +457,123 @@ function getDeferredPoolQueryAcquireStorage () {
   return deferredPoolQueryAcquireStorage
 }
 
+/**
+ * @returns {import('node:async_hooks').AsyncLocalStorage<PromisePoolQueryAcquire>}
+ */
+function getPromisePoolQueryAcquireStorage () {
+  if (promisePoolQueryAcquireStorage === undefined) {
+    const { AsyncLocalStorage } = require('node:async_hooks')
+    promisePoolQueryAcquireStorage = new AsyncLocalStorage()
+  }
+  return promisePoolQueryAcquireStorage
+}
+
+/**
+ * @param {DatabaseDriver} driver
+ * @returns {{ acquireStartCh: Channel, acquireFinishCh: Channel }}
+ */
+function getPromisePoolAcquireChannels (driver) {
+  let channels = promisePoolAcquireChannels[driver]
+  if (channels === undefined) {
+    const { channel } = require('./instrument')
+    channels = {
+      acquireStartCh: channel(`apm:${driver}:pool:acquire:start`),
+      acquireFinishCh: channel(`apm:${driver}:pool:acquire:finish`),
+    }
+    promisePoolAcquireChannels[driver] = channels
+  }
+  return channels
+}
+
+/**
+ * @param {DatabaseDriver} driver
+ * @param {Record<string, unknown>} config
+ * @returns {Record<string, unknown>}
+ */
+function createPromisePoolAcquireContext (driver, config) {
+  return driver === 'pg' ? { poolOptions: config } : { conf: config }
+}
+
+/**
+ * @param {DatabaseDriver} driver
+ * @param {{ acquireFinishCh: Channel }} channels
+ * @param {Record<string, unknown>|undefined} ctx
+ * @param {{
+ *   config?: Record<string, unknown>,
+ *   connectionParameters?: Record<string, unknown>,
+ *   opts?: Record<string, unknown>
+ * }} connection
+ * @param {number} waitTime
+ */
+function finishPromisePoolAcquireContext (driver, channels, ctx, connection, waitTime) {
+  ctx.poolWaitTime = waitTime
+  if (driver === 'pg') ctx.params = connection.connectionParameters
+  else if (driver === 'mariadb') ctx.connectionConfig = promisePoolConnectionConfigs.get(connection)
+  else ctx.connectionConfig = connection.config
+  channels.acquireFinishCh.publish(ctx)
+}
+
+/**
+ * @param {object} connection
+ * @param {Record<string, unknown>} config
+ */
+function setPromisePoolConnectionConfig (connection, config) {
+  promisePoolConnectionConfigs.set(connection, config)
+}
+
+/**
+ * @param {DatabaseDriver} driver
+ * @param {{ acquireStartCh: Channel, acquireFinishCh: Channel }} channels
+ * @param {Record<string, unknown>|undefined} ctx
+ * @param {number|undefined} start
+ * @param {PoolAcquireCapture} capture
+ * @param {unknown} error
+ * @param {Record<string, unknown>} config
+ */
+function finishPromisePoolAcquire (driver, channels, ctx, start, capture, error, config) {
+  const waitTime = capture.waitTime ?? acquireWait(start)
+  if (ctx === undefined) {
+    ctx = createPromisePoolAcquireContext(driver, config)
+    if (start !== undefined) ctx.startTime = performance.timeOrigin + start
+    channels.acquireStartCh.publish(ctx)
+  }
+  ctx.error = error
+  ctx.poolWaitTime = waitTime
+  channels.acquireFinishCh.publish(ctx)
+}
+
+/**
+ * @param {PoolAcquireCapture} capture
+ * @param {Function} method
+ * @param {object} receiver
+ * @param {ArgumentsLike} args
+ * @returns {unknown}
+ */
+function runWithPoolAcquireCapture (capture, method, receiver, args) {
+  const previous = currentPoolAcquireCapture
+  currentPoolAcquireCapture = capture
+  try {
+    return method.apply(receiver, args)
+  } finally {
+    currentPoolAcquireCapture = previous
+  }
+}
+
+/**
+ * @returns {PoolAcquireCapture|undefined}
+ */
+function getPoolAcquireCapture () {
+  return currentPoolAcquireCapture
+}
+
+/**
+ * @param {PoolAcquireCapture} capture
+ * @param {number} waitTime
+ */
+function setPoolAcquireCaptureWaitTime (capture, waitTime) {
+  capture.waitTime = waitTime
+}
+
 /**
  * @param {Function} method
  * @param {Record<string, unknown>} receiver
@@ -436,6 +685,15 @@ function clearPoolWaitTime (connection) {
   if (deferredPoolWaitCount !== 0 && deferredPoolWaitTimes.delete(connection)) deferredPoolWaitCount--
 }
 
+/**
+ * @param {object} connection
+ * @param {number} waitTime
+ */
+function setPoolWaitTime (connection, waitTime) {
+  clearPoolWaitTime(connection)
+  deferPoolWaitTime(deferredPromisePoolWaitTransfer, connection, waitTime)
+}
+
 /**
  * @param {PoolWaitTransfer} transfer
  * @param {object} connection
@@ -620,11 +878,15 @@ module.exports = {
   acquireWait,
   clearPoolWaitTime,
   dispatchesAcquireSynchronously,
+  getPoolAcquireCapture,
   isPoolQueryAcquire,
   reportPoolAcquireError,
   runOutsidePoolQueryAcquire,
   runPoolAcquireError,
   runWithPoolWait,
+  setPoolAcquireCaptureWaitTime,
+  setPoolWaitTime,
+  setPromisePoolConnectionConfig,
   takePoolWaitTime,
   wrapPoolAcquireCarrier,
   wrapPoolClusterGetConnection,
@@ -632,4 +894,7 @@ module.exports = {
   wrapPoolClusterQueryMethod,
   wrapPoolGetConnection,
   wrapPoolQueryMethod,
+  wrapPoolRelease,
+  wrapPromisePoolAcquire,
+  wrapPromisePoolQueryMethod,
 }

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

@@ -16,8 +16,7 @@ const runtimeRequire = typeof __webpack_require__ === 'function' ? __non_webpack
 const compiler = {
   parse: (sourceText, options) => {
     try {
-      // TODO: Figure out ESBuild `createRequire` issue and remove this hack.
-      const oxc = runtimeRequire(['oxc', 'parser'].join('-'))
+      const oxc = runtimeRequire('oxc-parser')
 
       compiler.parse = (sourceText, { range, isModule } = {}) => {
         const { program, errors } = oxc.parseSync('index.js', sourceText, {

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

@@ -4,40 +4,32 @@ const { readFileSync } = require('fs')
 const { join } = require('path')
 const { pathToFileURL } = require('url')
 const log = require('../../../../dd-trace/src/log')
-const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
 const instrumentations = require('./instrumentations')
 const { getRewriteTarget } = require('./targets')
-const { awaitContextCallback, waitForAsyncEnd } = require('./transforms')
-
-// `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
-// (ESM) statements that the transformer splices into the rewritten module.
-// `require()` accepts an absolute filesystem path; the ESM resolver rejects it
-// with `ERR_INVALID_MODULE_SPECIFIER` and needs a `file://` URL instead. We
-// pre-compute both forms here so each matcher hands the transformer a
-// specifier that is valid for the module type it is rewriting.
-let dcPolyfillCjs
-let dcPolyfillEsm
-
-try {
-  const resolved = require.resolve('dc-polyfill')
-  dcPolyfillCjs = resolved.replaceAll('\\', '/')
-  dcPolyfillEsm = pathToFileURL(resolved).href
-} catch {
-  // The `dc-polyfill` module is unavailable for some reason (like bundling).
-  // Let's just keep the default of using `diagnostics-channel` as a fallback
-  // which works for most Node versions.
-}
 
-/** @type {Record<string, string>} map of module base name to version */
+/**
+ * @typedef {object} InstrumentationMatcher
+ * @property {(name: string, transform: Function) => void} addTransform
+ * @property {(moduleName: string, version: string|undefined, filePath: string) => Transformer|undefined} getTransformer
+ *
+ * @typedef {object} Transformer
+ * @property {(source: string, moduleType: 'cjs'|'esm') => { code: string, map?: string }} transform
+ */
+
+/**
+ * @type {Record<string, string>} map of module base name to version
+ */
 const moduleVersions = {}
 const disabled = new Set()
-const matcherCjs = create(instrumentations, dcPolyfillCjs)
-const matcherEsm = create(instrumentations, dcPolyfillEsm)
 
-for (const matcher of [matcherCjs, matcherEsm]) {
-  matcher.addTransform('awaitContextCallback', awaitContextCallback)
-  matcher.addTransform('waitForAsyncEnd', waitForAsyncEnd)
-}
+// Matchers are built on the first module that actually needs rewriting. The
+// vendored transformer is a quarter megabyte of bundle that an application
+// without a rewrite target never needs to parse, and an application that loads
+// targets of only one module type never needs the other matcher.
+/** @type {InstrumentationMatcher|undefined} */
+let matcherCjs
+/** @type {InstrumentationMatcher|undefined} */
+let matcherEsm
 
 // Keep the marker split: source-map scanners can read a contiguous token in
 // string literals as this file's own inline map.
@@ -65,8 +57,7 @@ function rewrite (content, filename, format, target) {
 
   if (disabled.has(moduleName)) return content
 
-  const matcher = moduleType === 'esm' ? matcherEsm : matcherCjs
-  const transformer = matcher.getTransformer(moduleName, version, filePath)
+  const transformer = getMatcher(moduleType).getTransformer(moduleName, version, filePath)
 
   if (!transformer) return content
 
@@ -88,6 +79,80 @@ function rewrite (content, filename, format, target) {
   return content
 }
 
+/**
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {InstrumentationMatcher}
+ */
+function getMatcher (moduleType) {
+  if (moduleType === 'esm') {
+    matcherEsm ??= createMatcher(moduleType)
+
+    return matcherEsm
+  }
+
+  matcherCjs ??= createMatcher(moduleType)
+
+  return matcherCjs
+}
+
+/**
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {InstrumentationMatcher}
+ */
+function createMatcher (moduleType) {
+  const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
+  const {
+    awaitContextCallback,
+    awaitContextCallbackAtFunctionStart,
+    awaitContextCallbackAtTryStart,
+    configureGraphqlFastPath,
+    configureGraphqlJitCompileObject,
+    configureGraphqlJitDeferredField,
+    configureGraphqlJitExecute,
+    configureGraphqlJitRuntime,
+    configureMercuriusRequest,
+    waitForAsyncEnd,
+  } = require('./transforms')
+
+  const matcher = create(instrumentations, getDcPolyfillSpecifier(moduleType))
+
+  matcher.addTransform('awaitContextCallback', awaitContextCallback)
+  matcher.addTransform('awaitContextCallbackAtFunctionStart', awaitContextCallbackAtFunctionStart)
+  matcher.addTransform('awaitContextCallbackAtTryStart', awaitContextCallbackAtTryStart)
+  matcher.addTransform('waitForAsyncEnd', waitForAsyncEnd)
+  matcher.addTransform('configureGraphqlFastPath', configureGraphqlFastPath)
+  matcher.addTransform('configureGraphqlJitCompileObject', configureGraphqlJitCompileObject)
+  matcher.addTransform('configureGraphqlJitDeferredField', configureGraphqlJitDeferredField)
+  matcher.addTransform('configureGraphqlJitExecute', configureGraphqlJitExecute)
+  matcher.addTransform('configureGraphqlJitRuntime', configureGraphqlJitRuntime)
+  matcher.addTransform('configureMercuriusRequest', configureMercuriusRequest)
+
+  return matcher
+}
+
+/**
+ * `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
+ * (ESM) statements that the transformer splices into the rewritten module.
+ * `require()` accepts an absolute filesystem path; the ESM resolver rejects it
+ * with `ERR_INVALID_MODULE_SPECIFIER` and needs a `file://` URL instead. Each
+ * matcher therefore hands the transformer the form that is valid for the
+ * module type it is rewriting.
+ *
+ * @param {'cjs'|'esm'} moduleType
+ * @returns {string|undefined} `undefined` when `dc-polyfill` cannot be resolved
+ */
+function getDcPolyfillSpecifier (moduleType) {
+  try {
+    const resolved = require.resolve('dc-polyfill')
+
+    return moduleType === 'esm' ? pathToFileURL(resolved).href : resolved.replaceAll('\\', '/')
+  } catch {
+    // The `dc-polyfill` module is unavailable for some reason (like bundling).
+    // Let's just keep the default of using `diagnostics-channel` as a fallback
+    // which works for most Node versions.
+  }
+}
+
 /** @typedef {{ buffer: ArrayBuffer | SharedArrayBuffer, byteLength: number, byteOffset: number }} BufferView */
 
 /**

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

@@ -100,29 +100,4 @@ module.exports = [
     },
     channelName: 'resolveLanguageModel',
   },
-  // tool
-  {
-    module: {
-      name: 'ai',
-      versionRange: '>=4.0.0 <7.0.0',
-      filePath: 'dist/index.js',
-    },
-    functionQuery: {
-      functionName: 'tool',
-      kind: 'Sync',
-    },
-    channelName: 'tool',
-  },
-  {
-    module: {
-      name: 'ai',
-      versionRange: '>=4.0.0 <7.0.0',
-      filePath: 'dist/index.mjs',
-    },
-    functionQuery: {
-      functionName: 'tool',
-      kind: 'Sync',
-    },
-    channelName: 'tool',
-  },
 ]

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/graphql-jit.js

@@ -0,0 +1,83 @@
+'use strict'
+
+// Match the private closure only while its captured identifiers retain the verified
+// 0.8 contract. A future rewrite then misses safely instead of injecting stale names.
+const CREATE_BOUND_QUERY = 'FunctionDeclaration[id.name="createBoundQuery"]' +
+  '[params.0.name="compilationContext"][params.1.name="document"][params.4.name="operationName"]'
+const WRAP_QUERY = `${CREATE_BOUND_QUERY} ` +
+  'VariableDeclarator[id.name="ret"] > ObjectExpression > Property > FunctionExpression'
+const BUILD_COMPILATION_CONTEXT = 'FunctionDeclaration[id.name="buildCompilationContext"]' +
+  '[params.0.name="schema"][params.1.name="document"][params.2.name="options"][params.3.name="operationName"]'
+const COMPILE_OBJECT_TYPE = 'FunctionDeclaration[id.name="compileObjectType"]' +
+  '[params.0.name="context"][params.1.name="type"][params.2.name="fieldNodes"]' +
+  '[params.3.name="originPaths"][params.5.name="responsePath"][params.8.name="alwaysDefer"]' +
+  ':has(VariableDeclarator[id.name="field"])'
+const COMPILE_DEFERRED_FIELD = 'FunctionDeclaration[id.name="compileDeferredField"]' +
+  '[params.0.name="context"][params.1.name="deferredField"]'
+
+/**
+ * @typedef {{
+ *   module: { name: string, versionRange: string, filePath: string },
+ *   astQuery: string,
+ *   functionQuery?: { kind: 'Sync' },
+ *   transform?: string,
+ *   channelName?: string
+ * }} GraphqlJitInstrumentation
+ */
+
+/**
+ * @param {GraphqlJitInstrumentation[]} instrumentations
+ * @param {string} versionRange
+ * @param {string} filePath
+ */
+function addInstrumentations (instrumentations, versionRange, filePath) {
+  const moduleDefinition = { name: 'graphql-jit', versionRange, filePath }
+  instrumentations.push(
+    {
+      module: moduleDefinition,
+      astQuery: BUILD_COMPILATION_CONTEXT,
+      functionQuery: { kind: 'Sync' },
+      channelName: 'compile',
+    },
+    {
+      module: moduleDefinition,
+      astQuery: COMPILE_OBJECT_TYPE,
+      transform: 'configureGraphqlJitCompileObject',
+    },
+    {
+      module: moduleDefinition,
+      astQuery: COMPILE_DEFERRED_FIELD,
+      transform: 'configureGraphqlJitDeferredField',
+    },
+    {
+      module: moduleDefinition,
+      astQuery: CREATE_BOUND_QUERY,
+      transform: 'configureGraphqlJitRuntime',
+    },
+    {
+      module: moduleDefinition,
+      astQuery: WRAP_QUERY,
+      functionQuery: { kind: 'Sync' },
+      channelName: 'apm:graphql:execute',
+    },
+    {
+      module: moduleDefinition,
+      astQuery: WRAP_QUERY,
+      transform: 'configureGraphqlJitExecute',
+      channelName: 'apm:graphql:execute',
+    }
+  )
+}
+
+/**
+ * @type {GraphqlJitInstrumentation[]}
+ */
+const instrumentations = []
+// graphql-jit 0.7 and 0.8.0 use this layout and remain compatible with GraphQL 15 and 16.
+// GraphQL 17 compatibility starts at 0.8.1.
+addInstrumentations(instrumentations, '>=0.7.0 <0.8.5 || >=0.8.7 <0.9.0', 'dist/execution.js')
+addInstrumentations(instrumentations, '>=0.8.5 <0.8.7', 'dist/cjs/execution.js')
+addInstrumentations(instrumentations, '>=0.8.5 <0.8.7', 'dist/esm/execution.js')
+addInstrumentations(instrumentations, '>=0.8.7 <0.9.0', 'dist/execution.mjs')
+
+module.exports = instrumentations