From 62855c2ffe1db60406fc72f9091ea9a9233b4a76 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 18 Feb 2025 14:37:56 -0500 Subject: [PATCH 1/3] Update Step Functions Parent ID Generation --- src/trace/context/extractor.spec.ts | 2 +- .../context/extractors/step-function.spec.ts | 2 +- src/trace/step-function-service.spec.ts | 3 ++ src/trace/step-function-service.ts | 30 +++++++++++++------ 4 files changed, 26 insertions(+), 11 deletions(-) diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index 49729efd..e1e1c77f 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -1092,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","redrive_count":"0","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","retry_count":"2","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 609b5ffa..3447a324 100644 --- a/src/trace/context/extractors/step-function.spec.ts +++ b/src/trace/context/extractors/step-function.spec.ts @@ -72,7 +72,7 @@ describe("StepFunctionEventTraceExtractor", () => { expect(traceContext).not.toBeNull(); expect(traceContext?.toTraceId()).toBe("435175499815315247"); - expect(traceContext?.toSpanId()).toBe("5063839446130725204"); + 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 1eb6a710..6594419c 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -150,6 +150,7 @@ describe("StepFunctionContextService", () => { 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", + retry_count: "2", state_entered_time: "2022-12-08T21:08:19.224Z", state_name: "step-one", }); @@ -163,6 +164,7 @@ describe("StepFunctionContextService", () => { 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", + retry_count: "2", state_entered_time: "2022-12-08T21:08:19.224Z", state_name: "step-one", root_execution_id: @@ -179,6 +181,7 @@ describe("StepFunctionContextService", () => { 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", + retry_count: "2", 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 29510bc9..4ef926dc 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -6,6 +6,7 @@ import { Sha256 } from "@aws-crypto/sha256-js"; interface NestedStepFunctionContext { execution_id: string; redrive_count: string; + retry_count: string; state_entered_time: string; state_name: string; root_execution_id: string; @@ -15,6 +16,7 @@ interface NestedStepFunctionContext { interface LambdaRootStepFunctionContext { execution_id: string; redrive_count: string; + retry_count: string; state_entered_time: string; state_name: string; trace_id: string; @@ -25,6 +27,7 @@ interface LambdaRootStepFunctionContext { interface LegacyStepFunctionContext { execution_id: string; redrive_count: string; + retry_count: string; state_entered_time: string; state_name: string; } @@ -91,13 +94,14 @@ export class StepFunctionContextService { // Extract the common context variables const stateMachineContext = this.extractStateMachineContext(event); if (stateMachineContext === null) return; - const { execution_id, redrive_count, state_entered_time, state_name } = stateMachineContext; + const { execution_id, redrive_count, retry_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, + retry_count, state_entered_time, state_name, root_execution_id: event.RootExecutionId, @@ -107,6 +111,7 @@ export class StepFunctionContextService { this.context = { execution_id, redrive_count, + retry_count, state_entered_time, state_name, trace_id: event["x-datadog-trace-id"], @@ -115,7 +120,13 @@ export class StepFunctionContextService { } as LambdaRootStepFunctionContext; } } else { - this.context = { execution_id, redrive_count, state_entered_time, state_name } as LegacyStepFunctionContext; + this.context = { + execution_id, + redrive_count, + retry_count, + state_entered_time, + state_name, + } as LegacyStepFunctionContext; } } @@ -139,17 +150,16 @@ export class StepFunctionContextService { return null; } - const redrivePostfix = this.context.redrive_count === "0" ? "" : `#${this.context.redrive_count}`; + const countsSuffix = + this.context.retry_count !== "0" && this.context.redrive_count !== "0" + ? `#${this.context.retry_count}#${this.context.redrive_count}` + : ""; const parentId = this.deterministicSha256HashToBigIntString( - this.context.execution_id + - "#" + - this.context.state_name + - "#" + - this.context.state_entered_time + - redrivePostfix, + `${this.context.execution_id}#${this.context.state_name}#${this.context.state_entered_time}${countsSuffix}`, PARENT_ID, ); + const sampleMode = SampleMode.AUTO_KEEP; try { @@ -209,6 +219,7 @@ export class StepFunctionContextService { private extractStateMachineContext(event: any): { execution_id: string; redrive_count: string; + retry_count: string; state_entered_time: string; state_name: string; } | null { @@ -216,6 +227,7 @@ export class StepFunctionContextService { return { execution_id: event.Execution.Id, redrive_count: (event.Execution.RedriveCount ?? "0").toString(), + retry_count: (event.State.RetryCount ?? "0").toString(), state_entered_time: event.State.EnteredTime, state_name: event.State.Name, }; From 869c875d4021c68e3b71c8766fc971b88af30c67 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 18 Feb 2025 14:41:44 -0500 Subject: [PATCH 2/3] flip AND to OR --- src/trace/step-function-service.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index 4ef926dc..fe5fe548 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -151,7 +151,7 @@ export class StepFunctionContextService { } const countsSuffix = - this.context.retry_count !== "0" && this.context.redrive_count !== "0" + this.context.retry_count !== "0" || this.context.redrive_count !== "0" ? `#${this.context.retry_count}#${this.context.redrive_count}` : ""; From 2ec9ee1d39a3964e131c8ebdeb1e6de5bc4dc179 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 18 Feb 2025 14:45:48 -0500 Subject: [PATCH 3/3] update tests --- src/trace/context/extractor.spec.ts | 2 +- src/trace/context/extractors/step-function.spec.ts | 2 +- src/trace/step-function-service.spec.ts | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index e1e1c77f..cb47b6f6 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -690,7 +690,7 @@ describe("TraceContextExtractor", () => { expect(traceContext).not.toBeNull(); expect(traceContext?.toTraceId()).toBe("1139193989631387307"); - expect(traceContext?.toSpanId()).toBe("5892738536804826142"); + expect(traceContext?.toSpanId()).toBe("7747304477664363642"); expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); diff --git a/src/trace/context/extractors/step-function.spec.ts b/src/trace/context/extractors/step-function.spec.ts index 3447a324..1465ab61 100644 --- a/src/trace/context/extractors/step-function.spec.ts +++ b/src/trace/context/extractors/step-function.spec.ts @@ -72,7 +72,7 @@ describe("StepFunctionEventTraceExtractor", () => { expect(traceContext).not.toBeNull(); expect(traceContext?.toTraceId()).toBe("435175499815315247"); - expect(traceContext?.toSpanId()).toBe("3929055471293792800"); + expect(traceContext?.toSpanId()).toBe("8782364156266188026"); 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 6594419c..a96ed389 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -206,7 +206,7 @@ describe("StepFunctionContextService", () => { expect(spanContext).not.toBeNull(); expect(spanContext?.toTraceId()).toBe("1139193989631387307"); - expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.toSpanId()).toBe("7747304477664363642"); expect(spanContext?.sampleMode()).toBe("1"); expect(spanContext?.source).toBe("event"); }); @@ -221,7 +221,7 @@ describe("StepFunctionContextService", () => { expect(spanContext).not.toBeNull(); expect(spanContext?.toTraceId()).toBe("8676990472248253142"); - expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.toSpanId()).toBe("7747304477664363642"); expect(spanContext?.sampleMode()).toBe("1"); expect(spanContext?.source).toBe("event"); }); @@ -236,7 +236,7 @@ describe("StepFunctionContextService", () => { expect(spanContext).not.toBeNull(); expect(spanContext?.toTraceId()).toBe("10593586103637578129"); - expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.toSpanId()).toBe("7747304477664363642"); expect(spanContext?.sampleMode()).toBe("1"); expect(spanContext?.source).toBe("event"); }); @@ -261,7 +261,7 @@ describe("StepFunctionContextService", () => { expect(spanContext).not.toBeNull(); expect(spanContext?.toTraceId()).toBe("1139193989631387307"); - expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.toSpanId()).toBe("7747304477664363642"); expect(spanContext?.sampleMode()).toBe("1"); expect(spanContext?.source).toBe("event"); });