Skip to content

Commit 1e6a582

Browse files
kimi-pkimi-p
committed
Update Step Function Context Parsing
1 parent a891948 commit 1e6a582

File tree

3 files changed

+157
-74
lines changed

3 files changed

+157
-74
lines changed

‎src/index.spec.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ describe("datadog", () => {
331331

332332
expect(mockedIncrementInvocations).toBeCalledTimes(1);
333333
expect(mockedIncrementInvocations).toBeCalledWith(expect.anything(), mockContext);
334-
expect(logger.debug).toHaveBeenCalledTimes(9);
334+
expect(logger.debug).toHaveBeenCalledTimes(11);
335335
expect(logger.debug).toHaveBeenLastCalledWith('{"status":"debug","message":"datadog:Unpatching HTTP libraries"}');
336336
});
337337
});

‎src/trace/context.spec.ts

Lines changed: 84 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -667,31 +667,62 @@ describe("readTraceFromLambdaContext", () => {
667667
});
668668

669669
describe("readStepFunctionContextFromEvent", () => {
670-
const stepFunctionEvent = {
671-
dd: {
672-
Execution: {
673-
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
674-
StartTime: "2019-09-30T20:28:24.236Z",
675-
},
676-
State: {
677-
Name: "step-one",
678-
RetryCount: 2,
679-
},
680-
StateMachine: {
681-
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
682-
Name: "my-state-machine",
683-
},
670+
const stepFunctionEventWithoutInput = {
671+
Execution: {
672+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
673+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
674+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
675+
StartTime: "2022-12-08T21:08:17.924Z",
684676
},
677+
State: {
678+
Name: "step-one",
679+
EnteredTime: "2022-12-08T21:08:19.224Z",
680+
RetryCount: 2,
681+
},
682+
StateMachine: {
683+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
684+
Name: "my-state-machine",
685+
},
686+
} as const;
687+
688+
const stepFunctionEvent = {
689+
...stepFunctionEventWithoutInput,
690+
Execution: {
691+
...stepFunctionEventWithoutInput.Execution,
692+
Input: {
693+
MyInput: "MyValue"
694+
}
695+
}
685696
} as const;
686-
it("reads a trace from an execution id", () => {
697+
698+
it("reads a step function context from event without Execution.Input", () => {
699+
const result = readStepFunctionContextFromEvent(stepFunctionEventWithoutInput);
700+
expect(result).toEqual({
701+
"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",
702+
"step_function.execution_input": null,
703+
"step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be",
704+
"step_function.execution_role_arn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
705+
"step_function.execution_start_time": "2022-12-08T21:08:17.924Z",
706+
"step_function.state_entered_time": "2022-12-08T21:08:19.224Z",
707+
"step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
708+
"step_function.state_machine_name": "my-state-machine",
709+
"step_function.state_name": "step-one",
710+
"step_function.state_retry_count": 2,
711+
});
712+
});
713+
it("reads a step function context from event with Execution.Input", () => {
687714
const result = readStepFunctionContextFromEvent(stepFunctionEvent);
688715
expect(result).toEqual({
689-
"step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
690-
"step_function.retry_count": 2,
691-
"step_function.state_machine_arn":
692-
"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
716+
"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",
717+
"step_function.execution_input": { "MyInput": "MyValue" },
718+
"step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be",
719+
"step_function.execution_role_arn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
720+
"step_function.execution_start_time": "2022-12-08T21:08:17.924Z",
721+
"step_function.state_entered_time": "2022-12-08T21:08:19.224Z",
722+
"step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
693723
"step_function.state_machine_name": "my-state-machine",
694-
"step_function.step_name": "step-one",
724+
"step_function.state_name": "step-one",
725+
"step_function.state_retry_count": 2,
695726
});
696727
});
697728
it("returns undefined when event isn't an object", () => {
@@ -711,7 +742,7 @@ describe("readStepFunctionContextFromEvent", () => {
711742
it("returns undefined when Execution is missing Name field", () => {
712743
const result = readStepFunctionContextFromEvent({
713744
dd: {
714-
...stepFunctionEvent.dd,
745+
...stepFunctionEvent,
715746
Execution: {},
716747
},
717748
});
@@ -720,7 +751,7 @@ describe("readStepFunctionContextFromEvent", () => {
720751
it("returns undefined when Name isn't a string", () => {
721752
const result = readStepFunctionContextFromEvent({
722753
dd: {
723-
...stepFunctionEvent.dd,
754+
...stepFunctionEvent,
724755
Execution: {
725756
Name: 12345,
726757
},
@@ -731,7 +762,7 @@ describe("readStepFunctionContextFromEvent", () => {
731762
it("returns undefined when State isn't defined", () => {
732763
const result = readStepFunctionContextFromEvent({
733764
dd: {
734-
...stepFunctionEvent.dd,
765+
...stepFunctionEvent,
735766
State: undefined,
736767
},
737768
});
@@ -740,9 +771,9 @@ describe("readStepFunctionContextFromEvent", () => {
740771
it("returns undefined when try retry count isn't a number", () => {
741772
const result = readStepFunctionContextFromEvent({
742773
dd: {
743-
...stepFunctionEvent.dd,
774+
...stepFunctionEvent,
744775
State: {
745-
...stepFunctionEvent.dd.State,
776+
...stepFunctionEvent.State,
746777
RetryCount: "1",
747778
},
748779
},
@@ -752,9 +783,9 @@ describe("readStepFunctionContextFromEvent", () => {
752783
it("returns undefined when try step name isn't a string", () => {
753784
const result = readStepFunctionContextFromEvent({
754785
dd: {
755-
...stepFunctionEvent.dd,
786+
...stepFunctionEvent,
756787
State: {
757-
...stepFunctionEvent.dd.State,
788+
...stepFunctionEvent.State,
758789
Name: 1,
759790
},
760791
},
@@ -764,7 +795,7 @@ describe("readStepFunctionContextFromEvent", () => {
764795
it("returns undefined when StateMachine is undefined", () => {
765796
const result = readStepFunctionContextFromEvent({
766797
dd: {
767-
...stepFunctionEvent.dd,
798+
...stepFunctionEvent,
768799
StateMachine: undefined,
769800
},
770801
});
@@ -773,9 +804,9 @@ describe("readStepFunctionContextFromEvent", () => {
773804
it("returns undefined when StateMachineId isn't a string", () => {
774805
const result = readStepFunctionContextFromEvent({
775806
dd: {
776-
...stepFunctionEvent.dd,
807+
...stepFunctionEvent,
777808
StateMachine: {
778-
...stepFunctionEvent.dd.StateMachine,
809+
...stepFunctionEvent.StateMachine,
779810
Id: 1,
780811
},
781812
},
@@ -785,9 +816,9 @@ describe("readStepFunctionContextFromEvent", () => {
785816
it("returns undefined when StateMachineName isn't a string", () => {
786817
const result = readStepFunctionContextFromEvent({
787818
dd: {
788-
...stepFunctionEvent.dd,
819+
...stepFunctionEvent,
789820
StateMachine: {
790-
...stepFunctionEvent.dd.StateMachine,
821+
...stepFunctionEvent.StateMachine,
791822
Name: 1,
792823
},
793824
},
@@ -1062,21 +1093,26 @@ describe("extractTraceContext", () => {
10621093

10631094
it("adds step function metadata to xray", () => {
10641095
const stepFunctionEvent = {
1065-
dd: {
1066-
Execution: {
1067-
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
1068-
StartTime: "2019-09-30T20:28:24.236Z",
1069-
},
1070-
State: {
1071-
Name: "step-one",
1072-
RetryCount: 2,
1073-
},
1074-
StateMachine: {
1075-
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
1076-
Name: "my-state-machine",
1077-
},
1078-
},
1079-
} as const;
1096+
Execution: {
1097+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
1098+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
1099+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
1100+
StartTime: "2022-12-08T21:08:17.924Z",
1101+
Input: {
1102+
MyInput: "MyValue"
1103+
}
1104+
},
1105+
State: {
1106+
Name: "step-one",
1107+
EnteredTime: "2022-12-08T21:08:19.224Z",
1108+
RetryCount: 2,
1109+
},
1110+
StateMachine: {
1111+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
1112+
Name: "my-state-machine",
1113+
},
1114+
} as const;
1115+
10801116
jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
10811117
process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
10821118
process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
@@ -1089,7 +1125,7 @@ describe("extractTraceContext", () => {
10891125
const sentMessage = sentSegment.toString();
10901126
expect(sentMessage).toMatchInlineSnapshot(`
10911127
"{\\"format\\": \\"json\\", \\"version\\": 1}
1092-
{\\"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_id\\":\\"fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492\\",\\"step_function.retry_count\\":2,\\"step_function.state_machine_arn\\":\\"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ\\",\\"step_function.state_machine_name\\":\\"my-state-machine\\",\\"step_function.step_name\\":\\"step-one\\"}}}}"
1128+
{\\"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}}}}"
10931129
`);
10941130
});
10951131
});

‎src/trace/context.ts

Lines changed: 72 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,15 @@
1-
import { Context, EventBridgeEvent, KinesisStreamEvent, SNSEvent, SNSMessage, SQSEvent } from "aws-lambda";
2-
import { BigNumber } from "bignumber.js";
3-
import { randomBytes } from "crypto";
4-
import { createSocket, Socket } from "dgram";
5-
import { logDebug, logError } from "../utils";
1+
import {
2+
Context,
3+
EventBridgeEvent,
4+
KinesisStreamEvent,
5+
SNSEvent,
6+
SNSMessage,
7+
SQSEvent
8+
} from "aws-lambda";
9+
import {BigNumber} from "bignumber.js";
10+
import {randomBytes} from "crypto";
11+
import {createSocket, Socket} from "dgram";
12+
import {logDebug, logError} from "../utils";
613
import {
714
isAppSyncResolverEvent,
815
isEventBridgeEvent,
@@ -12,6 +19,7 @@ import {
1219
isSQSEvent,
1320
} from "../utils/event-type-guards";
1421
import {
22+
authorizingRequestIdHeader,
1523
awsXrayDaemonAddressEnvVar,
1624
parentIDHeader,
1725
SampleMode,
@@ -25,9 +33,8 @@ import {
2533
xraySubsegmentNamespace,
2634
xrayTraceEnvVar,
2735
} from "./constants";
28-
import { TraceExtractor } from "./listener";
29-
import { parseEventSourceSubType, eventSubTypes } from "./trigger";
30-
import { authorizingRequestIdHeader } from "./constants";
36+
import {TraceExtractor} from "./listener";
37+
import {eventSubTypes, parseEventSourceSubType} from "./trigger";
3138

3239
export interface XRayTraceHeader {
3340
traceID: string;
@@ -43,11 +50,16 @@ export interface TraceContext {
4350
}
4451

4552
export interface StepFunctionContext {
46-
"step_function.retry_count": number;
47-
"step_function.execution_id": string;
53+
"step_function.execution_name": string,
54+
"step_function.execution_id": string,
55+
"step_function.execution_input": string | object,
56+
"step_function.execution_role_arn": string,
57+
"step_function.execution_start_time": string,
4858
"step_function.state_machine_name": string;
4959
"step_function.state_machine_arn": string;
50-
"step_function.step_name": string;
60+
"step_function.state_entered_time": string;
61+
"step_function.state_name": string;
62+
"step_function.state_retry_count": number;
5163
}
5264

5365
/**
@@ -485,50 +497,85 @@ export function readStepFunctionContextFromEvent(event: any): StepFunctionContex
485497
if (typeof event !== "object") {
486498
return;
487499
}
488-
const { dd } = event;
489-
if (typeof dd !== "object") {
490-
return;
491-
}
492-
const execution = dd.Execution;
500+
501+
const execution = event.Execution;
493502
if (typeof execution !== "object") {
503+
logDebug("event.Execution is not an object.");
494504
return;
495505
}
496-
const executionID = execution.Name;
506+
const executionID = execution.Id;
497507
if (typeof executionID !== "string") {
508+
logDebug("event.Execution.Id is not a string.")
509+
return;
510+
}
511+
const executionInput = execution.Input;
512+
const executionName = execution.Name;
513+
if (typeof executionName !== "string") {
514+
logDebug("event.Execution.Name is not a string.")
515+
return;
516+
}
517+
const executionRoleArn = execution.RoleArn;
518+
if (typeof executionRoleArn !== "string") {
519+
logDebug("event.Execution.RoleArn is not a string.")
498520
return;
499521
}
500-
const state = dd.State;
522+
const executionStartTime = execution.StartTime;
523+
if (typeof executionStartTime !== "string") {
524+
logDebug("event.Execution.StartTime is not a string.")
525+
return;
526+
}
527+
528+
const state = event.State;
501529
if (typeof state !== "object") {
530+
logDebug("event.State is not an object.")
531+
return;
532+
}
533+
const stateRetryCount = state.RetryCount;
534+
if (typeof stateRetryCount !== "number") {
535+
logDebug("event.State.RetryCount is not a string.")
502536
return;
503537
}
504-
const retryCount = state.RetryCount;
505-
if (typeof retryCount !== "number") {
538+
const stateEnteredTime = state.EnteredTime;
539+
if (typeof stateEnteredTime !== "string") {
540+
logDebug("event.State.EnteredTime is not a string.")
506541
return;
507542
}
508-
const stepName = state.Name;
509-
if (typeof stepName !== "string") {
543+
const stateName = state.Name;
544+
if (typeof stateName !== "string") {
545+
logDebug("event.State.Name is not a string.")
510546
return;
511547
}
512-
const stateMachine = dd.StateMachine;
548+
549+
const stateMachine = event.StateMachine;
513550
if (typeof stateMachine !== "object") {
551+
logDebug("event.StateMachine is not an object.")
514552
return;
515553
}
516554
const stateMachineArn = stateMachine.Id;
517555
if (typeof stateMachineArn !== "string") {
556+
logDebug("event.StateMachine.Id is not a string.")
518557
return;
519558
}
520559
const stateMachineName = stateMachine.Name;
521560
if (typeof stateMachineName !== "string") {
561+
logDebug("event.StateMachine.Name is not a string.")
522562
return;
523563
}
564+
524565
return {
566+
"step_function.execution_name": executionName,
525567
"step_function.execution_id": executionID,
526-
"step_function.retry_count": retryCount,
568+
"step_function.execution_input": executionInput ?? null,
569+
"step_function.execution_role_arn": executionRoleArn,
570+
"step_function.execution_start_time": executionStartTime,
571+
"step_function.state_entered_time": stateEnteredTime,
527572
"step_function.state_machine_arn": stateMachineArn,
528573
"step_function.state_machine_name": stateMachineName,
529-
"step_function.step_name": stepName,
574+
"step_function.state_name": stateName,
575+
"step_function.state_retry_count": stateRetryCount
530576
};
531577
}
578+
532579
export function convertToSampleMode(xraySampled: number): SampleMode {
533580
return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT;
534581
}

0 commit comments

Comments
 (0)