feat: Add automatic trace propagation for Lambda Durable Functions

src/index.ts

@@ -27,6 +27,7 @@ import { getEnhancedMetricTags } from "./metrics/enhanced-metrics";
 import { DatadogTraceHeaders } from "./trace/context/extractor";
 import { SpanWrapper } from "./trace/span-wrapper";
 import { SpanOptions, TracerWrapper } from "./trace/tracer-wrapper";
+import { initDurableFunctionTracing } from "./trace/durable-function-patch";
 
 // Backwards-compatible export, TODO deprecate in next major
 export { DatadogTraceHeaders as TraceHeaders } from "./trace/context/extractor";
@@ -113,6 +114,9 @@ if (getEnvValue(coldStartTracingEnvVar, "true").toLowerCase() === "true" && !isM
   subscribeToDC();
 }
 
+// Initialize durable function tracing if SDK is present
+initDurableFunctionTracing();
+
 const initTime = Date.now();
 
 /**

src/trace/context/extractor.ts

@@ -14,8 +14,10 @@ import {
   SNSSQSEventTraceExtractor,
   SQSEventTraceExtractor,
   StepFunctionEventTraceExtractor,
+  DurableFunctionEventTraceExtractor,
 } from "./extractors";
 import { StepFunctionContextService } from "../step-function-service";
+import { DurableFunctionContextService } from "../durable-function-service";
 import { EventValidator } from "../../utils/event-validator";
 import { TracerWrapper } from "../tracer-wrapper";
 import { SpanContextWrapper } from "../span-context-wrapper";
@@ -37,6 +39,7 @@ export interface DatadogTraceHeaders {
 export class TraceContextExtractor {
   private xrayService: XrayService;
   private stepFunctionContextService?: StepFunctionContextService;
+  private durableFunctionContextService?: DurableFunctionContextService;
 
   constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {
     this.xrayService = new XrayService();
@@ -64,6 +67,14 @@ export class TraceContextExtractor {
       }
     }
 
+    if (spanContext === null) {
+      this.durableFunctionContextService = DurableFunctionContextService.instance(event);
+      if (this.durableFunctionContextService?.context) {
+        const extractor = new DurableFunctionEventTraceExtractor();
+        spanContext = extractor.extract(event);
+      }
+    }
+
     if (spanContext === null) {
       const contextExtractor = new LambdaContextTraceExtractor(this.tracerWrapper);
       spanContext = contextExtractor.extract(context);

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

@@ -0,0 +1,130 @@
+import { DurableFunctionEventTraceExtractor } from "./durable-function";
+import { DurableFunctionContextService } from "../../durable-function-service";
+
+describe("DurableFunctionEventTraceExtractor", () => {
+  const validDurableExecutionEvent = {
+    DurableExecutionArn:
+      "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/550e8400-e29b-41d4-a716-446655440001",
+    CheckpointToken: "some-token",
+    InitialExecutionState: {
+      Operations: [
+        {
+          Id: "step-1",
+          Type: "STEP",
+          Status: "SUCCEEDED",
+          ExecutionDetails: {
+            InputPayload: JSON.stringify({
+              headers: {
+                "x-datadog-trace-id": "12345678901234567890",
+                "x-datadog-parent-id": "9876543210987654321",
+              },
+            }),
+          },
+        },
+      ],
+    },
+  };
+
+  const regularLambdaEvent = {
+    body: '{"key": "value"}',
+    headers: {
+      "Content-Type": "application/json",
+    },
+  };
+
+  beforeEach(() => {
+    DurableFunctionContextService.reset();
+  });
+
+  describe("extract", () => {
+    it("extracts span context from valid durable execution event", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract(validDurableExecutionEvent);
+
+      expect(spanContext).not.toBeNull();
+      expect(spanContext?.toTraceId()).toBe("12345678901234567890");
+      expect(spanContext?.toSpanId()).toBe("9876543210987654321");
+      expect(spanContext?.source).toBe("event");
+    });
+
+    it("returns null for regular Lambda event", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract(regularLambdaEvent);
+
+      expect(spanContext).toBeNull();
+    });
+
+    it("returns null for null event", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract(null);
+
+      expect(spanContext).toBeNull();
+    });
+
+    it("returns null for undefined event", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract(undefined);
+
+      expect(spanContext).toBeNull();
+    });
+
+    it("returns null for event with invalid DurableExecutionArn", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract({
+        DurableExecutionArn: "invalid-arn",
+      });
+
+      expect(spanContext).toBeNull();
+    });
+
+    it("extracts deterministic span context when no preserved trace headers", () => {
+      const eventWithoutPreservedTrace = {
+        DurableExecutionArn:
+          "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-456/650e8400-e29b-41d4-a716-446655440002",
+        CheckpointToken: "some-token",
+        InitialExecutionState: {
+          Operations: [
+            {
+              Id: "step-1",
+              Type: "STEP",
+              Status: "SUCCEEDED",
+              ExecutionDetails: {
+                InputPayload: JSON.stringify({
+                  body: '{"key": "value"}',
+                }),
+              },
+            },
+          ],
+        },
+      };
+
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      const spanContext = extractor.extract(eventWithoutPreservedTrace);
+
+      expect(spanContext).not.toBeNull();
+      // Should have deterministic IDs based on execution ID
+      expect(spanContext?.toTraceId()).toBeDefined();
+      expect(spanContext?.toSpanId()).toBeDefined();
+    });
+
+    it("uses singleton service instance", () => {
+      const extractor = new DurableFunctionEventTraceExtractor();
+
+      // First extraction
+      extractor.extract(validDurableExecutionEvent);
+      const instance1 = DurableFunctionContextService.instance();
+
+      // Second extraction with same event
+      extractor.extract(validDurableExecutionEvent);
+      const instance2 = DurableFunctionContextService.instance();
+
+      expect(instance1).toBe(instance2);
+    });
+  });
+});

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

@@ -0,0 +1,19 @@
+import { SpanContextWrapper } from "../../span-context-wrapper";
+import { DurableFunctionContextService } from "../../durable-function-service";
+import { EventTraceExtractor } from "../extractor";
+
+export class DurableFunctionEventTraceExtractor implements EventTraceExtractor {
+  extract(event: any): SpanContextWrapper | null {
+    const durableFunctionInstance = DurableFunctionContextService.instance(event);
+    const durableFunctionContext = durableFunctionInstance.context;
+
+    if (durableFunctionContext !== undefined) {
+      const spanContext = durableFunctionInstance.spanContext;
+      if (spanContext !== null) {
+        return spanContext;
+      }
+    }
+
+    return null;
+  }
+}

src/trace/context/extractors/index.ts

@@ -7,5 +7,6 @@ export { SQSEventTraceExtractor } from "./sqs";
 export { SNSEventTraceExtractor } from "./sns";
 export { SNSSQSEventTraceExtractor } from "./sns-sqs";
 export { StepFunctionEventTraceExtractor } from "./step-function";
+export { DurableFunctionEventTraceExtractor } from "./durable-function";
 export { LambdaContextTraceExtractor } from "./lambda-context";
 export { CustomTraceExtractor } from "./custom";