Skip to content

Commit c94ecbb

Browse files
ericfirthericfirth
committed
Update the extraction functions to separate the setting of dsm checkpoints with the context stuff completely
1 parent 632126b commit c94ecbb

File tree

4 files changed

+205
-146
lines changed

4 files changed

+205
-146
lines changed

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

Lines changed: 34 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,35 +9,44 @@ export class KinesisEventTraceExtractor implements EventTraceExtractor {
99
constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {}
1010

1111
extract(event: KinesisStreamEvent): SpanContextWrapper | null {
12-
let context: SpanContextWrapper | null = null;
13-
14-
for (const record of event?.Records || []) {
15-
try {
16-
// If we already have a context and dsm is not enabled, we can break out of the loop early
17-
if (!this.config.dataStreamsEnabled && context) {
18-
break;
12+
// Set DSM consume checkpoints if enabled
13+
if (this.config.dataStreamsEnabled) {
14+
for (const record of event?.Records || []) {
15+
try {
16+
const kinesisData = record?.kinesis?.data;
17+
const decodedData = Buffer.from(kinesisData, "base64").toString("ascii");
18+
const parsedBody = JSON.parse(decodedData);
19+
const headers = parsedBody?._datadog ?? null;
20+
this.tracerWrapper.setConsumeCheckpoint(headers, "kinesis", record.eventSourceARN);
21+
} catch (error) {
22+
if (error instanceof Error) {
23+
logDebug("Unable to set DSM checkpoint for Kinesis event", error);
24+
}
1925
}
26+
}
27+
}
2028

21-
const kinesisData = record?.kinesis?.data;
22-
const decodedData = Buffer.from(kinesisData, "base64").toString("ascii");
23-
const parsedBody = JSON.parse(decodedData);
24-
const headers = parsedBody?._datadog ?? null;
25-
this.tracerWrapper.setConsumeCheckpoint(headers, "kinesis", record.eventSourceARN);
29+
// Extract trace context (original logic - unchanged from master)
30+
const kinesisData = event?.Records?.[0]?.kinesis.data;
31+
if (kinesisData === undefined) return null;
2632

27-
// if we already have a context, no need to extract again
28-
if (context) continue;
29-
if (headers) {
30-
context = this.tracerWrapper.extract(headers);
31-
if (context !== null) {
32-
logDebug(`Extracted trace context from Kinesis event`, { traceContext: context, headers });
33-
}
34-
}
35-
} catch (error) {
36-
if (error instanceof Error) {
37-
logDebug("Unable to extract trace context from Kinesis event", error);
38-
}
33+
try {
34+
const decodedData = Buffer.from(kinesisData, "base64").toString("ascii");
35+
const parsedBody = JSON.parse(decodedData);
36+
const headers = parsedBody?._datadog;
37+
if (headers) {
38+
const traceContext = this.tracerWrapper.extract(headers);
39+
if (traceContext === null) return null;
40+
41+
logDebug(`Extracted trace context from Kinesis event`, { traceContext, headers });
42+
return traceContext;
43+
}
44+
} catch (error) {
45+
if (error instanceof Error) {
46+
logDebug("Unable to extract trace context from Kinesis event", error);
3947
}
4048
}
41-
return context;
49+
50+
return null;
4251
}
4352
}

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

Lines changed: 71 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -10,65 +10,89 @@ export class SNSSQSEventTraceExtractor implements EventTraceExtractor {
1010
constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {}
1111

1212
extract(event: SQSEvent): SpanContextWrapper | null {
13-
logDebug("SNS-SQS Extractor Being Used");
14-
15-
let context: SpanContextWrapper | null = null;
16-
for (const record of event?.Records || []) {
17-
try {
18-
// If we already have a context and dsm is not enabled, we can break out of the loop early
19-
if (!this.config.dataStreamsEnabled && context) {
20-
break;
21-
}
22-
23-
let headers = null;
24-
// Try to extract trace context from SNS wrapped in SQS
25-
const body = record.body;
26-
if (body) {
27-
const parsedBody = JSON.parse(body);
28-
const snsMessageAttribute = parsedBody?.MessageAttributes?._datadog;
29-
if (snsMessageAttribute?.Value) {
30-
if (snsMessageAttribute.Type === "String") {
31-
headers = JSON.parse(snsMessageAttribute.Value);
32-
} else {
33-
// Try decoding base64 values
34-
const decodedValue = Buffer.from(snsMessageAttribute.Value, "base64").toString("ascii");
35-
headers = JSON.parse(decodedValue);
13+
// Set DSM consume checkpoints if enabled
14+
if (this.config.dataStreamsEnabled) {
15+
for (const record of event?.Records || []) {
16+
try {
17+
let headers = null;
18+
// Try to extract trace context from SNS wrapped in SQS
19+
const body = record.body;
20+
if (body) {
21+
const parsedBody = JSON.parse(body);
22+
const snsMessageAttribute = parsedBody?.MessageAttributes?._datadog;
23+
if (snsMessageAttribute?.Value) {
24+
if (snsMessageAttribute.Type === "String") {
25+
headers = JSON.parse(snsMessageAttribute.Value);
26+
} else {
27+
// Try decoding base64 values
28+
const decodedValue = Buffer.from(snsMessageAttribute.Value, "base64").toString("ascii");
29+
headers = JSON.parse(decodedValue);
30+
}
3631
}
3732
}
38-
}
3933

40-
// Check SQS message attributes as a fallback
41-
if (!headers) {
42-
const sqsMessageAttribute = record.messageAttributes?._datadog;
43-
if (sqsMessageAttribute?.stringValue) {
44-
headers = JSON.parse(sqsMessageAttribute.stringValue);
34+
// Check SQS message attributes as a fallback
35+
if (!headers) {
36+
const sqsMessageAttribute = record.messageAttributes?._datadog;
37+
if (sqsMessageAttribute?.stringValue) {
38+
headers = JSON.parse(sqsMessageAttribute.stringValue);
39+
}
4540
}
46-
}
4741

48-
// Set a checkpoint for the record, even if we don't have headers
49-
this.tracerWrapper.setConsumeCheckpoint(headers, "sqs", record.eventSourceARN);
42+
// Set a checkpoint for the record, even if we don't have headers
43+
this.tracerWrapper.setConsumeCheckpoint(headers, "sqs", record.eventSourceARN);
44+
} catch (error) {
45+
handleExtractionError(error, "SQS DSM checkpoint");
46+
}
47+
}
48+
}
5049

51-
// if we already have a context, no need to extract again
52-
if (context) continue;
50+
// Extract trace context (original logic - unchanged from master)
51+
logDebug("SNS-SQS Extractor Being Used");
52+
try {
53+
// Try to extract trace context from SNS wrapped in SQS
54+
const body = event?.Records?.[0]?.body;
55+
if (body) {
56+
const parsedBody = JSON.parse(body);
57+
const snsMessageAttribute = parsedBody?.MessageAttributes?._datadog;
58+
if (snsMessageAttribute?.Value) {
59+
let headers;
60+
if (snsMessageAttribute.Type === "String") {
61+
headers = JSON.parse(snsMessageAttribute.Value);
62+
} else {
63+
// Try decoding base64 values
64+
const decodedValue = Buffer.from(snsMessageAttribute.Value, "base64").toString("ascii");
65+
headers = JSON.parse(decodedValue);
66+
}
5367

54-
// Try to extract trace context from headers
55-
if (headers) {
56-
context = extractTraceContext(headers, this.tracerWrapper);
57-
} else {
68+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
69+
if (traceContext) {
70+
return traceContext;
71+
}
5872
logDebug("Failed to extract trace context from SNS-SQS event");
73+
}
74+
}
5975

60-
// Else try to extract trace context from attributes.AWSTraceHeader
61-
// (Upstream Java apps can pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case)
62-
const awsTraceHeader = record.attributes?.AWSTraceHeader;
63-
if (awsTraceHeader !== undefined) {
64-
context = extractFromAWSTraceHeader(awsTraceHeader, "SNS-SQS");
65-
}
76+
// Check SQS message attributes as a fallback
77+
const sqsMessageAttribute = event?.Records?.[0]?.messageAttributes?._datadog;
78+
if (sqsMessageAttribute?.stringValue) {
79+
const headers = JSON.parse(sqsMessageAttribute.stringValue);
80+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
81+
if (traceContext) {
82+
return traceContext;
6683
}
67-
} catch (error) {
68-
handleExtractionError(error, "SQS");
6984
}
85+
86+
// Else try to extract trace context from attributes.AWSTraceHeader
87+
// (Upstream Java apps can pass down Datadog trace context in the attributes.AWSTraceHeader in SQS case)
88+
const awsTraceHeader = event?.Records?.[0]?.attributes?.AWSTraceHeader;
89+
if (awsTraceHeader !== undefined) {
90+
return extractFromAWSTraceHeader(awsTraceHeader, "SNS-SQS");
91+
}
92+
} catch (error) {
93+
handleExtractionError(error, "SQS");
7094
}
7195

72-
return context;
96+
return null;
7397
}
7498
}

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

Lines changed: 48 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -11,51 +11,62 @@ export class SNSEventTraceExtractor implements EventTraceExtractor {
1111
constructor(private tracerWrapper: TracerWrapper, private config: TraceConfig) {}
1212

1313
extract(event: SNSEvent): SpanContextWrapper | null {
14-
let context: SpanContextWrapper | null = null;
15-
16-
for (const record of event?.Records || []) {
17-
try {
18-
// If we already have a context and dsm is not enabled, we can break out of the loop early
19-
if (!this.config.dataStreamsEnabled && context) {
20-
break;
21-
}
22-
23-
let headers = null;
24-
const sourceARN = record.Sns?.TopicArn;
25-
// First try to extract trace context from message attributes
26-
const messageAttribute = record.Sns?.MessageAttributes?._datadog;
27-
if (messageAttribute?.Value) {
28-
if (messageAttribute.Type === "String") {
29-
headers = JSON.parse(messageAttribute.Value);
30-
} else {
31-
// Try decoding base64 values
32-
const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii");
33-
headers = JSON.parse(decodedValue);
14+
// Set DSM consume checkpoints if enabled
15+
if (this.config.dataStreamsEnabled) {
16+
for (const record of event?.Records || []) {
17+
try {
18+
let headers = null;
19+
const sourceARN = record.Sns?.TopicArn;
20+
// First try to extract trace context from message attributes
21+
const messageAttribute = record.Sns?.MessageAttributes?._datadog;
22+
if (messageAttribute?.Value) {
23+
if (messageAttribute.Type === "String") {
24+
headers = JSON.parse(messageAttribute.Value);
25+
} else {
26+
// Try decoding base64 values
27+
const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii");
28+
headers = JSON.parse(decodedValue);
29+
}
3430
}
35-
}
3631

37-
// Set a checkpoint for the record, even if we don't have headers
38-
this.tracerWrapper.setConsumeCheckpoint(headers, "sns", sourceARN);
32+
// Set a checkpoint for the record, even if we don't have headers
33+
this.tracerWrapper.setConsumeCheckpoint(headers, "sns", sourceARN);
34+
} catch (error) {
35+
handleExtractionError(error, "SNS DSM checkpoint");
36+
}
37+
}
38+
}
3939

40-
// if we already have a context, no need to extract again
41-
if (context) continue;
42-
// Try to extract trace context from headers
43-
if (headers) {
44-
context = extractTraceContext(headers, this.tracerWrapper);
40+
// Extract trace context (original logic - unchanged from master)
41+
try {
42+
// First try to extract trace context from message attributes
43+
const messageAttribute = event?.Records?.[0]?.Sns?.MessageAttributes?._datadog;
44+
if (messageAttribute?.Value) {
45+
let headers;
46+
if (messageAttribute.Type === "String") {
47+
headers = JSON.parse(messageAttribute.Value);
4548
} else {
46-
logDebug("Failed to extract trace context from SNS event");
49+
// Try decoding base64 values
50+
const decodedValue = Buffer.from(messageAttribute.Value, "base64").toString("ascii");
51+
headers = JSON.parse(decodedValue);
52+
}
4753

48-
// Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can
49-
// pass down Datadog trace id (parent id wrong) in the env in SNS case)
50-
if (process.env[AMZN_TRACE_ID_ENV_VAR]) {
51-
context = extractFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR], "SNS");
52-
}
54+
const traceContext = extractTraceContext(headers, this.tracerWrapper);
55+
if (traceContext) {
56+
return traceContext;
5357
}
54-
} catch (error) {
55-
handleExtractionError(error, "SNS");
58+
logDebug("Failed to extract trace context from SNS event");
59+
}
60+
61+
// Then try to extract trace context from _X_AMZN_TRACE_ID header (Upstream Java apps can
62+
// pass down Datadog trace id (parent id wrong) in the env in SNS case)
63+
if (process.env[AMZN_TRACE_ID_ENV_VAR]) {
64+
return extractFromAWSTraceHeader(process.env[AMZN_TRACE_ID_ENV_VAR], "SNS");
5665
}
66+
} catch (error) {
67+
handleExtractionError(error, "SNS");
5768
}
5869

59-
return context;
70+
return null;
6071
}
6172
}

0 commit comments

Comments
 (0)