Propagate Step Function Trace Context through Managed Services

datadog_lambda/tracing.py

@@ -39,6 +39,7 @@
     _EventSource,
     parse_event_source,
     get_first_record,
+    is_step_function_event,
     EventTypes,
     EventSubtypes,
 )
@@ -313,19 +314,32 @@ def _extract_context_from_eventbridge_sqs_event(event):
     body = json.loads(body_str)
     detail = body.get("detail")
     dd_context = detail.get("_datadog")
+    if is_step_function_event(dd_context):
+        try:
+            return extract_context_from_step_functions(detail, None)
+        except Exception:
+            logger.debug(
+                "Failed to extract Step Functions context from EventBridge to SQS event."
+            )
+
     return propagator.extract(dd_context)
 
 
 def extract_context_from_eventbridge_event(event, lambda_context):
     """
     Extract datadog trace context from an EventBridge message's Details.
     This is only possible if Details is a JSON string.
+
+    If we find a Step Function context, try to extract the trace context from
+    that header.
     """
     try:
         detail = event.get("detail")
         dd_context = detail.get("_datadog")
         if not dd_context:
             return extract_context_from_lambda_context(lambda_context)
+        if is_step_function_event(dd_context):
+            return extract_context_from_step_functions(detail, lambda_context)
         return propagator.extract(dd_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
@@ -424,7 +438,7 @@ def _generate_sfn_trace_id(execution_id: str, part: str):
 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.
+    into lambda's event dict. Unwrap "Payload" if it exists to handle Legacy Lambda cases.
 
     If '_datadog' header is present, we have two cases:
       1. Root is a Lambda and we use its traceID
@@ -435,6 +449,8 @@ def extract_context_from_step_functions(event, lambda_context):
     object.
     """
     try:
+        event = event.get("Payload", event)
+
         meta = {}
         dd_data = event.get("_datadog")
 
@@ -472,20 +488,6 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
-def is_legacy_lambda_step_function(event):
-    """
-    Check if the event is a step function that called a legacy lambda
-    """
-    if not isinstance(event, dict) or "Payload" not in event:
-        return False
-
-    event = event.get("Payload")
-    return isinstance(event, dict) and (
-        "_datadog" in event
-        or ("Execution" in event and "StateMachine" in event and "State" in event)
-    )
-
-
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function
@@ -1309,8 +1311,15 @@ def create_inferred_span_from_eventbridge_event(event, context):
         synchronicity="async",
         tag_source="self",
     )
-    dt_format = "%Y-%m-%dT%H:%M:%SZ"
-    timestamp = event.get("time")
+
+    # Use more granular timestamp from upstream Step Function if possible
+    if is_step_function_event(event.get("detail")):
+        timestamp = event.get("detail").get("_datadog").get("State").get("EnteredTime")
+        dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
+    else:
+        timestamp = event.get("time")
+        dt_format = "%Y-%m-%dT%H:%M:%SZ"
+
     dt = datetime.strptime(timestamp, dt_format)
 
     tracer.set_tags(_dd_origin)
@@ -1320,6 +1329,10 @@ def create_inferred_span_from_eventbridge_event(event, context):
     if span:
         span.set_tags(tags)
     span.start = dt.replace(tzinfo=timezone.utc).timestamp()
+
+    # Since inferred span will later parent Lambda, preserve Lambda's current parent
+    span.parent_id = dd_trace_context.span_id
+
     return span
 
 

datadog_lambda/trigger.py

@@ -146,9 +146,7 @@ def parse_event_source(event: dict) -> _EventSource:
     if event.get("source") == "aws.events" or has_event_categories:
         event_source = _EventSource(EventTypes.CLOUDWATCH_EVENTS)
 
-    if (
-        "_datadog" in event and event.get("_datadog").get("serverless-version") == "v1"
-    ) or ("Execution" in event and "StateMachine" in event and "State" in event):
+    if is_step_function_event(event):
         event_source = _EventSource(EventTypes.STEPFUNCTIONS)
 
     event_record = get_first_record(event)
@@ -369,3 +367,29 @@ def extract_http_status_code_tag(trigger_tags, response):
         status_code = response.status_code
 
     return str(status_code)
+
+
+def is_step_function_event(event):
+    """
+    Check if the event is a step function that invoked the current lambda.
+
+    The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a
+    "_datadog" for JSONata style context propagation.
+
+    The actual event must contain "Execution", "StateMachine", and "State" fields.
+    """
+    event = event.get("Payload", event)
+
+    # JSONPath style
+    if all(field in event for field in ("Execution", "StateMachine", "State")):
+        return True
+
+    # JSONata style
+    if "_datadog" in event:
+        event = event["_datadog"]
+        return all(
+            field in event
+            for field in ("Execution", "StateMachine", "State", "serverless-version")
+        )
+
+    return False

datadog_lambda/wrapper.py

@@ -45,7 +45,6 @@
     is_authorizer_response,
     tracer,
     propagator,
-    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -279,8 +278,6 @@ def _before(self, event, context):
             self.response = None
             set_cold_start(init_timestamp_ns)
             submit_invocations_metric(context)
-            if is_legacy_lambda_step_function(event):
-                event = event["Payload"]
             self.trigger_tags = extract_trigger_tags(event, context)
             # Extract Datadog trace context and source from incoming requests
             dd_context, trace_context_source, event_source = extract_dd_trace_context(

tests/test_tracing.py

@@ -42,7 +42,6 @@
     service_mapping as global_service_mapping,
     propagator,
     emit_telemetry_on_exception_outside_of_handler,
-    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import EventTypes
 
@@ -836,6 +835,55 @@ def test_step_function_trace_data_sfn_root(self):
             expected_context,
         )
 
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_event_bridge(self):
+        lambda_ctx = get_mock_context()
+        sfn_event = {
+            "_datadog": {
+                "Execution": {
+                    "StartTime": "2025-03-11T01:16:31.408Z",
+                    "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+                    "RedriveCount": 0,
+                    "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
+                    "Name": "eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+                },
+                "StateMachine": {
+                    "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-inner-state-machine",
+                    "Name": "abhinav-inner-state-machine",
+                },
+                "State": {
+                    "EnteredTime": "2025-03-11T01:16:31.448Z",
+                    "RetryCount": 0,
+                    "Name": "EventBridge PutEvents",
+                },
+                "serverless-version": "v1",
+                "RootExecutionId": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-inner-state-machine:eb6298d0-93b5-4fe0-8af9-fefe2933b0ed",
+            }
+        }
+        ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
+        self.assertEqual(source, "event")
+        expected_context = Context(
+            trace_id=4728686021345621131,
+            span_id=2685222157636933868,
+            sampling_priority=1,
+            meta={"_dd.p.tid": "7683d2257c051fce"},
+        )
+        self.assertEqual(ctx, expected_context)
+        self.assertEqual(
+            get_dd_trace_context(),
+            {
+                TraceHeader.TRACE_ID: "4728686021345621131",
+                TraceHeader.PARENT_ID: "10713633173203262661",
+                TraceHeader.SAMPLING_PRIORITY: "1",
+                TraceHeader.TAGS: "_dd.p.tid=7683d2257c051fce",
+            },
+        )
+        create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
+        self.mock_send_segment.assert_called_with(
+            XraySubsegment.TRACE_KEY,
+            expected_context,
+        )
+
 
 class TestXRayContextConversion(unittest.TestCase):
     def test_convert_xray_trace_id(self):

tests/test_trigger.py

@@ -9,6 +9,7 @@
     get_event_source_arn,
     extract_trigger_tags,
     extract_http_status_code_tag,
+    is_step_function_event,
 )
 
 from tests.utils import get_mock_context
@@ -543,3 +544,68 @@ def test_extract_http_status_code_tag_from_response_object(self):
         response.status_code = 403
         status_code = extract_http_status_code_tag(trigger_tags, response)
         self.assertEqual(status_code, "403")
+
+
+class IsStepFunctionEvent(unittest.TestCase):
+    def test_is_step_function_event_jsonata(self):
+        event = {
+            "_datadog": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                    "RedriveCount": 0,
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                    "RetryCount": 0,
+                },
+                "x-datadog-trace-id": "5821803790426892636",
+                "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
+                "serverless-version": "v1",
+            }
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_jsonpath(self):
+        event = {
+            "Execution": {
+                "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                "RedriveCount": 0,
+            },
+            "StateMachine": {},
+            "State": {
+                "Name": "my-awesome-state",
+                "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                "RetryCount": 0,
+            },
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_legacy_lambda(self):
+        event = {
+            "Payload": {
+                "Execution": {
+                    "Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
+                    "RedriveCount": 0,
+                },
+                "StateMachine": {},
+                "State": {
+                    "Name": "my-awesome-state",
+                    "EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
+                    "RetryCount": 0,
+                },
+            }
+        }
+        self.assertTrue(is_step_function_event(event))
+
+    def test_is_step_function_event_dd_header(self):
+        event = {
+            "_datadog": {
+                "x-datadog-trace-id": "5821803790426892636",
+                "x-datadog-parent-id": "5821803790426892636",
+                "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
+                "x-datadog-sampling-priority": "1",
+            }
+        }
+        self.assertFalse(is_step_function_event(event))