using context object instead of pre-hashed values

Author: avedmala

datadog_lambda/tracing.py

@@ -356,11 +356,8 @@ def extract_context_from_kinesis_event(event, lambda_context):
     return extract_context_from_lambda_context(lambda_context)
 
 
-def _deterministic_sha256_hash(s: str, part: str) -> (int, int):
-    return _sha256_to_binary_part(hashlib.sha256(s.encode()).hexdigest(), part)
-
-
-def _sha256_to_binary_part(sha256_hash: str, part: str) -> (int, int):
+def _deterministic_sha256_hash(s: str, part: str) -> int:
+    sha256_hash = hashlib.sha256(s.encode()).hexdigest()
     # First two chars is '0b'. zfill to ensure 256 bits, but we only care about the first 128 bits
     binary_hash = bin(int(sha256_hash, 16))[2:].zfill(256)
     if part == HIGHER_64_BITS:
@@ -373,35 +370,56 @@ def _sha256_to_binary_part(sha256_hash: str, part: str) -> (int, int):
     return result
 
 
+def _parse_high_64_bits(trace_tags: str) -> str:
+    # todo: testme
+    if trace_tags:
+        for tag in trace_tags.split(","):
+            if "_dd.p.tid=" in tag:
+                return tag.split("=")[1]
+
+    return ""
+
+
+def _sfn_context_to_parent_id(context: dict) -> int:
+    # todo: testme
+    execution_id = context.get("Execution").get("Id")
+    state_name = context.get("State").get("Name")
+    state_entered_time = context.get("State").get("EnteredTime")
+
+    return _deterministic_sha256_hash(
+        f"{execution_id}#{state_name}#{state_entered_time}", HIGHER_64_BITS
+    )
+
+
 def extract_context_from_step_functions(event, lambda_context):
     """
     Only extract datadog trace context when Step Functions Context Object is injected
     into lambda's event dict.
     """
+    # todo: update docstring
     try:
         meta = {}
+        dd_data = event.get("_datadog")
 
-        if "_datadog" in event:
-            dd_data = event.get("_datadog")
-            parent_id = _sha256_to_binary_part(
-                dd_data.get("x-datadog-parent-id-hash"), HIGHER_64_BITS
-            )
+        if dd_data and dd_data.get("serverless-version") == "v2":
             if "x-datadog-trace-id" in dd_data:  # lambda root
-                dd_data["x-datadog-parent-id"] = str(parent_id)
-                return propagator.extract(dd_data)
+                trace_id = dd_data.get("x-datadog-trace-id")
+                high_64_bit_trace_id = _parse_high_64_bits(dd_data.get("x-datadog-tags"))
+                if high_64_bit_trace_id:
+                    meta["_dd.p.tid"] = high_64_bit_trace_id
             else:  # sfn root
-                trace_id = _sha256_to_binary_part(
-                    dd_data.get("x-datadog-trace-id-hash"), LOWER_64_BITS
+                trace_id = _deterministic_sha256_hash(
+                    dd_data.get("RootExecutionId"), LOWER_64_BITS
                 )
                 meta["_dd.p.tid"] = hex(
-                    _sha256_to_binary_part(
-                        dd_data.get("x-datadog-trace-id-hash"), HIGHER_64_BITS
+                    _deterministic_sha256_hash(
+                        dd_data.get("RootExecutionId"), HIGHER_64_BITS
                     )
                 )[2:]
+
+            parent_id = _sfn_context_to_parent_id(dd_data)
         else:
             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
@@ -412,9 +430,7 @@ def extract_context_from_step_functions(event, lambda_context):
                 _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
-            )
+            parent_id = _sfn_context_to_parent_id(event)
 
         sampling_priority = SamplingPriority.AUTO_KEEP
         return Context(

tests/test_tracing.py

@@ -656,22 +656,29 @@ def test_step_function_trace_data_lambda_root(self):
         lambda_ctx = get_mock_context()
         sfn_event = {
             "_datadog": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                },
                 "x-datadog-trace-id": "5821803790426892636",
                 "x-datadog-sampling-priority": "1",
                 "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
                 "traceparent": "00-672a7cb10000000050cb33b3c06ae95c-5fda9d8d1d1373f9-01",
                 "tracestate": "dd=p:5fda9d8d1d1373f9;s:1;t.dm:-0;t.tid:672a7cb100000000",
-                "x-datadog-parent-id-hash": "a926584eba705d6ec904c54db2ecc4d4a2c91e7dabe7ce87ac26edb43388fbc5",
                 "serverless-version": "v2",
             }
         }
         ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
         self.assertEqual(source, "event")
         expected_context = Context(
-            trace_id=137131089076080415507232535361568303452,
-            span_id=2965154499828669806,
+            trace_id=5821803790426892636,
+            span_id=6880978411788117524,
             sampling_priority=1,
-            meta={"_dd.p.dm": "-0", "_dd.p.tid": "672a7cb100000000"},
+            meta={"_dd.p.tid": "672a7cb100000000"},
         )
         self.assertEqual(ctx, expected_context)
         self.assertEqual(
@@ -680,7 +687,7 @@ def test_step_function_trace_data_lambda_root(self):
                 TraceHeader.TRACE_ID: "5821803790426892636",
                 TraceHeader.PARENT_ID: "10713633173203262661",
                 TraceHeader.SAMPLING_PRIORITY: "1",
-                TraceHeader.TAGS: "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
+                TraceHeader.TAGS: "_dd.p.tid=672a7cb100000000",
             },
         )
         create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
@@ -694,27 +701,34 @@ def test_step_function_trace_data_sfn_root(self):
         lambda_ctx = get_mock_context()
         sfn_event = {
             "_datadog": {
-                "x-datadog-trace-id-hash": "fed93f8c162880cb9aa90fcd1f8395383835841d5470d30215f3dd52906ebc58",
-                "x-datadog-parent-id-hash": "c5eb94cc9220ab5783e1db53debd54b8c93f6f2a3eae1c680d7b849f2d34e551",
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                },
+                "RootExecutionId": "4875aba4-ae31-4a4c-bf8a-63e9eee31dad",
                 "serverless-version": "v2",
             }
         }
         ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
         self.assertEqual(source, "event")
         expected_context = Context(
-            trace_id=1921084089721656632,
-            span_id=5038284214489885527,
+            trace_id=4521899030418994483,
+            span_id=6880978411788117524,
             sampling_priority=1,
-            meta={"_dd.p.tid": "7ed93f8c162880cb"},
+            meta={"_dd.p.tid": "12d1270d99cc5e03"},
         )
         self.assertEqual(ctx, expected_context)
         self.assertEqual(
             get_dd_trace_context(),
             {
-                TraceHeader.TRACE_ID: "1921084089721656632",
+                TraceHeader.TRACE_ID: "4521899030418994483",
                 TraceHeader.PARENT_ID: "10713633173203262661",
                 TraceHeader.SAMPLING_PRIORITY: "1",
-                TraceHeader.TAGS: "_dd.p.tid=7ed93f8c162880cb",
+                TraceHeader.TAGS: "_dd.p.tid=12d1270d99cc5e03",
             },
         )
         create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)