dd-trace 5.114.0 → 5.116.0

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

@@ -8,7 +8,8 @@ const { getCompileToRegexp } = require('./path-to-regexp')
 const {
   getRouterMountPaths,
   joinPath,
-  setLayerMatchers,
+  setLayerMeta,
+  getLayerMeta,
   isAppMounted,
   setRouterMountPath,
   extractMountPaths,
@@ -25,137 +26,219 @@ function isFastSlash (layer, matchers) {
   return layer.regexp?.fast_slash ?? matchers.hasSlashPath
 }
 
-// TODO: Move this function to a shared file between Express and Router
 /**
+ * Cache the per-layer dispatch metadata in a side table instead of replacing
+ * `layer.handle`. Phase-sorting hosts (loopback's `_findLayerByHandler`) map a
+ * layer back to the user handler by scanning the handle, so the handle has to
+ * stay the user's function.
+ *
+ * @param {{ handle: Function, name?: string, path?: string,
+ *   regexp?: { fast_star?: boolean, fast_slash?: boolean } }} layer
+ * @param {Array<{ path?: string, regex?: RegExp }> & { hasStarPath?: boolean, hasSlashPath?: boolean }} matchers
+ */
+function annotateLayer (layer, matchers) {
+  const handle = layer.handle
+  const name = handle._name || layer.name || handle.name
+
+  let captureRoute
+  let needMultiMatch = false
+  if (matchers.length !== 0 && !isFastStar(layer, matchers) && !isFastSlash(layer, matchers)) {
+    if (matchers.length === 1) {
+      captureRoute = matchers[0].path
+    } else {
+      needMultiMatch = true
+    }
+  }
+
+  setLayerMeta(layer, { name, captureRoute, needMultiMatch, matchers })
+}
+
+/**
+ * Resolve the route for a dispatched layer. Single-pattern layers carry a
+ * constant route; only multi-pattern stacks need the per-request `layer.path`
+ * match the host already computed.
+ *
+ * @param {{ captureRoute?: string, needMultiMatch: boolean,
+ *   matchers: Array<{ path?: string, regex?: RegExp }> }} meta
+ * @param {{ path?: string }} layer
+ * @returns {string | undefined}
+ */
+function resolveLayerRoute (meta, layer) {
+  if (!meta.needMultiMatch) return meta.captureRoute
+
+  for (const matcher of meta.matchers) {
+    if (matcher.regex?.test(layer.path)) return matcher.path
+  }
+}
+
+/**
+ * Build the request/error dispatch wrappers for one host (`express` / `router`).
+ * They wrap the layer's prototype dispatch and read the side-table metadata, so
+ * `layer.handle` is never replaced. The arity guard mirrors the host's own
+ * (`handle_request` skips 4-arg handlers, `handle_error` skips the rest), so a
+ * span is published only for the layer the host actually runs.
+ *
  * @param {string} name Channel namespace (`apm:<name>:middleware:*`).
- * @param {((pattern: string | RegExp) => RegExp | undefined) | undefined} compile
- *   Host-resolved path-to-regexp compile adapter, or undefined when the host
- *   instance ships no path-to-regexp. Captured here so each express/router
- *   instance keeps the dialect it actually loaded.
  */
-function createWrapRouterMethod (name, compile) {
+function createLayerDispatchWrappers (name) {
   const enterChannel = channel(`apm:${name}:middleware:enter`)
   const exitChannel = channel(`apm:${name}:middleware:exit`)
   const finishChannel = channel(`apm:${name}:middleware:finish`)
   const errorChannel = channel(`apm:${name}:middleware:error`)
   const nextChannel = channel(`apm:${name}:middleware:next`)
-  const routeAddedChannel = channel(`apm:${name}:route:added`)
+  const repeatChannel = channel(`apm:${name}:middleware:repeat`)
   // Bound per name so express and a bare router keep independent guards.
   const publishError = createErrorPublisher(errorChannel)
 
-  function wrapLayerHandle (layer, original, matchers) {
-    // Resolve `name` once at wrap time: cached on the original for any code
-    // that reads `_name`, captured in the closure so the per-call body avoids
-    // the property-lookup / `||` fallback.
-    const name = original._name || layer.name || original.name
-    original._name = name
-
-    // Wrap-time matcher analysis. The single-pattern case yields a constant
-    // route; only multi-pattern stacks need a per-request layer.path match.
-    let captureRoute
-    let needMultiMatch = false
-    if (matchers.length !== 0 && !isFastStar(layer, matchers) && !isFastSlash(layer, matchers)) {
-      if (matchers.length === 1) {
-        captureRoute = matchers[0].path
-      } else {
-        needMultiMatch = true
-      }
-    }
+  /**
+   * @param {import('node:http').IncomingMessage} req
+   * @param {string | undefined} layerName
+   * @param {(error?: unknown) => void} originalNext
+   */
+  function wrapNext (req, layerName, originalNext) {
+    // Per layer dispatch, N per request. Named `next`/arity-1 mirrors the
+    // router continuation so wrapCallback skips its name/length rewrite.
+    let calls = 0
+    return shimmer.wrapCallback(originalNext, original => function next (error) {
+      // A handler that calls `next()` and then rejects (`next(); await bg()`)
+      // makes the host call this continuation twice. Publish once so the second
+      // pass cannot tag the already-finished span's parent with a late error.
+      calls++
+      if (calls === 1) {
+        if (error && error !== 'route' && error !== 'router') {
+          publishError({ req, error })
+        }
 
-    // Split by arity: router only ever dispatches 3-arg request handlers
-    // through `Layer.handleRequest` and 4-arg error handlers through
-    // `Layer.handleError`. Specialising lets the per-call body use named
-    // parameters and `.call`, avoiding the rest-spread Array allocation that
-    // the unified shape forced on every middleware invocation.
-    const wrapped = original.length === 4
-      ? shimmer.wrapFunction(original, errorHandlerLayerWrap(layer, name, captureRoute, needMultiMatch, matchers))
-      : shimmer.wrapFunction(original, requestHandlerLayerWrap(layer, name, captureRoute, needMultiMatch, matchers))
-
-    // Workaround for loopback's phase-based middleware sorting. Its
-    // `_findLayerByHandler` maps a layer back to the user handler by scanning
-    // the layer handle's enumerable properties for the original function.
-    // Replacing `layer.handle` with this wrapper hides that handler, so without
-    // the back-reference loopback cannot tag the layer with its phase and
-    // `app.middleware(phase, ...)` handlers run in insertion order instead of
-    // phase order. The property name is part of that contract; keep it stable.
-    wrapped._datadog_orig = original
+        nextChannel.publish({ req })
+        finishChannel.publish({ req })
+      } else if (calls === 2) {
+        // Surface the repeat as a diagnostic on the still-live request span. The
+        // host cannot tell a legitimate `next(); await bg()` from a buggy double
+        // `next()`, so this only records that it happened, not that it is wrong.
+        repeatChannel.publish({ req, name: layerName, error })
+      }
 
-    return wrapped
+      original.apply(this, arguments)
+    })
   }
 
-  function requestHandlerLayerWrap (layer, name, captureRoute, needMultiMatch, matchers) {
-    return original => function (req, res, next) {
-      if (!enterChannel.hasSubscribers) return original.call(this, req, res, next)
-
-      const wrappedNext = typeof next === 'function' ? wrapNext(req, next) : next
+  // Every host dispatch turns a synchronous throw into `next(error)`, and the
+  // hosts that await the handler (router >=2, express 5, express 4 +
+  // express-async-errors) do the same for a rejected promise. Passing
+  // `wrappedNext` through captures both without a tracer-side try/catch; only
+  // `exit` needs the `finally`. express 4's native dispatch converts only the
+  // synchronous throw — exactly what the pre-refactor handle wrap caught.
+  function wrapLayerRequest (originalRequest) {
+    return function (req, res, next) {
+      if (!enterChannel.hasSubscribers) return originalRequest.call(this, req, res, next)
 
-      let route = captureRoute
-      if (needMultiMatch) {
-        for (const matcher of matchers) {
-          if (matcher.regex?.test(layer.path)) {
-            route = matcher.path
-            break
-          }
-        }
-      }
+      const meta = getLayerMeta(this)
+      if (meta === undefined || this.handle.length > 3) return originalRequest.call(this, req, res, next)
 
-      enterChannel.publish({ name, req, route, layer })
+      const wrappedNext = typeof next === 'function' ? wrapNext(req, meta.name, next) : next
+      enterChannel.publish({ name: meta.name, req, route: resolveLayerRoute(meta, this), layer: this })
 
       try {
-        return original.call(this, req, res, wrappedNext)
-      } catch (error) {
-        publishError({ req, error })
-        nextChannel.publish({ req })
-        finishChannel.publish({ req })
-
-        throw error
+        return originalRequest.call(this, req, res, wrappedNext)
       } finally {
         exitChannel.publish({ req })
       }
     }
   }
 
-  function errorHandlerLayerWrap (layer, name, captureRoute, needMultiMatch, matchers) {
-    return original => function (error, req, res, next) {
-      if (!enterChannel.hasSubscribers) return original.call(this, error, req, res, next)
+  function wrapLayerError (originalError) {
+    return function (error, req, res, next) {
+      if (!enterChannel.hasSubscribers) return originalError.call(this, error, req, res, next)
 
-      const wrappedNext = typeof next === 'function' ? wrapNext(req, next) : next
+      const meta = getLayerMeta(this)
+      if (meta === undefined || this.handle.length !== 4) return originalError.call(this, error, req, res, next)
 
-      let route = captureRoute
-      if (needMultiMatch) {
-        for (const matcher of matchers) {
-          if (matcher.regex?.test(layer.path)) {
-            route = matcher.path
-            break
-          }
-        }
+      const wrappedNext = typeof next === 'function' ? wrapNext(req, meta.name, next) : next
+      enterChannel.publish({ name: meta.name, req, route: resolveLayerRoute(meta, this), layer: this })
+
+      try {
+        return originalError.call(this, error, req, res, wrappedNext)
+      } finally {
+        exitChannel.publish({ req })
       }
+    }
+  }
 
-      enterChannel.publish({ name, req, route, layer })
+  // express <4.6.0 has no `Layer` prototype dispatch: the router invokes
+  // `layer.handle` directly and routes errors by its arity. There the handle is
+  // replaced in place, with the arity preserved so the host still routes
+  // correctly. Newer express, express 5, and the router package keep `handle`
+  // pristine and are traced through the prototype wraps above.
+  function wrapLegacyHandle (layer, original) {
+    // `annotateLayer` always runs first in `wrapStack`, so the captured meta is
+    // never undefined here (unlike the prototype wraps, where `this` can be any
+    // layer the host dispatches).
+    const meta = getLayerMeta(layer)
+    const wrapped = shimmer.wrapFunction(original, inner => function (...args) {
+      if (!enterChannel.hasSubscribers) return inner.apply(this, args)
+
+      const isErrorHandler = original.length === 4
+      const req = args[isErrorHandler ? 1 : 0]
+      const nextIndex = isErrorHandler ? 3 : 2
+      if (typeof args[nextIndex] === 'function') args[nextIndex] = wrapNext(req, meta.name, args[nextIndex])
+
+      enterChannel.publish({ name: meta.name, req, route: resolveLayerRoute(meta, layer), layer })
 
       try {
-        return original.call(this, error, req, res, wrappedNext)
-      } catch (caught) {
-        publishError({ req, error: caught })
+        return inner.apply(this, args)
+      } catch (error) {
+        // Unlike the prototype hosts, this router catches a synchronous throw
+        // outside the layer and calls its own `next(error)`, never `wrappedNext`.
+        // Mirror `wrapNext` here so the throwing layer still tags its error and
+        // finishes, rather than lingering on the stack until request finish.
+        publishError({ req, error })
         nextChannel.publish({ req })
         finishChannel.publish({ req })
 
-        throw caught
+        throw error
       } finally {
         exitChannel.publish({ req })
       }
-    }
+    })
+    Object.defineProperty(wrapped, 'length', { value: original.length, configurable: true })
+    return wrapped
   }
 
+  return { wrapLayerRequest, wrapLayerError, wrapLegacyHandle }
+}
+
+/**
+ * @param {{ handle_request?: unknown, handleRequest?: unknown }} layer
+ */
+function hasLayerDispatch (layer) {
+  return typeof layer.handle_request === 'function' || typeof layer.handleRequest === 'function'
+}
+
+// TODO: Move this function to a shared file between Express and Router
+/**
+ * @param {string} name Channel namespace (`apm:<name>:middleware:*`).
+ * @param {((pattern: string | RegExp) => RegExp | undefined) | undefined} compile
+ *   Host-resolved path-to-regexp compile adapter, or undefined when the host
+ *   instance ships no path-to-regexp. Captured here so each express/router
+ *   instance keeps the dialect it actually loaded.
+ * @param {((layer: object, original: Function) => Function) | undefined} [wrapLegacyHandle]
+ *   Fallback that replaces `layer.handle` for hosts without a `Layer` prototype
+ *   dispatch (express <4.6.0). Omitted for hosts that always ship one.
+ */
+function createWrapRouterMethod (name, compile, wrapLegacyHandle) {
+  const routeAddedChannel = channel(`apm:${name}:route:added`)
+
   function wrapStack (layers, matchers) {
     for (const layer of layers) {
-      if (layer.__handle) { // express-async-errors
-        layer.__handle = wrapLayerHandle(layer, layer.__handle, matchers)
-      } else {
-        layer.handle = wrapLayerHandle(layer, layer.handle, matchers)
-      }
+      annotateLayer(layer, matchers)
 
-      setLayerMatchers(layer, matchers)
+      if (wrapLegacyHandle !== undefined && !hasLayerDispatch(layer)) {
+        if (layer.__handle) { // express-async-errors
+          layer.__handle = wrapLegacyHandle(layer, layer.__handle)
+        } else {
+          layer.handle = wrapLegacyHandle(layer, layer.handle)
+        }
+      }
 
       if (layer.route) {
         for (const method of METHODS) {
@@ -169,21 +252,6 @@ function createWrapRouterMethod (name, compile) {
     }
   }
 
-  function wrapNext (req, originalNext) {
-    // Per layer dispatch, N per request. Named `next`/arity-1 mirrors the
-    // router continuation so wrapCallback skips its name/length rewrite.
-    return shimmer.wrapCallback(originalNext, original => function next (error) {
-      if (error && error !== 'route' && error !== 'router') {
-        publishError({ req, error })
-      }
-
-      nextChannel.publish({ req })
-      finishChannel.publish({ req })
-
-      original.apply(this, arguments)
-    })
-  }
-
   function extractMatchers (fn) {
     const arg = Array.isArray(fn) ? fn.flat(Infinity) : [fn]
 
@@ -304,6 +372,15 @@ addHook({ name: 'router', versions: ['>=1 <2'] }, Router => {
   return Router
 })
 
+addHook({ name: 'router', file: 'lib/layer.js', versions: ['>=1 <2'] }, Layer => {
+  const { wrapLayerRequest, wrapLayerError } = createLayerDispatchWrappers('router')
+
+  shimmer.wrap(Layer.prototype, 'handle_request', wrapLayerRequest)
+  shimmer.wrap(Layer.prototype, 'handle_error', wrapLayerError)
+
+  return Layer
+})
+
 const queryParserReadCh = channel('datadog:query:read:finish')
 
 addHook({ name: 'router', versions: ['>=2'] }, Router => {
@@ -365,6 +442,14 @@ function wrapHandleRequest (original) {
 addHook({
   name: 'router', file: 'lib/layer.js', versions: ['>=2'],
 }, Layer => {
+  const { wrapLayerRequest, wrapLayerError } = createLayerDispatchWrappers('router')
+
+  // `handleRequest` carries two concerns: the middleware dispatch span and the
+  // param-start publish (`wrapHandleRequest`). Wrap the dispatch first so it
+  // sits inner and param-start still fires before `middleware:enter`, matching
+  // the order from when the handle itself was wrapped.
+  shimmer.wrap(Layer.prototype, 'handleRequest', wrapLayerRequest)
+  shimmer.wrap(Layer.prototype, 'handleError', wrapLayerError)
   shimmer.wrap(Layer.prototype, 'handleRequest', wrapHandleRequest)
   return Layer
 })
@@ -404,4 +489,4 @@ addHook({
   return router
 })
 
-module.exports = { createWrapRouterMethod }
+module.exports = { createWrapRouterMethod, createLayerDispatchWrappers }

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

@@ -5,6 +5,10 @@ const realSetTimeout = setTimeout
 
 const shimmer = require('../../datadog-shimmer')
 const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
+const {
+  RUM_TEST_EXECUTION_ID_COOKIE_NAME: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
+} = require('../../dd-trace/src/ci-visibility/rum')
+const log = require('../../dd-trace/src/log')
 const { addHook, channel } = require('./helpers/instrument')
 
 const ciSeleniumDriverGetStartCh = channel('ci:selenium:driver:get')
@@ -20,7 +24,6 @@ if (window.DD_RUM && window.DD_RUM.stopSession) {
 const IS_RUM_ACTIVE_SCRIPT = 'return !!window.DD_RUM'
 
 const DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS = getValueFromEnvSources('DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS')
-const DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME = 'datadog-ci-visibility-test-execution-id'
 
 // TODO: can we increase the supported version range?
 addHook({
@@ -32,28 +35,47 @@ addHook({
     if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
       return get.apply(this, arguments)
     }
-    let traceId
-    const setTraceId = (inputTraceId) => {
-      traceId = inputTraceId
-    }
     const getResult = await get.apply(this, arguments)
 
-    const isRumActive = await this.executeScript(IS_RUM_ACTIVE_SCRIPT)
-    const capabilities = await this.getCapabilities()
+    let isRumActive
+    try {
+      isRumActive = await this.executeScript(IS_RUM_ACTIVE_SCRIPT)
+    } catch (error) {
+      log.error('Selenium RUM detection error', error)
+    }
+
+    let browserName
+    let browserVersion
+    try {
+      const capabilities = await this.getCapabilities()
+      browserName = capabilities.getBrowserName()
+      browserVersion = capabilities.getBrowserVersion()
+    } catch (error) {
+      log.error('Selenium browser metadata error', error)
+    }
 
-    ciSeleniumDriverGetStartCh.publish({
-      setTraceId,
+    const context = {
       seleniumVersion,
-      browserName: capabilities.getBrowserName(),
-      browserVersion: capabilities.getBrowserVersion(),
+      browserName,
+      browserVersion,
       isRumActive,
-    })
+      testExecutionId: undefined,
+    }
+    try {
+      ciSeleniumDriverGetStartCh.publish(context)
+    } catch (error) {
+      log.error('Selenium RUM correlation channel error', error)
+    }
 
-    if (traceId && isRumActive) {
-      await this.manage().addCookie({
-        name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
-        value: traceId,
-      })
+    if (context.testExecutionId && isRumActive) {
+      try {
+        await this.manage().addCookie({
+          name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
+          value: context.testExecutionId,
+        })
+      } catch (error) {
+        log.error('Selenium RUM correlation cookie error', error)
+      }
     }
 
     return getResult
@@ -63,16 +85,20 @@ addHook({
     if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
       return quit.apply(this, arguments)
     }
-    const isRumActive = await this.executeScript(RUM_STOP_SESSION_SCRIPT)
+    try {
+      const isRumActive = await this.executeScript(RUM_STOP_SESSION_SCRIPT)
 
-    if (isRumActive) {
-      // We'll have time for RUM to flush the events (there's no callback to know when it's done)
-      await new Promise(resolve => {
-        realSetTimeout(() => {
-          resolve()
-        }, DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS)
-      })
-      await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
+      if (isRumActive) {
+        // We'll have time for RUM to flush the events (there's no callback to know when it's done)
+        await new Promise(resolve => {
+          realSetTimeout(() => {
+            resolve()
+          }, DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS)
+        })
+        await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
+      }
+    } catch (error) {
+      log.error('Selenium RUM cleanup error', error)
     }
 
     return quit.apply(this, arguments)

package/packages/datadog-instrumentations/src/vitest-main-no-worker-init.js

@@ -784,7 +784,7 @@ function getRepeatedTestReport (task, testName, testSuiteAbsolutePath, testPrope
 
   return {
     errors,
-    finalAttempt: attempts[attempts.length - 1],
+    finalAttempt: attempts.at(-1),
     nonFinalAttempts: attempts.slice(0, -1),
     state,
     status,
@@ -880,7 +880,7 @@ function reportFinalTestAttempt (testReport) {
   const finalAttempt = testReport.finalAttempt
 
   if (status === 'fail') {
-    const error = errors[errors.length - 1] || errors[0]
+    const error = errors.at(-1) || errors[0]
     reportTestAttempt(testReport, finalAttempt || {
       error,
       finalStatus,

package/packages/datadog-instrumentations/src/vitest-main.js

@@ -1343,10 +1343,8 @@ function wrapTinyPoolRun (TinyPool) {
 
   shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
     // We have to do this before and after because the threads list gets recycled, that is, the processes are re-created
-    // eslint-disable-next-line unicorn/no-array-for-each
     this.threads.forEach(threadHandler)
     const runResult = await run.apply(this, arguments)
-    // eslint-disable-next-line unicorn/no-array-for-each
     this.threads.forEach(threadHandler)
     return runResult
   })

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

@@ -3,6 +3,7 @@
 const { tracingChannel } = /** @type {import('node:diagnostics_channel')} */ (require('dc-polyfill'))
 
 const shimmer = require('../../datadog-shimmer')
+const { getSegment } = require('../../dd-trace/src/util')
 const {
   addHook,
   channel,
@@ -106,7 +107,7 @@ function createWrapEmit (emit) {
     }
 
     const ctx = { req }
-    ctx.req.resStatus = headers[0].split(' ')[1]
+    ctx.req.resStatus = getSegment(headers[0], ' ', 1)
 
     emitCh.runStores(ctx, () => {
       try {

package/packages/datadog-plugin-aws-durable-execution-sdk-js/src/context.js

@@ -94,7 +94,7 @@ class RunInChildContextPlugin extends BaseContextPlugin {
   static retryable = false
 
   bindStart (ctx) {
-    if (SUPPRESSED_CHILD_CONTEXT_SUBTYPES.has(getRunInChildContextSubType(ctx))) {
+    if (SUPPRESSED_CHILD_CONTEXT_SUBTYPES.has(getRunInChildContextSubtype(ctx))) {
       // Pass the active store through unchanged so any nested spans
       // remain parented to the surrounding map/parallel span
       ctx.suppressed = true
@@ -105,7 +105,7 @@ class RunInChildContextPlugin extends BaseContextPlugin {
 }
 
 // runInChildContext has two overloads: `(name, fn, options)` and `(fn, options)`.
-function getRunInChildContextSubType (ctx) {
+function getRunInChildContextSubtype (ctx) {
   const args = ctx.arguments || []
   const opts = typeof args[0] === 'function' ? args[1] : args[2]
   return opts?.subType

package/packages/datadog-plugin-aws-durable-execution-sdk-js/src/util.js

@@ -16,7 +16,8 @@ const REPLAYED_STATUSES = new Set(['SUCCEEDED', 'FAILED'])
  * @returns {{ stepId: string | undefined, stepData: object | undefined }}
  */
 function getStepDataForNext (ctxImpl) {
-  const stepId = ctxImpl?.getNextStepId?.()
+  // @aws/durable-execution-sdk-js@2.2.0 renames .getNextStepId() to .peekStepId()
+  const stepId = ctxImpl?.peekStepId?.() ?? ctxImpl?.getNextStepId?.()
   const stepData = stepId ? ctxImpl?._executionContext?.getStepData?.(stepId) : undefined
   return { stepId, stepData }
 }

package/packages/datadog-plugin-aws-sdk/src/services/kinesis.js

@@ -2,7 +2,6 @@
 const { DsmPathwayCodec, getSizeOrZero } = require('../../../dd-trace/src/datastreams')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
-const { isEmpty } = require('../util')
 
 function recordDataAsString (data) {
   return Buffer.isBuffer(data) ? data.toString('utf8') : Buffer.from(data).toString('utf8')
@@ -13,6 +12,16 @@ function recordDataAsString (data) {
 // delete on consume, so their working set is ~the active shard count.
 const MAX_TRACKED_SHARD_ITERATORS = 1000
 
+// The default Kinesis record limit is 1 MiB. Streams configured for larger records
+// conservatively skip propagation above this point because the request does not expose that limit.
+const KINESIS_DEFAULT_MAX_RECORD_BYTES = 1_048_576
+
+// The DSM pathway field (`dd-pathway-ctx-base64`) always serializes to a fixed 55 bytes: a
+// 21-char key, a 28-char base64 value, and 6 bytes of JSON framing. Mirrors PATHWAY_HEADER_BYTES
+// in dd-trace/src/datastreams/processor.js. Reserved in the size gate so setDSMCheckpoint never
+// runs for a record that would ship over the cap once the pathway context is attached.
+const DSM_PATHWAY_FIELD_BYTES = 55
+
 class Kinesis extends BaseAwsSdkPlugin {
   static id = 'kinesis'
   static peerServicePrecursors = ['streamname']
@@ -237,6 +246,12 @@ class Kinesis extends BaseAwsSdkPlugin {
     }
   }
 
+  /**
+   * @param {import('../../../dd-trace/src/opentracing/span') | null} span
+   * @param {{ Data: Buffer|string|Uint8Array, PartitionKey?: string } | undefined} params
+   * @param {string} stream
+   * @param {boolean} injectTraceContext
+   */
   injectToMessage (span, params, stream, injectTraceContext) {
     if (!params) {
       return
@@ -251,30 +266,40 @@ class Kinesis extends BaseAwsSdkPlugin {
       }
     }
 
-    const ddInfo = {}
+    let ddInfo
     // For now we only inject to the first message; batches may change later.
     if (injectTraceContext) {
-      this.tracer.inject(span, 'text_map', ddInfo)
+      ddInfo = this.tracer.inject(span, 'text_map')
     }
 
-    if (this.config.dsmEnabled) {
-      parsedData._datadog = ddInfo
+    const dsmEnabled = this.config.dsmEnabled
+    if (!ddInfo && !dsmEnabled) return
+
+    parsedData._datadog = ddInfo ?? {}
+    // Gate before setDSMCheckpoint: a record we can't ship must not record a checkpoint.
+    // Reserve the pathway field that DSM appends after the gate.
+    let serialized = JSON.stringify(parsedData)
+    const reservedBytes = dsmEnabled ? DSM_PATHWAY_FIELD_BYTES : 0
+    const partitionKeyBytes = Buffer.byteLength(params.PartitionKey ?? '', 'utf8')
+    if (
+      Buffer.byteLength(serialized, 'utf8') + partitionKeyBytes + reservedBytes >
+      KINESIS_DEFAULT_MAX_RECORD_BYTES
+    ) {
+      log.info('Payload size too large to pass context')
+      return
+    }
+
+    if (dsmEnabled) {
       const dataStreamsContext = this.setDSMCheckpoint(span, params, stream)
-      if (dataStreamsContext) {
-        DsmPathwayCodec.encode(dataStreamsContext, ddInfo)
+      ddInfo = DsmPathwayCodec.encode(dataStreamsContext, ddInfo) ?? ddInfo
+      if (ddInfo) {
+        parsedData._datadog = ddInfo
+        serialized = JSON.stringify(parsedData)
       }
     }
 
-    if (isEmpty(ddInfo)) return
+    if (!ddInfo) return
 
-    parsedData._datadog = ddInfo
-    const serialized = JSON.stringify(parsedData)
-    const byteSize = Buffer.byteLength(serialized, 'utf8')
-    // Kinesis max payload size is 1 MiB; bail if our context push tipped us over.
-    if (byteSize >= 1_048_576) {
-      log.info('Payload size too large to pass context')
-      return
-    }
     params.Data = Buffer.from(serialized, 'utf8')
   }