wip

Author: ryanstrat

.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

@@ -12,7 +12,6 @@ import { XrayService } from "../xray-service";
  * 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
- * @param eventType The type of event (for logging)
  * @returns SpanContextWrapper or null
  */
 export function extractTraceContext(headers: any, tracerWrapper: TracerWrapper): SpanContextWrapper | null {

src/trace/context/extractor.spec.ts

@@ -694,6 +694,54 @@ describe("TraceContextExtractor", () => {
         expect(traceContext?.sampleMode()).toBe("1");
         expect(traceContext?.source).toBe("event");
       });
+
+      it("extracts and applies deterministic trace ID from StepFunction event end-to-end", async () => {
+        // This test verifies the complete flow from Step Function event to trace ID assignment
+        const event = {
+          Execution: {
+            Id: "arn:aws:states:us-east-1:123456789012:execution:MyStateMachine:test-execution-id",
+            Name: "test-execution-id",
+            StartTime: "2024-01-01T00:00:00.000Z",
+          },
+          State: {
+            Name: "ProcessData",
+            EnteredTime: "2024-01-01T00:00:01.000Z",
+            RetryCount: 0,
+          },
+          StateMachine: {
+            Id: "arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine",
+            Name: "MyStateMachine",
+          },
+        };
+
+        const tracerWrapper = new TracerWrapper();
+        const extractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig);
+
+        // Extract trace context through the full pipeline
+        const traceContext = await extractor.extract(event, {} as Context);
+        expect(traceContext).not.toBeNull();
+
+        // Verify the trace context was extracted from Step Function
+        expect(traceContext?.source).toBe("event");
+
+        // Verify the trace ID is deterministic based on execution ID
+        // The trace ID should be generated from SHA256 hash of the execution ID
+        const traceId = traceContext?.toTraceId();
+        expect(traceId).toBeDefined();
+        expect(traceId).not.toBe("0"); // Should not be zero
+        expect(traceId).toMatch(/^\d+$/); // Should be numeric string
+
+        // Verify the span ID is deterministic based on execution ID, state name, and entered time
+        const spanId = traceContext?.toSpanId();
+        expect(spanId).toBeDefined();
+        expect(spanId).not.toBe("0"); // Should not be zero
+        expect(spanId).toMatch(/^\d+$/); // Should be numeric string
+
+        // Verify that extracting the same event produces the same trace IDs (deterministic)
+        const traceContext2 = await extractor.extract(event, {} as Context);
+        expect(traceContext2?.toTraceId()).toBe(traceId);
+        expect(traceContext2?.toSpanId()).toBe(spanId);
+      });
     });
 
     describe("lambda context", () => {

src/trace/context/extractor.ts

@@ -43,8 +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) {
       const customExtractor = new CustomTraceExtractor(this.config.traceExtractor);
@@ -57,7 +55,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 +94,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/sns-sqs.ts

@@ -9,6 +9,7 @@ export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
 
   extract(event: SQSEvent): SpanContextWrapper | null {
+    logDebug("SNS-SQS Extractor Being Used");
     try {
       // Try to extract trace context from SNS wrapped in SQS
       const body = event?.Records?.[0]?.body;

src/trace/context/extractors/sqs.ts

@@ -9,6 +9,7 @@ export class SQSEventTraceExtractor implements EventTraceExtractor {
   constructor(private tracerWrapper: TracerWrapper) {}
 
   extract(event: SQSEvent): SpanContextWrapper | null {
+    logDebug("SQS Extractor Being Used");
     try {
       // First try to extract trace context from message attributes
       let headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
@@ -25,6 +26,7 @@ export class SQSEventTraceExtractor implements EventTraceExtractor {
 
       if (headers) {
         const parsedHeaders = JSON.parse(headers);
+
         const traceContext = extractTraceContext(parsedHeaders, this.tracerWrapper);
         if (traceContext) {
           return traceContext;

src/trace/context/extractors/step-function.spec.ts

@@ -1,5 +1,6 @@
 import { StepFunctionContextService } from "../../step-function-service";
 import { StepFunctionEventTraceExtractor } from "./step-function";
+import { TracerWrapper } from "../../tracer-wrapper";
 
 describe("StepFunctionEventTraceExtractor", () => {
   beforeEach(() => {
@@ -112,5 +113,77 @@ describe("StepFunctionEventTraceExtractor", () => {
       const traceContext = extractor.extract({});
       expect(traceContext).toBeNull();
     });
+
+    it("extracts trace context end-to-end with deterministic IDs for v1 nested context", () => {
+      const nestedPayload = {
+        _datadog: {
+          Execution: {
+            Id: "arn:aws:states:us-east-1:123456789012:execution:parent-machine:parent-execution-id",
+            Name: "parent-execution-id",
+            StartTime: "2024-01-01T00:00:00.000Z",
+          },
+          State: {
+            Name: "InvokeNestedStateMachine",
+            EnteredTime: "2024-01-01T00:00:05.000Z",
+            RetryCount: 0,
+          },
+          "serverless-version": "v1",
+          RootExecutionId: "arn:aws:states:us-east-1:123456789012:execution:root-machine:root-execution-id",
+        },
+      };
+
+      const extractor = new StepFunctionEventTraceExtractor();
+      const traceContext = extractor.extract(nestedPayload);
+
+      expect(traceContext).not.toBeNull();
+      expect(traceContext?.source).toBe("event");
+
+      // Verify trace ID is based on root execution ID (deterministic)
+      const traceId = traceContext?.toTraceId();
+      expect(traceId).toBeDefined();
+      expect(traceId).not.toBe("0");
+
+      // Extract again to verify deterministic behavior
+      const traceContext2 = extractor.extract(nestedPayload);
+      expect(traceContext2?.toTraceId()).toBe(traceId);
+    });
+
+    it("extracts trace context end-to-end with propagated IDs for v1 lambda root context", () => {
+      const lambdaRootPayload = {
+        _datadog: {
+          Execution: {
+            Id: "arn:aws:states:us-east-1:123456789012:execution:machine:execution-id",
+            Name: "execution-id",
+            StartTime: "2024-01-01T00:00:00.000Z",
+          },
+          State: {
+            Name: "ProcessLambda",
+            EnteredTime: "2024-01-01T00:00:02.000Z",
+            RetryCount: 0,
+          },
+          "serverless-version": "v1",
+          "x-datadog-trace-id": "1234567890123456789",
+          "x-datadog-tags": "_dd.p.tid=fedcba9876543210,_dd.p.dm=-0",
+        },
+      };
+
+      const extractor = new StepFunctionEventTraceExtractor();
+      const traceContext = extractor.extract(lambdaRootPayload);
+
+      expect(traceContext).not.toBeNull();
+      expect(traceContext?.source).toBe("event");
+
+      // Verify trace ID comes from propagated headers, not deterministic hash
+      expect(traceContext?.toTraceId()).toBe("1234567890123456789");
+
+      // Verify deterministic span ID based on execution details
+      const spanId = traceContext?.toSpanId();
+      expect(spanId).toBeDefined();
+      expect(spanId).not.toBe("0");
+
+      // Extract again to verify deterministic span ID
+      const traceContext2 = extractor.extract(lambdaRootPayload);
+      expect(traceContext2?.toSpanId()).toBe(spanId);
+    });
   });
 });

src/trace/context/extractors/step-function.ts

@@ -5,11 +5,16 @@ import { EventTraceExtractor } from "../extractor";
 export class StepFunctionEventTraceExtractor implements EventTraceExtractor {
   extract(event: any): SpanContextWrapper | null {
     // Probably StepFunctionContextService hasn't been called
-    const instance = StepFunctionContextService.instance(event);
-    const context = instance.context;
+    const stepFunctionInstance = StepFunctionContextService.instance(event);
+    const stepFunctionContext = stepFunctionInstance.context;
 
-    if (context === undefined) return null;
+    if (stepFunctionContext !== undefined) {
+      const spanContext = stepFunctionInstance.spanContext;
+      if (spanContext !== null) {
+        return spanContext;
+      }
+    }
 
-    return instance.spanContext;
+    return null;
   }
 }

src/trace/listener.spec.ts

@@ -33,6 +33,15 @@ jest.mock("./tracer-wrapper", () => {
       return mockExtract(event);
     }
 
+    startSpan(name: any, options: any): any {
+      return {
+        toSpanId: () => "mockSpanId",
+        toTraceId: () => "mockTraceId",
+        finish: jest.fn(),
+        setTag: jest.fn(),
+      };
+    }
+
     injectSpan(span: any): any {
       return {
         [DATADOG_PARENT_ID_HEADER]: span.toSpanId(),
@@ -304,6 +313,88 @@ describe("TraceListener", () => {
     );
   });
 
+  it("wraps dd-trace span around invocation with Step Function context", async () => {
+    const listener = new TraceListener(defaultConfig);
+    mockTraceSource = TraceSource.Event;
+
+    // Mock Step Function context with deterministic trace IDs
+    mockSpanContext = {
+      toTraceId: () => "512d06a10e5e34cb", // Hex converted to decimal would be different
+      toSpanId: () => "7069a031ef9ad2cc",
+      _sampling: {
+        priority: "1",
+      },
+    };
+    mockSpanContextWrapper = {
+      spanContext: mockSpanContext,
+    };
+
+    const stepFunctionSQSEvent = {
+      Records: [
+        {
+          messageId: "4ead33f3-51c8-4094-87bd-5325dc143cbd",
+          receiptHandle:
+            "AQEBrGtLZCUS1POUEZtdZRoB0zXgT14OQC48A4Xk4Qbnv/v4d0ib5rFI1wEah823t2hE9haPm6nNN1aGsJmYkqa9Y8qaBQscp9f7HKJyybT5hpdKEn07fY0VRv/Of63u1RN1YdFdY5uhI8XGWRc4w7t62lQwMMFY5Ahy7XLVwnav81KRjGFdgxzITrtx3YKxmISNvXzPiiHNKb7jT+ClfXi91bEYHi3Od3ji5xGajAofgYrj2VBDULyohsfMkwlvAanD2wfj2x++wL5LSpFEtMFnvThzt7Dh5FEZChVMzWV+fRFpljivHX58ZeuGv4yIIjLVuuDGn5uAY5ES4CsdINrBAru6K5gDSPUajRzE3TktNgAq5Niqfky1x0srLRAJjTDdmZK8/CXU0sRT/MCT99vkCHa0bC17S/9au5bCbrB4k/T9J8W39AA6kIYhebkq3IQr",
+          body: '{"testData":"Hello from Step Functions to SQS"}',
+          attributes: {
+            ApproximateReceiveCount: "1",
+            SentTimestamp: "1752594520503",
+            SenderId: "AROAWGCM4HXU73A4V34AJ:EcGTcmgJbwwOwXPbloVwgSaDOmwhYBLH",
+            ApproximateFirstReceiveTimestamp: "1752594520516",
+          },
+          messageAttributes: {
+            _datadog: {
+              stringValue:
+                '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","StartTime":"2025-07-15T15:48:40.302Z","Name":"a4912895-93a3-4803-a712-69fecb55c025","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sqs-demo-dev-StepFunctionsExecutionRole-s6ozc2dVrvLH","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sqs-demo-dev-state-machine","Name":"rstrat-sfn-sqs-demo-dev-state-machine"},"State":{"Name":"SendToSQS","EnteredTime":"2025-07-15T15:48:40.333Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","serverless-version":"v1"}',
+              stringListValues: [],
+              binaryListValues: [],
+              dataType: "String",
+            },
+          },
+          md5OfMessageAttributes: "5469b8f90bb6ab27e95816c1fa178680",
+          md5OfBody: "f0c0ddb2ed09a09e8791013f142e8d7e",
+          eventSource: "aws:sqs",
+          eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sqs-demo-dev-process-event-queue",
+          awsRegion: "sa-east-1",
+        },
+      ],
+    };
+
+    await listener.onStartInvocation(stepFunctionSQSEvent, context as any);
+    const unwrappedFunc = () => {};
+    const wrappedFunc = listener.onWrap(unwrappedFunc);
+    wrappedFunc();
+    await listener.onCompleteInvocation();
+
+    expect(mockWrap).toHaveBeenCalledWith(
+      "aws.lambda",
+      {
+        resource: "my-Lambda",
+        service: "my-Lambda",
+        tags: {
+          cold_start: true,
+          function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
+          function_version: "$LATEST",
+          request_id: "1234",
+          resource_names: "my-Lambda",
+          functionname: "my-lambda",
+          "_dd.parent_source": "event",
+          "function_trigger.event_source": "sqs",
+          "function_trigger.event_source_arn":
+            "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sqs-demo-dev-process-event-queue",
+          datadog_lambda: datadogLambdaVersion,
+          dd_trace: ddtraceVersion,
+        },
+        type: "serverless",
+        childOf: expect.objectContaining({
+          toSpanId: expect.any(Function),
+          toTraceId: expect.any(Function),
+        }),
+      },
+      unwrappedFunc,
+    );
+  });
+
   it("injects authorizer context if it exists", async () => {
     const listener = new TraceListener(defaultConfig);
     mockTraceSource = TraceSource.Event;

src/trace/tracer-wrapper.ts

@@ -75,7 +75,9 @@ export class TracerWrapper {
   }
 
   public startSpan<T = (...args: any[]) => any>(name: string, options: TraceOptions): T | null {
+    logDebug("Starting Span named:", { name });
     if (!this.isTracerAvailable) {
+      logDebug("No Tracer available, cannot start span");
       return null;
     }
     return this.tracer.startSpan(name, options);