From faeb8d3b1f5c7572e5ec36756d846d9814d81df4 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 15 Nov 2024 15:04:17 -0500 Subject: [PATCH 1/4] initial impl for setting and using the new contexts --- src/trace/step-function-service.ts | 229 ++++++++++++++++------------- 1 file changed, 129 insertions(+), 100 deletions(-) diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index 5f2e3d9f..018a3021 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 StepFunctionRootContext { + execution_id: string; + state_entered_time: string; + state_name: string; + root_execution_id: string; + serverless_version: string; } +interface LambdaRootContext { + execution_id: string; + state_entered_time: string; + state_name: string; + trace_id: string; + dd_p_tid: string; + serverless_version: string; +} + +interface LegacyContext { + execution_id: string; + state_entered_time: string; + state_name: string; +} + +export type StepFunctionContext = StepFunctionRootContext | LambdaRootContext | LegacyContext; + 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 StepFunctionRootContext { + return typeof obj?.root_execution_id === "string" && typeof obj?.serverless_version === "string"; +} + +function isLambdaRootContext(obj: any): obj is LambdaRootContext { + return ( + typeof obj?.trace_id === "string" && + typeof obj?.dd_p_tid === "string" && + typeof obj?.serverless_version === "string" + ); +} + +function isLegacyContext(obj: any): obj is LegacyContext { + 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,74 @@ 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 (event.serverless_version === "string" && event.serverless_version == "v1") { + const serverless_version = event.serverless_version; + + if (event.RootExecutionId === "string") { + const root_execution_id = event.RootExecutionId; + + this.context = { + execution_id, + state_entered_time, + state_name, + root_execution_id, + serverless_version, + } as StepFunctionRootContext; + } else if (event.trace_id === "string" && event.dd_p_tid === "string") { + const trace_id = event.trace_id; + const dd_p_tid = event.dd_p_tid; + + this.context = { + execution_id, + state_entered_time, + state_name, + trace_id, + dd_p_tid, + serverless_version, + } as LambdaRootContext; + } + } else { + this.context = { execution_id, state_entered_time, state_name } as LegacyContext; } - - 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 +158,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 +178,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 +200,30 @@ 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; + } } From 1f0e7621dc8a3c61e3794821cbd6ae8e857ef8f9 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Fri, 15 Nov 2024 15:49:31 -0500 Subject: [PATCH 2/4] fix current tests and lint --- src/trace/context/extractor.spec.ts | 6 +- src/trace/step-function-service.spec.ts | 116 ++++-------------------- src/trace/step-function-service.ts | 17 ++-- 3 files changed, 29 insertions(+), 110 deletions(-) diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index a4dcd966..fde44e79 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -946,7 +946,7 @@ describe("TraceContextExtractor", () => { }); }); - describe("addTraceContexToXray", () => { + describe("addTraceContextToXray", () => { beforeEach(() => { StepFunctionContextService["_instance"] = undefined as any; sentSegment = undefined; @@ -955,7 +955,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 +1006,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..93b9e467 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: { @@ -40,75 +40,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 +76,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 +89,15 @@ 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", }); }); }); @@ -185,10 +107,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; @@ -213,7 +135,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 018a3021..e49948d5 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -90,29 +90,26 @@ export class StepFunctionContextService { if (stateMachineContext === null) return; const { execution_id, state_entered_time, state_name } = stateMachineContext; - if (event.serverless_version === "string" && event.serverless_version == "v1") { + if (event.serverless_version === "string" && event.serverless_version === "v1") { const serverless_version = event.serverless_version; if (event.RootExecutionId === "string") { - const root_execution_id = event.RootExecutionId; - this.context = { execution_id, state_entered_time, state_name, - root_execution_id, + root_execution_id: event.RootExecutionId, serverless_version, } as StepFunctionRootContext; } else if (event.trace_id === "string" && event.dd_p_tid === "string") { - const trace_id = event.trace_id; - const dd_p_tid = event.dd_p_tid; + const ptid = event["x-datadog-tags"]; // todo: parse me properly this.context = { execution_id, state_entered_time, state_name, - trace_id, - dd_p_tid, + trace_id: event["x-datadog-trace-id"], + dd_p_tid: ptid, serverless_version, } as LambdaRootContext; } @@ -131,8 +128,8 @@ export class StepFunctionContextService { 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"]; + 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); From 185ecb922cc81ac5b20b06430282178ab051ff5c Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Sun, 17 Nov 2024 13:02:25 -0500 Subject: [PATCH 3/4] fix JSON parsing in StepFunctionContext creation --- src/trace/context/extractor.spec.ts | 88 ++++++++++++++++++- src/trace/step-function-service.spec.ts | 110 ++++++++++++++++++++++++ src/trace/step-function-service.ts | 51 ++++++----- 3 files changed, 225 insertions(+), 24 deletions(-) diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index fde44e79..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); }); diff --git a/src/trace/step-function-service.spec.ts b/src/trace/step-function-service.spec.ts index 93b9e467..391fb7d7 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -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(); @@ -100,6 +150,36 @@ describe("StepFunctionContextService", () => { 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", + }); + }); }); describe("spanContext", () => { @@ -122,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 diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index e49948d5..7aa63626 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -3,7 +3,7 @@ import { SampleMode, TraceSource } from "./trace-context-service"; import { SpanContextWrapper } from "./span-context-wrapper"; import { Sha256 } from "@aws-crypto/sha256-js"; -interface StepFunctionRootContext { +interface NestedStepFunctionContext { execution_id: string; state_entered_time: string; state_name: string; @@ -11,7 +11,7 @@ interface StepFunctionRootContext { serverless_version: string; } -interface LambdaRootContext { +interface LambdaRootStepFunctionContext { execution_id: string; state_entered_time: string; state_name: string; @@ -20,24 +20,24 @@ interface LambdaRootContext { serverless_version: string; } -interface LegacyContext { +interface LegacyStepFunctionContext { execution_id: string; state_entered_time: string; state_name: string; } -export type StepFunctionContext = StepFunctionRootContext | LambdaRootContext | LegacyContext; +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 StepFunctionRootContext { +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 LambdaRootContext { +function isLambdaRootContext(obj: any): obj is LambdaRootStepFunctionContext { return ( typeof obj?.trace_id === "string" && typeof obj?.dd_p_tid === "string" && @@ -45,7 +45,7 @@ function isLambdaRootContext(obj: any): obj is LambdaRootContext { ); } -function isLegacyContext(obj: any): obj is LegacyContext { +function isLegacyContext(obj: any): obj is LegacyStepFunctionContext { return ( typeof obj?.execution_id === "string" && typeof obj?.state_entered_time === "string" && @@ -90,31 +90,27 @@ export class StepFunctionContextService { if (stateMachineContext === null) return; const { execution_id, state_entered_time, state_name } = stateMachineContext; - if (event.serverless_version === "string" && event.serverless_version === "v1") { - const serverless_version = event.serverless_version; - - if (event.RootExecutionId === "string") { + 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, - } as StepFunctionRootContext; - } else if (event.trace_id === "string" && event.dd_p_tid === "string") { - const ptid = event["x-datadog-tags"]; // todo: parse me properly - + 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: ptid, - serverless_version, - } as LambdaRootContext; + 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 LegacyContext; + this.context = { execution_id, state_entered_time, state_name } as LegacyStepFunctionContext; } } @@ -223,4 +219,19 @@ export class StepFunctionContextService { 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 ""; + } } From d7c70ebe4a309de5164c7caf5221845be2ffad27 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Sun, 17 Nov 2024 13:18:01 -0500 Subject: [PATCH 4/4] add state machine case to trigger tag parsing --- event_samples/states.json | 22 ++++++++++++++++++++++ src/trace/trigger.spec.ts | 8 ++++++++ src/trace/trigger.ts | 23 ++++++++++++++++++++++- src/utils/event-type-guards.ts | 14 ++++++++++++++ 4 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 event_samples/states.json 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/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