From 006b1f7bffe734f5a155801100ff47d7314c9d67 Mon Sep 17 00:00:00 2001 From: Yiming Luo <10097700+lym953@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:29:37 -0400 Subject: [PATCH] [SLES-2971] feat(traces): gate X-Ray-driven sampling behind DD_MERGE_XRAY_TRACES Adopting an AWS-generated X-Ray `Sampled` flag as a Datadog sampling decision is opt-in behavior in the tracer libraries, where `DD_MERGE_XRAY_TRACES` defaults to false. bottlecap never read it, so there was no way to decline. Add the config and gate the flag on it. A header a Datadog library planted is never gated -- its `Sampled` is a Datadog decision regardless of the setting, and gating it would break dd-trace-java's SQS propagation, which uses this same header and has nothing to do with X-Ray. Scope is the sampling verdict only. Trace and parent IDs are still taken from an AWS-generated header either way, so correlation is unchanged by this flag; that is pre-existing v88+ behavior and gating it is a separate argument. The flag is read off `DatadogCompositePropagator`, which already carries the config to every extraction site -- threading `Arc` through the listener, LWA proxy, and interceptor state instead would have touched five more files. --- bottlecap/src/config/mod.rs | 10 ++ .../src/lifecycle/invocation/span_inferrer.rs | 21 ++-- .../invocation/triggers/sqs_event.rs | 112 ++++++++++++++---- bottlecap/src/traces/propagation/mod.rs | 8 ++ 4 files changed, 121 insertions(+), 30 deletions(-) diff --git a/bottlecap/src/config/mod.rs b/bottlecap/src/config/mod.rs index db0a1f192..d2794217a 100644 --- a/bottlecap/src/config/mod.rs +++ b/bottlecap/src/config/mod.rs @@ -80,6 +80,11 @@ pub struct LambdaConfig { /// without durable execution context enrichment. Defaults to 0 until the tracer-side /// durable execution support is released; set to 50 to re-enable enrichment. pub lambda_durable_function_log_buffer_size: usize, + + /// `DD_MERGE_XRAY_TRACES` — opt in to letting X-Ray's own sampling decision drive Datadog's. + /// Off by default, matching the tracer libraries. Only affects headers X-Ray generated: a + /// header a Datadog library planted carries a Datadog decision and is never gated by this. + pub merge_xray_traces: bool, } impl Default for LambdaConfig { @@ -104,6 +109,7 @@ impl Default for LambdaConfig { api_security_sample_delay: Duration::from_secs(30), custom_metrics_exclude_tags: Vec::new(), lambda_durable_function_log_buffer_size: 0, + merge_xray_traces: false, } } } @@ -180,6 +186,9 @@ pub struct LambdaConfigSource { /// 0 (hold mechanism disabled). #[serde(deserialize_with = "deser_opt_lossless")] pub lambda_durable_function_log_buffer_size: Option, + + #[serde(deserialize_with = "deser_opt_bool")] + pub merge_xray_traces: Option, } impl DatadogConfigExtension for LambdaConfig { @@ -204,6 +213,7 @@ impl DatadogConfigExtension for LambdaConfig { api_security_enabled, api_security_sample_delay, lambda_durable_function_log_buffer_size, + merge_xray_traces, ], option: [span_dedup_timeout, api_key_secret_reload_interval, appsec_rules], ); diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs index 50e80ba5a..2a722510f 100644 --- a/bottlecap/src/lifecycle/invocation/span_inferrer.rs +++ b/bottlecap/src/lifecycle/invocation/span_inferrer.rs @@ -199,14 +199,17 @@ impl SpanInferrer { } #[must_use] - fn should_skip_inferred_span(identified_trigger: &IdentifiedTrigger) -> bool { + fn should_skip_inferred_span( + identified_trigger: &IdentifiedTrigger, + merge_xray_traces: bool, + ) -> bool { match identified_trigger { // There is no inferred span for ALB events IdentifiedTrigger::ALBEvent(_) => true, // There is no inferred span for Step Functions events // if the `SpanContext` is generated IdentifiedTrigger::StepFunctionEvent(_) => { - extract_generated_span_context(identified_trigger).is_some() + extract_generated_span_context(identified_trigger, merge_xray_traces).is_some() } _ => false, } @@ -231,7 +234,8 @@ impl SpanInferrer { let identified_trigger = IdentifiedTrigger::from_value(payload_value); let should_enrich_span = Self::should_enrich_span(&identified_trigger); - let should_skip_inferred_span = Self::should_skip_inferred_span(&identified_trigger); + let should_skip_inferred_span = + Self::should_skip_inferred_span(&identified_trigger, self.config.ext.merge_xray_traces); let wrapped_inferred_span = Self::get_wrapped_inferred_span(&identified_trigger, &mut inferred_span, &self.config); let span_pointers = Self::get_span_pointers(&identified_trigger); @@ -396,7 +400,8 @@ pub fn extract_span_context( propagator: Arc, ) -> Option { let identified_trigger = IdentifiedTrigger::from_value(payload_value); - let generated_span_context = extract_generated_span_context(&identified_trigger); + let generated_span_context = + extract_generated_span_context(&identified_trigger, propagator.merge_xray_traces()); let trigger = SpanInferrer::get_trigger_type(identified_trigger); // Order matters here: check inferred span for trace context first, then fallback to generated span context. @@ -421,12 +426,14 @@ pub fn extract_span_context( #[must_use] pub fn extract_generated_span_context( identified_trigger: &IdentifiedTrigger, + merge_xray_traces: bool, ) -> Option { match identified_trigger { IdentifiedTrigger::StepFunctionEvent(t) => Some(t.get_span_context()), - IdentifiedTrigger::SqsRecord(t) => { - extract_trace_context_from_aws_trace_header(t.attributes.aws_trace_header.clone()) - } + IdentifiedTrigger::SqsRecord(t) => extract_trace_context_from_aws_trace_header( + t.attributes.aws_trace_header.clone(), + merge_xray_traces, + ), _ => None, } } diff --git a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs index d4a45e65a..f99d4b8e2 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs @@ -238,6 +238,7 @@ impl ServiceNameResolver for SqsRecord { // not be passed to the tracer.Propagator, instead extracting context directly. pub(crate) fn extract_trace_context_from_aws_trace_header( headers_string: Option, + merge_xray_traces: bool, ) -> Option { let value = headers_string?; if !value.starts_with("Root=") { @@ -277,11 +278,13 @@ pub(crate) fn extract_trace_context_from_aws_trace_header( return None; } - // Whether `Sampled` is a decision we can act on depends on who wrote the header. - let sampling_priority = match (datadog_planted, sampled.as_str()) { + // Whether `Sampled` is a decision we can act on depends on who wrote the header, and on + // whether the user opted into X-Ray driving their sampling. + let honor_sampled_flag = datadog_planted || merge_xray_traces; + let sampling_priority = match (honor_sampled_flag, sampled.as_str()) { // A keep is safe to honor whoever wrote it: the cost is over-retention, not lost spans. (_, "1") => Some("1"), - // Datadog planted this header, so the drop is Datadog's own decision. + // Datadog planted this header, or the user asked for X-Ray's decisions to count. (true, _) => Some("0"), // X-Ray's drop, or no decision at all when its tracing is off. Leaving the priority // unset lets the tracer sample; sending 0 would drop its spans on X-Ray's behalf. @@ -683,9 +686,10 @@ mod tests { let event = SqsRecord::new(payload).expect("Failed to deserialize EventBridgeEvent"); assert_eq!( - extract_trace_context_from_aws_trace_header(Some( - event.attributes.aws_trace_header.unwrap().clone() - )) + extract_trace_context_from_aws_trace_header( + Some(event.attributes.aws_trace_header.unwrap().clone()), + false + ) .unwrap(), SpanContext { trace_id: 130_944_522_478_755_159, @@ -705,10 +709,13 @@ 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(), - )) + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0" + .to_string(), + ), + false, + ) .expect("failed to extract context"); // X-Ray's decision, not Datadog's: keep the IDs, let the tracer sample. @@ -718,10 +725,13 @@ mod tests { #[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(), - )) + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac;Sampled=0" + .to_string(), + ), + false, + ) .expect("failed to extract context"); // Zeroed high bits mean a Datadog library wrote this, so `Sampled=0` is our own decision. @@ -730,28 +740,84 @@ mod tests { #[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(), - )) + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac;Sampled=1" + .to_string(), + ), + false, + ) .expect("failed to extract context"); assert_eq!(context.sampling.priority, "1".parse().ok()); } + #[test] + fn aws_generated_root_id_not_sampled_drops_when_merging_is_on() { + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0" + .to_string(), + ), + true, + ) + .expect("failed to extract context"); + + // Opting in means X-Ray's decision counts, drops included. + assert_eq!(context.sampling.priority, "0".parse().ok()); + } + + #[test] + fn merging_does_not_change_datadog_planted_headers() { + for merge in [false, true] { + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-68029e8a-0000000035578e774943fd9d;Parent=76c040bdc454a7ac;Sampled=0" + .to_string(), + ), + merge, + ) + .expect("failed to extract context"); + + assert_eq!(context.sampling.priority, "0".parse().ok(), "merge={merge}"); + } + } + + #[test] + fn merging_leaves_correlation_ids_alone() { + for merge in [false, true] { + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=0" + .to_string(), + ), + merge, + ) + .expect("failed to extract context"); + + // The flag gates the sampling verdict only; the IDs are taken either way. + assert_eq!(context.trace_id, 130_944_522_478_755_159, "merge={merge}"); + assert_eq!(context.span_id, 9_032_698_535_745_367_362, "merge={merge}"); + } + } + #[test] fn truncated_root_does_not_panic() { assert!( - extract_trace_context_from_aws_trace_header(Some("Root=1-64cc".to_string())).is_none() + extract_trace_context_from_aws_trace_header(Some("Root=1-64cc".to_string()), false) + .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(), - )) + let context = extract_trace_context_from_aws_trace_header( + Some( + "Root=1-64cc2edd-112fbf1701d1355973a11d57;Parent=7d5a9776024b2d42;Sampled=1" + .to_string(), + ), + false, + ) .expect("failed to extract context"); assert_eq!(context.sampling.priority, "1".parse().ok()); diff --git a/bottlecap/src/traces/propagation/mod.rs b/bottlecap/src/traces/propagation/mod.rs index 0d591a303..82b1b1997 100644 --- a/bottlecap/src/traces/propagation/mod.rs +++ b/bottlecap/src/traces/propagation/mod.rs @@ -43,6 +43,14 @@ impl DatadogCompositePropagator { Self { inner, config } } + /// Whether the user opted into X-Ray's own sampling decision driving Datadog's + /// (`DD_MERGE_XRAY_TRACES`). Read from here because this is the config-bearing object already + /// threaded to every trace-extraction site. + #[must_use] + pub fn merge_xray_traces(&self) -> bool { + self.config.ext.merge_xray_traces + } + pub fn extract(&self, carrier: &dyn Extractor) -> Option { let mut context = match self.inner.extract(carrier) { ExtractResult::Continue(context) => context,