dd-trace 5.108.0 → 5.109.0

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

@@ -0,0 +1,224 @@
+'use strict'
+
+const crypto = require('crypto')
+const log = require('../../dd-trace/src/log')
+const TextMapPropagator = require('../../dd-trace/src/opentracing/propagation/text_map')
+
+const CHECKPOINT_NAME_PREFIX = '_datadog_'
+
+// Propagator that injects only Datadog-style headers (`x-datadog-*`) regardless of the user's
+// `DD_TRACE_PROPAGATION_STYLE_INJECT`. Checkpoints are written and read entirely by Datadog code,
+// so honoring user style preferences would only complicate the payload contract. AWS runs a single
+// tracer, so one lazily-built propagator suffices.
+let datadogOnlyPropagator
+
+function getDatadogOnlyPropagator (tracer) {
+  if (datadogOnlyPropagator) return datadogOnlyPropagator
+  const config = tracer._config
+  // Shadow `tracePropagationStyle.inject` while inheriting every other field (x-datadog-tags length
+  // cap, etc.) from the live config. Disable `legacyBaggageEnabled` only to keep `ot-baggage-*` out
+  // of the checkpoint payload we persist (sensitive-data concern) — not a serverless-wide change.
+  const shadowConfig = Object.create(config)
+  shadowConfig.tracePropagationStyle = {
+    ...config.tracePropagationStyle,
+    inject: ['datadog'],
+  }
+  shadowConfig.legacyBaggageEnabled = false
+  datadogOnlyPropagator = new TextMapPropagator(shadowConfig)
+  return datadogOnlyPropagator
+}
+
+/**
+ * Build the Datadog-format headers dict from a span context.
+ * Mirrors ddtrace-py HTTPPropagator.inject output so the same payload
+ * can be consumed by either language's datadog-lambda wrapper.
+ * @param {object} tracer
+ * @param {object} span
+ * @returns {Record<string, string>}
+ */
+function injectHeaders (tracer, span) {
+  const headers = {}
+  const propagator = getDatadogOnlyPropagator(tracer)
+  const ctx = typeof span?.context === 'function' ? span.context() : span
+  propagator.inject(ctx, headers)
+  return headers
+}
+
+/**
+ * Mutate headers in-place to set parent_id to the provided value.
+ * @param {Record<string, string>} headers
+ * @param {string | number | undefined} parentId
+ */
+function overrideParentId (headers, parentId) {
+  if (!parentId) return
+  if (headers['x-datadog-trace-id']) {
+    headers['x-datadog-parent-id'] = String(parentId)
+  }
+}
+
+/**
+ * Whether the current trace context warrants a new checkpoint over the previously-saved one.
+ * @param {Record<string, string>} currentHeaders
+ * @param {Record<string, string>} previousHeaders
+ * @returns {boolean}
+ */
+function needsCheckpointUpdate (currentHeaders, previousHeaders) {
+  for (const key of Object.keys(currentHeaders)) {
+    if (currentHeaders[key] !== previousHeaders[key]) return true
+  }
+  return false
+}
+
+/**
+ * Find the checkpoint with the highest N for _datadog_{N} in the event's operations.
+ * @param {unknown} event
+ * @returns {{ checkpointNumber: number, operation: object } | undefined}
+ */
+function findLastCheckpoint (event) {
+  if (!event || typeof event !== 'object') return
+
+  const operations = event.InitialExecutionState?.Operations
+  if (!Array.isArray(operations)) return
+
+  let highest
+  for (const op of operations) {
+    const name = op?.Name
+    if (typeof name !== 'string' || !name.startsWith(CHECKPOINT_NAME_PREFIX)) continue
+    const suffix = name.slice(CHECKPOINT_NAME_PREFIX.length)
+    const checkpointNumber = Number.parseInt(suffix, 10)
+    if (Number.isNaN(checkpointNumber) || String(checkpointNumber) !== suffix) continue
+
+    if (!highest || checkpointNumber > highest.checkpointNumber) {
+      highest = { checkpointNumber, operation: op }
+    }
+  }
+
+  return highest
+}
+
+/**
+ * Parse the JSON payload from a checkpoint STEP operation's Payload or StepDetails.Result.
+ * @param {object} op
+ * @returns {Record<string, string> | undefined}
+ */
+function parseCheckpointPayload (op) {
+  try {
+    const raw = op?.Payload ?? op?.StepDetails?.Result
+    if (!raw || typeof raw !== 'string') return
+    const parsed = JSON.parse(raw)
+    return parsed && typeof parsed === 'object' ? parsed : undefined
+  } catch {
+    log.debug('Failed to parse checkpoint payload')
+  }
+}
+
+/**
+ * Save a _datadog_{number} STEP operation via the SDK's checkpoint manager.
+ * Uses a deterministic blake2b hash of (name:arn) as the stepId so the save is
+ * idempotent within an execution.
+ * @param {object} checkpointManager
+ * @param {string} executionArn
+ * @param {number} number
+ * @param {Record<string, string>} headers
+ * @returns {Promise<void>}
+ */
+async function saveCheckpoint (checkpointManager, executionArn, number, headers) {
+  const name = `${CHECKPOINT_NAME_PREFIX}${number}`
+  const stepId = crypto
+    .createHash('blake2b512')
+    .update(`${name}:${executionArn}`)
+    .digest('hex')
+    .slice(0, 64)
+  const payload = JSON.stringify(headers)
+
+  // Queue START and SUCCEED back-to-back before awaiting. This allows callers
+  // to trigger save right before termination without losing the second update.
+  const startPromise = checkpointManager.checkpoint(stepId, {
+    Id: stepId,
+    Action: 'START',
+    Type: 'STEP',
+    SubType: 'STEP',
+    Name: name,
+  })
+  const succeedPromise = checkpointManager.checkpoint(stepId, {
+    Id: stepId,
+    Action: 'SUCCEED',
+    Type: 'STEP',
+    SubType: 'STEP',
+    Name: name,
+    Payload: payload,
+  })
+  await Promise.all([
+    startPromise,
+    succeedPromise,
+  ])
+  log.debug('Saved trace context checkpoint %s', name)
+}
+
+/**
+ * Save a new trace-context checkpoint when the current context differs from
+ * the most recent `_datadog_{N}` operation already in the event.
+ *
+ * Every checkpoint across the durable execution carries the same
+ * `x-datadog-parent-id` so all resumed invocations attach to the same anchor:
+ * - First checkpoint (no previous): anchor at `firstExecutionSpanId`.
+ * - Subsequent: reuse the prior checkpoint's `x-datadog-parent-id` verbatim —
+ *   that value originated from the first save and is the anchor we've been
+ *   carrying forward.
+ *
+ * Caller is responsible for invoking this only when a save is appropriate — i.e.
+ * the SDK is about to return Status: PENDING (see PENDING_TERMINATION_REASONS in
+ * handler.js). This function does not re-check that.
+ *
+ * Errors (including a rejected checkpoint-manager call) propagate to the caller, which swallows
+ * them at the fire-and-forget boundary.
+ *
+ * @param {object} tracer
+ * @param {object} span - aws.durable.execute span
+ * @param {object} durableContext - SDK's DurableContextImpl
+ * @param {string | undefined} firstExecutionSpanId - span id of the first
+ *   invocation's `aws.durable.execute` span. Only consulted on the very
+ *   first save; ignored once a prior `_datadog_{N}` exists. We anchor at
+ *   this span (which this integration owns) rather than its parent so the
+ *   anchor doesn't depend on whatever upstream context happens to be
+ *   active when `bindStart` fires.
+ * @param {unknown} event - raw invocation event (has InitialExecutionState)
+ * @returns {Promise<void>}
+ */
+async function saveTraceContextCheckpointIfUpdated (
+  tracer, span, durableContext, firstExecutionSpanId, event,
+) {
+  const checkpointManager = durableContext.checkpoint ?? durableContext.checkpointManager
+  if (typeof checkpointManager?.checkpoint !== 'function') return
+
+  const currentHeaders = injectHeaders(tracer, span)
+  if (currentHeaders['x-datadog-trace-id'] === undefined) return
+
+  const latest = findLastCheckpoint(event)
+
+  let newNumber
+  if (latest) {
+    const previousHeaders = parseCheckpointPayload(latest.operation)
+    if (!previousHeaders) return
+
+    // x-datadog-parent-id reflects the active span at save time and always differs, so exclude it
+    // from the comparison. Capture the previous anchor first to carry it forward on a real update.
+    // needsCheckpointUpdate only reads currentHeaders' keys, so deleting it from there is enough.
+    const anchoredSpanId = previousHeaders['x-datadog-parent-id']
+    delete currentHeaders['x-datadog-parent-id']
+    if (!needsCheckpointUpdate(currentHeaders, previousHeaders)) return
+
+    newNumber = latest.checkpointNumber + 1
+    overrideParentId(currentHeaders, anchoredSpanId)
+  } else {
+    newNumber = 0
+    if (firstExecutionSpanId) overrideParentId(currentHeaders, firstExecutionSpanId)
+  }
+
+  const executionArn = event?.DurableExecutionArn || ''
+  await saveCheckpoint(checkpointManager, executionArn, newNumber, currentHeaders)
+}
+
+module.exports = {
+  saveTraceContextCheckpointIfUpdated,
+}

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

