diff --git a/src/trace/durable-function-context.spec.ts b/src/trace/durable-function-context.spec.ts index 689caa5d..81939e99 100644 --- a/src/trace/durable-function-context.spec.ts +++ b/src/trace/durable-function-context.spec.ts @@ -1,4 +1,8 @@ -import { parseDurableExecutionArn, extractDurableFunctionContext } from "./durable-function-context"; +import { + parseDurableExecutionArn, + extractDurableFunctionContext, + extractDurableExecutionStatus, +} from "./durable-function-context"; describe("durable-function-context", () => { describe("parseDurableExecutionArn", () => { @@ -132,4 +136,41 @@ describe("durable-function-context", () => { expect(result).toBeUndefined(); }); }); + + describe("extractDurableExecutionStatus", () => { + const durableEvent = { + DurableExecutionArn: + "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + }; + + it.each(["SUCCEEDED", "FAILED", "PENDING"])("returns %s when result.Status is %s", (status) => { + const result = extractDurableExecutionStatus(durableEvent, { Status: status }); + expect(result).toBe(status); + }); + + it("returns undefined when result.Status is not a valid status", () => { + const result = extractDurableExecutionStatus(durableEvent, { Status: "UNKNOWN" }); + expect(result).toBeUndefined(); + }); + + it("returns undefined when result has no Status field", () => { + const result = extractDurableExecutionStatus(durableEvent, {}); + expect(result).toBeUndefined(); + }); + + it("returns undefined when result is null", () => { + const result = extractDurableExecutionStatus(durableEvent, null); + expect(result).toBeUndefined(); + }); + + it("returns undefined when event has no DurableExecutionArn", () => { + const result = extractDurableExecutionStatus({ body: "{}" }, { Status: "SUCCEEDED" }); + expect(result).toBeUndefined(); + }); + + it("returns undefined when event is null", () => { + const result = extractDurableExecutionStatus(null, { Status: "SUCCEEDED" }); + expect(result).toBeUndefined(); + }); + }); }); diff --git a/src/trace/durable-function-context.ts b/src/trace/durable-function-context.ts index 29d4c53c..19630c81 100644 --- a/src/trace/durable-function-context.ts +++ b/src/trace/durable-function-context.ts @@ -6,6 +6,8 @@ export interface DurableFunctionContext { "aws_lambda.durable_function.first_invocation"?: string; } +const VALID_DURABLE_EXECUTION_STATUSES = new Set(["SUCCEEDED", "FAILED", "PENDING"]); + export function extractDurableFunctionContext(event: any): DurableFunctionContext | undefined { const durableExecutionArn = event?.DurableExecutionArn; @@ -33,6 +35,23 @@ export function extractDurableFunctionContext(event: any): DurableFunctionContex return context; } +/** + * Extracts the durable function execution status from the handler result. + * Only applies when the event contains a DurableExecutionArn. + */ +export function extractDurableExecutionStatus(event: any, result: any): string | undefined { + if (typeof event?.DurableExecutionArn !== "string") { + return undefined; + } + + const status = result?.Status; + if (typeof status !== "string" || !VALID_DURABLE_EXECUTION_STATUSES.has(status)) { + return undefined; + } + + return status; +} + /** * Parses a DurableExecutionArn to extract execution name and ID. * ARN format: arn:aws:lambda:{region}:{account}:function:{func}:{version}/durable-execution/{name}/{id} diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index f70865fa..246a907f 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -568,4 +568,44 @@ describe("TraceListener", () => { currentSpanSpy.mockRestore(); } }); + + it("sets execution_status tag on the aws.lambda span when result.Status is valid", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const durableEvent = { + DurableExecutionArn: + "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + }; + await listener.onStartInvocation(durableEvent, context as any); + listener.onEndingInvocation(durableEvent, { Status: "SUCCEEDED" }, false); + + expect(mockSetTag).toHaveBeenCalledWith("aws_lambda.durable_function.execution_status", "SUCCEEDED"); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("does not set execution_status tag when result.Status is invalid", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const durableEvent = { + DurableExecutionArn: + "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + }; + await listener.onStartInvocation(durableEvent, context as any); + listener.onEndingInvocation(durableEvent, { Status: "UNKNOWN" }, false); + + expect(mockSetTag).not.toHaveBeenCalledWith("aws_lambda.durable_function.execution_status", expect.anything()); + } finally { + currentSpanSpy.mockRestore(); + } + }); }); diff --git a/src/trace/listener.ts b/src/trace/listener.ts index 217ba7a8..f068c535 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -20,7 +20,11 @@ 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 { + DurableFunctionContext, + extractDurableFunctionContext, + extractDurableExecutionStatus, +} 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"; @@ -234,6 +238,10 @@ export class TraceListener { } } } + const executionStatus = extractDurableExecutionStatus(event, result); + if (executionStatus !== undefined) { + this.tracerWrapper.currentSpan.setTag("aws_lambda.durable_function.execution_status", executionStatus); + } let rootSpan = this.inferredSpan; if (!rootSpan) {