From c0d13794784cff4bc9d3853442db4d6828fb9ea6 Mon Sep 17 00:00:00 2001 From: Ishara Shanmugasundaram Date: Thu, 11 Jun 2026 13:27:21 -0400 Subject: [PATCH] fix: honor upstream trace context in Step Functions Execution.Input._datadog Legacy Step Functions invocations were ignoring upstream Datadog trace headers nested under Execution.Input._datadog and always hashing Execution.Id instead. Extract shared helper and read _datadog before falling back to deterministic trace id generation. APMSVLS-513 Co-authored-by: Cursor --- datadog_lambda/tracing.py | 30 ++++++++--- tests/test_tracing.py | 101 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 7 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 562258f7..b3f79a96 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -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 + + 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 @@ -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) @@ -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) + 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 diff --git a/tests/test_tracing.py b/tests/test_tracing.py index e0f5fb28..fc18f6e5 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -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"""