dd-trace 6.13.0 → 6.14.0

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

@@ -1,5 +1,7 @@
 'use strict'
 
+const { isNativeError } = require('node:util').types
+
 const shimmer = require('../../datadog-shimmer')
 const nomenclature = require('../../dd-trace/src/service-naming')
 const spanEndingHook = require('../../dd-trace/src/opentelemetry/span-ending-hook')
@@ -12,11 +14,44 @@ const errorChannel = channel('apm:next:request:error')
 const pageLoadChannel = channel('apm:next:page:load')
 const bodyParsedChannel = channel('apm:next:body-parsed')
 const queryParsedChannel = channel('apm:next:query-parsed')
+const backgroundRevalidationChannel = channel('apm:next:request:background-revalidation')
+
+/**
+ * @typedef {import('node:http').IncomingMessage & {
+ *   error?: unknown,
+ *   originalRequest?: import('node:http').IncomingMessage
+ * }} NextNodeRequest
+ *
+ * @typedef {import('node:http').ServerResponse & {
+ *   originalResponse?: import('node:http').ServerResponse
+ * }} NextNodeResponse
+ *
+ * @typedef {object} NextRequest
+ * @property {unknown} [error]
+ *
+ * @typedef {object} NextRequestContext
+ * @property {NextNodeRequest} req
+ * @property {NextNodeResponse} res
+ * @property {boolean} [finishOnResponse]
+ * @property {boolean} [handlerFinished]
+ * @property {boolean} [responseFinished]
+ * @property {unknown} [error]
+ * @property {NextRequest} [nextRequest]
+ *
+ * @typedef {object} ResponseGeneratorContext
+ * @property {boolean} [hasResolved]
+ *
+ * @typedef {object} HandleResponseOptions
+ * @property {NextNodeRequest} req
+ * @property {(context?: ResponseGeneratorContext) => Promise<unknown>} responseGenerator
+ *
+ * @typedef {object} ResponseCacheEntry
+ * @property {{ status?: number }} [value]
+ */
 
 const requests = new WeakSet()
 const nodeNextRequestsToNextRequests = new WeakMap()
-const requestErrors = new WeakMap()
-
+const noFallbackErrorPrototypes = new WeakSet()
 // Next.js <= 14.2.6
 const MIDDLEWARE_HEADER = 'x-middleware-invoke'
 
@@ -149,7 +184,16 @@ function getRequestMeta (req, key) {
   return typeof key === 'string' ? meta[key] : meta
 }
 
