Skip to content

Commit 276a115

Browse files
ryanstratryanstrat
committed
Refactor common extraction logic
1 parent 71b9447 commit 276a115

6 files changed

Lines changed: 137 additions & 196 deletions

File tree

‎src/trace/context/extractor-utils.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import { logDebug } from "../../utils";
2+
import { StepFunctionContextService } from "../step-function-service";
3+
import { SpanContextWrapper } from "../span-context-wrapper";
4+
import { TracerWrapper } from "../tracer-wrapper";
5+
import { XrayService } from "../xray-service";
6+
7+
/**
8+
* Common utility functions for trace context extraction
9+
*/
10+
11+
/**
12+
* Attempts to extract trace context from headers, falling back to Step Function context if needed
13+
* @param headers The headers object to extract from
14+
* @param tracerWrapper The tracer wrapper instance
15+
* @param eventType The type of event (for logging)
16+
* @returns SpanContextWrapper or null
17+
*/
18+
export function extractTraceContext(headers: any, tracerWrapper: TracerWrapper): SpanContextWrapper | null {
19+
// First try to extract as regular trace headers
20+
const traceContext = tracerWrapper.extract(headers);
21+
if (traceContext) {
22+
return traceContext;
23+
}
24+
25+
// If that fails, check if this is a Step Function context
26+
const stepFunctionInstance = StepFunctionContextService.instance(headers);
27+
const stepFunctionContext = stepFunctionInstance.context;
28+
29+
if (stepFunctionContext !== undefined) {
30+
const spanContext = stepFunctionInstance.spanContext;
31+
if (spanContext !== null) {
32+
return spanContext;
33+
}
34+
}
35+
36+
return null;
37+
}
38+
39+
/**
40+
* Extracts trace context from AWS Trace Header
41+
* @param awsTraceHeader The AWS trace header string
42+
* @param eventType The type of event (for logging)
43+
* @returns SpanContextWrapper or null
44+
*/
45+
export function extractFromAWSTraceHeader(awsTraceHeader: string, eventType: string): SpanContextWrapper | null {
46+
const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceHeader);
47+
if (traceContext) {
48+
logDebug(`Extracted trace context from ${eventType} event attributes AWSTraceHeader`);
49+
return traceContext;
50+
} else {
51+
logDebug(`No Datadog trace context found from ${eventType} event attributes AWSTraceHeader`);
52+
return null;
53+
}
54+
}
55+
56+
/**
57+
* Common error handler for extraction operations
58+
* @param error The error that occurred
59+
* @param eventType The type of event (for logging)
60+
*/
61+
export function handleExtractionError(error: unknown, eventType: string): void {
62+
if (error instanceof Error) {
63+
logDebug(`Unable to extract trace context from ${eventType} event`, error);
64+
}
65+
}

‎src/trace/context/extractors/event-bridge-sqs.ts‎

Lines changed: 10 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,24 @@
1-
import { EventBridgeEvent, SQSEvent } from "aws-lambda";
2-
import { EventTraceExtractor } from "../extractor";
3-
import { logDebug } from "../../../utils";
1+
import { SQSEvent } from "aws-lambda";
42
import { TracerWrapper } from "../../tracer-wrapper";
3+
import { EventTraceExtractor } from "../extractor";
54
import { SpanContextWrapper } from "../../span-context-wrapper";
6-
import { StepFunctionContextService } from "../../step-function-service";
5+
import { extractTraceContext, handleExtractionError } from "../extractor-utils";
76

