Skip to content

Commit 33b7bfc

Browse files
authored
[SVLS-4298] Java upstream sqs trace context propagation (#454)
* java upstream sqs trace context propagation
1 parent 8e79a70 commit 33b7bfc

File tree

4 files changed

+80
-22
lines changed

4 files changed

+80
-22
lines changed

datadog_lambda/tracing.py

Lines changed: 47 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@
6868

6969
propagator = HTTPPropagator()
7070

71+
DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000"
72+
7173

7274
def _convert_xray_trace_id(xray_trace_id):
7375
"""
@@ -248,28 +250,52 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
248250
first_record.get("Sns", {}).get("MessageAttributes", {}),
249251
)
250252
dd_payload = msg_attributes.get("_datadog", {})
251-
# SQS uses dataType and binaryValue/stringValue
252-
# SNS uses Type and Value
253-
dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", ""))
254-
if dd_json_data_type == "Binary":
255-
dd_json_data = dd_payload.get(
256-
"binaryValue",
257-
dd_payload.get("Value", r"{}"),
258-
)
259-
dd_json_data = base64.b64decode(dd_json_data)
260-
elif dd_json_data_type == "String":
261-
dd_json_data = dd_payload.get(
262-
"stringValue",
263-
dd_payload.get("Value", r"{}"),
264-
)
253+
if dd_payload:
254+
# SQS uses dataType and binaryValue/stringValue
255+
# SNS uses Type and Value
256+
dd_json_data = None
257+
dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", ""))
258+
if dd_json_data_type == "Binary":
259+
dd_json_data = dd_payload.get(
260+
"binaryValue",
261+
dd_payload.get("Value", r"{}"),
262+
)
263+
dd_json_data = base64.b64decode(dd_json_data)
264+
elif dd_json_data_type == "String":
265+
dd_json_data = dd_payload.get(
266+
"stringValue",
267+
dd_payload.get("Value", r"{}"),
268+
)
269+
else:
270+
logger.debug(
271+
"Datadog Lambda Python only supports extracting trace"
272+
"context from String or Binary SQS/SNS message attributes"
273+
)
274+
275+
if dd_json_data:
276+
dd_data = json.loads(dd_json_data)
277+
return propagator.extract(dd_data)
265278
else:
266-
logger.debug(
267-
"Datadog Lambda Python only supports extracting trace"
268-
"context from String or Binary SQS/SNS message attributes"
269-
)
270-
return extract_context_from_lambda_context(lambda_context)
271-
dd_data = json.loads(dd_json_data)
272-
return propagator.extract(dd_data)
279+
# Handle case where trace context is injected into attributes.AWSTraceHeader
280+
# example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
281+
x_ray_header = first_record.get("attributes", {}).get("AWSTraceHeader")
282+
if x_ray_header:
283+
x_ray_context = parse_xray_header(x_ray_header)
284+
trace_id_parts = x_ray_context.get("trace_id", "").split("-")
285+
if len(trace_id_parts) > 2 and trace_id_parts[2].startswith(
286+
DD_TRACE_JAVA_TRACE_ID_PADDING
287+
):
288+
# If it starts with eight 0's padding,
289+
# then this AWSTraceHeader contains Datadog injected trace context
290+
logger.debug(
291+
"Found dd-trace injected trace context from AWSTraceHeader"
292+
)
293+
return Context(
294+
trace_id=int(trace_id_parts[2][8:], 16),
295+
span_id=int(int(x_ray_context["parent_id"], 16)),
296+
sampling_priority=float(x_ray_context["sampled"]),
297+
)
298+
return extract_context_from_lambda_context(lambda_context)
273299
except Exception as e:
274300
logger.debug("The trace extractor returned with error %s", e)
275301
return extract_context_from_lambda_context(lambda_context)
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"Records": [
3+
{
4+
"messageId": "f7e888aa-1368-484c-8e15-fc3f0f7c6fea",
5+
"receiptHandle": "AQEBN1aYTQ1c5huZh9bkhBYqcMMnqTUMRh8MfUPyGXkEolcn23rvM9saGEg3wTK/7JnJ1s3Uk107uLjaP6yV6+zS3oQRU0vMG2LfyTgHovWhYQ8TnrpC7XpYL+Uf+oc9KoILQopiYi4wsFnOWQqy82yQmlOA3W+CZ3Rvq8N6rNcmyaZEXVdozHG+FyMCMQ8QdTcCHhzR9YKnkZ87Y40+LhysUR57VNPVtRwENI8H1uMEfmxaCkW+CAkdCGoXeX+KioT7pHJDZaEutXM3VRmGXDDzCXvfUJQ9JQIlP5xe66JO8/cpCyl5sDoHsCjLy6X/XCmfG2+XclPObGHBzcMSjG1RQtHsEGTOAJrLREucqf/oj0Ab4svpxz6lR4UXrICygZ2x0NZcNFXcZx3GV2QL9nHmJxzrO2lnNTEOMuYB4SnqtIhsaDTcmkYHumaAJdRHl5BksFcU5qpS7BQrnRvXn5Sz3hYdR2KuYKN5Oq6W1vuT16o=",
6+
"body": "{\"hello\":\"world\"}",
7+
"attributes": {
8+
"ApproximateReceiveCount": "1",
9+
"AWSTraceHeader": "Root=1-65eb7350-000000006dfd06bf489aa4e5;Parent=48cc02b6aafae897;Sampled=1",
10+
"SentTimestamp": "1709929297382",
11+
"SenderId": "AROAWGCM4HXUTHYJKOM7M:DdTraceXLambda-sqsjavache-sqsjavaproducerforPython-moi7s7Hu7Ppy",
12+
"ApproximateFirstReceiveTimestamp": "1709929297387"
13+
},
14+
"messageAttributes": {},
15+
"md5OfMessageAttributes": "",
16+
"md5OfBody": "fbc24bcc7a1794758fc1327fcfebdaf6",
17+
"eventSource": "aws:sqs",
18+
"eventSourceARN": "arn:aws:sqs:us-west-2:444442222111:DdTraceXLambda-sqsjavachecksNeste-sqsJavaProducer2sqsjavaproducerfo-dwpHQF6fcZT4",
19+
"awsRegion": "us-west-2"
20+
}
21+
]
22+
}

tests/test_cold_start.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,6 @@ def test_trace_ignore_libs(self):
240240

241241

242242
def test_lazy_loaded_package_imports(monkeypatch):
243-
244243
spans = []
245244

246245
def finish(span):

tests/test_tracing.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1963,6 +1963,17 @@ def test_extract_context_from_sqs_batch_event(self):
19631963
self.assertEqual(context.span_id, 7431398482019833808)
19641964
self.assertEqual(context.sampling_priority, 1)
19651965

1966+
def test_extract_context_from_sqs_java_upstream_event(self):
1967+
event_sample_source = "sqs-java-upstream"
1968+
test_file = event_samples + event_sample_source + ".json"
1969+
with open(test_file, "r") as event:
1970+
event = json.load(event)
1971+
ctx = get_mock_context()
1972+
context, source, event_type = extract_dd_trace_context(event, ctx)
1973+
self.assertEqual(context.trace_id, 7925498337868555493)
1974+
self.assertEqual(context.span_id, 5245570649555658903)
1975+
self.assertEqual(context.sampling_priority, 1)
1976+
19661977
def test_extract_context_from_sns_event_with_string_msg_attr(self):
19671978
event_sample_source = "sns-string-msg-attribute"
19681979
test_file = event_samples + event_sample_source + ".json"

0 commit comments

Comments
 (0)