Propagate Step Function Trace Context through Managed Services

src/trace/context/extractor-utils.spec.ts

@@ -0,0 +1,164 @@
+import { TracerWrapper } from "../tracer-wrapper";
+import { extractTraceContext, extractFromAWSTraceHeader } from "./extractor-utils";
+import { StepFunctionContextService } from "../step-function-service";
+
+describe("extractor-utils", () => {
+  beforeEach(() => {
+    StepFunctionContextService["_instance"] = undefined as any;
+  });
+  describe("extractTraceContext", () => {
+    it("returns span context when tracer wrapper successfully extracts from headers", () => {
+      const legacyStepFunctionEvent = {
+        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",
+          RedriveCount: 0,
+          StartTime: "2022-12-08T21:08:17.924Z",
+        },
+        State: {
+          Name: "step-one",
+          EnteredTime: "2022-12-08T21:08:19.224Z",
+          RetryCount: 2,
+        },
+        StateMachine: {
+          Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+          Name: "my-state-machine",
+        },
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const result = extractTraceContext(legacyStepFunctionEvent, tracerWrapper);
+
+      // Should return a span context from Step Function context since headers extraction fails
+      expect(result).not.toBeNull();
+    });
+
+    it("returns null when no trace context can be extracted", () => {
+      const emptyEvent = {
+        someOtherProperty: "value",
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const result = extractTraceContext(emptyEvent, tracerWrapper);
+
+      expect(result).toBeNull();
+    });
+
+    it("extracts context from LambdaRootStepFunctionContext", () => {
+      const lambdaRootStepFunctionEvent = {
+        _datadog: {
+          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",
+            RedriveCount: 0,
+            StartTime: "2022-12-08T21:08:17.924Z",
+          },
+          State: {
+            Name: "step-one",
+            EnteredTime: "2022-12-08T21:08:19.224Z",
+            RetryCount: 2,
+          },
+          StateMachine: {
+            Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+            Name: "my-state-machine",
+          },
+          "x-datadog-trace-id": "10593586103637578129",
+          "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
+          "serverless-version": "v1",
+        },
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const result = extractTraceContext(lambdaRootStepFunctionEvent, tracerWrapper);
+
+      expect(result).not.toBeNull();
+    });
+
+    it("extracts context from NestedStepFunctionContext", () => {
+      const nestedStepFunctionEvent = {
+        _datadog: {
+          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",
+            RedriveCount: 0,
+            StartTime: "2022-12-08T21:08:17.924Z",
+          },
+          State: {
+            Name: "step-one",
+            EnteredTime: "2022-12-08T21:08:19.224Z",
+            RetryCount: 2,
+          },
+          StateMachine: {
+            Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+            Name: "my-state-machine",
+          },
+          RootExecutionId:
+            "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc",
+          "serverless-version": "v1",
+        },
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const result = extractTraceContext(nestedStepFunctionEvent, tracerWrapper);
+
+      expect(result).not.toBeNull();
+    });
+
+    it("extracts context from legacy lambda StepFunctionContext", () => {
+      const event = {
+        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",
+            RedriveCount: 0,
+            StartTime: "2022-12-08T21:08:17.924Z",
+          },
+          State: {
+            Name: "step-one",
+            EnteredTime: "2022-12-08T21:08:19.224Z",
+            RetryCount: 2,
+          },
+          StateMachine: {
+            Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+            Name: "my-state-machine",
+          },
+        },
+      };
+
+      const tracerWrapper = new TracerWrapper();
+      const result = extractTraceContext(event, tracerWrapper);
+
+      expect(result).not.toBeNull();
+    });
+  });
+
+  describe("extractFromAWSTraceHeader", () => {
+    it("returns null when AWS trace header is invalid", () => {
+      const invalidHeader = "invalid-header";
+      const eventType = "SQS";
+
+      const result = extractFromAWSTraceHeader(invalidHeader, eventType);
+
+      expect(result).toBeNull();
+    });
+  });
+});

src/trace/context/extractor-utils.ts

@@ -0,0 +1,62 @@
+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);
+
+  if (stepFunctionInstance.context !== undefined) {
+    if (stepFunctionInstance.spanContext !== null) {
+      return stepFunctionInstance.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);
+  }
+}