diff --git a/event_samples/states.json b/event_samples/states.json new file mode 100644 index 00000000..a65ad73c --- /dev/null +++ b/event_samples/states.json @@ -0,0 +1,22 @@ +{ + "_datadog": { + "Execution": { + "Id": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c", + "StartTime": "2024-11-13T16:46:47.715Z", + "Name": "1356a963-42a5-48b0-ba3f-73bde559a50c", + "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-MyStateMachine-wsx8chv4d-role-1su0fkfd3", + "RedriveCount": 0 + }, + "StateMachine": { + "Id": "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d", + "Name": "MyStateMachine-wsx8chv4d" + }, + "State": { + "Name": "Lambda Invoke", + "EnteredTime": "2024-11-13T16:46:47.740Z", + "RetryCount": 0 + }, + "RootExecutionId": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c", + "serverless-version": "v1" + } + } diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index a4dcd966..ee5860f8 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -874,8 +874,8 @@ describe("TraceContextExtractor", () => { expect(extractor).toBeInstanceOf(_class); }); - it("returns StepFunctionEventTraceExtractor when event contains StepFunctionContext", () => { - const event = { + it("returns StepFunctionEventTraceExtractor when event contains LegacyStepFunctionContext", () => { + const legacyStepFunctionEvent = { Execution: { Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", Input: { @@ -900,10 +900,90 @@ describe("TraceContextExtractor", () => { const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); // Mimick TraceContextService.extract initialization - const instance = StepFunctionContextService.instance(event); + const instance = StepFunctionContextService.instance(legacyStepFunctionEvent); traceContextExtractor["stepFunctionContextService"] = instance; - const extractor = traceContextExtractor["getTraceEventExtractor"](event); + const extractor = traceContextExtractor["getTraceEventExtractor"](legacyStepFunctionEvent); + + expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); + }); + + it("returns StepFunctionEventTraceExtractor when event contains LambdaRootStepFunctionContext", () => { + const lambdaRootStepFunctionEvent = { + _datadog: { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: + "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + "x-datadog-trace-id": "10593586103637578129", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000", + "serverless-version": "v1", + }, + }; + + const tracerWrapper = new TracerWrapper(); + const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + + // Mimick TraceContextService.extract initialization + const instance = StepFunctionContextService.instance(lambdaRootStepFunctionEvent); + traceContextExtractor["stepFunctionContextService"] = instance; + + const extractor = traceContextExtractor["getTraceEventExtractor"](lambdaRootStepFunctionEvent); + + expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); + }); + + it("returns StepFunctionEventTraceExtractor when event contains NestedStepFunctionContext", () => { + const nestedStepFunctionEvent = { + _datadog: { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: + "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + RootExecutionId: + "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc", + "serverless-version": "v1", + }, + }; + + const tracerWrapper = new TracerWrapper(); + const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + + // Mimick TraceContextService.extract initialization + const instance = StepFunctionContextService.instance(nestedStepFunctionEvent); + traceContextExtractor["stepFunctionContextService"] = instance; + + const extractor = traceContextExtractor["getTraceEventExtractor"](nestedStepFunctionEvent); expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); }); @@ -946,7 +1026,7 @@ describe("TraceContextExtractor", () => { }); }); - describe("addTraceContexToXray", () => { + describe("addTraceContextToXray", () => { beforeEach(() => { StepFunctionContextService["_instance"] = undefined as any; sentSegment = undefined; @@ -955,7 +1035,7 @@ describe("TraceContextExtractor", () => { process.env["AWS_XRAY_DAEMON_ADDRESS"] = undefined; }); - it("adds StepFunction context when present over metadata", () => { + it("adds legacy StepFunction context when present over metadata", () => { jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); process.env["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; @@ -1006,7 +1086,7 @@ describe("TraceContextExtractor", () => { const sentMessage = sentSegment.toString(); expect(sentMessage).toEqual( - '{"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}}}}', + '{"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"}}}}', ); }); diff --git a/src/trace/step-function-service.spec.ts b/src/trace/step-function-service.spec.ts index c8a888d0..391fb7d7 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -1,7 +1,7 @@ import { PARENT_ID, StepFunctionContextService } from "./step-function-service"; describe("StepFunctionContextService", () => { - const stepFunctionEvent = { + const legacyStepFunctionEvent = { Execution: { Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", Input: { @@ -21,6 +21,56 @@ describe("StepFunctionContextService", () => { Name: "my-state-machine", }, } as const; + const lambdaRootStepFunctionEvent = { + _datadog: { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + "x-datadog-trace-id": "10593586103637578129", + "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000", + "serverless-version": "v1", + }, + } as const; + const nestedStepFunctionEvent = { + _datadog: { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + RootExecutionId: + "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc", + "serverless-version": "v1", + }, + } as const; describe("instance", () => { it("returns the same instance every time", () => { const instance1 = StepFunctionContextService.instance(); @@ -40,75 +90,35 @@ describe("StepFunctionContextService", () => { ["event is not an object", "event"], ["event is missing Execution property", {}], [ - "Execution is missing Name field", + "Execution is not defined", { - ...stepFunctionEvent, - Execution: {}, + ...legacyStepFunctionEvent, + Execution: undefined, }, ], [ "Execution Id is not a string", { - ...stepFunctionEvent, + ...legacyStepFunctionEvent, Execution: { - ...stepFunctionEvent.Execution, + ...legacyStepFunctionEvent.Execution, Id: 1, }, }, ], - [ - "Execution Name isn't a string", - { - ...stepFunctionEvent, - Execution: { - ...stepFunctionEvent.Execution, - Name: 12345, - }, - }, - ], - [ - "Execution RoleArn isn't a string", - { - ...stepFunctionEvent, - Execution: { - ...stepFunctionEvent.Execution, - RoleArn: 12345, - }, - }, - ], - [ - "Execution StartTime isn't a string", - { - ...stepFunctionEvent, - Execution: { - ...stepFunctionEvent.Execution, - StartTime: 12345, - }, - }, - ], [ "State is not defined", { - ...stepFunctionEvent, + ...legacyStepFunctionEvent, State: undefined, }, ], - [ - "State RetryCount is not a number", - { - ...stepFunctionEvent, - State: { - ...stepFunctionEvent.State, - RetryCount: "1", - }, - }, - ], [ "State EnteredTime is not a string", { - ...stepFunctionEvent, + ...legacyStepFunctionEvent, State: { - ...stepFunctionEvent.State, + ...legacyStepFunctionEvent.State, EnteredTime: 12345, }, }, @@ -116,36 +126,9 @@ describe("StepFunctionContextService", () => { [ "State Name is not a string", { - ...stepFunctionEvent, + ...legacyStepFunctionEvent, State: { - ...stepFunctionEvent, - Name: 1, - }, - }, - ], - [ - "StateMachine is undefined", - { - ...stepFunctionEvent, - StateMachine: undefined, - }, - ], - [ - "StateMachine Id is not a string", - { - ...stepFunctionEvent, - StateMachine: { - ...stepFunctionEvent.StateMachine, - Id: 1, - }, - }, - ], - [ - "StateMachine Name is not a string", - { - ...stepFunctionEvent, - StateMachine: { - ...stepFunctionEvent.StateMachine, + ...legacyStepFunctionEvent, Name: 1, }, }, @@ -156,26 +139,45 @@ describe("StepFunctionContextService", () => { expect(instance.context).toBeUndefined(); }); - it("sets context from valid event", () => { + it("sets context from valid legacy event", () => { const instance = StepFunctionContextService.instance(); // Force setting event - instance["setContext"](stepFunctionEvent); + instance["setContext"](legacyStepFunctionEvent); expect(instance.context).toEqual({ - "step_function.execution_id": + 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_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be", - "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, + state_entered_time: "2022-12-08T21:08:19.224Z", + state_name: "step-one", + }); + }); + + it("sets context from valid nested event", () => { + const instance = StepFunctionContextService.instance(); + // Force setting event + instance["setContext"](nestedStepFunctionEvent); + expect(instance.context).toEqual({ + 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", + root_execution_id: + "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc", + serverless_version: "v1", + }); + }); + + it("sets context from valid Lambda root event", () => { + const instance = StepFunctionContextService.instance(); + // Force setting event + instance["setContext"](lambdaRootStepFunctionEvent); + expect(instance.context).toEqual({ + 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", + trace_id: "10593586103637578129", + dd_p_tid: "6734e7c300000000", + serverless_version: "v1", }); }); }); @@ -185,10 +187,10 @@ describe("StepFunctionContextService", () => { jest.resetModules(); StepFunctionContextService["_instance"] = undefined as any; }); - it("returns a SpanContextWrapper when event is valid", () => { + it("returns a SpanContextWrapper when legacy event is valid", () => { const instance = StepFunctionContextService.instance(); // Force setting event - instance["setContext"](stepFunctionEvent); + instance["setContext"](legacyStepFunctionEvent); const spanContext = instance.spanContext; @@ -200,6 +202,36 @@ describe("StepFunctionContextService", () => { expect(spanContext?.source).toBe("event"); }); + it("returns a SpanContextWrapper when nested event is valid", () => { + const instance = StepFunctionContextService.instance(); + // Force setting event + instance["setContext"](nestedStepFunctionEvent); + + const spanContext = instance.spanContext; + + expect(spanContext).not.toBeNull(); + + expect(spanContext?.toTraceId()).toBe("8676990472248253142"); + expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.sampleMode()).toBe("1"); + expect(spanContext?.source).toBe("event"); + }); + + it("returns a SpanContextWrapper when Lambda root event is valid", () => { + const instance = StepFunctionContextService.instance(); + // Force setting event + instance["setContext"](lambdaRootStepFunctionEvent); + + const spanContext = instance.spanContext; + + expect(spanContext).not.toBeNull(); + + expect(spanContext?.toTraceId()).toBe("10593586103637578129"); + expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.sampleMode()).toBe("1"); + expect(spanContext?.source).toBe("event"); + }); + it("returns null when context is not set", () => { const instance = StepFunctionContextService.instance(); // Force setting event @@ -213,7 +245,7 @@ describe("StepFunctionContextService", () => { it("returns a SpanContextWrapper when event is from legacy lambda", () => { const instance = StepFunctionContextService.instance(); // Force setting event - instance["setContext"]({ Payload: stepFunctionEvent }); + instance["setContext"]({ Payload: legacyStepFunctionEvent }); const spanContext = instance.spanContext; diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index 5f2e3d9f..7aa63626 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -3,23 +3,57 @@ import { SampleMode, TraceSource } from "./trace-context-service"; import { SpanContextWrapper } from "./span-context-wrapper"; import { Sha256 } from "@aws-crypto/sha256-js"; -export interface StepFunctionContext { - "step_function.execution_name": string; - "step_function.execution_id": string; - "step_function.execution_input": object; - "step_function.execution_role_arn": string; - "step_function.execution_start_time": string; - "step_function.state_machine_name": string; - "step_function.state_machine_arn": string; - "step_function.state_entered_time": string; - "step_function.state_name": string; - "step_function.state_retry_count": number; +interface NestedStepFunctionContext { + execution_id: string; + state_entered_time: string; + state_name: string; + root_execution_id: string; + serverless_version: string; } +interface LambdaRootStepFunctionContext { + execution_id: string; + state_entered_time: string; + state_name: string; + trace_id: string; + dd_p_tid: string; + serverless_version: string; +} + +interface LegacyStepFunctionContext { + execution_id: string; + state_entered_time: string; + state_name: string; +} + +export type StepFunctionContext = NestedStepFunctionContext | LambdaRootStepFunctionContext | LegacyStepFunctionContext; + export const TRACE_ID = "traceId"; export const PARENT_ID = "spanId"; export const DD_P_TID = "_dd.p.tid"; +// Type Guard Functions +function isStepFunctionRootContext(obj: any): obj is NestedStepFunctionContext { + return typeof obj?.root_execution_id === "string" && typeof obj?.serverless_version === "string"; +} + +function isLambdaRootContext(obj: any): obj is LambdaRootStepFunctionContext { + return ( + typeof obj?.trace_id === "string" && + typeof obj?.dd_p_tid === "string" && + typeof obj?.serverless_version === "string" + ); +} + +function isLegacyContext(obj: any): obj is LegacyStepFunctionContext { + return ( + typeof obj?.execution_id === "string" && + typeof obj?.state_entered_time === "string" && + typeof obj?.state_name === "string" && + obj?.serverless_version === undefined + ); +} + export class StepFunctionContextService { private static _instance: StepFunctionContextService; public context?: StepFunctionContext; @@ -41,104 +75,67 @@ export class StepFunctionContextService { // always triggered by the same event. if (typeof event !== "object") return; - // Legacy lambda parsing + // Extract Payload if available (Legacy lambda parsing) if (typeof event.Payload === "object") { event = event.Payload; } - // Execution - const execution = event.Execution; - if (typeof execution !== "object") { - logDebug("event.Execution is not an object."); - return; - } - const executionID = execution.Id; - if (typeof executionID !== "string") { - logDebug("event.Execution.Id is not a string."); - return; - } - const executionInput = execution.Input; - const executionName = execution.Name; - if (typeof executionName !== "string") { - logDebug("event.Execution.Name is not a string."); - return; - } - const executionRoleArn = execution.RoleArn; - if (typeof executionRoleArn !== "string") { - logDebug("event.Execution.RoleArn is not a string."); - return; - } - const executionStartTime = execution.StartTime; - if (typeof executionStartTime !== "string") { - logDebug("event.Execution.StartTime is not a string."); - return; - } - - // State - const state = event.State; - if (typeof state !== "object") { - logDebug("event.State is not an object."); - return; - } - const stateRetryCount = state.RetryCount; - if (typeof stateRetryCount !== "number") { - logDebug("event.State.RetryCount is not a number."); - return; - } - const stateEnteredTime = state.EnteredTime; - if (typeof stateEnteredTime !== "string") { - logDebug("event.State.EnteredTime is not a string."); - return; - } - const stateName = state.Name; - if (typeof stateName !== "string") { - logDebug("event.State.Name is not a string."); - return; - } - - // StateMachine - const stateMachine = event.StateMachine; - if (typeof stateMachine !== "object") { - logDebug("event.StateMachine is not an object."); - return; - } - const stateMachineArn = stateMachine.Id; - if (typeof stateMachineArn !== "string") { - logDebug("event.StateMachine.Id is not a string."); - return; - } - const stateMachineName = stateMachine.Name; - if (typeof stateMachineName !== "string") { - logDebug("event.StateMachine.Name is not a string."); - return; + // Extract _datadog if available (JSONata v1 parsing) + if (typeof event._datadog === "object") { + event = event._datadog; + } + + // Extract the common context variables + const stateMachineContext = this.extractStateMachineContext(event); + if (stateMachineContext === null) return; + const { execution_id, state_entered_time, state_name } = stateMachineContext; + + if (typeof event["serverless-version"] === "string" && event["serverless-version"] === "v1") { + if (typeof event.RootExecutionId === "string") { + this.context = { + execution_id, + state_entered_time, + state_name, + root_execution_id: event.RootExecutionId, + serverless_version: event["serverless-version"], + } as NestedStepFunctionContext; + } else if (typeof event["x-datadog-trace-id"] === "string" && typeof event["x-datadog-tags"] === "string") { + this.context = { + execution_id, + state_entered_time, + state_name, + trace_id: event["x-datadog-trace-id"], + dd_p_tid: this.parsePTid(event["x-datadog-tags"]), + serverless_version: event["serverless-version"], + } as LambdaRootStepFunctionContext; + } + } else { + this.context = { execution_id, state_entered_time, state_name } as LegacyStepFunctionContext; } - - const context = { - "step_function.execution_name": executionName, - "step_function.execution_id": executionID, - "step_function.execution_input": executionInput ?? {}, - "step_function.execution_role_arn": executionRoleArn, - "step_function.execution_start_time": executionStartTime, - "step_function.state_entered_time": stateEnteredTime, - "step_function.state_machine_arn": stateMachineArn, - "step_function.state_machine_name": stateMachineName, - "step_function.state_name": stateName, - "step_function.state_retry_count": stateRetryCount, - }; - - this.context = context; } public get spanContext(): SpanContextWrapper | null { if (this.context === undefined) return null; - const traceId = this.deterministicSha256HashToBigIntString(this.context["step_function.execution_id"], TRACE_ID); + let traceId: string; + let ptid: string; + + if (isStepFunctionRootContext(this.context)) { + traceId = this.deterministicSha256HashToBigIntString(this.context.root_execution_id, TRACE_ID); + ptid = this.deterministicSha256HashToBigIntString(this.context.root_execution_id, DD_P_TID); + } else if (isLambdaRootContext(this.context)) { + traceId = this.context.trace_id; + ptid = this.context.dd_p_tid; + } else if (isLegacyContext(this.context)) { + traceId = this.deterministicSha256HashToBigIntString(this.context.execution_id, TRACE_ID); + ptid = this.deterministicSha256HashToBigIntString(this.context.execution_id, DD_P_TID); + } else { + logDebug("StepFunctionContext doesn't match any known formats!"); + return null; + } + const parentId = this.deterministicSha256HashToBigIntString( - this.context["step_function.execution_id"] + - "#" + - this.context["step_function.state_name"] + - "#" + - this.context["step_function.state_entered_time"], + this.context.execution_id + "#" + this.context.state_name + "#" + this.context.state_entered_time, PARENT_ID, ); const sampleMode = SampleMode.AUTO_KEEP; @@ -154,7 +151,6 @@ export class StepFunctionContextService { sampling: { priority: sampleMode.toString(2) }, }); - const ptid = this.deterministicSha256HashToBigIntString(this.context["step_function.execution_id"], DD_P_TID); ddSpanContext._trace.tags["_dd.p.tid"] = id(ptid, 10).toString(16); if (ddSpanContext === null) return null; @@ -175,7 +171,7 @@ export class StepFunctionContextService { } private deterministicSha256Hash(s: string, type: string): string { - // returns 128 bits hash unless mostSignificant64Bits options is set to true. + // returns upper or lower 64 bits of the hash const hash = new Sha256(); hash.update(s); @@ -197,4 +193,45 @@ export class StepFunctionContextService { private numberToBinaryString(num: number): string { return num.toString(2).padStart(8, "0"); } + + private extractStateMachineContext(event: any): { + execution_id: string; + state_entered_time: string; + state_name: string; + } | null { + const execution = event.Execution; + const state = event.State; + + if ( + typeof execution === "object" && + typeof execution.Id === "string" && + typeof state === "object" && + typeof state.EnteredTime === "string" && + typeof state.Name === "string" + ) { + return { + execution_id: execution.Id, + state_entered_time: state.EnteredTime, + state_name: state.Name, + }; + } + + logDebug("Cannot extract StateMachine context! Invalid execution or state data."); + return null; + } + + /** + * Parse a list of trace tags such as [_dd.p.tid=66bcb5eb00000000,_dd.p.dm=-0] and return the + * value of the _dd.p.tid tag or an empty string if not found. + */ + private parsePTid(traceTags: string): string { + if (traceTags) { + for (const tag of traceTags.split(",")) { + if (tag.includes("_dd.p.tid=")) { + return tag.split("=")[1]; + } + } + } + return ""; + } } diff --git a/src/trace/trigger.spec.ts b/src/trace/trigger.spec.ts index 29b64892..e8a47758 100644 --- a/src/trace/trigger.spec.ts +++ b/src/trace/trigger.spec.ts @@ -109,6 +109,14 @@ describe("parseEventSource", () => { }, file: "sqs.json", }, + { + result: { + "function_trigger.event_source": "states", + "function_trigger.event_source_arn": + "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d", + }, + file: "states.json", + }, ]; const bufferedResponses = [ diff --git a/src/trace/trigger.ts b/src/trace/trigger.ts index e23e0fb6..a820a02e 100644 --- a/src/trace/trigger.ts +++ b/src/trace/trigger.ts @@ -92,6 +92,18 @@ function extractEventBridgeARN(event: EventBridgeEvent) { return event.source; } +function extractStateMachineARN(event: any) { + // Extract Payload if available (Legacy lambda parsing) + if (typeof event.Payload === "object") { + event = event.Payload; + } + // Extract _datadog if available (JSONata v1 parsing) + if (typeof event._datadog === "object") { + event = event._datadog; + } + return event.StateMachine.Id; +} + export enum eventTypes { apiGateway = "api-gateway", applicationLoadBalancer = "application-load-balancer", @@ -106,6 +118,7 @@ export enum eventTypes { s3 = "s3", sns = "sns", sqs = "sqs", + stepFunctions = "states", } export enum eventSubTypes { @@ -134,7 +147,7 @@ export function parseEventSourceSubType(event: any): eventSubTypes { * parseEventSource parses the triggering event to determine the source * Possible Returns: * api-gateway | application-load-balancer | cloudwatch-logs | - * cloudwatch-events | cloudfront | dynamodb | kinesis | s3 | sns | sqs + * cloudwatch-events | cloudfront | dynamodb | kinesis | s3 | sns | sqs | states */ export function parseEventSource(event: any) { if (eventType.isLambdaUrlEvent(event)) { @@ -186,6 +199,10 @@ export function parseEventSource(event: any) { if (eventType.isEventBridgeEvent(event)) { return eventTypes.eventBridge; } + + if (eventType.isStepFunctionsEvent(event)) { + return eventTypes.stepFunctions; + } } /** @@ -256,6 +273,10 @@ export function parseEventSourceARN(source: string | undefined, event: any, cont eventSourceARN = extractEventBridgeARN(event); } + if (source === "states") { + eventSourceARN = extractStateMachineARN(event); + } + return eventSourceARN; } diff --git a/src/utils/event-type-guards.ts b/src/utils/event-type-guards.ts index 3c99e6ec..b26bde02 100644 --- a/src/utils/event-type-guards.ts +++ b/src/utils/event-type-guards.ts @@ -106,3 +106,17 @@ export function isEventBridgeEvent(event: any): event is EventBridgeEvent