Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 97 additions & 0 deletions src/trace/durable-function-context.spec.ts
Original file line number Diff line number Diff line change
@@ -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();
});
});
});
37 changes: 37 additions & 0 deletions src/trace/durable-function-context.ts
Original file line number Diff line number Diff line change
@@ -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 };
}
11 changes: 11 additions & 0 deletions src/trace/listener.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -288,6 +291,7 @@ export class TraceListener {

// Reset singletons and trace context
this.stepFunctionContext = undefined;
this.durableFunctionContext = undefined;
StepFunctionContextService.reset();
this.contextService.reset();
}
Expand Down Expand Up @@ -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;
}
Expand Down
Loading