Propagate Step Function Trace Context through Managed Services

.gitlab/scripts/publish_layers.sh

@@ -87,6 +87,12 @@ if [[ ! ${STAGES[@]} =~ $STAGE ]]; then
 fi
 
 layer="${LAYERS[$index]}"
+if [ -z "$LAYER_NAME_SUFFIX" ]; then
+    echo "No layer name suffix"
+else
+    layer="${layer}-${LAYER_NAME_SUFFIX}"
+fi
+echo "layer name: $layer"
 
 if [[ "$STAGE" =~ ^(staging|sandbox)$ ]]; then
     # Deploy latest version

src/trace/context/extractor-utils.ts

@@ -0,0 +1,64 @@
+import { logDebug } from "../../utils";
+import { StepFunctionContextService } from "../step-function-service";
+import { SpanContextWrapper } from "../span-context-wrapper";
+import { TracerWrapper } from "../tracer-wrapper";
+import { XrayService } from "../xray-service";
+
+/**
+ * Common utility functions for trace context extraction
+ */
+
+/**
+ * Attempts to extract trace context from headers, falling back to Step Function context if needed
+ * @param headers The headers object to extract from
+ * @param tracerWrapper The tracer wrapper instance
+ * @returns SpanContextWrapper or null
+ */
+export function extractTraceContext(headers: any, tracerWrapper: TracerWrapper): SpanContextWrapper | null {
+  // First try to extract as regular trace headers
+  const traceContext = tracerWrapper.extract(headers);
+  if (traceContext) {
+    return traceContext;
+  }
+
+  // If that fails, check if this is a Step Function context
+  const stepFunctionInstance = StepFunctionContextService.instance(headers);
+  const stepFunctionContext = stepFunctionInstance.context;
+
+  if (stepFunctionContext !== undefined) {
+    const spanContext = stepFunctionInstance.spanContext;
+    if (spanContext !== null) {
+      return spanContext;
+    }
+  }
+
+  return null;
+}
+
+/**
+ * Extracts trace context from AWS Trace Header
+ * @param awsTraceHeader The AWS trace header string
+ * @param eventType The type of event (for logging)
+ * @returns SpanContextWrapper or null
+ */
+export function extractFromAWSTraceHeader(awsTraceHeader: string, eventType: string): SpanContextWrapper | null {
+  const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceHeader);
+  if (traceContext) {
+    logDebug(`Extracted trace context from ${eventType} event attributes AWSTraceHeader`);
+    return traceContext;
+  } else {
+    logDebug(`No Datadog trace context found from ${eventType} event attributes AWSTraceHeader`);
+    return null;
+  }
+}
+
+/**
+ * Common error handler for extraction operations
+ * @param error The error that occurred
+ * @param eventType The type of event (for logging)
+ */
+export function handleExtractionError(error: unknown, eventType: string): void {
+  if (error instanceof Error) {
+    logDebug(`Unable to extract trace context from ${eventType} event`, error);
+  }
+}

src/trace/context/extractor.ts

