durable: extract trace context from checkpoints

README.md

@@ -30,6 +30,7 @@ Besides the environment variables supported by dd-trace-js, the datadog-lambda-j
 | DD_COLD_START_TRACE_SKIP_LIB | optionally skip creating Cold Start Spans for a comma-separated list of libraries. Useful to limit depth or skip known libraries. | `./opentracing/tracer` |
 | DD_CAPTURE_LAMBDA_PAYLOAD | [Captures incoming and outgoing AWS Lambda payloads][1] in the Datadog APM spans for Lambda invocations. | `false` |
 | DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH | Determines the level of detail captured from AWS Lambda payloads, which are then assigned as tags for the `aws.lambda` span. It specifies the nesting depth of the JSON payload structure to process. Once the specified maximum depth is reached, the tag's value is set to the stringified value of any nested elements beyond this level.  <br> For example, given the input payload: <pre>{<br>  "lv1" : {<br>    "lv2": {<br>      "lv3": "val"<br>    }<br>  }<br>}</pre> If the depth is set to `2`, the resulting tag's key is set to `function.request.lv1.lv2` and the value is `{\"lv3\": \"val\"}`. <br> If the depth is set to `0`, the resulting tag's key is set to `function.request` and value is `{\"lv1\":{\"lv2\":{\"lv3\": \"val\"}}}` | `10` |
+| DD_DURABLE_CROSS_INVOCATION_TRACING_ENABLED | For AWS Durable functions, the tracer creates extra checkpoints named `_datadog_{N}` to propagate trace context across function invocations, keeping spans from multiple invocations in one intact trace for each durable execution. | `true` |
 
 
 ## Lambda Profiling Beta

src/trace/context/extractor.ts

