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(