Skip to content

fix: honor upstream trace context in Step Functions Execution.Input._datadog (APMSVLS-513) #830

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
chemystery09 merged 1 commit into from
Jun 12, 2026
+124 −7
Merged

fix: honor upstream trace context in Step Functions Execution.Input._datadog (APMSVLS-513) #830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 23 additions & 7 deletions datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,21 @@ def _generate_sfn_parent_id(context: dict) -> int:
)


def _extract_dd_trace_id_from_dd_data(dd_data, meta):
"""
Read native Datadog trace headers from a `_datadog` dict.
Returns the 64-bit trace_id and populates meta["_dd.p.tid"] when present.
Returns None if no usable trace id is found.
"""
if not dd_data or "x-datadog-trace-id" not in dd_data:
return None
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 high_64_bit_trace_id:
meta["_dd.p.tid"] = high_64_bit_trace_id
return trace_id
Comment on lines +485 to +497

@rithikanarayan rithikanarayan Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think this is fine. You're reusing the existing code so should be ok



def _generate_sfn_trace_id(execution_id: str, part: str):
"""
Take the SHA-256 hash of the execution_id to calculate the trace ID. If the high 64 bits are
Expand Down Expand Up @@ -517,10 +532,7 @@ def extract_context_from_step_functions(event, lambda_context):

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
trace_id = _extract_dd_trace_id_from_dd_data(event, meta)
else: # sfn root
root_execution_id = event.get("RootExecutionId")
trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS)
Expand All @@ -530,9 +542,13 @@ def extract_context_from_step_functions(event, lambda_context):

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)
meta["_dd.p.tid"] = _generate_sfn_trace_id(execution_id, HIGHER_64_BITS)
execution = event.get("Execution", {})
dd_input = execution.get("Input", {}).get("_datadog")
trace_id = _extract_dd_trace_id_from_dd_data(dd_input, meta)
Comment on lines +545 to +547

@rithikanarayan rithikanarayan Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Same here, should be fine

if trace_id is None:
execution_id = execution.get("Id")
trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
meta["_dd.p.tid"] = _generate_sfn_trace_id(execution_id, HIGHER_64_BITS)
parent_id = _generate_sfn_parent_id(event)

sampling_priority = SamplingPriority.AUTO_KEEP
Expand Down
101 changes: 101 additions & 0 deletions tests/test_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -751,6 +751,107 @@ def test_step_function_trace_data(self):
sfn_event, 435175499815315247, 3929055471293792800, "3e7a89d1b7310603"
)

@with_trace_propagation_style("datadog")
def test_step_function_trace_data_with_input_datadog(self):
"""Legacy SFN context with upstream trace in Execution.Input._datadog"""
sfn_event = {
"Execution": {
"Id": "arn:aws:states:sa-east-1:425362996713:execution:rstrat-sfn-lambda-jsonpath-demo-dev-state-machine:7d67bd1e-57df-4021-8b52-f55d5cad1169",
"Input": {
"hello": "world",
"_datadog": {
"x-datadog-trace-id": "12766418539254701015",
"x-datadog-parent-id": "5497030307431673011",
"x-datadog-sampling-priority": "1",
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=69f4ed3a00000000",
"traceparent": "00-69f4ed3a00000000b12b6e05a63cdbd7-4c495ff0aa056cb3-01",
"tracestate": "dd=p:4c495ff0aa056cb3;s:1;t.dm:-0;t.tid:69f4ed3a00000000",
},
},
"StartTime": "2026-05-01T18:13:14.130Z",
"Name": "7d67bd1e-57df-4021-8b52-f55d5cad1169",
"RoleArn": "arn:aws:iam::425362996713:role/rstrat-sfn-lambda-jsonpat-StepFunctionsExecutionRol-mtfEz4KEzBEE",
"RedriveCount": 0,
},
"State": {
"Name": "InvokeDownstreamSecond",
"EnteredTime": "2026-05-01T18:13:16.310Z",
"RetryCount": 0,
},
"StateMachine": {
"Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:rstrat-sfn-lambda-jsonpath-demo-dev-state-machine",
"Name": "rstrat-sfn-lambda-jsonpath-demo-dev-state-machine",
},
}
self._test_step_function_trace_data_common(
sfn_event, 12766418539254701015, 5718818197702795437, "69f4ed3a00000000"
)

@with_trace_propagation_style("datadog")
def test_step_function_trace_data_input_datadog_without_tid(self):
"""Execution.Input._datadog without _dd.p.tid still preserves upstream trace id"""
sfn_event = {
"Execution": {
"Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"Input": {
"_datadog": {
"x-datadog-trace-id": "5821803790426892636",
},
},
"Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
"StartTime": "2024-12-04T19:38:04.069Z",
"RedriveCount": 0,
},
"State": {
"Name": "Lambda Invoke",
"EnteredTime": "2024-12-04T19:38:04.118Z",
"RetryCount": 0,
},
"StateMachine": {
"Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
"Name": "abhinav-activity-state-machine",
},
}
lambda_ctx = get_mock_context()
expected_context = Context(
trace_id=5821803790426892636,
span_id=3929055471293792800,
sampling_priority=1,
meta={},
)

ctx, source, _ = extract_dd_trace_context(sfn_event, lambda_ctx)

self.assertEqual(source, "event")
self.assertEqual(ctx, expected_context)

@with_trace_propagation_style("datadog")
def test_step_function_trace_data_input_without_datadog(self):
"""Execution.Input without _datadog falls back to deterministic trace id"""
sfn_event = {
"Execution": {
"Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"Input": {"hello": "world"},
"Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
"StartTime": "2024-12-04T19:38:04.069Z",
"RedriveCount": 0,
},
"State": {
"Name": "Lambda Invoke",
"EnteredTime": "2024-12-04T19:38:04.118Z",
"RetryCount": 0,
},
"StateMachine": {
"Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
"Name": "abhinav-activity-state-machine",
},
}
self._test_step_function_trace_data_common(
sfn_event, 435175499815315247, 3929055471293792800, "3e7a89d1b7310603"
)

@with_trace_propagation_style("datadog")
def test_step_function_trace_data_retry(self):
"""Test step function trace data extraction with non-zero retry count"""
Expand Down
Loading