From d724af00e74df19bf7ce1712bf2306860e1a9539 Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Fri, 4 Sep 2026 15:19:42 -0400
Subject: [PATCH 1/6] [APMSVLS-469] fix(traces): filter sampled-out chunks when
 tracer or extension computes stats

Trace intake now drops sampled-out chunks whenever stats are computed
before the backend, by either the tracer (Datadog-Client-Computed-Stats)
or the extension, instead of only when extension-side stats are enabled.
The backend still receives them when it owns stats computation, since it
needs them to compute stats. The error sampler now rescues eligible
errored drops on both non-backend paths.
---
 bottlecap/src/config/mod.rs             |  12 +-
 bottlecap/src/traces/trace_processor.rs | 321 +++++++++++++++---------
 bottlecap/tests/apm_integration_test.rs | 165 ++++++++----
 3 files changed, 324 insertions(+), 174 deletions(-)

diff --git a/bottlecap/src/config/mod.rs b/bottlecap/src/config/mod.rs
index 26119d613..d1a9ed258 100644
--- a/bottlecap/src/config/mod.rs
+++ b/bottlecap/src/config/mod.rs
@@ -70,12 +70,12 @@ pub struct LambdaConfig {
     pub lambda_extension_compute_stats: bool,
 
     /// `DD_SERVERLESS_ERROR_SAMPLER_ENABLED`: rescue errored trace chunks that would
-    /// otherwise be dropped, on the `lambda_extension_compute_stats` path. The
-    /// sampler runs in `AlwaysKeep` mode, so this is a plain on/off switch with
-    /// no volume ceiling: enabled rescues every errored chunk, whatever the
-    /// tracer's sampling rate. A function that errors on most invocations under
-    /// `DD_TRACE_SAMPLE_RATE=0.01` will ingest close to every trace, not 1%.
-    /// See APMSVLS-469.
+    /// otherwise be dropped, whenever stats are computed before the backend, by
+    /// either the tracer or the extension. The sampler runs in `AlwaysKeep` mode,
+    /// so this is a plain on/off switch with no volume ceiling: enabled rescues
+    /// every errored chunk, whatever the tracer's sampling rate. A function that
+    /// errors on most invocations under `DD_TRACE_SAMPLE_RATE=0.01` will ingest
+    /// close to every trace, not 1%. See APMSVLS-469.
     pub serverless_error_sampler_enabled: bool,
 
     pub span_dedup_timeout: Option<Duration>,
diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index 265748858..63860d59c 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -93,10 +93,10 @@ pub fn error_rare_sample_tracer_drop_enabled(config: &config::Config) -> bool {
 #[allow(clippy::module_name_repetitions)]
 pub struct ServerlessTraceProcessor {
     pub obfuscation_config: Arc<obfuscation_config::ObfuscationConfig>,
-    /// Rescues errored chunks on the `lambda_extension_compute_stats`
-    /// path. Shared across invocations so a stateful mode (`RateLimited`) would
-    /// keep its rolling window; a std Mutex rather than tokio because
-    /// `process_traces` is synchronous.
+    /// Rescues errored chunks whenever stats are computed before the backend,
+    /// by either the tracer or the extension. Shared across invocations so a
+    /// stateful mode (`RateLimited`) would keep its rolling window; a std Mutex
+    /// rather than tokio because `process_traces` is synchronous.
     pub error_sampler: Arc<std::sync::Mutex<ErrorsSampler>>,
     /// Whether `DD_APM_FEATURES` enables [`ERROR_RARE_SAMPLE_TRACER_DROP`]. When
     /// true, chunks with a negative priority are rescued like `AutoDrop` chunks
@@ -634,8 +634,14 @@ impl TraceProcessor for ServerlessTraceProcessor {
         };
 
         // Sampled-out chunks are preserved in payloads_for_stats above, so their
-        // stats are still counted after they are removed here.
-        if config.ext.lambda_extension_compute_stats
+        // stats are still counted after they are removed here. Filter whenever
+        // stats are computed before the backend, by either the tracer or the
+        // extension: the backend only needs these chunks when it owns stats.
+        let stats_computed_by = StatsComputedBy::resolve(
+            config.ext.lambda_extension_compute_stats,
+            header_tags.generic.client_computed_stats,
+        );
+        if stats_computed_by != StatsComputedBy::Backend
             && let TracerPayloadCollection::V07(ref mut tracer_payloads) = payload
         {
             self.drop_sampled_out_chunks(tracer_payloads);
@@ -1413,139 +1419,210 @@ mod tests {
         );
     }
 
-    /// Verifies that when `lambda_extension_compute_stats` is true, `process_traces`
-    /// filters sampled-out chunks from the backend payload while preserving them in the
-    /// stats collection.
+    /// Verifies that `process_traces` filters sampled-out chunks from the backend
+    /// payload whenever stats are computed before the backend, by either the tracer
+    /// or the extension, while preserving them in the stats collection. When the
+    /// backend owns stats, all chunks are kept so the backend can compute stats.
     #[test]
     #[allow(clippy::unwrap_used)]
     fn test_process_traces_filters_sampled_out_chunks() {
-        let config = create_compute_stats_config();
-        let (tags_provider, processor) =
-            create_test_processor(&config, enabled_error_sampler(), false);
-        let header_tags = create_test_header_tags();
+        // (lambda_extension_compute_stats, client_computed_stats) -> filtering expected
+        let cases = [
+            (false, false, false), // backend owns stats: keep all chunks
+            (true, false, true),   // extension owns stats: filter
+            (false, true, true),   // tracer owns stats: filter
+            (true, true, true),    // tracer owns stats even when extension configured: filter
+        ];
 
-        let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
-            let mut metrics = HashMap::new();
-            if let Some(p) = priority {
-                metrics.insert("_sampling_priority_v1".to_string(), p);
-            }
-            pb::Span {
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                metrics,
-                service: "svc".to_string(),
-                name: "op".to_string(),
-                resource: "res".to_string(),
-                ..Default::default()
-            }
-        };
+        for (compute_on_extension, client_computed_stats, expect_filtering) in cases {
+            let config = Arc::new(Config {
+                apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
+                ext: crate::config::LambdaConfig {
+                    lambda_extension_compute_stats: compute_on_extension,
+                    ..Default::default()
+                },
+                ..Config::default()
+            });
+            let (tags_provider, processor) =
+                create_test_processor(&config, enabled_error_sampler(), false);
+            let header_tags = tracer_header_tags::TracerHeaderTags {
+                generic: tracer_header_tags::TracerGenericTags {
+                    client_computed_stats,
+                    ..Default::default()
+                },
+                ..create_test_header_tags()
+            };
 
-        // Three traces: kept (priority 1), dropped (priority 0), dropped (priority -1)
-        let traces = vec![
-            vec![make_span(1, Some(1.0))],
-            vec![make_span(2, Some(0.0))],
-            vec![make_span(3, Some(-1.0))],
-        ];
+            let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
+                let mut metrics = HashMap::new();
+                if let Some(p) = priority {
+                    metrics.insert("_sampling_priority_v1".to_string(), p);
+                }
+                pb::Span {
+                    trace_id,
+                    span_id: trace_id,
+                    parent_id: 0,
+                    metrics,
+                    service: "svc".to_string(),
+                    name: "op".to_string(),
+                    resource: "res".to_string(),
+                    ..Default::default()
+                }
+            };
 
-        let (payload_info, stats_collection) =
-            processor.process_traces(config, tags_provider, header_tags, traces, 0, None);
-        let payload_info = payload_info.expect("expected Some payload");
+            // Three traces: kept (priority 1), dropped (priority 0), dropped (priority -1)
+            let traces = vec![
+                vec![make_span(1, Some(1.0))],
+                vec![make_span(2, Some(0.0))],
+                vec![make_span(3, Some(-1.0))],
+            ];
 
-        // Stats collection must include all three traces
-        let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else {
-            panic!("expected V07");
-        };
-        let stats_span_count: usize = stats_payloads
-            .iter()
-            .flat_map(|tp| tp.chunks.iter())
-            .map(|c| c.spans.len())
-            .sum();
-        assert_eq!(stats_span_count, 3, "stats must include all traces");
+            let (payload_info, stats_collection) =
+                processor.process_traces(config, tags_provider, header_tags, traces, 0, None);
+            let payload_info = payload_info.expect("expected Some payload");
 
-        // Backend payload must only contain the kept trace (priority 1)
-        let backend_send_data = payload_info.builder.build();
-        let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
-        else {
-            panic!("expected V07");
-        };
-        let backend_span_count: usize = backend_payloads
-            .iter()
-            .flat_map(|tp| tp.chunks.iter())
-            .map(|c| c.spans.len())
-            .sum();
-        assert_eq!(
-            backend_span_count, 1,
-            "backend must only include kept traces"
-        );
+            // Stats collection must include all three traces
+            let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else {
+                panic!("expected V07");
+            };
+            let stats_span_count: usize = stats_payloads
+                .iter()
+                .flat_map(|tp| tp.chunks.iter())
+                .map(|c| c.spans.len())
+                .sum();
+            assert_eq!(
+                stats_span_count, 3,
+                "stats must include all traces (compute_on_extension={compute_on_extension}, \
+                 client_computed_stats={client_computed_stats})"
+            );
+
+            let backend_send_data = payload_info.builder.build();
+            let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
+            else {
+                panic!("expected V07");
+            };
+            let backend_span_count: usize = backend_payloads
+                .iter()
+                .flat_map(|tp| tp.chunks.iter())
+                .map(|c| c.spans.len())
+                .sum();
+            let expected_backend_span_count = if expect_filtering { 1 } else { 3 };
+            assert_eq!(
+                backend_span_count, expected_backend_span_count,
+                "compute_on_extension={compute_on_extension}, \
+                 client_computed_stats={client_computed_stats}"
+            );
+        }
     }
 
-    /// On the compute-stats path, the error sampler rescues errored chunks that
-    /// would otherwise be dropped (`AutoDrop`), stamping `_dd.errors_sr`, while
-    /// non-errored P0 chunks and explicit user drops are still dropped.
+    /// On any stats path where stats are computed before the backend (extension or
+    /// tracer), the error sampler rescues errored chunks that would otherwise be
+    /// dropped (`AutoDrop`), stamping `_dd.errors_sr`, while non-errored P0 chunks
+    /// and explicit user drops are still dropped. Priorities are never rewritten.
     #[test]
+    #[allow(clippy::too_many_lines)]
     fn test_error_sampler_rescues_errored_p0_chunks() {
-        let config = create_compute_stats_config();
-        let (tags_provider, processor) =
-            create_test_processor(&config, enabled_error_sampler(), false);
-        let header_tags = create_test_header_tags();
+        // (owner label, lambda_extension_compute_stats, client_computed_stats)
+        let cases = [("extension", true, false), ("tracer", false, true)];
 
-        let make_span = |trace_id: u64, priority: f64, error: i32| -> pb::Span {
-            let mut metrics = HashMap::new();
-            metrics.insert("_sampling_priority_v1".to_string(), priority);
-            pb::Span {
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                error,
-                metrics,
-                service: "svc".to_string(),
-                name: "op".to_string(),
-                resource: "res".to_string(),
-                ..Default::default()
-            }
-        };
+        for (owner, compute_on_extension, client_computed_stats) in cases {
+            let config = Arc::new(Config {
+                apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
+                ext: crate::config::LambdaConfig {
+                    lambda_extension_compute_stats: compute_on_extension,
+                    ..Default::default()
+                },
+                ..Config::default()
+            });
+            let (tags_provider, processor) =
+                create_test_processor(&config, enabled_error_sampler(), false);
 
-        // trace 1: kept normally (priority 1). trace 2: errored P0 (rescued).
-        // trace 3: non-errored P0 (dropped). trace 4: errored user drop (dropped).
-        let traces = vec![
-            vec![make_span(1, 1.0, 0)],
-            vec![make_span(2, 0.0, 1)],
-            vec![make_span(3, 0.0, 0)],
-            vec![make_span(4, -1.0, 1)],
-        ];
+            let header_tags = tracer_header_tags::TracerHeaderTags {
+                lang: "rust",
+                lang_version: "1.0",
+                lang_interpreter: "",
+                lang_vendor: "",
+                tracer_version: "1.0",
+                container_id: "",
+                generic: tracer_header_tags::TracerGenericTags {
+                    client_computed_stats,
+                    ..Default::default()
+                },
+            };
 
-        let (payload_info, _stats) =
-            processor.process_traces(config, tags_provider, header_tags, traces, 0, None);
-        let payload_info = payload_info.expect("expected Some payload");
-        let backend_send_data = payload_info.builder.build();
-        let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
-        else {
-            panic!("expected V07");
-        };
+            let make_span = |trace_id: u64, priority: f64, error: i32| -> pb::Span {
+                let mut metrics = HashMap::new();
+                metrics.insert("_sampling_priority_v1".to_string(), priority);
+                pb::Span {
+                    trace_id,
+                    span_id: trace_id,
+                    parent_id: 0,
+                    error,
+                    metrics,
+                    service: "svc".to_string(),
+                    name: "op".to_string(),
+                    resource: "res".to_string(),
+                    ..Default::default()
+                }
+            };
 
-        let kept: Vec<u64> = backend_payloads
-            .iter()
-            .flat_map(|tp| tp.chunks.iter())
-            .flat_map(|c| c.spans.iter())
-            .map(|s| s.trace_id)
-            .collect();
-        assert_eq!(kept.len(), 2, "kept normal trace + rescued errored trace");
-        assert!(kept.contains(&1), "priority-1 trace kept");
-        assert!(kept.contains(&2), "errored P0 trace rescued");
-        assert!(!kept.contains(&3), "non-errored P0 trace dropped");
-        assert!(!kept.contains(&4), "errored user-drop trace not rescued");
+            // trace 1: kept normally (priority 1). trace 2: errored P0 (rescued).
+            // trace 3: non-errored P0 (dropped). trace 4: errored user drop (dropped).
+            let traces = vec![
+                vec![make_span(1, 1.0, 0)],
+                vec![make_span(2, 0.0, 1)],
+                vec![make_span(3, 0.0, 0)],
+                vec![make_span(4, -1.0, 1)],
+            ];
+
+            let (payload_info, _stats) =
+                processor.process_traces(config, tags_provider, header_tags, traces, 0, None);
+            let payload_info = payload_info.expect("expected Some payload");
+            let backend_send_data = payload_info.builder.build();
+            let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
+            else {
+                panic!("expected V07");
+            };
 
-        let rescued_root = backend_payloads
-            .iter()
-            .flat_map(|tp| tp.chunks.iter())
-            .flat_map(|c| c.spans.iter())
-            .find(|s| s.trace_id == 2)
-            .expect("rescued trace present");
-        assert!(
-            rescued_root.metrics.contains_key("_dd.errors_sr"),
-            "_dd.errors_sr stamped on rescued root"
-        );
+            let kept: Vec<(u64, f64)> = backend_payloads
+                .iter()
+                .flat_map(|tp| tp.chunks.iter())
+                .flat_map(|c| c.spans.iter())
+                .map(|s| (s.trace_id, s.metrics["_sampling_priority_v1"]))
+                .collect();
+            assert_eq!(
+                kept.len(),
+                2,
+                "kept normal trace + rescued errored trace (owner={owner})"
+            );
+            assert!(
+                kept.contains(&(1, 1.0)),
+                "priority-1 trace kept (owner={owner})"
+            );
+            assert!(
+                kept.contains(&(2, 0.0)),
+                "errored P0 trace rescued with its original priority (owner={owner})"
+            );
+            assert!(
+                !kept.iter().any(|(id, _)| *id == 3),
+                "non-errored P0 trace dropped (owner={owner})"
+            );
+            assert!(
+                !kept.iter().any(|(id, _)| *id == 4),
+                "errored user-drop trace not rescued (owner={owner})"
+            );
+
+            let rescued_root = backend_payloads
+                .iter()
+                .flat_map(|tp| tp.chunks.iter())
+                .flat_map(|c| c.spans.iter())
+                .find(|s| s.trace_id == 2)
+                .expect("rescued trace present");
+            assert!(
+                rescued_root.metrics.contains_key("_dd.errors_sr"),
+                "_dd.errors_sr stamped on rescued root (owner={owner})"
+            );
+        }
     }
 
     /// With `error_rare_sample_tracer_drop` enabled, an errored chunk the tracer
diff --git a/bottlecap/tests/apm_integration_test.rs b/bottlecap/tests/apm_integration_test.rs
index ed85ada51..e9ae9a589 100644
--- a/bottlecap/tests/apm_integration_test.rs
+++ b/bottlecap/tests/apm_integration_test.rs
@@ -493,59 +493,132 @@ fn captured_compute_stats(traces: &[pb::AgentPayload]) -> Option<String> {
     span.meta.get(COMPUTE_STATS_KEY).cloned()
 }
 
-/// An errored P0 trace is rescued and reaches the intake, while a non-errored P0 trace is dropped.
+/// An errored P0 trace is rescued and reaches the intake, while a non-errored P0 trace is
+/// dropped. Applies whenever stats are computed before the backend, by either the
+/// extension or the tracer.
 #[tokio::test]
 async fn e2e_error_sampler_rescues_only_errored_p0_traces() {
-    let make_span = |trace_id: u64, error: i32| {
-        let mut span = pb::Span {
-            service: "fake-intake-trace-service".to_string(),
-            name: "web.request".to_string(),
-            resource: "GET /fake".to_string(),
-            trace_id,
-            span_id: trace_id,
-            parent_id: 0,
-            start: STATS_SPAN_START_NS,
-            duration: 5_000_000,
-            error,
-            r#type: "web".to_string(),
-            ..pb::Span::default()
+    // (owner label, lambda_extension_compute_stats, client_computed_stats)
+    let cases = [("extension", true, false), ("tracer", false, true)];
+
+    for (owner, compute_on_extension, client_computed_stats) in cases {
+        let make_span = |trace_id: u64, error: i32| {
+            let mut span = pb::Span {
+                service: "fake-intake-trace-service".to_string(),
+                name: "web.request".to_string(),
+                resource: "GET /fake".to_string(),
+                trace_id,
+                span_id: trace_id,
+                parent_id: 0,
+                start: STATS_SPAN_START_NS,
+                duration: 5_000_000,
+                error,
+                r#type: "web".to_string(),
+                ..pb::Span::default()
+            };
+            span.metrics
+                .insert("_sampling_priority_v1".to_string(), 0.0);
+            span
         };
-        span.metrics
-            .insert("_sampling_priority_v1".to_string(), 0.0);
-        span
-    };
 
-    let rescued_trace_id = 1;
-    let dropped_trace_id = 2;
-    let outcome = run_processor_pipeline_with_traces(
-        true,
-        false,
-        vec![
-            vec![make_span(rescued_trace_id, 1)],
-            vec![make_span(dropped_trace_id, 0)],
-        ],
-    )
-    .await;
+        let rescued_trace_id = 1;
+        let dropped_trace_id = 2;
+        let outcome = run_processor_pipeline_with_traces(
+            compute_on_extension,
+            client_computed_stats,
+            vec![
+                vec![make_span(rescued_trace_id, 1)],
+                vec![make_span(dropped_trace_id, 0)],
+            ],
+        )
+        .await;
 
-    let captured_spans: Vec<&pb::Span> = captured_trace_spans(&outcome.traces).collect();
-    assert_eq!(
-        captured_spans.len(),
-        1,
-        "only the rescued trace reaches intake"
-    );
+        let captured_spans: Vec<&pb::Span> = captured_trace_spans(&outcome.traces).collect();
+        assert_eq!(
+            captured_spans.len(),
+            1,
+            "only the rescued trace reaches intake (owner={owner})"
+        );
 
-    let rescued_span = captured_spans[0];
-    assert_eq!(rescued_span.trace_id, rescued_trace_id);
-    assert!(
-        rescued_span.metrics.contains_key("_dd.errors_sr"),
-        "the rescued root span must carry its error sampling rate",
-    );
-    assert!(
-        captured_spans
+        let rescued_span = captured_spans[0];
+        assert_eq!(rescued_span.trace_id, rescued_trace_id);
+        assert!(
+            rescued_span.metrics.contains_key("_dd.errors_sr"),
+            "the rescued root span must carry its error sampling rate (owner={owner})",
+        );
+        assert!(
+            captured_spans
+                .iter()
+                .all(|span| span.trace_id != dropped_trace_id),
+            "the non-errored P0 trace must not reach intake (owner={owner})",
+        );
+    }
+}
+
+/// Sampled-out chunks are filtered from the trace intake whenever stats are computed
+/// before the backend, by either the extension or the tracer, while backend-owned
+/// stats keep them so the backend can compute stats.
+#[tokio::test]
+async fn e2e_sampled_out_chunks_filtered_by_stats_owner() {
+    // (owner label, lambda_extension_compute_stats, client_computed_stats) -> filtering expected
+    let cases = [
+        ("backend", false, false, false),
+        ("extension", true, false, true),
+        ("tracer", false, true, true),
+    ];
+
+    for (owner, compute_on_extension, client_computed_stats, expect_filtering) in cases {
+        let make_span = |trace_id: u64, priority: f64| {
+            let mut span = pb::Span {
+                service: "fake-intake-trace-service".to_string(),
+                name: "web.request".to_string(),
+                resource: "GET /fake".to_string(),
+                trace_id,
+                span_id: trace_id,
+                parent_id: 0,
+                start: STATS_SPAN_START_NS,
+                duration: 5_000_000,
+                error: 0,
+                r#type: "web".to_string(),
+                ..pb::Span::default()
+            };
+            span.metrics
+                .insert("_sampling_priority_v1".to_string(), priority);
+            span
+        };
+
+        // kept (priority 1), auto-dropped (priority 0), explicit drop (priority -1)
+        let outcome = run_processor_pipeline_with_traces(
+            compute_on_extension,
+            client_computed_stats,
+            vec![
+                vec![make_span(1, 1.0)],
+                vec![make_span(2, 0.0)],
+                vec![make_span(3, -1.0)],
+            ],
+        )
+        .await;
+
+        let mut captured_trace_ids: Vec<u64> = outcome
+            .traces
             .iter()
-            .all(|span| span.trace_id != dropped_trace_id),
-        "the non-errored P0 trace must not reach intake",
-    );
+            .flat_map(|payload| &payload.tracer_payloads)
+            .flat_map(|payload| &payload.chunks)
+            .flat_map(|chunk| &chunk.spans)
+            .map(|span| span.trace_id)
+            .collect();
+        captured_trace_ids.sort_unstable();
+
+        let expected_trace_ids: Vec<u64> = if expect_filtering {
+            vec![1]
+        } else {
+            vec![1, 2, 3]
+        };
+        assert_eq!(
+            captured_trace_ids, expected_trace_ids,
+            "unexpected intake contents (owner={owner})"
+        );
+    }
 }
 
 /// A sampling decision maker in span meta is promoted into the chunk tags that

From e7803c1deeb747569d0ff64cdf02a37f86006f7e Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Tue, 8 Sep 2026 16:46:59 -0400
Subject: [PATCH 2/6] Reuse trace stats ownership decision

---
 bottlecap/src/traces/trace_processor.rs | 38 +++++++++++++------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index 63860d59c..8dbbe6810 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -248,9 +248,10 @@ struct ChunkProcessor {
     obfuscation_config: Arc<obfuscation_config::ObfuscationConfig>,
     tags_provider: Arc<provider::Provider>,
     span_pointers: Option<Vec<SpanPointer>>,
-    /// Whether the tracer signaled (via `Datadog-Client-Computed-Stats`) that it already
-    /// computed trace stats client-side. Used to decide whether to stamp `_dd.compute_stats`.
-    client_computed_stats: bool,
+    /// Which party computes trace stats for this request, resolved once by the caller from
+    /// config and the tracer's `Datadog-Client-Computed-Stats` signal. Used to decide
+    /// whether to stamp `_dd.compute_stats`.
+    stats_computed_by: StatsComputedBy,
 }
 
 impl TraceChunkProcessor for ChunkProcessor {
@@ -283,12 +284,7 @@ impl TraceChunkProcessor for ChunkProcessor {
             .spans
             .retain(|span| !filter_span_from_lambda_library_or_runtime(span));
 
-        // The stats-routing decision depends only on config and the per-request
-        // client_computed_stats flag, so resolve it once instead of per span.
-        let stamp_compute_stats = StatsComputedBy::resolve(
-            self.config.ext.lambda_extension_compute_stats,
-            self.client_computed_stats,
-        ) == StatsComputedBy::Backend;
+        let stamp_compute_stats = self.stats_computed_by == StatsComputedBy::Backend;
 
         for span in &mut chunk.spans {
             // Service name could be incorrectly set to 'aws.lambda'
@@ -587,6 +583,12 @@ impl TraceProcessor for ServerlessTraceProcessor {
         body_size: usize,
         span_pointers: Option<Vec<SpanPointer>>,
     ) -> (Option<SendDataBuilderInfo>, TracerPayloadCollection) {
+        // Who computes stats drives both the `_dd.compute_stats` stamp below and the
+        // sampled-out filtering further down, so resolve it once for this request.
+        let stats_computed_by = StatsComputedBy::resolve(
+            config.ext.lambda_extension_compute_stats,
+            header_tags.generic.client_computed_stats,
+        );
         let mut payload = trace_utils::collect_pb_trace_chunks(
             traces,
             &header_tags,
@@ -595,7 +597,7 @@ impl TraceProcessor for ServerlessTraceProcessor {
                 obfuscation_config: self.obfuscation_config.clone(),
                 tags_provider: tags_provider.clone(),
                 span_pointers,
-                client_computed_stats: header_tags.generic.client_computed_stats,
+                stats_computed_by,
             },
             true, // send agentless since we are the agent
         )
@@ -637,10 +639,6 @@ impl TraceProcessor for ServerlessTraceProcessor {
         // stats are still counted after they are removed here. Filter whenever
         // stats are computed before the backend, by either the tracer or the
         // extension: the backend only needs these chunks when it owns stats.
-        let stats_computed_by = StatsComputedBy::resolve(
-            config.ext.lambda_extension_compute_stats,
-            header_tags.generic.client_computed_stats,
-        );
         if stats_computed_by != StatsComputedBy::Backend
             && let TracerPayloadCollection::V07(ref mut tracer_payloads) = payload
         {
@@ -1238,7 +1236,7 @@ mod tests {
                 )]),
             )),
             span_pointers: None,
-            client_computed_stats: false,
+            stats_computed_by: StatsComputedBy::Backend,
         };
 
         processor.process(&mut chunk, 0);
@@ -1323,7 +1321,7 @@ mod tests {
                 )]),
             )),
             span_pointers: None,
