diff --git a/src/trace/context/extractor-utils.spec.ts b/src/trace/context/extractor-utils.spec.ts new file mode 100644 index 00000000..afe577e5 --- /dev/null +++ b/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(); + }); + }); +}); diff --git a/src/trace/context/extractor-utils.ts b/src/trace/context/extractor-utils.ts new file mode 100644 index 00000000..5f0f1e17 --- /dev/null +++ b/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); + } +} diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index cb47b6f6..37cae4a3 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -17,7 +17,6 @@ import { SNSEventTraceExtractor, SNSSQSEventTraceExtractor, SQSEventTraceExtractor, - StepFunctionEventTraceExtractor, } from "./extractors"; import { StepFunctionContextService } from "../step-function-service"; import { SpanContextWrapper } from "../span-context-wrapper"; @@ -694,6 +693,54 @@ describe("TraceContextExtractor", () => { expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); + + it("extracts and applies deterministic trace ID from StepFunction event end-to-end", async () => { + // This test verifies the complete flow from Step Function event to trace ID assignment + const event = { + Execution: { + Id: "arn:aws:states:us-east-1:123456789012:execution:MyStateMachine:test-execution-id", + Name: "test-execution-id", + StartTime: "2024-01-01T00:00:00.000Z", + }, + State: { + Name: "ProcessData", + EnteredTime: "2024-01-01T00:00:01.000Z", + RetryCount: 0, + }, + StateMachine: { + Id: "arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine", + Name: "MyStateMachine", + }, + }; + + const tracerWrapper = new TracerWrapper(); + const extractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + + // Extract trace context through the full pipeline + const traceContext = await extractor.extract(event, {} as Context); + expect(traceContext).not.toBeNull(); + + // Verify the trace context was extracted from Step Function + expect(traceContext?.source).toBe("event"); + + // Verify the trace ID is deterministic based on execution ID + // The trace ID should be generated from SHA256 hash of the execution ID + const traceId = traceContext?.toTraceId(); + expect(traceId).toBeDefined(); + expect(traceId).not.toBe("0"); // Should not be zero + expect(traceId).toMatch(/^\d+$/); // Should be numeric string + + // Verify the span ID is deterministic based on execution ID, state name, and entered time + const spanId = traceContext?.toSpanId(); + expect(spanId).toBeDefined(); + expect(spanId).not.toBe("0"); // Should not be zero + expect(spanId).toMatch(/^\d+$/); // Should be numeric string + + // Verify that extracting the same event produces the same trace IDs (deterministic) + const traceContext2 = await extractor.extract(event, {} as Context); + expect(traceContext2?.toTraceId()).toBe(traceId); + expect(traceContext2?.toSpanId()).toBe(spanId); + }); }); describe("lambda context", () => { @@ -734,301 +781,146 @@ describe("TraceContextExtractor", () => { expect(traceContext?.source).toBe("event"); }); }); - - describe("xray", () => { - // Fallback, event and context don't contain any trace context - it("extracts trace context from Xray", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const tracerWrapper = new TracerWrapper(); - const extractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - - const traceContext = await extractor.extract({}, {} as Context); - expect(traceContext).not.toBeNull(); - - expect(traceContext?.toTraceId()).toBe("4110911582297405557"); - expect(traceContext?.toSpanId()).toBe("797643193680388254"); - expect(traceContext?.sampleMode()).toBe("2"); - expect(traceContext?.source).toBe("xray"); - }); - }); }); - describe("getTraceEventExtractor", () => { - beforeEach(() => { - StepFunctionContextService["_instance"] = undefined as any; - }); - it.each([ - ["null", null], - ["undefined", undefined], - ["a string", "some-value"], - ["a number", 1234], - ["an object which doesn't match any expected event", { custom: "event" }], - ])("returns undefined when event is '%s'", (_, event) => { - const tracerWrapper = new TracerWrapper(); - const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + describe("xray", () => { + // Fallback, event and context don't contain any trace context + it("extracts trace context from Xray", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - const extractor = traceContextExtractor["getTraceEventExtractor"](event); - - expect(extractor).toBeUndefined(); - }); - - // Returns the expected event extractor when payload is from that event - it.each([ - [ - "HTTPEventTraceExtractor", - "headers", - HTTPEventTraceExtractor, - { - headers: {}, - }, - ], - [ - "SNSEventTraceExtractor", - "SNS event", - SNSEventTraceExtractor, - { - Records: [ - { - Sns: {}, - }, - ], - }, - ], - [ - "SNSSQSventTraceExtractor", - "SNS to SQS event", - SNSSQSEventTraceExtractor, - { - Records: [ - { - eventSource: "aws:sqs", - body: '{"Type":"Notification", "TopicArn":"some-topic"}', - }, - ], - }, - ], - [ - "EventBridgeSQSTraceExtractor", - "EventBridge to SQS event", - EventBridgeSQSEventTraceExtractor, - { - Records: [ - { - eventSource: "aws:sqs", - body: '{"detail-type":"some-detail-type"}', - }, - ], - }, - ], - [ - "AppSyncEventTraceExtractor", - "AppSync event", - AppSyncEventTraceExtractor, - { - info: { - selectionSetGraphQL: "some-selection", - }, - request: { - headers: {}, - }, - }, - ], - [ - "SQSEventTraceExtractor", - "SQS event", - SQSEventTraceExtractor, - { - Records: [ - { - eventSource: "aws:sqs", - }, - ], - }, - ], - [ - "KinesisEventTraceExtractor", - "Kinesis stream event", - KinesisEventTraceExtractor, - { - Records: [ - { - kinesis: {}, - }, - ], - }, - ], - [ - "EventBridgeEventTraceExtractor", - "EventBridge event", - EventBridgeEventTraceExtractor, - { - "detail-type": "some-detail-type", - }, - ], - ])("returns %s when event contains %s", (_, __, _class, event) => { const tracerWrapper = new TracerWrapper(); - const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + const extractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - const extractor = traceContextExtractor["getTraceEventExtractor"](event); + const traceContext = await extractor.extract({}, {} as Context); + expect(traceContext).not.toBeNull(); - expect(extractor).toBeInstanceOf(_class); + expect(traceContext?.toTraceId()).toBe("4110911582297405557"); + expect(traceContext?.toSpanId()).toBe("797643193680388254"); + expect(traceContext?.sampleMode()).toBe("2"); + expect(traceContext?.source).toBe("xray"); }); + }); +}); - it("returns StepFunctionEventTraceExtractor when event contains LegacyStepFunctionContext", () => { - 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 traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - - // Mimick TraceContextService.extract initialization - const instance = StepFunctionContextService.instance(legacyStepFunctionEvent); - traceContextExtractor["stepFunctionContextService"] = instance; - - const extractor = traceContextExtractor["getTraceEventExtractor"](legacyStepFunctionEvent); - - expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); - }); +describe("getTraceEventExtractor", () => { + beforeEach(() => { + StepFunctionContextService["_instance"] = undefined as any; + }); + it.each([ + ["null", null], + ["undefined", undefined], + ["a string", "some-value"], + ["a number", 1234], + ["an object which doesn't match any expected event", { custom: "event" }], + ])("returns undefined when event is '%s'", (_, event) => { + const tracerWrapper = new TracerWrapper(); + const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + + const extractor = traceContextExtractor["getTraceEventExtractor"](event); + + expect(extractor).toBeUndefined(); + }); - it("returns StepFunctionEventTraceExtractor when event contains 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", + // Returns the expected event extractor when payload is from that event + it.each([ + [ + "HTTPEventTraceExtractor", + "headers", + HTTPEventTraceExtractor, + { + headers: {}, + }, + ], + [ + "SNSEventTraceExtractor", + "SNS event", + SNSEventTraceExtractor, + { + Records: [ + { + Sns: {}, }, - State: { - Name: "step-one", - EnteredTime: "2022-12-08T21:08:19.224Z", - RetryCount: 2, + ], + }, + ], + [ + "SNSSQSventTraceExtractor", + "SNS to SQS event", + SNSSQSEventTraceExtractor, + { + Records: [ + { + eventSource: "aws:sqs", + body: '{"Type":"Notification", "TopicArn":"some-topic"}', }, - StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", + ], + }, + ], + [ + "EventBridgeSQSTraceExtractor", + "EventBridge to SQS event", + EventBridgeSQSEventTraceExtractor, + { + Records: [ + { + eventSource: "aws:sqs", + body: '{"detail-type":"some-detail-type"}', }, - "x-datadog-trace-id": "10593586103637578129", - "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000", - "serverless-version": "v1", + ], + }, + ], + [ + "AppSyncEventTraceExtractor", + "AppSync event", + AppSyncEventTraceExtractor, + { + info: { + selectionSetGraphQL: "some-selection", }, - }; - - const tracerWrapper = new TracerWrapper(); - const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - - // Mimick TraceContextService.extract initialization - const instance = StepFunctionContextService.instance(lambdaRootStepFunctionEvent); - traceContextExtractor["stepFunctionContextService"] = instance; - - const extractor = traceContextExtractor["getTraceEventExtractor"](lambdaRootStepFunctionEvent); - - expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); - }); - - it("returns StepFunctionEventTraceExtractor when event contains 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", + request: { + headers: {}, }, - }; - - const tracerWrapper = new TracerWrapper(); - const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - - // Mimick TraceContextService.extract initialization - const instance = StepFunctionContextService.instance(nestedStepFunctionEvent); - traceContextExtractor["stepFunctionContextService"] = instance; - - const extractor = traceContextExtractor["getTraceEventExtractor"](nestedStepFunctionEvent); - - expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); - }); - - it("returns StepFunctionEventTraceExtractor when event contains 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, + }, + ], + [ + "SQSEventTraceExtractor", + "SQS event", + SQSEventTraceExtractor, + { + Records: [ + { + eventSource: "aws:sqs", }, - StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", + ], + }, + ], + [ + "KinesisEventTraceExtractor", + "Kinesis stream event", + KinesisEventTraceExtractor, + { + Records: [ + { + kinesis: {}, }, - }, - }; - - const tracerWrapper = new TracerWrapper(); - const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - - // Mimick TraceContextService.extract initialization - const instance = StepFunctionContextService.instance(event); - traceContextExtractor["stepFunctionContextService"] = instance; + ], + }, + ], + [ + "EventBridgeEventTraceExtractor", + "EventBridge event", + EventBridgeEventTraceExtractor, + { + "detail-type": "some-detail-type", + }, + ], + ])("returns %s when event contains %s", (_, __, _class, event) => { + const tracerWrapper = new TracerWrapper(); + const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); - const extractor = traceContextExtractor["getTraceEventExtractor"](event); + const extractor = traceContextExtractor["getTraceEventExtractor"](event); - expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); - }); + expect(extractor).toBeInstanceOf(_class); }); describe("addTraceContextToXray", () => { diff --git a/src/trace/context/extractor.ts b/src/trace/context/extractor.ts index fb690835..35150a2c 100644 --- a/src/trace/context/extractor.ts +++ b/src/trace/context/extractor.ts @@ -43,8 +43,6 @@ export class TraceContextExtractor { } async extract(event: any, context: Context): Promise { - this.stepFunctionContextService = StepFunctionContextService.instance(event); - let spanContext: SpanContextWrapper | null = null; if (this.config.traceExtractor) { const customExtractor = new CustomTraceExtractor(this.config.traceExtractor); @@ -58,6 +56,14 @@ export class TraceContextExtractor { } } + if (spanContext === null) { + this.stepFunctionContextService = StepFunctionContextService.instance(event); + if (this.stepFunctionContextService?.context) { + const extractor = new StepFunctionEventTraceExtractor(); + spanContext = extractor?.extract(event); + } + } + if (spanContext === null) { const contextExtractor = new LambdaContextTraceExtractor(this.tracerWrapper); spanContext = contextExtractor.extract(context); @@ -88,8 +94,6 @@ export class TraceContextExtractor { if (EventValidator.isKinesisStreamEvent(event)) return new KinesisEventTraceExtractor(this.tracerWrapper); if (EventValidator.isEventBridgeEvent(event)) return new EventBridgeEventTraceExtractor(this.tracerWrapper); - if (this.stepFunctionContextService?.context) return new StepFunctionEventTraceExtractor(); - return; } diff --git a/src/trace/context/extractors/event-bridge-sqs.spec.ts b/src/trace/context/extractors/event-bridge-sqs.spec.ts index 70aecac5..24ca16fa 100644 --- a/src/trace/context/extractors/event-bridge-sqs.spec.ts +++ b/src/trace/context/extractors/event-bridge-sqs.spec.ts @@ -1,5 +1,6 @@ import { TracerWrapper } from "../../tracer-wrapper"; import { EventBridgeSQSEventTraceExtractor } from "./event-bridge-sqs"; +import { StepFunctionContextService } from "../../step-function-service"; let mockSpanContext: any = null; @@ -42,7 +43,7 @@ describe("EventBridgeSQSEventTraceExtractor", () => { messageId: "e995e54f-1724-41fa-82c0-8b81821f854e", receiptHandle: "AQEB4mIfRcyqtzn1X5Ss+ConhTejVGc+qnAcmu3/Z9ZvbNkaPcpuDLX/bzvPD/ZkAXJUXZcemGSJmd7L3snZHKMP2Ck8runZiyl4mubiLb444pZvdiNPuGRJ6a3FvgS/GQPzho/9nNMyOi66m8Viwh70v4EUCPGO4JmD3TTDAUrrcAnqU4WSObjfC/NAp9bI6wH2CEyAYEfex6Nxplbl/jBf9ZUG0I3m3vQd0Q4l4gd4jIR4oxQUglU2Tldl4Kx5fMUAhTRLAENri6HsY81avBkKd9FAuxONlsITB5uj02kOkvLlRGEcalqsKyPJ7AFaDLrOLaL3U+yReroPEJ5R5nwhLOEbeN5HROlZRXeaAwZOIN8BjqdeooYTIOrtvMEVb7a6OPLMdH1XB+ddevtKAH8K9Tm2ZjpaA7dtBGh1zFVHzBk=", - body: '{"version":"0","id":"af718b2a-b987-e8c0-7a2b-a188fad2661a","detail-type":"my.Detail","source":"my.Source","account":"425362996713","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"7379586022458917877","x-datadog-parent-id":"2644033662113726488","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}', + body: '{"version":"0","id":"af718b2a-b987-e8c0-7a2b-a188fad2661a","detail-type":"my.Detail","source":"my.Source","account":"123456123456","time":"2023-08-03T22:49:03Z","region":"us-east-1","resources":[],"detail":{"text":"Hello, world!","_datadog":{"x-datadog-trace-id":"7379586022458917877","x-datadog-parent-id":"2644033662113726488","x-datadog-sampling-priority":"1","x-datadog-tags":"_dd.p.dm=-0"}}}', attributes: { ApproximateReceiveCount: "1", AWSTraceHeader: "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0", @@ -124,5 +125,45 @@ describe("EventBridgeSQSEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).toBeNull(); }); + + it("extracts trace context from Step Function EventBridge-SQS event", () => { + // Reset StepFunctionContextService instance + StepFunctionContextService["_instance"] = undefined as any; + + const tracerWrapper = new TracerWrapper(); + + const payload = { + Records: [ + { + messageId: "0fc0e02f-ab25-4fde-b5ff-22aba9a9f20e", + receiptHandle: + "AQEBROlXUgqqRdo/j0GcfxBNldIKy8FO6Ee0ZnP5YeAp4pwQ+v9XovX47vSzNHAZooCa0r8D7Uoow0y4bhGiH/Tt5HXAseDUlvWHB6bULonzAdvRmLd1W1OCY9D1uH3TpHZYn6JdoQd6Koxndx5wDwhv5UKxcbOwDjlc6X/30OKkTm4gcr7Otzu4GxCt6N/FmDxcRIDogZk80UE1kN6Q5EHI9LB6V+oleqqCbQwg5FYmbVc+DjwPBY4/5NI6x1/XZLFZA0TdezOdOuNq4+4DGK8e35Bafg4hXp+06zg8E5XPdMQV5V4iDzJhenPEdXusGL36byBHyC4aDunTSpeIND0/0ctqyH1vEJHo09LJ1jztPj05hBQeDU5QXCIKRpuo5+nEHE+Jm1ZLrUWUoIg1uAIamDzQ0CWNtPjGkNn3POiTGpD2e0aqrE5VpXZ8N30HKKFM", + body: '{"version":"0","id":"1aec4f0c-35e7-934c-9928-a5db3e526bca","detail-type":"ProcessEvent","source":"demo.stepfunction","account":"123456123456","time":"2025-07-14T19:33:03Z","region":"sa-east-1","resources":["arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sqs-demo-dev-state-machine","arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d"],"detail":{"message":"Event from Step Functions","timestamp":"2025-07-14T19:33:03.483Z","executionName":"aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","stateMachineName":"rstrat-sfn-evb-sqs-demo-dev-state-machine","input":{"testData":"Hello with SQS integration"},"_datadog":{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","StartTime":"2025-07-14T19:33:03.446Z","Name":"aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-evb-sqs-demo-d-StepFunctionsExecutionRol-mAumgN9x07FQ","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-evb-sqs-demo-dev-state-machine","Name":"rstrat-sfn-evb-sqs-demo-dev-state-machine"},"State":{"Name":"PublishToEventBridge","EnteredTime":"2025-07-14T19:33:03.483Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-evb-sqs-demo-dev-state-machine:aa2f4ded-196f-4d6b-b41d-aa64f3193f2d","serverless-version":"v1"}}}', + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1752521583745", + SenderId: "AIDAIELDKKY42PBA6I2NG", + ApproximateFirstReceiveTimestamp: "1752521583758", + }, + messageAttributes: {}, + md5OfBody: "957cded00b7b10a6e1b79864f24a7b5f", + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-evb-sqs-demo-dev-process-event-queue", + awsRegion: "sa-east-1", + }, + ], + }; + + const extractor = new EventBridgeSQSEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // The trace IDs are deterministically generated from the Step Function execution context + expect(traceContext?.toTraceId()).toBe("7858567057595668526"); + expect(traceContext?.toSpanId()).toBe("3674709292670593712"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/event-bridge-sqs.ts b/src/trace/context/extractors/event-bridge-sqs.ts index 1fd93201..45428478 100644 --- a/src/trace/context/extractors/event-bridge-sqs.ts +++ b/src/trace/context/extractors/event-bridge-sqs.ts @@ -1,30 +1,24 @@ -import { EventBridgeEvent, SQSEvent } from "aws-lambda"; -import { EventTraceExtractor } from "../extractor"; -import { logDebug } from "../../../utils"; +import { SQSEvent } from "aws-lambda"; import { TracerWrapper } from "../../tracer-wrapper"; +import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; +import { extractTraceContext, handleExtractionError } from "../extractor-utils"; export class EventBridgeSQSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} extract(event: SQSEvent): SpanContextWrapper | null { - const body = event?.Records?.[0]?.body; - if (body === undefined) return null; - try { - const parsedBody = JSON.parse(body) as EventBridgeEvent; - const headers = parsedBody?.detail?._datadog; - if (headers === undefined) return null; - - const traceContext = this.tracerWrapper.extract(headers); - if (traceContext === null) return null; - - logDebug("Extracted trace context from EventBridge-SQS event", { traceContext, event }); - return traceContext; - } catch (error) { - if (error instanceof Error) { - logDebug("Unable to extract trace context from EventBridge-SQS event", error); + const body = event?.Records?.[0]?.body; + if (body) { + const parsedBody = JSON.parse(body); + const headers = parsedBody?.detail?._datadog; + if (headers) { + return extractTraceContext(headers, this.tracerWrapper); + } } + } catch (error) { + handleExtractionError(error, "EventBridge-SQS"); } return null; diff --git a/src/trace/context/extractors/event-bridge.spec.ts b/src/trace/context/extractors/event-bridge.spec.ts index 27e10306..346b89c4 100644 --- a/src/trace/context/extractors/event-bridge.spec.ts +++ b/src/trace/context/extractors/event-bridge.spec.ts @@ -1,5 +1,6 @@ import { TracerWrapper } from "../../tracer-wrapper"; import { EventBridgeEventTraceExtractor } from "./event-bridge"; +import { StepFunctionContextService } from "../../step-function-service"; let mockSpanContext: any = null; @@ -106,5 +107,66 @@ describe("EventBridgeEventTraceExtractor", () => { const traceContext = extractor.extract(payload); expect(traceContext).toBeNull(); }); + + it("extracts trace context from Step Function EventBridge event", () => { + // Reset StepFunctionContextService instance + StepFunctionContextService["_instance"] = undefined as any; + + const tracerWrapper = new TracerWrapper(); + + const payload = { + version: "0", + id: "af718b2a-b987-e8c0-7a2b-a188fad2661a", + "detail-type": "ProcessEvent", + source: "demo.stepfunction", + account: "123456123456", + time: "2025-07-11T14:59:35Z", + region: "us-east-2", + resources: [ + "arn:aws:states:us-east-2:123456123456:stateMachine:rstrat-sfn-evb-demo-dev-state-machine", + "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105", + ], + detail: { + message: "Event from Step Functions", + timestamp: "2025-07-11T14:59:35.830Z", + executionName: "6c190e7b-eb77-46db-af26-9066d353b105", + stateMachineName: "rstrat-sfn-evb-demo-dev-state-machine", + input: { + testData: "Hello with Datadog tracing", + }, + _datadog: { + Execution: { + Id: "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105", + StartTime: "2025-07-11T14:59:35.806Z", + Name: "6c190e7b-eb77-46db-af26-9066d353b105", + RoleArn: "arn:aws:iam::123456123456:role/rstrat-sfn-evb-demo-dev-StepFunctionsExecutionRole-8maJHu01fhZZ", + RedriveCount: 0, + }, + StateMachine: { + Id: "arn:aws:states:us-east-2:123456123456:stateMachine:rstrat-sfn-evb-demo-dev-state-machine", + Name: "rstrat-sfn-evb-demo-dev-state-machine", + }, + State: { + Name: "PublishToEventBridge", + EnteredTime: "2025-07-11T14:59:35.830Z", + RetryCount: 0, + }, + RootExecutionId: + "arn:aws:states:us-east-2:123456123456:execution:rstrat-sfn-evb-demo-dev-state-machine:6c190e7b-eb77-46db-af26-9066d353b105", + "serverless-version": "v1", + }, + }, + }; + + const extractor = new EventBridgeEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + expect(traceContext?.toTraceId()).toBe("1503104665848096006"); + expect(traceContext?.toSpanId()).toBe("159267866761498620"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/event-bridge.ts b/src/trace/context/extractors/event-bridge.ts index c589c929..41086f6a 100644 --- a/src/trace/context/extractors/event-bridge.ts +++ b/src/trace/context/extractors/event-bridge.ts @@ -1,26 +1,20 @@ -import { EventTraceExtractor } from "../extractor"; import { EventBridgeEvent } from "aws-lambda"; -import { logDebug } from "../../../utils"; import { TracerWrapper } from "../../tracer-wrapper"; +import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; +import { extractTraceContext, handleExtractionError } from "../extractor-utils"; export class EventBridgeEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} - extract(event: EventBridgeEvent): SpanContextWrapper | null { - const headers = event?.detail?._datadog; - if (headers === undefined) return null; - + extract(event: EventBridgeEvent): SpanContextWrapper | null { try { - const traceContext = this.tracerWrapper.extract(headers); - if (traceContext === null) return null; - - logDebug(`Extracted trace context from Eventbridge event`, { traceContext, event }); - return traceContext; - } catch (error) { - if (error instanceof Error) { - logDebug("Unable to extract trace context from EventBridge event", error); + const headers = event?.detail?._datadog; + if (headers) { + return extractTraceContext(headers, this.tracerWrapper); } + } catch (error) { + handleExtractionError(error, "EventBridge"); } return null; diff --git a/src/trace/context/extractors/sns-sqs.spec.ts b/src/trace/context/extractors/sns-sqs.spec.ts index 33d673eb..82b39c32 100644 --- a/src/trace/context/extractors/sns-sqs.spec.ts +++ b/src/trace/context/extractors/sns-sqs.spec.ts @@ -1,5 +1,6 @@ import { TracerWrapper } from "../../tracer-wrapper"; import { SNSSQSEventTraceExtractor } from "./sns-sqs"; +import { StepFunctionContextService } from "../../step-function-service"; let mockSpanContext: any = null; @@ -212,5 +213,54 @@ describe("SNSSQSEventTraceExtractor", () => { expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); + + it("extracts trace context from Step Function SNS-SQS event", () => { + // Reset StepFunctionContextService instance + StepFunctionContextService["_instance"] = undefined as any; + + const tracerWrapper = new TracerWrapper(); + + const payload = { + Records: [ + { + messageId: "43a5f138-f166-40f1-b7e4-a7e0af9d633d", + receiptHandle: + "AQEBrhzl4RiITHp/ui07Y0DlDIdrmYHveKjqDIsx2gG7Z3fvrDohnfnpy/r4esh/ZsilJUR/C3uohYe6HUqvixymhx+io9S/MYNoA1zjmSVd1V4ZKe6saMs6L7aSW5TgrLpuxOtNGWvmNijdlQlOoYW1xRlkzkBywFkELfazExJHrThbxpxXcHbcAoh1Vz77EvlcAQNbc11vTccoUcMcdczvoLd/wgyrsIf0z8qdUQHaspWYoWOlZOsoflDMddYwqWO3LNRphAGMp5ISTDVbVqo1/U+wOqBj+b3dOYP9k0vS9Mj+36t+EJ8+KETFXRPNk4mZ+7hvG+UCYBN582gT502MnQitxylHKWOlH77nIokfk43FjhjsybLE48KdWdO49O2WKslXwCpPLQWnbKWlUl05/12tIk41MolVyfiWywW9R/S7hgcSr51tEBcjZTW8GR8r", + body: '{"testData":"Hello from SQS integration test","timestamp":"2025-07-15T18:16:27Z"}', + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1752603390964", + SenderId: "AIDAIOA2GYWSHW4E2VXIO", + ApproximateFirstReceiveTimestamp: "1752603390980", + }, + messageAttributes: { + _datadog: { + stringValue: + '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-sqs-demo-dev-state-machine:c363b975-c342-4e40-815a-8dd2496f5e81","StartTime":"2025-07-15T18:16:30.746Z","Name":"c363b975-c342-4e40-815a-8dd2496f5e81","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sns-sqs-demo-d-StepFunctionsExecutionRol-T2O3igeuSihu","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sns-sqs-demo-dev-state-machine","Name":"rstrat-sfn-sns-sqs-demo-dev-state-machine"},"State":{"Name":"PublishToSNS","EnteredTime":"2025-07-15T18:16:30.776Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-sqs-demo-dev-state-machine:c363b975-c342-4e40-815a-8dd2496f5e81","serverless-version":"v1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + md5OfBody: "1e832c0d0aa5188dc5e3f2e85c9cb5e7", + md5OfMessageAttributes: "64e36d01aec95ca5a2160c13299e9c3b", + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sns-sqs-demo-dev-process-event-queue", + awsRegion: "sa-east-1", + }, + ], + }; + + const extractor = new SNSSQSEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // The StepFunctionContextService generates deterministic trace IDs + expect(traceContext?.toTraceId()).toBe("1657966791618574655"); + expect(traceContext?.toSpanId()).toBe("5100002956473485303"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/sns-sqs.ts b/src/trace/context/extractors/sns-sqs.ts index e68ee6a2..fd23ad9e 100644 --- a/src/trace/context/extractors/sns-sqs.ts +++ b/src/trace/context/extractors/sns-sqs.ts @@ -1,52 +1,57 @@ -import { SNSMessage, SQSEvent } from "aws-lambda"; -import { EventTraceExtractor } from "../extractor"; +import { SQSEvent } from "aws-lambda"; import { TracerWrapper } from "../../tracer-wrapper"; import { logDebug } from "../../../utils"; +import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; -import { XrayService } from "../../xray-service"; +import { extractTraceContext, extractFromAWSTraceHeader, handleExtractionError } from "../extractor-utils"; export class SNSSQSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} extract(event: SQSEvent): SpanContextWrapper | null { + logDebug("SNS-SQS Extractor Being Used"); try { - // First try to extract trace context from message attributes - if (event?.Records?.[0]?.body) { - const parsedBody = JSON.parse(event?.Records?.[0]?.body) as SNSMessage; - const messageAttribute = parsedBody?.MessageAttributes?._datadog; - if (messageAttribute?.Value) { + // Try to extract trace context from SNS wrapped in SQS + const body = event?.Records?.[0]?.body; + if (body) { + const parsedBody = JSON.parse(body); + const snsMessageAttribute = parsedBody?.MessageAttributes?._datadog; + if (snsMessageAttribute?.Value) { let headers; - if (messageAttribute.Type === "String") { - headers = JSON.parse(messageAttribute.Value); + if (snsMessageAttribute.Type === "String") { + headers = JSON.parse(snsMessageAttribute.Value); } else { - const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii"); + // Try decoding base64 values + const decodedValue = Buffer.from(snsMessageAttribute.Value, "base64").toString("ascii"); headers = JSON.parse(decodedValue); } - const traceContext = this.tracerWrapper.extract(headers); + const traceContext = extractTraceContext(headers, this.tracerWrapper); if (traceContext) { - logDebug("Extracted trace context from SNS-SQS event"); return traceContext; - } else { - logDebug("Failed to extract trace context from SNS-SQS event"); } + logDebug("Failed to extract trace context from SNS-SQS event"); } } - // Then try to extract trace context from attributes.AWSTraceHeader. (Upstream Java apps can - // pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) - if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) { - const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader); + + // Check SQS message attributes as a fallback + const sqsMessageAttribute = event?.Records?.[0]?.messageAttributes?._datadog; + if (sqsMessageAttribute?.stringValue) { + const headers = JSON.parse(sqsMessageAttribute.stringValue); + const traceContext = extractTraceContext(headers, this.tracerWrapper); if (traceContext) { - logDebug("Extracted trace context from SNS-SQS event attributes.AWSTraceHeader"); return traceContext; - } else { - logDebug("No Datadog trace context found from SNS-SQS event attributes.AWSTraceHeader"); } } - } catch (error) { - if (error instanceof Error) { - logDebug("Unable to extract trace context from SNS-SQS event", error); + + // Else try to extract trace context from attributes.AWSTraceHeader + // (Upstream Java apps can pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) + const awsTraceHeader = event?.Records?.[0]?.attributes?.AWSTraceHeader; + if (awsTraceHeader !== undefined) { + return extractFromAWSTraceHeader(awsTraceHeader, "SNS-SQS"); } + } catch (error) { + handleExtractionError(error, "SQS"); } return null; diff --git a/src/trace/context/extractors/sns.spec.ts b/src/trace/context/extractors/sns.spec.ts index aadc29f6..a3ef1159 100644 --- a/src/trace/context/extractors/sns.spec.ts +++ b/src/trace/context/extractors/sns.spec.ts @@ -1,6 +1,7 @@ import { SNSEvent } from "aws-lambda"; import { TracerWrapper } from "../../tracer-wrapper"; import { SNSEventTraceExtractor } from "./sns"; +import { StepFunctionContextService } from "../../step-function-service"; let mockSpanContext: any = null; @@ -256,5 +257,57 @@ describe("SNSEventTraceExtractor", () => { expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); + + it("extracts trace context from Step Function SNS event", () => { + // Reset StepFunctionContextService instance + StepFunctionContextService["_instance"] = undefined as any; + + const tracerWrapper = new TracerWrapper(); + + const payload: SNSEvent = { + Records: [ + { + EventSource: "aws:sns", + EventVersion: "1.0", + EventSubscriptionArn: + "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic:f18241f8-a4f7-4586-80db-97bd1939a557", + Sns: { + Type: "Notification", + MessageId: "46d2665c-7ee2-50ba-a4cd-06acf35f5d5f", + TopicArn: "arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic", + Message: + '{"source":"demo.stepfunction","detailType":"ProcessEvent","message":"Test event from Step Functions","customData":{"userId":"12345","action":"test"}}', + Timestamp: "2025-07-15T17:10:21.503Z", + SignatureVersion: "1", + Signature: + "fHeJta0GWCs/lHhI6wesXiT+66i1SZ+XH58pyd8mKHKD8bepXsnWvfQdDsOkO2AVv2CqPBF58sAWQae6yob2aMawe/vo8eeahJCaguK8a/3HLj7kP+nXGjgSGvzQm4CdYEyAUco453/mfE/BSf0SkdctxW0rjMs27T2l964Lt2Y/vJeiXVibs/AqEIu3ImekbM8+EIfNMOLBdRBVE47650vawazMkcpPtg5o/8LCA/jNUNj9VCTJrvzep8/vVJEcuHbZ3pcmajA9UJmP3000G0+to0cXwZ5YaakOxQTv81I+cfC99yQJoogLklbgiu+4bqEeNWbwW9KdQz1U+79NgA==", + SigningCertUrl: + "https://sns.sa-east-1.amazonaws.com/SimpleNotificationService-9c6465fa7f48f5cacd23014631ec1136.pem", + Subject: "Event from Step Functions", + UnsubscribeUrl: + "https://sns.sa-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:sa-east-1:123456123456:rstrat-sfn-sns-demo-dev-process-event-topic:f18241f8-a4f7-4586-80db-97bd1939a557", + MessageAttributes: { + _datadog: { + Type: "String", + Value: + '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-demo-dev-state-machine:79049e80-5cc6-49da-9dc0-f19ba2921772","StartTime":"2025-07-15T17:10:21.328Z","Name":"79049e80-5cc6-49da-9dc0-f19ba2921772","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sns-demo-dev-StepFunctionsExecutionRole-LrsdDm6wMmBh","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sns-demo-dev-state-machine","Name":"rstrat-sfn-sns-demo-dev-state-machine"},"State":{"Name":"PublishToSNS","EnteredTime":"2025-07-15T17:10:21.354Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sns-demo-dev-state-machine:79049e80-5cc6-49da-9dc0-f19ba2921772","serverless-version":"v1"}', + }, + }, + }, + }, + ], + }; + + const extractor = new SNSEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // The StepFunctionContextService generates deterministic trace IDs + expect(traceContext?.toTraceId()).toBe("3995810302240690842"); + expect(traceContext?.toSpanId()).toBe("8347071195300897803"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/sns.ts b/src/trace/context/extractors/sns.ts index 33a009db..7336f552 100644 --- a/src/trace/context/extractors/sns.ts +++ b/src/trace/context/extractors/sns.ts @@ -3,7 +3,8 @@ import { TracerWrapper } from "../../tracer-wrapper"; import { logDebug } from "../../../utils"; import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; -import { XrayService, AMZN_TRACE_ID_ENV_VAR } from "../../xray-service"; +import { AMZN_TRACE_ID_ENV_VAR } from "../../xray-service"; +import { extractTraceContext, extractFromAWSTraceHeader, handleExtractionError } from "../extractor-utils"; export class SNSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} @@ -22,29 +23,20 @@ export class SNSEventTraceExtractor implements EventTraceExtractor { headers = JSON.parse(decodedValue); } - const traceContext = this.tracerWrapper.extract(headers); + const traceContext = extractTraceContext(headers, this.tracerWrapper); if (traceContext) { - logDebug("Extracted trace context from SNS event"); return traceContext; - } else { - logDebug("Failed to extract trace context from SNS event"); } + logDebug("Failed to extract trace context from SNS event"); } + // Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can // pass down Datadog trace id (parent id wrong) in the env in SNS case) if (process.env[AMZN_TRACE_ID_ENV_VAR]) { - const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR]); - if (traceContext) { - logDebug("Extracted Datadog trace context from _X_AMZN_TRACE_ID"); - return traceContext; - } else { - logDebug("No Datadog trace context found from _X_AMZN_TRACE_ID"); - } + return extractFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR], "SNS"); } } catch (error) { - if (error instanceof Error) { - logDebug("Unable to extract trace context from SNS event", error); - } + handleExtractionError(error, "SNS"); } return null; diff --git a/src/trace/context/extractors/sqs.spec.ts b/src/trace/context/extractors/sqs.spec.ts index 9c820a7e..4a1d41cd 100644 --- a/src/trace/context/extractors/sqs.spec.ts +++ b/src/trace/context/extractors/sqs.spec.ts @@ -1,6 +1,7 @@ import { SQSEvent } from "aws-lambda"; import { TracerWrapper } from "../../tracer-wrapper"; import { SQSEventTraceExtractor } from "./sqs"; +import { StepFunctionContextService } from "../../step-function-service"; let mockSpanContext: any = null; @@ -223,5 +224,54 @@ describe("SQSEventTraceExtractor", () => { expect(traceContext?.sampleMode()).toBe("1"); expect(traceContext?.source).toBe("event"); }); + + it("extracts trace context from Step Function SQS event", () => { + // Reset StepFunctionContextService instance + StepFunctionContextService["_instance"] = undefined as any; + + const tracerWrapper = new TracerWrapper(); + + const payload: SQSEvent = { + Records: [ + { + messageId: "4ead33f3-51c8-4094-87bd-5325dc143cbd", + receiptHandle: + "AQEBrGtLZCUS1POUEZtdZRoB0zXgT14OQC48A4Xk4Qbnv/v4d0ib5rFI1wEah823t2hE9haPm6nNN1aGsJmYkqa9Y8qaBQscp9f7HKJyybT5hpdKEn07fY0VRv/Of63u1RN1YdFdY5uhI8XGWRc4w7t62lQwMMFY5Ahy7XLVwnav81KRjGFdgxzITrtx3YKxmISNvXzPiiHNKb7jT+ClfXi91bEYHi3Od3ji5xGajAofgYrj2VBDULyohsfMkwlvAanD2wfj2x++wL5LSpFEtMFnvThzt7Dh5FEZChVMzWV+fRFpljivHX58ZeuGv4yIIjLVuuDGn5uAY5ES4CsdINrBAru6K5gDSPUajRzE3TktNgAq5Niqfky1x0srLRAJjTDdmZK8/CXU0sRT/MCT99vkCHa0bC17S/9au5bCbrB4k/T9J8W39AA6kIYhebkq3IQr", + body: '{"testData":"Hello from Step Functions to SQS"}', + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1752594520503", + SenderId: "AROAWGCM4HXU73A4V34AJ:EcGTcmgJbwwOwXPbloVwgSaDOmwhYBLH", + ApproximateFirstReceiveTimestamp: "1752594520516", + }, + messageAttributes: { + _datadog: { + stringValue: + '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","StartTime":"2025-07-15T15:48:40.302Z","Name":"a4912895-93a3-4803-a712-69fecb55c025","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sqs-demo-dev-StepFunctionsExecutionRole-s6ozc2dVrvLH","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sqs-demo-dev-state-machine","Name":"rstrat-sfn-sqs-demo-dev-state-machine"},"State":{"Name":"SendToSQS","EnteredTime":"2025-07-15T15:48:40.333Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","serverless-version":"v1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + md5OfMessageAttributes: "5469b8f90bb6ab27e95816c1fa178680", + md5OfBody: "f0c0ddb2ed09a09e8791013f142e8d7e", + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sqs-demo-dev-process-event-queue", + awsRegion: "sa-east-1", + }, + ], + }; + + const extractor = new SQSEventTraceExtractor(tracerWrapper); + + const traceContext = extractor.extract(payload); + expect(traceContext).not.toBeNull(); + + // The StepFunctionContextService generates deterministic trace IDs + expect(traceContext?.toTraceId()).toBe("7148114900282288397"); + expect(traceContext?.toSpanId()).toBe("6711327198021343353"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); }); }); diff --git a/src/trace/context/extractors/sqs.ts b/src/trace/context/extractors/sqs.ts index 01d1472f..f4b883e4 100644 --- a/src/trace/context/extractors/sqs.ts +++ b/src/trace/context/extractors/sqs.ts @@ -1,14 +1,15 @@ import { SQSEvent } from "aws-lambda"; -import { EventTraceExtractor } from "../extractor"; import { TracerWrapper } from "../../tracer-wrapper"; import { logDebug } from "../../../utils"; +import { EventTraceExtractor } from "../extractor"; import { SpanContextWrapper } from "../../span-context-wrapper"; -import { XrayService } from "../../xray-service"; +import { extractTraceContext, extractFromAWSTraceHeader, handleExtractionError } from "../extractor-utils"; export class SQSEventTraceExtractor implements EventTraceExtractor { constructor(private tracerWrapper: TracerWrapper) {} extract(event: SQSEvent): SpanContextWrapper | null { + logDebug("SQS Extractor Being Used"); try { // First try to extract trace context from message attributes let headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue; @@ -23,30 +24,24 @@ export class SQSEventTraceExtractor implements EventTraceExtractor { } } - if (headers !== undefined) { - const traceContext = this.tracerWrapper.extract(JSON.parse(headers)); + if (headers) { + const parsedHeaders = JSON.parse(headers); + + const traceContext = extractTraceContext(parsedHeaders, this.tracerWrapper); if (traceContext) { - logDebug("Extracted trace context from SQS event messageAttributes"); return traceContext; - } else { - logDebug("Failed to extract trace context from messageAttributes"); } + logDebug("Failed to extract trace context from SQS event"); } - // Then try to extract trace context from attributes.AWSTraceHeader. (Upstream Java apps can - // pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) - if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) { - const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader); - if (traceContext) { - logDebug("Extracted trace context from SQS event attributes AWSTraceHeader"); - return traceContext; - } else { - logDebug("No Datadog trace context found from SQS event attributes AWSTraceHeader"); - } + + // Else try to extract trace context from attributes.AWSTraceHeader + // (Upstream Java apps can pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case) + const awsTraceHeader = event?.Records?.[0]?.attributes?.AWSTraceHeader; + if (awsTraceHeader !== undefined) { + return extractFromAWSTraceHeader(awsTraceHeader, "SQS"); } } catch (error) { - if (error instanceof Error) { - logDebug("Unable to extract trace context from SQS event", error); - } + handleExtractionError(error, "SQS"); } return null; diff --git a/src/trace/context/extractors/step-function.spec.ts b/src/trace/context/extractors/step-function.spec.ts index 1465ab61..6007daeb 100644 --- a/src/trace/context/extractors/step-function.spec.ts +++ b/src/trace/context/extractors/step-function.spec.ts @@ -5,6 +5,7 @@ describe("StepFunctionEventTraceExtractor", () => { beforeEach(() => { StepFunctionContextService["_instance"] = undefined as any; }); + describe("extract", () => { const payload = { Execution: { @@ -112,5 +113,77 @@ describe("StepFunctionEventTraceExtractor", () => { const traceContext = extractor.extract({}); expect(traceContext).toBeNull(); }); + + it("extracts trace context end-to-end with deterministic IDs for v1 nested context", () => { + const nestedPayload = { + _datadog: { + Execution: { + Id: "arn:aws:states:us-east-1:123456789012:execution:parent-machine:parent-execution-id", + Name: "parent-execution-id", + StartTime: "2024-01-01T00:00:00.000Z", + }, + State: { + Name: "InvokeNestedStateMachine", + EnteredTime: "2024-01-01T00:00:05.000Z", + RetryCount: 0, + }, + "serverless-version": "v1", + RootExecutionId: "arn:aws:states:us-east-1:123456789012:execution:root-machine:root-execution-id", + }, + }; + + const extractor = new StepFunctionEventTraceExtractor(); + const traceContext = extractor.extract(nestedPayload); + + expect(traceContext).not.toBeNull(); + expect(traceContext?.source).toBe("event"); + + // Verify trace ID is based on root execution ID (deterministic) + const traceId = traceContext?.toTraceId(); + expect(traceId).toBeDefined(); + expect(traceId).not.toBe("0"); + + // Extract again to verify deterministic behavior + const traceContext2 = extractor.extract(nestedPayload); + expect(traceContext2?.toTraceId()).toBe(traceId); + }); + + it("extracts trace context end-to-end with propagated IDs for v1 lambda root context", () => { + const lambdaRootPayload = { + _datadog: { + Execution: { + Id: "arn:aws:states:us-east-1:123456789012:execution:machine:execution-id", + Name: "execution-id", + StartTime: "2024-01-01T00:00:00.000Z", + }, + State: { + Name: "ProcessLambda", + EnteredTime: "2024-01-01T00:00:02.000Z", + RetryCount: 0, + }, + "serverless-version": "v1", + "x-datadog-trace-id": "1234567890123456789", + "x-datadog-tags": "_dd.p.tid=fedcba9876543210,_dd.p.dm=-0", + }, + }; + + const extractor = new StepFunctionEventTraceExtractor(); + const traceContext = extractor.extract(lambdaRootPayload); + + expect(traceContext).not.toBeNull(); + expect(traceContext?.source).toBe("event"); + + // Verify trace ID comes from propagated headers, not deterministic hash + expect(traceContext?.toTraceId()).toBe("1234567890123456789"); + + // Verify deterministic span ID based on execution details + const spanId = traceContext?.toSpanId(); + expect(spanId).toBeDefined(); + expect(spanId).not.toBe("0"); + + // Extract again to verify deterministic span ID + const traceContext2 = extractor.extract(lambdaRootPayload); + expect(traceContext2?.toSpanId()).toBe(spanId); + }); }); }); diff --git a/src/trace/context/extractors/step-function.ts b/src/trace/context/extractors/step-function.ts index 6197e5c1..3fb0b0f9 100644 --- a/src/trace/context/extractors/step-function.ts +++ b/src/trace/context/extractors/step-function.ts @@ -5,11 +5,16 @@ import { EventTraceExtractor } from "../extractor"; export class StepFunctionEventTraceExtractor implements EventTraceExtractor { extract(event: any): SpanContextWrapper | null { // Probably StepFunctionContextService hasn't been called - const instance = StepFunctionContextService.instance(event); - const context = instance.context; + const stepFunctionInstance = StepFunctionContextService.instance(event); + const stepFunctionContext = stepFunctionInstance.context; - if (context === undefined) return null; + if (stepFunctionContext !== undefined) { + const spanContext = stepFunctionInstance.spanContext; + if (spanContext !== null) { + return spanContext; + } + } - return instance.spanContext; + return null; } } diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 2b3f538d..5b27fb6c 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -33,6 +33,15 @@ jest.mock("./tracer-wrapper", () => { return mockExtract(event); } + startSpan(name: any, options: any): any { + return { + toSpanId: () => "mockSpanId", + toTraceId: () => "mockTraceId", + finish: jest.fn(), + setTag: jest.fn(), + }; + } + injectSpan(span: any): any { return { [DATADOG_PARENT_ID_HEADER]: span.toSpanId(), @@ -312,6 +321,88 @@ describe("TraceListener", () => { ); }); + it("wraps dd-trace span around invocation with Step Function context", async () => { + const listener = new TraceListener(defaultConfig); + mockTraceSource = TraceSource.Event; + + // Mock Step Function context with deterministic trace IDs + mockSpanContext = { + toTraceId: () => "512d06a10e5e34cb", // Hex converted to decimal would be different + toSpanId: () => "7069a031ef9ad2cc", + _sampling: { + priority: "1", + }, + }; + mockSpanContextWrapper = { + spanContext: mockSpanContext, + }; + + const stepFunctionSQSEvent = { + Records: [ + { + messageId: "4ead33f3-51c8-4094-87bd-5325dc143cbd", + receiptHandle: + "AQEBrGtLZCUS1POUEZtdZRoB0zXgT14OQC48A4Xk4Qbnv/v4d0ib5rFI1wEah823t2hE9haPm6nNN1aGsJmYkqa9Y8qaBQscp9f7HKJyybT5hpdKEn07fY0VRv/Of63u1RN1YdFdY5uhI8XGWRc4w7t62lQwMMFY5Ahy7XLVwnav81KRjGFdgxzITrtx3YKxmISNvXzPiiHNKb7jT+ClfXi91bEYHi3Od3ji5xGajAofgYrj2VBDULyohsfMkwlvAanD2wfj2x++wL5LSpFEtMFnvThzt7Dh5FEZChVMzWV+fRFpljivHX58ZeuGv4yIIjLVuuDGn5uAY5ES4CsdINrBAru6K5gDSPUajRzE3TktNgAq5Niqfky1x0srLRAJjTDdmZK8/CXU0sRT/MCT99vkCHa0bC17S/9au5bCbrB4k/T9J8W39AA6kIYhebkq3IQr", + body: '{"testData":"Hello from Step Functions to SQS"}', + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1752594520503", + SenderId: "AROAWGCM4HXU73A4V34AJ:EcGTcmgJbwwOwXPbloVwgSaDOmwhYBLH", + ApproximateFirstReceiveTimestamp: "1752594520516", + }, + messageAttributes: { + _datadog: { + stringValue: + '{"Execution":{"Id":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","StartTime":"2025-07-15T15:48:40.302Z","Name":"a4912895-93a3-4803-a712-69fecb55c025","RoleArn":"arn:aws:iam::123456123456:role/rstrat-sfn-sqs-demo-dev-StepFunctionsExecutionRole-s6ozc2dVrvLH","RedriveCount":0},"StateMachine":{"Id":"arn:aws:states:sa-east-1:123456123456:stateMachine:rstrat-sfn-sqs-demo-dev-state-machine","Name":"rstrat-sfn-sqs-demo-dev-state-machine"},"State":{"Name":"SendToSQS","EnteredTime":"2025-07-15T15:48:40.333Z","RetryCount":0},"RootExecutionId":"arn:aws:states:sa-east-1:123456123456:execution:rstrat-sfn-sqs-demo-dev-state-machine:a4912895-93a3-4803-a712-69fecb55c025","serverless-version":"v1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + md5OfMessageAttributes: "5469b8f90bb6ab27e95816c1fa178680", + md5OfBody: "f0c0ddb2ed09a09e8791013f142e8d7e", + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sqs-demo-dev-process-event-queue", + awsRegion: "sa-east-1", + }, + ], + }; + + await listener.onStartInvocation(stepFunctionSQSEvent, context as any); + const unwrappedFunc = () => {}; + const wrappedFunc = listener.onWrap(unwrappedFunc); + wrappedFunc(); + await listener.onCompleteInvocation(); + + expect(mockWrap).toHaveBeenCalledWith( + "aws.lambda", + { + resource: "my-Lambda", + service: "my-Lambda", + tags: { + cold_start: "true", + function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda", + function_version: "$LATEST", + request_id: "1234", + resource_names: "my-Lambda", + functionname: "my-lambda", + "_dd.parent_source": "event", + "function_trigger.event_source": "sqs", + "function_trigger.event_source_arn": + "arn:aws:sqs:sa-east-1:123456123456:rstrat-sfn-sqs-demo-dev-process-event-queue", + datadog_lambda: datadogLambdaVersion, + dd_trace: ddtraceVersion, + }, + type: "serverless", + childOf: expect.objectContaining({ + toSpanId: expect.any(Function), + toTraceId: expect.any(Function), + }), + }, + unwrappedFunc, + ); + }); + it("injects authorizer context if it exists", async () => { const listener = new TraceListener(defaultConfig); mockTraceSource = TraceSource.Event; diff --git a/src/trace/tracer-wrapper.ts b/src/trace/tracer-wrapper.ts index 78cb6b49..97097b6e 100644 --- a/src/trace/tracer-wrapper.ts +++ b/src/trace/tracer-wrapper.ts @@ -76,6 +76,7 @@ export class TracerWrapper { public startSpan any>(name: string, options: TraceOptions): T | null { if (!this.isTracerAvailable) { + logDebug("No Tracer available, cannot start span"); return null; } return this.tracer.startSpan(name, options);