Skip to content

Commit 1619bce

Browse files
ryanstratduncanista
ryanstrat
and
duncanista
authored
Propagate Step Function Trace Context through Managed Services (#667)
* Add SFN -> EventBridge -> Lambda use case * Simplify code * Remove additional classes * Add implementation and test for SFN -> EventBridge -> SQS -> Lambda case * Add support for Eventbridge -> SNS -> Lambda pattern for both header and SFN context * Revert "Add support for Eventbridge -> SNS -> Lambda pattern for both header and SFN context" This reverts commit a146e03. * Add implementation and test for SFN -> SQS -> Lambda * Add implementation and test for SFN -> SNS -> Lambda * Add implementation and test for SFN -> SNS -> SQS -> Lambda Run code formatting * Refactor common extraction logic * Remove unrequired changes to SQS binary format handling * wip * wip * revent yarn.lock changes * simplify extractor-util logic Co-authored-by: jordan gonzález <[email protected]> * Revert changes to layer publishing script Changes to CI should not be made in the same PR as adding features * Fixes test failure from merging main --------- Co-authored-by: jordan gonzález <[email protected]>
1 parent 53a093f commit 1619bce

18 files changed

+907
-379
lines changed

‎src/trace/context/extractor-utils.spec.ts

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import { TracerWrapper } from "../tracer-wrapper";
2+
import { extractTraceContext, extractFromAWSTraceHeader } from "./extractor-utils";
3+
import { StepFunctionContextService } from "../step-function-service";
4+
5+
describe("extractor-utils", () => {
6+
beforeEach(() => {
7+
StepFunctionContextService["_instance"] = undefined as any;
8+
});
9+
describe("extractTraceContext", () => {
10+
it("returns span context when tracer wrapper successfully extracts from headers", () => {
11+
const legacyStepFunctionEvent = {
12+
Execution: {
13+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
14+
Input: {
15+
MyInput: "MyValue",
16+
},
17+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
18+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
19+
RedriveCount: 0,
20+
StartTime: "2022-12-08T21:08:17.924Z",
21+
},
22+
State: {
23+
Name: "step-one",
24+
EnteredTime: "2022-12-08T21:08:19.224Z",
25+
RetryCount: 2,
26+
},
27+
StateMachine: {
28+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
29+
Name: "my-state-machine",
30+
},
31+
};
32+
33+
const tracerWrapper = new TracerWrapper();
34+
const result = extractTraceContext(legacyStepFunctionEvent, tracerWrapper);
35+
36+
// Should return a span context from Step Function context since headers extraction fails
37+
expect(result).not.toBeNull();
38+
});
39+
40+
it("returns null when no trace context can be extracted", () => {
41+
const emptyEvent = {
42+
someOtherProperty: "value",
43+
};
44+
45+
const tracerWrapper = new TracerWrapper();
46+
const result = extractTraceContext(emptyEvent, tracerWrapper);
47+
48+
expect(result).toBeNull();
49+
});
50+
51+
it("extracts context from LambdaRootStepFunctionContext", () => {
52+
const lambdaRootStepFunctionEvent = {
53+
_datadog: {
54+
Execution: {
55+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
56+
Input: {
57+
MyInput: "MyValue",
58+
},
59+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
60+
RoleArn:
61+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
62+
RedriveCount: 0,
63+
StartTime: "2022-12-08T21:08:17.924Z",
64+
},
65+
State: {
66+
Name: "step-one",
67+
EnteredTime: "2022-12-08T21:08:19.224Z",
68+
RetryCount: 2,
69+
},
70+
StateMachine: {
71+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
72+
Name: "my-state-machine",
73+
},
74+
"x-datadog-trace-id": "10593586103637578129",
75+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
76+
"serverless-version": "v1",
77+
},
78+
};
79+
80+
const tracerWrapper = new TracerWrapper();
81+
const result = extractTraceContext(lambdaRootStepFunctionEvent, tracerWrapper);
82+
83+
expect(result).not.toBeNull();
84+
});
85+
86+
it("extracts context from NestedStepFunctionContext", () => {
87+
const nestedStepFunctionEvent = {
88+
_datadog: {
89+
Execution: {
90+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
91+
Input: {
92+
MyInput: "MyValue",
93+
},
94+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
95+
RoleArn:
96+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
97+
RedriveCount: 0,
98+
StartTime: "2022-12-08T21:08:17.924Z",
99+
},
100+
State: {
101+
Name: "step-one",
102+
EnteredTime: "2022-12-08T21:08:19.224Z",
103+
RetryCount: 2,
104+
},
105+
StateMachine: {
106+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
107+
Name: "my-state-machine",
108+
},
109+
RootExecutionId:
110+
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc",
111+
"serverless-version": "v1",
112+
},
113+
};
114+
115+
const tracerWrapper = new TracerWrapper();
116+
const result = extractTraceContext(nestedStepFunctionEvent, tracerWrapper);
117+
118+
expect(result).not.toBeNull();
119+
});
120+
121+
it("extracts context from legacy lambda StepFunctionContext", () => {
122+
const event = {
123+
Payload: {
124+
Execution: {
125+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
126+
Input: {
127+
MyInput: "MyValue",
128+
},
129+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
130+
RoleArn:
131+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
132+
RedriveCount: 0,
133+
StartTime: "2022-12-08T21:08:17.924Z",
134+
},
135+
State: {
136+
Name: "step-one",
137+
EnteredTime: "2022-12-08T21:08:19.224Z",
138+
RetryCount: 2,
139+
},
140+
StateMachine: {
141+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
142+
Name: "my-state-machine",
143+
},
144+
},
145+
};
146+
147+
const tracerWrapper = new TracerWrapper();
148+
const result = extractTraceContext(event, tracerWrapper);
149+
150+
expect(result).not.toBeNull();
151+
});
152+
});
153+
154+
describe("extractFromAWSTraceHeader", () => {
155+
it("returns null when AWS trace header is invalid", () => {
156+
const invalidHeader = "invalid-header";
157+
const eventType = "SQS";
158+
159+
const result = extractFromAWSTraceHeader(invalidHeader, eventType);
160+
161+
expect(result).toBeNull();
162+
});
163+
});
164+
});

‎src/trace/context/extractor-utils.ts

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import { logDebug } from "../../utils";
2+
import { StepFunctionContextService } from "../step-function-service";
3+
import { SpanContextWrapper } from "../span-context-wrapper";
4+
import { TracerWrapper } from "../tracer-wrapper";
5+
import { XrayService } from "../xray-service";
6+
7+
/**
8+
* Common utility functions for trace context extraction
9+
*/
10+
11+
/**
12+
* Attempts to extract trace context from headers, falling back to Step Function context if needed
13+
* @param headers The headers object to extract from
14+
* @param tracerWrapper The tracer wrapper instance
15+
* @returns SpanContextWrapper or null
16+
*/
17+
export function extractTraceContext(headers: any, tracerWrapper: TracerWrapper): SpanContextWrapper | null {
18+
// First try to extract as regular trace headers
19+
const traceContext = tracerWrapper.extract(headers);
20+
if (traceContext) {
21+
return traceContext;
22+
}
23+
24+
// If that fails, check if this is a Step Function context
25+
const stepFunctionInstance = StepFunctionContextService.instance(headers);
26+
27+
if (stepFunctionInstance.context !== undefined) {
28+
if (stepFunctionInstance.spanContext !== null) {
29+
return stepFunctionInstance.spanContext;
30+
}
31+
}
32+
33+
return null;
34+
}
35+
36+
/**
37+
* Extracts trace context from AWS Trace Header
38+
* @param awsTraceHeader The AWS trace header string
39+
* @param eventType The type of event (for logging)
40+
* @returns SpanContextWrapper or null
41+
*/
42+
export function extractFromAWSTraceHeader(awsTraceHeader: string, eventType: string): SpanContextWrapper | null {
43+
const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceHeader);
44+
if (traceContext) {
45+
logDebug(`Extracted trace context from ${eventType} event attributes AWSTraceHeader`);
46+
return traceContext;
47+
} else {
48+
logDebug(`No Datadog trace context found from ${eventType} event attributes AWSTraceHeader`);
49+
return null;
50+
}
51+
}
52+
53+
/**
54+
* Common error handler for extraction operations
55+
* @param error The error that occurred
56+
* @param eventType The type of event (for logging)
57+
*/
58+
export function handleExtractionError(error: unknown, eventType: string): void {
59+
if (error instanceof Error) {
60+
logDebug(`Unable to extract trace context from ${eventType} event`, error);
61+
}
62+
}

0 commit comments

Comments
 (0)