-function instrument (req, res, handler, error) {
+/**
+ * @template T
+ * @param {NextNodeRequest} req
+ * @param {NextNodeResponse} res
+ * @param {(ctx: NextRequestContext) => Promise<T>} handler
+ * @param {unknown} [error]
+ * @param {boolean} [finishOnResponse]
+ * @returns {Promise<T>}
+ */
+function instrument (req, res, handler, error, finishOnResponse = false) {
   req = req.originalRequest || req
   res = res.originalResponse || res
 
@@ -165,7 +209,17 @@ function instrument (req, res, handler, error) {
 
   requests.add(req)
 
-  const ctx = { req, res }
+  const ctx = finishOnResponse ? { req, res, finishOnResponse } : { req, res }
+  if (finishOnResponse) {
+    const onResponseFinish = () => {
+      res.removeListener('finish', onResponseFinish)
+      res.removeListener('close', onResponseFinish)
+      ctx.responseFinished = true
+      if (ctx.handlerFinished) publishFinish(ctx, ctx.error)
+    }
+    res.once('finish', onResponseFinish)
+    res.once('close', onResponseFinish)
+  }
   if (queryParsedChannel.hasSubscribers && req.url) {
     const queryIndex = req.url.indexOf('?')
     if (queryIndex !== -1) {
@@ -213,9 +267,8 @@ function wrapServeStatic (serveStatic) {
   }
 }
 
-function finish (ctx, result, err) {
-  publishError(ctx.req, ctx, err)
-  requestErrors.delete(ctx.req)
+function publishFinish (ctx, error) {
+  publishError(ctx.req, ctx, error)
 
   const maybeNextRequest = nodeNextRequestsToNextRequests.get(ctx.req)
   if (maybeNextRequest) {
@@ -225,24 +278,35 @@ function finish (ctx, result, err) {
   finishChannel.publish(ctx)
 }
 
-function publishError (req, ctx, error) {
-  if (!error) return
+function finish (ctx, result, error) {
+  if (ctx.finishOnResponse) {
+    ctx.handlerFinished = true
+    if (error && !isNoFallbackError(error)) ctx.error = error
+    if (!ctx.responseFinished) return
+  }
 
-  req = req.originalRequest || req
-  let errors = requestErrors.get(req)
-  if (errors?.has(error)) return
+  publishFinish(ctx, error)
+}
 
-  if (!errors) {
-    errors = new Set()
-    requestErrors.set(req, errors)
-  }
-  errors.add(error)
+/** @param {unknown} error */
+function isNoFallbackError (error) {
+  return isNativeError(error) && noFallbackErrorPrototypes.has(Object.getPrototypeOf(error))
+}
+
+/**
+ * @param {NextNodeRequest} req
+ * @param {NextRequestContext | undefined} ctx
+ * @param {unknown} error
+ */
+function publishError (req, ctx, error) {
+  if (!error || isNoFallbackError(error)) return
 
   if (ctx) {
     ctx.error = error
     errorChannel.publish(ctx)
   } else {
-    errorChannel.publish({ error })
+    req = req.originalRequest || req
+    errorChannel.publish({ req, error })
   }
 }
 
@@ -267,13 +331,29 @@ addHook({
 // From Next 15.4.1, route modules execute through precompiled runtime bundles that bypass the
 // classic server hooks above. Match bundler and experimental filename variants without enumerating
 // them so App Routes, Pages APIs, and App Pages reuse the existing Next request lifecycle.
+// The route module classes and inherited methods are selected inside those bundles, so there is no
+// stable source function for Orchestrion to rewrite.
 const patchedRouteModules = new WeakSet()
+const routeResponses = new WeakMap()
+const activeRouteRequests = new WeakMap()
 const COMPILED_RUNTIME_PATH = 'dist/compiled/next-server/'
+
+/** @param {{ NoFallbackError: typeof Error }} noFallbackError */
+function captureNoFallbackError (noFallbackError) {
+  noFallbackErrorPrototypes.add(noFallbackError.NoFallbackError.prototype)
+  return noFallbackError
+}
+
+// The sentinel identity only exists at runtime, so there is no source function for Orchestrion to rewrite.
+addHook({
+  name: 'next',
+  versions: ['>=15.4.1'],
+  file: 'dist/shared/lib/no-fallback-error.external.js',
+}, captureNoFallbackError)
+
 function wrapOnRequestError (onRequestError) {
   return function (req, error) {
-    if (error) {
-      errorChannel.publish({ error })
-    }
+    if (error) publishError(req, undefined, error)
     return onRequestError.apply(this, arguments)
   }
 }
@@ -296,37 +376,127 @@ function publishRoutePage (ctx, routeModule, fallbackPage, isAppPath) {
 
   const pageData = getRoutePage(routeModule, fallbackPage)
   if (pageData.page) {
-    pageLoadChannel.publish(isAppPath ? { ...pageData, isAppPath: true } : pageData)
+    if (isAppPath) pageData.isAppPath = true
+    pageLoadChannel.publish(pageData)
   }
 }
 
 function wrapAppRouteHandle (handle) {
   return function (req, context) {
+    if (finishChannel.hasSubscribers) {
+      const nodeRequest = activeRouteRequests.get(this)
+      if (nodeRequest) {
+        nodeNextRequestsToNextRequests.set(nodeRequest, req)
+      }
+    }
+
+    return handle.apply(this, arguments)
+  }
+}
+
+function wrapPrepare (prepare) {
+  return function (req, res) {
+    if (startChannel.hasSubscribers || queryParsedChannel.hasSubscribers) {
+      routeResponses.set(req, res)
+    }
+
+    return prepare.apply(this, arguments)
+  }
+}
+
+/**
+ * @param {(context?: ResponseGeneratorContext) => unknown} responseGenerator
+ * @param {object | undefined} routeModule
+ * @param {NextNodeRequest} req
+ * @returns {(context?: ResponseGeneratorContext) => unknown}
+ */
+function wrapResponseGenerator (responseGenerator, routeModule, req) {
+  return function (context) {
+    if (context?.hasResolved) {
+      const nodeRequest = req.originalRequest || req
+      return backgroundRevalidationChannel.runStores(nodeRequest, () => responseGenerator.apply(this, arguments))
+    }
+    if (!routeModule) {
+      return responseGenerator.apply(this, arguments)
+    }
+
+    // Next calls the route module before the generator's first await. Limit the association to that
+    // synchronous call so concurrent requests can share the route module instance safely.
+    activeRouteRequests.set(routeModule, req)
+    try {
+      return responseGenerator.apply(this, arguments)
+    } finally {
+      activeRouteRequests.delete(routeModule)
+    }
+  }
+}
+
+/**
+ * @param {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>} handleResponse
+ * @param {boolean} associateNextRequest
+ * @returns {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>}
+ */
+function wrapHandleResponse (handleResponse, associateNextRequest) {
+  return function (options) {
+    const req = options.req
     if (!startChannel.hasSubscribers && !queryParsedChannel.hasSubscribers) {
-      return handle.apply(this, arguments)
+      return handleResponse.apply(this, arguments)
     }
 
-    const res = { statusCode: 500 }
-    nodeNextRequestsToNextRequests.set(req, req)
+    const res = routeResponses.get(req)
+    routeResponses.delete(req)
+
+    const routeModule = res && associateNextRequest ? this : undefined
+    options.responseGenerator = wrapResponseGenerator(options.responseGenerator, routeModule, req)
+
+    if (!res) return handleResponse.apply(this, arguments)
 
     return instrument(req, res, ctx => {
       publishRoutePage(ctx, this, undefined, true)
 
-      return handle.apply(this, arguments).then(response => {
-        if (ctx) ctx.res.statusCode = response?.status || 200
-        return response
+      return handleResponse.apply(this, arguments).then(result => {
+        const statusCode = result?.value?.status
+        if (ctx && typeof statusCode === 'number') {
+          ctx.res.statusCode = statusCode
+        }
+
+        return result
+      }, error => {
+        if (ctx && (typeof ctx.res.statusCode !== 'number' || ctx.res.statusCode < 400)) {
+          ctx.res.statusCode = 500
+        }
+
+        throw error
       })
-    })
+    }, undefined, finishChannel.hasSubscribers)
   }
 }
 
-function instrumentRouteModule (RouteModule, method, wrapper, handleErrors) {
+/**
+ * @param {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>} handleResponse
+ * @returns {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>}
+ */
+function wrapAppRouteHandleResponse (handleResponse) {
+  return wrapHandleResponse(handleResponse, true)
+}
+
+/**
+ * @param {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>} handleResponse
+ * @returns {(options: HandleResponseOptions) => Promise<ResponseCacheEntry | undefined>}
+ */
+function wrapAppPageHandleResponse (handleResponse) {
+  return wrapHandleResponse(handleResponse, false)
+}
+
+function instrumentRouteModule (RouteModule, wrappers, handleErrors) {
   const proto = RouteModule?.prototype
   if (!proto || patchedRouteModules.has(RouteModule)) return
 
   patchedRouteModules.add(RouteModule)
-  if (typeof proto[method] === 'function') {
-    shimmer.wrap(proto, method, wrapper)
+  for (const [method, wrapper] of wrappers) {
+    if (typeof proto[method] === 'function') {
+      shimmer.wrap(proto, method, wrapper)
+    }
   }
   if (handleErrors && typeof proto.onRequestError === 'function') {
     shimmer.wrap(proto, 'onRequestError', wrapOnRequestError)
@@ -334,7 +504,11 @@ function instrumentRouteModule (RouteModule, method, wrapper, handleErrors) {
 }
 
 function instrumentAppRouteRuntime (runtime) {
-  instrumentRouteModule(runtime.AppRouteRouteModule, 'handle', wrapAppRouteHandle, true)
+  instrumentRouteModule(runtime.AppRouteRouteModule, [
+    ['prepare', wrapPrepare],
+    ['handleResponse', wrapAppRouteHandleResponse],
+    ['handle', wrapAppRouteHandle],
+  ], true)
   return runtime
 }
 
@@ -361,40 +535,16 @@ function wrapPagesApiRender (render) {
 
 function instrumentPagesApiRuntime (runtime) {
   const PagesAPIRouteModule = runtime.PagesAPIRouteModule || runtime.default
-  instrumentRouteModule(PagesAPIRouteModule, 'render', wrapPagesApiRender, false)
+  instrumentRouteModule(PagesAPIRouteModule, [['render', wrapPagesApiRender]], false)
   return runtime
 }
 
-function wrapAppPageRender (render) {
-  return function (req, res, context = {}) {
-    if (!startChannel.hasSubscribers && !queryParsedChannel.hasSubscribers) {
-      return render.apply(this, arguments)
-    }
-
-    return instrument(req, res, ctx => {
-      publishRoutePage(ctx, this, context.page, true)
-
-      return render.apply(this, arguments).then(result => {
-        const statusCode = result?.metadata?.statusCode
-        if (ctx && typeof statusCode === 'number') {
-          ctx.res.statusCode = statusCode
-        }
-
-        return result
-      }, error => {
-        if (ctx && (typeof ctx.res.statusCode !== 'number' || ctx.res.statusCode < 400)) {
-          ctx.res.statusCode = 500
-        }
-
-        throw error
-      })
-    })
-  }
-}
-
 function instrumentAppPageRuntime (runtime) {
   const AppPageRouteModule = runtime.AppPageRouteModule || runtime.default
-  instrumentRouteModule(AppPageRouteModule, 'render', wrapAppPageRender, true)
+  instrumentRouteModule(AppPageRouteModule, [
+    ['prepare', wrapPrepare],
+    ['handleResponse', wrapAppPageHandleResponse],
+  ], true)
   return runtime
 }
 

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

@@ -6,6 +6,7 @@ const { addHook } = require('./helpers/instrument')
 
 const ch = dc.tracingChannel('apm:openai:request')
 const onStreamedChunkCh = dc.channel('apm:openai:request:chunk')
+const responsePromiseContexts = new WeakMap()
 
 // Provider lifecycle channels. Payloads stay OpenAI-native:
 // before { args, parentSpan, abortController, pending }
@@ -68,7 +69,7 @@ const V4_PACKAGE_SHIMS = [
     file: 'resources/responses/responses',
     targetClass: 'Responses',
     baseResource: 'responses',
-    methods: ['create'],
+    methods: ['create'], // implicitly covers `parse` as well which calls `create` under the hood
     streamedResponse: true,
     versions: ['>=4.87.0'],
   },
@@ -264,6 +265,42 @@ function wrapStreamIterator (response, options, ctx) {
 
 const extensions = ['.js', '.mjs']
 
+for (const extension of extensions) {
+  const apiPromiseShims = [
+    { file: `core${extension}`, versions: ['>=4 <5'] },
+    { file: `core/api-promise${extension}`, versions: ['>=5'] },
+  ]
+
+  for (const { file, versions } of apiPromiseShims) {
+    // APIPromise is a Promise subclass whose methods are dynamically replaced on pagination instances,
+    // so preserving the returned object's identity requires shimmer instead of Orchestrion.
+    addHook({ name: 'openai', file, versions }, exports => {
+      const { prototype } = exports.APIPromise
+
+      shimmer.wrap(prototype, 'parse', parse => function (...args) {
+        const state = responsePromiseContexts.get(this.responsePromise)
+        if (!state) return parse.apply(this, args)
+
+        const parsedPromise = parse.apply(this, args)
+          .then(body => Promise.all([this.responsePromise, body]))
+
+        return handleUnwrappedAPIPromise(parsedPromise, state)
+      })
+
+      shimmer.wrap(prototype, 'asResponse', asResponse => function (...args) {
+        const state = responsePromiseContexts.get(this.responsePromise)
+        const responsePromise = asResponse.apply(this, args)
+
+        if (!state?.getBeforeVerdict) return responsePromise
+
+        return Promise.all([state.getBeforeVerdict(), responsePromise]).then(([, response]) => response)
+      })
+
+      return exports
+    })
+  }
+}
+
 for (const extension of extensions) {
   for (const shim of V4_PACKAGE_SHIMS) {
     const { file, targetClass, baseResource, methods, versions, streamedResponse } = shim
@@ -309,44 +346,8 @@ for (const extension of extensions) {
               }
               : null
 
-            if (baseResource === 'chat.completions' && typeof apiProm._thenUnwrap === 'function') {
-              // this should only ever be invoked from a client.beta.chat.completions.parse call
-              shimmer.wrap(apiProm, '_thenUnwrap', origApiPromThenUnwrap => function (...args) {
-                // TODO(sam.brenner): I wonder if we can patch the APIPromise prototype instead, although
-                // we might not have access to everything we need...
-
-                // this is a new apipromise instance
-                const unwrappedPromise = origApiPromThenUnwrap.apply(this, args)
-
-                shimmer.wrap(unwrappedPromise, 'parse', origApiPromParse => function (...args) {
-                  const parsedPromise = origApiPromParse.apply(this, args)
-                    .then(body => Promise.all([this.responsePromise, body]))
-
-                  return handleUnwrappedAPIPromise(
-                    parsedPromise, ctx, stream, getBeforeVerdict, afterChannel, parentSpan
-                  )
-                })
-
-                return unwrappedPromise
-              })
-            }
-
-            // wrapping `parse` avoids problematic wrapping of `then` when trying to call
-            // `withResponse` in userland code after. This way, we can return the whole `APIPromise`
-            shimmer.wrap(apiProm, 'parse', origApiPromParse => function (...args) {
-              const parsedPromise = origApiPromParse.apply(this, args)
-                .then(body => Promise.all([this.responsePromise, body]))
-
-              return handleUnwrappedAPIPromise(parsedPromise, ctx, stream, getBeforeVerdict, afterChannel, parentSpan)
-            })
-
-            // Gate `.asResponse()` callers on the before verdict so raw-response paths still block.
-            if (beforeChannel && typeof apiProm.asResponse === 'function') {
-              shimmer.wrap(apiProm, 'asResponse', origAsResponse => function (...args) {
-                const responsePromise = origAsResponse.apply(this, args)
-                return Promise.all([getBeforeVerdict(), responsePromise]).then(([, response]) => response)
-              })
-            }
+            const state = { ctx, stream, getBeforeVerdict, afterChannel, parentSpan }
+            responsePromiseContexts.set(apiProm.responsePromise, state)
 
             ch.end.publish(ctx)
 
@@ -359,7 +360,8 @@ for (const extension of extensions) {
   }
 }
 
-function handleUnwrappedAPIPromise (apiProm, ctx, stream, getBeforeVerdict, afterChannel, parentSpan) {
+function handleUnwrappedAPIPromise (apiProm, state) {
+  const { ctx, stream, getBeforeVerdict, afterChannel, parentSpan } = state
   const gatedApiProm = getBeforeVerdict
     ? Promise.all([getBeforeVerdict(), apiProm]).then(([, result]) => result)
     : apiProm

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

@@ -11,10 +11,12 @@ const {
 const {
   clearPoolWaitTime,
   dispatchesAcquireSynchronously,
+  getPoolAcquireCapture,
   isPoolQueryAcquire,
   runOutsidePoolQueryAcquire,
   runPoolAcquireError,
   runWithPoolWait,
+  setPoolAcquireCaptureWaitTime,
   takePoolWaitTime,
   wrapPoolAcquireCarrier,
   wrapPoolQueryMethod,
@@ -69,6 +71,19 @@ addHook({ name: 'pg', versions: ['>=8.0.3'] }, pg => {
      * @returns {unknown}
      */
     const wrappedConnect = function (callback) {
+      const capture = getPoolAcquireCapture()
+      if (capture !== undefined && typeof callback !== 'function') {
+        const start = acquireStart(this)
+        const result = connect.apply(this, arguments)
+        return result.then(client => {
+          setPoolAcquireCaptureWaitTime(capture, acquireWait(start))
+          return client
+        }, error => {
+          setPoolAcquireCaptureWaitTime(capture, acquireWait(start))
+          throw error
+        })
+      }
+
       // Pool.query always supplies a callback, so a missing callback is an explicit acquire.
       if (typeof callback !== 'function') {
         if (!poolAcquireStartCh.hasSubscribers) {