use redrive count to generate step functions parent ID

Author: avedmala

datadog_lambda/tracing.py

@@ -384,12 +384,21 @@ def _parse_high_64_bits(trace_tags: str) -> str:
 
 
 def _generate_sfn_parent_id(context: dict) -> int:
+    """
+    The upstream Step Function can propagate its execution context to downstream Lambdas. The Lambda can use these
+    details to share the same traceID and infer its parent's spanID.
+
+    Excluding redriveCount when its 0 to account for cases where customers are using an old version of the Lambda layer
+    that doesn't use this value for its parentID generation.
+    """
     execution_id = context.get("Execution").get("Id")
+    redrive_count = context.get("Execution").get("RedriveCount")
     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
+        f"{execution_id}#{state_name}#{state_entered_time}#{redrive_count}",
+        HIGHER_64_BITS,
     )
 
 

tests/test_tracing.py

@@ -631,7 +631,7 @@ def test_step_function_trace_data(self):
         self.assertEqual(source, "event")
         expected_context = Context(
             trace_id=3675572987363469717,
-            span_id=6880978411788117524,
+            span_id=4929949072763648481,
             sampling_priority=1,
             meta={"_dd.p.tid": "e987c84b36b11ab"},
         )
@@ -673,7 +673,7 @@ def test_step_function_trace_data_lambda_root(self):
         self.assertEqual(source, "event")
         expected_context = Context(
             trace_id=5821803790426892636,
-            span_id=6880978411788117524,
+            span_id=4929949072763648481,
             sampling_priority=1,
             meta={"_dd.p.tid": "672a7cb100000000"},
         )
@@ -714,7 +714,7 @@ def test_step_function_trace_data_sfn_root(self):
         self.assertEqual(source, "event")
         expected_context = Context(
             trace_id=4521899030418994483,
-            span_id=6880978411788117524,
+            span_id=4929949072763648481,
             sampling_priority=1,
             meta={"_dd.p.tid": "12d1270d99cc5e03"},
         )

tests/test_xray.py

@@ -34,9 +34,9 @@ def test_get_xray_host_port_success(self):
 
     def test_send_segment_sampled_out(self):
         os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
-        os.environ[
-            "_X_AMZN_TRACE_ID"
-        ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0;Lineage=c6c5b1b9:0"
+        os.environ["_X_AMZN_TRACE_ID"] = (
+            "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0;Lineage=c6c5b1b9:0"
+        )
 
         with patch(
             "datadog_lambda.xray.sock.send", MagicMock(return_value=None)
@@ -47,9 +47,9 @@ def test_send_segment_sampled_out(self):
 
     def test_send_segment_sampled(self):
         os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
-        os.environ[
-            "_X_AMZN_TRACE_ID"
-        ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
+        os.environ["_X_AMZN_TRACE_ID"] = (
+            "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
+        )
         with patch(
             "datadog_lambda.xray.sock.send", MagicMock(return_value=None)
         ) as mock_send: