durable: extract trace context from checkpoints and input payload

datadog_lambda/tracing.py

@@ -61,6 +61,7 @@
 DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000"
 HIGHER_64_BITS = "HIGHER_64_BITS"
 LOWER_64_BITS = "LOWER_64_BITS"
+_TRACE_CHECKPOINT_PREFIX = "_datadog_"
 
 
 def _dsm_set_checkpoint(context_json, event_type, arn):
@@ -546,6 +547,54 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
+def _extract_context_from_durable_checkpoint(operation):
+    if not isinstance(operation, dict):
+        return None
+
+    step_details = operation.get("StepDetails")
+    if not isinstance(step_details, dict):
+        return None
+
+    result = step_details.get("Result")
+    if isinstance(result, str):
+        try:
+            result = json.loads(result)
+        except Exception:
+            return None
+
+    if not isinstance(result, dict):
+        return None
+
+    return propagator.extract(result)
+
+
+def extract_context_from_durable_execution(event):
+    operations = event.get("InitialExecutionState", {}).get("Operations")
+    if isinstance(operations, dict):
+        operations = list(operations.values())
+    if not isinstance(operations, list) or not operations:
+        return None
+
+    highest = -1
+    best_operation = None
+    for operation in operations:
+        if not isinstance(operation, dict):
+            continue
+        name = operation.get("Name")
+        if not isinstance(name, str) or not name.startswith(_TRACE_CHECKPOINT_PREFIX):
+            continue
+        suffix = name[len(_TRACE_CHECKPOINT_PREFIX) :]
+        try:
+            number = int(suffix)
+        except (TypeError, ValueError):
+            continue
+        if number > highest:
+            highest = number
+            best_operation = operation
+
+    return _extract_context_from_durable_checkpoint(best_operation)
+
+
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function
@@ -633,9 +682,12 @@ def extract_dd_trace_context(
     global dd_trace_context
     trace_context_source = None
     event_source = parse_event_source(event)
+    context = None
 
     if extractor is not None:
         context = extract_context_custom_extractor(extractor, event, lambda_context)
+    elif isinstance(event, dict) and "DurableExecutionArn" in event:
+        context = extract_context_from_durable_execution(event)
     elif isinstance(event, (set, dict)) and "request" in event:
         context = extract_context_from_request_header_or_context(
             event, lambda_context, event_source

tests/test_tracing.py

@@ -394,6 +394,44 @@ def test_with_w3c_trace_headers(self):
             headers, {"headers": headers}
         )
 
+    @with_trace_propagation_style("datadog")
+    def test_extracts_durable_trace_context_from_latest_checkpoint_operation_map(self):
+        lambda_ctx = get_mock_context()
+        headers = {
+            TraceHeader.TRACE_ID: "123",
+            TraceHeader.PARENT_ID: "321",
+            TraceHeader.SAMPLING_PRIORITY: "1",
+        }
+
+        event = {
+            "DurableExecutionArn": "arn:aws:lambda:us-east-2:123456789012:function:demo:1/durable-execution/demo/abc",
+            "CheckpointToken": "token",
+            "InitialExecutionState": {
+                "Operations": {
+                    "0": {"Type": "EXECUTION"},
+                    "1": {
+                        "Name": "_datadog_0",
+                        "StepDetails": {
+                            "Result": {
+                                TraceHeader.TRACE_ID: "999",
+                                TraceHeader.PARENT_ID: "888",
+                                TraceHeader.SAMPLING_PRIORITY: "1",
+                            }
+                        },
+                    },
+                    "2": {
+                        "Name": "_datadog_1",
+                        "StepDetails": {"Result": headers},
+                    },
+                }
+            },
+        }
+
+        ctx, source, _ = extract_dd_trace_context(event, lambda_ctx)
+
+        self.assertEqual(source, "event")
+        self.assertEqual(ctx, Context(trace_id=123, span_id=321, sampling_priority=1))
+
     @with_trace_propagation_style("datadog")
     def test_with_extractor_function(self):
         def extractor_foo(event, context):