add logic to extract traceID from _datadog header

Author: avedmala

datadog_lambda/constants.py

@@ -3,9 +3,8 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
-# Datadog trace sampling priority
-
 
+# Datadog trace sampling priority
 class SamplingPriority(object):
     USER_REJECT = -1
     AUTO_REJECT = 0
@@ -18,6 +17,7 @@ class TraceHeader(object):
     TRACE_ID = "x-datadog-trace-id"
     PARENT_ID = "x-datadog-parent-id"
     SAMPLING_PRIORITY = "x-datadog-sampling-priority"
+    TAGS = "x-datadog-tags"
 
 
 # X-Ray subsegment to save Datadog trace metadata

datadog_lambda/tracing.py

@@ -24,6 +24,7 @@
     TraceContextSource,
     XrayDaemon,
     Headers,
+    TraceHeader,
 )
 from datadog_lambda.xray import (
     send_segment,
@@ -380,10 +381,26 @@ def extract_context_from_step_functions(event, lambda_context):
         execution_id = event.get("Execution").get("Id")
         state_name = event.get("State").get("Name")
         state_entered_time = event.get("State").get("EnteredTime")
-        # returning 128 bits since 128bit traceId will be break up into
-        # traditional traceId and _dd.p.tid tag
-        # https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
-        trace_id = _deterministic_sha256_hash(execution_id, LOWER_64_BITS)
+        meta = {}
+
+        if "_datadog" in event:
+            # use the trace ID from the top-most parent when it exists
+            trace_header = event.get("_datadog")
+            trace_id = trace_header.get(TraceHeader.TRACE_ID)
+            tags = trace_header.get(TraceHeader.TAGS)
+            for tag in tags.split(","):
+                tag_key, tag_val = tag.split("=")
+                meta[tag_key] = tag_val
+        else:
+            # returning 128 bits since 128bit traceId will be break up into
+            # traditional traceId and _dd.p.tid tag
+            # https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
+            trace_id = _deterministic_sha256_hash(execution_id, LOWER_64_BITS)
+            # take the higher 64 bits as _dd.p.tid tag and use hex to encode
+            # [2:] to remove '0x' in the hex str
+            meta["_dd.p.tid"] = hex(
+                _deterministic_sha256_hash(execution_id, HIGHER_64_BITS)
+            )[2:]
 
         parent_id = _deterministic_sha256_hash(
             f"{execution_id}#{state_name}#{state_entered_time}", HIGHER_64_BITS
@@ -394,13 +411,7 @@ def extract_context_from_step_functions(event, lambda_context):
             trace_id=trace_id,
             span_id=parent_id,
             sampling_priority=sampling_priority,
-            # take the higher 64 bits as _dd.p.tid tag and use hex to encode
-            # [2:] to remove '0x' in the hex str
-            meta={
-                "_dd.p.tid": hex(
-                    _deterministic_sha256_hash(execution_id, HIGHER_64_BITS)
-                )[2:]
-            },
+            meta=meta,
         )
     except Exception as e:
         logger.debug("The Step Functions trace extractor returned with error %s", e)

tests/test_tracing.py

@@ -617,7 +617,7 @@ def test_with_complete_datadog_trace_headers_with_trigger_tags(self):
     @with_trace_propagation_style("datadog")
     def test_step_function_trace_data(self):
         lambda_ctx = get_mock_context()
-        sqs_event = {
+        sf_event = {
             "Execution": {
                 "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
             },
@@ -627,7 +627,7 @@ def test_step_function_trace_data(self):
                 "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
             },
         }
-        ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx)
+        ctx, source, event_source = extract_dd_trace_context(sf_event, lambda_ctx)
         self.assertEqual(source, "event")
         expected_context = Context(
             trace_id=3675572987363469717,
@@ -642,7 +642,48 @@ def test_step_function_trace_data(self):
                 TraceHeader.TRACE_ID: "3675572987363469717",
                 TraceHeader.PARENT_ID: "10713633173203262661",
                 TraceHeader.SAMPLING_PRIORITY: "1",
-                "x-datadog-tags": "_dd.p.tid=e987c84b36b11ab",
+                TraceHeader.TAGS: "_dd.p.tid=e987c84b36b11ab",
+            },
+        )
+        create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
+        self.mock_send_segment.assert_called_with(
+            XraySubsegment.TRACE_KEY,
+            expected_context,
+        )
+
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data(self):
+        lambda_ctx = get_mock_context()
+        sf_event = {
+            "Execution": {
+                "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+            },
+            "StateMachine": {},
+            "State": {
+                "Name": "my-awesome-state",
+                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+            },
+            "_datadog": {
+                "x-datadog-trace-id": "4061173386180447114",
+                "x-datadog-tags": "_dd.p.tid=aac3639aa724716",
+            },
+        }
+        ctx, source, event_source = extract_dd_trace_context(sf_event, lambda_ctx)
+        self.assertEqual(source, "event")
+        expected_context = Context(
+            trace_id=4061173386180447114,
+            span_id=6880978411788117524,
+            sampling_priority=1,
+            meta={"_dd.p.tid": "aac3639aa724716"},
+        )
+        self.assertEqual(ctx, expected_context)
+        self.assertEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "4061173386180447114",
+                TraceHeader.PARENT_ID: "10713633173203262661",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+                TraceHeader.TAGS: "_dd.p.tid=aac3639aa724716",
             },
         )
         create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)