@@ -0,0 +1,43 @@
+'use strict'
+
+const { createHash } = require('node:crypto')
+
+/**
+ * Populates the replay and operation_id tags for the op the DurableContextImpl is about to
+ * run, deriving both from a single `getNextStepId()` call. `aws.durable.replayed` is always
+ * set ('true' when the next stepId already has a SUCCEEDED checkpoint entry, i.e. the op will
+ * be served from the SDK's checkpoint). `aws.durable.operation_id` — the 16-hex-char MD5 of
+ * the stepId, mirroring the SDK's internal calculation — is only added when a stepId exists.
+ * @param {Record<string, string>} meta - The span meta/tags object to populate.
+ * @param {object} [ctxImpl] - The DurableContextImpl about to run the op.
+ * @returns {void}
+ */
+function addOpMeta (meta, ctxImpl) {
+  const stepId = ctxImpl?.getNextStepId?.()
+  if (!stepId) {
+    meta['aws.durable.replayed'] = 'false'
+    return
+  }
+  const stepData = ctxImpl?._executionContext?.getStepData?.(stepId)
+  meta['aws.durable.replayed'] = String(stepData?.Status === 'SUCCEEDED')
+  meta['aws.durable.operation_id'] = createHash('md5').update(stepId).digest('hex').slice(0, 16)
+}
+
+/**
+ * The SDK wraps user errors in typed classes (StepError, ChildContextError, etc.); we follow the
+ * `.cause` chain to recover the user's original Error. SDK wrappers expose a string `errorType`
+ * field, so the loop stops once we leave the wrapper hierarchy.
+ * @param {{ error?: unknown }} ctx
+ * @returns {unknown} the unwrapped error, or `ctx.error` unchanged when it isn't an Error
+ */
+function unwrapDurableError (ctx) {
+  let err = ctx?.error
+  if (!(err instanceof Error)) return err
+
+  while (typeof err.errorType === 'string' && err.cause instanceof Error) {
+    err = err.cause
+  }
+  return err
+}
+
+module.exports = { addOpMeta, unwrapDurableError }

