Skip to content

Propagate Step Function Trace Context through Managed Services #573

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Mar 19, 2025
Merged

Propagate Step Function Trace Context through Managed Services #573

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
a38c77b
support step function context in eventbridge extraction
avedmala Mar 11, 2025
0a51475
fix circular import
avedmala Mar 11, 2025
5700d22
fix lint
avedmala Mar 11, 2025
f524509
fix parent_id value in unit test
avedmala Mar 11, 2025
55ff075
use provided timestamp from step function
avedmala Mar 11, 2025
d7daf0d
extract context in eventbridge sqs case
avedmala Mar 12, 2025
e490610
Update datadog_lambda/trigger.py
avedmala Mar 13, 2025
e31c8f8
add sqs/sns support
avedmala Mar 13, 2025
dc4a283
add more test cases
avedmala Mar 13, 2025
f59de19
fix expected _dd.p.tid vals
avedmala Mar 13, 2025
6222016
add lru cache and remove usage of all()
avedmala Mar 13, 2025
45ba945
exception handling for timestamp parsing
avedmala Mar 13, 2025
dc43c9e
remove memo
avedmala Mar 13, 2025
eec4743
refactor tests to use common helper
avedmala Mar 17, 2025
1b350b9
removed event_type assertion
avedmala Mar 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 56 additions & 25 deletions datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
_EventSource,
parse_event_source,
get_first_record,
is_step_function_event,
EventTypes,
EventSubtypes,
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am curious about how the concatenation of two queues (e.g., SFN → EventBridge → SQS → Lambda) is handled. Is it achieved by extracting two different contexts in the Python tracer? Does this mean that it also supports SFN → EventBridge → SQS → SNS → Lambda?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SFN → EventBridge → SQS → Lambda is handled the following way

  • Final event that enters the Lambda is an SQS event
  • Our context extractor for SQS events checks if there's an EventBridge event within and uses that if its valid

SFN → SNS → SQS → Lambda is handled very similarly with another explicitly check in the SQS extractor looking for SNS events nested

We don't handle SFN → SQS → SNS → Lambda AFAIK but we wouldn't be able to handle SFN → EventBridge → SQS → SNS → Lambda out of the box either

But this is only because it's not explicitly handled. The current python layer implementation is messy because it relies on explicit handling. I think a perfect solution would be one where it's all handled recursively and customers can nest an arbitrary number of supported services without explicit handling

I think AWS team would like to do something like this in bottlecap

Copy link

@nine5two7 nine5two7 Mar 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@avedmala Thanks for the explanation. Very informative. I am guessing that a recursive solution should not be that complicated? @purple4reina @joeyzhao2018

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Regarding the "recursive solution", is it written down in any RFC? it sounds interesting and might be able to solve some other problems.

"""
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)
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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


Expand Down
30 changes: 27 additions & 3 deletions datadog_lambda/trigger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way we can memoize this function? It looks like it can potentially be called several times in the course of a single invocation.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, or it looks like the function can be called multiple times per invocation, but with different "events" each time? If that's true, then we can probably leave it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a great idea!

Correct me if I'm wrong but does the layer only handle one event per invocation? Or if it's a busy Lambda does it stay alive and potentially handle hundreds of events?

Just wondering to get an idea of how large to make the cache. I guess it can be pretty small anyway since each event is new and we don't repeat

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each runtime instance will only ever handle one event at a time. It never handles two events concurrently.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah just realized we can't memoize it because event is a dict and mutable types are unhashable

We could serialize the dict and use that but I'm thinking that'd be much slower

"""
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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really like these comments. For someone who hasn't work on step functions for a while, these comments help me recollect these historical context. It'll help future maintenance of the code as well.

"""
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
)
3 changes: 0 additions & 3 deletions datadog_lambda/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@
is_authorizer_response,
tracer,
propagator,
is_legacy_lambda_step_function,
)
from datadog_lambda.trigger import (
extract_trigger_tags,
Expand Down Expand Up @@ -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"]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved this unwrapping to happen inside of tracing.extract_context_from_step_functions()

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(
Expand Down
Loading
Loading