[SLES-2971] feat(traces): gate X-Ray-driven sampling behind DD_MERGE_XRAY_TRACES

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<usize>,
+
+    #[serde(deserialize_with = "deser_opt_bool")]
+    pub merge_xray_traces: Option<bool>,
 }
 
 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],
         );

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<DatadogCompositePropagator>,
 ) -> Option<SpanContext> {
     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<SpanContext> {
     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,
     }
 }

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<String>,
+    merge_xray_traces: bool,
 ) -> Option<SpanContext> {
     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());

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<SpanContext> {
         let mut context = match self.inner.extract(carrier) {
             ExtractResult::Continue(context) => context,