[SVLS-8538] add durable function tags (#730)

* [SVLS-8538] add durable function tags

* format

* add a  tiny comment to re-trigger the pipeline

---------

Co-authored-by: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com>

Author: jchrostek-dd

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

@@ -0,0 +1,97 @@
+import { parseDurableExecutionArn, extractDurableFunctionContext } from "./durable-function-context";
+
+describe("durable-function-context", () => {
+  describe("parseDurableExecutionArn", () => {
+    it("returns execution name and ID for a valid ARN", () => {
+      const arn =
+        "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/550e8400-e29b-41d4-a716-446655440001";
+      const result = parseDurableExecutionArn(arn);
+
+      expect(result).toEqual({
+        executionName: "order-123",
+        executionId: "550e8400-e29b-41d4-a716-446655440001",
+      });
+    });
+
+    it("returns undefined for ARN without durable-execution marker", () => {
+      const arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST";
+      const result = parseDurableExecutionArn(arn);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined for malformed ARN with only execution name", () => {
+      const arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123";
+      const result = parseDurableExecutionArn(arn);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined for malformed ARN with empty execution name", () => {
+      const arn =
+        "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution//550e8400-e29b-41d4-a716-446655440002";
+      const result = parseDurableExecutionArn(arn);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined for malformed ARN with empty execution ID", () => {
+      const arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/";
+      const result = parseDurableExecutionArn(arn);
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe("extractDurableFunctionContext", () => {
+    it("extracts context from event.DurableExecutionArn", () => {
+      const event = {
+        DurableExecutionArn:
+          "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+        CheckpointToken: "some-token",
+        InitialExecutionState: {
+          Operations: [],
+        },
+      };
+      const result = extractDurableFunctionContext(event);
+
+      expect(result).toEqual({
+        durable_function_execution_name: "my-execution",
+        durable_function_execution_id: "550e8400-e29b-41d4-a716-446655440004",
+      });
+    });
+
+    it("returns undefined for regular Lambda event without DurableExecutionArn", () => {
+      const event = {
+        body: '{"key": "value"}',
+        headers: {
+          "Content-Type": "application/json",
+        },
+      };
+      const result = extractDurableFunctionContext(event);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined when event is null", () => {
+      const result = extractDurableFunctionContext(null);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined when event is undefined", () => {
+      const result = extractDurableFunctionContext(undefined);
+
+      expect(result).toBeUndefined();
+    });
+
+    it("returns undefined when DurableExecutionArn cannot be parsed", () => {
+      const event = {
+        DurableExecutionArn: "invalid-arn-without-durable-execution-marker",
+      };
+      const result = extractDurableFunctionContext(event);
+
+      expect(result).toBeUndefined();
+    });
+  });
+});

src/trace/durable-function-context.ts

@@ -0,0 +1,37 @@
+import { logDebug } from "../utils";
+
+export interface DurableFunctionContext {
+  durable_function_execution_name: string;
+  durable_function_execution_id: string;
+}
+
+export function extractDurableFunctionContext(event: any): DurableFunctionContext | undefined {
+  const durableExecutionArn = event?.DurableExecutionArn;
+
+  if (typeof durableExecutionArn !== "string") {
+    return undefined;
+  }
+
+  const parsed = parseDurableExecutionArn(durableExecutionArn);
+  if (!parsed) {
+    logDebug("Failed to parse DurableExecutionArn", { arn: durableExecutionArn });
+    return undefined;
+  }
+
+  return {
+    durable_function_execution_name: parsed.executionName,
+    durable_function_execution_id: parsed.executionId,
+  };
+}
+
+/**
+ * Parses a DurableExecutionArn to extract execution name and ID.
+ * ARN format: arn:aws:lambda:{region}:{account}:function:{func}:{version}/durable-execution/{name}/{id}
+ */
+export function parseDurableExecutionArn(arn: string): { executionName: string; executionId: string } | undefined {
+  // Match only the trailing durable execution segment.
+  const match = arn.match(/\/durable-execution\/([^/]+)\/([^/]+)$/);
+  if (!match) return undefined;
+  const [, executionName, executionId] = match;
+  return { executionName, executionId };
+}

src/trace/listener.ts

@@ -15,6 +15,7 @@ import { SpanWrapper } from "./span-wrapper";
 import { getTraceTree, clearTraceTree } from "../runtime/index";
 import { TraceContext, TraceContextService, TraceSource } from "./trace-context-service";
 import { StepFunctionContext, StepFunctionContextService } from "./step-function-service";
+import { DurableFunctionContext, extractDurableFunctionContext } from "./durable-function-context";
 import { XrayService } from "./xray-service";
 import { AUTHORIZING_REQUEST_ID_HEADER } from "./context/extractors/http";
 import { getSpanPointerAttributes, SpanPointerAttributes } from "../utils/span-pointers";
@@ -85,6 +86,7 @@ export class TraceListener {
   private contextService: TraceContextService;
   private context?: Context;
   private stepFunctionContext?: StepFunctionContext;
+  private durableFunctionContext?: DurableFunctionContext;
   private tracerWrapper: TracerWrapper;
   private inferrer: SpanInferrer;
   private inferredSpan?: SpanWrapper;
@@ -146,6 +148,7 @@ export class TraceListener {
     const eventSource = parseEventSource(event);
     this.triggerTags = extractTriggerTags(event, context, eventSource);
     this.stepFunctionContext = StepFunctionContextService.instance().context;
+    this.durableFunctionContext = extractDurableFunctionContext(event);
 
     if (this.config.addSpanPointers) {
       this.spanPointerAttributesList = getSpanPointerAttributes(eventSource, event);
@@ -288,6 +291,7 @@ export class TraceListener {
 
     // Reset singletons and trace context
     this.stepFunctionContext = undefined;
+    this.durableFunctionContext = undefined;
     StepFunctionContextService.reset();
     this.contextService.reset();
   }
@@ -328,6 +332,13 @@ export class TraceListener {
         ...this.stepFunctionContext,
       };
     }
+    if (this.durableFunctionContext) {
+      logDebug("Applying durable function context to the aws.lambda span");
+      options.tags = {
+        ...options.tags,
+        ...this.durableFunctionContext,
+      };
+    }
     if (this.lambdaSpanParentContext) {
       options.childOf = this.lambdaSpanParentContext;
     }