Skip to content

Commit 6319b1d

Browse files
jchrostek-ddclaude
andcommitted
feat: Add automatic trace propagation for Lambda Durable Functions
This adds automatic tracing support for AWS Lambda Durable Functions: - Auto-patches @aws/durable-execution-sdk-js when datadog-lambda-js loads - Extracts trace context from preserved InputPayload (HTTP headers, SNS/SQS attributes) - Falls back to deterministic trace ID from execution ARN when no parent exists - Creates child spans for each operation (step, wait, invoke, parallel, map) - Detects replay operations to avoid duplicate spans Zero user code changes required - existing code with datadog() wrapper works automatically. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 1030376 commit 6319b1d

12 files changed

Lines changed: 1654 additions & 0 deletions

src/index.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import { getEnhancedMetricTags } from "./metrics/enhanced-metrics";
2727
import { DatadogTraceHeaders } from "./trace/context/extractor";
2828
import { SpanWrapper } from "./trace/span-wrapper";
2929
import { SpanOptions, TracerWrapper } from "./trace/tracer-wrapper";
30+
import { initDurableFunctionTracing } from "./trace/durable-function-patch";
3031

3132
// Backwards-compatible export, TODO deprecate in next major
3233
export { DatadogTraceHeaders as TraceHeaders } from "./trace/context/extractor";
@@ -113,6 +114,9 @@ if (getEnvValue(coldStartTracingEnvVar, "true").toLowerCase() === "true" && !isM
113114
subscribeToDC();
114115
}
115116

117+
// Initialize durable function tracing if SDK is present
118+
initDurableFunctionTracing();
119+
116120
const initTime = Date.now();
117121

