diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index ee5860f8..49729efd 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -669,6 +669,7 @@ describe("TraceContextExtractor", () => { Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -883,6 +884,7 @@ describe("TraceContextExtractor", () => { }, Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -919,6 +921,7 @@ describe("TraceContextExtractor", () => { Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -959,6 +962,7 @@ describe("TraceContextExtractor", () => { Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -999,6 +1003,7 @@ describe("TraceContextExtractor", () => { Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -1049,6 +1054,7 @@ describe("TraceContextExtractor", () => { }, Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -1086,7 +1092,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":{"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"}}}}', + '{"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","redrive_count":"0","state_entered_time":"2022-12-08T21:08:19.224Z","state_name":"step-one"}}}}', ); }); diff --git a/src/trace/context/extractors/step-function.spec.ts b/src/trace/context/extractors/step-function.spec.ts index 37847cf4..723bab67 100644 --- a/src/trace/context/extractors/step-function.spec.ts +++ b/src/trace/context/extractors/step-function.spec.ts @@ -8,22 +8,41 @@ describe("StepFunctionEventTraceExtractor", () => { describe("extract", () => { const payload = { 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", + Id: "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", + Name: "72a7ca3e-901c-41bb-b5a3-5f279b92a316", + RoleArn: + "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", + StartTime: "2024-12-04T19:38:04.069Z", + RedriveCount: 0, }, State: { - Name: "step-one", - EnteredTime: "2022-12-08T21:08:19.224Z", - RetryCount: 2, + Name: "Lambda Invoke", + EnteredTime: "2024-12-04T19:38:04.118Z", + RetryCount: 0, }, StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine", + Name: "abhinav-activity-state-machine", + }, + }; + + const redrivePayload = { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316", + Name: "72a7ca3e-901c-41bb-b5a3-5f279b92a316", + RoleArn: + "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j", + StartTime: "2024-12-04T19:38:04.069Z", + RedriveCount: 1, + }, + State: { + Name: "Lambda Invoke", + EnteredTime: "2024-12-04T19:38:04.118Z", + RetryCount: 0, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine", + Name: "abhinav-activity-state-machine", }, }; it("extracts trace context with valid payload", () => { @@ -36,8 +55,25 @@ describe("StepFunctionEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).not.toBeNull(); - expect(traceContext?.toTraceId()).toBe("1139193989631387307"); - expect(traceContext?.toSpanId()).toBe("5892738536804826142"); + expect(traceContext?.toTraceId()).toBe("435175499815315247"); + expect(traceContext?.toSpanId()).toBe("3929055471293792800"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); + + // https://github.com/DataDog/logs-backend/blob/c17618cb552fc369ca40282bae0a65803f82f694/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L46 + it("extracts trace context with valid redriven payload", () => { + // Mimick TraceContextService.extract initialization + StepFunctionContextService.instance(redrivePayload); + + const extractor = new StepFunctionEventTraceExtractor(); + + // Payload is sent again for safety in case the instance wasn't previously initialized + const traceContext = extractor.extract(redrivePayload); + expect(traceContext).not.toBeNull(); + + expect(traceContext?.toTraceId()).toBe("435175499815315247"); + expect(traceContext?.toSpanId()).toBe("5063839446130725204"); expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); @@ -49,8 +85,8 @@ describe("StepFunctionEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).not.toBeNull(); - expect(traceContext?.toTraceId()).toBe("1139193989631387307"); - expect(traceContext?.toSpanId()).toBe("5892738536804826142"); + expect(traceContext?.toTraceId()).toBe("435175499815315247"); + expect(traceContext?.toSpanId()).toBe("3929055471293792800"); expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); @@ -65,8 +101,8 @@ describe("StepFunctionEventTraceExtractor", () => { const traceContext = extractor.extract({ Payload: payload }); expect(traceContext).not.toBeNull(); - expect(traceContext?.toTraceId()).toBe("1139193989631387307"); - expect(traceContext?.toSpanId()).toBe("5892738536804826142"); + expect(traceContext?.toTraceId()).toBe("435175499815315247"); + expect(traceContext?.toSpanId()).toBe("3929055471293792800"); expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); diff --git a/src/trace/step-function-service.spec.ts b/src/trace/step-function-service.spec.ts index 391fb7d7..faf9b081 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -9,6 +9,7 @@ describe("StepFunctionContextService", () => { }, Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -30,6 +31,7 @@ describe("StepFunctionContextService", () => { }, Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -55,6 +57,7 @@ describe("StepFunctionContextService", () => { }, Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + RedriveCount: 0, StartTime: "2022-12-08T21:08:17.924Z", }, State: { @@ -106,6 +109,16 @@ describe("StepFunctionContextService", () => { }, }, ], + [ + "Execution RedriveCount is not a number", + { + ...legacyStepFunctionEvent, + Execution: { + ...legacyStepFunctionEvent.Execution, + RedriveCount: "0", + }, + }, + ], [ "State is not defined", { @@ -146,6 +159,7 @@ describe("StepFunctionContextService", () => { 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", + redrive_count: "0", state_entered_time: "2022-12-08T21:08:19.224Z", state_name: "step-one", }); @@ -158,6 +172,7 @@ describe("StepFunctionContextService", () => { 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", + redrive_count: "0", state_entered_time: "2022-12-08T21:08:19.224Z", state_name: "step-one", root_execution_id: @@ -173,6 +188,7 @@ describe("StepFunctionContextService", () => { 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", + redrive_count: "0", state_entered_time: "2022-12-08T21:08:19.224Z", state_name: "step-one", trace_id: "10593586103637578129", diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index 92e6725f..be987b67 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -5,6 +5,7 @@ import { Sha256 } from "@aws-crypto/sha256-js"; interface NestedStepFunctionContext { execution_id: string; + redrive_count: string; state_entered_time: string; state_name: string; root_execution_id: string; @@ -13,6 +14,7 @@ interface NestedStepFunctionContext { interface LambdaRootStepFunctionContext { execution_id: string; + redrive_count: string; state_entered_time: string; state_name: string; trace_id: string; @@ -22,6 +24,7 @@ interface LambdaRootStepFunctionContext { interface LegacyStepFunctionContext { execution_id: string; + redrive_count: string; state_entered_time: string; state_name: string; } @@ -88,12 +91,13 @@ export class StepFunctionContextService { // Extract the common context variables const stateMachineContext = this.extractStateMachineContext(event); if (stateMachineContext === null) return; - const { execution_id, state_entered_time, state_name } = stateMachineContext; + const { execution_id, redrive_count, 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, + redrive_count, state_entered_time, state_name, root_execution_id: event.RootExecutionId, @@ -102,6 +106,7 @@ export class StepFunctionContextService { } else if (typeof event["x-datadog-trace-id"] === "string" && typeof event["x-datadog-tags"] === "string") { this.context = { execution_id, + redrive_count, state_entered_time, state_name, trace_id: event["x-datadog-trace-id"], @@ -110,7 +115,7 @@ export class StepFunctionContextService { } as LambdaRootStepFunctionContext; } } else { - this.context = { execution_id, state_entered_time, state_name } as LegacyStepFunctionContext; + this.context = { execution_id, redrive_count, state_entered_time, state_name } as LegacyStepFunctionContext; } } @@ -134,8 +139,15 @@ export class StepFunctionContextService { return null; } + const redrivePostfix = this.context.redrive_count === "0" ? "" : `#${this.context.redrive_count}`; + const parentId = this.deterministicSha256HashToBigIntString( - this.context.execution_id + "#" + this.context.state_name + "#" + this.context.state_entered_time, + this.context.execution_id + + "#" + + this.context.state_name + + "#" + + this.context.state_entered_time + + redrivePostfix, PARENT_ID, ); const sampleMode = SampleMode.AUTO_KEEP; @@ -196,12 +208,14 @@ export class StepFunctionContextService { private extractStateMachineContext(event: any): { execution_id: string; + redrive_count: string; state_entered_time: string; state_name: string; } | null { if (this.isValidContextObject(event)) { return { execution_id: event.Execution.Id, + redrive_count: event.Execution.RedriveCount.toString(), state_entered_time: event.State.EnteredTime, state_name: event.State.Name, }; @@ -214,6 +228,7 @@ export class StepFunctionContextService { private isValidContextObject(context: any): boolean { return ( typeof context?.Execution?.Id === "string" && + typeof context?.Execution?.RedriveCount === "number" && typeof context?.State?.EnteredTime === "string" && typeof context?.State?.Name === "string" );