87
export class EventBridgeSQSEventTraceExtractor implements EventTraceExtractor {
98
constructor(private tracerWrapper: TracerWrapper) {}
109

1110
extract(event: SQSEvent): SpanContextWrapper | null {
12-
const body = event?.Records?.[0]?.body;
13-
if (body === undefined) return null;
14-
1511
try {
16-
const parsedBody = JSON.parse(body) as EventBridgeEvent<any, any>;
17-
const headers = parsedBody?.detail?._datadog;
18-
if (headers === undefined) return null;
19-
20-
// First try to extract as regular trace headers
21-
const traceContext = this.tracerWrapper.extract(headers);
22-
if (traceContext !== null) {
23-
logDebug("Extracted trace context from EventBridge-SQS event", { traceContext, event });
24-
return traceContext;
25-
}
26-
27-
// If that fails, check if this is a Step Function context
28-
// The StepFunctionContextService can handle the Step Function format
29-
const stepFunctionInstance = StepFunctionContextService.instance(headers);
30-
const stepFunctionContext = stepFunctionInstance.context;
31-
32-
if (stepFunctionContext !== undefined) {
33-
const spanContext = stepFunctionInstance.spanContext;
34-
if (spanContext !== null) {
35-
logDebug("Extracted Step Function trace context from EventBridge-SQS event", { spanContext, event });
36-
return spanContext;
12+
const body = event?.Records?.[0]?.body;
13+
if (body) {
14+
const parsedBody = JSON.parse(body);
15+
const headers = parsedBody?.detail?._datadog;
16+
if (headers) {
17+
return extractTraceContext(headers, this.tracerWrapper);
3718
}
3819
}
3920
} catch (error) {
40-
if (error instanceof Error) {
41-
logDebug("Unable to extract trace context from EventBridge-SQS event", error);
42-
}
21+
handleExtractionError(error, "EventBridge-SQS");
4322
}
4423

4524
return null;

‎src/trace/context/extractors/event-bridge.ts‎

Lines changed: 7 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,20 @@
1-
import { EventTraceExtractor } from "../extractor";
21
import { EventBridgeEvent } from "aws-lambda";
3-
import { logDebug } from "../../../utils";
42
import { TracerWrapper } from "../../tracer-wrapper";
3+
import { EventTraceExtractor } from "../extractor";
54
import { SpanContextWrapper } from "../../span-context-wrapper";
6-
import { StepFunctionContextService } from "../../step-function-service";
5+
import { extractTraceContext, handleExtractionError } from "../extractor-utils";
76

87
export class EventBridgeEventTraceExtractor implements EventTraceExtractor {
98
constructor(private tracerWrapper: TracerWrapper) {}
109

11-
extract(event: EventBridgeEvent<any, any>): SpanContextWrapper | null {
12-
const headers = event?.detail?._datadog;
13-
if (headers === undefined) return null;
14-
10+
extract(event: EventBridgeEvent<string, any>): SpanContextWrapper | null {
1511
try {
16-
// First try to extract as regular trace headers
17-
const traceContext = this.tracerWrapper.extract(headers);
18-
if (traceContext !== null) {
19-
logDebug(`Extracted trace context from Eventbridge event`, { traceContext, event });
20-
return traceContext;
21-
}
22-
23-
// If that fails, check if this is a Step Function context
24-
// The StepFunctionContextService can handle the Step Function format
25-
const stepFunctionInstance = StepFunctionContextService.instance(headers);
26-
const stepFunctionContext = stepFunctionInstance.context;
27-
28-
if (stepFunctionContext !== undefined) {
29-
const spanContext = stepFunctionInstance.spanContext;
30-
if (spanContext !== null) {
31-
logDebug(`Extracted Step Function trace context from Eventbridge event`, { spanContext, event });
32-
return spanContext;
33-
}
12+
const headers = event?.detail?._datadog;
13+
if (headers) {
14+
return extractTraceContext(headers, this.tracerWrapper);
3415
}
3516
} catch (error) {
36-
if (error instanceof Error) {
37-
logDebug("Unable to extract trace context from EventBridge event", error);
38-
}
17+
handleExtractionError(error, "EventBridge");
3918
}
4019

4120
return null;

‎src/trace/context/extractors/sns-sqs.ts‎

Lines changed: 25 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,92 +1,56 @@
1-
import { SNSMessage, SQSEvent } from "aws-lambda";
2-
import { EventTraceExtractor } from "../extractor";
1+
import { SQSEvent } from "aws-lambda";
32
import { TracerWrapper } from "../../tracer-wrapper";
43
import { logDebug } from "../../../utils";
4+
import { EventTraceExtractor } from "../extractor";
55
import { SpanContextWrapper } from "../../span-context-wrapper";
6-
import { XrayService } from "../../xray-service";
7-
import { StepFunctionContextService } from "../../step-function-service";
6+
import { extractTraceContext, extractFromAWSTraceHeader, handleExtractionError } from "../extractor-utils";
87

98
export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
109
constructor(private tracerWrapper: TracerWrapper) {}
1110

1211
extract(event: SQSEvent): SpanContextWrapper | null {
1312
try {
14-
// First try to extract trace context from message attributes
15-
if (event?.Records?.[0]?.body) {
16-
const parsedBody = JSON.parse(event?.Records?.[0]?.body) as SNSMessage;
17-
const messageAttribute = parsedBody?.MessageAttributes?._datadog;
18-
if (messageAttribute?.Value) {
13+
// Try to extract trace context from SNS wrapped in SQS
14+
const body = event?.Records?.[0]?.body;
15+
if (body) {
16+
const parsedBody = JSON.parse(body);
17+
const snsMessageAttribute = parsedBody?.MessageAttributes?._datadog;
18+
if (snsMessageAttribute?.Value) {
1919
let headers;
20-
if (messageAttribute.Type === "String") {
21-
headers = JSON.parse(messageAttribute.Value);
20+
if (snsMessageAttribute.Type === "String") {
21+
headers = JSON.parse(snsMessageAttribute.Value);
2222
} else {
23-
const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii");
23+
// Try decoding base64 values
24+
const decodedValue = Buffer.from(snsMessageAttribute.Value, "base64").toString("ascii");
2425
headers = JSON.parse(decodedValue);
2526
}
2627

27-
// First try to extract as regular trace headers
28-
const traceContext = this.tracerWrapper.extract(headers);
28+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
2929
if (traceContext) {
30-
logDebug("Extracted trace context from SNS-SQS event");
3130
return traceContext;
3231
}
33-
34-
// If that fails, check if this is a Step Function context
35-
const stepFunctionInstance = StepFunctionContextService.instance(headers);
36-
const stepFunctionContext = stepFunctionInstance.context;
37-
38-
if (stepFunctionContext !== undefined) {
39-
const spanContext = stepFunctionInstance.spanContext;
40-
if (spanContext !== null) {
41-
logDebug("Extracted Step Function trace context from SNS-SQS event", { spanContext, event });
42-
return spanContext;
43-
}
44-
}
45-
4632
logDebug("Failed to extract trace context from SNS-SQS event");
4733
}
4834
}
4935

50-
// Check SQS message attributes for Step Function context
51-
const sqsMessageAttribute = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
52-
if (sqsMessageAttribute) {
53-
const parsedHeaders = JSON.parse(sqsMessageAttribute);
54-
55-
// First try to extract as regular trace headers
56-
const traceContext = this.tracerWrapper.extract(parsedHeaders);
36+
// Check SQS message attributes as a fallback
37+
const sqsMessageAttribute = event?.Records?.[0]?.messageAttributes?._datadog;
38+
if (sqsMessageAttribute?.stringValue) {
39+
const headers = JSON.parse(sqsMessageAttribute.stringValue);
40+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
5741
if (traceContext) {
58-
logDebug("Extracted trace context from SQS messageAttributes");
5942
return traceContext;
6043
}
61-
62-
// If that fails, check if this is a Step Function context
63-
const stepFunctionInstance = StepFunctionContextService.instance(parsedHeaders);
64-
const stepFunctionContext = stepFunctionInstance.context;
65-
66-
if (stepFunctionContext !== undefined) {
67-
const spanContext = stepFunctionInstance.spanContext;
68-
if (spanContext !== null) {
69-
logDebug("Extracted Step Function trace context from SQS messageAttributes", { spanContext, event });
70-
return spanContext;
71-
}
72-
}
7344
}
7445

75-
// Then try to extract trace context from attributes.AWSTraceHeader. (Upstream Java apps can
76-
// pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case)
77-
if (event?.Records?.[0]?.attributes?.AWSTraceHeader !== undefined) {
78-
const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader);
79-
if (traceContext) {
80-
logDebug("Extracted trace context from SNS-SQS event attributes.AWSTraceHeader");
81-
return traceContext;
82-
} else {
83-
logDebug("No Datadog trace context found from SNS-SQS event attributes.AWSTraceHeader");
84-
}
46+
// Else try to extract trace context from attributes.AWSTraceHeader
47+
// (Upstream Java apps can pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case)
48+
const awsTraceHeader = event?.Records?.[0]?.attributes?.AWSTraceHeader;
49+
if (awsTraceHeader !== undefined) {
50+
return extractFromAWSTraceHeader(awsTraceHeader, "SNS-SQS");
8551
}
8652
} catch (error) {
87-
if (error instanceof Error) {
88-
logDebug("Unable to extract trace context from SNS-SQS event", error);
89-
}
53+
handleExtractionError(error, "SQS");
9054
}
9155

9256
return null;

‎src/trace/context/extractors/sns.ts‎

Lines changed: 6 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ import { TracerWrapper } from "../../tracer-wrapper";
33
import { logDebug } from "../../../utils";
44
import { EventTraceExtractor } from "../extractor";
55
import { SpanContextWrapper } from "../../span-context-wrapper";
6-
import { XrayService, AMZN_TRACE_ID_ENV_VAR } from "../../xray-service";
7-
import { StepFunctionContextService } from "../../step-function-service";
6+
import { AMZN_TRACE_ID_ENV_VAR } from "../../xray-service";
7+
import { extractTraceContext, extractFromAWSTraceHeader, handleExtractionError } from "../extractor-utils";
88

99
export class SNSEventTraceExtractor implements EventTraceExtractor {
1010
constructor(private tracerWrapper: TracerWrapper) {}
@@ -23,42 +23,20 @@ export class SNSEventTraceExtractor implements EventTraceExtractor {
2323
headers = JSON.parse(decodedValue);
2424
}
2525

26-
// First try to extract as regular trace headers
27-
const traceContext = this.tracerWrapper.extract(headers);
26+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
2827
if (traceContext) {
29-
logDebug("Extracted trace context from SNS event");
3028
return traceContext;
3129
}
32-
33-
// If that fails, check if this is a Step Function context
34-
const stepFunctionInstance = StepFunctionContextService.instance(headers);
35-
const stepFunctionContext = stepFunctionInstance.context;
36-
37-
if (stepFunctionContext !== undefined) {
38-
const spanContext = stepFunctionInstance.spanContext;
39-
if (spanContext !== null) {
40-
logDebug("Extracted Step Function trace context from SNS event", { spanContext, event });
41-
return spanContext;
42-
}
43-
}
44-
4530
logDebug("Failed to extract trace context from SNS event");
4631
}
32+
4733
// Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can
4834
// pass down Datadog trace id (parent id wrong) in the env in SNS case)
4935
if (process.env[AMZN_TRACE_ID_ENV_VAR]) {
50-
const traceContext = XrayService.extraceDDContextFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR]);
51-
if (traceContext) {
52-
logDebug("Extracted Datadog trace context from _X_AMZN_TRACE_ID");
53-
return traceContext;
54-
} else {
55-
logDebug("No Datadog trace context found from _X_AMZN_TRACE_ID");
56-
}
36+
return extractFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR], "SNS");
5737
}
5838
} catch (error) {
59-
if (error instanceof Error) {
60-
logDebug("Unable to extract trace context from SNS event", error);
61-
}
39+
handleExtractionError(error, "SNS");
6240
}
6341

6442
return null;

0 commit comments

Comments
 (0)