Skip to content

Commit 56965ac

Browse files
avedmalaavedmala
authored
Step Functions Trace Merging Support (#588)
* initial impl for setting and using the new contexts * fix current tests and lint * fix JSON parsing in StepFunctionContext creation * add state machine case to trigger tag parsing
1 parent 8087c87 commit 56965ac

File tree

7 files changed

+419
-205
lines changed

7 files changed

+419
-205
lines changed

‎event_samples/states.json

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"_datadog": {
3+
"Execution": {
4+
"Id": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c",
5+
"StartTime": "2024-11-13T16:46:47.715Z",
6+
"Name": "1356a963-42a5-48b0-ba3f-73bde559a50c",
7+
"RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-MyStateMachine-wsx8chv4d-role-1su0fkfd3",
8+
"RedriveCount": 0
9+
},
10+
"StateMachine": {
11+
"Id": "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d",
12+
"Name": "MyStateMachine-wsx8chv4d"
13+
},
14+
"State": {
15+
"Name": "Lambda Invoke",
16+
"EnteredTime": "2024-11-13T16:46:47.740Z",
17+
"RetryCount": 0
18+
},
19+
"RootExecutionId": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c",
20+
"serverless-version": "v1"
21+
}
22+
}

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

Lines changed: 87 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -874,8 +874,8 @@ describe("TraceContextExtractor", () => {
874874
expect(extractor).toBeInstanceOf(_class);
875875
});
876876

877-
it("returns StepFunctionEventTraceExtractor when event contains StepFunctionContext", () => {
878-
const event = {
877+
it("returns StepFunctionEventTraceExtractor when event contains LegacyStepFunctionContext", () => {
878+
const legacyStepFunctionEvent = {
879879
Execution: {
880880
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
881881
Input: {
@@ -900,10 +900,90 @@ describe("TraceContextExtractor", () => {
900900
const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
901901

902902
// Mimick TraceContextService.extract initialization
903-
const instance = StepFunctionContextService.instance(event);
903+
const instance = StepFunctionContextService.instance(legacyStepFunctionEvent);
904904
traceContextExtractor["stepFunctionContextService"] = instance;
905905

906-
const extractor = traceContextExtractor["getTraceEventExtractor"](event);
906+
const extractor = traceContextExtractor["getTraceEventExtractor"](legacyStepFunctionEvent);
907+
908+
expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor);
909+
});
910+
911+
it("returns StepFunctionEventTraceExtractor when event contains LambdaRootStepFunctionContext", () => {
912+
const lambdaRootStepFunctionEvent = {
913+
_datadog: {
914+
Execution: {
915+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
916+
Input: {
917+
MyInput: "MyValue",
918+
},
919+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
920+
RoleArn:
921+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
922+
StartTime: "2022-12-08T21:08:17.924Z",
923+
},
924+
State: {
925+
Name: "step-one",
926+
EnteredTime: "2022-12-08T21:08:19.224Z",
927+
RetryCount: 2,
928+
},
929+
StateMachine: {
930+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
931+
Name: "my-state-machine",
932+
},
933+
"x-datadog-trace-id": "10593586103637578129",
934+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
935+
"serverless-version": "v1",
936+
},
937+
};
938+
939+
const tracerWrapper = new TracerWrapper();
940+
const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
941+
942+
// Mimick TraceContextService.extract initialization
943+
const instance = StepFunctionContextService.instance(lambdaRootStepFunctionEvent);
944+
traceContextExtractor["stepFunctionContextService"] = instance;
945+
946+
const extractor = traceContextExtractor["getTraceEventExtractor"](lambdaRootStepFunctionEvent);
947+
948+
expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor);
949+
});
950+
951+
it("returns StepFunctionEventTraceExtractor when event contains NestedStepFunctionContext", () => {
952+
const nestedStepFunctionEvent = {
953+
_datadog: {
954+
Execution: {
955+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
956+
Input: {
957+
MyInput: "MyValue",
958+
},
959+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
960+
RoleArn:
961+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
962+
StartTime: "2022-12-08T21:08:17.924Z",
963+
},
964+
State: {
965+
Name: "step-one",
966+
EnteredTime: "2022-12-08T21:08:19.224Z",
967+
RetryCount: 2,
968+
},
969+
StateMachine: {
970+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
971+
Name: "my-state-machine",
972+
},
973+
RootExecutionId:
974+
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc",
975+
"serverless-version": "v1",
976+
},
977+
};
978+
979+
const tracerWrapper = new TracerWrapper();
980+
const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
981+
982+
// Mimick TraceContextService.extract initialization
983+
const instance = StepFunctionContextService.instance(nestedStepFunctionEvent);
984+
traceContextExtractor["stepFunctionContextService"] = instance;
985+
986+
const extractor = traceContextExtractor["getTraceEventExtractor"](nestedStepFunctionEvent);
907987

908988
expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor);
909989
});
@@ -946,7 +1026,7 @@ describe("TraceContextExtractor", () => {
9461026
});
9471027
});
9481028

949-
describe("addTraceContexToXray", () => {
1029+
describe("addTraceContextToXray", () => {
9501030
beforeEach(() => {
9511031
StepFunctionContextService["_instance"] = undefined as any;
9521032
sentSegment = undefined;
@@ -955,7 +1035,7 @@ describe("TraceContextExtractor", () => {
9551035
process.env["AWS_XRAY_DAEMON_ADDRESS"] = undefined;
9561036
});
9571037

958-
it("adds StepFunction context when present over metadata", () => {
1038+
it("adds legacy StepFunction context when present over metadata", () => {
9591039
jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
9601040

9611041
process.env["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
@@ -1006,7 +1086,7 @@ describe("TraceContextExtractor", () => {
10061086

10071087
const sentMessage = sentSegment.toString();
10081088
expect(sentMessage).toEqual(
1009-
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"root_span_metadata":{"step_function.execution_name":"85a9933e-9e11-83dc-6a61-b92367b6c3be","step_function.execution_id":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf","step_function.execution_input":{"MyInput":"MyValue"},"step_function.execution_role_arn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03","step_function.execution_start_time":"2022-12-08T21:08:17.924Z","step_function.state_entered_time":"2022-12-08T21:08:19.224Z","step_function.state_machine_arn":"arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential","step_function.state_machine_name":"my-state-machine","step_function.state_name":"step-one","step_function.state_retry_count":2}}}}',
1089+
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"root_span_metadata":{"execution_id":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf","state_entered_time":"2022-12-08T21:08:19.224Z","state_name":"step-one"}}}}',
10101090
);
10111091
});
10121092

0 commit comments

Comments
 (0)