@@ -43,7 +43,6 @@ export class TraceContextExtractor {
   }
 
   async extract(event: any, context: Context): Promise<SpanContextWrapper | null> {
-    this.stepFunctionContextService = StepFunctionContextService.instance(event);
 
     let spanContext: SpanContextWrapper | null = null;
     if (this.config.traceExtractor) {
@@ -57,7 +56,15 @@ export class TraceContextExtractor {
         spanContext = eventExtractor.extract(event);
       }
     }
-
+/*
+    if (spanContext === null) {
+      this.stepFunctionContextService = StepFunctionContextService.instance(event)
+      if (this.stepFunctionContextService?.context) {
+        const extractor = new StepFunctionEventTraceExtractor();
+        spanContext = extractor?.extract(event);
+      }
+    }
+*/
     if (spanContext === null) {
       const contextExtractor = new LambdaContextTraceExtractor(this.tracerWrapper);
       spanContext = contextExtractor.extract(context);
@@ -88,6 +95,7 @@ export class TraceContextExtractor {
     if (EventValidator.isKinesisStreamEvent(event)) return new KinesisEventTraceExtractor(this.tracerWrapper);
     if (EventValidator.isEventBridgeEvent(event)) return new EventBridgeEventTraceExtractor(this.tracerWrapper);
 
+    this.stepFunctionContextService = StepFunctionContextService.instance(event);
     if (this.stepFunctionContextService?.context) return new StepFunctionEventTraceExtractor();
 
     return;

src/trace/context/extractors/event-bridge-sqs.spec.ts

@@ -1,5 +1,6 @@
 import { TracerWrapper } from "../../tracer-wrapper";
 import { EventBridgeSQSEventTraceExtractor } from "./event-bridge-sqs";
+import { StepFunctionContextService } from "../../step-function-service";
 
 let mockSpanContext: any = null;
 
@@ -42,7 +43,7 @@ describe("EventBridgeSQSEventTraceExtractor", () => {
             messageId: "e995e54f-1724-41fa-82c0-8b81821f854e",
             receiptHandle:
               "AQEB4mIfRcyqtzn1X5Ss+ConhTejVGc+qnAcmu3/Z9ZvbNkaPcpuDLX/bzvPD/ZkAXJUXZcemGSJmd7L3snZHKMP2Ck8runZiyl4mubiLb444pZvdiNPuGRJ6a3FvgS/GQPzho/9nNMyOi66m8Viwh70v4EUCPGO4JmD3TTDAUrrcAnqU4WSObjfC/NAp9bI6wH2CEyAYEfex6Nxplbl/jBf9ZUG0I3m3vQd0Q4l4gd4jIR4oxQUglU2Tldl4Kx5fMUAhTRLAENri6HsY81avBkKd9FAuxONlsITB5uj02kOkvLlRGEcalqsKyPJ7AFaDLrOLaL3U+yReroPEJ5R5nwhLOEbeN5HROlZRXeaAwZOIN8BjqdeooYTIOrtvMEVb7a6OPLMdH1XB+ddevtKAH8K9Tm2ZjpaA7dtBGh1zFVHzBk=",
-            body: '{"version":"0","id":"af718b2a-b987-e8c0-7a2b-a188fad2661a","detail-type":"my.Detail","source":"my.Source","account":"425362996713","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"7379586022458917877","x-datadog-parent-id":"2644033662113726488","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}',
+            body: '{"version":"0","id":"af718b2a-b987-e8c0-7a2b-a188fad2661a","detail-type":"my.Detail","source":"my.Source","account":"123456123456","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"7379586022458917877","x-datadog-parent-id":"2644033662113726488","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}',
             attributes: {
               ApproximateReceiveCount: "1",
               AWSTraceHeader: "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0",
@@ -124,5 +125,45 @@ describe("EventBridgeSQSEventTraceExtractor", () => {
       const traceContext = extractor.extract(payload);
       expect(traceContext).toBeNull();
     });
+
+    it("extracts trace context from Step Function EventBridge-SQS event", () => {
+      // Reset StepFunctionContextService instance
+      StepFunctionContextService["_instance"] = undefined as any;
+
+      const tracerWrapper = new TracerWrapper();
+
+      const payload = {
+        Records: [
+          {
+            messageId: "0fc0e02f-ab25-4fde-b5ff-22aba9a9f20e",
+            receiptHandle:
+              "AQEBROlXUgqqRdo/j0GcfxBNldIKy8FO6Ee0ZnP5YeAp4pwQ+v9XovX47vSzNHAZooCa0r8D7Uoow0y4bhGiH/Tt5HXAseDUlvWHB6bULonzAdvRmLd1W1OCY9D1uH3TpHZYn6JdoQd6Koxndx5wDwhv5UKxcbOwDjlc6X/30OKkTm4gcr7Otzu4GxCt6N/FmDxcRIDogZk80UE1kN6Q5EHI9LB6V+oleqqCbQwg5FYmbVc+DjwPBY4/5NI6x1/XZLFZA0TdezOdOuNq4+4DGK8e35Bafg4hXp+06zg8E5XPdMQV5V4iDzJhenPEdXusGL36byBHyC4aDunTSpeIND0/0ctqyH1vEJHo09LJ1jztPj05hBQeDU5QXCIKRpuo5+nEHE+Jm1ZLrUWUoIg1uAIamDzQ0CWNtPjGkNn3POiTGpD2e0aqrE5VpXZ8N30HKKFM",
+            body: '{"version":"0","id":"1aec4f0c-35e7-934c-9928-a5db3e526bca","detail-type":"ProcessEvent","source":"demo.stepfunction","account":"123456123456","time":"2025-07-14T19:33:03Z","region":"sa-east-1","resources":["arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sqs-demo-dev-state-machine","arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d"],"detail":{"message":"Event from Step Functions","timestamp":"2025-07-14T19:33:03.483Z","executionName":"aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","stateMachineName":"rstrat-sfn-evb-sqs-demo-dev-state-machine","input":{"testData":"Hello with SQS integration"},"_datadog":{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","StartTime":"2025-07-14T19:33:03.446Z","Name":"aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-evb-sqs-demo-d-StepFunctionsExecutionRol-mAumgN9x07FQ","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sqs-demo-dev-state-machine","Name":"rstrat-sfn-evb-sqs-demo-dev-state-machine"},"State":{"Name":"PublishToEventBridge","EnteredTime":"2025-07-14T19:33:03.483Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","serverless-version":"v1"}}}',
+            attributes: {
+              ApproximateReceiveCount: "1",
+              SentTimestamp: "1752521583745",
+              SenderId: "AIDAIELDKKY42PBA6I2NG",
+              ApproximateFirstReceiveTimestamp: "1752521583758",
+            },
+            messageAttributes: {},
+            md5OfBody: "957cded00b7b10a6e1b79864f24a7b5f",
+            eventSource: "aws:sqs",
+            eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-evb-sqs-demo-dev-process-event-queue",
+            awsRegion: "sa-east-1",
+          },
+        ],
+      };
+
+      const extractor = new EventBridgeSQSEventTraceExtractor(tracerWrapper);
+
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      // The trace IDs are deterministically generated from the Step Function execution context
+      expect(traceContext?.toTraceId()).toBe("7858567057595668526");
+      expect(traceContext?.toSpanId()).toBe("3674709292670593712");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
   });
 });

src/trace/context/extractors/event-bridge-sqs.ts

@@ -1,30 +1,24 @@
-import { EventBridgeEvent, SQSEvent } from "aws-lambda";
-import { EventTraceExtractor } from "../extractor";
-import { logDebug } from "../../../utils";
+import { SQSEvent } from "aws-lambda";
 import { TracerWrapper } from "../../tracer-wrapper";
+import { EventTraceExtractor } from "../extractor";
 import { SpanContextWrapper } from "../../span-context-wrapper";
+import { extractTraceContext, handleExtractionError } from "../extractor-utils";
 
 export class EventBridgeSQSEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
 
   extract(event: SQSEvent): SpanContextWrapper | null {
-    const body = event?.Records?.[0]?.body;
-    if (body === undefined) return null;
-
     try {
-      const parsedBody = JSON.parse(body) as EventBridgeEvent<any, any>;
-      const headers = parsedBody?.detail?._datadog;
-      if (headers === undefined) return null;
-
-      const traceContext = this.tracerWrapper.extract(headers);
-      if (traceContext === null) return null;
-
-      logDebug("Extracted trace context from EventBridge-SQS event", { traceContext, event });
-      return traceContext;
-    } catch (error) {
-      if (error instanceof Error) {
-        logDebug("Unable to extract trace context from EventBridge-SQS event", error);
+      const body = event?.Records?.[0]?.body;
+      if (body) {
+        const parsedBody = JSON.parse(body);
+        const headers = parsedBody?.detail?._datadog;
+        if (headers) {
+          return extractTraceContext(headers, this.tracerWrapper);
+        }
       }
+    } catch (error) {
+      handleExtractionError(error, "EventBridge-SQS");
     }
 
     return null;

src/trace/context/extractors/event-bridge.spec.ts

@@ -1,5 +1,6 @@
 import { TracerWrapper } from "../../tracer-wrapper";
 import { EventBridgeEventTraceExtractor } from "./event-bridge";
+import { StepFunctionContextService } from "../../step-function-service";
 
 let mockSpanContext: any = null;
 
@@ -106,5 +107,66 @@ describe("EventBridgeEventTraceExtractor", () => {
       const traceContext = extractor.extract(payload);
       expect(traceContext).toBeNull();
     });
+
+    it("extracts trace context from Step Function EventBridge event", () => {
+      // Reset StepFunctionContextService instance
+      StepFunctionContextService["_instance"] = undefined as any;
+
+      const tracerWrapper = new TracerWrapper();
+
+      const payload = {
+        version: "0",
+        id: "af718b2a-b987-e8c0-7a2b-a188fad2661a",
+        "detail-type": "ProcessEvent",
+        source: "demo.stepfunction",
+        account: "123456123456",
+        time: "2025-07-11T14:59:35Z",
+        region: "us-east-2",
+        resources: [
+          "arn:aws:states:us-east-2:123456123456:stateMachine:rstrat-sfn-evb-demo-dev-state-machine",
+          "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105",
+        ],
+        detail: {
+          message: "Event from Step Functions",
+          timestamp: "2025-07-11T14:59:35.830Z",
+          executionName: "6c190e7b-eb77-46db-af26-9066d353b105",
+          stateMachineName: "rstrat-sfn-evb-demo-dev-state-machine",
+          input: {
+            testData: "Hello with Datadog tracing",
+          },
+          _datadog: {
+            Execution: {
+              Id: "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105",
+              StartTime: "2025-07-11T14:59:35.806Z",
+              Name: "6c190e7b-eb77-46db-af26-9066d353b105",
+              RoleArn: "arn:aws:iam::123456123456:role/rstrat-sfn-evb-demo-dev-StepFunctionsExecutionRole-8maJHu01fhZZ",
+              RedriveCount: 0,
+            },
+            StateMachine: {
+              Id: "arn:aws:states:us-east-2:123456123456:stateMachine:rstrat-sfn-evb-demo-dev-state-machine",
+              Name: "rstrat-sfn-evb-demo-dev-state-machine",
+            },
+            State: {
+              Name: "PublishToEventBridge",
+              EnteredTime: "2025-07-11T14:59:35.830Z",
+              RetryCount: 0,
+            },
+            RootExecutionId:
+              "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105",
+            "serverless-version": "v1",
+          },
+        },
+      };
+
+      const extractor = new EventBridgeEventTraceExtractor(tracerWrapper);
+
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      expect(traceContext?.toTraceId()).toBe("1503104665848096006");
+      expect(traceContext?.toSpanId()).toBe("159267866761498620");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
   });
 });

src/trace/context/extractors/event-bridge.ts

@@ -1,26 +1,20 @@
-import { EventTraceExtractor } from "../extractor";
 import { EventBridgeEvent } from "aws-lambda";
-import { logDebug } from "../../../utils";
 import { TracerWrapper } from "../../tracer-wrapper";
+import { EventTraceExtractor } from "../extractor";
 import { SpanContextWrapper } from "../../span-context-wrapper";
+import { extractTraceContext, handleExtractionError } from "../extractor-utils";
 
 export class EventBridgeEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
 
-  extract(event: EventBridgeEvent<any, any>): SpanContextWrapper | null {
-    const headers = event?.detail?._datadog;
-    if (headers === undefined) return null;
-
+  extract(event: EventBridgeEvent<string, any>): SpanContextWrapper | null {
     try {
-      const traceContext = this.tracerWrapper.extract(headers);
-      if (traceContext === null) return null;
-
-      logDebug(`Extracted trace context from Eventbridge event`, { traceContext, event });
-      return traceContext;
-    } catch (error) {
-      if (error instanceof Error) {
-        logDebug("Unable to extract trace context from EventBridge event", error);
+      const headers = event?.detail?._datadog;
+      if (headers) {
+        return extractTraceContext(headers, this.tracerWrapper);
       }
+    } catch (error) {
+      handleExtractionError(error, "EventBridge");
     }
 
     return null;

src/trace/context/extractors/sns-sqs.spec.ts

@@ -1,5 +1,6 @@
 import { TracerWrapper } from "../../tracer-wrapper";
 import { SNSSQSEventTraceExtractor } from "./sns-sqs";
+import { StepFunctionContextService } from "../../step-function-service";
 
 let mockSpanContext: any = null;
 
@@ -212,5 +213,54 @@ describe("SNSSQSEventTraceExtractor", () => {
       expect(traceContext?.sampleMode()).toBe("1");
       expect(traceContext?.source).toBe("event");
     });
+
+    it("extracts trace context from Step Function SNS-SQS event", () => {
+      // Reset StepFunctionContextService instance
+      StepFunctionContextService["_instance"] = undefined as any;
+
+      const tracerWrapper = new TracerWrapper();
+
+      const payload = {
+        Records: [
+          {
+            messageId: "43a5f138-f166-40f1-b7e4-a7e0af9d633d",
+            receiptHandle:
+              "AQEBrhzl4RiITHp/ui07Y0DlDIdrmYHveKjqDIsx2gG7Z3fvrDohnfnpy/r4esh/ZsilJUR/C3uohYe6HUqvixymhx+io9S/MYNoA1zjmSVd1V4ZKe6saMs6L7aSW5TgrLpuxOtNGWvmNijdlQlOoYW1xRlkzkBywFkELfazExJHrThbxpxXcHbcAoh1Vz77EvlcAQNbc11vTccoUcMcdczvoLd/wgyrsIf0z8qdUQHaspWYoWOlZOsoflDMddYwqWO3LNRphAGMp5ISTDVbVqo1/U+wOqBj+b3dOYP9k0vS9Mj+36t+EJ8+KETFXRPNk4mZ+7hvG+UCYBN582gT502MnQitxylHKWOlH77nIokfk43FjhjsybLE48KdWdO49O2WKslXwCpPLQWnbKWlUl05/12tIk41MolVyfiWywW9R/S7hgcSr51tEBcjZTW8GR8r",
+            body: '{"testData":"Hello from SQS integration test","timestamp":"2025-07-15T18:16:27Z"}',
+            attributes: {
+              ApproximateReceiveCount: "1",
+              SentTimestamp: "1752603390964",
+              SenderId: "AIDAIOA2GYWSHW4E2VXIO",
+              ApproximateFirstReceiveTimestamp: "1752603390980",
+            },
+            messageAttributes: {
+              _datadog: {
+                stringValue:
+                  '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-sqs-demo-dev-state-machine:c363b975-c342-4e40-815a-8dd2496f5e81","StartTime":"2025-07-15T18:16:30.746Z","Name":"c363b975-c342-4e40-815a-8dd2496f5e81","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sns-sqs-demo-d-StepFunctionsExecutionRol-T2O3igeuSihu","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sns-sqs-demo-dev-state-machine","Name":"rstrat-sfn-sns-sqs-demo-dev-state-machine"},"State":{"Name":"PublishToSNS","EnteredTime":"2025-07-15T18:16:30.776Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-sqs-demo-dev-state-machine:c363b975-c342-4e40-815a-8dd2496f5e81","serverless-version":"v1"}',
+                stringListValues: [],
+                binaryListValues: [],
+                dataType: "String",
+              },
+            },
+            md5OfBody: "1e832c0d0aa5188dc5e3f2e85c9cb5e7",
+            md5OfMessageAttributes: "64e36d01aec95ca5a2160c13299e9c3b",
+            eventSource: "aws:sqs",
+            eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sns-sqs-demo-dev-process-event-queue",
+            awsRegion: "sa-east-1",
+          },
+        ],
+      };
+
+      const extractor = new SNSSQSEventTraceExtractor(tracerWrapper);
+
+      const traceContext = extractor.extract(payload);
+      expect(traceContext).not.toBeNull();
+
+      // The StepFunctionContextService generates deterministic trace IDs
+      expect(traceContext?.toTraceId()).toBe("1657966791618574655");
+      expect(traceContext?.toSpanId()).toBe("5100002956473485303");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
   });
 });