[SVLS-4298] Java upstream sqs trace context propagation (#454)

* java upstream sqs trace context propagation

Author: joeyzhao2018

datadog_lambda/tracing.py

@@ -68,6 +68,8 @@
 
 propagator = HTTPPropagator()
 
+DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000"
+
 
 def _convert_xray_trace_id(xray_trace_id):
  """
@@ -248,28 +250,52 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
  first_record.get("Sns", {}).get("MessageAttributes", {}),
  )
  dd_payload = msg_attributes.get("_datadog", {})
- # SQS uses dataType and binaryValue/stringValue
- # SNS uses Type and Value
- dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", ""))
- if dd_json_data_type == "Binary":
- dd_json_data = dd_payload.get(
- "binaryValue",
- dd_payload.get("Value", r"{}"),
- )
- dd_json_data = base64.b64decode(dd_json_data)
- elif dd_json_data_type == "String":
- dd_json_data = dd_payload.get(
- "stringValue",
- dd_payload.get("Value", r"{}"),
- )
+ if dd_payload:
+ # SQS uses dataType and binaryValue/stringValue
+ # SNS uses Type and Value
+ dd_json_data = None
+ dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", ""))
+ if dd_json_data_type == "Binary":
+ dd_json_data = dd_payload.get(
+ "binaryValue",
+ dd_payload.get("Value", r"{}"),
+ )
+ dd_json_data = base64.b64decode(dd_json_data)
+ elif dd_json_data_type == "String":
+ dd_json_data = dd_payload.get(
+ "stringValue",
+ dd_payload.get("Value", r"{}"),
+ )
+ else:
+ logger.debug(
+ "Datadog Lambda Python only supports extracting trace"
+ "context from String or Binary SQS/SNS message attributes"
+ )
+
+ if dd_json_data:
+ dd_data = json.loads(dd_json_data)
+ return propagator.extract(dd_data)
  else:
- logger.debug(
- "Datadog Lambda Python only supports extracting trace"
- "context from String or Binary SQS/SNS message attributes"
- )
- return extract_context_from_lambda_context(lambda_context)
- dd_data = json.loads(dd_json_data)
- return propagator.extract(dd_data)
+ # Handle case where trace context is injected into attributes.AWSTraceHeader
+ # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
+ x_ray_header = first_record.get("attributes", {}).get("AWSTraceHeader")
+ if x_ray_header:
+ x_ray_context = parse_xray_header(x_ray_header)
+ trace_id_parts = x_ray_context.get("trace_id", "").split("-")
+ if len(trace_id_parts) > 2 and trace_id_parts[2].startswith(
+ DD_TRACE_JAVA_TRACE_ID_PADDING
+ ):
+ # If it starts with eight 0's padding,
+ # then this AWSTraceHeader contains Datadog injected trace context
+ logger.debug(
+ "Found dd-trace injected trace context from AWSTraceHeader"
+ )
+ return Context(
+ trace_id=int(trace_id_parts[2][8:], 16),
+ span_id=int(int(x_ray_context["parent_id"], 16)),
+ sampling_priority=float(x_ray_context["sampled"]),
+ )
+ return extract_context_from_lambda_context(lambda_context)
  except Exception as e:
  logger.debug("The trace extractor returned with error %s", e)
  return extract_context_from_lambda_context(lambda_context)

tests/event_samples/sqs-java-upstream.json

@@ -0,0 +1,22 @@
+{
+ "Records": [
+ {
+ "messageId": "f7e888aa-1368-484c-8e15-fc3f0f7c6fea",
+ "receiptHandle": "AQEBN1aYTQ1c5huZh9bkhBYqcMMnqTUMRh8MfUPyGXkEolcn23rvM9saGEg3wTK/7JnJ1s3Uk107uLjaP6yV6+zS3oQRU0vMG2LfyTgHovWhYQ8TnrpC7XpYL+Uf+oc9KoILQopiYi4wsFnOWQqy82yQmlOA3W+CZ3Rvq8N6rNcmyaZEXVdozHG+FyMCMQ8QdTcCHhzR9YKnkZ87Y40+LhysUR57VNPVtRwENI8H1uMEfmxaCkW+CAkdCGoXeX+KioT7pHJDZaEutXM3VRmGXDDzCXvfUJQ9JQIlP5xe66JO8/cpCyl5sDoHsCjLy6X/XCmfG2+XclPObGHBzcMSjG1RQtHsEGTOAJrLREucqf/oj0Ab4svpxz6lR4UXrICygZ2x0NZcNFXcZx3GV2QL9nHmJxzrO2lnNTEOMuYB4SnqtIhsaDTcmkYHumaAJdRHl5BksFcU5qpS7BQrnRvXn5Sz3hYdR2KuYKN5Oq6W1vuT16o=",
+ "body": "{\"hello\":\"world\"}",
+ "attributes": {
+ "ApproximateReceiveCount": "1",
+ "AWSTraceHeader": "Root=1-65eb7350-000000006dfd06bf489aa4e5;Parent=48cc02b6aafae897;Sampled=1",
+ "SentTimestamp": "1709929297382",
+ "SenderId": "AROAWGCM4HXUTHYJKOM7M:DdTraceXLambda-sqsjavache-sqsjavaproducerforPython-moi7s7Hu7Ppy",
+ "ApproximateFirstReceiveTimestamp": "1709929297387"
+ },
+ "messageAttributes": {},
+ "md5OfMessageAttributes": "",
+ "md5OfBody": "fbc24bcc7a1794758fc1327fcfebdaf6",
+ "eventSource": "aws:sqs",
+ "eventSourceARN": "arn:aws:sqs:us-west-2:444442222111:DdTraceXLambda-sqsjavachecksNeste-sqsJavaProducer2sqsjavaproducerfo-dwpHQF6fcZT4",
+ "awsRegion": "us-west-2"
+ }
+ ]
+}

tests/test_cold_start.py

@@ -240,7 +240,6 @@ def test_trace_ignore_libs(self):
 
 
 def test_lazy_loaded_package_imports(monkeypatch):
-
  spans = []
 
  def finish(span):

tests/test_tracing.py

@@ -1963,6 +1963,17 @@ def test_extract_context_from_sqs_batch_event(self):
  self.assertEqual(context.span_id, 7431398482019833808)
  self.assertEqual(context.sampling_priority, 1)
 
+ def test_extract_context_from_sqs_java_upstream_event(self):
+ event_sample_source = "sqs-java-upstream"
+ test_file = event_samples + event_sample_source + ".json"
+ with open(test_file, "r") as event:
+ event = json.load(event)
+ ctx = get_mock_context()
+ context, source, event_type = extract_dd_trace_context(event, ctx)
+ self.assertEqual(context.trace_id, 7925498337868555493)
+ self.assertEqual(context.span_id, 5245570649555658903)
+ self.assertEqual(context.sampling_priority, 1)
+
  def test_extract_context_from_sns_event_with_string_msg_attr(self):
  event_sample_source = "sns-string-msg-attribute"
  test_file = event_samples + event_sample_source + ".json"