-            client_computed_stats: false,
+            stats_computed_by: StatsComputedBy::Backend,
         };
 
         processor.process(&mut chunk, 0);
@@ -1407,7 +1405,7 @@ mod tests {
                 )]),
             )),
             span_pointers: None,
-            client_computed_stats: false,
+            stats_computed_by: StatsComputedBy::Backend,
         };
 
         processor.process(&mut chunk, 0);
@@ -2224,6 +2222,10 @@ mod tests {
         client_computed_stats: bool,
     ) -> ChunkProcessor {
         let tags_provider = create_tags_provider(config.clone());
+        let stats_computed_by = StatsComputedBy::resolve(
+            config.ext.lambda_extension_compute_stats,
+            client_computed_stats,
+        );
         ChunkProcessor {
             config,
             obfuscation_config: Arc::new(
@@ -2231,7 +2233,7 @@ mod tests {
             ),
             tags_provider,
             span_pointers: None,
-            client_computed_stats,
+            stats_computed_by,
         }
     }
 

From 493b4bd040f1116b73345a6ee7d011cf07dd1372 Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Thu, 1 Oct 2026 13:40:15 -0400
Subject: [PATCH 3/6] Simplify stats ownership tests

---
 bottlecap/src/traces/trace_processor.rs | 103 ++++++++++--------------
 bottlecap/tests/apm_integration_test.rs |  68 ++++++----------
 2 files changed, 68 insertions(+), 103 deletions(-)

diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index 8dbbe6810..cdc0c690e 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -248,9 +248,8 @@ struct ChunkProcessor {
     obfuscation_config: Arc<obfuscation_config::ObfuscationConfig>,
     tags_provider: Arc<provider::Provider>,
     span_pointers: Option<Vec<SpanPointer>>,
-    /// Which party computes trace stats for this request, resolved once by the caller from
-    /// config and the tracer's `Datadog-Client-Computed-Stats` signal. Used to decide
-    /// whether to stamp `_dd.compute_stats`.
+    /// Resolved once per request by the caller; used to decide whether to stamp
+    /// `_dd.compute_stats`.
     stats_computed_by: StatsComputedBy,
 }
 
@@ -845,10 +844,14 @@ mod tests {
     }
 
     fn create_compute_stats_config() -> Arc<Config> {
+        create_compute_stats_config_with(true)
+    }
+
+    fn create_compute_stats_config_with(lambda_extension_compute_stats: bool) -> Arc<Config> {
         Arc::new(Config {
             apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
             ext: crate::config::LambdaConfig {
-                lambda_extension_compute_stats: true,
+                lambda_extension_compute_stats,
                 ..Default::default()
             },
             ..Config::default()
@@ -1432,15 +1435,25 @@ mod tests {
             (true, true, true),    // tracer owns stats even when extension configured: filter
         ];
 
+        let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
+            let mut metrics = HashMap::new();
+            if let Some(p) = priority {
+                metrics.insert("_sampling_priority_v1".to_string(), p);
+            }
+            pb::Span {
+                trace_id,
+                span_id: trace_id,
+                parent_id: 0,
+                metrics,
+                service: "svc".to_string(),
+                name: "op".to_string(),
+                resource: "res".to_string(),
+                ..Default::default()
+            }
+        };
+
         for (compute_on_extension, client_computed_stats, expect_filtering) in cases {
-            let config = Arc::new(Config {
-                apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
-                ext: crate::config::LambdaConfig {
-                    lambda_extension_compute_stats: compute_on_extension,
-                    ..Default::default()
-                },
-                ..Config::default()
-            });
+            let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
                 create_test_processor(&config, enabled_error_sampler(), false);
             let header_tags = tracer_header_tags::TracerHeaderTags {
@@ -1451,23 +1464,6 @@ mod tests {
                 ..create_test_header_tags()
             };
 
-            let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
-                let mut metrics = HashMap::new();
-                if let Some(p) = priority {
-                    metrics.insert("_sampling_priority_v1".to_string(), p);
-                }
-                pb::Span {
-                    trace_id,
-                    span_id: trace_id,
-                    parent_id: 0,
-                    metrics,
-                    service: "svc".to_string(),
-                    name: "op".to_string(),
-                    resource: "res".to_string(),
-                    ..Default::default()
-                }
-            };
-
             // Three traces: kept (priority 1), dropped (priority 0), dropped (priority -1)
             let traces = vec![
                 vec![make_span(1, Some(1.0))],
@@ -1518,50 +1514,37 @@ mod tests {
     /// dropped (`AutoDrop`), stamping `_dd.errors_sr`, while non-errored P0 chunks
     /// and explicit user drops are still dropped. Priorities are never rewritten.
     #[test]
-    #[allow(clippy::too_many_lines)]
     fn test_error_sampler_rescues_errored_p0_chunks() {
         // (owner label, lambda_extension_compute_stats, client_computed_stats)
         let cases = [("extension", true, false), ("tracer", false, true)];
 
+        let make_span = |trace_id: u64, priority: f64, error: i32| -> pb::Span {
+            let mut metrics = HashMap::new();
+            metrics.insert("_sampling_priority_v1".to_string(), priority);
+            pb::Span {
+                trace_id,
+                span_id: trace_id,
+                parent_id: 0,
+                error,
+                metrics,
+                service: "svc".to_string(),
+                name: "op".to_string(),
+                resource: "res".to_string(),
+                ..Default::default()
+            }
+        };
+
         for (owner, compute_on_extension, client_computed_stats) in cases {
-            let config = Arc::new(Config {
-                apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
-                ext: crate::config::LambdaConfig {
-                    lambda_extension_compute_stats: compute_on_extension,
-                    ..Default::default()
-                },
-                ..Config::default()
-            });
+            let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
                 create_test_processor(&config, enabled_error_sampler(), false);
 
             let header_tags = tracer_header_tags::TracerHeaderTags {
-                lang: "rust",
-                lang_version: "1.0",
-                lang_interpreter: "",
-                lang_vendor: "",
-                tracer_version: "1.0",
-                container_id: "",
                 generic: tracer_header_tags::TracerGenericTags {
                     client_computed_stats,
                     ..Default::default()
                 },
-            };
-
-            let make_span = |trace_id: u64, priority: f64, error: i32| -> pb::Span {
-                let mut metrics = HashMap::new();
-                metrics.insert("_sampling_priority_v1".to_string(), priority);
-                pb::Span {
-                    trace_id,
-                    span_id: trace_id,
-                    parent_id: 0,
-                    error,
-                    metrics,
-                    service: "svc".to_string(),
-                    name: "op".to_string(),
-                    resource: "res".to_string(),
-                    ..Default::default()
-                }
+                ..create_test_header_tags()
             };
 
             // trace 1: kept normally (priority 1). trace 2: errored P0 (rescued).
diff --git a/bottlecap/tests/apm_integration_test.rs b/bottlecap/tests/apm_integration_test.rs
index e9ae9a589..a22e3f88b 100644
--- a/bottlecap/tests/apm_integration_test.rs
+++ b/bottlecap/tests/apm_integration_test.rs
@@ -493,6 +493,26 @@ fn captured_compute_stats(traces: &[pb::AgentPayload]) -> Option<String> {
     span.meta.get(COMPUTE_STATS_KEY).cloned()
 }
 
+/// Single-span trace root with the given sampling priority and error flag.
+fn sampled_root_span(trace_id: u64, priority: f64, error: i32) -> pb::Span {
+    let mut span = pb::Span {
+        service: "fake-intake-trace-service".to_string(),
+        name: "web.request".to_string(),
+        resource: "GET /fake".to_string(),
+        trace_id,
+        span_id: trace_id,
+        parent_id: 0,
+        start: STATS_SPAN_START_NS,
+        duration: 5_000_000,
+        error,
+        r#type: "web".to_string(),
+        ..pb::Span::default()
+    };
+    span.metrics
+        .insert("_sampling_priority_v1".to_string(), priority);
+    span
+}
+
 /// An errored P0 trace is rescued and reaches the intake, while a non-errored P0 trace is
 /// dropped. Applies whenever stats are computed before the backend, by either the
 /// extension or the tracer.
@@ -502,33 +522,14 @@ async fn e2e_error_sampler_rescues_only_errored_p0_traces() {
     let cases = [("extension", true, false), ("tracer", false, true)];
 
     for (owner, compute_on_extension, client_computed_stats) in cases {
-        let make_span = |trace_id: u64, error: i32| {
-            let mut span = pb::Span {
-                service: "fake-intake-trace-service".to_string(),
-                name: "web.request".to_string(),
-                resource: "GET /fake".to_string(),
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                start: STATS_SPAN_START_NS,
-                duration: 5_000_000,
-                error,
-                r#type: "web".to_string(),
-                ..pb::Span::default()
-            };
-            span.metrics
-                .insert("_sampling_priority_v1".to_string(), 0.0);
-            span
-        };
-
         let rescued_trace_id = 1;
         let dropped_trace_id = 2;
         let outcome = run_processor_pipeline_with_traces(
             compute_on_extension,
             client_computed_stats,
             vec![
-                vec![make_span(rescued_trace_id, 1)],
-                vec![make_span(dropped_trace_id, 0)],
+                vec![sampled_root_span(rescued_trace_id, 0.0, 1)],
+                vec![sampled_root_span(dropped_trace_id, 0.0, 0)],
             ],
         )
         .await;
@@ -568,33 +569,14 @@ async fn e2e_sampled_out_chunks_filtered_by_stats_owner() {
     ];
 
     for (owner, compute_on_extension, client_computed_stats, expect_filtering) in cases {
-        let make_span = |trace_id: u64, priority: f64| {
-            let mut span = pb::Span {
-                service: "fake-intake-trace-service".to_string(),
-                name: "web.request".to_string(),
-                resource: "GET /fake".to_string(),
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                start: STATS_SPAN_START_NS,
-                duration: 5_000_000,
-                error: 0,
-                r#type: "web".to_string(),
-                ..pb::Span::default()
-            };
-            span.metrics
-                .insert("_sampling_priority_v1".to_string(), priority);
-            span
-        };
-
         // kept (priority 1), auto-dropped (priority 0), explicit drop (priority -1)
         let outcome = run_processor_pipeline_with_traces(
             compute_on_extension,
             client_computed_stats,
             vec![
-                vec![make_span(1, 1.0)],
-                vec![make_span(2, 0.0)],
-                vec![make_span(3, -1.0)],
+                vec![sampled_root_span(1, 1.0, 0)],
+                vec![sampled_root_span(2, 0.0, 0)],
+                vec![sampled_root_span(3, -1.0, 0)],
             ],
         )
         .await;

From 0136be75c66dc61dbecb56c814f010cc75a208d7 Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Thu, 22 Oct 2026 14:04:39 -0400
Subject: [PATCH 4/6] [APMSVLS-469] test: cover rare tracer-drop rescue on both
 stats owners

Run the error_rare_sample_tracer_drop rescue tests for both the
extension-owned and tracer-owned stats paths, since the tracer-owned
path now also filters sampled-out chunks. Stats-count assertions stay
extension-only: the extension suppresses its own stats generation when
the tracer computed stats.
---
 bottlecap/src/traces/trace_processor.rs |  59 +++++++-----
 bottlecap/tests/apm_integration_test.rs | 118 ++++++++++++++----------
 2 files changed, 106 insertions(+), 71 deletions(-)

diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index cdc0c690e..f5a17b574 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -1608,33 +1608,46 @@ mod tests {
 
     /// With `error_rare_sample_tracer_drop` enabled, an errored chunk the tracer
     /// dropped itself (priority -1, no decision maker) is rescued like an
-    /// `AutoDrop` chunk, matching the Go agent.
+    /// `AutoDrop` chunk, matching the Go agent. Runs for both stats owners: the
+    /// tracer-owned path reaches the rescue only when filtering also applies to
+    /// tracer-computed stats.
     #[test]
     fn test_error_rare_sample_tracer_drop_rescues_errored_user_drop() {
-        let config = create_compute_stats_config();
-        let (tags_provider, processor) =
-            create_test_processor(&config, enabled_error_sampler(), true);
-        let header_tags = create_test_header_tags();
+        // (owner label, lambda_extension_compute_stats, client_computed_stats)
+        let cases = [("extension", true, false), ("tracer", false, true)];
 
-        let span = make_tracer_drop_span(7, -1.0, 1, None);
+        for (owner, compute_on_extension, client_computed_stats) in cases {
+            let config = create_compute_stats_config_with(compute_on_extension);
+            let (tags_provider, processor) =
+                create_test_processor(&config, enabled_error_sampler(), true);
+            let header_tags = tracer_header_tags::TracerHeaderTags {
+                generic: tracer_header_tags::TracerGenericTags {
+                    client_computed_stats,
+                    ..Default::default()
+                },
+                ..create_test_header_tags()
+            };
 
-        let (payload_info, _stats) = processor.process_traces(
-            config,
-            tags_provider,
-            header_tags,
-            vec![vec![span]],
-            0,
-            None,
-        );
-        let payload_info = payload_info.expect("errored tracer drop must be rescued");
-        let backend_send_data = payload_info.builder.build();
-        let rescued = trace_spans(backend_send_data.get_payloads())
-            .find(|span| span.trace_id == 7)
-            .expect("rescued trace present");
-        assert!(
-            rescued.metrics.contains_key("_dd.errors_sr"),
-            "_dd.errors_sr stamped on rescued root"
-        );
+            let span = make_tracer_drop_span(7, -1.0, 1, None);
+
+            let (payload_info, _stats) = processor.process_traces(
+                config,
+                tags_provider,
+                header_tags,
+                vec![vec![span]],
+                0,
+                None,
+            );
+            let payload_info = payload_info.expect("errored tracer drop must be rescued");
+            let backend_send_data = payload_info.builder.build();
+            let rescued = trace_spans(backend_send_data.get_payloads())
+                .find(|span| span.trace_id == 7)
+                .expect("rescued trace present");
+            assert!(
+                rescued.metrics.contains_key("_dd.errors_sr"),
+                "_dd.errors_sr stamped on rescued root (owner={owner})"
+            );
+        }
     }
 
     /// Span factory for tracer-drop tests: a single-span trace with an explicit
diff --git a/bottlecap/tests/apm_integration_test.rs b/bottlecap/tests/apm_integration_test.rs
index a22e3f88b..02e94e82e 100644
--- a/bottlecap/tests/apm_integration_test.rs
+++ b/bottlecap/tests/apm_integration_test.rs
@@ -695,60 +695,82 @@ fn captured_trace_spans(traces: &[pb::AgentPayload]) -> impl Iterator<Item = &pb
 /// rule drop (priority -1, no decision maker) is rescued and reaches the intake
 /// with `_dd.errors_sr`, while an errored manual drop (decision maker "-4") and
 /// a nonerrored drop are excluded. Dropped inputs remain represented in stats.
+/// Runs for both stats owners: the tracer-owned path reaches the rescue only
+/// when filtering also applies to tracer-computed stats.
 #[tokio::test]
 async fn e2e_error_rare_sample_tracer_drop_rescues_tracer_drops_selectively() {
-    let rescued_trace_id = 1; // errored tracer rule drop
-    let manual_drop_trace_id = 2; // errored manual tracer drop
-    let nonerrored_trace_id = 3; // nonerrored tracer rule drop
-    let kept_trace_id = 4; // normal priority-1 trace
+    // (owner label, lambda_extension_compute_stats, client_computed_stats)
+    let cases = [("extension", true, false), ("tracer", false, true)];
 
-    let outcome = run_processor_pipeline_with_settings(
-        true,
-        false,
-        true,
-        vec![
-            vec![error_sampler_span(rescued_trace_id, -1.0, 1, None)],
-            vec![error_sampler_span(
-                manual_drop_trace_id,
-                -1.0,
-                1,
-                Some("-4"),
-            )],
-            vec![error_sampler_span(nonerrored_trace_id, -1.0, 0, None)],
-            vec![error_sampler_span(kept_trace_id, 1.0, 0, None)],
-        ],
-    )
-    .await;
+    for (owner, compute_on_extension, client_computed_stats) in cases {
+        let rescued_trace_id = 1; // errored tracer rule drop
+        let manual_drop_trace_id = 2; // errored manual tracer drop
+        let nonerrored_trace_id = 3; // nonerrored tracer rule drop
+        let kept_trace_id = 4; // normal priority-1 trace
 
-    let captured_spans: Vec<_> = captured_trace_spans(&outcome.traces).collect();
-    assert_eq!(
-        captured_spans.len(),
-        2,
-        "only the rescued tracer drop and the normally kept trace reach intake"
-    );
-    assert!(
-        captured_spans
-            .iter()
-            .any(|span| span.trace_id == kept_trace_id),
-        "normally kept trace must reach intake"
-    );
+        let outcome = run_processor_pipeline_with_settings(
+            compute_on_extension,
+            client_computed_stats,
+            true,
+            vec![
+                vec![error_sampler_span(rescued_trace_id, -1.0, 1, None)],
+                vec![error_sampler_span(
+                    manual_drop_trace_id,
+                    -1.0,
+                    1,
+                    Some("-4"),
+                )],
+                vec![error_sampler_span(nonerrored_trace_id, -1.0, 0, None)],
+                vec![error_sampler_span(kept_trace_id, 1.0, 0, None)],
+            ],
+        )
+        .await;
 
-    let rescued_root = captured_spans
-        .iter()
-        .find(|span| span.trace_id == rescued_trace_id)
-        .expect("rescued trace present");
-    assert!(
-        rescued_root.metrics.contains_key("_dd.errors_sr"),
-        "the rescued root span must carry its error sampling rate",
-    );
+        let captured_spans: Vec<_> = captured_trace_spans(&outcome.traces).collect();
+        assert_eq!(
+            captured_spans.len(),
+            2,
+            "only the rescued tracer drop and the normally kept trace reach intake \
+             (owner={owner})"
+        );
+        assert!(
+            captured_spans
+                .iter()
+                .any(|span| span.trace_id == kept_trace_id),
+            "normally kept trace must reach intake (owner={owner})"
+        );
 
-    // Dropped inputs must still be accounted for in stats: four hits, two of
-    // them errored.
-    let grouped = captured_grouped_stats(&outcome.stats);
-    let total_hits: u64 = grouped.iter().map(|g| g.hits).sum();
-    let total_errors: u64 = grouped.iter().map(|g| g.errors).sum();
-    assert_eq!(total_hits, 4, "all traces must be represented in stats");
-    assert_eq!(total_errors, 2, "both errored traces must be counted");
+        let rescued_root = captured_spans
+            .iter()
+            .find(|span| span.trace_id == rescued_trace_id)
+            .expect("rescued trace present");
+        assert!(
+            rescued_root.metrics.contains_key("_dd.errors_sr"),
+            "the rescued root span must carry its error sampling rate (owner={owner})",
+        );
+
+        // Dropped inputs must still be accounted for in stats, but only when the
+        // extension computes them: with tracer-owned stats the extension
+        // suppresses its own stats generation entirely.
+        if compute_on_extension {
+            let grouped = captured_grouped_stats(&outcome.stats);
+            let total_hits: u64 = grouped.iter().map(|g| g.hits).sum();
+            let total_errors: u64 = grouped.iter().map(|g| g.errors).sum();
+            assert_eq!(
+                total_hits, 4,
+                "all traces must be represented in stats (owner={owner})"
+            );
+            assert_eq!(
+                total_errors, 2,
+                "both errored traces must be counted (owner={owner})"
+            );
+        } else {
+            assert!(
+                outcome.stats.is_empty(),
+                "extension must not emit stats when the tracer computed them (owner={owner})"
+            );
+        }
+    }
 }
 
 /// APMSVLS-674 control: with the feature disabled (shipping default), the

From dd4ce8f711538c4f0d79cfa45e1238981d4cb21e Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Thu, 22 Oct 2026 14:18:29 -0400
Subject: [PATCH 5/6] Simplify sampled-trace filtering tests

---
 bottlecap/src/traces/trace_processor.rs | 83 +++++--------------------
 bottlecap/tests/apm_integration_test.rs | 29 ++-------
 2 files changed, 22 insertions(+), 90 deletions(-)

diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index f5a17b574..e50119d80 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -1427,12 +1427,12 @@ mod tests {
     #[test]
     #[allow(clippy::unwrap_used)]
     fn test_process_traces_filters_sampled_out_chunks() {
-        // (lambda_extension_compute_stats, client_computed_stats) -> filtering expected
+        // (lambda_extension_compute_stats, client_computed_stats, expected_backend_span_count)
         let cases = [
-            (false, false, false), // backend owns stats: keep all chunks
-            (true, false, true),   // extension owns stats: filter
-            (false, true, true),   // tracer owns stats: filter
-            (true, true, true),    // tracer owns stats even when extension configured: filter
+            (false, false, 3), // backend owns stats: keep all chunks
+            (true, false, 1),  // extension owns stats: filter
+            (false, true, 1),  // tracer owns stats: filter
+            (true, true, 1),   // tracer owns stats even when extension configured: filter
         ];
 
         let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
@@ -1452,17 +1452,12 @@ mod tests {
             }
         };
 
-        for (compute_on_extension, client_computed_stats, expect_filtering) in cases {
+        for (compute_on_extension, client_computed_stats, expected_backend_span_count) in cases {
             let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
                 create_test_processor(&config, enabled_error_sampler(), false);
-            let header_tags = tracer_header_tags::TracerHeaderTags {
-                generic: tracer_header_tags::TracerGenericTags {
-                    client_computed_stats,
-                    ..Default::default()
-                },
-                ..create_test_header_tags()
-            };
+            let mut header_tags = create_test_header_tags();
+            header_tags.generic.client_computed_stats = client_computed_stats;
 
             // Three traces: kept (priority 1), dropped (priority 0), dropped (priority -1)
             let traces = vec![
@@ -1476,14 +1471,7 @@ mod tests {
             let payload_info = payload_info.expect("expected Some payload");
 
             // Stats collection must include all three traces
-            let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else {
-                panic!("expected V07");
-            };
-            let stats_span_count: usize = stats_payloads
-                .iter()
-                .flat_map(|tp| tp.chunks.iter())
-                .map(|c| c.spans.len())
-                .sum();
+            let stats_span_count = trace_spans(&stats_collection).count();
             assert_eq!(
                 stats_span_count, 3,
                 "stats must include all traces (compute_on_extension={compute_on_extension}, \
@@ -1491,16 +1479,7 @@ mod tests {
             );
 
             let backend_send_data = payload_info.builder.build();
-            let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
-            else {
-                panic!("expected V07");
-            };
-            let backend_span_count: usize = backend_payloads
-                .iter()
-                .flat_map(|tp| tp.chunks.iter())
-                .map(|c| c.spans.len())
-                .sum();
-            let expected_backend_span_count = if expect_filtering { 1 } else { 3 };
+            let backend_span_count = trace_spans(backend_send_data.get_payloads()).count();
             assert_eq!(
                 backend_span_count, expected_backend_span_count,
                 "compute_on_extension={compute_on_extension}, \
@@ -1538,14 +1517,8 @@ mod tests {
             let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
                 create_test_processor(&config, enabled_error_sampler(), false);
-
-            let header_tags = tracer_header_tags::TracerHeaderTags {
-                generic: tracer_header_tags::TracerGenericTags {
-                    client_computed_stats,
-                    ..Default::default()
-                },
-                ..create_test_header_tags()
-            };
+            let mut header_tags = create_test_header_tags();
+            header_tags.generic.client_computed_stats = client_computed_stats;
 
             // trace 1: kept normally (priority 1). trace 2: errored P0 (rescued).
             // trace 3: non-errored P0 (dropped). trace 4: errored user drop (dropped).
@@ -1560,15 +1533,7 @@ mod tests {
                 processor.process_traces(config, tags_provider, header_tags, traces, 0, None);
             let payload_info = payload_info.expect("expected Some payload");
             let backend_send_data = payload_info.builder.build();
-            let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads()
-            else {
-                panic!("expected V07");
-            };
-
-            let kept: Vec<(u64, f64)> = backend_payloads
-                .iter()
-                .flat_map(|tp| tp.chunks.iter())
-                .flat_map(|c| c.spans.iter())
+            let kept: Vec<(u64, f64)> = trace_spans(backend_send_data.get_payloads())
                 .map(|s| (s.trace_id, s.metrics["_sampling_priority_v1"]))
                 .collect();
             assert_eq!(
@@ -1584,19 +1549,8 @@ mod tests {
                 kept.contains(&(2, 0.0)),
                 "errored P0 trace rescued with its original priority (owner={owner})"
             );
-            assert!(
-                !kept.iter().any(|(id, _)| *id == 3),
-                "non-errored P0 trace dropped (owner={owner})"
-            );
-            assert!(
-                !kept.iter().any(|(id, _)| *id == 4),
-                "errored user-drop trace not rescued (owner={owner})"
-            );
 
-            let rescued_root = backend_payloads
-                .iter()
-                .flat_map(|tp| tp.chunks.iter())
-                .flat_map(|c| c.spans.iter())
+            let rescued_root = trace_spans(backend_send_data.get_payloads())
                 .find(|s| s.trace_id == 2)
                 .expect("rescued trace present");
             assert!(
@@ -1620,13 +1574,8 @@ mod tests {
             let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
                 create_test_processor(&config, enabled_error_sampler(), true);
-            let header_tags = tracer_header_tags::TracerHeaderTags {
-                generic: tracer_header_tags::TracerGenericTags {
-                    client_computed_stats,
-                    ..Default::default()
-                },
-                ..create_test_header_tags()
-            };
+            let mut header_tags = create_test_header_tags();
+            header_tags.generic.client_computed_stats = client_computed_stats;
 
             let span = make_tracer_drop_span(7, -1.0, 1, None);
 
diff --git a/bottlecap/tests/apm_integration_test.rs b/bottlecap/tests/apm_integration_test.rs
index 02e94e82e..d5b1903b7 100644
--- a/bottlecap/tests/apm_integration_test.rs
+++ b/bottlecap/tests/apm_integration_test.rs
@@ -547,12 +547,6 @@ async fn e2e_error_sampler_rescues_only_errored_p0_traces() {
             rescued_span.metrics.contains_key("_dd.errors_sr"),
             "the rescued root span must carry its error sampling rate (owner={owner})",
         );
-        assert!(
-            captured_spans
-                .iter()
-                .all(|span| span.trace_id != dropped_trace_id),
-            "the non-errored P0 trace must not reach intake (owner={owner})",
-        );
     }
 }
 
@@ -561,14 +555,14 @@ async fn e2e_error_sampler_rescues_only_errored_p0_traces() {
 /// stats keep them so the backend can compute stats.
 #[tokio::test]
 async fn e2e_sampled_out_chunks_filtered_by_stats_owner() {
-    // (owner label, lambda_extension_compute_stats, client_computed_stats) -> filtering expected
+    // (owner label, lambda_extension_compute_stats, client_computed_stats, expected_trace_ids)
     let cases = [
-        ("backend", false, false, false),
-        ("extension", true, false, true),
-        ("tracer", false, true, true),
+        ("backend", false, false, &[1, 2, 3][..]),
+        ("extension", true, false, &[1][..]),
+        ("tracer", false, true, &[1][..]),
     ];
 
-    for (owner, compute_on_extension, client_computed_stats, expect_filtering) in cases {
+    for (owner, compute_on_extension, client_computed_stats, expected_trace_ids) in cases {
         // kept (priority 1), auto-dropped (priority 0), explicit drop (priority -1)
         let outcome = run_processor_pipeline_with_traces(
             compute_on_extension,
@@ -581,21 +575,10 @@ async fn e2e_sampled_out_chunks_filtered_by_stats_owner() {
         )
         .await;
 
-        let mut captured_trace_ids: Vec<u64> = outcome
-            .traces
-            .iter()
-            .flat_map(|payload| &payload.tracer_payloads)
-            .flat_map(|payload| &payload.chunks)
-            .flat_map(|chunk| &chunk.spans)
+        let mut captured_trace_ids: Vec<u64> = captured_trace_spans(&outcome.traces)
             .map(|span| span.trace_id)
             .collect();
         captured_trace_ids.sort_unstable();
-
-        let expected_trace_ids: Vec<u64> = if expect_filtering {
-            vec![1]
-        } else {
-            vec![1, 2, 3]
-        };
         assert_eq!(
             captured_trace_ids, expected_trace_ids,
             "unexpected intake contents (owner={owner})"

From ccd4137cf00c4bf724a124d1757eb66af2ed4385 Mon Sep 17 00:00:00 2001
From: Lucas Pimentel <lucas.pimentel@datadoghq.com>
Date: Thu, 22 Oct 2026 14:28:55 -0400
Subject: [PATCH 6/6] Simplify trace sampling test setup

---
 bottlecap/src/traces/trace_processor.rs | 47 ++++---------------------
 bottlecap/tests/apm_integration_test.rs | 24 +++----------
 2 files changed, 11 insertions(+), 60 deletions(-)

diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs
index e50119d80..d98737497 100644
--- a/bottlecap/src/traces/trace_processor.rs
+++ b/bottlecap/src/traces/trace_processor.rs
@@ -1435,23 +1435,6 @@ mod tests {
             (true, true, 1),   // tracer owns stats even when extension configured: filter
         ];
 
-        let make_span = |trace_id: u64, priority: Option<f64>| -> pb::Span {
-            let mut metrics = HashMap::new();
-            if let Some(p) = priority {
-                metrics.insert("_sampling_priority_v1".to_string(), p);
-            }
-            pb::Span {
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                metrics,
-                service: "svc".to_string(),
-                name: "op".to_string(),
-                resource: "res".to_string(),
-                ..Default::default()
-            }
-        };
-
         for (compute_on_extension, client_computed_stats, expected_backend_span_count) in cases {
             let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
@@ -1461,9 +1444,9 @@ mod tests {
 
             // Three traces: kept (priority 1), dropped (priority 0), dropped (priority -1)
             let traces = vec![
-                vec![make_span(1, Some(1.0))],
-                vec![make_span(2, Some(0.0))],
-                vec![make_span(3, Some(-1.0))],
+                vec![make_tracer_drop_span(1, 1.0, 0, None)],
+                vec![make_tracer_drop_span(2, 0.0, 0, None)],
+                vec![make_tracer_drop_span(3, -1.0, 0, None)],
             ];
 
             let (payload_info, stats_collection) =
@@ -1497,22 +1480,6 @@ mod tests {
         // (owner label, lambda_extension_compute_stats, client_computed_stats)
         let cases = [("extension", true, false), ("tracer", false, true)];
 
-        let make_span = |trace_id: u64, priority: f64, error: i32| -> pb::Span {
-            let mut metrics = HashMap::new();
-            metrics.insert("_sampling_priority_v1".to_string(), priority);
-            pb::Span {
-                trace_id,
-                span_id: trace_id,
-                parent_id: 0,
-                error,
-                metrics,
-                service: "svc".to_string(),
-                name: "op".to_string(),
-                resource: "res".to_string(),
-                ..Default::default()
-            }
-        };
-
         for (owner, compute_on_extension, client_computed_stats) in cases {
             let config = create_compute_stats_config_with(compute_on_extension);
             let (tags_provider, processor) =
@@ -1523,10 +1490,10 @@ mod tests {
             // trace 1: kept normally (priority 1). trace 2: errored P0 (rescued).
             // trace 3: non-errored P0 (dropped). trace 4: errored user drop (dropped).
             let traces = vec![
-                vec![make_span(1, 1.0, 0)],
-                vec![make_span(2, 0.0, 1)],
-                vec![make_span(3, 0.0, 0)],
-                vec![make_span(4, -1.0, 1)],
+                vec![make_tracer_drop_span(1, 1.0, 0, None)],
+                vec![make_tracer_drop_span(2, 0.0, 1, None)],
+                vec![make_tracer_drop_span(3, 0.0, 0, None)],
+                vec![make_tracer_drop_span(4, -1.0, 1, None)],
             ];
 
             let (payload_info, _stats) =
diff --git a/bottlecap/tests/apm_integration_test.rs b/bottlecap/tests/apm_integration_test.rs
index d5b1903b7..11417d540 100644
--- a/bottlecap/tests/apm_integration_test.rs
+++ b/bottlecap/tests/apm_integration_test.rs
@@ -495,7 +495,7 @@ fn captured_compute_stats(traces: &[pb::AgentPayload]) -> Option<String> {
 
 /// Single-span trace root with the given sampling priority and error flag.
 fn sampled_root_span(trace_id: u64, priority: f64, error: i32) -> pb::Span {
-    let mut span = pb::Span {
+    pb::Span {
         service: "fake-intake-trace-service".to_string(),
         name: "web.request".to_string(),
         resource: "GET /fake".to_string(),
@@ -506,11 +506,9 @@ fn sampled_root_span(trace_id: u64, priority: f64, error: i32) -> pb::Span {
         duration: 5_000_000,
         error,
         r#type: "web".to_string(),
+        metrics: HashMap::from([("_sampling_priority_v1".to_string(), priority)]),
         ..pb::Span::default()
-    };
-    span.metrics
-        .insert("_sampling_priority_v1".to_string(), priority);
-    span
+    }
 }
 
 /// An errored P0 trace is rescued and reaches the intake, while a non-errored P0 trace is
@@ -643,22 +641,8 @@ fn error_sampler_span(
     error: i32,
     decision_maker: Option<&str>,
 ) -> pb::Span {
-    let mut span = pb::Span {
-        service: "fake-intake-trace-service".to_string(),
-        name: "web.request".to_string(),
-        resource: "GET /fake".to_string(),
-        trace_id,
-        span_id: trace_id,
-        parent_id: 0,
-        start: STATS_SPAN_START_NS,
-        duration: 5_000_000,
-        error,
-        r#type: "web".to_string(),
-        ..pb::Span::default()
-    };
+    let mut span = sampled_root_span(trace_id, priority, error);
     span.metrics.insert("_top_level".to_string(), 1.0);
-    span.metrics
-        .insert("_sampling_priority_v1".to_string(), priority);
     if let Some(decision_maker) = decision_maker {
         span.meta
             .insert("_dd.p.dm".to_string(), decision_maker.to_string());