diff --git a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs index f65174287..45a725581 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs @@ -247,10 +247,18 @@ pub(crate) fn extract_trace_context_from_aws_trace_header( let mut trace_id = String::new(); let mut parent_id = String::new(); let mut sampled = String::new(); + // `dd-trace-java` is the only tracer that plants context here, and it zeroes the high 32 bits + // of the ID section; an AWS-generated ID has random hex there. It ignores headers without + // that padding on extract (`XRayHttpCodec`), and so did the Go extension (`rootRegex`). + let mut datadog_planted = false; for part in value.split(';') { if part.starts_with("Root=") { - trace_id = part[24..].to_string(); + // `Root=1-<8 hex epoch>-<24 hex id>`: the ID section starts at 16, and the trace ID + // Datadog uses is its low 16 hex digits. + let id_section = part.get(16..)?; + datadog_planted = id_section.starts_with("00000000"); + trace_id = id_section.get(8..).unwrap_or_default().to_string(); } else if let Some(parent_part) = part.strip_prefix("Parent=") { parent_id = parent_part.to_string(); } else if part.starts_with("Sampled=") && sampled.is_empty() { @@ -269,7 +277,18 @@ pub(crate) fn extract_trace_context_from_aws_trace_header( return None; } - let sampling_priority = i8::from(sampled == "1"); + // Whether `Sampled` is a decision we can act on depends on who wrote the header. + let sampling_priority = match (datadog_planted, sampled.as_str()) { + // A keep is safe to honor whoever wrote it: the cost is over-retention, not lost spans. + (_, "1") => Some("1"), + // `dd-trace-java` planted this header, so the drop is Datadog's own decision. + (true, "0") => Some("0"), + // X-Ray's drop, or no decision at all: X-Ray tracing is off, or `dd-trace-java` wrote the + // header before its priority was set, so it omitted `Sampled`. Leaving the priority unset + // lets the tracer sample; sending 0 would drop its spans on someone else's behalf. + _ => None, + } + .and_then(|priority| priority.parse().ok()); Some(SpanContext { // the context from AWS Header is used by Datadog only and does not contain the upper @@ -277,7 +296,7 @@ pub(crate) fn extract_trace_context_from_aws_trace_header( trace_id: u128::from(trace_id), span_id: parent_id, sampling: Sampling { - priority: sampling_priority.to_string().parse().ok(), + priority: sampling_priority, mechanism: None, }, origin: None, @@ -672,8 +691,10 @@ mod tests { SpanContext { trace_id: 130_944_522_478_755_159, span_id: 9_032_698_535_745_367_362, + // This fixture's root ID has non-zero high bits, i.e. AWS generated it rather + // than a Datadog library, so its `Sampled=0` is not propagated as a drop. sampling: Sampling { - priority: "0".parse().ok(), + priority: None, mechanism: None, }, origin: None, @@ -682,4 +703,71 @@ mod tests { } ); } + + #[test] + fn aws_generated_root_id_not_sampled_leaves_priority_unset() { + let context = extract_trace_context_from_aws_trace_header(Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0" + .to_string(), + )) + .expect("failed to extract context"); + + // X-Ray's decision, not Datadog's: keep the IDs, let the tracer sample. + assert_eq!(context.sampling.priority, None); + assert_eq!(context.trace_id, 130_944_522_478_755_159); + } + + #[test] + fn datadog_planted_root_id_not_sampled_drops() { + let context = extract_trace_context_from_aws_trace_header(Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac;Sampled=0" + .to_string(), + )) + .expect("failed to extract context"); + + // Zeroed high bits mean a Datadog library wrote this, so `Sampled=0` is our own decision. + assert_eq!(context.sampling.priority, "0".parse().ok()); + } + + #[test] + fn datadog_planted_root_id_without_sampled_leaves_priority_unset() { + let context = extract_trace_context_from_aws_trace_header(Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac".to_string(), + )) + .expect("failed to extract context"); + + // `dd-trace-java` omits `Sampled` when the priority isn't set yet, so there's no + // decision to carry. + assert_eq!(context.sampling.priority, None); + } + + #[test] + fn datadog_planted_root_id_sampled_keeps() { + let context = extract_trace_context_from_aws_trace_header(Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac;Sampled=1" + .to_string(), + )) + .expect("failed to extract context"); + + assert_eq!(context.sampling.priority, "1".parse().ok()); + } + + #[test] + fn truncated_root_does_not_panic() { + assert!( + extract_trace_context_from_aws_trace_header(Some("Root=1-64cc".to_string())).is_none() + ); + } + + #[test] + fn aws_generated_root_id_sampled_keeps_priority() { + let context = extract_trace_context_from_aws_trace_header(Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=1" + .to_string(), + )) + .expect("failed to extract context"); + + assert_eq!(context.sampling.priority, "1".parse().ok()); + assert_eq!(context.trace_id, 130_944_522_478_755_159); + } }