@@ -5,6 +5,7 @@ import { XrayService } from "../xray-service";
 import {
   AppSyncEventTraceExtractor,
   CustomTraceExtractor,
+  DurableExecutionEventTraceExtractor,
   EventBridgeEventTraceExtractor,
   EventBridgeSQSEventTraceExtractor,
   HTTPEventTraceExtractor,
@@ -81,6 +82,9 @@ export class TraceContextExtractor {
   private getTraceEventExtractor(event: any): EventTraceExtractor | undefined {
     if (!event || typeof event !== "object") return;
 
+    if (EventValidator.isDurableExecutionEvent(event))
+      return new DurableExecutionEventTraceExtractor(this.tracerWrapper);
+
     const headers = event.headers ?? event.multiValueHeaders;
     if (headers !== null && typeof headers === "object") {
       return new HTTPEventTraceExtractor(this.tracerWrapper, this.config.decodeAuthorizerContext);

src/trace/context/extractors/durable-execution.spec.ts

@@ -0,0 +1,65 @@
+import { DurableExecutionEventTraceExtractor } from "./durable-execution";
+import { TracerWrapper } from "../../tracer-wrapper";
+
+function makeTracerWrapper(extractReturn: any = null): TracerWrapper {
+  return {
+    extract: jest.fn().mockReturnValue(extractReturn),
+  } as unknown as TracerWrapper;
+}
+
+describe("DurableExecutionEventTraceExtractor", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("extracts checkpoint headers via the standard propagator", () => {
+    const executionArn = "arn:aws:lambda:us-east-2:123456789012:function:demo:$LATEST/durable-execution/demo/abc";
+
+    const checkpointHeaders = {
+      "x-datadog-trace-id": "149750110124521191",
+      "x-datadog-parent-id": "987654321012345678",
+      "x-datadog-sampling-priority": "1",
+    };
+
+    const event = {
+      DurableExecutionArn: executionArn,
+      CheckpointToken: "t-1",
+      InitialExecutionState: {
+        Operations: [
+          {
+            Id: "op-1",
+            Name: "_datadog_0",
+            Status: "SUCCEEDED",
+            StepDetails: {
+              Result: JSON.stringify(checkpointHeaders),
+            },
+          },
+        ],
+      },
+    };
+
+    const sentinelContext = { sentinel: true };
+    const tracerWrapper = makeTracerWrapper(sentinelContext);
+    const extractor = new DurableExecutionEventTraceExtractor(tracerWrapper);
+    const context = extractor.extract(event);
+
+    // Checkpoint headers are Datadog-style; the default extract list includes
+    // `datadog`, so the standard extract path picks them up.
+    expect(tracerWrapper.extract).toHaveBeenCalledWith(checkpointHeaders);
+    expect(context).toBe(sentinelContext);
+  });
+
+  it("returns null when no checkpoint exists", () => {
+    const tracerWrapper = makeTracerWrapper();
+    const extractor = new DurableExecutionEventTraceExtractor(tracerWrapper);
+
+    const context = extractor.extract({
+      DurableExecutionArn: "arn:aws:lambda:us-east-2:123:function:demo",
+      CheckpointToken: "t-empty",
+      InitialExecutionState: { Operations: [] },
+    });
+
+    expect(context).toBeNull();
+    expect(tracerWrapper.extract).not.toHaveBeenCalled();
+  });
+});

src/trace/context/extractors/durable-execution.ts

@@ -0,0 +1,105 @@
+/**
+ * Durable Execution Trace Extractor — Checkpoint Approach
+ *
+ * Strategy:
+ * 1. Look for trace context in the latest `_datadog_{N}` checkpoint.
+ * 2. If no trace checkpoint exists, return null and let the default extraction
+ *    path create the context.
+ *
+ * The extracted context becomes the parent of the `aws.lambda` span (and any
+ * downstream spans created by dd-trace-js, including `aws.durable.execute`).
+ * Therefore all `aws.lambda` spans will be anchored to the first
+ * `aws.durable.execute` span for a durable execution.
+ *
+ * Checkpoint data will be written by the dd-trace-js plugin in Datadog style
+ * (`x-datadog-*`). Extraction goes through the standard `TracerWrapper.extract`
+ * path, which honors `DD_TRACE_PROPAGATION_STYLE_EXTRACT`. The default extract
+ * list (`datadog, tracecontext, baggage`) already includes `datadog`. Customers
+ *  who override the extract list MUST keep `datadog` in it.
+ */
+
+import { logDebug } from "../../../utils";
+import { SpanContextWrapper } from "../../span-context-wrapper";
+import { TracerWrapper } from "../../tracer-wrapper";
+import { EventTraceExtractor } from "../extractor";
+
+const TRACE_CHECKPOINT_NAME_PREFIX = "_datadog_";
+
+interface CheckpointOperation {
+  Name?: string;
+  Payload?: string;
+  StepDetails?: { Result?: string };
+}
+
+interface DurableExecutionEventShape {
+  DurableExecutionArn?: string;
+  InitialExecutionState?: { Operations?: CheckpointOperation[] };
+}
+
+function parseTraceCheckpointNumber(name: unknown): number | null {
+  if (typeof name !== "string") return null;
+
+  if (!name.startsWith(TRACE_CHECKPOINT_NAME_PREFIX)) return null;
+  const suffix = name.slice(TRACE_CHECKPOINT_NAME_PREFIX.length);
+  const n = Number.parseInt(suffix, 10);
+  if (Number.isNaN(n) || String(n) !== suffix) return null;
+  return n;
+}
+
+/**
+ * Find the highest-numbered `_datadog_{N}` checkpoint in the event and return
+ * its parsed header dict.
+ *
+ * Each invocation that changes trace context saves a new checkpoint with N+1;
+ * the one with the highest N is the most recent. Headers are written by the
+ * dd-trace-js plugin via `tracer.inject(span, 'http_headers', headers)` so the
+ * payload is a standard HTTP-style header dict.
+ *
+ */
+function findLatestCheckpointHeaders(event: DurableExecutionEventShape): Record<string, string> | null {
+  const operations = event.InitialExecutionState?.Operations;
+  if (!operations || operations.length === 0) return null;
+
+  let best: { number: number; op: CheckpointOperation } | null = null;
+  for (const op of operations) {
+    const n = parseTraceCheckpointNumber(op?.Name);
+    if (n === null) continue;
+    if (best === null || n > best.number) {
+      best = { number: n, op };
+    }
+  }
+  if (best === null) return null;
+
+  const raw = best.op.Payload ?? best.op.StepDetails?.Result;
+  if (!raw || typeof raw !== "string") return null;
+  try {
+    const parsed = JSON.parse(raw);
+    if (parsed && typeof parsed === "object") {
+      return parsed as Record<string, string>;
+    }
+  } catch (e) {
+    logDebug(`Failed to parse trace checkpoint payload: ${e}`);
+  }
+  return null;
+}
+
+export class DurableExecutionEventTraceExtractor implements EventTraceExtractor {
+  constructor(private tracerWrapper: TracerWrapper) {}
+
+  extract(event: unknown): SpanContextWrapper | null {
+    const e = event as DurableExecutionEventShape | undefined;
+    if (!e?.DurableExecutionArn) {
+      logDebug("No DurableExecutionArn in event");
+      return null;
+    }
+
+    const checkpointHeaders = findLatestCheckpointHeaders(e);
+    if (checkpointHeaders) {
+      logDebug("Extracting trace context from durable checkpoint");
+      return this.tracerWrapper.extract(checkpointHeaders);
+    }
+
+    logDebug("No durable trace context found; deferring to default extraction");
+    return null;
+  }
+}

src/trace/context/extractors/index.ts

@@ -9,3 +9,4 @@ export { SNSSQSEventTraceExtractor } from "./sns-sqs";
 export { StepFunctionEventTraceExtractor } from "./step-function";
 export { LambdaContextTraceExtractor } from "./lambda-context";
 export { CustomTraceExtractor } from "./custom";
+export { DurableExecutionEventTraceExtractor } from "./durable-execution";

src/utils/event-validator.ts

@@ -74,4 +74,8 @@ export class EventValidator {
   static isKinesisStreamEvent(event: any): event is KinesisStreamEvent {
     return Array.isArray(event.Records) && event.Records.length > 0 && event.Records[0].kinesis !== undefined;
   }
+
+  static isDurableExecutionEvent(event: any): boolean {
+    return typeof event.DurableExecutionArn === "string";
+  }
 }