Update Step Functions Parent ID Generation (#622)

Updates the parent ID generation using upstream Step Functions context object to be able to handle retries

Old behavior would result in the same parent ID for all retries, this new behavior will use the `retry_count` when it's non-zero. This way we don't break backwards compatibility with old layers that don't use this value in the common case of `redrive_count:0`

Author: avedmala

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");
       });
@@ -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"}}}}',
       );
     });
 

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("8782364156266188026");
       expect(traceContext?.sampleMode()).toBe("1");
       expect(traceContext?.source).toBe("event");
     });

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",
@@ -203,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");
     });
@@ -218,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");
     });
@@ -233,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");
     });
@@ -258,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");
     });

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,13 +219,15 @@ export class StepFunctionContextService {
   private extractStateMachineContext(event: any): {
     execution_id: string;
     redrive_count: string;
+    retry_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 ?? "0").toString(),
+        retry_count: (event.State.RetryCount ?? "0").toString(),
         state_entered_time: event.State.EnteredTime,
         state_name: event.State.Name,
       };