feat(bottlecap): Update Step Functions trace extraction#591

Merged

avedmala merged 9 commits intomainfrom

avedmala/sfn-trace-context-update

Mar 10, 2025

Contributor

avedmala commented Mar 6, 2025 •

edited

Loading

What

Updates the creation of trace context from a Step Function execution's context object to...

Make use of State.RetryCount and Execution.RedriveCount for inferable parent ID generation in a backwards compatible manner (related to Update Step Functions Parent ID Generation datadog-lambda-python#559)
Support the new trace context propagation for multi-level trace merging. These are cases where we're receiving a Step Function event but that Step Function has another parent service which can be either a Lambda or another Step Function (related Explicit trace ID propagation for SFN w/o Hashing datadog-lambda-python#537)

Why

Parity of DataDog/datadog-agent#34264

Using these new values for parent ID generation will prevent collisions with "retry spans" which are Step Function spans that parent a Lambda. Without using the State.RetryCount the other values we use for the hash are identical across tries
Multi-level trace merging is the future, it lets us merge an arbitrary number of Lambda and Step Function traces. The previous approach only lets us do a max depth of 2 services, losing the context after that. We're able to always preserve some about the top-most service or the root service to keep the Trace ID intact while using the context object to infer the parent ID.

Testing

Expanded unit tests to account for new cases

Live trace of the nested step function trace merging (SFN -> SFN -> Lambda)

Live trace of the lambda root trace merging (Lambda -> SFN -> SFN -> Lambda)

avedmala commented

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs

-              #[allow(clippy::module_name_repetitions)]
-              #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
-              pub struct LegacyStepFunctionEvent {

Contributor Author

avedmala Mar 6, 2025 •

edited

Loading

This was never used as we were just unpacking "Payload": {...} if it existed rather than treating them as separate events

Contributor

duncanista Mar 7, 2025

Oh good catch, it was being covered, just not through this struct

avedmala commented

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs

-                          format!("{hi_tid:x}"),
-                      )]);
+                      let (lo_tid, tags) =
+                          if let (Some(trace_id), Some(trace_tags)) = (&self.trace_id, &self.trace_tags) {

Contributor Author

avedmala Mar 6, 2025

We wouldn't need this if-else checking if we defined separate events rather than including all fields in one and making some optional

But, this way we're keeping all the complexity in this file, otherwise we'd have to match on all the different types of events in the span_inferrer

This logic will need infrequent edits so it made sense to me to pull it down similar to what we do in the python layer

avedmala marked this pull request as ready for review

March 6, 2025 22:46

avedmala requested a review from a team as a code owner

March 6, 2025 22:46

astuyve approved these changes

View reviewed changes

Contributor

astuyve left a comment

Looks pretty solid and well tested!

duncanista reviewed

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs Outdated Show resolved Hide resolved

duncanista reviewed

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs Outdated

Comment on lines 64 to 68

+                      let p = payload
+                          .get("Payload")
+                          .unwrap_or(&payload)
+                          .get("_datadog")
+                          .unwrap_or(payload.get("Payload").unwrap_or(&payload));

Contributor

duncanista Mar 7, 2025

This can be simplified

duncanista reviewed

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs Outdated Show resolved Hide resolved

duncanista reviewed

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs Outdated Show resolved Hide resolved

duncanista reviewed

View reviewed changes

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs Outdated Show resolved Hide resolved

duncanista reviewed

View reviewed changes

Contributor

duncanista left a comment

Left some comments!

avedmala added 9 commits

March 10, 2025 13:47


          add new test data

94b4ecb


          handling for retry_count + redrive_count and new optionals

c19e8c9


          implement get_span_context for new cases

cb10a31


          fix lint

935f8e4


          re-used constants and existing tag extractor

4b23b18


          remove super call

d53417e


          fix import

f50b5a8


          empty commit

df7af65


          empty commit

e8cfc56

avedmala force-pushed the avedmala/sfn-trace-context-update branch from ad55c16 to e8cfc56 Compare

March 10, 2025 17:48

Contributor Author

avedmala commented Mar 10, 2025

/merge

dd-devflow bot commented Mar 10, 2025 •

edited

Loading

View all feedbacks in Devflow UI.
2025-03-10 17:48:26 UTC ℹ️ Start processing command /merge

2025-03-10 17:48:32 UTC ℹ️ MergeQueue: pull request added to the queue

The median merge time in main is 0s.

2025-03-10 17:50:07 UTC ⚠️ MergeQueue: This merge request build was cancelled

abhinav.vedmala@datadoghq.com cancelled this merge request build

dd-devflow bot added mergequeue-status: queued mergequeue-status: in_progress and removed mergequeue-status: queued labels

Contributor Author

avedmala commented Mar 10, 2025

/remove

dd-devflow bot commented Mar 10, 2025 •

edited

Loading

View all feedbacks in Devflow UI.
2025-03-10 17:49:58 UTC ℹ️ Start processing command /remove

2025-03-10 17:50:04 UTC ℹ️ Devflow: /remove

dd-devflow bot added mergequeue-status: removed and removed mergequeue-status: in_progress labels

avedmala merged commit f3c3d2e into main

28 of 33 checks passed

avedmala deleted the avedmala/sfn-trace-context-update branch

March 10, 2025 17:54

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mergequeue-status: removed