[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"