Propagate Step Function Trace Context through Managed Services (#573)

Allows us to extract Step Function trace context in the following cases
1. SFN -> EventBridge -> Lambda
2. SFN -> EventBridge -> SQS -> Lambda
3. SFN -> SQS -> Lambda
4. SFN -> SNS -> Lambda
5. SFN -> SNS -> SQS -> Lambda

Author: avedmala

datadog_lambda/tracing.py

@@ -39,6 +39,7 @@
     _EventSource,
     parse_event_source,
     get_first_record,
+    is_step_function_event,
     EventTypes,
     EventSubtypes,
 )
@@ -271,6 +272,15 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
 
             if dd_json_data:
                 dd_data = json.loads(dd_json_data)
+
+                if is_step_function_event(dd_data):
+                    try:
+                        return extract_context_from_step_functions(dd_data, None)
+                    except Exception:
+                        logger.debug(
+                            "Failed to extract Step Functions context from SQS/SNS event."
+                        )
+
                 return propagator.extract(dd_data)
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
@@ -313,19 +323,39 @@ 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(dd_context, 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)
+
+        try:
+            return extract_context_from_step_functions(dd_context, None)
+        except Exception:
+            logger.debug(
+                "Failed to extract Step Functions context from EventBridge event."
+            )
+
         return propagator.extract(dd_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
@@ -424,7 +454,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,25 +465,25 @@ def extract_context_from_step_functions(event, lambda_context):
     object.
     """
     try:
+        event = event.get("Payload", event)
+        event = event.get("_datadog", event)
+
         meta = {}
-        dd_data = event.get("_datadog")
 
-        if dd_data and dd_data.get("serverless-version") == "v1":
-            if "x-datadog-trace-id" in dd_data:  # lambda root
-                trace_id = int(dd_data.get("x-datadog-trace-id"))
-                high_64_bit_trace_id = _parse_high_64_bits(
-                    dd_data.get("x-datadog-tags")
-                )
+        if event.get("serverless-version") == "v1":
+            if "x-datadog-trace-id" in event:  # lambda root
+                trace_id = int(event.get("x-datadog-trace-id"))
+                high_64_bit_trace_id = _parse_high_64_bits(event.get("x-datadog-tags"))
                 if high_64_bit_trace_id:
                     meta["_dd.p.tid"] = high_64_bit_trace_id
             else:  # sfn root
-                root_execution_id = dd_data.get("RootExecutionId")
+                root_execution_id = event.get("RootExecutionId")
                 trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS)
                 meta["_dd.p.tid"] = _generate_sfn_trace_id(
                     root_execution_id, HIGHER_64_BITS
                 )
 
-            parent_id = _generate_sfn_parent_id(dd_data)
+            parent_id = _generate_sfn_parent_id(event)
         else:
             execution_id = event.get("Execution").get("Id")
             trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
@@ -472,20 +502,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 +1325,18 @@ 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")
+    dt_format = "%Y-%m-%dT%H:%M:%SZ"
+
+    # Use more granular timestamp from upstream Step Function if possible
+    try:
+        if is_step_function_event(event.get("detail")):
+            timestamp = event["detail"]["_datadog"]["State"]["EnteredTime"]
+            dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
+    except (TypeError, KeyError, AttributeError):
+        logger.debug("Error parsing timestamp from Step Functions event")
+
     dt = datetime.strptime(timestamp, dt_format)
 
     tracer.set_tags(_dd_origin)
@@ -1320,6 +1346,11 @@ 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
+    if dd_trace_context.span_id:
+        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 "Execution" in event and "StateMachine" in event and "State" in event:
+        return True
+
+    # JSONata style
+    dd_context = event.get("_datadog")
+    return (
+        dd_context
+        and "Execution" in dd_context
+        and "StateMachine" in dd_context
+        and "State" in dd_context
+        and "serverless-version" in dd_context
+    )

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,
@@ -286,8 +285,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(