118122
/**

src/trace/context/extractor.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@ import {
1414
SNSSQSEventTraceExtractor,
1515
SQSEventTraceExtractor,
1616
StepFunctionEventTraceExtractor,
17+
DurableFunctionEventTraceExtractor,
1718
} from "./extractors";
1819
import { StepFunctionContextService } from "../step-function-service";
20+
import { DurableFunctionContextService } from "../durable-function-service";
1921
import { EventValidator } from "../../utils/event-validator";
2022
import { TracerWrapper } from "../tracer-wrapper";
2123
import { SpanContextWrapper } from "../span-context-wrapper";
@@ -37,6 +39,7 @@ export interface DatadogTraceHeaders {
3739
export class TraceContextExtractor {
3840
private xrayService: XrayService;
3941
private stepFunctionContextService?: StepFunctionContextService;
42+
private durableFunctionContextService?: DurableFunctionContextService;
4043

4144
constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {
4245
this.xrayService = new XrayService();
@@ -64,6 +67,14 @@ export class TraceContextExtractor {
6467
}
6568
}
6669

70+
if (spanContext === null) {
71+
this.durableFunctionContextService = DurableFunctionContextService.instance(event);
72+
if (this.durableFunctionContextService?.context) {
73+
const extractor = new DurableFunctionEventTraceExtractor();
74+
spanContext = extractor.extract(event);
75+
}
76+
}
77+
6778
if (spanContext === null) {
6879
const contextExtractor = new LambdaContextTraceExtractor(this.tracerWrapper);
6980
spanContext = contextExtractor.extract(context);
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import { DurableFunctionEventTraceExtractor } from "./durable-function";
2+
import { DurableFunctionContextService } from "../../durable-function-service";
3+
4+
describe("DurableFunctionEventTraceExtractor", () => {
5+
const validDurableExecutionEvent = {
6+
DurableExecutionArn:
7+
"arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/550e8400-e29b-41d4-a716-446655440001",
8+
CheckpointToken: "some-token",
9+
InitialExecutionState: {
10+
Operations: [
11+
{
12+
Id: "step-1",
13+
Type: "STEP",
14+
Status: "SUCCEEDED",
15+
ExecutionDetails: {
16+
InputPayload: JSON.stringify({
17+
headers: {
18+
"x-datadog-trace-id": "12345678901234567890",
19+
"x-datadog-parent-id": "9876543210987654321",
20+
},
21+
}),
22+
},
23+
},
24+
],
25+
},
26+
};
27+
28+
const regularLambdaEvent = {
29+
body: '{"key": "value"}',
30+
headers: {
31+
"Content-Type": "application/json",
32+
},
33+
};
34+
35+
beforeEach(() => {
36+
DurableFunctionContextService.reset();
37+
});
38+
39+
describe("extract", () => {
40+
it("extracts span context from valid durable execution event", () => {
41+
const extractor = new DurableFunctionEventTraceExtractor();
42+
43+
const spanContext = extractor.extract(validDurableExecutionEvent);
44+
45+
expect(spanContext).not.toBeNull();
46+
expect(spanContext?.toTraceId()).toBe("12345678901234567890");
47+
expect(spanContext?.toSpanId()).toBe("9876543210987654321");
48+
expect(spanContext?.source).toBe("event");
49+
});
50+
51+
it("returns null for regular Lambda event", () => {
52+
const extractor = new DurableFunctionEventTraceExtractor();
53+
54+
const spanContext = extractor.extract(regularLambdaEvent);
55+
56+
expect(spanContext).toBeNull();
57+
});
58+
59+
it("returns null for null event", () => {
60+
const extractor = new DurableFunctionEventTraceExtractor();
61+
62+
const spanContext = extractor.extract(null);
63+
64+
expect(spanContext).toBeNull();
65+
});
66+
67+
it("returns null for undefined event", () => {
68+
const extractor = new DurableFunctionEventTraceExtractor();
69+
70+
const spanContext = extractor.extract(undefined);
71+
72+
expect(spanContext).toBeNull();
73+
});
74+
75+
it("returns null for event with invalid DurableExecutionArn", () => {
76+
const extractor = new DurableFunctionEventTraceExtractor();
77+
78+
const spanContext = extractor.extract({
79+
DurableExecutionArn: "invalid-arn",
80+
});
81+
82+
expect(spanContext).toBeNull();
83+
});
84+
85+
it("extracts deterministic span context when no preserved trace headers", () => {
86+
const eventWithoutPreservedTrace = {
87+
DurableExecutionArn:
88+
"arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-456/650e8400-e29b-41d4-a716-446655440002",
89+
CheckpointToken: "some-token",
90+
InitialExecutionState: {
91+
Operations: [
92+
{
93+
Id: "step-1",
94+
Type: "STEP",
95+
Status: "SUCCEEDED",
96+
ExecutionDetails: {
97+
InputPayload: JSON.stringify({
98+
body: '{"key": "value"}',
99+
}),
100+
},
101+
},
102+
],
103+
},
104+
};
105+
106+
const extractor = new DurableFunctionEventTraceExtractor();
107+
108+
const spanContext = extractor.extract(eventWithoutPreservedTrace);
109+
110+
expect(spanContext).not.toBeNull();
111+
// Should have deterministic IDs based on execution ID
112+
expect(spanContext?.toTraceId()).toBeDefined();
113+
expect(spanContext?.toSpanId()).toBeDefined();
114+
});
115+
116+
it("uses singleton service instance", () => {
117+
const extractor = new DurableFunctionEventTraceExtractor();
118+
119+
// First extraction
120+
extractor.extract(validDurableExecutionEvent);
121+
const instance1 = DurableFunctionContextService.instance();
122+
123+
// Second extraction with same event
124+
extractor.extract(validDurableExecutionEvent);
125+
const instance2 = DurableFunctionContextService.instance();
126+
127+
expect(instance1).toBe(instance2);
128+
});
129+
});
130+
});
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import { SpanContextWrapper } from "../../span-context-wrapper";
2+
import { DurableFunctionContextService } from "../../durable-function-service";
3+
import { EventTraceExtractor } from "../extractor";
4+
5+
export class DurableFunctionEventTraceExtractor implements EventTraceExtractor {
6+
extract(event: any): SpanContextWrapper | null {
7+
const durableFunctionInstance = DurableFunctionContextService.instance(event);
8+
const durableFunctionContext = durableFunctionInstance.context;
9+
10+
if (durableFunctionContext !== undefined) {
11+
const spanContext = durableFunctionInstance.spanContext;
12+
if (spanContext !== null) {
13+
return spanContext;
14+
}
15+
}
16+
17+
return null;
18+
}
19+
}

src/trace/context/extractors/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,6 @@ export { SQSEventTraceExtractor } from "./sqs";
77
export { SNSEventTraceExtractor } from "./sns";
88
export { SNSSQSEventTraceExtractor } from "./sns-sqs";
99
export { StepFunctionEventTraceExtractor } from "./step-function";
10+
export { DurableFunctionEventTraceExtractor } from "./durable-function";
1011
export { LambdaContextTraceExtractor } from "./lambda-context";
1112
export { CustomTraceExtractor } from "./custom";

0 commit comments

Comments
 (0)