|
| 1 | +/** |
| 2 | + * Durable Execution Trace Extractor — Checkpoint Approach |
| 3 | + * |
| 4 | + * Strategy: |
| 5 | + * 1. Look for trace context in the latest `_datadog_{N}` checkpoint. |
| 6 | + * 2. If no trace checkpoint exists, return null and let the default extraction |
| 7 | + * path create the context. |
| 8 | + * |
| 9 | + * The extracted context becomes the parent of the `aws.lambda` span (and any |
| 10 | + * downstream spans created by dd-trace-js, including `aws.durable.execute`). |
| 11 | + * Therefore all `aws.lambda` spans will be anchored to the first |
| 12 | + * `aws.durable.execute` span for a durable execution. |
| 13 | + * |
| 14 | + * Checkpoint data will be written by the dd-trace-js plugin in Datadog style |
| 15 | + * (`x-datadog-*`). Extraction goes through the standard `TracerWrapper.extract` |
| 16 | + * path, which honors `DD_TRACE_PROPAGATION_STYLE_EXTRACT`. The default extract |
| 17 | + * list (`datadog, tracecontext, baggage`) already includes `datadog`. Customers |
| 18 | + * who override the extract list MUST keep `datadog` in it. |
| 19 | + */ |
| 20 | + |
| 21 | +import { logDebug } from "../../../utils"; |
| 22 | +import { SpanContextWrapper } from "../../span-context-wrapper"; |
| 23 | +import { TracerWrapper } from "../../tracer-wrapper"; |
| 24 | +import { EventTraceExtractor } from "../extractor"; |
| 25 | + |
| 26 | +const TRACE_CHECKPOINT_NAME_PREFIX = "_datadog_"; |
| 27 | + |
| 28 | +interface CheckpointOperation { |
| 29 | + Name?: string; |
| 30 | + Payload?: string; |
| 31 | + StepDetails?: { Result?: string }; |
| 32 | +} |
| 33 | + |
| 34 | +interface DurableExecutionEventShape { |
| 35 | + DurableExecutionArn?: string; |
| 36 | + InitialExecutionState?: { Operations?: CheckpointOperation[] }; |
| 37 | +} |
| 38 | + |
| 39 | +function parseTraceCheckpointNumber(name: unknown): number | null { |
| 40 | + if (typeof name !== "string") return null; |
| 41 | + |
| 42 | + if (!name.startsWith(TRACE_CHECKPOINT_NAME_PREFIX)) return null; |
| 43 | + const suffix = name.slice(TRACE_CHECKPOINT_NAME_PREFIX.length); |
| 44 | + const n = Number.parseInt(suffix, 10); |
| 45 | + if (Number.isNaN(n) || String(n) !== suffix) return null; |
| 46 | + return n; |
| 47 | +} |
| 48 | + |
| 49 | +/** |
| 50 | + * Find the highest-numbered `_datadog_{N}` checkpoint in the event and return |
| 51 | + * its parsed header dict. |
| 52 | + * |
| 53 | + * Each invocation that changes trace context saves a new checkpoint with N+1; |
| 54 | + * the one with the highest N is the most recent. Headers are written by the |
| 55 | + * dd-trace-js plugin via `tracer.inject(span, 'http_headers', headers)` so the |
| 56 | + * payload is a standard HTTP-style header dict. |
| 57 | + * |
| 58 | + */ |
| 59 | +function findLatestCheckpointHeaders(event: DurableExecutionEventShape): Record<string, string> | null { |
| 60 | + const operations = event.InitialExecutionState?.Operations; |
| 61 | + if (!operations || operations.length === 0) return null; |
| 62 | + |
| 63 | + let best: { number: number; op: CheckpointOperation } | null = null; |
| 64 | + for (const op of operations) { |
| 65 | + const n = parseTraceCheckpointNumber(op?.Name); |
| 66 | + if (n === null) continue; |
| 67 | + if (best === null || n > best.number) { |
| 68 | + best = { number: n, op }; |
| 69 | + } |
| 70 | + } |
| 71 | + if (best === null) return null; |
| 72 | + |
| 73 | + const raw = best.op.Payload ?? best.op.StepDetails?.Result; |
| 74 | + if (!raw || typeof raw !== "string") return null; |
| 75 | + try { |
| 76 | + const parsed = JSON.parse(raw); |
| 77 | + if (parsed && typeof parsed === "object") { |
| 78 | + return parsed as Record<string, string>; |
| 79 | + } |
| 80 | + } catch (e) { |
| 81 | + logDebug(`Failed to parse trace checkpoint payload: ${e}`); |
| 82 | + } |
| 83 | + return null; |
| 84 | +} |
| 85 | + |
| 86 | +export class DurableExecutionEventTraceExtractor implements EventTraceExtractor { |
| 87 | + constructor(private tracerWrapper: TracerWrapper) {} |
| 88 | + |
| 89 | + extract(event: unknown): SpanContextWrapper | null { |
| 90 | + const e = event as DurableExecutionEventShape | undefined; |
| 91 | + if (!e?.DurableExecutionArn) { |
| 92 | + logDebug("No DurableExecutionArn in event"); |
| 93 | + return null; |
| 94 | + } |
| 95 | + |
| 96 | + const checkpointHeaders = findLatestCheckpointHeaders(e); |
| 97 | + if (checkpointHeaders) { |
| 98 | + logDebug("Extracting trace context from durable checkpoint"); |
| 99 | + return this.tracerWrapper.extract(checkpointHeaders); |
| 100 | + } |
| 101 | + |
| 102 | + logDebug("No durable trace context found; deferring to default extraction"); |
| 103 | + return null; |
| 104 | + } |
| 105 | +} |
0 commit comments