Skip to content

Commit 21eadee

Browse files
andreiborzaandreiborza
committed
Use getTraceData and uniq from core for message propagation
getTraceData({ span }) replaces the hand-rolled sentry-trace/baggage pair (matching the kafkajs/amqplib producer idiom and gaining its enabled/validity guards), callers pass the serialized headers in so batch sends compute them once, and addPropagationFieldsToAttributeNames uses core's uniq.
1 parent e28f4e4 commit 21eadee

4 files changed

Lines changed: 30 additions & 39 deletions

File tree

‎packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/MessageAttributes.ts‎

Lines changed: 15 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,5 @@
1-
import type { Span } from '@sentry/core';
2-
import {
3-
debug,
4-
dynamicSamplingContextToSentryBaggageHeader,
5-
getDynamicSamplingContextFromSpan,
6-
spanToTraceHeader,
7-
} from '@sentry/core';
1+
import type { SerializedTraceData } from '@sentry/core';
2+
import { debug, uniq } from '@sentry/core';
83
import { DEBUG_BUILD } from '../../../../debug-build';
94
import type { SNS, SQS } from '../aws-sdk.types';
105

@@ -23,36 +18,27 @@ export interface AwsSdkContextObject {
2318
};
2419
}
2520

26-
/** Build the `sentry-trace`/`baggage` header pair carrying the span's trace context. */
27-
export function getPropagationHeaders(span: Span): Record<string, string> {
28-
const headers: Record<string, string> = {
29-
[SENTRY_TRACE_HEADER]: spanToTraceHeader(span),
30-
};
31-
const baggage = dynamicSamplingContextToSentryBaggageHeader(getDynamicSamplingContextFromSpan(span));
32-
if (baggage) {
33-
headers[BAGGAGE_HEADER] = baggage;
34-
}
35-
return headers;
36-
}
37-
3821
/**
39-
* Inject the span's trace-propagation headers into an SQS/SNS message-attribute map, so the consumer
40-
* can continue the trace. Respects the SQS 10-attribute quota. Mirrors the OTel integration's
41-
* `injectPropagationContext`, but writes Sentry's `sentry-trace`/`baggage` instead of W3C headers.
22+
* Inject trace-propagation headers (from `getTraceData({ span })`) into an SQS/SNS message-attribute
23+
* map, so the consumer can continue the trace. Respects the SQS 10-attribute quota. Mirrors the OTel
24+
* integration's `injectPropagationContext`, but writes Sentry's `sentry-trace`/`baggage` instead of
25+
* W3C headers. Callers pass the precomputed headers so batch sends serialize them only once.
4226
*/
4327
export function injectPropagationContext(
4428
attributesMap: SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap | undefined,
45-
span: Span,
29+
traceData: SerializedTraceData,
4630
): SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap {
4731
const attributes = attributesMap ?? {};
48-
const headers = getPropagationHeaders(span);
49-
const headerKeys = Object.keys(headers);
32+
const headerKeys = Object.keys(traceData) as (keyof SerializedTraceData)[];
5033

5134
if (Object.keys(attributes).length + headerKeys.length <= MAX_MESSAGE_ATTRIBUTES) {
5235
for (const key of headerKeys) {
53-
// Index-assigning into the SQS/SNS map union needs one concrete map type; the written value
54-
// shape is valid for both.
55-
(attributes as SQS.MessageBodyAttributeMap)[key] = { DataType: 'String', StringValue: headers[key] };
36+
const value = traceData[key];
37+
if (value) {
38+
// Index-assigning into the SQS/SNS map union needs one concrete map type; the written value
39+
// shape is valid for both.
40+
(attributes as SQS.MessageBodyAttributeMap)[key] = { DataType: 'String', StringValue: value };
41+
}
5642
}
5743
} else {
5844
DEBUG_BUILD &&
@@ -79,7 +65,5 @@ export function extractPropagationHeaders(
7965
}
8066

8167
export function addPropagationFieldsToAttributeNames(messageAttributeNames: string[] = []): string[] {
82-
return messageAttributeNames.length
83-
? Array.from(new Set([...messageAttributeNames, ...PROPAGATION_FIELDS]))
84-
: [...PROPAGATION_FIELDS];
68+
return uniq([...messageAttributeNames, ...PROPAGATION_FIELDS]);
8569
}

‎packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/lambda.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { Span } from '@sentry/core';
2-
import { debug, SPAN_KIND } from '@sentry/core';
2+
import { debug, getTraceData, SPAN_KIND } from '@sentry/core';
33
import { DEBUG_BUILD } from '../../../../debug-build';
44
import {
55
ATTR_FAAS_EXECUTION,
@@ -8,7 +8,6 @@ import {
88
ATTR_FAAS_INVOKED_REGION,
99
} from '../constants';
1010
import type { NormalizedRequest, NormalizedResponse } from '../types';
11-
import { getPropagationHeaders } from './MessageAttributes';
1211
import type { RequestMetadata, ServiceExtension } from './ServiceExtension';
1312

1413
const INVOKE_COMMAND = 'Invoke';
@@ -54,7 +53,7 @@ export class LambdaServiceExtension implements ServiceExtension {
5453

5554
function injectLambdaPropagationContext(clientContext: string | undefined, span: Span): string | undefined {
5655
try {
57-
const propagatedContext = getPropagationHeaders(span);
56+
const propagatedContext = getTraceData({ span });
5857

5958
const parsedClientContext = clientContext ? JSON.parse(Buffer.from(clientContext, 'base64').toString('utf8')) : {};
6059

‎packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/sns.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { Span, SpanKindValue } from '@sentry/core';
2-
import { SPAN_KIND } from '@sentry/core';
2+
import { getTraceData, SPAN_KIND } from '@sentry/core';
33
import { MESSAGING_DESTINATION_NAME, MESSAGING_SYSTEM } from '@sentry/conventions/attributes';
44
import {
55
ATTR_AWS_SNS_TOPIC_ARN,
@@ -46,7 +46,7 @@ export class SnsServiceExtension implements ServiceExtension {
4646
public requestPostSpanHook(request: NormalizedRequest, span: Span): void {
4747
if (request.commandName === 'Publish') {
4848
const origMessageAttributes = request.commandInput.MessageAttributes ?? {};
49-
request.commandInput.MessageAttributes = injectPropagationContext(origMessageAttributes, span);
49+
request.commandInput.MessageAttributes = injectPropagationContext(origMessageAttributes, getTraceData({ span }));
5050
}
5151
}
5252

‎packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/sqs.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { Span, SpanKindValue } from '@sentry/core';
2-
import { propagationContextFromHeaders, SPAN_KIND } from '@sentry/core';
2+
import { getTraceData, propagationContextFromHeaders, SPAN_KIND } from '@sentry/core';
33
import {
44
MESSAGING_BATCH_MESSAGE_COUNT,
55
MESSAGING_DESTINATION_NAME,
@@ -62,16 +62,24 @@ export class SqsServiceExtension implements ServiceExtension {
6262
case 'SendMessage':
6363
{
6464
const origMessageAttributes = request.commandInput.MessageAttributes ?? {};
65-
request.commandInput.MessageAttributes = injectPropagationContext(origMessageAttributes, span);
65+
request.commandInput.MessageAttributes = injectPropagationContext(
66+
origMessageAttributes,
67+
getTraceData({ span }),
68+
);
6669
}
6770
break;
6871

6972
case 'SendMessageBatch':
7073
{
7174
const entries = request.commandInput?.Entries;
7275
if (Array.isArray(entries)) {
76+
// Serialized once; the headers are identical for every entry of the batch.
77+
const traceData = getTraceData({ span });
7378
entries.forEach((messageParams: { MessageAttributes: SQS.MessageBodyAttributeMap }) => {
74-
messageParams.MessageAttributes = injectPropagationContext(messageParams.MessageAttributes ?? {}, span);
79+
messageParams.MessageAttributes = injectPropagationContext(
80+
messageParams.MessageAttributes ?? {},
81+
traceData,
82+
);
7583
});
7684
}
7785
}

0 commit comments

Comments
 (0)