package/packages/datadog-plugin-aws-sdk/src/base.js

@@ -172,18 +172,12 @@ class BaseAwsSdkPlugin extends ClientPlugin {
     })
 
     this.addSub(`apm:aws:request:complete:${this.serviceIdentifier}`, ctx => {
-      const { response, cbExists = false, currentStore } = ctx
+      const { response, currentStore } = ctx
       if (!currentStore) return
       const { span } = currentStore
       if (!span) return
 
       storage('legacy').run(currentStore, () => {
-        // try to extract DSM context from response if no callback exists as extraction normally happens in CB
-        if (!cbExists && this.serviceIdentifier === 'sqs') {
-          const params = response.request.params
-          const operation = response.request.operation
-          this.responseExtractDSMContext(operation, params, response.data ?? response, span)
-        }
         this.addResponseTags(span, response)
 
         if (this._tracerConfig?.DD_TRACE_AWS_ADD_SPAN_POINTERS) {

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

@@ -8,11 +8,18 @@ function recordDataAsString (data) {
   return Buffer.isBuffer(data) ? data.toString('utf8') : Buffer.from(data).toString('utf8')
 }
 
+// Caps the promise-path iterator→stream cache so abandoned shard iterators
+// (AWS expires them after 5 minutes) can't grow it without bound. Polling loops
+// delete on consume, so their working set is ~the active shard count.
+const MAX_TRACKED_SHARD_ITERATORS = 1000
+
 class Kinesis extends BaseAwsSdkPlugin {
   static id = 'kinesis'
   static peerServicePrecursors = ['streamname']
   static isPayloadReporter = true
 
+  #shardIteratorStreams = new Map()
+
   constructor (...args) {
     super(...args)
 
@@ -20,51 +27,107 @@ class Kinesis extends BaseAwsSdkPlugin {
     // in the base class
     this.requestTags = new WeakMap()
 
-    this.addBind('apm:aws:response:start:kinesis', ctx => {
-      const { request, response } = ctx
-      const plugin = this
+    this.addBind('apm:aws:response:start:kinesis', ctx => this.#startResponseSpan(ctx))
+
+    // Promise / event-emitter calls never publish response:start, so create and finish the
+    // consumer span from request:complete instead. Callback calls handle it via the bind above.
+    this.addSub('apm:aws:request:complete:kinesis', ctx => {
+      if (ctx.cbExists) return
+      // v2 nests the SDK payload under response.data; v3 spreads the output onto response.
+      const response = ctx.response?.data ?? ctx.response
+      const responseCtx = { request: ctx.request, response }
+      this.#startResponseSpan(responseCtx)
+      if (responseCtx.needsFinish) this.finish(responseCtx)
+      // The async store that carries streamName to getRecords on the callback path is
+      // absent here, so map each shard iterator to its stream for the DSM topic tag.
+      if (this.config.dsmEnabled) this.#trackShardStream(ctx.request, response)
+    })
 
-      let store = this._parentMap.get(request)
+    this.addSub('apm:aws:response:finish:kinesis', ctx => {
+      if (!ctx.needsFinish) return
+      this.finish(ctx)
+    })
+  }
 
-      // if we have either of these operations, we want to store the streamName param
-      // since it is not typically available during get/put records requests
-      if (request.operation === 'getShardIterator' || request.operation === 'listShards') {
-        return this.storeStreamName(request.params, request.operation, store)
-      }
+  /**
+   * @param {object} ctx Completion context carrying the SDK request and response.
+   */
+  #startResponseSpan (ctx) {
+    const { request, response } = ctx
 
-      if (request.operation === 'getRecords') {
-        let span
-        const responseExtraction = this.responseExtract(request.params, request.operation, response)
-        if (responseExtraction && responseExtraction.maybeChildOf) {
-          ctx.needsFinish = true
-          const options = {
-            childOf: responseExtraction.maybeChildOf,
-            meta: {
-              ...this.requestTags.get(request),
-              'span.kind': 'server',
-            },
-            integrationName: 'aws-sdk',
-          }
-          span = plugin.startSpan('aws.response', options, ctx)
-          store = ctx.currentStore
-        }
+    let store = this._parentMap.get(request)
 
-        // get the stream name that should have been stored previously
-        const { streamName } = store
+    // if we have either of these operations, we want to store the streamName param
+    // since it is not typically available during get/put records requests
+    if (request.operation === 'getShardIterator' || request.operation === 'listShards') {
+      return this.storeStreamName(request.params, request.operation, store)
+    }
 
-        // extract DSM context after as we might not have a parent-child but may have a DSM context
-        this.responseExtractDSMContext(
-          request.operation, request.params, response, span || null, { streamName }
-        )
+    if (request.operation === 'getRecords') {
+      let span
+      const responseExtraction = this.responseExtract(request.params, request.operation, response)
+      if (responseExtraction && responseExtraction.maybeChildOf) {
+        ctx.needsFinish = true
+        const options = {
+          childOf: responseExtraction.maybeChildOf,
+          meta: {
+            ...this.requestTags.get(request),
+            'span.kind': 'server',
+          },
+          integrationName: 'aws-sdk',
+        }
+        span = this.startSpan('aws.response', options, ctx)
+        store = ctx.currentStore
       }
 
-      return store
-    })
+      if (this.config.dsmEnabled) {
+        // streamName rides the async store on the callback path; the promise path has no
+        // such link, so fall back to the iterator the producer returned.
+        const streamName = store?.streamName ?? this.#shardIteratorStreams.get(request.params.ShardIterator)
+        this.responseExtractDSMContext(request.operation, request.params, response, span || null, { streamName })
+      }
+    }
 
-    this.addSub('apm:aws:response:finish:kinesis', ctx => {
-      if (!ctx.needsFinish) return
-      this.finish(ctx)
-    })
+    return store
+  }
+
+  /**
+   * @param {object} request SDK request; reads `operation` and `params`.
+   * @param {object} response SDK output; reads `ShardIterator` / `NextShardIterator`.
+   */
+  #trackShardStream (request, response) {
+    if (request.operation === 'getShardIterator') {
+      this.#rememberShardStream(response?.ShardIterator, request.params?.StreamName)
+    } else if (request.operation === 'getRecords') {
+      this.#advanceShardStream(request.params?.ShardIterator, response?.NextShardIterator)
+    }
+  }
+
+  /**
+   * @param {string} [iterator] Shard iterator the producer returned.
+   * @param {string} [streamName] Stream the iterator belongs to.
+   */
+  #rememberShardStream (iterator, streamName) {
+    if (!iterator || streamName === undefined) return
+    // FIFO-evict the oldest entry (Map keeps insertion order) when the cap is hit; only
+    // abandoned iterators get here, so no realistic test drives the cap (eviction ignored).
+    /* istanbul ignore if */
+    if (this.#shardIteratorStreams.size >= MAX_TRACKED_SHARD_ITERATORS) {
+      this.#shardIteratorStreams.delete(this.#shardIteratorStreams.keys().next().value)
+    }
+    this.#shardIteratorStreams.set(iterator, streamName)
+  }
+
+  /**
+   * @param {string} [consumedIterator] Iterator just passed to getRecords.
+   * @param {string} [nextIterator] NextShardIterator for the following poll.
+   */
+  #advanceShardStream (consumedIterator, nextIterator) {
+    const streamName = this.#shardIteratorStreams.get(consumedIterator)
+    if (streamName === undefined) return
+    this.#shardIteratorStreams.delete(consumedIterator)
+    // carry the stream onto the next iterator so the polling loop keeps its topic
+    if (nextIterator) this.#rememberShardStream(nextIterator, streamName)
   }
 
   generateTags (params, operation, response) {

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

@@ -54,47 +54,64 @@ class Sqs extends BaseAwsSdkPlugin {
     // in the base class
     this.requestTags = new WeakMap()
 
-    this.addBind('apm:aws:response:start:sqs', ctx => {
-      const { request, response } = ctx
-      const contextExtraction = this.responseExtract(request.params, request.operation, response)
-
-      let store = this._parentMap.get(request)
-      let span
-      let parsedMessageAttributes
-      let parsedFirstBody
-      let firstBodyChecked = false
-      if (contextExtraction !== undefined) {
-        parsedFirstBody = contextExtraction.parsedBody
-        firstBodyChecked = contextExtraction.bodyChecked === true
-        if (contextExtraction.datadogContext !== undefined) {
+    this.addBind('apm:aws:response:start:sqs', ctx => this.#startResponseSpan(ctx))
+
+    // No-callback receives (promises, event emitters) never publish response:start, so link and
+    // finish the consumer span here instead. Callback paths reach the same logic via the bind above.
+    this.addSub('apm:aws:request:complete:sqs', ctx => {
+      if (ctx.cbExists) return
+      // v2 nests the SDK payload under response.data; v3 spreads the output onto response.
+      const responseCtx = { request: ctx.request, response: ctx.response?.data ?? ctx.response }
+      this.#startResponseSpan(responseCtx)
+      if (responseCtx.needsFinish) this.finish(responseCtx)
+    })
+
+    this.addSub('apm:aws:response:finish:sqs', ctx => {
+      if (!ctx.needsFinish) return
+      this.finish(ctx)
+    })
+  }
+
+  #startResponseSpan (ctx) {
+    const { request, response } = ctx
+    const contextExtraction = this.responseExtract(request.params, request.operation, response)
+
+    let store = this._parentMap.get(request)
+    let span
+    let parsedMessageAttributes
+    let parsedFirstBody
+    let firstBodyChecked = false
+    if (contextExtraction !== undefined) {
+      parsedFirstBody = contextExtraction.parsedBody
+      firstBodyChecked = contextExtraction.bodyChecked === true
+      if (contextExtraction.datadogContext !== undefined) {
+        parsedMessageAttributes = contextExtraction.parsedAttributes
+        // request:start records requestTags only after the isEnabled gate, so an absent entry
+        // means this consumer is disabled — gate on it instead of paying isEnabled again here.
+        const requestTags = this.requestTags.get(request)
+        if (requestTags !== undefined) {
           ctx.needsFinish = true
           const options = {
             childOf: contextExtraction.datadogContext,
             meta: {
-              ...this.requestTags.get(request),
+              ...requestTags,
               'span.kind': 'server',
             },
             integrationName: 'aws-sdk',
           }
-          parsedMessageAttributes = contextExtraction.parsedAttributes
           span = this.startSpan('aws.response', options, ctx)
           store = ctx.currentStore
         }
       }
+    }
 
-      // Extract DSM context after, as we might not have a parent-child but may have a DSM context.
-      this.responseExtractDSMContext(
-        request.operation, request.params, response, span ?? null,
-        { parsedAttributes: parsedMessageAttributes, parsedFirstBody, firstBodyChecked }
-      )
-
-      return store
-    })
+    // Extract DSM context after, as we might not have a parent-child but may have a DSM context.
+    this.responseExtractDSMContext(
+      request.operation, request.params, response, span ?? null,
+      { parsedAttributes: parsedMessageAttributes, parsedFirstBody, firstBodyChecked }
+    )
 
-    this.addSub('apm:aws:response:finish:sqs', ctx => {
-      if (!ctx.needsFinish) return
-      this.finish(ctx)
-    })
+    return store
   }
 
   operationFromRequest (request) {

package/packages/datadog-plugin-bullmq/src/filter.js

@@ -0,0 +1,35 @@
+'use strict'
+
+const log = require('../../dd-trace/src/log')
+
+/**
+ * @typedef {object} BullmqJobShape
+ * @property {string} [name]
+ * @property {unknown} [data]
+ * @property {unknown} [opts]
+ * @property {string} [queueName]
+ */
+
+/**
+ * @typedef {(job: BullmqJobShape) => boolean} BullmqFilter
+ */
+
+/**
+ * Resolve a user-provided filter from plugin config. Returns `undefined` when
+ * no filter is configured so callers can short-circuit the filtering path on a
+ * cheap truthy check. If `producerFilter` is present but not a function, log
+ * an error and fall back to no filter.
+ *
+ * @param {{ producerFilter?: unknown }} config Plugin config that may carry a `producerFilter`.
+ * @returns {BullmqFilter | undefined}
+ */
+function getFilter (config) {
+  if (typeof config?.producerFilter === 'function') {
+    return /** @type {BullmqFilter} */ (config.producerFilter)
+  }
+  if (config?.producerFilter !== undefined) {
+    log.error('Expected `producerFilter` to be a function. Ignoring.')
+  }
+}
+
+module.exports = { getFilter }