From 8b61cdc216512fd967715b930dd3cf99bc936ced Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 14:23:19 -0400 Subject: [PATCH 01/49] Rescue errored P0 trace chunks from backend drop Wire the shared datadog-agent-trace-sampler into the SCL trace processor so error chunks with automatic-drop priority (0) get a second look before being forwarded. On a keep, the sampler's rate is stamped as a positive `_dd.errors_sr` metric on the chunk's root span, which the APM backend treats as an `error` ingestion reason and retains without any priority promotion. Chunks are never dropped locally: unrescued chunks are forwarded unchanged for the backend to discard. - Default to rate-limited rescue at 10 error TPS per process, with a shared budget across requests and processor clones - Configure via DD_APM_ERROR_SAMPLER_MODE (rate_limited | always_keep) and DD_APM_ERROR_TPS (<= 0 disables rescue in either mode); invalid values warn and fall back to defaults - Gate rescue on agent stats computation and submit all chunks to the stats concentrator before any rescue stamping --- Cargo.lock | 1 + crates/datadog-serverless-compat/src/main.rs | 1 + crates/datadog-trace-agent/Cargo.toml | 1 + crates/datadog-trace-agent/src/config.rs | 173 ++++ .../src/stats_processor.rs | 1 + .../src/trace_processor.rs | 952 +++++++++++++++++- .../tests/integration_test.rs | 225 ++++- 7 files changed, 1340 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 71307f53..d28ba062 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -647,6 +647,7 @@ dependencies = [ "anyhow", "async-trait", "bytes", + "datadog-agent-trace-sampler", "datadog-fips", "datadog-mock-intake", "duplicate", diff --git a/crates/datadog-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs index 4f5ecdc5..8254a248 100644 --- a/crates/datadog-serverless-compat/src/main.rs +++ b/crates/datadog-serverless-compat/src/main.rs @@ -208,6 +208,7 @@ pub async fn main() { let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor::new( stats_concentrator.as_ref().map(|c| c.handle.clone()), + trace_processor::ServerlessTraceProcessor::new_error_sampler(&config), )); let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher { diff --git a/crates/datadog-trace-agent/Cargo.toml b/crates/datadog-trace-agent/Cargo.toml index 406a30c8..a9b73ded 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -33,6 +33,7 @@ libdd-trace-stats = { version = "13.0.0", features = ["stats-obfuscation"] } libdd-common = { workspace = true, features = ["https"] } libdd-trace-obfuscation = { workspace = true, features = ["https"] } libdd-trace-utils = { workspace = true, features = ["https", "mini_agent"] } +datadog-agent-trace-sampler = { path = "../datadog-agent-trace-sampler" } datadog-fips = { path = "../datadog-fips" } reqwest = { version = "0.12.23", features = [ "json", diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index c97384d7..a9c9d2bc 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -9,14 +9,60 @@ use std::env; use std::str::FromStr; use std::sync::OnceLock; +use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; use libdd_trace_obfuscation::obfuscation_config; use libdd_trace_utils::config_utils::{ read_cloud_env, trace_intake_url, trace_intake_url_prefixed, trace_stats_url, trace_stats_url_prefixed, }; use libdd_trace_utils::trace_utils; +use tracing::warn; const DEFAULT_APM_RECEIVER_PORT: u16 = 8126; +/// Default error rescue budget: error traces per second, matching the Go agent's `ErrorTPS`. +const DEFAULT_ERROR_SAMPLER_TPS: f64 = 10.0; + +/// Parses the error rescue sampler settings from the environment. +/// +/// - `DD_APM_ERROR_SAMPLER_MODE`: `rate_limited` (default) or `always_keep`. +/// Surrounding whitespace and casing are normalized; any other value warns +/// and falls back to `rate_limited`. +/// - `DD_APM_ERROR_TPS`: target error traces per second, default 10. Parsed as +/// `f64`; non-finite or unparsable values warn and fall back to the default. +/// A finite value <= 0 disables rescue in either mode. +/// +/// `extra_sample_rate` is fixed at 1.0 and intentionally not configurable. +/// Startup stays operational when either optional setting is malformed. +fn parse_error_sampler_config() -> ErrorSamplerConfig { + let mut config = ErrorSamplerConfig { + mode: ErrorSamplerMode::RateLimited, + target_tps: DEFAULT_ERROR_SAMPLER_TPS, + extra_sample_rate: 1.0, + }; + + if let Ok(raw) = env::var("DD_APM_ERROR_SAMPLER_MODE") { + match raw.trim().to_lowercase().as_str() { + "rate_limited" => config.mode = ErrorSamplerMode::RateLimited, + "always_keep" => config.mode = ErrorSamplerMode::AlwaysKeep, + _ => { + warn!("Invalid DD_APM_ERROR_SAMPLER_MODE {raw:?}; using default mode rate_limited"); + } + } + } + + if let Ok(raw) = env::var("DD_APM_ERROR_TPS") { + match raw.trim().parse::() { + Ok(tps) if tps.is_finite() => config.target_tps = tps, + _ => { + warn!( + "Invalid DD_APM_ERROR_TPS {raw:?}; using default {DEFAULT_ERROR_SAMPLER_TPS}" + ); + } + } + } + + config +} const DEFAULT_DOGSTATSD_PORT: u16 = 8125; const DSM_PIPELINE_STATS_ROUTE: &str = "/api/v0.1/pipeline_stats"; @@ -128,6 +174,10 @@ pub struct Config { pub additional_metric_tags_cardinality_limit: Option, /// Whether the agent should compute trace stats pub agent_stats_computation_enabled: bool, + /// Error rescue sampler settings, parsed from `DD_APM_ERROR_SAMPLER_MODE` + /// and `DD_APM_ERROR_TPS`. `extra_sample_rate` is fixed at 1.0 and not + /// configurable. See `parse_error_sampler_config` for the defaults. + pub error_sampler: ErrorSamplerConfig, } impl Config { @@ -308,6 +358,7 @@ impl Config { agent_stats_computation_enabled: env::var("DD_AGENT_STATS_COMPUTATION_ENABLED") .map(|val| val.to_lowercase() == "true") .unwrap_or(true), + error_sampler: parse_error_sampler_config(), }) } } @@ -319,6 +370,7 @@ mod tests { use std::collections::HashMap; use crate::config; + use datadog_agent_trace_sampler::ErrorSamplerMode; #[test] #[serial] @@ -915,6 +967,126 @@ mod tests { }, ); } + + fn assert_rate_limited_default(error_sampler: &config::ErrorSamplerConfig) { + assert!( + matches!(error_sampler.mode, ErrorSamplerMode::RateLimited), + "expected default mode RateLimited" + ); + assert_eq!(error_sampler.target_tps, 10.0); + assert_eq!(error_sampler.extra_sample_rate, 1.0); + } + + #[test] + #[serial] + fn test_error_sampler_defaults() { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ], + || { + let config = config::Config::new().unwrap(); + assert_rate_limited_default(&config.error_sampler); + }, + ); + } + + #[test] + #[serial] + fn test_error_sampler_mode_normalization() { + for (raw, expected) in [ + ("always_keep", ErrorSamplerMode::AlwaysKeep), + ("ALWAYS_KEEP", ErrorSamplerMode::AlwaysKeep), + (" Always_Keep ", ErrorSamplerMode::AlwaysKeep), + ("rate_limited", ErrorSamplerMode::RateLimited), + ("RATE_LIMITED", ErrorSamplerMode::RateLimited), + (" Rate_Limited ", ErrorSamplerMode::RateLimited), + ] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), + ], + || { + let config = config::Config::new().unwrap(); + assert!( + matches!(config.error_sampler.mode, m if std::mem::discriminant(&m) == std::mem::discriminant(&expected)), + "mode {raw:?} should parse to {expected:?}, got {:?}", + config.error_sampler.mode + ); + }, + ); + } + } + + #[test] + #[serial] + fn test_error_sampler_invalid_mode_falls_back_to_default() { + for raw in ["bogus", "", "always-keep"] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), + ], + || { + let config = config::Config::new().unwrap(); + assert_rate_limited_default(&config.error_sampler); + }, + ); + } + } + + #[test] + #[serial] + fn test_error_sampler_tps_parsing() { + for (raw, expected_tps) in [("5.5", 5.5), ("0", 0.0), ("-1", -1.0), (" 10 ", 10.0)] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_ERROR_TPS", Some(raw)), + ], + || { + let config = config::Config::new().unwrap(); + assert_eq!( + config.error_sampler.target_tps, expected_tps, + "DD_APM_ERROR_TPS {raw:?} should parse to {expected_tps}" + ); + }, + ); + } + } + + #[test] + #[serial] + fn test_error_sampler_invalid_tps_falls_back_to_default() { + for raw in ["not_a_number", "inf", "-inf", "NaN", "1e400", ""] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_ERROR_TPS", Some(raw)), + ], + || { + let config = config::Config::new().unwrap(); + assert_rate_limited_default(&config.error_sampler); + }, + ); + } + } } /// Test helpers for creating Config instances in tests @@ -957,6 +1129,7 @@ pub mod test_helpers { additional_metric_tags: vec![], additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled: true, + error_sampler: ErrorSamplerConfig::default(), } } } diff --git a/crates/datadog-trace-agent/src/stats_processor.rs b/crates/datadog-trace-agent/src/stats_processor.rs index 9b7f5394..104eaea3 100644 --- a/crates/datadog-trace-agent/src/stats_processor.rs +++ b/crates/datadog-trace-agent/src/stats_processor.rs @@ -164,6 +164,7 @@ mod tests { additional_metric_tags: vec![], additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled, + error_sampler: datadog_agent_trace_sampler::ErrorSamplerConfig::default(), } } diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index b1f2ed65..caabc916 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -4,10 +4,13 @@ use std::sync::Arc; use async_trait::async_trait; +use datadog_agent_trace_sampler::{ErrorsSampler, SampleDecision, SpanView, TraceView}; use http_body_util::BodyExt; use hyper::{StatusCode, http}; use libdd_common::http_common; use libdd_library_config::tracer_metadata::TracerMetadata; +use std::sync::{Mutex, MutexGuard, PoisonError}; +use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc::Sender; use tracing::{debug, error, warn}; @@ -27,6 +30,12 @@ use crate::{ const TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY: &str = "_dd.tags.function"; +/// The root-span metric the backend uses to rescue traces the ordinary P0 drop +/// would discard: a positive `_dd.errors_sr` resolves the chunk's ingestion +/// reason to `error`, which is retained at low priority without any priority +/// promotion. +const ERRORS_SR_METRIC_KEY: &str = "_dd.errors_sr"; + /// Rough upper bound on the protobuf framing overhead added when a V07 `TracerPayload` is /// wrapped in the outer `AgentPayload` envelope before being sent const V07_ENVELOPE_OVERHEAD_BYTES: usize = 64; @@ -124,14 +133,75 @@ const MAX_IN_FLIGHT_ENQUEUES: usize = 10; pub struct ServerlessTraceProcessor { pub stats_concentrator: Option, enqueue_permits: Arc, + /// Shared error rescue sampler. The `Arc` means processor clones (one per + /// connection) share a single sampler state and TPS budget for the whole + /// process lifetime. + error_sampler: Arc>, } impl ServerlessTraceProcessor { #[allow(clippy::must_use_candidate)] - pub fn new(stats_concentrator: Option) -> Self { + pub fn new( + stats_concentrator: Option, + error_sampler: Arc>, + ) -> Self { ServerlessTraceProcessor { stats_concentrator, enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), + error_sampler, + } + } + + /// Builds the error rescue sampler from parsed config settings. + #[must_use] + pub fn new_error_sampler(config: &Config) -> Arc> { + Arc::new(Mutex::new(ErrorsSampler::new(config.error_sampler))) + } + + /// Locks the shared sampler, recovering from a poisoned guard (a panic in + /// another thread while holding the lock) instead of panicking on unlock. + fn lock_sampler(&self) -> MutexGuard<'_, ErrorsSampler> { + self.error_sampler + .lock() + .unwrap_or_else(PoisonError::into_inner) + } + + /// Applies the error rescue pass to a payload collection. + /// + /// For every V07 chunk that would be dropped by the backend's ordinary P0 + /// drop (chunk priority is exactly 0, i.e. an automatic drop, not an + /// explicit user drop or the no-priority sentinel) and contains at least + /// one span with a non-zero error flag, consults the shared error sampler. + /// On a keep, stamps the sampler's `_dd.errors_sr` on the chunk's root + /// span; the backend keeps such chunks without any priority promotion. + /// Chunks are never removed or reordered: unrescued chunks stay in the + /// payload and the backend discards them. + /// + /// `now_unix_secs` is passed in explicitly so tests can exercise the + /// rolling window without sleeping. + fn apply_error_rescue( + &self, + payload: &mut TracerPayloadCollection, + config: &Config, + now_unix_secs: i64, + ) { + let TracerPayloadCollection::V07(tracer_payloads) = payload else { + return; + }; + let mut sampler = self.lock_sampler(); + // Skip all view construction when the sampler is disabled by config + // (target_tps <= 0): nothing can be rescued. + if sampler.is_disabled() { + return; + } + for tracer_payload in tracer_payloads.iter_mut() { + // The sampler keys its per-signature rate limits on the env the + // tracer reported for this payload, falling back to the agent's + // configured env, consistent with how stats are flushed. + let env: &str = resolve_payload_env(&tracer_payload.env, &config.env); + for chunk in tracer_payload.chunks.iter_mut() { + sample_and_stamp(&mut sampler, chunk, env, now_unix_secs); + } } } @@ -180,6 +250,86 @@ impl ServerlessTraceProcessor { } } +/// Chooses the env used to key the error sampler's per-signature rates: the +/// tracer payload's own env when nonempty, otherwise the agent's configured +/// env. This matches how stats prefer the payload env and fall back to the +/// agent config env. +fn resolve_payload_env<'a>(tracer_payload_env: &'a str, config_env: &'a str) -> &'a str { + if tracer_payload_env.is_empty() { + config_env + } else { + tracer_payload_env + } +} + +/// Builds the sampler's read-only view of a span from its already enriched and +/// obfuscated fields. +fn span_view(span: &pb::Span) -> SpanView<'_> { + SpanView { + service: &span.service, + name: &span.name, + resource: &span.resource, + error: span.error != 0, + http_status_code: span.meta.get("http.status_code").map(String::as_str), + error_type: span.meta.get("error.type").map(String::as_str), + } +} + +/// Consults the error sampler for one chunk and, on a keep, stamps the +/// sampler's `_dd.errors_sr` value on the chunk's root span. +/// +/// Only automatic-drop chunks are candidates: chunk priority must be exactly 0. +/// Explicit user drops (-1), other negative priorities, positive priorities, +/// and the no-priority sentinel (`i8::MIN`) are all left untouched. The chunk +/// must contain at least one span with a non-zero error flag; HTTP status or +/// error metadata alone does not qualify. The root is resolved with +/// `get_root_span_index`; empty or rootless chunks are left unchanged. Chunks +/// are never removed here: on a Drop decision the unrescued chunk is forwarded +/// as-is and the backend's ordinary P0 drop handles it. +fn sample_and_stamp( + sampler: &mut ErrorsSampler, + chunk: &mut pb::TraceChunk, + env: &str, + now_unix_secs: i64, +) { + if chunk.priority != 0 { + return; + } + // An error anywhere in the chunk makes it a rescue candidate, not just an + // error on the root span. Checked before the root-span search because + // non-errored chunks are the common case on this path. + if !chunk.spans.iter().any(|span| span.error != 0) { + return; + } + let Ok(root_index) = trace_utils::get_root_span_index(&chunk.spans) else { + // No identifiable root span: leave the chunk unchanged rather than + // guessing an index. + return; + }; + + let Some(root) = chunk.spans.get(root_index) else { + return; + }; + let views: Vec = chunk.spans.iter().map(span_view).collect(); + let trace = TraceView { + env, + trace_id: root.trace_id, + root_index, + // The raw `_sample_rate` wire value is passed through: the shared + // sampler sanitizes non-finite or out-of-range rates to 1.0 itself. + root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), + spans: &views, + }; + let decision = sampler.sample(now_unix_secs, &trace); + + if let SampleDecision::Keep { errors_sr } = decision + && let Some(root) = chunk.spans.get_mut(root_index) + { + root.metrics + .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); + } +} + #[async_trait] impl TraceProcessor for ServerlessTraceProcessor { async fn process_traces( @@ -283,6 +433,17 @@ impl TraceProcessor for ServerlessTraceProcessor { Self::send_to_concentrator(concentrator, &payload); } + // Error rescue runs after stats submission so the concentrator observes every + // submitted chunk exactly as the tracer sent it, and before payload splitting so + // the newly inserted metric is included in the recomputed outbound size. It is + // gated on agent stats computation: without it, P0 chunks are not expected here. + if config.agent_stats_computation_enabled { + let now_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_secs() as i64); + self.apply_error_rescue(&mut payload, &config, now_unix_secs); + } + let pieces: Vec<(TracerPayloadCollection, usize)> = match payload { TracerPayloadCollection::V07(payloads) => { let split_budget = @@ -361,6 +522,9 @@ mod tests { encoded_size, split_oversized_payloads, }, }; + use datadog_agent_trace_sampler::{ + ErrorSamplerConfig, ErrorSamplerMode, ErrorsSampler, SampleDecision, SpanView, TraceView, + }; use libdd_common::{Endpoint, http_common}; use libdd_trace_protobuf::pb; use libdd_trace_utils::test_utils::{create_test_gcp_json_span, create_test_gcp_span}; @@ -425,9 +589,16 @@ mod tests { additional_metric_tags: vec![], additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled: false, + error_sampler: ErrorSamplerConfig::default(), } } + fn default_error_sampler() -> Arc> { + Arc::new(std::sync::Mutex::new(ErrorsSampler::new( + ErrorSamplerConfig::default(), + ))) + } + fn create_test_metadata() -> MiniAgentMetadata { MiniAgentMetadata { azure_spring_app_hostname: Default::default(), @@ -543,7 +714,8 @@ mod tests { .body(http_common::Body::from(bytes)) .unwrap(); - let trace_processor = trace_processor::ServerlessTraceProcessor::new(None); + let trace_processor = + trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -616,7 +788,8 @@ mod tests { .body(http_common::Body::from(bytes)) .unwrap(); - let trace_processor = trace_processor::ServerlessTraceProcessor::new(None); + let trace_processor = + trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -702,7 +875,8 @@ mod tests { .body(http_common::Body::from(bytes)) .unwrap(); - let trace_processor = trace_processor::ServerlessTraceProcessor::new(None); + let trace_processor = + trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -760,7 +934,8 @@ mod tests { Receiver, ) = mpsc::channel(1); - let trace_processor = trace_processor::ServerlessTraceProcessor::new(None); + let trace_processor = + trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); let config = Arc::new(Config { enqueue_permit_timeout_secs: 0, ..create_test_config() @@ -828,7 +1003,8 @@ mod tests { Receiver, ) = mpsc::channel(MAX_IN_FLIGHT_ENQUEUES + 1); - let trace_processor = trace_processor::ServerlessTraceProcessor::new(None); + let trace_processor = + trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); // Uses the default (non-zero) enqueue_permit_timeout_secs, so the 11th request's permit // acquire has real time to succeed once an earlier request's send completes and // releases its permit, rather than racing a 0-second timeout against task scheduling. @@ -869,4 +1045,768 @@ mod tests { allowing all of them to succeed rather than shedding load" ); } + + // ---- Error rescue tests ---- + + const RESCUE_NOW: i64 = 1_700_000_000; + + fn test_span(trace_id: u64, span_id: u64, parent_id: u64, error: i32) -> pb::Span { + pb::Span { + service: "test-service".to_string(), + name: "test-operation".to_string(), + resource: "GET /test".to_string(), + trace_id, + span_id, + parent_id, + error, + ..Default::default() + } + } + + fn test_chunk(spans: Vec, priority: i32) -> pb::TraceChunk { + pb::TraceChunk { + spans, + priority, + ..Default::default() + } + } + + fn errored_root_chunk(trace_id: u64, priority: i32) -> pb::TraceChunk { + test_chunk(vec![test_span(trace_id, trace_id + 1, 0, 1)], priority) + } + + fn rescue_config(error_sampler: ErrorSamplerConfig) -> Config { + Config { + agent_stats_computation_enabled: true, + env: "agent-env".to_string(), + error_sampler, + ..create_test_config() + } + } + + fn always_keep_sampler_config() -> ErrorSamplerConfig { + ErrorSamplerConfig { + mode: ErrorSamplerMode::AlwaysKeep, + target_tps: 10.0, + extra_sample_rate: 1.0, + } + } + + fn rate_limited_sampler_config(target_tps: f64) -> ErrorSamplerConfig { + ErrorSamplerConfig { + mode: ErrorSamplerMode::RateLimited, + target_tps, + extra_sample_rate: 1.0, + } + } + + fn sampler_for(config: &ErrorSamplerConfig) -> Arc> { + Arc::new(std::sync::Mutex::new(ErrorsSampler::new(*config))) + } + + fn run_rescue( + processor: &trace_processor::ServerlessTraceProcessor, + config: &Config, + chunks: Vec, + now_unix_secs: i64, + ) -> Vec { + let mut payload = TracerPayloadCollection::V07(vec![pb::TracerPayload { + chunks, + ..Default::default() + }]); + processor.apply_error_rescue(&mut payload, config, now_unix_secs); + match payload { + TracerPayloadCollection::V07(mut payloads) => payloads.remove(0).chunks, + _ => unreachable!(), + } + } + + fn root(chunk: &pb::TraceChunk) -> &pb::Span { + // All fixtures in this section place the root first or resolve it the + // same way the processor does. + &chunk.spans[0] + } + + fn errors_sr(span: &pb::Span) -> Option { + span.metrics.get("_dd.errors_sr").copied() + } + + /// Builds the equivalent standalone `TraceView` for a fixture chunk, used to + /// compare adapter decisions with direct shared-sampler usage. + fn equivalent_trace_view<'a>( + chunk: &pb::TraceChunk, + env: &'a str, + views: &'a [SpanView<'a>], + ) -> TraceView<'a> { + let root_index = trace_utils::get_root_span_index(&chunk.spans).unwrap(); + let root = &chunk.spans[root_index]; + TraceView { + env, + trace_id: root.trace_id, + root_index, + root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), + spans: views, + } + } + + #[test] + fn test_rescue_stamps_errored_p0_root_and_keeps_priority() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let chunks = run_rescue( + &processor, + &config, + vec![errored_root_chunk(0xdead_beef, 0)], + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 1, "rescued chunk must still be forwarded"); + assert_eq!(chunks[0].priority, 0, "priority must not be promoted"); + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + } + + #[test] + fn test_rescue_stamps_actual_root_when_error_is_on_child() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // Healthy root first, errored child second: only the root may be stamped. + let spans = vec![ + test_span(0xfeed, 0x101, 0, 0), + test_span(0xfeed, 0x102, 0x101, 1), + ]; + let chunks = run_rescue(&processor, &config, vec![test_chunk(spans, 0)], RESCUE_NOW); + + assert_eq!(chunks.len(), 1); + assert_eq!(errors_sr(&chunks[0].spans[0]), Some(1.0), "root stamped"); + assert_eq!( + errors_sr(&chunks[0].spans[1]), + None, + "errored child must not be stamped" + ); + } + + #[test] + fn test_rescue_uses_resolved_root_not_first_span() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // Root (parent_id 0) appears last; the first span is a child. + let spans = vec![ + test_span(0xbeef, 0x201, 0x203, 0), + test_span(0xbeef, 0x202, 0x203, 1), + test_span(0xbeef, 0x203, 0, 0), + ]; + let chunks = run_rescue(&processor, &config, vec![test_chunk(spans, 0)], RESCUE_NOW); + + assert_eq!(chunks.len(), 1); + assert_eq!( + errors_sr(&chunks[0].spans[2]), + Some(1.0), + "the resolved root (last span) must be stamped" + ); + assert_eq!(errors_sr(&chunks[0].spans[0]), None); + assert_eq!(errors_sr(&chunks[0].spans[1]), None); + } + + #[test] + fn test_no_error_flags_leaves_chunk_unchanged_and_spends_no_budget() { + // With RateLimited, counting is global across signatures via + // `all_sigs_seen`, so if the healthy chunks below were wrongly fed to + // the sampler, the final errored chunk's rate would drop below 1.0 and + // it would not be stamped with 1.0. Correct behavior: only errored + // chunks reach the sampler, so the errored chunk is the first count and + // is kept at rate 1.0. + let config = rescue_config(rate_limited_sampler_config(1.0)); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let mut chunks = Vec::new(); + for i in 0..9_u64 { + chunks.push(errored_root_chunk(0xa000 + i, 0)); + // Strip the error flag: healthy P0 chunk with a distinct signature. + chunks.last_mut().unwrap().spans[0].error = 0; + } + chunks.push(errored_root_chunk(0xb000, 0)); + + let chunks = run_rescue(&processor, &config, chunks, RESCUE_NOW); + + for (i, chunk) in chunks.iter().enumerate() { + let is_last = i == chunks.len() - 1; + assert_eq!( + errors_sr(root(chunk)), + if is_last { Some(1.0) } else { None }, + "chunk {i}: only the errored chunk may be rescued" + ); + } + } + + #[test] + fn test_http_500_metadata_alone_is_not_an_error() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let mut span = test_span(0xc0de, 0xc0de + 1, 0, 0); + span.meta + .insert("http.status_code".to_string(), "500".to_string()); + span.meta + .insert("error.type".to_string(), "Error".to_string()); + + let chunks = run_rescue( + &processor, + &config, + vec![test_chunk(vec![span], 0)], + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 1); + assert_eq!(errors_sr(root(&chunks[0])), None); + } + + #[test] + fn test_non_automatic_drop_priorities_are_never_rescued() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // -1 (explicit user drop), other negatives, positive priorities, and + // the no-priority sentinel (i8::MIN) are all out of scope. + let priorities = [-1_i32, -5, 1, 2, i8::MIN as i32]; + let chunks = run_rescue( + &processor, + &config, + priorities + .iter() + .map(|p| errored_root_chunk(0x1000_u64 + (*p).unsigned_abs() as u64, *p)) + .collect(), + RESCUE_NOW, + ); + + for chunk in &chunks { + assert_eq!( + errors_sr(root(chunk)), + None, + "priority {} must not be rescued", + chunk.priority + ); + } + } + + #[test] + fn test_empty_chunk_does_not_panic_and_stays_unchanged() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // An empty chunk cannot be scored: root resolution fails and it must be + // left unchanged without panicking. For non-empty chunks, + // `get_root_span_index` always resolves (falling back to the last span + // when no span has parent_id 0), so the cyclic chunk below exercises + // the fallback path and gets stamped on the resolved fallback root. + let cyclic = test_chunk( + vec![ + test_span(0xd00d, 0xd01, 0xd02, 1), + test_span(0xd00d, 0xd02, 0xd01, 0), + ], + 0, + ); + + let chunks = run_rescue( + &processor, + &config, + vec![test_chunk(vec![], 0), cyclic], + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 2); + assert!(chunks[0].spans.is_empty(), "empty chunk unchanged"); + assert_eq!( + errors_sr(&chunks[1].spans[1]), + Some(1.0), + "fallback root (last span) stamped" + ); + assert_eq!(errors_sr(&chunks[1].spans[0]), None); + } + + #[tokio::test] + async fn test_rescue_is_noop_when_agent_stats_disabled() { + let config = rescue_config(always_keep_sampler_config()); + let config = Config { + agent_stats_computation_enabled: false, + ..config + }; + let (tx, mut rx): ( + Sender, + tokio::sync::mpsc::Receiver, + ) = mpsc::channel(1); + + let start = get_current_timestamp_nanos(); + let mut json_span = create_test_json_span(11, 222, 333, start, true); + // Root span with an error and an automatic-drop priority: eligible. + json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); + let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); + let request = Request::builder() + .header("datadog-meta-tracer-version", "4.0.0") + .header("datadog-meta-lang", "nodejs") + .header("datadog-meta-lang-version", "v19.7.0") + .header("datadog-meta-lang-interpreter", "v8") + .header("content-length", "100") + .body(http_common::Body::from(bytes)) + .unwrap(); + + let trace_processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + let res = trace_processor + .process_traces( + Arc::new(config), + request, + tx, + Arc::new(create_test_metadata()), + ) + .await; + assert!(res.is_ok()); + + let send_data = rx.recv().await.expect("payload forwarded"); + let payloads = send_data.get_payloads(); + let TracerPayloadCollection::V07(tracer_payloads) = payloads else { + panic!("expected V07 payload"); + }; + let chunk = &tracer_payloads[0].chunks[0]; + assert_eq!(chunk.priority, 0); + for span in &chunk.spans { + assert_eq!( + errors_sr(span), + None, + "rescue must be a no-op when agent stats computation is disabled" + ); + } + } + + #[test] + fn test_disabled_tps_disables_rescue_in_both_modes() { + for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { + for tps in [0.0_f64, -3.0] { + let sampler_config = ErrorSamplerConfig { + mode, + target_tps: tps, + extra_sample_rate: 1.0, + }; + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let chunks = run_rescue( + &processor, + &config, + vec![errored_root_chunk(0xe000, 0)], + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 1, "chunk still forwarded"); + assert_eq!( + errors_sr(root(&chunks[0])), + None, + "mode {mode:?} tps {tps}: disabled sampler must not rescue" + ); + } + } + } + + #[test] + fn test_always_keep_rescues_every_eligible_chunk() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let chunks = run_rescue( + &processor, + &config, + (0..10_u64) + .map(|i| errored_root_chunk(0xf000 + i, 0)) + .collect(), + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 10); + for chunk in &chunks { + assert_eq!(errors_sr(root(chunk)), Some(1.0)); + assert_eq!(chunk.priority, 0); + } + } + + #[test] + fn test_rate_limited_under_sustained_load_keeps_and_rejects() { + let sampler_config = rate_limited_sampler_config(10.0); + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // 100 chunks with the same signature in one 5-second bucket: the + // default rate is 10 / (100 / 5) = 0.5, so a deterministic mix of keeps + // and rejections is expected. No exact keep count is asserted: the + // sampler is an adaptive rolling-window sampler, not a hard token + // bucket. + let ids: Vec = (0..100_u64).map(|i| 0x11_0000 + i).collect(); + let input: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); + let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); + + let keeps = chunks + .iter() + .filter(|c| errors_sr(root(c)).is_some()) + .count(); + let drops = chunks.len() - keeps; + assert!(keeps > 0, "expected some keeps, got none"); + assert!(drops > 0, "expected some drops, got none"); + assert_eq!(chunks.len(), 100, "every chunk is still forwarded"); + for chunk in &chunks { + assert_eq!(chunk.priority, 0, "priority unchanged on rescue decision"); + } + + // The same input through a standalone shared sampler must produce the + // identical decision sequence, validating the adapter wiring (env, + // sample rate, views) against direct crate usage. + let mut standalone = ErrorsSampler::new(sampler_config); + for (chunk, id) in chunks.iter().zip(&ids) { + let views: Vec = chunk + .spans + .iter() + .map(|s| SpanView { + service: &s.service, + name: &s.name, + resource: &s.resource, + error: s.error != 0, + http_status_code: s.meta.get("http.status_code").map(String::as_str), + error_type: s.meta.get("error.type").map(String::as_str), + }) + .collect(); + let trace = equivalent_trace_view(chunk, "test-env", &views); + let expected = standalone.sample(RESCUE_NOW, &trace); + let actual = match errors_sr(root(chunk)) { + Some(errors_sr) => SampleDecision::Keep { errors_sr }, + None => SampleDecision::Drop, + }; + assert_eq!(actual, expected, "decision mismatch for trace id {id}"); + } + } + + #[test] + fn test_rate_limited_bucket_transitions_and_steady_state() { + let sampler_config = rate_limited_sampler_config(10.0); + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + // First bucket: 100 distinct signatures push the default rate to 0.5. + let ids: Vec = (0..100_u64).map(|i| 0x22_0000 + i).collect(); + let first: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); + let first_chunks = run_rescue(&processor, &config, first, RESCUE_NOW); + let first_keeps = first_chunks + .iter() + .filter(|c| errors_sr(root(c)).is_some()) + .count(); + assert!( + first_keeps > 0 && first_keeps < 100, + "mixed decisions in the first bucket" + ); + + // A full window later (the rolling window is 6 buckets of 5 seconds), + // the same IDs in a fresh bucket still produce mixed decisions, and + // every chunk is still forwarded. + let second: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); + let second_chunks = run_rescue(&processor, &config, second, RESCUE_NOW + 40); + assert_eq!(second_chunks.len(), 100); + let second_keeps = second_chunks + .iter() + .filter(|c| errors_sr(root(c)).is_some()) + .count(); + assert!( + second_keeps > 0 && second_keeps < 100, + "mixed decisions after window rotation" + ); + } + + #[test] + fn test_processor_clones_share_one_budget() { + let sampler_config = rate_limited_sampler_config(1.0); + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + let processor_clone = processor.clone(); + + // Ten errored chunks with the same signature but distinct IDs: five + // through the original, five through the clone. All must draw from the + // same budget, matching a standalone sampler fed the same sequence. + let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); + let mut observed = Vec::new(); + for (i, id) in ids.iter().enumerate() { + let target = if i < 5 { &processor } else { &processor_clone }; + let chunks = run_rescue( + target, + &config, + vec![errored_root_chunk(*id, 0)], + RESCUE_NOW, + ); + observed.push(errors_sr(root(&chunks[0])).is_some()); + } + + let mut standalone = ErrorsSampler::new(sampler_config); + let spans = [SpanView { + service: "test-service", + name: "test-operation", + resource: "GET /test", + error: true, + http_status_code: None, + error_type: None, + }]; + for (i, id) in ids.iter().enumerate() { + let trace = TraceView { + env: "test-env", + trace_id: *id, + root_index: 0, + root_global_sample_rate: 1.0, + spans: &spans, + }; + let expected = matches!( + standalone.sample(RESCUE_NOW, &trace), + SampleDecision::Keep { .. } + ); + assert_eq!( + observed[i], expected, + "clone {i} decision diverged from the shared-budget standalone sampler" + ); + } + assert!( + observed.iter().any(|kept| *kept), + "expected at least one keep" + ); + assert!( + !observed.iter().all(|kept| *kept), + "expected at least one drop" + ); + } + + #[test] + fn test_rescue_uses_payload_env_with_config_fallback() { + assert_eq!( + super::resolve_payload_env("tracer-env", "agent-env"), + "tracer-env", + "nonempty payload env wins" + ); + assert_eq!( + super::resolve_payload_env("", "agent-env"), + "agent-env", + "empty payload env falls back to the agent config env" + ); + assert_eq!( + super::resolve_payload_env("", ""), + "", + "both empty stays empty" + ); + } + + #[test] + fn test_rescue_preserves_chunk_metadata_and_span_order() { + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let mut root_span = test_span(0x501d, 0x502, 0, 0); + root_span + .metrics + .insert("_sampling_priority_v1".to_string(), 0.0); + root_span.metrics.insert("_sample_rate".to_string(), 0.25); + root_span + .metrics + .insert("_dd.span_sampling.rule".to_string(), 1.0); + let mut child = test_span(0x501d, 0x503, 0x502, 1); + child.meta.insert("keep".to_string(), "me".to_string()); + + let mut chunk = test_chunk(vec![root_span, child], 0); + chunk.tags.insert("_dd.p.dm".to_string(), "-4".to_string()); + chunk + .tags + .insert("origin".to_string(), "synthetics".to_string()); + chunk.origin = "synthetics".to_string(); + + let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); + + assert_eq!(chunks.len(), 1); + let rescued = &chunks[0]; + assert_eq!(rescued.priority, 0, "priority preserved"); + assert_eq!(rescued.origin, "synthetics", "origin preserved"); + assert_eq!( + rescued.tags.get("_dd.p.dm").map(String::as_str), + Some("-4"), + "decision maker preserved" + ); + assert_eq!( + rescued.tags.get("origin").map(String::as_str), + Some("synthetics") + ); + assert_eq!( + rescued.spans[0] + .metrics + .get("_dd.span_sampling.rule") + .copied(), + Some(1.0), + "single-span sampling metrics preserved" + ); + assert_eq!( + rescued.spans[0] + .metrics + .get("_sampling_priority_v1") + .copied(), + Some(0.0), + "sampling priority metric preserved" + ); + assert_eq!( + errors_sr(&rescued.spans[0]), + Some(1.0), + "root stamped despite existing metrics" + ); + assert_eq!( + rescued.spans[1].meta.get("keep").map(String::as_str), + Some("me"), + "span metadata and order preserved" + ); + assert_eq!(rescued.spans.len(), 2, "no spans added or removed"); + } + + #[test] + fn test_probabilistic_decision_maker_chunk_gets_no_workaround() { + // A priority-0 chunk carrying the chunk-level probabilistic decision + // maker (`_dd.p.dm = "-9"`) is rescued like any other automatic-drop + // chunk: priority and decision maker are preserved untouched. The + // backend resolves such chunks to the `probabilistic` ingestion reason + // before checking `_dd.errors_sr`, so it may still drop them; SCL does + // not work around that limitation. + let config = rescue_config(always_keep_sampler_config()); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let mut chunk = errored_root_chunk(0x600d, 0); + chunk.tags.insert("_dd.p.dm".to_string(), "-9".to_string()); + + let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].priority, 0, "priority preserved, no promotion"); + assert_eq!( + chunks[0].tags.get("_dd.p.dm").map(String::as_str), + Some("-9"), + "probabilistic decision maker preserved, no rewrite" + ); + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + } + + #[test] + fn test_rescue_passes_raw_sample_rate_to_sampler() { + // The raw `_sample_rate` wire value is passed to the shared sampler, + // which sanitizes it. A value outside (0, 1] falls back to 1.0, so a + // bogus rate does not change the stamped rescue rate. + let sampler_config = rate_limited_sampler_config(10.0); + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + + let mut chunk = errored_root_chunk(0x77_00, 0); + chunk.spans[0] + .metrics + .insert("_sample_rate".to_string(), f64::NAN); + + let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); + + // Single signature, well under budget: sanitized rate 1.0 keeps it. + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + } + + #[tokio::test] + async fn test_stats_concentrator_observes_pre_rescue_chunks() { + let config = rescue_config(always_keep_sampler_config()); + let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); + let concentrator = super::StatsConcentratorHandle::new(stats_tx); + let processor = trace_processor::ServerlessTraceProcessor::new( + Some(concentrator), + sampler_for(&config.error_sampler), + ); + + let start = get_current_timestamp_nanos(); + let mut json_span = create_test_json_span(11, 222, 333, start, true); + json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); + let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); + let request = Request::builder() + .header("datadog-meta-tracer-version", "4.0.0") + .header("datadog-meta-lang", "nodejs") + .header("datadog-meta-lang-version", "v19.7.0") + .header("datadog-meta-lang-interpreter", "v8") + .header("content-length", "100") + .body(http_common::Body::from(bytes)) + .unwrap(); + + let res = processor + .process_traces( + Arc::new(config), + request, + mpsc::channel(1).0, + Arc::new(create_test_metadata()), + ) + .await; + assert!(res.is_ok()); + + // The concentrator must receive the chunk exactly as the tracer sent + // it: before rescue stamping, with no `_dd.errors_sr` anywhere. + let (chunk, _metadata) = match stats_rx.try_recv() { + Ok(crate::stats_concentrator_service::ConcentratorCommand::AddChunk( + chunk, + metadata, + )) => (*chunk, metadata), + Ok(_) => panic!("expected an AddChunk command"), + Err(err) => panic!("expected an AddChunk command, got {err}"), + }; + assert_eq!(chunk.priority, 0); + for span in &chunk.spans { + assert_eq!( + errors_sr(span), + None, + "stats must observe the pre-rescue chunk" + ); + } + } } diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 57de50ef..60883155 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -22,10 +22,11 @@ use datadog_trace_agent::{ }; use http_body_util::BodyExt; use hyper::StatusCode; -use serde_json::Value; +use libdd_trace_utils::test_utils::create_test_json_span; +use serde_json::{Value, json}; use serial_test::serial; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, UNIX_EPOCH}; #[cfg(all(windows, feature = "windows-pipes"))] use common::helpers::send_named_pipe_request; @@ -98,9 +99,10 @@ pub fn create_mini_agent_with_real_flushers( let aggregator = Arc::new(tokio::sync::Mutex::new(TraceAggregator::default())); let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new(Some( - stats_concentrator_handle.clone(), - ))), + trace_processor: Arc::new(ServerlessTraceProcessor::new( + Some(stats_concentrator_handle.clone()), + ServerlessTraceProcessor::new_error_sampler(&config), + )), trace_flusher: Arc::new(ServerlessTraceFlusher::new( aggregator.clone(), config.clone(), @@ -294,7 +296,10 @@ async fn test_mini_agent_tcp_handles_requests() { let test_port = config.dd_apm_receiver_port; let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new(None)), + trace_processor: Arc::new(ServerlessTraceProcessor::new( + None, + ServerlessTraceProcessor::new_error_sampler(&config), + )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -408,7 +413,10 @@ async fn test_mini_agent_named_pipe_handles_requests() { let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new(None)), + trace_processor: Arc::new(ServerlessTraceProcessor::new( + None, + ServerlessTraceProcessor::new_error_sampler(&config), + )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -578,7 +586,10 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new(None)), + trace_processor: Arc::new(ServerlessTraceProcessor::new( + None, + ServerlessTraceProcessor::new_error_sampler(&config), + )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -1278,3 +1289,201 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { let _ = agent_handle.await; verify_stats_request(&mock_intake).await; } + +/// Builds a mixed four-trace payload for the error rescue wire test: an +/// eligible errored automatic-drop chunk, a healthy automatic-drop chunk, an +/// explicit user drop, and a positive-priority chunk. +fn create_error_rescue_test_payload() -> Vec { + let start = UNIX_EPOCH.elapsed().unwrap().as_nanos() as i64; + let span = |trace_id: u64, name: &str, error: i64, priority: f64| { + let mut span = create_test_json_span(trace_id, trace_id + 1, 0, start, false); + span["name"] = json!(name); + span["error"] = json!(error); + span["metrics"]["_sampling_priority_v1"] = json!(priority); + span + }; + let traces = vec![ + vec![span(700, "rescued_error_p0", 1, 0.0)], + vec![span(701, "healthy_p0", 0, 0.0)], + vec![span(702, "user_drop_p0", 1, -1.0)], + vec![span(703, "priority_1", 1, 1.0)], + ]; + rmp_serde::to_vec(&traces).expect("Failed to serialize error rescue test trace") +} + +/// End-to-end error rescue verification against the outbound wire payload and +/// agent-computed stats: an eligible errored P0 chunk is rescued (positive +/// `_dd.errors_sr` on the root, priority still 0) while all other chunks are +/// forwarded unchanged, and every submitted span still contributes to stats. +/// This proves the wire contract only; backend retention of rescued chunks is +/// established by the backend implementation, not by this fake intake. +#[cfg(test)] +#[tokio::test] +#[serial] +async fn test_error_rescue_outbound_payload_and_stats() { + use libdd_trace_protobuf::pb::AgentPayload; + use prost::Message as _; + + let mock_server: MockServer = MockServer::start().await; + tokio::time::sleep(Duration::from_millis(50)).await; + + let mut config = create_tcp_test_config(8137); // use different port to avoid race condition with other tests + configure_mock_endpoints(&mut config, &mock_server.url()); + config.agent_stats_computation_enabled = true; + // Deterministic rescue settings: every eligible chunk is kept at 1.0. + config.error_sampler = datadog_agent_trace_sampler::ErrorSamplerConfig { + mode: datadog_agent_trace_sampler::ErrorSamplerMode::AlwaysKeep, + target_tps: 1.0, + extra_sample_rate: 1.0, + }; + let config = Arc::new(config); + let test_port = config.dd_apm_receiver_port; + + let (mini_agent, stats_concentrator_service_handle) = + create_mini_agent_with_real_flushers(config); + + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + let agent_handle = tokio::spawn(async move { + let _ = mini_agent + .start_mini_agent(shutdown_rx, Some(stats_concentrator_service_handle)) + .await; + }); + + let mut server_ready = false; + for _ in 0..20 { + tokio::time::sleep(Duration::from_millis(50)).await; + if let Ok(response) = send_tcp_request(test_port, "/info", "GET", None, &[]).await + && response.status().is_success() + { + server_ready = true; + break; + } + } + assert!( + server_ready, + "Mini agent server failed to start within timeout" + ); + + let trace_response = send_tcp_request( + test_port, + "/v0.4/traces", + "POST", + Some(create_error_rescue_test_payload()), + &[], + ) + .await + .expect("Failed to send /v0.4/traces request"); + assert_eq!(trace_response.status(), StatusCode::OK); + + verify_trace_request(&mock_server).await; + + // Decode the actual outbound protobuf payload(s) and locate each submitted + // chunk by its root span name. + let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); + let mut chunks_by_root_name: std::collections::HashMap< + String, + libdd_trace_protobuf::pb::TraceChunk, + > = std::collections::HashMap::new(); + for req in &trace_reqs { + let agent_payload = AgentPayload::decode(&req.body[..]) + .expect("Failed to decode outbound AgentPayload protobuf"); + for tracer_payload in agent_payload.tracer_payloads { + for chunk in tracer_payload.chunks { + let root_name = chunk + .spans + .iter() + .find(|s| s.parent_id == 0) + .map(|s| s.name.clone()) + .expect("chunk has a root span"); + chunks_by_root_name.insert(root_name, chunk); + } + } + } + + let names = [ + "rescued_error_p0", + "healthy_p0", + "user_drop_p0", + "priority_1", + ]; + for name in names { + assert!( + chunks_by_root_name.contains_key(name), + "expected chunk {name} to be forwarded to the backend, got: {:?}", + chunks_by_root_name.keys().collect::>() + ); + } + + let errors_sr_of = |name: &str| -> Option { + chunks_by_root_name[name] + .spans + .iter() + .find(|s| s.parent_id == 0) + .and_then(|root| root.metrics.get("_dd.errors_sr").copied()) + }; + + // The eligible errored automatic-drop chunk is rescued: positive + // `_dd.errors_sr` on the root, chunk priority still 0. + assert_eq!( + errors_sr_of("rescued_error_p0"), + Some(1.0), + "rescued root must carry a positive _dd.errors_sr" + ); + assert_eq!( + chunks_by_root_name["rescued_error_p0"].priority, 0, + "rescued chunk priority must remain 0, no promotion" + ); + + // Non-candidate chunks receive no rescue metric and keep their priority. + assert_eq!(errors_sr_of("healthy_p0"), None, "healthy P0 not rescued"); + assert_eq!(chunks_by_root_name["healthy_p0"].priority, 0); + assert_eq!( + errors_sr_of("user_drop_p0"), + None, + "explicit user drop never rescued" + ); + assert_eq!(chunks_by_root_name["user_drop_p0"].priority, -1); + assert_eq!( + errors_sr_of("priority_1"), + None, + "positive priority not rescued" + ); + assert_eq!(chunks_by_root_name["priority_1"].priority, 1); + + // Wait for the stats flush, then assert agent-computed stats include all + // four submitted spans, independent of rescue decisions. + tokio::time::sleep(FLUSH_WAIT_DURATION).await; + let _ = shutdown_tx.send(true); + let _ = agent_handle.await; + + let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); + assert!( + !stats_reqs.is_empty(), + "Expected at least one stats request" + ); + + let all_groups: Vec<_> = stats_reqs + .iter() + .map(|req| decode_stats_payload(&req.body)) + .flat_map(|payload| { + payload + .stats + .into_iter() + .flat_map(|csp| csp.stats.into_iter()) + .flat_map(|bucket| bucket.stats.into_iter()) + }) + .collect(); + + for name in names { + let group = all_groups.iter().find(|g| g.name == name); + assert!( + group.is_some(), + "expected span {name} to contribute to agent-computed stats, got: {:?}", + all_groups.iter().map(|g| &g.name).collect::>() + ); + assert!( + group.unwrap().hits > 0, + "expected span {name} to have a positive hit count in stats" + ); + } +} From 711b81254f095aaf8d89318e5bd1f7db4daa5741 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 14:46:44 -0400 Subject: [PATCH 02/49] Exercise error rescue in stats regression tests --- crates/datadog-trace-agent/src/trace_processor.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index caabc916..89c57f48 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1362,6 +1362,7 @@ mod tests { let start = get_current_timestamp_nanos(); let mut json_span = create_test_json_span(11, 222, 333, start, true); // Root span with an error and an automatic-drop priority: eligible. + json_span["error"] = serde_json::json!(1); json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); let request = Request::builder() @@ -1769,6 +1770,7 @@ mod tests { let start = get_current_timestamp_nanos(); let mut json_span = create_test_json_span(11, 222, 333, start, true); + json_span["error"] = serde_json::json!(1); json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); let request = Request::builder() From 0816af3d3577f41916e2c1d2203ff9573a428743 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 18 Sep 2026 15:07:37 -0400 Subject: [PATCH 03/49] Read sampler clock while holding the sampler lock --- .../src/trace_processor.rs | 53 ++++++++++++++----- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 89c57f48..0ce3409b 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -177,30 +177,58 @@ impl ServerlessTraceProcessor { /// Chunks are never removed or reordered: unrescued chunks stay in the /// payload and the backend discards them. /// - /// `now_unix_secs` is passed in explicitly so tests can exercise the - /// rolling window without sleeping. - fn apply_error_rescue( + fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { + let mut sampler = self.lock_sampler(); + // Skip all view construction when the sampler is disabled by config + // (target_tps <= 0): nothing can be rescued. + if sampler.is_disabled() { + return; + } + // Read the clock while holding the sampler lock so that concurrent + // requests deliver timestamps in lock-acquisition order, which is + // monotonic. Reading the clock before acquiring the lock could hand the + // sampler an older time after a newer one, moving the rolling window + // backwards and clearing counts from the current bucket, which would + // undercount TPS and rescue too many error chunks. + let now_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_secs() as i64); + self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); + } + + /// Test-only variant that injects a synthetic timestamp so tests can + /// exercise the rolling window without sleeping. + #[cfg(test)] + fn apply_error_rescue_at( &self, payload: &mut TracerPayloadCollection, config: &Config, now_unix_secs: i64, ) { - let TracerPayloadCollection::V07(tracer_payloads) = payload else { - return; - }; let mut sampler = self.lock_sampler(); - // Skip all view construction when the sampler is disabled by config - // (target_tps <= 0): nothing can be rescued. if sampler.is_disabled() { return; } + self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); + } + + fn rescue_with_sampler( + &self, + payload: &mut TracerPayloadCollection, + config: &Config, + now_unix_secs: i64, + sampler: &mut ErrorsSampler, + ) { + let TracerPayloadCollection::V07(tracer_payloads) = payload else { + return; + }; for tracer_payload in tracer_payloads.iter_mut() { // The sampler keys its per-signature rate limits on the env the // tracer reported for this payload, falling back to the agent's // configured env, consistent with how stats are flushed. let env: &str = resolve_payload_env(&tracer_payload.env, &config.env); for chunk in tracer_payload.chunks.iter_mut() { - sample_and_stamp(&mut sampler, chunk, env, now_unix_secs); + sample_and_stamp(sampler, chunk, env, now_unix_secs); } } } @@ -438,10 +466,7 @@ impl TraceProcessor for ServerlessTraceProcessor { // the newly inserted metric is included in the recomputed outbound size. It is // gated on agent stats computation: without it, P0 chunks are not expected here. if config.agent_stats_computation_enabled { - let now_unix_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0, |d| d.as_secs() as i64); - self.apply_error_rescue(&mut payload, &config, now_unix_secs); + self.apply_error_rescue(&mut payload, &config); } let pieces: Vec<(TracerPayloadCollection, usize)> = match payload { @@ -1114,7 +1139,7 @@ mod tests { chunks, ..Default::default() }]); - processor.apply_error_rescue(&mut payload, config, now_unix_secs); + processor.apply_error_rescue_at(&mut payload, config, now_unix_secs); match payload { TracerPayloadCollection::V07(mut payloads) => payloads.remove(0).chunks, _ => unreachable!(), From 5f5b4833e30702558826f5db6a0c8b80de24711b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 13:08:31 -0400 Subject: [PATCH 04/49] Clamp sampler timestamps against backward clock steps --- .../src/trace_processor.rs | 44 ++++++++++++++++--- 1 file changed, 39 insertions(+), 5 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 0ce3409b..21e84588 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; +use std::sync::atomic::{AtomicI64, Ordering}; use async_trait::async_trait; use datadog_agent_trace_sampler::{ErrorsSampler, SampleDecision, SpanView, TraceView}; @@ -137,6 +138,10 @@ pub struct ServerlessTraceProcessor { /// connection) share a single sampler state and TPS budget for the whole /// process lifetime. error_sampler: Arc>, + /// Last timestamp handed to the sampler, shared across clones. Clamps the + /// clock so a backward wall-clock step cannot move sampler bucket IDs + /// backwards. Updated only under the sampler lock. + last_sampler_timestamp: Arc, } impl ServerlessTraceProcessor { @@ -149,6 +154,7 @@ impl ServerlessTraceProcessor { stats_concentrator, enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), error_sampler, + last_sampler_timestamp: Arc::new(AtomicI64::new(i64::MIN)), } } @@ -185,14 +191,13 @@ impl ServerlessTraceProcessor { return; } // Read the clock while holding the sampler lock so that concurrent - // requests deliver timestamps in lock-acquisition order, which is - // monotonic. Reading the clock before acquiring the lock could hand the - // sampler an older time after a newer one, moving the rolling window - // backwards and clearing counts from the current bucket, which would - // undercount TPS and rescue too many error chunks. + // requests deliver timestamps in lock-acquisition order, and clamp so a + // backward wall-clock step cannot move the sampler's rolling window + // backwards, which would undercount TPS and rescue too many chunks. let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs() as i64); + let now_unix_secs = self.clamp_sampler_timestamp(now_unix_secs); self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); } @@ -212,6 +217,17 @@ impl ServerlessTraceProcessor { self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); } + /// Clamps the timestamp so it never moves backwards relative to the last + /// one handed to the shared sampler, and records it as the new floor. Must + /// be called while holding the sampler lock, alongside the clock read, so + /// concurrent requests cannot interleave a read with the clamp. + fn clamp_sampler_timestamp(&self, now_unix_secs: i64) -> i64 { + let previous = self + .last_sampler_timestamp + .fetch_max(now_unix_secs, Ordering::Relaxed); + now_unix_secs.max(previous) + } + fn rescue_with_sampler( &self, payload: &mut TracerPayloadCollection, @@ -1642,6 +1658,24 @@ mod tests { ); } + #[test] + fn test_sampler_timestamp_clamp_never_moves_backwards() { + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&rate_limited_sampler_config(1.0)), + ); + // A clone shares the clamp floor with the original. + let clone = processor.clone(); + + assert_eq!(processor.clamp_sampler_timestamp(100), 100); + // Backward timestamps, through either instance, are clamped to the + // last value seen; forward ones update the floor. + assert_eq!(clone.clamp_sampler_timestamp(50), 100); + assert_eq!(processor.clamp_sampler_timestamp(100), 100); + assert_eq!(clone.clamp_sampler_timestamp(200), 200); + assert_eq!(processor.clamp_sampler_timestamp(199), 200); + } + #[test] fn test_rescue_uses_payload_env_with_config_fallback() { assert_eq!( From c71ef792da3354afa35a7049bd0b464e9f267b77 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 13:12:17 -0400 Subject: [PATCH 05/49] Match test oracle env with processor fallback env --- crates/datadog-trace-agent/src/trace_processor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 21e84588..acd52461 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1548,7 +1548,7 @@ mod tests { error_type: s.meta.get("error.type").map(String::as_str), }) .collect(); - let trace = equivalent_trace_view(chunk, "test-env", &views); + let trace = equivalent_trace_view(chunk, &config.env, &views); let expected = standalone.sample(RESCUE_NOW, &trace); let actual = match errors_sr(root(chunk)) { Some(errors_sr) => SampleDecision::Keep { errors_sr }, From 1296884f07ef0ad63931e3632d50d4116863eae6 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 13:30:40 -0400 Subject: [PATCH 06/49] Match clone-budget oracle env with processor fallback env --- crates/datadog-trace-agent/src/trace_processor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index acd52461..7f9f6ac5 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1633,7 +1633,7 @@ mod tests { }]; for (i, id) in ids.iter().enumerate() { let trace = TraceView { - env: "test-env", + env: &config.env, trace_id: *id, root_index: 0, root_global_sample_rate: 1.0, From 9f741d96b0fb3e565d7b730a2032652256722533 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 13:35:09 -0400 Subject: [PATCH 07/49] Document root-resolution fallback in sampler doc comment --- crates/datadog-trace-agent/src/trace_processor.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 7f9f6ac5..59b67af9 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -327,9 +327,11 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { /// and the no-priority sentinel (`i8::MIN`) are all left untouched. The chunk /// must contain at least one span with a non-zero error flag; HTTP status or /// error metadata alone does not qualify. The root is resolved with -/// `get_root_span_index`; empty or rootless chunks are left unchanged. Chunks -/// are never removed here: on a Drop decision the unrescued chunk is forwarded -/// as-is and the backend's ordinary P0 drop handles it. +/// `get_root_span_index`: empty chunks are left unchanged, and non-empty +/// chunks always resolve a root, falling back to the last span when no span +/// has `parent_id 0` (e.g. a cyclic chunk). Chunks are never removed here: on +/// a Drop decision the unrescued chunk is forwarded as-is and the backend's +/// ordinary P0 drop handles it. fn sample_and_stamp( sampler: &mut ErrorsSampler, chunk: &mut pb::TraceChunk, From 50a8e476ccee1ebd6f3c7bcdc2c77b68bdfd437b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 14:18:58 -0400 Subject: [PATCH 08/49] Share the payload-env fallback rule between stats and error rescue Extracts the "prefer tracer payload env, fall back to agent config env" rule into one function reused by stats flushing and error-rescue sampling, and removes duplicated sampler lock/disabled-check and SpanView-construction code in the error sampler and its tests. --- .../src/stats_concentrator_service.rs | 11 +++-- .../src/trace_processor.rs | 47 ++++++------------- 2 files changed, 21 insertions(+), 37 deletions(-) diff --git a/crates/datadog-trace-agent/src/stats_concentrator_service.rs b/crates/datadog-trace-agent/src/stats_concentrator_service.rs index 01d5164f..46897746 100644 --- a/crates/datadog-trace-agent/src/stats_concentrator_service.rs +++ b/crates/datadog-trace-agent/src/stats_concentrator_service.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use tokio::sync::{mpsc, oneshot}; use crate::config::Config; +use crate::trace_processor::resolve_payload_env; use libdd_library_config::tracer_metadata::TracerMetadata; use libdd_trace_protobuf::pb::{ClientStatsPayload, TraceChunk}; use libdd_trace_stats::span_concentrator::{CardinalityLimitConfig, SpanConcentrator}; @@ -191,11 +192,11 @@ impl StatsConcentratorService { // Do not set hostname so the trace stats backend can aggregate stats properly hostname: String::new(), // Prefer env from the tracer payload, fall back to agent config - env: metadata - .service_env - .clone() - .filter(|s| !s.is_empty()) - .unwrap_or_else(|| self.config.env.clone()), + env: resolve_payload_env( + metadata.service_env.as_deref().unwrap_or(""), + &self.config.env, + ) + .to_string(), version: metadata.service_version.clone().unwrap_or_default(), lang: metadata.tracer_language.clone(), tracer_version: metadata.tracer_version.clone(), diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 59b67af9..0be3af51 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -184,12 +184,6 @@ impl ServerlessTraceProcessor { /// payload and the backend discards them. /// fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { - let mut sampler = self.lock_sampler(); - // Skip all view construction when the sampler is disabled by config - // (target_tps <= 0): nothing can be rescued. - if sampler.is_disabled() { - return; - } // Read the clock while holding the sampler lock so that concurrent // requests deliver timestamps in lock-acquisition order, and clamp so a // backward wall-clock step cannot move the sampler's rolling window @@ -198,12 +192,14 @@ impl ServerlessTraceProcessor { .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs() as i64); let now_unix_secs = self.clamp_sampler_timestamp(now_unix_secs); - self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); + self.apply_error_rescue_at(payload, config, now_unix_secs); } - /// Test-only variant that injects a synthetic timestamp so tests can - /// exercise the rolling window without sleeping. - #[cfg(test)] + /// Applies the rescue pass for an explicit timestamp, taking the sampler + /// lock and skipping all view construction when the sampler is disabled + /// by config (target_tps <= 0). Used by `apply_error_rescue` with the + /// clamped wall clock, and directly by tests with a synthetic timestamp + /// so they can exercise the rolling window without sleeping. fn apply_error_rescue_at( &self, payload: &mut TracerPayloadCollection, @@ -294,15 +290,15 @@ impl ServerlessTraceProcessor { } } -/// Chooses the env used to key the error sampler's per-signature rates: the -/// tracer payload's own env when nonempty, otherwise the agent's configured -/// env. This matches how stats prefer the payload env and fall back to the -/// agent config env. -fn resolve_payload_env<'a>(tracer_payload_env: &'a str, config_env: &'a str) -> &'a str { - if tracer_payload_env.is_empty() { +/// Chooses the env used to key per-signature state: the payload-reported env +/// when nonempty, otherwise the agent's configured env. Shared by stats +/// flushing and error-rescue sampling so the fallback rule can't drift +/// between the two. +pub(crate) fn resolve_payload_env<'a>(payload_env: &'a str, config_env: &'a str) -> &'a str { + if payload_env.is_empty() { config_env } else { - tracer_payload_env + payload_env } } @@ -637,9 +633,7 @@ mod tests { } fn default_error_sampler() -> Arc> { - Arc::new(std::sync::Mutex::new(ErrorsSampler::new( - ErrorSamplerConfig::default(), - ))) + sampler_for(&ErrorSamplerConfig::default()) } fn create_test_metadata() -> MiniAgentMetadata { @@ -1538,18 +1532,7 @@ mod tests { // sample rate, views) against direct crate usage. let mut standalone = ErrorsSampler::new(sampler_config); for (chunk, id) in chunks.iter().zip(&ids) { - let views: Vec = chunk - .spans - .iter() - .map(|s| SpanView { - service: &s.service, - name: &s.name, - resource: &s.resource, - error: s.error != 0, - http_status_code: s.meta.get("http.status_code").map(String::as_str), - error_type: s.meta.get("error.type").map(String::as_str), - }) - .collect(); + let views: Vec = chunk.spans.iter().map(super::span_view).collect(); let trace = equivalent_trace_view(chunk, &config.env, &views); let expected = standalone.sample(RESCUE_NOW, &trace); let actual = match errors_sr(root(chunk)) { From 54925ad7fe39588ae1ff6a9fc38c9eb82f9fe82c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 24 Sep 2026 14:40:53 -0400 Subject: [PATCH 09/49] Fix error-rescue clock read to happen under the sampler lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The clock read and timestamp clamp were happening before the sampler lock was acquired, letting concurrent requests deliver timestamps out of lock-acquisition order and undercount TPS in the rolling window. 🤖 --- .../datadog-trace-agent/src/trace_processor.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 0be3af51..1a2ded35 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -184,6 +184,12 @@ impl ServerlessTraceProcessor { /// payload and the backend discards them. /// fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { + let mut sampler = self.lock_sampler(); + // Skip all view construction when the sampler is disabled by config + // (target_tps <= 0): nothing can be rescued. + if sampler.is_disabled() { + return; + } // Read the clock while holding the sampler lock so that concurrent // requests deliver timestamps in lock-acquisition order, and clamp so a // backward wall-clock step cannot move the sampler's rolling window @@ -192,14 +198,12 @@ impl ServerlessTraceProcessor { .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs() as i64); let now_unix_secs = self.clamp_sampler_timestamp(now_unix_secs); - self.apply_error_rescue_at(payload, config, now_unix_secs); + self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); } - /// Applies the rescue pass for an explicit timestamp, taking the sampler - /// lock and skipping all view construction when the sampler is disabled - /// by config (target_tps <= 0). Used by `apply_error_rescue` with the - /// clamped wall clock, and directly by tests with a synthetic timestamp - /// so they can exercise the rolling window without sleeping. + /// Test-only variant that injects a synthetic timestamp so tests can + /// exercise the rolling window without sleeping. + #[cfg(test)] fn apply_error_rescue_at( &self, payload: &mut TracerPayloadCollection, From bf785df757bd47ed0ba1ad931965b14aa9d6d2ab Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 16:57:09 -0400 Subject: [PATCH 10/49] Simplify error rescue tests Extract shared fixtures for rescue processor setup, errored P0 requests, and keep counts; replace redundant assertions and fully-qualified paths. No behavior change. --- crates/datadog-trace-agent/src/config.rs | 8 +- .../src/trace_processor.rs | 228 ++++++------------ .../tests/integration_test.rs | 30 +-- 3 files changed, 100 insertions(+), 166 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index a9c9d2bc..d5972b58 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -63,6 +63,7 @@ fn parse_error_sampler_config() -> ErrorSamplerConfig { config } + const DEFAULT_DOGSTATSD_PORT: u16 = 8125; const DSM_PIPELINE_STATS_ROUTE: &str = "/api/v0.1/pipeline_stats"; @@ -1015,10 +1016,9 @@ mod tests { ], || { let config = config::Config::new().unwrap(); - assert!( - matches!(config.error_sampler.mode, m if std::mem::discriminant(&m) == std::mem::discriminant(&expected)), - "mode {raw:?} should parse to {expected:?}, got {:?}", - config.error_sampler.mode + assert_eq!( + config.error_sampler.mode, expected, + "mode {raw:?} should parse to {expected:?}" ); }, ); diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 1a2ded35..43c55eb4 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -182,7 +182,6 @@ impl ServerlessTraceProcessor { /// span; the backend keeps such chunks without any priority promotion. /// Chunks are never removed or reordered: unrescued chunks stay in the /// payload and the backend discards them. - /// fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { let mut sampler = self.lock_sampler(); // Skip all view construction when the sampler is disabled by config @@ -242,7 +241,7 @@ impl ServerlessTraceProcessor { // The sampler keys its per-signature rate limits on the env the // tracer reported for this payload, falling back to the agent's // configured env, consistent with how stats are flushed. - let env: &str = resolve_payload_env(&tracer_payload.env, &config.env); + let env = resolve_payload_env(&tracer_payload.env, &config.env); for chunk in tracer_payload.chunks.iter_mut() { sample_and_stamp(sampler, chunk, env, now_unix_secs); } @@ -1145,6 +1144,19 @@ mod tests { Arc::new(std::sync::Mutex::new(ErrorsSampler::new(*config))) } + /// Builds the standard rescue test fixture: a config with the given + /// sampler settings and a processor whose shared sampler uses them. + fn rescue_setup( + sampler_config: ErrorSamplerConfig, + ) -> (Config, trace_processor::ServerlessTraceProcessor) { + let config = rescue_config(sampler_config); + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + sampler_for(&config.error_sampler), + ); + (config, processor) + } + fn run_rescue( processor: &trace_processor::ServerlessTraceProcessor, config: &Config, @@ -1190,13 +1202,34 @@ mod tests { } } + fn rescued_count(chunks: &[pb::TraceChunk]) -> usize { + chunks + .iter() + .filter(|c| errors_sr(root(c)).is_some()) + .count() + } + + /// Builds a `/v0.4/traces` request carrying a single root span with an + /// error and an automatic-drop priority: eligible for rescue. + fn errored_p0_request() -> http_common::HttpRequest { + let start = get_current_timestamp_nanos(); + let mut json_span = create_test_json_span(11, 222, 333, start, true); + json_span["error"] = serde_json::json!(1); + json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); + let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); + Request::builder() + .header("datadog-meta-tracer-version", "4.0.0") + .header("datadog-meta-lang", "nodejs") + .header("datadog-meta-lang-version", "v19.7.0") + .header("datadog-meta-lang-interpreter", "v8") + .header("content-length", "100") + .body(http_common::Body::from(bytes)) + .unwrap() + } + #[test] fn test_rescue_stamps_errored_p0_root_and_keeps_priority() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); let chunks = run_rescue( &processor, @@ -1212,11 +1245,7 @@ mod tests { #[test] fn test_rescue_stamps_actual_root_when_error_is_on_child() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); // Healthy root first, errored child second: only the root may be stamped. let spans = vec![ @@ -1236,11 +1265,7 @@ mod tests { #[test] fn test_rescue_uses_resolved_root_not_first_span() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); // Root (parent_id 0) appears last; the first span is a child. let spans = vec![ @@ -1268,18 +1293,15 @@ mod tests { // it would not be stamped with 1.0. Correct behavior: only errored // chunks reach the sampler, so the errored chunk is the first count and // is kept at rate 1.0. - let config = rescue_config(rate_limited_sampler_config(1.0)); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); - let mut chunks = Vec::new(); - for i in 0..9_u64 { - chunks.push(errored_root_chunk(0xa000 + i, 0)); - // Strip the error flag: healthy P0 chunk with a distinct signature. - chunks.last_mut().unwrap().spans[0].error = 0; - } + // Healthy P0 chunks with distinct trace IDs, then one errored chunk. + let mut chunks: Vec = (0..9_u64) + .map(|i| { + let trace_id = 0xa000 + i; + test_chunk(vec![test_span(trace_id, trace_id + 1, 0, 0)], 0) + }) + .collect(); chunks.push(errored_root_chunk(0xb000, 0)); let chunks = run_rescue(&processor, &config, chunks, RESCUE_NOW); @@ -1296,11 +1318,7 @@ mod tests { #[test] fn test_http_500_metadata_alone_is_not_an_error() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); let mut span = test_span(0xc0de, 0xc0de + 1, 0, 0); span.meta @@ -1321,11 +1339,7 @@ mod tests { #[test] fn test_non_automatic_drop_priorities_are_never_rescued() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); // -1 (explicit user drop), other negatives, positive priorities, and // the no-priority sentinel (i8::MIN) are all out of scope. @@ -1334,8 +1348,8 @@ mod tests { &processor, &config, priorities - .iter() - .map(|p| errored_root_chunk(0x1000_u64 + (*p).unsigned_abs() as u64, *p)) + .into_iter() + .map(|p| errored_root_chunk(0x1000 + u64::from(p.unsigned_abs()), p)) .collect(), RESCUE_NOW, ); @@ -1352,11 +1366,7 @@ mod tests { #[test] fn test_empty_chunk_does_not_panic_and_stays_unchanged() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); // An empty chunk cannot be scored: root resolution fails and it must be // left unchanged without panicking. For non-empty chunks, @@ -1390,30 +1400,15 @@ mod tests { #[tokio::test] async fn test_rescue_is_noop_when_agent_stats_disabled() { - let config = rescue_config(always_keep_sampler_config()); let config = Config { agent_stats_computation_enabled: false, - ..config + ..rescue_config(always_keep_sampler_config()) }; let (tx, mut rx): ( Sender, - tokio::sync::mpsc::Receiver, + Receiver, ) = mpsc::channel(1); - - let start = get_current_timestamp_nanos(); - let mut json_span = create_test_json_span(11, 222, 333, start, true); - // Root span with an error and an automatic-drop priority: eligible. - json_span["error"] = serde_json::json!(1); - json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); - let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); - let request = Request::builder() - .header("datadog-meta-tracer-version", "4.0.0") - .header("datadog-meta-lang", "nodejs") - .header("datadog-meta-lang-version", "v19.7.0") - .header("datadog-meta-lang-interpreter", "v8") - .header("content-length", "100") - .body(http_common::Body::from(bytes)) - .unwrap(); + let request = errored_p0_request(); let trace_processor = trace_processor::ServerlessTraceProcessor::new( None, @@ -1449,16 +1444,11 @@ mod tests { fn test_disabled_tps_disables_rescue_in_both_modes() { for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { for tps in [0.0_f64, -3.0] { - let sampler_config = ErrorSamplerConfig { + let (config, processor) = rescue_setup(ErrorSamplerConfig { mode, target_tps: tps, extra_sample_rate: 1.0, - }; - let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + }); let chunks = run_rescue( &processor, @@ -1479,11 +1469,7 @@ mod tests { #[test] fn test_always_keep_rescues_every_eligible_chunk() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); let chunks = run_rescue( &processor, @@ -1504,11 +1490,7 @@ mod tests { #[test] fn test_rate_limited_under_sustained_load_keeps_and_rejects() { let sampler_config = rate_limited_sampler_config(10.0); - let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(sampler_config); // 100 chunks with the same signature in one 5-second bucket: the // default rate is 10 / (100 / 5) = 0.5, so a deterministic mix of keeps @@ -1519,10 +1501,7 @@ mod tests { let input: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); - let keeps = chunks - .iter() - .filter(|c| errors_sr(root(c)).is_some()) - .count(); + let keeps = rescued_count(&chunks); let drops = chunks.len() - keeps; assert!(keeps > 0, "expected some keeps, got none"); assert!(drops > 0, "expected some drops, got none"); @@ -1549,21 +1528,13 @@ mod tests { #[test] fn test_rate_limited_bucket_transitions_and_steady_state() { - let sampler_config = rate_limited_sampler_config(10.0); - let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); // First bucket: 100 distinct signatures push the default rate to 0.5. let ids: Vec = (0..100_u64).map(|i| 0x22_0000 + i).collect(); - let first: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); - let first_chunks = run_rescue(&processor, &config, first, RESCUE_NOW); - let first_keeps = first_chunks - .iter() - .filter(|c| errors_sr(root(c)).is_some()) - .count(); + let make_chunks = || ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); + let first_chunks = run_rescue(&processor, &config, make_chunks(), RESCUE_NOW); + let first_keeps = rescued_count(&first_chunks); assert!( first_keeps > 0 && first_keeps < 100, "mixed decisions in the first bucket" @@ -1572,13 +1543,9 @@ mod tests { // A full window later (the rolling window is 6 buckets of 5 seconds), // the same IDs in a fresh bucket still produce mixed decisions, and // every chunk is still forwarded. - let second: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); - let second_chunks = run_rescue(&processor, &config, second, RESCUE_NOW + 40); + let second_chunks = run_rescue(&processor, &config, make_chunks(), RESCUE_NOW + 40); assert_eq!(second_chunks.len(), 100); - let second_keeps = second_chunks - .iter() - .filter(|c| errors_sr(root(c)).is_some()) - .count(); + let second_keeps = rescued_count(&second_chunks); assert!( second_keeps > 0 && second_keeps < 100, "mixed decisions after window rotation" @@ -1588,11 +1555,7 @@ mod tests { #[test] fn test_processor_clones_share_one_budget() { let sampler_config = rate_limited_sampler_config(1.0); - let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(sampler_config); let processor_clone = processor.clone(); // Ten errored chunks with the same signature but distinct IDs: five @@ -1649,10 +1612,7 @@ mod tests { #[test] fn test_sampler_timestamp_clamp_never_moves_backwards() { - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&rate_limited_sampler_config(1.0)), - ); + let (_config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); // A clone shares the clamp floor with the original. let clone = processor.clone(); @@ -1666,7 +1626,7 @@ mod tests { } #[test] - fn test_rescue_uses_payload_env_with_config_fallback() { + fn test_resolve_payload_env_prefers_payload_env() { assert_eq!( super::resolve_payload_env("tracer-env", "agent-env"), "tracer-env", @@ -1686,11 +1646,7 @@ mod tests { #[test] fn test_rescue_preserves_chunk_metadata_and_span_order() { - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); let mut root_span = test_span(0x501d, 0x502, 0, 0); root_span @@ -1762,11 +1718,7 @@ mod tests { // backend resolves such chunks to the `probabilistic` ingestion reason // before checking `_dd.errors_sr`, so it may still drop them; SCL does // not work around that limitation. - let config = rescue_config(always_keep_sampler_config()); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); let mut chunk = errored_root_chunk(0x600d, 0); chunk.tags.insert("_dd.p.dm".to_string(), "-9".to_string()); @@ -1788,12 +1740,7 @@ mod tests { // The raw `_sample_rate` wire value is passed to the shared sampler, // which sanitizes it. A value outside (0, 1] falls back to 1.0, so a // bogus rate does not change the stamped rescue rate. - let sampler_config = rate_limited_sampler_config(10.0); - let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); let mut chunk = errored_root_chunk(0x77_00, 0); chunk.spans[0] @@ -1816,24 +1763,10 @@ mod tests { sampler_for(&config.error_sampler), ); - let start = get_current_timestamp_nanos(); - let mut json_span = create_test_json_span(11, 222, 333, start, true); - json_span["error"] = serde_json::json!(1); - json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); - let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); - let request = Request::builder() - .header("datadog-meta-tracer-version", "4.0.0") - .header("datadog-meta-lang", "nodejs") - .header("datadog-meta-lang-version", "v19.7.0") - .header("datadog-meta-lang-interpreter", "v8") - .header("content-length", "100") - .body(http_common::Body::from(bytes)) - .unwrap(); - let res = processor .process_traces( Arc::new(config), - request, + errored_p0_request(), mpsc::channel(1).0, Arc::new(create_test_metadata()), ) @@ -1842,11 +1775,10 @@ mod tests { // The concentrator must receive the chunk exactly as the tracer sent // it: before rescue stamping, with no `_dd.errors_sr` anywhere. - let (chunk, _metadata) = match stats_rx.try_recv() { - Ok(crate::stats_concentrator_service::ConcentratorCommand::AddChunk( - chunk, - metadata, - )) => (*chunk, metadata), + let chunk = match stats_rx.try_recv() { + Ok(crate::stats_concentrator_service::ConcentratorCommand::AddChunk(chunk, _)) => { + *chunk + } Ok(_) => panic!("expected an AddChunk command"), Err(err) => panic!("expected an AddChunk command, got {err}"), }; diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 60883155..4e02ad73 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1321,8 +1321,10 @@ fn create_error_rescue_test_payload() -> Vec { #[tokio::test] #[serial] async fn test_error_rescue_outbound_payload_and_stats() { - use libdd_trace_protobuf::pb::AgentPayload; + use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; + use libdd_trace_protobuf::pb::{AgentPayload, TraceChunk}; use prost::Message as _; + use std::collections::HashMap; let mock_server: MockServer = MockServer::start().await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -1331,8 +1333,8 @@ async fn test_error_rescue_outbound_payload_and_stats() { configure_mock_endpoints(&mut config, &mock_server.url()); config.agent_stats_computation_enabled = true; // Deterministic rescue settings: every eligible chunk is kept at 1.0. - config.error_sampler = datadog_agent_trace_sampler::ErrorSamplerConfig { - mode: datadog_agent_trace_sampler::ErrorSamplerMode::AlwaysKeep, + config.error_sampler = ErrorSamplerConfig { + mode: ErrorSamplerMode::AlwaysKeep, target_tps: 1.0, extra_sample_rate: 1.0, }; @@ -1380,10 +1382,7 @@ async fn test_error_rescue_outbound_payload_and_stats() { // Decode the actual outbound protobuf payload(s) and locate each submitted // chunk by its root span name. let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); - let mut chunks_by_root_name: std::collections::HashMap< - String, - libdd_trace_protobuf::pb::TraceChunk, - > = std::collections::HashMap::new(); + let mut chunks_by_root_name: HashMap = HashMap::new(); for req in &trace_reqs { let agent_payload = AgentPayload::decode(&req.body[..]) .expect("Failed to decode outbound AgentPayload protobuf"); @@ -1475,14 +1474,17 @@ async fn test_error_rescue_outbound_payload_and_stats() { .collect(); for name in names { - let group = all_groups.iter().find(|g| g.name == name); - assert!( - group.is_some(), - "expected span {name} to contribute to agent-computed stats, got: {:?}", - all_groups.iter().map(|g| &g.name).collect::>() - ); + let group = all_groups + .iter() + .find(|g| g.name == name) + .unwrap_or_else(|| { + panic!( + "expected span {name} to contribute to agent-computed stats, got: {:?}", + all_groups.iter().map(|g| &g.name).collect::>() + ) + }); assert!( - group.unwrap().hits > 0, + group.hits > 0, "expected span {name} to have a positive hit count in stats" ); } From d9d694d26cf62384580a03e4b13b71c9dd0c9d11 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 19 Oct 2026 12:15:48 -0400 Subject: [PATCH 11/49] Add opt-in rescue for tracer-dropped errors --- Cargo.lock | 1 + crates/datadog-trace-agent/Cargo.toml | 1 + crates/datadog-trace-agent/src/config.rs | 137 ++++++ .../src/stats_processor.rs | 1 + .../src/trace_processor.rs | 405 +++++++++++++++--- .../tests/integration_test.rs | 177 +++++--- 6 files changed, 602 insertions(+), 120 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d28ba062..44f9c5b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -674,6 +674,7 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tracing", + "tracing-test", ] [[package]] diff --git a/crates/datadog-trace-agent/Cargo.toml b/crates/datadog-trace-agent/Cargo.toml index a9b73ded..295e07b8 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -44,6 +44,7 @@ prost = "0.14.1" [dev-dependencies] datadog-mock-intake = { path = "../datadog-mock-intake" } +tracing-test = { version = "0.2.5", default-features = false } rmp-serde = "1.1.1" serial_test = "2.0.0" duplicate = "2.0.1" diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index d5972b58..c061cedd 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -64,6 +64,15 @@ fn parse_error_sampler_config() -> ErrorSamplerConfig { config } +// The trace agent already constructs its own configuration. Adopting +// datadog-agent-config solely for this flag would introduce broader loader coupling. +fn error_rare_sample_tracer_drop_enabled(raw: &str) -> bool { + let delimiter = if raw.contains(',') { ',' } else { ' ' }; + raw.split(delimiter) + .map(str::trim) + .any(|entry| entry == "error_rare_sample_tracer_drop") +} + const DEFAULT_DOGSTATSD_PORT: u16 = 8125; const DSM_PIPELINE_STATS_ROUTE: &str = "/api/v0.1/pipeline_stats"; @@ -179,6 +188,8 @@ pub struct Config { /// and `DD_APM_ERROR_TPS`. `extra_sample_rate` is fixed at 1.0 and not /// configurable. See `parse_error_sampler_config` for the defaults. pub error_sampler: ErrorSamplerConfig, + /// Opt-in eligibility for tracer-dropped errors, independent of sampler enablement. + pub error_rare_sample_tracer_drop: bool, } impl Config { @@ -360,6 +371,8 @@ impl Config { .map(|val| val.to_lowercase() == "true") .unwrap_or(true), error_sampler: parse_error_sampler_config(), + error_rare_sample_tracer_drop: env::var("DD_APM_FEATURES") + .is_ok_and(|raw| error_rare_sample_tracer_drop_enabled(&raw)), }) } } @@ -969,6 +982,103 @@ mod tests { ); } + #[test] + fn test_apm_features_grammar() { + let feature = "error_rare_sample_tracer_drop"; + for (raw, expected) in [ + ("".to_string(), false), + (" ".to_string(), false), + ("other".to_string(), false), + (feature.to_string(), true), + (format!(" {feature} "), true), + (format!("other {feature}"), true), + (format!("other,{feature}"), true), + (format!(" other ,, {feature} "), true), + (format!("other,{feature} extra"), false), + (format!("other {feature},extra"), false), + (format!("other\t{feature}"), false), + (format!("\t{feature}\t"), true), + (feature.to_uppercase(), false), + (format!("prefix_{feature}"), false), + (format!("{feature}_suffix"), false), + ] { + assert_eq!( + super::error_rare_sample_tracer_drop_enabled(&raw), + expected, + "{raw:?}" + ); + } + } + + #[test] + #[serial] + fn test_tracer_drop_feature_does_not_enable_sampler_or_stats() { + for feature in [None, Some("error_rare_sample_tracer_drop"), Some("other")] { + for (tps, expected_tps) in [(None, 10.0), (Some("0"), 0.0), (Some("-1"), -1.0)] { + for mode in [None, Some("always_keep")] { + for stats in ["true", "false"] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("test")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("test")), + ("DD_APM_FEATURES", feature), + ("DD_APM_ERROR_TPS", tps), + ("DD_APM_ERROR_SAMPLER_MODE", mode), + ("DD_AGENT_STATS_COMPUTATION_ENABLED", Some(stats)), + ], + || { + let config = config::Config::new().unwrap(); + assert_eq!( + config.error_rare_sample_tracer_drop, + feature == Some("error_rare_sample_tracer_drop") + ); + assert_eq!(config.error_sampler.target_tps, expected_tps); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); + assert_eq!( + config.error_sampler.mode, + if mode.is_some() { + ErrorSamplerMode::AlwaysKeep + } else { + ErrorSamplerMode::RateLimited + } + ); + assert_eq!(config.agent_stats_computation_enabled, stats == "true"); + }, + ); + } + } + } + } + } + + #[cfg(unix)] + #[test] + #[serial] + fn test_unreadable_apm_features_is_disabled() { + use std::os::unix::ffi::OsStringExt; + temp_env::with_var( + "DD_APM_FEATURES", + Some(std::ffi::OsString::from_vec(vec![0xff])), + || { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("test")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("test")), + ("DD_APM_ERROR_TPS", None), + ("DD_APM_ERROR_SAMPLER_MODE", None), + ], + || { + assert!(!config::Config::new().unwrap().error_rare_sample_tracer_drop); + }, + ); + }, + ); + } + fn assert_rate_limited_default(error_sampler: &config::ErrorSamplerConfig) { assert!( matches!(error_sampler.mode, ErrorSamplerMode::RateLimited), @@ -987,6 +1097,9 @@ mod tests { ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_TPS", None), + ("DD_APM_ERROR_SAMPLER_MODE", None), ], || { let config = config::Config::new().unwrap(); @@ -1012,6 +1125,8 @@ mod tests { ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_TPS", None), ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), ], || { @@ -1035,6 +1150,8 @@ mod tests { ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_TPS", None), ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), ], || { @@ -1055,6 +1172,8 @@ mod tests { ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_SAMPLER_MODE", None), ("DD_APM_ERROR_TPS", Some(raw)), ], || { @@ -1070,6 +1189,7 @@ mod tests { #[test] #[serial] + #[tracing_test::traced_test] fn test_error_sampler_invalid_tps_falls_back_to_default() { for raw in ["not_a_number", "inf", "-inf", "NaN", "1e400", ""] { temp_env::with_vars( @@ -1078,14 +1198,30 @@ mod tests { ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), + ("DD_APM_ERROR_SAMPLER_MODE", None), ("DD_APM_ERROR_TPS", Some(raw)), ], || { let config = config::Config::new().unwrap(); assert_rate_limited_default(&config.error_sampler); + let _ = crate::trace_processor::ServerlessTraceProcessor::new_error_sampler( + &config, + ); }, ); } + assert!(!logs_contain("DD_APM_ERROR_TPS <= 0 disables error rescue")); + logs_assert(|lines| { + if lines + .iter() + .any(|line| line.contains("DD_APM_ERROR_TPS <= 0 disables error rescue")) + { + Err("fallback TPS must not warn that rescue is disabled".into()) + } else { + Ok(()) + } + }); } } @@ -1130,6 +1266,7 @@ pub mod test_helpers { additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled: true, error_sampler: ErrorSamplerConfig::default(), + error_rare_sample_tracer_drop: false, } } } diff --git a/crates/datadog-trace-agent/src/stats_processor.rs b/crates/datadog-trace-agent/src/stats_processor.rs index 104eaea3..6c2e1677 100644 --- a/crates/datadog-trace-agent/src/stats_processor.rs +++ b/crates/datadog-trace-agent/src/stats_processor.rs @@ -165,6 +165,7 @@ mod tests { additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled, error_sampler: datadog_agent_trace_sampler::ErrorSamplerConfig::default(), + error_rare_sample_tracer_drop: false, } } diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 43c55eb4..ed27a7f3 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -31,11 +31,12 @@ use crate::{ const TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY: &str = "_dd.tags.function"; -/// The root-span metric the backend uses to rescue traces the ordinary P0 drop -/// would discard: a positive `_dd.errors_sr` resolves the chunk's ingestion -/// reason to `error`, which is retained at low priority without any priority -/// promotion. +/// The root-span metric used for error rescue without priority promotion. +/// A positive value signals error sampling; backend retention also depends on +/// its ingestion-reason handling (for example, probabilistic decision makers). const ERRORS_SR_METRIC_KEY: &str = "_dd.errors_sr"; +const DECISION_MAKER_KEY: &str = "_dd.p.dm"; +const MANUAL_DECISION_MAKER: &str = "-4"; /// Rough upper bound on the protobuf framing overhead added when a V07 `TracerPayload` is /// wrapped in the outer `AgentPayload` envelope before being sent @@ -104,6 +105,19 @@ struct ChunkProcessor { impl TraceChunkProcessor for ChunkProcessor { fn process(&mut self, chunk: &mut pb::TraceChunk, root_span_index: usize) { + // Promote before local enrichment or any future sampled-out filtering. + // Existing chunk values, including empty strings, take precedence. + if !chunk.tags.contains_key(DECISION_MAKER_KEY) + && let Some(value) = chunk + .spans + .iter() + .find_map(|span| span.meta.get(DECISION_MAKER_KEY)) + { + chunk + .tags + .insert(DECISION_MAKER_KEY.to_string(), value.clone()); + } + // Clone app_name once instead of once per span let app_name = self.config.app_name.clone(); @@ -161,6 +175,11 @@ impl ServerlessTraceProcessor { /// Builds the error rescue sampler from parsed config settings. #[must_use] pub fn new_error_sampler(config: &Config) -> Arc> { + if config.error_rare_sample_tracer_drop && config.error_sampler.target_tps <= 0.0 { + warn!( + "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue" + ); + } Arc::new(Mutex::new(ErrorsSampler::new(config.error_sampler))) } @@ -174,14 +193,13 @@ impl ServerlessTraceProcessor { /// Applies the error rescue pass to a payload collection. /// - /// For every V07 chunk that would be dropped by the backend's ordinary P0 - /// drop (chunk priority is exactly 0, i.e. an automatic drop, not an - /// explicit user drop or the no-priority sentinel) and contains at least - /// one span with a non-zero error flag, consults the shared error sampler. + /// For every eligible V07 chunk containing a non-zero error flag, consults + /// the shared sampler. P0 is always eligible; the opt-in feature adds negative + /// priorities except the no-priority sentinel and exact manual drops. /// On a keep, stamps the sampler's `_dd.errors_sr` on the chunk's root - /// span; the backend keeps such chunks without any priority promotion. + /// span without priority promotion. Retention is decided by the backend. /// Chunks are never removed or reordered: unrescued chunks stay in the - /// payload and the backend discards them. + /// payload for the backend to handle. fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { let mut sampler = self.lock_sampler(); // Skip all view construction when the sampler is disabled by config @@ -243,7 +261,13 @@ impl ServerlessTraceProcessor { // configured env, consistent with how stats are flushed. let env = resolve_payload_env(&tracer_payload.env, &config.env); for chunk in tracer_payload.chunks.iter_mut() { - sample_and_stamp(sampler, chunk, env, now_unix_secs); + sample_and_stamp( + sampler, + chunk, + env, + now_unix_secs, + config.error_rare_sample_tracer_drop, + ); } } } @@ -321,23 +345,30 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { /// Consults the error sampler for one chunk and, on a keep, stamps the /// sampler's `_dd.errors_sr` value on the chunk's root span. /// -/// Only automatic-drop chunks are candidates: chunk priority must be exactly 0. -/// Explicit user drops (-1), other negative priorities, positive priorities, -/// and the no-priority sentinel (`i8::MIN`) are all left untouched. The chunk +/// P0 chunks are candidates. With the feature enabled, negative priorities also +/// qualify except the no-priority sentinel and priority -1 with chunk DM -4. +/// Only promoted chunk tags are consulted, never span metadata. The chunk /// must contain at least one span with a non-zero error flag; HTTP status or /// error metadata alone does not qualify. The root is resolved with /// `get_root_span_index`: empty chunks are left unchanged, and non-empty /// chunks always resolve a root, falling back to the last span when no span /// has `parent_id 0` (e.g. a cyclic chunk). Chunks are never removed here: on -/// a Drop decision the unrescued chunk is forwarded as-is and the backend's -/// ordinary P0 drop handles it. +/// a Drop decision the unrescued chunk is forwarded without rescue mutation. fn sample_and_stamp( sampler: &mut ErrorsSampler, chunk: &mut pb::TraceChunk, env: &str, now_unix_secs: i64, + tracer_drop_enabled: bool, ) { - if chunk.priority != 0 { + let eligible = chunk.priority == 0 + || (tracer_drop_enabled + && chunk.priority < 0 + && chunk.priority != i32::from(i8::MIN) + && !(chunk.priority == -1 + && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) + == Some(MANUAL_DECISION_MAKER))); + if !eligible { return; } // An error anywhere in the chunk makes it a rescue candidate, not just an @@ -479,9 +510,9 @@ impl TraceProcessor for ServerlessTraceProcessor { } // Error rescue runs after stats submission so the concentrator observes every - // submitted chunk exactly as the tracer sent it, and before payload splitting so + // submitted chunk before rescue stamping, and before payload splitting so // the newly inserted metric is included in the recomputed outbound size. It is - // gated on agent stats computation: without it, P0 chunks are not expected here. + // gated on agent stats computation, independently of tracer-drop eligibility. if config.agent_stats_computation_enabled { self.apply_error_rescue(&mut payload, &config); } @@ -632,6 +663,7 @@ mod tests { additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled: false, error_sampler: ErrorSamplerConfig::default(), + error_rare_sample_tracer_drop: false, } } @@ -1210,12 +1242,15 @@ mod tests { } /// Builds a `/v0.4/traces` request carrying a single root span with an - /// error and an automatic-drop priority: eligible for rescue. - fn errored_p0_request() -> http_common::HttpRequest { + /// error and the requested sampling priority and decision maker. + fn errored_request(priority: i32, dm: Option<&str>) -> http_common::HttpRequest { let start = get_current_timestamp_nanos(); let mut json_span = create_test_json_span(11, 222, 333, start, true); json_span["error"] = serde_json::json!(1); - json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(0.0); + json_span["metrics"]["_sampling_priority_v1"] = serde_json::json!(priority); + if let Some(dm) = dm { + json_span["meta"]["_dd.p.dm"] = serde_json::json!(dm); + } let bytes = rmp_serde::to_vec(&vec![vec![json_span]]).unwrap(); Request::builder() .header("datadog-meta-tracer-version", "4.0.0") @@ -1227,6 +1262,179 @@ mod tests { .unwrap() } + #[test] + fn test_decision_maker_promotion_is_unconditional_and_preserves_metadata() { + use libdd_trace_utils::tracer_payload::TraceChunkProcessor as _; + for feature in [false, true] { + for (chunk_dm, first, second, child_first, expected) in [ + (None, None, None, false, None), + (None, Some("-4"), None, false, Some("-4")), + (None, None, Some("-3"), false, Some("-3")), + (None, Some("-3"), Some("-4"), true, Some("-3")), + (None, Some("-4"), Some("-3"), false, Some("-4")), + (None, Some("-3"), Some("-4"), false, Some("-3")), + (Some("-9"), Some("-4"), Some("-3"), false, Some("-9")), + (Some(""), Some("-4"), None, false, Some("")), + (None, Some(""), Some("-4"), true, Some("")), + ] { + let mut config = create_test_config(); + config.error_rare_sample_tracer_drop = feature; + config.agent_stats_computation_enabled = false; + let mut spans = vec![test_span(1, 2, 0, 1), test_span(1, 3, 2, 0)]; + if child_first { + spans.swap(0, 1); + } + for (span, dm) in spans.iter_mut().zip([first, second]) { + span.meta.insert("keep".into(), "metadata".into()); + span.metrics.insert("_sampling_priority_v1".into(), -1.0); + if let Some(dm) = dm { + span.meta + .insert(super::DECISION_MAKER_KEY.into(), dm.into()); + } + } + let mut chunk = test_chunk(spans, -1); + chunk.origin = "synthetics".into(); + chunk.tags.insert("other".into(), "value".into()); + if let Some(dm) = chunk_dm { + chunk + .tags + .insert(super::DECISION_MAKER_KEY.into(), dm.into()); + } + let original = chunk.clone(); + let root_index = usize::from(child_first); + super::ChunkProcessor { + config: Arc::new(config), + mini_agent_metadata: Arc::new(create_test_metadata()), + } + .process(&mut chunk, root_index); + assert_eq!( + chunk + .tags + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + expected + ); + assert_eq!(chunk.tags.get("other"), original.tags.get("other")); + assert_eq!(chunk.priority, original.priority); + assert_eq!(chunk.origin, original.origin); + for (span, before) in chunk.spans.iter().zip(&original.spans) { + assert_eq!(span.span_id, before.span_id); + assert_eq!( + span.metrics.get("_sampling_priority_v1"), + before.metrics.get("_sampling_priority_v1") + ); + for (key, value) in &before.meta { + assert_eq!(span.meta.get(key), Some(value)); + } + } + } + } + } + + #[test] + fn test_tracer_drop_rescue_eligibility_matrix() { + for feature in [false, true] { + for priority in [0, -1, -2, -127, i32::from(i8::MIN), 1, 2] { + for dm in [ + None, + Some(""), + Some("-4"), + Some("-3"), + Some("-9"), + Some("-40"), + ] { + for error in [0, 1, -1] { + let (mut config, processor) = rescue_setup(always_keep_sampler_config()); + config.error_rare_sample_tracer_drop = feature; + let mut chunk = test_chunk( + vec![test_span(1, 2, 0, 0), test_span(1, 3, 2, error)], + priority, + ); + // Span DM deliberately disagrees: rescue must read only chunk tags. + chunk.spans[0] + .meta + .insert(super::DECISION_MAKER_KEY.into(), "-4".into()); + chunk.spans[0] + .meta + .insert("http.status_code".into(), "500".into()); + if let Some(dm) = dm { + chunk + .tags + .insert(super::DECISION_MAKER_KEY.into(), dm.into()); + } + let original = chunk.clone(); + let mut chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); + let expected = error != 0 + && (priority == 0 + || (feature + && priority < 0 + && priority != -128 + && !(priority == -1 && dm == Some("-4")))); + assert_eq!( + errors_sr(&chunks[0].spans[0]), + expected.then_some(1.0), + "feature={feature} priority={priority} dm={dm:?} error={error}" + ); + assert_eq!(errors_sr(&chunks[0].spans[1]), None); + chunks[0].spans[0] + .metrics + .remove(super::ERRORS_SR_METRIC_KEY); + assert_eq!(chunks[0], original); + } + } + } + } + } + + #[test] + #[tracing_test::traced_test] + fn test_disabled_sampler_initialization_warning() { + let warning = "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue"; + let mut expected_warnings = 0; + for feature in [false, true] { + for tps in [0.0, -1.0, 10.0] { + let mut config = rescue_config(always_keep_sampler_config()); + config.error_rare_sample_tracer_drop = feature; + config.error_sampler.target_tps = tps; + let processor = trace_processor::ServerlessTraceProcessor::new( + None, + trace_processor::ServerlessTraceProcessor::new_error_sampler(&config), + ); + expected_warnings += usize::from(feature && tps <= 0.0); + logs_assert(|lines| { + let count = lines.iter().filter(|line| line.contains(warning)).count(); + if count == expected_warnings { + Ok(()) + } else { + Err(format!( + "expected {expected_warnings} warnings, got {count}" + )) + } + }); + for _ in 0..3 { + run_rescue( + &processor.clone(), + &config, + vec![errored_root_chunk(1, -1)], + RESCUE_NOW, + ); + } + } + } + assert!(logs_contain(warning)); + logs_assert(|lines| { + let count = lines + .iter() + .filter(|line| line.contains("DD_APM_ERROR_TPS <= 0 disables error rescue")) + .count(); + if count == 2 { + Ok(()) + } else { + Err(format!("expected two initialization warnings, got {count}")) + } + }); + } + #[test] fn test_rescue_stamps_errored_p0_root_and_keeps_priority() { let (config, processor) = rescue_setup(always_keep_sampler_config()); @@ -1341,8 +1549,8 @@ mod tests { fn test_non_automatic_drop_priorities_are_never_rescued() { let (config, processor) = rescue_setup(always_keep_sampler_config()); - // -1 (explicit user drop), other negatives, positive priorities, and - // the no-priority sentinel (i8::MIN) are all out of scope. + // With the feature off, all negative and positive priorities, + // including the no-priority sentinel, are out of scope. let priorities = [-1_i32, -5, 1, 2, i8::MIN as i32]; let chunks = run_rescue( &processor, @@ -1364,6 +1572,26 @@ mod tests { } } + #[test] + fn test_tracer_drop_empty_and_fallback_root() { + let (mut config, processor) = rescue_setup(always_keep_sampler_config()); + config.error_rare_sample_tracer_drop = true; + let chunks = run_rescue( + &processor, + &config, + vec![ + test_chunk(vec![], -1), + test_chunk(vec![test_span(1, 2, 3, 1), test_span(1, 3, 2, 0)], -1), + ], + RESCUE_NOW, + ); + assert_eq!(chunks.len(), 2); + assert!(chunks[0].spans.is_empty()); + assert_eq!(errors_sr(&chunks[1].spans[0]), None); + assert_eq!(errors_sr(&chunks[1].spans[1]), Some(1.0)); + assert_eq!(chunks[1].priority, -1); + } + #[test] fn test_empty_chunk_does_not_panic_and_stays_unchanged() { let (config, processor) = rescue_setup(always_keep_sampler_config()); @@ -1402,13 +1630,14 @@ mod tests { async fn test_rescue_is_noop_when_agent_stats_disabled() { let config = Config { agent_stats_computation_enabled: false, + error_rare_sample_tracer_drop: true, ..rescue_config(always_keep_sampler_config()) }; let (tx, mut rx): ( Sender, Receiver, ) = mpsc::channel(1); - let request = errored_p0_request(); + let request = errored_request(-1, Some("-3")); let trace_processor = trace_processor::ServerlessTraceProcessor::new( None, @@ -1430,7 +1659,21 @@ mod tests { panic!("expected V07 payload"); }; let chunk = &tracer_payloads[0].chunks[0]; - assert_eq!(chunk.priority, 0); + assert_eq!(chunk.priority, -1); + assert_eq!( + chunk + .tags + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + Some("-3") + ); + assert_eq!( + chunk.spans[0] + .meta + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + Some("-3") + ); for span in &chunk.spans { assert_eq!( errors_sr(span), @@ -1444,20 +1687,25 @@ mod tests { fn test_disabled_tps_disables_rescue_in_both_modes() { for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { for tps in [0.0_f64, -3.0] { - let (config, processor) = rescue_setup(ErrorSamplerConfig { + let (mut config, processor) = rescue_setup(ErrorSamplerConfig { mode, target_tps: tps, extra_sample_rate: 1.0, }); + config.error_rare_sample_tracer_drop = true; let chunks = run_rescue( &processor, &config, - vec![errored_root_chunk(0xe000, 0)], + vec![ + errored_root_chunk(0xe000, 0), + errored_root_chunk(0xe001, -1), + ], RESCUE_NOW, ); - assert_eq!(chunks.len(), 1, "chunk still forwarded"); + assert_eq!(chunks.len(), 2, "chunks still forwarded"); + assert_eq!(errors_sr(root(&chunks[1])), None); assert_eq!( errors_sr(root(&chunks[0])), None, @@ -1555,7 +1803,8 @@ mod tests { #[test] fn test_processor_clones_share_one_budget() { let sampler_config = rate_limited_sampler_config(1.0); - let (config, processor) = rescue_setup(sampler_config); + let (mut config, processor) = rescue_setup(sampler_config); + config.error_rare_sample_tracer_drop = true; let processor_clone = processor.clone(); // Ten errored chunks with the same signature but distinct IDs: five @@ -1565,10 +1814,28 @@ mod tests { let mut observed = Vec::new(); for (i, id) in ids.iter().enumerate() { let target = if i < 5 { &processor } else { &processor_clone }; + let mut manual = errored_root_chunk(*id, -1); + manual + .tags + .insert(super::DECISION_MAKER_KEY.into(), "-4".into()); + let mut healthy = errored_root_chunk(*id, -1); + healthy.spans[0].error = 0; + let excluded = run_rescue( + target, + &config, + vec![ + manual, + healthy, + errored_root_chunk(*id, -128), + errored_root_chunk(*id, 1), + ], + RESCUE_NOW, + ); + assert_eq!(rescued_count(&excluded), 0); let chunks = run_rescue( target, &config, - vec![errored_root_chunk(*id, 0)], + vec![errored_root_chunk(*id, if i % 2 == 0 { 0 } else { -1 })], RESCUE_NOW, ); observed.push(errors_sr(root(&chunks[0])).is_some()); @@ -1755,39 +2022,59 @@ mod tests { #[tokio::test] async fn test_stats_concentrator_observes_pre_rescue_chunks() { - let config = rescue_config(always_keep_sampler_config()); + let mut config = rescue_config(always_keep_sampler_config()); + config.error_rare_sample_tracer_drop = true; + let config = Arc::new(config); let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); - let concentrator = super::StatsConcentratorHandle::new(stats_tx); let processor = trace_processor::ServerlessTraceProcessor::new( - Some(concentrator), + Some(super::StatsConcentratorHandle::new(stats_tx)), sampler_for(&config.error_sampler), ); - - let res = processor - .process_traces( - Arc::new(config), - errored_p0_request(), - mpsc::channel(1).0, - Arc::new(create_test_metadata()), - ) - .await; - assert!(res.is_ok()); - - // The concentrator must receive the chunk exactly as the tracer sent - // it: before rescue stamping, with no `_dd.errors_sr` anywhere. - let chunk = match stats_rx.try_recv() { - Ok(crate::stats_concentrator_service::ConcentratorCommand::AddChunk(chunk, _)) => { - *chunk - } - Ok(_) => panic!("expected an AddChunk command"), - Err(err) => panic!("expected an AddChunk command, got {err}"), - }; - assert_eq!(chunk.priority, 0); - for span in &chunk.spans { + for (priority, dm, rescued) in [ + (0, None, true), + (-1, Some("-3"), true), + (-1, Some("-4"), false), + (-128, None, false), + ] { + let (tx, mut rx) = mpsc::channel(1); + let res = processor + .process_traces( + config.clone(), + errored_request(priority, dm), + tx, + Arc::new(create_test_metadata()), + ) + .await; + assert!(res.is_ok()); + let chunk = match stats_rx.try_recv().expect("stats chunk") { + crate::stats_concentrator_service::ConcentratorCommand::AddChunk(chunk, _) => { + *chunk + } + _ => panic!("expected AddChunk"), + }; + assert_eq!(chunk.priority, priority); assert_eq!( - errors_sr(span), - None, - "stats must observe the pre-rescue chunk" + chunk + .tags + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + dm + ); + assert!( + chunk.spans.iter().all(|span| errors_sr(span).is_none()), + "stats precede stamping" + ); + let send_data = rx.recv().await.expect("every chunk forwarded"); + let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { + panic!("expected V07"); + }; + let forwarded = &payloads[0].chunks[0]; + assert_eq!(forwarded.priority, priority); + assert_eq!( + errors_sr( + &forwarded.spans[trace_utils::get_root_span_index(&forwarded.spans).unwrap()] + ), + rescued.then_some(1.0) ); } } diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 4e02ad73..6d281120 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1290,9 +1290,23 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { verify_stats_request(&mock_intake).await; } -/// Builds a mixed four-trace payload for the error rescue wire test: an -/// eligible errored automatic-drop chunk, a healthy automatic-drop chunk, an -/// explicit user drop, and a positive-priority chunk. +/// Identifiable stats-eligible roots covering automatic, tracer, and manual drops. +fn error_rescue_cases() -> Vec<(&'static str, i32, i64, Option<&'static str>, bool)> { + vec![ + ("rescued_error_p0", 0, 1, None, false), + ("healthy_p0", 0, 0, None, false), + ("negative_absent", -1, 1, None, false), + ("negative_nonmanual", -1, 1, Some("-3"), false), + ("negative_probabilistic", -1, 1, Some("-9"), false), + ("negative_empty", -1, 1, Some(""), false), + ("manual_root", -1, 1, Some("-4"), false), + ("manual_child", -1, 1, Some("-4"), true), + ("other_negative_manual_dm", -2, 1, Some("-4"), false), + ("no_priority", -128, 1, Some("-3"), false), + ("priority_1", 1, 1, None, false), + ("priority_2", 2, 1, None, false), + ] +} fn create_error_rescue_test_payload() -> Vec { let start = UNIX_EPOCH.elapsed().unwrap().as_nanos() as i64; let span = |trace_id: u64, name: &str, error: i64, priority: f64| { @@ -1302,36 +1316,61 @@ fn create_error_rescue_test_payload() -> Vec { span["metrics"]["_sampling_priority_v1"] = json!(priority); span }; - let traces = vec![ - vec![span(700, "rescued_error_p0", 1, 0.0)], - vec![span(701, "healthy_p0", 0, 0.0)], - vec![span(702, "user_drop_p0", 1, -1.0)], - vec![span(703, "priority_1", 1, 1.0)], - ]; + let traces: Vec<_> = error_rescue_cases() + .into_iter() + .enumerate() + .map(|(index, (name, priority, error, dm, on_child))| { + let id = 700 + index as u64 * 10; + let mut root = span(id, name, error, f64::from(priority)); + if let Some(dm) = dm + && !on_child + { + root["meta"]["_dd.p.dm"] = json!(dm); + } + if on_child { + let mut child = create_test_json_span(id, id + 2, id + 1, start, false); + child["name"] = json!("manual_child_span"); + child["meta"]["span.kind"] = json!("server"); + child["meta"]["_dd.p.dm"] = json!(dm.unwrap()); + child["metrics"]["_sampling_priority_v1"] = json!(priority); + vec![root, child] + } else { + vec![root] + } + }) + .collect(); rmp_serde::to_vec(&traces).expect("Failed to serialize error rescue test trace") } -/// End-to-end error rescue verification against the outbound wire payload and -/// agent-computed stats: an eligible errored P0 chunk is rescued (positive -/// `_dd.errors_sr` on the root, priority still 0) while all other chunks are -/// forwarded unchanged, and every submitted span still contributes to stats. -/// This proves the wire contract only; backend retention of rescued chunks is -/// established by the backend implementation, not by this fake intake. -#[cfg(test)] +/// Checks unconditional promotion, opt-in rescue, unchanged priorities, forwarding, +/// and stats against the actual outbound payload. Fake intake proves the wire +/// contract, not backend retention. +#[tokio::test] +#[serial] +async fn test_error_rescue_feature_off_outbound_payload_and_stats() { + verify_error_rescue_outbound_payload_and_stats(false).await; +} + #[tokio::test] #[serial] -async fn test_error_rescue_outbound_payload_and_stats() { +async fn test_error_rescue_feature_on_outbound_payload_and_stats() { + verify_error_rescue_outbound_payload_and_stats(true).await; +} + +async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; use libdd_trace_protobuf::pb::{AgentPayload, TraceChunk}; use prost::Message as _; use std::collections::HashMap; - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; - let mut config = create_tcp_test_config(8137); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + // nextest runs tests in separate processes, so #[serial] alone cannot + // prevent the paired tests from competing for the same listener. + let mut config = create_tcp_test_config(if feature { 8138 } else { 8137 }); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; + config.error_rare_sample_tracer_drop = feature; // Deterministic rescue settings: every eligible chunk is kept at 1.0. config.error_sampler = ErrorSamplerConfig { mode: ErrorSamplerMode::AlwaysKeep, @@ -1377,11 +1416,17 @@ async fn test_error_rescue_outbound_payload_and_stats() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; + tokio::time::sleep(FLUSH_WAIT_DURATION).await; + let _ = shutdown_tx.send(true); + tokio::time::timeout(Duration::from_secs(10), agent_handle) + .await + .expect("bounded agent shutdown") + .expect("agent task"); // Decode the actual outbound protobuf payload(s) and locate each submitted // chunk by its root span name. - let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); + let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); let mut chunks_by_root_name: HashMap = HashMap::new(); for req in &trace_reqs { let agent_payload = AgentPayload::decode(&req.body[..]) @@ -1399,13 +1444,11 @@ async fn test_error_rescue_outbound_payload_and_stats() { } } - let names = [ - "rescued_error_p0", - "healthy_p0", - "user_drop_p0", - "priority_1", - ]; - for name in names { + let names: Vec<_> = error_rescue_cases() + .into_iter() + .map(|case| case.0) + .collect(); + for name in names.iter().copied() { assert!( chunks_by_root_name.contains_key(name), "expected chunk {name} to be forwarded to the backend, got: {:?}", @@ -1433,37 +1476,49 @@ async fn test_error_rescue_outbound_payload_and_stats() { "rescued chunk priority must remain 0, no promotion" ); - // Non-candidate chunks receive no rescue metric and keep their priority. - assert_eq!(errors_sr_of("healthy_p0"), None, "healthy P0 not rescued"); - assert_eq!(chunks_by_root_name["healthy_p0"].priority, 0); - assert_eq!( - errors_sr_of("user_drop_p0"), - None, - "explicit user drop never rescued" - ); - assert_eq!(chunks_by_root_name["user_drop_p0"].priority, -1); - assert_eq!( - errors_sr_of("priority_1"), - None, - "positive priority not rescued" - ); - assert_eq!(chunks_by_root_name["priority_1"].priority, 1); - - // Wait for the stats flush, then assert agent-computed stats include all - // four submitted spans, independent of rescue decisions. - tokio::time::sleep(FLUSH_WAIT_DURATION).await; - let _ = shutdown_tx.send(true); - let _ = agent_handle.await; - - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); - assert!( - !stats_reqs.is_empty(), - "Expected at least one stats request" - ); + for (name, priority, error, dm, on_child) in error_rescue_cases() { + let chunk = &chunks_by_root_name[name]; + let eligible = error != 0 + && (priority == 0 + || (feature + && priority < 0 + && priority != -128 + && !(priority == -1 && dm == Some("-4")))); + assert_eq!( + errors_sr_of(name), + eligible.then_some(1.0), + "{name}, feature={feature}" + ); + assert_eq!(chunk.priority, priority, "{name}"); + assert_eq!(chunk.tags.get("_dd.p.dm").map(String::as_str), dm, "{name}"); + assert_eq!(chunk.spans.len(), if on_child { 2 } else { 1 }); + for span in &chunk.spans { + assert_eq!( + span.metrics.get("_sampling_priority_v1").copied(), + Some(f64::from(priority)), + "{name}" + ); + let expected_dm = if (span.parent_id != 0) == on_child { + dm + } else { + None + }; + assert_eq!( + span.meta.get("_dd.p.dm").map(String::as_str), + expected_dm, + "{name}" + ); + if span.parent_id != 0 { + assert!(!span.metrics.contains_key("_dd.errors_sr")); + } + } + } - let all_groups: Vec<_> = stats_reqs - .iter() - .map(|req| decode_stats_payload(&req.body)) + // Agent-computed stats include every submitted span, independent of rescue. + verify_stats_request(&mock_intake).await; + let all_groups: Vec<_> = mock_intake + .stats_payloads() + .into_iter() .flat_map(|payload| { payload .stats @@ -1473,7 +1528,7 @@ async fn test_error_rescue_outbound_payload_and_stats() { }) .collect(); - for name in names { + for name in names.into_iter().chain(["manual_child_span"]) { let group = all_groups .iter() .find(|g| g.name == name) From a104da06d5e1d7e17b44b6014df1b913c75bc6e3 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 19 Oct 2026 15:22:25 -0400 Subject: [PATCH 12/49] Simplify error rescue test assertions --- crates/datadog-trace-agent/src/config.rs | 1 - .../src/trace_processor.rs | 15 +++++++------ .../tests/integration_test.rs | 22 ++++--------------- 3 files changed, 12 insertions(+), 26 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index c061cedd..bef85099 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -1211,7 +1211,6 @@ mod tests { }, ); } - assert!(!logs_contain("DD_APM_ERROR_TPS <= 0 disables error rescue")); logs_assert(|lines| { if lines .iter() diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index ed27a7f3..33b05c98 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1421,7 +1421,6 @@ mod tests { } } } - assert!(logs_contain(warning)); logs_assert(|lines| { let count = lines .iter() @@ -1705,12 +1704,14 @@ mod tests { ); assert_eq!(chunks.len(), 2, "chunks still forwarded"); - assert_eq!(errors_sr(root(&chunks[1])), None); - assert_eq!( - errors_sr(root(&chunks[0])), - None, - "mode {mode:?} tps {tps}: disabled sampler must not rescue" - ); + for chunk in &chunks { + assert_eq!( + errors_sr(root(chunk)), + None, + "mode {mode:?} tps {tps} priority {}: disabled sampler must not rescue", + chunk.priority + ); + } } } } diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 6d281120..b4013eda 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1307,6 +1307,7 @@ fn error_rescue_cases() -> Vec<(&'static str, i32, i64, Option<&'static str>, bo ("priority_2", 2, 1, None, false), ] } + fn create_error_rescue_test_payload() -> Vec { let start = UNIX_EPOCH.elapsed().unwrap().as_nanos() as i64; let span = |trace_id: u64, name: &str, error: i64, priority: f64| { @@ -1464,18 +1465,6 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { .and_then(|root| root.metrics.get("_dd.errors_sr").copied()) }; - // The eligible errored automatic-drop chunk is rescued: positive - // `_dd.errors_sr` on the root, chunk priority still 0. - assert_eq!( - errors_sr_of("rescued_error_p0"), - Some(1.0), - "rescued root must carry a positive _dd.errors_sr" - ); - assert_eq!( - chunks_by_root_name["rescued_error_p0"].priority, 0, - "rescued chunk priority must remain 0, no promotion" - ); - for (name, priority, error, dm, on_child) in error_rescue_cases() { let chunk = &chunks_by_root_name[name]; let eligible = error != 0 @@ -1498,17 +1487,14 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { Some(f64::from(priority)), "{name}" ); - let expected_dm = if (span.parent_id != 0) == on_child { - dm - } else { - None - }; + let is_child = span.parent_id != 0; + let expected_dm = if is_child == on_child { dm } else { None }; assert_eq!( span.meta.get("_dd.p.dm").map(String::as_str), expected_dm, "{name}" ); - if span.parent_id != 0 { + if is_child { assert!(!span.metrics.contains_key("_dd.errors_sr")); } } From 134039d959a01fb12984f510ede986b220c94cad Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 16:08:19 -0400 Subject: [PATCH 13/49] Remove SCL error sampler mode setting --- crates/datadog-agent-trace-sampler/README.md | 9 +- crates/datadog-trace-agent/src/config.rs | 176 ++++++++----------- 2 files changed, 74 insertions(+), 111 deletions(-) diff --git a/crates/datadog-agent-trace-sampler/README.md b/crates/datadog-agent-trace-sampler/README.md index 55d4a653..cbb2d618 100644 --- a/crates/datadog-agent-trace-sampler/README.md +++ b/crates/datadog-agent-trace-sampler/README.md @@ -15,12 +15,11 @@ rescue strategies are available, selected by `ErrorSamplerMode`: `ScoreSampler` (`ScoreSampler` targeting `ErrorTPS`, from `pkg/trace/sampler/` in `DataDog/datadog-agent`): keep up to `target_tps` error traces per second, distributed fairly across distinct trace signatures. Suits continuous - processes that can hit error storms (the Serverless Compatibility Layer's - default). + processes that can hit error storms (the Serverless Compatibility Layer uses + this strategy). -Per-platform defaults are chosen by each consumer's config layer (e.g. via a -`DD_APM_ERROR_SAMPLER_MODE` env var), not this crate — it has no notion of -which platform it runs on. +Each consumer chooses its strategy through its configuration or integration +layer, not this crate, which has no notion of which platform it runs on. ## Why dependency-free diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index bef85099..e2bc9a15 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -22,17 +22,14 @@ const DEFAULT_APM_RECEIVER_PORT: u16 = 8126; /// Default error rescue budget: error traces per second, matching the Go agent's `ErrorTPS`. const DEFAULT_ERROR_SAMPLER_TPS: f64 = 10.0; -/// Parses the error rescue sampler settings from the environment. +/// Parses the rate-limited error rescue budget from the environment. /// -/// - `DD_APM_ERROR_SAMPLER_MODE`: `rate_limited` (default) or `always_keep`. -/// Surrounding whitespace and casing are normalized; any other value warns -/// and falls back to `rate_limited`. -/// - `DD_APM_ERROR_TPS`: target error traces per second, default 10. Parsed as -/// `f64`; non-finite or unparsable values warn and fall back to the default. -/// A finite value <= 0 disables rescue in either mode. +/// `DD_APM_ERROR_TPS` sets the target error traces per second, default 10. Parsed +/// as `f64`; non-finite or unparsable values warn and fall back to the default. +/// A finite value <= 0 disables rescue. /// -/// `extra_sample_rate` is fixed at 1.0 and intentionally not configurable. -/// Startup stays operational when either optional setting is malformed. +/// The strategy is fixed to `RateLimited`, and `extra_sample_rate` is fixed at +/// 1.0. Startup stays operational when the optional TPS setting is malformed. fn parse_error_sampler_config() -> ErrorSamplerConfig { let mut config = ErrorSamplerConfig { mode: ErrorSamplerMode::RateLimited, @@ -40,16 +37,6 @@ fn parse_error_sampler_config() -> ErrorSamplerConfig { extra_sample_rate: 1.0, }; - if let Ok(raw) = env::var("DD_APM_ERROR_SAMPLER_MODE") { - match raw.trim().to_lowercase().as_str() { - "rate_limited" => config.mode = ErrorSamplerMode::RateLimited, - "always_keep" => config.mode = ErrorSamplerMode::AlwaysKeep, - _ => { - warn!("Invalid DD_APM_ERROR_SAMPLER_MODE {raw:?}; using default mode rate_limited"); - } - } - } - if let Ok(raw) = env::var("DD_APM_ERROR_TPS") { match raw.trim().parse::() { Ok(tps) if tps.is_finite() => config.target_tps = tps, @@ -184,8 +171,8 @@ pub struct Config { pub additional_metric_tags_cardinality_limit: Option, /// Whether the agent should compute trace stats pub agent_stats_computation_enabled: bool, - /// Error rescue sampler settings, parsed from `DD_APM_ERROR_SAMPLER_MODE` - /// and `DD_APM_ERROR_TPS`. `extra_sample_rate` is fixed at 1.0 and not + /// Rate-limited error rescue settings, with the budget parsed from + /// `DD_APM_ERROR_TPS`. `extra_sample_rate` is fixed at 1.0 and not /// configurable. See `parse_error_sampler_config` for the defaults. pub error_sampler: ErrorSamplerConfig, /// Opt-in eligibility for tracer-dropped errors, independent of sampler enablement. @@ -1015,39 +1002,29 @@ mod tests { fn test_tracer_drop_feature_does_not_enable_sampler_or_stats() { for feature in [None, Some("error_rare_sample_tracer_drop"), Some("other")] { for (tps, expected_tps) in [(None, 10.0), (Some("0"), 0.0), (Some("-1"), -1.0)] { - for mode in [None, Some("always_keep")] { - for stats in ["true", "false"] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("test")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("test")), - ("DD_APM_FEATURES", feature), - ("DD_APM_ERROR_TPS", tps), - ("DD_APM_ERROR_SAMPLER_MODE", mode), - ("DD_AGENT_STATS_COMPUTATION_ENABLED", Some(stats)), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_rare_sample_tracer_drop, - feature == Some("error_rare_sample_tracer_drop") - ); - assert_eq!(config.error_sampler.target_tps, expected_tps); - assert_eq!(config.error_sampler.extra_sample_rate, 1.0); - assert_eq!( - config.error_sampler.mode, - if mode.is_some() { - ErrorSamplerMode::AlwaysKeep - } else { - ErrorSamplerMode::RateLimited - } - ); - assert_eq!(config.agent_stats_computation_enabled, stats == "true"); - }, - ); - } + for stats in ["true", "false"] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("test")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("test")), + ("DD_APM_FEATURES", feature), + ("DD_APM_ERROR_TPS", tps), + ("DD_AGENT_STATS_COMPUTATION_ENABLED", Some(stats)), + ], + || { + let config = config::Config::new().unwrap(); + assert_eq!( + config.error_rare_sample_tracer_drop, + feature == Some("error_rare_sample_tracer_drop") + ); + assert_eq!(config.error_sampler.target_tps, expected_tps); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); + assert_eq!(config.error_sampler.mode, ErrorSamplerMode::RateLimited); + assert_eq!(config.agent_stats_computation_enabled, stats == "true"); + }, + ); } } } @@ -1069,7 +1046,6 @@ mod tests { ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("test")), ("DD_APM_ERROR_TPS", None), - ("DD_APM_ERROR_SAMPLER_MODE", None), ], || { assert!(!config::Config::new().unwrap().error_rare_sample_tracer_drop); @@ -1099,7 +1075,6 @@ mod tests { ("WEBSITE_SITE_NAME", Some("my-azure-function")), ("DD_APM_FEATURES", None), ("DD_APM_ERROR_TPS", None), - ("DD_APM_ERROR_SAMPLER_MODE", None), ], || { let config = config::Config::new().unwrap(); @@ -1110,55 +1085,46 @@ mod tests { #[test] #[serial] - fn test_error_sampler_mode_normalization() { - for (raw, expected) in [ - ("always_keep", ErrorSamplerMode::AlwaysKeep), - ("ALWAYS_KEEP", ErrorSamplerMode::AlwaysKeep), - (" Always_Keep ", ErrorSamplerMode::AlwaysKeep), - ("rate_limited", ErrorSamplerMode::RateLimited), - ("RATE_LIMITED", ErrorSamplerMode::RateLimited), - (" Rate_Limited ", ErrorSamplerMode::RateLimited), + fn test_error_sampler_mode_env_is_ignored() { + for raw in [ + "always_keep", + "ALWAYS_KEEP", + " Always_Keep ", + "rate_limited", + "RATE_LIMITED", + " Rate_Limited ", + "bogus", + "", + "always-keep", ] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", None), - ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_sampler.mode, expected, - "mode {raw:?} should parse to {expected:?}" - ); - }, - ); - } - } - - #[test] - #[serial] - fn test_error_sampler_invalid_mode_falls_back_to_default() { - for raw in ["bogus", "", "always-keep"] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", None), - ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), - ], - || { - let config = config::Config::new().unwrap(); - assert_rate_limited_default(&config.error_sampler); - }, - ); + for (tps, expected_tps) in [ + (None, 10.0), + (Some("5.5"), 5.5), + (Some("0"), 0.0), + (Some("-1"), -1.0), + ] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_TPS", tps), + ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), + ], + || { + let config = config::Config::new().unwrap(); + assert_eq!( + config.error_sampler.mode, + ErrorSamplerMode::RateLimited, + "mode setting {raw:?} must not change the strategy with TPS {tps:?}" + ); + assert_eq!(config.error_sampler.target_tps, expected_tps); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); + }, + ); + } } } @@ -1173,7 +1139,6 @@ mod tests { ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_SAMPLER_MODE", None), ("DD_APM_ERROR_TPS", Some(raw)), ], || { @@ -1199,7 +1164,6 @@ mod tests { ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), - ("DD_APM_ERROR_SAMPLER_MODE", None), ("DD_APM_ERROR_TPS", Some(raw)), ], || { From d1241ae6ac0c39a0ee90d88ddd656a15071636df Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 16:52:56 -0400 Subject: [PATCH 14/49] Share error sampler floor under one mutex --- .../src/trace_processor.rs | 115 ++++++++++-------- 1 file changed, 67 insertions(+), 48 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 33b05c98..015e0e78 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -2,10 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; -use std::sync::atomic::{AtomicI64, Ordering}; use async_trait::async_trait; -use datadog_agent_trace_sampler::{ErrorsSampler, SampleDecision, SpanView, TraceView}; +use datadog_agent_trace_sampler::{ + ErrorSamplerConfig, ErrorsSampler, SampleDecision, SpanView, TraceView, +}; use http_body_util::BodyExt; use hyper::{StatusCode, http}; use libdd_common::http_common; @@ -144,49 +145,64 @@ impl TraceChunkProcessor for ChunkProcessor { /// handing their pieces to the flusher) at once const MAX_IN_FLIGHT_ENQUEUES: usize = 10; +/// Error rescue sampler and its monotonic timestamp floor, shared under one mutex. +pub struct ErrorRescueState { + sampler: ErrorsSampler, + last_sampler_timestamp: i64, +} + +impl ErrorRescueState { + fn new(config: ErrorSamplerConfig) -> Self { + Self { + sampler: ErrorsSampler::new(config), + last_sampler_timestamp: i64::MIN, + } + } + + /// Records a timestamp without moving the sampler's rolling window backwards. + fn clamp_timestamp(&mut self, now_unix_secs: i64) -> i64 { + self.last_sampler_timestamp = self.last_sampler_timestamp.max(now_unix_secs); + self.last_sampler_timestamp + } +} + #[derive(Clone)] pub struct ServerlessTraceProcessor { pub stats_concentrator: Option, enqueue_permits: Arc, - /// Shared error rescue sampler. The `Arc` means processor clones (one per - /// connection) share a single sampler state and TPS budget for the whole - /// process lifetime. - error_sampler: Arc>, - /// Last timestamp handed to the sampler, shared across clones. Clamps the - /// clock so a backward wall-clock step cannot move sampler bucket IDs - /// backwards. Updated only under the sampler lock. - last_sampler_timestamp: Arc, + /// Processor clones and independently constructed processors using this + /// `Arc` share one TPS budget and timestamp floor under the same lock. + error_rescue_state: Arc>, } impl ServerlessTraceProcessor { #[allow(clippy::must_use_candidate)] pub fn new( stats_concentrator: Option, - error_sampler: Arc>, + error_rescue_state: Arc>, ) -> Self { ServerlessTraceProcessor { stats_concentrator, enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), - error_sampler, - last_sampler_timestamp: Arc::new(AtomicI64::new(i64::MIN)), + error_rescue_state, } } - /// Builds the error rescue sampler from parsed config settings. + /// Builds the shared error rescue state from parsed config settings. #[must_use] - pub fn new_error_sampler(config: &Config) -> Arc> { + pub fn new_error_sampler(config: &Config) -> Arc> { if config.error_rare_sample_tracer_drop && config.error_sampler.target_tps <= 0.0 { warn!( "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue" ); } - Arc::new(Mutex::new(ErrorsSampler::new(config.error_sampler))) + Arc::new(Mutex::new(ErrorRescueState::new(config.error_sampler))) } - /// Locks the shared sampler, recovering from a poisoned guard (a panic in - /// another thread while holding the lock) instead of panicking on unlock. - fn lock_sampler(&self) -> MutexGuard<'_, ErrorsSampler> { - self.error_sampler + /// Locks the shared sampler and timestamp floor, recovering from a poisoned + /// guard (a panic in another thread while holding the lock). + fn lock_sampler(&self) -> MutexGuard<'_, ErrorRescueState> { + self.error_rescue_state .lock() .unwrap_or_else(PoisonError::into_inner) } @@ -201,10 +217,10 @@ impl ServerlessTraceProcessor { /// Chunks are never removed or reordered: unrescued chunks stay in the /// payload for the backend to handle. fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { - let mut sampler = self.lock_sampler(); + let mut state = self.lock_sampler(); // Skip all view construction when the sampler is disabled by config // (target_tps <= 0): nothing can be rescued. - if sampler.is_disabled() { + if state.sampler.is_disabled() { return; } // Read the clock while holding the sampler lock so that concurrent @@ -214,8 +230,8 @@ impl ServerlessTraceProcessor { let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs() as i64); - let now_unix_secs = self.clamp_sampler_timestamp(now_unix_secs); - self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); + let now_unix_secs = state.clamp_timestamp(now_unix_secs); + self.rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); } /// Test-only variant that injects a synthetic timestamp so tests can @@ -227,22 +243,11 @@ impl ServerlessTraceProcessor { config: &Config, now_unix_secs: i64, ) { - let mut sampler = self.lock_sampler(); - if sampler.is_disabled() { + let mut state = self.lock_sampler(); + if state.sampler.is_disabled() { return; } - self.rescue_with_sampler(payload, config, now_unix_secs, &mut sampler); - } - - /// Clamps the timestamp so it never moves backwards relative to the last - /// one handed to the shared sampler, and records it as the new floor. Must - /// be called while holding the sampler lock, alongside the clock read, so - /// concurrent requests cannot interleave a read with the clamp. - fn clamp_sampler_timestamp(&self, now_unix_secs: i64) -> i64 { - let previous = self - .last_sampler_timestamp - .fetch_max(now_unix_secs, Ordering::Relaxed); - now_unix_secs.max(previous) + self.rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); } fn rescue_with_sampler( @@ -591,8 +596,8 @@ mod tests { config::{Config, Tags}, peer_tags::peer_tag_keys, trace_processor::{ - self, MAX_IN_FLIGHT_ENQUEUES, TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY, TraceProcessor, - encoded_size, split_oversized_payloads, + self, ErrorRescueState, MAX_IN_FLIGHT_ENQUEUES, TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY, + TraceProcessor, encoded_size, split_oversized_payloads, }, }; use datadog_agent_trace_sampler::{ @@ -667,7 +672,7 @@ mod tests { } } - fn default_error_sampler() -> Arc> { + fn default_error_sampler() -> Arc> { sampler_for(&ErrorSamplerConfig::default()) } @@ -1172,8 +1177,8 @@ mod tests { } } - fn sampler_for(config: &ErrorSamplerConfig) -> Arc> { - Arc::new(std::sync::Mutex::new(ErrorsSampler::new(*config))) + fn sampler_for(config: &ErrorSamplerConfig) -> Arc> { + Arc::new(std::sync::Mutex::new(ErrorRescueState::new(*config))) } /// Builds the standard rescue test fixture: a config with the given @@ -1878,19 +1883,33 @@ mod tests { ); } + #[test] + fn test_independent_processors_share_sampler_timestamp_floor() { + let sampler = sampler_for(&rate_limited_sampler_config(1.0)); + let first = trace_processor::ServerlessTraceProcessor::new(None, Arc::clone(&sampler)); + let second = trace_processor::ServerlessTraceProcessor::new(None, sampler); + + assert_eq!(first.lock_sampler().last_sampler_timestamp, i64::MIN); + assert_eq!(first.lock_sampler().clamp_timestamp(200), 200); + assert_eq!(second.lock_sampler().clamp_timestamp(50), 200); + assert_eq!(second.lock_sampler().clamp_timestamp(300), 300); + assert_eq!(first.lock_sampler().clamp_timestamp(250), 300); + } + #[test] fn test_sampler_timestamp_clamp_never_moves_backwards() { let (_config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); // A clone shares the clamp floor with the original. let clone = processor.clone(); - assert_eq!(processor.clamp_sampler_timestamp(100), 100); + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + assert_eq!(processor.lock_sampler().clamp_timestamp(100), 100); // Backward timestamps, through either instance, are clamped to the // last value seen; forward ones update the floor. - assert_eq!(clone.clamp_sampler_timestamp(50), 100); - assert_eq!(processor.clamp_sampler_timestamp(100), 100); - assert_eq!(clone.clamp_sampler_timestamp(200), 200); - assert_eq!(processor.clamp_sampler_timestamp(199), 200); + assert_eq!(clone.lock_sampler().clamp_timestamp(50), 100); + assert_eq!(processor.lock_sampler().clamp_timestamp(100), 100); + assert_eq!(clone.lock_sampler().clamp_timestamp(200), 200); + assert_eq!(processor.lock_sampler().clamp_timestamp(199), 200); } #[test] From 045cc1d682ff8bc4ebb6f7cfc05fc3fc4da8457a Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 17:15:20 -0400 Subject: [PATCH 15/49] Route error rescue through shared clock helper --- .../src/trace_processor.rs | 102 ++++++++++++++++-- 1 file changed, 91 insertions(+), 11 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 015e0e78..8defc3e8 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -217,6 +217,19 @@ impl ServerlessTraceProcessor { /// Chunks are never removed or reordered: unrescued chunks stay in the /// payload for the backend to handle. fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { + self.apply_error_rescue_with_clock(payload, config, || { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_secs() as i64) + }); + } + + fn apply_error_rescue_with_clock i64>( + &self, + payload: &mut TracerPayloadCollection, + config: &Config, + read_clock: F, + ) { let mut state = self.lock_sampler(); // Skip all view construction when the sampler is disabled by config // (target_tps <= 0): nothing can be rescued. @@ -227,15 +240,12 @@ impl ServerlessTraceProcessor { // requests deliver timestamps in lock-acquisition order, and clamp so a // backward wall-clock step cannot move the sampler's rolling window // backwards, which would undercount TPS and rescue too many chunks. - let now_unix_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0, |d| d.as_secs() as i64); - let now_unix_secs = state.clamp_timestamp(now_unix_secs); + let now_unix_secs = state.clamp_timestamp(read_clock()); self.rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); } - /// Test-only variant that injects a synthetic timestamp so tests can - /// exercise the rolling window without sleeping. + /// Test-only variant that injects a synthetic clock reading through the + /// production clamp path, exercising the rolling window without sleeping. #[cfg(test)] fn apply_error_rescue_at( &self, @@ -243,11 +253,7 @@ impl ServerlessTraceProcessor { config: &Config, now_unix_secs: i64, ) { - let mut state = self.lock_sampler(); - if state.sampler.is_disabled() { - return; - } - self.rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); + self.apply_error_rescue_with_clock(payload, config, || now_unix_secs); } fn rescue_with_sampler( @@ -1883,6 +1889,80 @@ mod tests { ); } + #[test] + fn test_error_rescue_clamps_backward_timestamps_across_clones() { + let (config, processor) = rescue_setup(always_keep_sampler_config()); + let clone = processor.clone(); + + let chunks = run_rescue( + &processor, + &config, + vec![errored_root_chunk(0x10000, 0)], + RESCUE_NOW, + ); + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); + + let chunks = run_rescue( + &clone, + &config, + vec![errored_root_chunk(0x10001, 0)], + RESCUE_NOW - 40, + ); + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + assert_eq!(clone.lock_sampler().last_sampler_timestamp, RESCUE_NOW); + } + + #[test] + fn test_error_rescue_reads_clock_under_lock() { + let (config, processor) = rescue_setup(always_keep_sampler_config()); + let mut payload = TracerPayloadCollection::V07(vec![pb::TracerPayload { + chunks: vec![errored_root_chunk(0x10000, 0)], + ..Default::default() + }]); + let mut clock_called = false; + + processor.apply_error_rescue_with_clock(&mut payload, &config, || { + clock_called = true; + assert!(matches!( + processor.error_rescue_state.try_lock(), + Err(std::sync::TryLockError::WouldBlock) + )); + RESCUE_NOW + }); + + assert!(clock_called); + assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); + } + + #[test] + fn test_disabled_error_rescue_does_not_read_clock() { + for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { + for target_tps in [0.0, -3.0] { + let (config, processor) = rescue_setup(ErrorSamplerConfig { + mode, + target_tps, + extra_sample_rate: 1.0, + }); + let original = pb::TracerPayload { + chunks: vec![errored_root_chunk(0x10000, 0)], + ..Default::default() + }; + let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); + + processor.apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("disabled sampler must not read the clock"); + }); + + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + let TracerPayloadCollection::V07(payloads) = payload else { + panic!("disabled rescue must preserve the payload format"); + }; + assert_eq!(payloads, vec![original]); + } + } + } + #[test] fn test_independent_processors_share_sampler_timestamp_floor() { let sampler = sampler_for(&rate_limited_sampler_config(1.0)); From cd727cea2c12e5ca8af59166ec2af1adf30b5d98 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 17:46:21 -0400 Subject: [PATCH 16/49] Rescue errored no-priority chunks --- .../src/trace_processor.rs | 40 +++++++++---------- .../tests/integration_test.rs | 6 +-- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 8defc3e8..bb976434 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -210,8 +210,8 @@ impl ServerlessTraceProcessor { /// Applies the error rescue pass to a payload collection. /// /// For every eligible V07 chunk containing a non-zero error flag, consults - /// the shared sampler. P0 is always eligible; the opt-in feature adds negative - /// priorities except the no-priority sentinel and exact manual drops. + /// the shared sampler. P0 and no-priority (-128) chunks are always eligible; + /// the opt-in feature adds other negative priorities except exact manual drops. /// On a keep, stamps the sampler's `_dd.errors_sr` on the chunk's root /// span without priority promotion. Retention is decided by the backend. /// Chunks are never removed or reordered: unrescued chunks stay in the @@ -356,8 +356,10 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { /// Consults the error sampler for one chunk and, on a keep, stamps the /// sampler's `_dd.errors_sr` value on the chunk's root span. /// -/// P0 chunks are candidates. With the feature enabled, negative priorities also -/// qualify except the no-priority sentinel and priority -1 with chunk DM -4. +/// P0 and no-priority (-128) chunks are candidates. No priority means a missing +/// sampling decision, not an explicit drop, matching Go's baseline eligibility. +/// With the feature enabled, other negative priorities also qualify except +/// priority -1 with chunk DM -4. /// Only promoted chunk tags are consulted, never span metadata. The chunk /// must contain at least one span with a non-zero error flag; HTTP status or /// error metadata alone does not qualify. The root is resolved with @@ -373,9 +375,9 @@ fn sample_and_stamp( tracer_drop_enabled: bool, ) { let eligible = chunk.priority == 0 + || chunk.priority == i32::from(i8::MIN) || (tracer_drop_enabled && chunk.priority < 0 - && chunk.priority != i32::from(i8::MIN) && !(chunk.priority == -1 && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) == Some(MANUAL_DECISION_MAKER))); @@ -1377,9 +1379,9 @@ mod tests { let mut chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); let expected = error != 0 && (priority == 0 + || priority == -128 || (feature && priority < 0 - && priority != -128 && !(priority == -1 && dm == Some("-4")))); assert_eq!( errors_sr(&chunks[0].spans[0]), @@ -1559,9 +1561,9 @@ mod tests { fn test_non_automatic_drop_priorities_are_never_rescued() { let (config, processor) = rescue_setup(always_keep_sampler_config()); - // With the feature off, all negative and positive priorities, - // including the no-priority sentinel, are out of scope. - let priorities = [-1_i32, -5, 1, 2, i8::MIN as i32]; + // With the feature off, explicit negative drops and positive keeps + // are out of scope. The no-priority sentinel is not an explicit drop. + let priorities = [-1_i32, -5, 1, 2]; let chunks = run_rescue( &processor, &config, @@ -1710,11 +1712,12 @@ mod tests { vec![ errored_root_chunk(0xe000, 0), errored_root_chunk(0xe001, -1), + errored_root_chunk(0xe002, i32::from(i8::MIN)), ], RESCUE_NOW, ); - assert_eq!(chunks.len(), 2, "chunks still forwarded"); + assert_eq!(chunks.len(), 3, "chunks still forwarded"); for chunk in &chunks { assert_eq!( errors_sr(root(chunk)), @@ -1820,8 +1823,9 @@ mod tests { let processor_clone = processor.clone(); // Ten errored chunks with the same signature but distinct IDs: five - // through the original, five through the clone. All must draw from the - // same budget, matching a standalone sampler fed the same sequence. + // through the original, five through the clone. Mix P0, explicit drops, + // and no-priority chunks. All must draw from the same budget, matching + // a standalone sampler fed the same sequence. let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); let mut observed = Vec::new(); for (i, id) in ids.iter().enumerate() { @@ -1835,19 +1839,15 @@ mod tests { let excluded = run_rescue( target, &config, - vec![ - manual, - healthy, - errored_root_chunk(*id, -128), - errored_root_chunk(*id, 1), - ], + vec![manual, healthy, errored_root_chunk(*id, 1)], RESCUE_NOW, ); assert_eq!(rescued_count(&excluded), 0); + let priority = [0, -1, i32::from(i8::MIN)][i % 3]; let chunks = run_rescue( target, &config, - vec![errored_root_chunk(*id, if i % 2 == 0 { 0 } else { -1 })], + vec![errored_root_chunk(*id, priority)], RESCUE_NOW, ); observed.push(errors_sr(root(&chunks[0])).is_some()); @@ -2134,7 +2134,7 @@ mod tests { (0, None, true), (-1, Some("-3"), true), (-1, Some("-4"), false), - (-128, None, false), + (-128, None, true), ] { let (tx, mut rx) = mpsc::channel(1); let res = processor diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index b4013eda..873437ac 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1469,10 +1469,8 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { let chunk = &chunks_by_root_name[name]; let eligible = error != 0 && (priority == 0 - || (feature - && priority < 0 - && priority != -128 - && !(priority == -1 && dm == Some("-4")))); + || priority == -128 + || (feature && priority < 0 && !(priority == -1 && dm == Some("-4")))); assert_eq!( errors_sr_of(name), eligible.then_some(1.0), From 5ffcd60006cbfa7be697d67fa16c8238d6aeb2db Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 20:25:21 -0400 Subject: [PATCH 17/49] Decouple error rescue from stats computation --- .../src/trace_processor.rs | 129 ++++++++++-------- 1 file changed, 70 insertions(+), 59 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index bb976434..59348114 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -522,13 +522,12 @@ impl TraceProcessor for ServerlessTraceProcessor { Self::send_to_concentrator(concentrator, &payload); } - // Error rescue runs after stats submission so the concentrator observes every - // submitted chunk before rescue stamping, and before payload splitting so - // the newly inserted metric is included in the recomputed outbound size. It is - // gated on agent stats computation, independently of tracer-drop eligibility. - if config.agent_stats_computation_enabled { - self.apply_error_rescue(&mut payload, &config); - } + // Error rescue is independent of agent stats computation: eligible errored + // chunks can arrive whether or not the agent computes stats. Run it after + // conditional stats submission so the concentrator observes chunks before + // rescue stamping, and before payload splitting so the inserted metric is + // included in the recomputed outbound size. + self.apply_error_rescue(&mut payload, &config); let pieces: Vec<(TracerPayloadCollection, usize)> = match payload { TracerPayloadCollection::V07(payloads) => { @@ -1639,59 +1638,71 @@ mod tests { } #[tokio::test] - async fn test_rescue_is_noop_when_agent_stats_disabled() { - let config = Config { - agent_stats_computation_enabled: false, - error_rare_sample_tracer_drop: true, - ..rescue_config(always_keep_sampler_config()) - }; - let (tx, mut rx): ( - Sender, - Receiver, - ) = mpsc::channel(1); - let request = errored_request(-1, Some("-3")); - - let trace_processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); - let res = trace_processor - .process_traces( - Arc::new(config), - request, - tx, - Arc::new(create_test_metadata()), - ) - .await; - assert!(res.is_ok()); - - let send_data = rx.recv().await.expect("payload forwarded"); - let payloads = send_data.get_payloads(); - let TracerPayloadCollection::V07(tracer_payloads) = payloads else { - panic!("expected V07 payload"); - }; - let chunk = &tracer_payloads[0].chunks[0]; - assert_eq!(chunk.priority, -1); - assert_eq!( - chunk - .tags - .get(super::DECISION_MAKER_KEY) - .map(String::as_str), - Some("-3") - ); - assert_eq!( - chunk.spans[0] - .meta - .get(super::DECISION_MAKER_KEY) - .map(String::as_str), - Some("-3") - ); - for span in &chunk.spans { - assert_eq!( - errors_sr(span), - None, - "rescue must be a no-op when agent stats computation is disabled" + async fn test_rescue_stamps_eligible_chunks_when_agent_stats_disabled() { + for feature in [false, true] { + let config = Arc::new(Config { + agent_stats_computation_enabled: false, + error_rare_sample_tracer_drop: feature, + ..rescue_config(always_keep_sampler_config()) + }); + let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); + let processor = trace_processor::ServerlessTraceProcessor::new( + Some(super::StatsConcentratorHandle::new(stats_tx)), + sampler_for(&config.error_sampler), ); + for (priority, dm, rescued) in [ + (0, None, true), + (-128, None, true), + (-1, None, feature), + (-1, Some("-3"), feature), + (-1, Some("-4"), false), + (1, None, false), + (2, None, false), + ] { + let (tx, mut rx) = mpsc::channel(1); + let res = processor + .process_traces( + config.clone(), + errored_request(priority, dm), + tx, + Arc::new(create_test_metadata()), + ) + .await; + assert!(res.is_ok()); + assert!( + matches!( + stats_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + ), + "stats must not be submitted when agent stats computation is disabled" + ); + + let send_data = rx.recv().await.expect("payload forwarded"); + let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { + panic!("expected V07 payload"); + }; + let chunk = &payloads[0].chunks[0]; + assert_eq!( + errors_sr(&chunk.spans[0]), + rescued.then_some(1.0), + "feature={feature} priority={priority} dm={dm:?}" + ); + assert_eq!(chunk.priority, priority, "priority must not be promoted"); + assert_eq!( + chunk + .tags + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + dm + ); + assert_eq!( + chunk.spans[0] + .meta + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + dm + ); + } } } From 71f136ae36b91425b8140b9457334e9f1d0ee615 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 20:39:05 -0400 Subject: [PATCH 18/49] Consolidate error rescue eligibility tests --- .../src/trace_processor.rs | 175 ++---------------- 1 file changed, 15 insertions(+), 160 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 59348114..1a77c0c9 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1346,7 +1346,7 @@ mod tests { #[test] fn test_tracer_drop_rescue_eligibility_matrix() { for feature in [false, true] { - for priority in [0, -1, -2, -127, i32::from(i8::MIN), 1, 2] { + for priority in [0, -1, -2, -127, -5, i32::from(i8::MIN), 1, 2] { for dm in [ None, Some(""), @@ -1369,6 +1369,19 @@ mod tests { chunk.spans[0] .meta .insert("http.status_code".into(), "500".into()); + chunk.spans[0] + .meta + .insert("error.type".into(), "Error".into()); + chunk.spans[0] + .metrics + .insert("_sampling_priority_v1".into(), f64::from(priority)); + chunk.spans[0].metrics.insert("_sample_rate".into(), 0.25); + chunk.spans[0] + .metrics + .insert("_dd.span_sampling.rule".into(), 1.0); + chunk.spans[1].meta.insert("keep".into(), "me".into()); + chunk.origin = "synthetics".into(); + chunk.tags.insert("origin".into(), "synthetics".into()); if let Some(dm) = dm { chunk .tags @@ -1376,6 +1389,7 @@ mod tests { } let original = chunk.clone(); let mut chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); + assert_eq!(chunks.len(), 1, "chunk must still be forwarded"); let expected = error != 0 && (priority == 0 || priority == -128 @@ -1462,26 +1476,6 @@ mod tests { assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); } - #[test] - fn test_rescue_stamps_actual_root_when_error_is_on_child() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - // Healthy root first, errored child second: only the root may be stamped. - let spans = vec![ - test_span(0xfeed, 0x101, 0, 0), - test_span(0xfeed, 0x102, 0x101, 1), - ]; - let chunks = run_rescue(&processor, &config, vec![test_chunk(spans, 0)], RESCUE_NOW); - - assert_eq!(chunks.len(), 1); - assert_eq!(errors_sr(&chunks[0].spans[0]), Some(1.0), "root stamped"); - assert_eq!( - errors_sr(&chunks[0].spans[1]), - None, - "errored child must not be stamped" - ); - } - #[test] fn test_rescue_uses_resolved_root_not_first_span() { let (config, processor) = rescue_setup(always_keep_sampler_config()); @@ -1535,54 +1529,6 @@ mod tests { } } - #[test] - fn test_http_500_metadata_alone_is_not_an_error() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - let mut span = test_span(0xc0de, 0xc0de + 1, 0, 0); - span.meta - .insert("http.status_code".to_string(), "500".to_string()); - span.meta - .insert("error.type".to_string(), "Error".to_string()); - - let chunks = run_rescue( - &processor, - &config, - vec![test_chunk(vec![span], 0)], - RESCUE_NOW, - ); - - assert_eq!(chunks.len(), 1); - assert_eq!(errors_sr(root(&chunks[0])), None); - } - - #[test] - fn test_non_automatic_drop_priorities_are_never_rescued() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - // With the feature off, explicit negative drops and positive keeps - // are out of scope. The no-priority sentinel is not an explicit drop. - let priorities = [-1_i32, -5, 1, 2]; - let chunks = run_rescue( - &processor, - &config, - priorities - .into_iter() - .map(|p| errored_root_chunk(0x1000 + u64::from(p.unsigned_abs()), p)) - .collect(), - RESCUE_NOW, - ); - - for chunk in &chunks { - assert_eq!( - errors_sr(root(chunk)), - None, - "priority {} must not be rescued", - chunk.priority - ); - } - } - #[test] fn test_tracer_drop_empty_and_fallback_root() { let (mut config, processor) = rescue_setup(always_keep_sampler_config()); @@ -2022,97 +1968,6 @@ mod tests { ); } - #[test] - fn test_rescue_preserves_chunk_metadata_and_span_order() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - let mut root_span = test_span(0x501d, 0x502, 0, 0); - root_span - .metrics - .insert("_sampling_priority_v1".to_string(), 0.0); - root_span.metrics.insert("_sample_rate".to_string(), 0.25); - root_span - .metrics - .insert("_dd.span_sampling.rule".to_string(), 1.0); - let mut child = test_span(0x501d, 0x503, 0x502, 1); - child.meta.insert("keep".to_string(), "me".to_string()); - - let mut chunk = test_chunk(vec![root_span, child], 0); - chunk.tags.insert("_dd.p.dm".to_string(), "-4".to_string()); - chunk - .tags - .insert("origin".to_string(), "synthetics".to_string()); - chunk.origin = "synthetics".to_string(); - - let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); - - assert_eq!(chunks.len(), 1); - let rescued = &chunks[0]; - assert_eq!(rescued.priority, 0, "priority preserved"); - assert_eq!(rescued.origin, "synthetics", "origin preserved"); - assert_eq!( - rescued.tags.get("_dd.p.dm").map(String::as_str), - Some("-4"), - "decision maker preserved" - ); - assert_eq!( - rescued.tags.get("origin").map(String::as_str), - Some("synthetics") - ); - assert_eq!( - rescued.spans[0] - .metrics - .get("_dd.span_sampling.rule") - .copied(), - Some(1.0), - "single-span sampling metrics preserved" - ); - assert_eq!( - rescued.spans[0] - .metrics - .get("_sampling_priority_v1") - .copied(), - Some(0.0), - "sampling priority metric preserved" - ); - assert_eq!( - errors_sr(&rescued.spans[0]), - Some(1.0), - "root stamped despite existing metrics" - ); - assert_eq!( - rescued.spans[1].meta.get("keep").map(String::as_str), - Some("me"), - "span metadata and order preserved" - ); - assert_eq!(rescued.spans.len(), 2, "no spans added or removed"); - } - - #[test] - fn test_probabilistic_decision_maker_chunk_gets_no_workaround() { - // A priority-0 chunk carrying the chunk-level probabilistic decision - // maker (`_dd.p.dm = "-9"`) is rescued like any other automatic-drop - // chunk: priority and decision maker are preserved untouched. The - // backend resolves such chunks to the `probabilistic` ingestion reason - // before checking `_dd.errors_sr`, so it may still drop them; SCL does - // not work around that limitation. - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - let mut chunk = errored_root_chunk(0x600d, 0); - chunk.tags.insert("_dd.p.dm".to_string(), "-9".to_string()); - - let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); - - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].priority, 0, "priority preserved, no promotion"); - assert_eq!( - chunks[0].tags.get("_dd.p.dm").map(String::as_str), - Some("-9"), - "probabilistic decision maker preserved, no rewrite" - ); - assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - } - #[test] fn test_rescue_passes_raw_sample_rate_to_sampler() { // The raw `_sample_rate` wire value is passed to the shared sampler, From ce5b24c48e4a00884c61cf027ab2656fcce800b2 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 20:59:40 -0400 Subject: [PATCH 19/49] Merge empty-chunk rescue tests --- .../src/trace_processor.rs | 89 ++++++++----------- 1 file changed, 39 insertions(+), 50 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 1a77c0c9..5fd15217 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1530,57 +1530,46 @@ mod tests { } #[test] - fn test_tracer_drop_empty_and_fallback_root() { - let (mut config, processor) = rescue_setup(always_keep_sampler_config()); - config.error_rare_sample_tracer_drop = true; - let chunks = run_rescue( - &processor, - &config, - vec![ - test_chunk(vec![], -1), - test_chunk(vec![test_span(1, 2, 3, 1), test_span(1, 3, 2, 0)], -1), - ], - RESCUE_NOW, - ); - assert_eq!(chunks.len(), 2); - assert!(chunks[0].spans.is_empty()); - assert_eq!(errors_sr(&chunks[1].spans[0]), None); - assert_eq!(errors_sr(&chunks[1].spans[1]), Some(1.0)); - assert_eq!(chunks[1].priority, -1); - } - - #[test] - fn test_empty_chunk_does_not_panic_and_stays_unchanged() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - // An empty chunk cannot be scored: root resolution fails and it must be - // left unchanged without panicking. For non-empty chunks, - // `get_root_span_index` always resolves (falling back to the last span - // when no span has parent_id 0), so the cyclic chunk below exercises - // the fallback path and gets stamped on the resolved fallback root. - let cyclic = test_chunk( - vec![ - test_span(0xd00d, 0xd01, 0xd02, 1), - test_span(0xd00d, 0xd02, 0xd01, 0), - ], - 0, - ); - - let chunks = run_rescue( - &processor, - &config, - vec![test_chunk(vec![], 0), cyclic], - RESCUE_NOW, - ); + fn test_empty_chunk_and_fallback_root() { + for (feature, priority) in [(false, 0), (true, -1)] { + let (mut config, processor) = rescue_setup(always_keep_sampler_config()); + config.error_rare_sample_tracer_drop = feature; + let context = format!("feature={feature} priority={priority}"); + + // Eligible empty chunks return at the no-error guard before root resolution. + // The cyclic chunk has no parent_id 0, so the last span is the fallback root. + let cyclic = test_chunk( + vec![ + test_span(0xd00d, 0xd01, 0xd02, 1), + test_span(0xd00d, 0xd02, 0xd01, 0), + ], + priority, + ); + let chunks = run_rescue( + &processor, + &config, + vec![test_chunk(vec![], priority), cyclic], + RESCUE_NOW, + ); - assert_eq!(chunks.len(), 2); - assert!(chunks[0].spans.is_empty(), "empty chunk unchanged"); - assert_eq!( - errors_sr(&chunks[1].spans[1]), - Some(1.0), - "fallback root (last span) stamped" - ); - assert_eq!(errors_sr(&chunks[1].spans[0]), None); + assert_eq!(chunks.len(), 2, "{context}"); + assert!( + chunks[0].spans.is_empty(), + "empty chunk unchanged: {context}" + ); + assert_eq!(chunks[0].priority, priority, "{context}"); + assert_eq!(chunks[1].priority, priority, "{context}"); + assert_eq!( + errors_sr(&chunks[1].spans[0]), + None, + "non-root span not stamped: {context}" + ); + assert_eq!( + errors_sr(&chunks[1].spans[1]), + Some(1.0), + "fallback root (last span) stamped: {context}" + ); + } } #[tokio::test] From 8fb7f79d8e1d127ab9e04a936fce2b1a04a4c5db Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 21:19:17 -0400 Subject: [PATCH 20/49] Consolidate error sampler adapter tests --- .../src/trace_processor.rs | 152 +++++++----------- 1 file changed, 55 insertions(+), 97 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 5fd15217..8347073c 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1676,91 +1676,67 @@ mod tests { } } - #[test] - fn test_always_keep_rescues_every_eligible_chunk() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - let chunks = run_rescue( - &processor, - &config, - (0..10_u64) - .map(|i| errored_root_chunk(0xf000 + i, 0)) - .collect(), - RESCUE_NOW, - ); - - assert_eq!(chunks.len(), 10); - for chunk in &chunks { - assert_eq!(errors_sr(root(chunk)), Some(1.0)); - assert_eq!(chunk.priority, 0); - } - } - #[test] fn test_rate_limited_under_sustained_load_keeps_and_rejects() { - let sampler_config = rate_limited_sampler_config(10.0); - let (config, processor) = rescue_setup(sampler_config); - - // 100 chunks with the same signature in one 5-second bucket: the - // default rate is 10 / (100 / 5) = 0.5, so a deterministic mix of keeps - // and rejections is expected. No exact keep count is asserted: the - // sampler is an adaptive rolling-window sampler, not a hard token - // bucket. - let ids: Vec = (0..100_u64).map(|i| 0x11_0000 + i).collect(); - let input: Vec = ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); - let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); - - let keeps = rescued_count(&chunks); - let drops = chunks.len() - keeps; - assert!(keeps > 0, "expected some keeps, got none"); - assert!(drops > 0, "expected some drops, got none"); - assert_eq!(chunks.len(), 100, "every chunk is still forwarded"); - for chunk in &chunks { - assert_eq!(chunk.priority, 0, "priority unchanged on rescue decision"); - } + for sample_rate in [None, Some(0.25_f64)] { + let sampler_config = rate_limited_sampler_config(10.0); + let (config, processor) = rescue_setup(sampler_config); + let context = format!("sample_rate={sample_rate:?}"); + + // With the default client rate, 100 chunks with the same signature + // in one 5-second bucket give a default rate of 10 / (100 / 5) = 0.5. + // A finite client rate affects both weighting and applied probability. + // No exact keep count is asserted: this is an adaptive rolling-window + // sampler, not a hard token bucket. + let ids: Vec = (0..100_u64).map(|i| 0x11_0000 + i).collect(); + let input: Vec = ids + .iter() + .map(|id| { + let mut chunk = errored_root_chunk(*id, 0); + if let Some(rate) = sample_rate { + chunk.spans[0].metrics.insert("_sample_rate".into(), rate); + } + chunk + }) + .collect(); + let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); + + let keeps = rescued_count(&chunks); + let drops = chunks.len() - keeps; + if sample_rate.is_none() { + assert!(keeps > 0, "{context}: expected some keeps, got none"); + } + assert!(drops > 0, "{context}: expected some drops, got none"); + assert_eq!( + chunks.len(), + 100, + "{context}: every chunk is still forwarded" + ); - // The same input through a standalone shared sampler must produce the - // identical decision sequence, validating the adapter wiring (env, - // sample rate, views) against direct crate usage. - let mut standalone = ErrorsSampler::new(sampler_config); - for (chunk, id) in chunks.iter().zip(&ids) { - let views: Vec = chunk.spans.iter().map(super::span_view).collect(); - let trace = equivalent_trace_view(chunk, &config.env, &views); - let expected = standalone.sample(RESCUE_NOW, &trace); - let actual = match errors_sr(root(chunk)) { - Some(errors_sr) => SampleDecision::Keep { errors_sr }, - None => SampleDecision::Drop, - }; - assert_eq!(actual, expected, "decision mismatch for trace id {id}"); + // Compare every decision against direct crate usage. Set the expected + // client rate from the input case, independently of the output metric. + let mut standalone = ErrorsSampler::new(sampler_config); + for (chunk, id) in chunks.iter().zip(&ids) { + assert_eq!( + chunk.priority, 0, + "{context}: priority changed for trace id {id}" + ); + let views: Vec = chunk.spans.iter().map(super::span_view).collect(); + let mut trace = equivalent_trace_view(chunk, &config.env, &views); + trace.root_global_sample_rate = sample_rate.unwrap_or(1.0); + let expected = standalone.sample(RESCUE_NOW, &trace); + let actual = match errors_sr(root(chunk)) { + Some(errors_sr) => SampleDecision::Keep { errors_sr }, + None => SampleDecision::Drop, + }; + assert_eq!( + actual, expected, + "{context}: decision mismatch for trace id {id}" + ); + } } } - #[test] - fn test_rate_limited_bucket_transitions_and_steady_state() { - let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); - - // First bucket: 100 distinct signatures push the default rate to 0.5. - let ids: Vec = (0..100_u64).map(|i| 0x22_0000 + i).collect(); - let make_chunks = || ids.iter().map(|id| errored_root_chunk(*id, 0)).collect(); - let first_chunks = run_rescue(&processor, &config, make_chunks(), RESCUE_NOW); - let first_keeps = rescued_count(&first_chunks); - assert!( - first_keeps > 0 && first_keeps < 100, - "mixed decisions in the first bucket" - ); - - // A full window later (the rolling window is 6 buckets of 5 seconds), - // the same IDs in a fresh bucket still produce mixed decisions, and - // every chunk is still forwarded. - let second_chunks = run_rescue(&processor, &config, make_chunks(), RESCUE_NOW + 40); - assert_eq!(second_chunks.len(), 100); - let second_keeps = rescued_count(&second_chunks); - assert!( - second_keeps > 0 && second_keeps < 100, - "mixed decisions after window rotation" - ); - } - #[test] fn test_processor_clones_share_one_budget() { let sampler_config = rate_limited_sampler_config(1.0); @@ -1957,24 +1933,6 @@ mod tests { ); } - #[test] - fn test_rescue_passes_raw_sample_rate_to_sampler() { - // The raw `_sample_rate` wire value is passed to the shared sampler, - // which sanitizes it. A value outside (0, 1] falls back to 1.0, so a - // bogus rate does not change the stamped rescue rate. - let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); - - let mut chunk = errored_root_chunk(0x77_00, 0); - chunk.spans[0] - .metrics - .insert("_sample_rate".to_string(), f64::NAN); - - let chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); - - // Single signature, well under budget: sanitized rate 1.0 keeps it. - assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - } - #[tokio::test] async fn test_stats_concentrator_observes_pre_rescue_chunks() { let mut config = rescue_config(always_keep_sampler_config()); From b9218ca459f3a7dd33d92f5cc178abc1fd907140 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 21:35:31 -0400 Subject: [PATCH 21/49] Simplify shared error sampler budget test --- .../src/trace_processor.rs | 27 +++++-------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 8347073c..3190a47a 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1750,6 +1750,7 @@ mod tests { // a standalone sampler fed the same sequence. let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); let mut observed = Vec::new(); + let mut standalone = ErrorsSampler::new(sampler_config); for (i, id) in ids.iter().enumerate() { let target = if i < 5 { &processor } else { &processor_clone }; let mut manual = errored_root_chunk(*id, -1); @@ -1772,34 +1773,18 @@ mod tests { vec![errored_root_chunk(*id, priority)], RESCUE_NOW, ); - observed.push(errors_sr(root(&chunks[0])).is_some()); - } - - let mut standalone = ErrorsSampler::new(sampler_config); - let spans = [SpanView { - service: "test-service", - name: "test-operation", - resource: "GET /test", - error: true, - http_status_code: None, - error_type: None, - }]; - for (i, id) in ids.iter().enumerate() { - let trace = TraceView { - env: &config.env, - trace_id: *id, - root_index: 0, - root_global_sample_rate: 1.0, - spans: &spans, - }; + let actual = errors_sr(root(&chunks[0])).is_some(); + let views: Vec = chunks[0].spans.iter().map(super::span_view).collect(); + let trace = equivalent_trace_view(&chunks[0], &config.env, &views); let expected = matches!( standalone.sample(RESCUE_NOW, &trace), SampleDecision::Keep { .. } ); assert_eq!( - observed[i], expected, + actual, expected, "clone {i} decision diverged from the shared-budget standalone sampler" ); + observed.push(actual); } assert!( observed.iter().any(|kept| *kept), From c707b2bc1fca395d5ba1c1a0a890daecfeb6270d Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 27 Oct 2026 21:51:22 -0400 Subject: [PATCH 22/49] Simplify error sampler config tests --- crates/datadog-trace-agent/src/config.rs | 155 ++++++++++++----------- 1 file changed, 79 insertions(+), 76 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index e2bc9a15..9cd9794b 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -22,7 +22,7 @@ const DEFAULT_APM_RECEIVER_PORT: u16 = 8126; /// Default error rescue budget: error traces per second, matching the Go agent's `ErrorTPS`. const DEFAULT_ERROR_SAMPLER_TPS: f64 = 10.0; -/// Parses the rate-limited error rescue budget from the environment. +/// Parses the rate-limited error rescue budget from an optional raw TPS value. /// /// `DD_APM_ERROR_TPS` sets the target error traces per second, default 10. Parsed /// as `f64`; non-finite or unparsable values warn and fall back to the default. @@ -30,14 +30,14 @@ const DEFAULT_ERROR_SAMPLER_TPS: f64 = 10.0; /// /// The strategy is fixed to `RateLimited`, and `extra_sample_rate` is fixed at /// 1.0. Startup stays operational when the optional TPS setting is malformed. -fn parse_error_sampler_config() -> ErrorSamplerConfig { +fn parse_error_sampler_config(raw: Option<&str>) -> ErrorSamplerConfig { let mut config = ErrorSamplerConfig { mode: ErrorSamplerMode::RateLimited, target_tps: DEFAULT_ERROR_SAMPLER_TPS, extra_sample_rate: 1.0, }; - if let Ok(raw) = env::var("DD_APM_ERROR_TPS") { + if let Some(raw) = raw { match raw.trim().parse::() { Ok(tps) if tps.is_finite() => config.target_tps = tps, _ => { @@ -357,7 +357,7 @@ impl Config { agent_stats_computation_enabled: env::var("DD_AGENT_STATS_COMPUTATION_ENABLED") .map(|val| val.to_lowercase() == "true") .unwrap_or(true), - error_sampler: parse_error_sampler_config(), + error_sampler: parse_error_sampler_config(env::var("DD_APM_ERROR_TPS").ok().as_deref()), error_rare_sample_tracer_drop: env::var("DD_APM_FEATURES") .is_ok_and(|raw| error_rare_sample_tracer_drop_enabled(&raw)), }) @@ -1066,21 +1066,33 @@ mod tests { #[test] #[serial] - fn test_error_sampler_defaults() { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", None), - ], - || { - let config = config::Config::new().unwrap(); - assert_rate_limited_default(&config.error_sampler); - }, - ); + fn test_error_sampler_env_is_wired() { + for (raw, expected_tps) in [ + (None, 10.0), + (Some("5.5"), 5.5), + (Some("0"), 0.0), + (Some("-1"), -1.0), + ] { + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", None), + ("DD_APM_ERROR_TPS", raw), + ], + || { + let config = config::Config::new().unwrap(); + assert_eq!( + config.error_sampler.target_tps, expected_tps, + "DD_APM_ERROR_TPS {raw:?} should reach the config" + ); + assert_eq!(config.error_sampler.mode, ErrorSamplerMode::RateLimited); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); + }, + ); + } } #[test] @@ -1097,41 +1109,6 @@ mod tests { "", "always-keep", ] { - for (tps, expected_tps) in [ - (None, 10.0), - (Some("5.5"), 5.5), - (Some("0"), 0.0), - (Some("-1"), -1.0), - ] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", tps), - ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_sampler.mode, - ErrorSamplerMode::RateLimited, - "mode setting {raw:?} must not change the strategy with TPS {tps:?}" - ); - assert_eq!(config.error_sampler.target_tps, expected_tps); - assert_eq!(config.error_sampler.extra_sample_rate, 1.0); - }, - ); - } - } - } - - #[test] - #[serial] - fn test_error_sampler_tps_parsing() { - for (raw, expected_tps) in [("5.5", 5.5), ("0", 0.0), ("-1", -1.0), (" 10 ", 10.0)] { temp_env::with_vars( [ ("DD_API_KEY", Some("_not_a_real_key_")), @@ -1139,42 +1116,68 @@ mod tests { ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), ("WEBSITE_SITE_NAME", Some("my-azure-function")), ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", Some(raw)), + ("DD_APM_ERROR_TPS", Some("5.5")), + ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), ], || { let config = config::Config::new().unwrap(); assert_eq!( - config.error_sampler.target_tps, expected_tps, - "DD_APM_ERROR_TPS {raw:?} should parse to {expected_tps}" + config.error_sampler.mode, + ErrorSamplerMode::RateLimited, + "mode setting {raw:?} must not change the strategy" ); + assert_eq!(config.error_sampler.target_tps, 5.5); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); }, ); } } + #[test] + fn test_error_sampler_tps_parsing() { + for (raw, expected_tps) in [ + (None, 10.0), + (Some("5.5"), 5.5), + (Some("0"), 0.0), + (Some("-1"), -1.0), + (Some(" 10 "), 10.0), + (Some("not_a_number"), 10.0), + (Some("inf"), 10.0), + (Some("-inf"), 10.0), + (Some("NaN"), 10.0), + (Some("1e400"), 10.0), + (Some(""), 10.0), + ] { + let error_sampler = super::parse_error_sampler_config(raw); + assert_eq!( + error_sampler.target_tps, expected_tps, + "DD_APM_ERROR_TPS {raw:?} should parse to {expected_tps}" + ); + assert_eq!(error_sampler.mode, ErrorSamplerMode::RateLimited); + assert_eq!(error_sampler.extra_sample_rate, 1.0); + } + } + #[test] #[serial] #[tracing_test::traced_test] fn test_error_sampler_invalid_tps_falls_back_to_default() { - for raw in ["not_a_number", "inf", "-inf", "NaN", "1e400", ""] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), - ("DD_APM_ERROR_TPS", Some(raw)), - ], - || { - let config = config::Config::new().unwrap(); - assert_rate_limited_default(&config.error_sampler); - let _ = crate::trace_processor::ServerlessTraceProcessor::new_error_sampler( - &config, - ); - }, - ); - } + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), + ("DD_APM_ERROR_TPS", Some("not_a_number")), + ], + || { + let config = config::Config::new().unwrap(); + assert_rate_limited_default(&config.error_sampler); + let _ = + crate::trace_processor::ServerlessTraceProcessor::new_error_sampler(&config); + }, + ); logs_assert(|lines| { if lines .iter() From d02d28661b2b5980b3e81979429aa3d36deb10cf Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 09:19:17 -0400 Subject: [PATCH 23/49] Simplify error rescue processing and tests --- .../src/trace_processor.rs | 3 +- .../tests/integration_test.rs | 37 ++++++++----------- 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 3190a47a..e44dd4f9 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -241,7 +241,7 @@ impl ServerlessTraceProcessor { // backward wall-clock step cannot move the sampler's rolling window // backwards, which would undercount TPS and rescue too many chunks. let now_unix_secs = state.clamp_timestamp(read_clock()); - self.rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); + Self::rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); } /// Test-only variant that injects a synthetic clock reading through the @@ -257,7 +257,6 @@ impl ServerlessTraceProcessor { } fn rescue_with_sampler( - &self, payload: &mut TracerPayloadCollection, config: &Config, now_unix_secs: i64, diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 873437ac..fc0b89fd 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1445,34 +1445,25 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { } } - let names: Vec<_> = error_rescue_cases() - .into_iter() - .map(|case| case.0) - .collect(); - for name in names.iter().copied() { - assert!( - chunks_by_root_name.contains_key(name), - "expected chunk {name} to be forwarded to the backend, got: {:?}", - chunks_by_root_name.keys().collect::>() - ); - } - - let errors_sr_of = |name: &str| -> Option { - chunks_by_root_name[name] + let cases = error_rescue_cases(); + for &(name, priority, error, dm, on_child) in &cases { + let chunk = chunks_by_root_name.get(name).unwrap_or_else(|| { + panic!( + "expected chunk {name} to be forwarded to the backend, got: {:?}", + chunks_by_root_name.keys().collect::>() + ) + }); + let errors_sr = chunk .spans .iter() .find(|s| s.parent_id == 0) - .and_then(|root| root.metrics.get("_dd.errors_sr").copied()) - }; - - for (name, priority, error, dm, on_child) in error_rescue_cases() { - let chunk = &chunks_by_root_name[name]; + .and_then(|root| root.metrics.get("_dd.errors_sr").copied()); let eligible = error != 0 && (priority == 0 || priority == -128 || (feature && priority < 0 && !(priority == -1 && dm == Some("-4")))); assert_eq!( - errors_sr_of(name), + errors_sr, eligible.then_some(1.0), "{name}, feature={feature}" ); @@ -1512,7 +1503,11 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { }) .collect(); - for name in names.into_iter().chain(["manual_child_span"]) { + for name in cases + .into_iter() + .map(|case| case.0) + .chain(["manual_child_span"]) + { let group = all_groups .iter() .find(|g| g.name == name) From eb5a37b2e660f1cc6fe78b51112e30ace29f85e7 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 11:53:58 -0400 Subject: [PATCH 24/49] Skip error sampler lock without rescue candidates --- .../src/trace_processor.rs | 214 +++++++++++------- 1 file changed, 138 insertions(+), 76 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index e44dd4f9..4ffd4a15 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -230,18 +230,63 @@ impl ServerlessTraceProcessor { config: &Config, read_clock: F, ) { - let mut state = self.lock_sampler(); - // Skip all view construction when the sampler is disabled by config - // (target_tps <= 0): nothing can be rescued. - if state.sampler.is_disabled() { + let TracerPayloadCollection::V07(tracer_payloads) = payload else { + return; + }; + // Find candidates and build their sampler views before taking the shared + // lock, so requests without eligible errored chunks never contend on it + // and the lock is held only for the sampling decisions. + let candidates = collect_rescue_candidates(tracer_payloads, config); + if candidates.is_empty() { return; } - // Read the clock while holding the sampler lock so that concurrent - // requests deliver timestamps in lock-acquisition order, and clamp so a - // backward wall-clock step cannot move the sampler's rolling window - // backwards, which would undercount TPS and rescue too many chunks. - let now_unix_secs = state.clamp_timestamp(read_clock()); - Self::rescue_with_sampler(payload, config, now_unix_secs, &mut state.sampler); + let keeps: Vec<(usize, usize, usize, f64)> = { + let mut state = self.lock_sampler(); + // Nothing can be rescued when the sampler is disabled by config + // (target_tps <= 0). + if state.sampler.is_disabled() { + return; + } + // Read the clock while holding the sampler lock so that concurrent + // requests deliver timestamps in lock-acquisition order, and clamp so a + // backward wall-clock step cannot move the sampler's rolling window + // backwards, which would undercount TPS and rescue too many chunks. + let now_unix_secs = state.clamp_timestamp(read_clock()); + candidates + .iter() + .filter_map(|candidate| { + let trace = TraceView { + env: candidate.env, + trace_id: candidate.trace_id, + root_index: candidate.root_index, + root_global_sample_rate: candidate.root_global_sample_rate, + spans: &candidate.spans, + }; + match state.sampler.sample(now_unix_secs, &trace) { + SampleDecision::Keep { errors_sr } => Some(( + candidate.payload_index, + candidate.chunk_index, + candidate.root_index, + errors_sr, + )), + SampleDecision::Drop => None, + } + }) + .collect() + }; + // On a keep, stamp the sampler's `_dd.errors_sr` on the chunk's root span. + // Chunks are never removed here: on a Drop decision the unrescued chunk is + // forwarded without rescue mutation. + for (payload_index, chunk_index, root_index, errors_sr) in keeps { + if let Some(root) = tracer_payloads + .get_mut(payload_index) + .and_then(|tp| tp.chunks.get_mut(chunk_index)) + .and_then(|chunk| chunk.spans.get_mut(root_index)) + { + root.metrics + .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); + } + } } /// Test-only variant that injects a synthetic clock reading through the @@ -256,32 +301,6 @@ impl ServerlessTraceProcessor { self.apply_error_rescue_with_clock(payload, config, || now_unix_secs); } - fn rescue_with_sampler( - payload: &mut TracerPayloadCollection, - config: &Config, - now_unix_secs: i64, - sampler: &mut ErrorsSampler, - ) { - let TracerPayloadCollection::V07(tracer_payloads) = payload else { - return; - }; - for tracer_payload in tracer_payloads.iter_mut() { - // The sampler keys its per-signature rate limits on the env the - // tracer reported for this payload, falling back to the agent's - // configured env, consistent with how stats are flushed. - let env = resolve_payload_env(&tracer_payload.env, &config.env); - for chunk in tracer_payload.chunks.iter_mut() { - sample_and_stamp( - sampler, - chunk, - env, - now_unix_secs, - config.error_rare_sample_tracer_drop, - ); - } - } - } - fn send_to_concentrator( concentrator: &StatsConcentratorHandle, payload: &TracerPayloadCollection, @@ -352,8 +371,53 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { } } -/// Consults the error sampler for one chunk and, on a keep, stamps the -/// sampler's `_dd.errors_sr` value on the chunk's root span. +/// A chunk selected for an error sampler decision, with the sampler's +/// read-only view of its spans. +struct RescueCandidate<'a> { + payload_index: usize, + chunk_index: usize, + root_index: usize, + env: &'a str, + trace_id: u64, + root_global_sample_rate: f64, + spans: Vec>, +} + +/// Collects the chunks that need an error sampler decision, in payload order. +fn collect_rescue_candidates<'a>( + tracer_payloads: &'a [pb::TracerPayload], + config: &'a Config, +) -> Vec> { + let mut candidates = Vec::new(); + for (payload_index, tracer_payload) in tracer_payloads.iter().enumerate() { + // The sampler keys its per-signature rate limits on the env the + // tracer reported for this payload, falling back to the agent's + // configured env, consistent with how stats are flushed. + let env = resolve_payload_env(&tracer_payload.env, &config.env); + for (chunk_index, chunk) in tracer_payload.chunks.iter().enumerate() { + let Some(root_index) = rescue_root_index(chunk, config.error_rare_sample_tracer_drop) + else { + continue; + }; + let root = &chunk.spans[root_index]; + candidates.push(RescueCandidate { + payload_index, + chunk_index, + root_index, + env, + trace_id: root.trace_id, + // The raw `_sample_rate` wire value is passed through: the shared + // sampler sanitizes non-finite or out-of-range rates to 1.0 itself. + root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), + spans: chunk.spans.iter().map(span_view).collect(), + }); + } + } + candidates +} + +/// Returns the root span index of a chunk that needs an error sampler decision, +/// or `None` when the chunk is not a rescue candidate. /// /// P0 and no-priority (-128) chunks are candidates. No priority means a missing /// sampling decision, not an explicit drop, matching Go's baseline eligibility. @@ -364,15 +428,8 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { /// error metadata alone does not qualify. The root is resolved with /// `get_root_span_index`: empty chunks are left unchanged, and non-empty /// chunks always resolve a root, falling back to the last span when no span -/// has `parent_id 0` (e.g. a cyclic chunk). Chunks are never removed here: on -/// a Drop decision the unrescued chunk is forwarded without rescue mutation. -fn sample_and_stamp( - sampler: &mut ErrorsSampler, - chunk: &mut pb::TraceChunk, - env: &str, - now_unix_secs: i64, - tracer_drop_enabled: bool, -) { +/// has `parent_id 0` (e.g. a cyclic chunk). +fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Option { let eligible = chunk.priority == 0 || chunk.priority == i32::from(i8::MIN) || (tracer_drop_enabled @@ -381,41 +438,19 @@ fn sample_and_stamp( && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) == Some(MANUAL_DECISION_MAKER))); if !eligible { - return; + return None; } // An error anywhere in the chunk makes it a rescue candidate, not just an // error on the root span. Checked before the root-span search because // non-errored chunks are the common case on this path. if !chunk.spans.iter().any(|span| span.error != 0) { - return; - } - let Ok(root_index) = trace_utils::get_root_span_index(&chunk.spans) else { - // No identifiable root span: leave the chunk unchanged rather than - // guessing an index. - return; - }; - - let Some(root) = chunk.spans.get(root_index) else { - return; - }; - let views: Vec = chunk.spans.iter().map(span_view).collect(); - let trace = TraceView { - env, - trace_id: root.trace_id, - root_index, - // The raw `_sample_rate` wire value is passed through: the shared - // sampler sanitizes non-finite or out-of-range rates to 1.0 itself. - root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), - spans: &views, - }; - let decision = sampler.sample(now_unix_secs, &trace); - - if let SampleDecision::Keep { errors_sr } = decision - && let Some(root) = chunk.spans.get_mut(root_index) - { - root.metrics - .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); + return None; } + // No identifiable root span: leave the chunk unchanged rather than + // guessing an index. + trace_utils::get_root_span_index(&chunk.spans) + .ok() + .filter(|&root_index| root_index < chunk.spans.len()) } #[async_trait] @@ -1841,6 +1876,33 @@ mod tests { assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); } + #[test] + fn test_error_rescue_without_candidates_skips_sampler() { + let (config, processor) = rescue_setup(always_keep_sampler_config()); + let original = pb::TracerPayload { + chunks: vec![ + // Kept priority: not eligible. + errored_root_chunk(0x10000, 1), + // Eligible priority but no errored span. + test_chunk(vec![test_span(0x20000, 0x20001, 0, 0)], 0), + // Negative priority with the tracer-drop feature disabled. + errored_root_chunk(0x30000, -1), + ], + ..Default::default() + }; + let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); + + processor.apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("a payload without rescue candidates must not consult the sampler"); + }); + + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + let TracerPayloadCollection::V07(payloads) = payload else { + panic!("rescue must preserve the payload format"); + }; + assert_eq!(payloads, vec![original]); + } + #[test] fn test_disabled_error_rescue_does_not_read_clock() { for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { From 650b34d477241f6a1a720a7df1dc34b62afa35ce Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 15:44:32 -0400 Subject: [PATCH 25/49] Focus error rescue tests on distinct guarantees --- .../src/trace_processor.rs | 234 ++++++++---------- .../tests/integration_test.rs | 19 +- 2 files changed, 110 insertions(+), 143 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 4ffd4a15..08a859c6 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1379,73 +1379,82 @@ mod tests { #[test] fn test_tracer_drop_rescue_eligibility_matrix() { - for feature in [false, true] { - for priority in [0, -1, -2, -127, -5, i32::from(i8::MIN), 1, 2] { - for dm in [ - None, - Some(""), - Some("-4"), - Some("-3"), - Some("-9"), - Some("-40"), - ] { - for error in [0, 1, -1] { - let (mut config, processor) = rescue_setup(always_keep_sampler_config()); - config.error_rare_sample_tracer_drop = feature; - let mut chunk = test_chunk( - vec![test_span(1, 2, 0, 0), test_span(1, 3, 2, error)], - priority, - ); - // Span DM deliberately disagrees: rescue must read only chunk tags. - chunk.spans[0] - .meta - .insert(super::DECISION_MAKER_KEY.into(), "-4".into()); - chunk.spans[0] - .meta - .insert("http.status_code".into(), "500".into()); - chunk.spans[0] - .meta - .insert("error.type".into(), "Error".into()); - chunk.spans[0] - .metrics - .insert("_sampling_priority_v1".into(), f64::from(priority)); - chunk.spans[0].metrics.insert("_sample_rate".into(), 0.25); - chunk.spans[0] - .metrics - .insert("_dd.span_sampling.rule".into(), 1.0); - chunk.spans[1].meta.insert("keep".into(), "me".into()); - chunk.origin = "synthetics".into(); - chunk.tags.insert("origin".into(), "synthetics".into()); - if let Some(dm) = dm { - chunk - .tags - .insert(super::DECISION_MAKER_KEY.into(), dm.into()); - } - let original = chunk.clone(); - let mut chunks = run_rescue(&processor, &config, vec![chunk], RESCUE_NOW); - assert_eq!(chunks.len(), 1, "chunk must still be forwarded"); - let expected = error != 0 - && (priority == 0 - || priority == -128 - || (feature - && priority < 0 - && !(priority == -1 && dm == Some("-4")))); - assert_eq!( - errors_sr(&chunks[0].spans[0]), - expected.then_some(1.0), - "feature={feature} priority={priority} dm={dm:?} error={error}" - ); - assert_eq!(errors_sr(&chunks[0].spans[1]), None); - chunks[0].spans[0] - .metrics - .remove(super::ERRORS_SR_METRIC_KEY); - assert_eq!(chunks[0], original); + // Expected eligibility with the feature off and on, respectively. + for (priority, dm, expected) in [ + (0, None, [true, true]), + (0, Some("-4"), [true, true]), + (-128, None, [true, true]), + (-128, Some("-4"), [true, true]), + (-1, None, [false, true]), + (-1, Some(""), [false, true]), + (-1, Some("-4"), [false, false]), + (-1, Some("-3"), [false, true]), + (-1, Some("-9"), [false, true]), + (-1, Some("-40"), [false, true]), + (-2, Some("-4"), [false, true]), + (-127, None, [false, true]), + (-5, None, [false, true]), + (1, None, [false, false]), + (2, None, [false, false]), + ] { + for (feature, eligible) in [false, true].into_iter().zip(expected) { + for (error, has_error) in [(0, false), (1, true), (-1, true)] { + let mut chunk = test_chunk( + vec![test_span(1, 2, 0, 0), test_span(1, 3, 2, error)], + priority, + ); + // Span metadata alone must not determine eligibility. + chunk.spans[0] + .meta + .insert(super::DECISION_MAKER_KEY.into(), "-4".into()); + chunk.spans[0] + .meta + .insert("http.status_code".into(), "500".into()); + chunk.spans[0] + .meta + .insert("error.type".into(), "Error".into()); + if let Some(dm) = dm { + chunk + .tags + .insert(super::DECISION_MAKER_KEY.into(), dm.into()); } + assert_eq!( + super::rescue_root_index(&chunk, feature), + (eligible && has_error).then_some(0), + "feature={feature} priority={priority} dm={dm:?} error={error}" + ); } } } } + #[test] + fn test_rescue_preserves_chunks_except_root_metric() { + let (config, processor) = rescue_setup(always_keep_sampler_config()); + let mut chunk = test_chunk(vec![test_span(1, 2, 0, 0), test_span(1, 3, 2, 1)], 0); + chunk.spans[0].metrics.extend([ + ("_sampling_priority_v1".into(), 0.0), + ("_sample_rate".into(), 0.25), + ("_dd.span_sampling.rule".into(), 1.0), + ]); + chunk.spans[0] + .meta + .insert(super::DECISION_MAKER_KEY.into(), "-3".into()); + chunk.spans[1].meta.insert("keep".into(), "me".into()); + chunk.origin = "synthetics".into(); + chunk.tags.insert("origin".into(), "synthetics".into()); + let healthy = test_chunk(vec![test_span(2, 4, 0, 0)], 0); + let original = vec![chunk, healthy]; + let mut chunks = run_rescue(&processor, &config, original.clone(), RESCUE_NOW); + assert_eq!(chunks.len(), 2, "every chunk must still be forwarded"); + assert_eq!(errors_sr(&chunks[0].spans[0]), Some(1.0)); + assert_eq!(errors_sr(&chunks[0].spans[1]), None); + chunks[0].spans[0] + .metrics + .remove(super::ERRORS_SR_METRIC_KEY); + assert_eq!(chunks, original); + } + #[test] #[tracing_test::traced_test] fn test_disabled_sampler_initialization_warning() { @@ -1608,71 +1617,37 @@ mod tests { #[tokio::test] async fn test_rescue_stamps_eligible_chunks_when_agent_stats_disabled() { - for feature in [false, true] { - let config = Arc::new(Config { - agent_stats_computation_enabled: false, - error_rare_sample_tracer_drop: feature, - ..rescue_config(always_keep_sampler_config()) - }); - let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); - let processor = trace_processor::ServerlessTraceProcessor::new( - Some(super::StatsConcentratorHandle::new(stats_tx)), - sampler_for(&config.error_sampler), - ); - for (priority, dm, rescued) in [ - (0, None, true), - (-128, None, true), - (-1, None, feature), - (-1, Some("-3"), feature), - (-1, Some("-4"), false), - (1, None, false), - (2, None, false), - ] { - let (tx, mut rx) = mpsc::channel(1); - let res = processor - .process_traces( - config.clone(), - errored_request(priority, dm), - tx, - Arc::new(create_test_metadata()), - ) - .await; - assert!(res.is_ok()); - assert!( - matches!( - stats_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - ), - "stats must not be submitted when agent stats computation is disabled" - ); - - let send_data = rx.recv().await.expect("payload forwarded"); - let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { - panic!("expected V07 payload"); - }; - let chunk = &payloads[0].chunks[0]; - assert_eq!( - errors_sr(&chunk.spans[0]), - rescued.then_some(1.0), - "feature={feature} priority={priority} dm={dm:?}" - ); - assert_eq!(chunk.priority, priority, "priority must not be promoted"); - assert_eq!( - chunk - .tags - .get(super::DECISION_MAKER_KEY) - .map(String::as_str), - dm - ); - assert_eq!( - chunk.spans[0] - .meta - .get(super::DECISION_MAKER_KEY) - .map(String::as_str), - dm - ); - } - } + let config = Arc::new(Config { + agent_stats_computation_enabled: false, + ..rescue_config(always_keep_sampler_config()) + }); + let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); + let processor = trace_processor::ServerlessTraceProcessor::new( + Some(super::StatsConcentratorHandle::new(stats_tx)), + sampler_for(&config.error_sampler), + ); + let (tx, mut rx) = mpsc::channel(1); + let res = processor + .process_traces( + config, + errored_request(0, None), + tx, + Arc::new(create_test_metadata()), + ) + .await; + assert!(res.is_ok()); + assert!( + matches!( + stats_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + ), + "stats must not be submitted when agent stats computation is disabled" + ); + let send_data = rx.recv().await.expect("payload forwarded"); + let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { + panic!("expected V07 payload"); + }; + assert_eq!(errors_sr(&payloads[0].chunks[0].spans[0]), Some(1.0)); } #[test] @@ -1989,12 +1964,7 @@ mod tests { Some(super::StatsConcentratorHandle::new(stats_tx)), sampler_for(&config.error_sampler), ); - for (priority, dm, rescued) in [ - (0, None, true), - (-1, Some("-3"), true), - (-1, Some("-4"), false), - (-128, None, true), - ] { + for (priority, dm, rescued) in [(0, None, true), (-1, Some("-4"), false)] { let (tx, mut rx) = mpsc::channel(1); let res = processor .process_traces( diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index fc0b89fd..3105962f 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1295,16 +1295,11 @@ fn error_rescue_cases() -> Vec<(&'static str, i32, i64, Option<&'static str>, bo vec![ ("rescued_error_p0", 0, 1, None, false), ("healthy_p0", 0, 0, None, false), - ("negative_absent", -1, 1, None, false), ("negative_nonmanual", -1, 1, Some("-3"), false), - ("negative_probabilistic", -1, 1, Some("-9"), false), - ("negative_empty", -1, 1, Some(""), false), - ("manual_root", -1, 1, Some("-4"), false), ("manual_child", -1, 1, Some("-4"), true), ("other_negative_manual_dm", -2, 1, Some("-4"), false), ("no_priority", -128, 1, Some("-3"), false), ("priority_1", 1, 1, None, false), - ("priority_2", 2, 1, None, false), ] } @@ -1446,7 +1441,7 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { } let cases = error_rescue_cases(); - for &(name, priority, error, dm, on_child) in &cases { + for &(name, priority, _, dm, on_child) in &cases { let chunk = chunks_by_root_name.get(name).unwrap_or_else(|| { panic!( "expected chunk {name} to be forwarded to the backend, got: {:?}", @@ -1458,13 +1453,15 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { .iter() .find(|s| s.parent_id == 0) .and_then(|root| root.metrics.get("_dd.errors_sr").copied()); - let eligible = error != 0 - && (priority == 0 - || priority == -128 - || (feature && priority < 0 && !(priority == -1 && dm == Some("-4")))); + let rescued = match name { + "rescued_error_p0" | "no_priority" => true, + "negative_nonmanual" | "other_negative_manual_dm" => feature, + "healthy_p0" | "manual_child" | "priority_1" => false, + _ => panic!("missing expected rescue outcome for {name}"), + }; assert_eq!( errors_sr, - eligible.then_some(1.0), + rescued.then_some(1.0), "{name}, feature={feature}" ); assert_eq!(chunk.priority, priority, "{name}"); From e8dd83bcb9a5fab0d40fcb5ba9a209d98f499e18 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 15:57:39 -0400 Subject: [PATCH 26/49] Simplify error rescue with lazy locking --- .../src/trace_processor.rs | 160 +++++++----------- 1 file changed, 60 insertions(+), 100 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 08a859c6..8c53bf83 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -233,58 +233,48 @@ impl ServerlessTraceProcessor { let TracerPayloadCollection::V07(tracer_payloads) = payload else { return; }; - // Find candidates and build their sampler views before taking the shared - // lock, so requests without eligible errored chunks never contend on it - // and the lock is held only for the sampling decisions. - let candidates = collect_rescue_candidates(tracer_payloads, config); - if candidates.is_empty() { - return; - } - let keeps: Vec<(usize, usize, usize, f64)> = { - let mut state = self.lock_sampler(); - // Nothing can be rescued when the sampler is disabled by config - // (target_tps <= 0). - if state.sampler.is_disabled() { - return; - } - // Read the clock while holding the sampler lock so that concurrent - // requests deliver timestamps in lock-acquisition order, and clamp so a - // backward wall-clock step cannot move the sampler's rolling window - // backwards, which would undercount TPS and rescue too many chunks. - let now_unix_secs = state.clamp_timestamp(read_clock()); - candidates - .iter() - .filter_map(|candidate| { - let trace = TraceView { - env: candidate.env, - trace_id: candidate.trace_id, - root_index: candidate.root_index, - root_global_sample_rate: candidate.root_global_sample_rate, - spans: &candidate.spans, - }; - match state.sampler.sample(now_unix_secs, &trace) { - SampleDecision::Keep { errors_sr } => Some(( - candidate.payload_index, - candidate.chunk_index, - candidate.root_index, - errors_sr, - )), - SampleDecision::Drop => None, + let mut state = None; + let mut read_clock = Some(read_clock); + for tracer_payload in tracer_payloads.iter_mut() { + let env = resolve_payload_env(&tracer_payload.env, &config.env); + for chunk in &mut tracer_payload.chunks { + let Some(root_index) = + rescue_root_index(chunk, config.error_rare_sample_tracer_drop) + else { + continue; + }; + // Requests without candidates never take the shared lock. + let state = state.get_or_insert_with(|| self.lock_sampler()); + if let Some(read_clock) = read_clock.take() { + if state.sampler.is_disabled() { + return; } - }) - .collect() - }; - // On a keep, stamp the sampler's `_dd.errors_sr` on the chunk's root span. - // Chunks are never removed here: on a Drop decision the unrescued chunk is - // forwarded without rescue mutation. - for (payload_index, chunk_index, root_index, errors_sr) in keeps { - if let Some(root) = tracer_payloads - .get_mut(payload_index) - .and_then(|tp| tp.chunks.get_mut(chunk_index)) - .and_then(|chunk| chunk.spans.get_mut(root_index)) - { - root.metrics - .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); + // Read once under the lock, and clamp backward wall-clock steps + // so the sampler's rolling window cannot move backwards. + state.clamp_timestamp(read_clock()); + } + let now_unix_secs = state.last_sampler_timestamp; + let root = &chunk.spans[root_index]; + let views: Vec = chunk.spans.iter().map(span_view).collect(); + let trace = TraceView { + env, + trace_id: root.trace_id, + root_index, + // The shared sampler sanitizes the raw wire rate. + root_global_sample_rate: root + .metrics + .get("_sample_rate") + .copied() + .unwrap_or(1.0), + spans: &views, + }; + if let SampleDecision::Keep { errors_sr } = + state.sampler.sample(now_unix_secs, &trace) + && let Some(root) = chunk.spans.get_mut(root_index) + { + root.metrics + .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); + } } } } @@ -371,51 +361,6 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { } } -/// A chunk selected for an error sampler decision, with the sampler's -/// read-only view of its spans. -struct RescueCandidate<'a> { - payload_index: usize, - chunk_index: usize, - root_index: usize, - env: &'a str, - trace_id: u64, - root_global_sample_rate: f64, - spans: Vec>, -} - -/// Collects the chunks that need an error sampler decision, in payload order. -fn collect_rescue_candidates<'a>( - tracer_payloads: &'a [pb::TracerPayload], - config: &'a Config, -) -> Vec> { - let mut candidates = Vec::new(); - for (payload_index, tracer_payload) in tracer_payloads.iter().enumerate() { - // The sampler keys its per-signature rate limits on the env the - // tracer reported for this payload, falling back to the agent's - // configured env, consistent with how stats are flushed. - let env = resolve_payload_env(&tracer_payload.env, &config.env); - for (chunk_index, chunk) in tracer_payload.chunks.iter().enumerate() { - let Some(root_index) = rescue_root_index(chunk, config.error_rare_sample_tracer_drop) - else { - continue; - }; - let root = &chunk.spans[root_index]; - candidates.push(RescueCandidate { - payload_index, - chunk_index, - root_index, - env, - trace_id: root.trace_id, - // The raw `_sample_rate` wire value is passed through: the shared - // sampler sanitizes non-finite or out-of-range rates to 1.0 itself. - root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), - spans: chunk.spans.iter().map(span_view).collect(), - }); - } - } - candidates -} - /// Returns the root span index of a chunk that needs an error sampler decision, /// or `None` when the chunk is not a rescue candidate. /// @@ -1832,10 +1777,19 @@ mod tests { #[test] fn test_error_rescue_reads_clock_under_lock() { let (config, processor) = rescue_setup(always_keep_sampler_config()); - let mut payload = TracerPayloadCollection::V07(vec![pb::TracerPayload { - chunks: vec![errored_root_chunk(0x10000, 0)], - ..Default::default() - }]); + let mut payload = TracerPayloadCollection::V07(vec![ + pb::TracerPayload { + chunks: vec![ + errored_root_chunk(0x10000, 1), + errored_root_chunk(0x10001, 0), + ], + ..Default::default() + }, + pb::TracerPayload { + chunks: vec![errored_root_chunk(0x20000, 0)], + ..Default::default() + }, + ]); let mut clock_called = false; processor.apply_error_rescue_with_clock(&mut payload, &config, || { @@ -1849,6 +1803,12 @@ mod tests { assert!(clock_called); assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); + let TracerPayloadCollection::V07(payloads) = payload else { + panic!("expected V07 payload"); + }; + assert_eq!(errors_sr(root(&payloads[0].chunks[0])), None); + assert_eq!(errors_sr(root(&payloads[0].chunks[1])), Some(1.0)); + assert_eq!(errors_sr(root(&payloads[1].chunks[0])), Some(1.0)); } #[test] From 40d313ef6c8ec307da97f7e4a6ed7dc79cf1edd6 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 16:00:42 -0400 Subject: [PATCH 27/49] Consolidate error rescue clock regression tests --- .../src/trace_processor.rs | 67 ++++++------------- 1 file changed, 21 insertions(+), 46 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 8c53bf83..c92c5a74 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1754,24 +1754,28 @@ mod tests { fn test_error_rescue_clamps_backward_timestamps_across_clones() { let (config, processor) = rescue_setup(always_keep_sampler_config()); let clone = processor.clone(); + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); - let chunks = run_rescue( - &processor, - &config, - vec![errored_root_chunk(0x10000, 0)], - RESCUE_NOW, - ); - assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); - - let chunks = run_rescue( - &clone, - &config, - vec![errored_root_chunk(0x10001, 0)], - RESCUE_NOW - 40, - ); - assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - assert_eq!(clone.lock_sampler().last_sampler_timestamp, RESCUE_NOW); + for (i, (now, expected)) in [ + (RESCUE_NOW, RESCUE_NOW), + (RESCUE_NOW - 40, RESCUE_NOW), + (RESCUE_NOW, RESCUE_NOW), + (RESCUE_NOW + 40, RESCUE_NOW + 40), + (RESCUE_NOW + 20, RESCUE_NOW + 40), + ] + .into_iter() + .enumerate() + { + let target = if i % 2 == 0 { &processor } else { &clone }; + let chunks = run_rescue( + target, + &config, + vec![errored_root_chunk(0x10000 + i as u64, 0)], + now, + ); + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + assert_eq!(target.lock_sampler().last_sampler_timestamp, expected); + } } #[test] @@ -1866,35 +1870,6 @@ mod tests { } } - #[test] - fn test_independent_processors_share_sampler_timestamp_floor() { - let sampler = sampler_for(&rate_limited_sampler_config(1.0)); - let first = trace_processor::ServerlessTraceProcessor::new(None, Arc::clone(&sampler)); - let second = trace_processor::ServerlessTraceProcessor::new(None, sampler); - - assert_eq!(first.lock_sampler().last_sampler_timestamp, i64::MIN); - assert_eq!(first.lock_sampler().clamp_timestamp(200), 200); - assert_eq!(second.lock_sampler().clamp_timestamp(50), 200); - assert_eq!(second.lock_sampler().clamp_timestamp(300), 300); - assert_eq!(first.lock_sampler().clamp_timestamp(250), 300); - } - - #[test] - fn test_sampler_timestamp_clamp_never_moves_backwards() { - let (_config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); - // A clone shares the clamp floor with the original. - let clone = processor.clone(); - - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); - assert_eq!(processor.lock_sampler().clamp_timestamp(100), 100); - // Backward timestamps, through either instance, are clamped to the - // last value seen; forward ones update the floor. - assert_eq!(clone.lock_sampler().clamp_timestamp(50), 100); - assert_eq!(processor.lock_sampler().clamp_timestamp(100), 100); - assert_eq!(clone.lock_sampler().clamp_timestamp(200), 200); - assert_eq!(processor.lock_sampler().clamp_timestamp(199), 200); - } - #[test] fn test_resolve_payload_env_prefers_payload_env() { assert_eq!( From 1c5bd08e9e6c5322e641942dae8a4e5e24704d74 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 16:03:00 -0400 Subject: [PATCH 28/49] Consolidate disabled error rescue tests --- .../src/trace_processor.rs | 56 +++++-------------- 1 file changed, 13 insertions(+), 43 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index c92c5a74..3f7cf03f 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1606,26 +1606,24 @@ mod tests { }); config.error_rare_sample_tracer_drop = true; - let chunks = run_rescue( - &processor, - &config, - vec![ + let original = pb::TracerPayload { + chunks: vec![ errored_root_chunk(0xe000, 0), errored_root_chunk(0xe001, -1), errored_root_chunk(0xe002, i32::from(i8::MIN)), ], - RESCUE_NOW, - ); + ..Default::default() + }; + let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); + processor.apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("disabled sampler must not read the clock"); + }); - assert_eq!(chunks.len(), 3, "chunks still forwarded"); - for chunk in &chunks { - assert_eq!( - errors_sr(root(chunk)), - None, - "mode {mode:?} tps {tps} priority {}: disabled sampler must not rescue", - chunk.priority - ); - } + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + let TracerPayloadCollection::V07(payloads) = payload else { + panic!("disabled rescue must preserve the payload format"); + }; + assert_eq!(payloads, vec![original], "mode={mode:?} tps={tps}"); } } } @@ -1842,34 +1840,6 @@ mod tests { assert_eq!(payloads, vec![original]); } - #[test] - fn test_disabled_error_rescue_does_not_read_clock() { - for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { - for target_tps in [0.0, -3.0] { - let (config, processor) = rescue_setup(ErrorSamplerConfig { - mode, - target_tps, - extra_sample_rate: 1.0, - }); - let original = pb::TracerPayload { - chunks: vec![errored_root_chunk(0x10000, 0)], - ..Default::default() - }; - let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); - - processor.apply_error_rescue_with_clock(&mut payload, &config, || { - panic!("disabled sampler must not read the clock"); - }); - - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); - let TracerPayloadCollection::V07(payloads) = payload else { - panic!("disabled rescue must preserve the payload format"); - }; - assert_eq!(payloads, vec![original]); - } - } - } - #[test] fn test_resolve_payload_env_prefers_payload_env() { assert_eq!( From ad7a8118509849c50e02ccb583bc53760c0e2516 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 16:04:19 -0400 Subject: [PATCH 29/49] Remove tests for unused error sampler setting --- crates/datadog-trace-agent/src/config.rs | 38 ------------------------ 1 file changed, 38 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 9cd9794b..4e027dae 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -1095,44 +1095,6 @@ mod tests { } } - #[test] - #[serial] - fn test_error_sampler_mode_env_is_ignored() { - for raw in [ - "always_keep", - "ALWAYS_KEEP", - " Always_Keep ", - "rate_limited", - "RATE_LIMITED", - " Rate_Limited ", - "bogus", - "", - "always-keep", - ] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", Some("5.5")), - ("DD_APM_ERROR_SAMPLER_MODE", Some(raw)), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_sampler.mode, - ErrorSamplerMode::RateLimited, - "mode setting {raw:?} must not change the strategy" - ); - assert_eq!(config.error_sampler.target_tps, 5.5); - assert_eq!(config.error_sampler.extra_sample_rate, 1.0); - }, - ); - } - } - #[test] fn test_error_sampler_tps_parsing() { for (raw, expected_tps) in [ From 99a3223425b02c38c123552e3e78f2a530e60d3e Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 16:14:01 -0400 Subject: [PATCH 30/49] Focus error rescue config tests on wiring --- crates/datadog-trace-agent/src/config.rs | 117 ++++------------------- 1 file changed, 17 insertions(+), 100 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 4e027dae..1b072f01 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -997,39 +997,6 @@ mod tests { } } - #[test] - #[serial] - fn test_tracer_drop_feature_does_not_enable_sampler_or_stats() { - for feature in [None, Some("error_rare_sample_tracer_drop"), Some("other")] { - for (tps, expected_tps) in [(None, 10.0), (Some("0"), 0.0), (Some("-1"), -1.0)] { - for stats in ["true", "false"] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("test")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("test")), - ("DD_APM_FEATURES", feature), - ("DD_APM_ERROR_TPS", tps), - ("DD_AGENT_STATS_COMPUTATION_ENABLED", Some(stats)), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_rare_sample_tracer_drop, - feature == Some("error_rare_sample_tracer_drop") - ); - assert_eq!(config.error_sampler.target_tps, expected_tps); - assert_eq!(config.error_sampler.extra_sample_rate, 1.0); - assert_eq!(config.error_sampler.mode, ErrorSamplerMode::RateLimited); - assert_eq!(config.agent_stats_computation_enabled, stats == "true"); - }, - ); - } - } - } - } - #[cfg(unix)] #[test] #[serial] @@ -1055,44 +1022,26 @@ mod tests { ); } - fn assert_rate_limited_default(error_sampler: &config::ErrorSamplerConfig) { - assert!( - matches!(error_sampler.mode, ErrorSamplerMode::RateLimited), - "expected default mode RateLimited" - ); - assert_eq!(error_sampler.target_tps, 10.0); - assert_eq!(error_sampler.extra_sample_rate, 1.0); - } - #[test] #[serial] fn test_error_sampler_env_is_wired() { - for (raw, expected_tps) in [ - (None, 10.0), - (Some("5.5"), 5.5), - (Some("0"), 0.0), - (Some("-1"), -1.0), - ] { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", None), - ("DD_APM_ERROR_TPS", raw), - ], - || { - let config = config::Config::new().unwrap(); - assert_eq!( - config.error_sampler.target_tps, expected_tps, - "DD_APM_ERROR_TPS {raw:?} should reach the config" - ); - assert_eq!(config.error_sampler.mode, ErrorSamplerMode::RateLimited); - assert_eq!(config.error_sampler.extra_sample_rate, 1.0); - }, - ); - } + temp_env::with_vars( + [ + ("DD_API_KEY", Some("_not_a_real_key_")), + ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), + ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), + ("WEBSITE_SITE_NAME", Some("my-azure-function")), + ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), + ("DD_APM_ERROR_TPS", Some("5.5")), + ], + || { + let config = config::Config::new().unwrap(); + assert!(config.error_rare_sample_tracer_drop); + assert_eq!(config.error_sampler.target_tps, 5.5); + assert_eq!(config.error_sampler.mode, ErrorSamplerMode::RateLimited); + assert_eq!(config.error_sampler.extra_sample_rate, 1.0); + }, + ); } #[test] @@ -1119,38 +1068,6 @@ mod tests { assert_eq!(error_sampler.extra_sample_rate, 1.0); } } - - #[test] - #[serial] - #[tracing_test::traced_test] - fn test_error_sampler_invalid_tps_falls_back_to_default() { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("_not_a_real_key_")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("my-azure-function")), - ("DD_APM_FEATURES", Some("error_rare_sample_tracer_drop")), - ("DD_APM_ERROR_TPS", Some("not_a_number")), - ], - || { - let config = config::Config::new().unwrap(); - assert_rate_limited_default(&config.error_sampler); - let _ = - crate::trace_processor::ServerlessTraceProcessor::new_error_sampler(&config); - }, - ); - logs_assert(|lines| { - if lines - .iter() - .any(|line| line.contains("DD_APM_ERROR_TPS <= 0 disables error rescue")) - { - Err("fallback TPS must not warn that rescue is disabled".into()) - } else { - Ok(()) - } - }); - } } /// Test helpers for creating Config instances in tests From 3d216c07ab7ec54e966957ddbd354f83842fa38f Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 16:23:00 -0400 Subject: [PATCH 31/49] Focus decision maker tests on precedence --- .../src/trace_processor.rs | 98 +++++++------------ 1 file changed, 37 insertions(+), 61 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 3f7cf03f..36f575f8 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1254,71 +1254,47 @@ mod tests { } #[test] - fn test_decision_maker_promotion_is_unconditional_and_preserves_metadata() { + fn test_decision_maker_promotion_precedence() { use libdd_trace_utils::tracer_payload::TraceChunkProcessor as _; - for feature in [false, true] { - for (chunk_dm, first, second, child_first, expected) in [ - (None, None, None, false, None), - (None, Some("-4"), None, false, Some("-4")), - (None, None, Some("-3"), false, Some("-3")), - (None, Some("-3"), Some("-4"), true, Some("-3")), - (None, Some("-4"), Some("-3"), false, Some("-4")), - (None, Some("-3"), Some("-4"), false, Some("-3")), - (Some("-9"), Some("-4"), Some("-3"), false, Some("-9")), - (Some(""), Some("-4"), None, false, Some("")), - (None, Some(""), Some("-4"), true, Some("")), - ] { - let mut config = create_test_config(); - config.error_rare_sample_tracer_drop = feature; - config.agent_stats_computation_enabled = false; - let mut spans = vec![test_span(1, 2, 0, 1), test_span(1, 3, 2, 0)]; - if child_first { - spans.swap(0, 1); - } - for (span, dm) in spans.iter_mut().zip([first, second]) { - span.meta.insert("keep".into(), "metadata".into()); - span.metrics.insert("_sampling_priority_v1".into(), -1.0); - if let Some(dm) = dm { - span.meta - .insert(super::DECISION_MAKER_KEY.into(), dm.into()); - } - } - let mut chunk = test_chunk(spans, -1); - chunk.origin = "synthetics".into(); - chunk.tags.insert("other".into(), "value".into()); - if let Some(dm) = chunk_dm { - chunk - .tags + for (chunk_dm, first, second, child_first, expected) in [ + (None, None, None, false, None), + (None, Some("-4"), None, false, Some("-4")), + (None, None, Some("-3"), false, Some("-3")), + (None, Some("-3"), Some("-4"), true, Some("-3")), + (None, Some("-4"), Some("-3"), false, Some("-4")), + (None, Some("-3"), Some("-4"), false, Some("-3")), + (Some("-9"), Some("-4"), Some("-3"), false, Some("-9")), + (Some(""), Some("-4"), None, false, Some("")), + (None, Some(""), Some("-4"), true, Some("")), + ] { + let mut spans = vec![test_span(1, 2, 0, 1), test_span(1, 3, 2, 0)]; + if child_first { + spans.swap(0, 1); + } + for (span, dm) in spans.iter_mut().zip([first, second]) { + if let Some(dm) = dm { + span.meta .insert(super::DECISION_MAKER_KEY.into(), dm.into()); } - let original = chunk.clone(); - let root_index = usize::from(child_first); - super::ChunkProcessor { - config: Arc::new(config), - mini_agent_metadata: Arc::new(create_test_metadata()), - } - .process(&mut chunk, root_index); - assert_eq!( - chunk - .tags - .get(super::DECISION_MAKER_KEY) - .map(String::as_str), - expected - ); - assert_eq!(chunk.tags.get("other"), original.tags.get("other")); - assert_eq!(chunk.priority, original.priority); - assert_eq!(chunk.origin, original.origin); - for (span, before) in chunk.spans.iter().zip(&original.spans) { - assert_eq!(span.span_id, before.span_id); - assert_eq!( - span.metrics.get("_sampling_priority_v1"), - before.metrics.get("_sampling_priority_v1") - ); - for (key, value) in &before.meta { - assert_eq!(span.meta.get(key), Some(value)); - } - } } + let mut chunk = test_chunk(spans, -1); + if let Some(dm) = chunk_dm { + chunk + .tags + .insert(super::DECISION_MAKER_KEY.into(), dm.into()); + } + super::ChunkProcessor { + config: Arc::new(create_test_config()), + mini_agent_metadata: Arc::new(create_test_metadata()), + } + .process(&mut chunk, usize::from(child_first)); + assert_eq!( + chunk + .tags + .get(super::DECISION_MAKER_KEY) + .map(String::as_str), + expected + ); } } From cd0e4b938ca860ed81d709679ccb3c9092c4b968 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 17:22:56 -0400 Subject: [PATCH 32/49] Simplify disabled error rescue warning test --- .../src/trace_processor.rs | 43 +++---------------- 1 file changed, 6 insertions(+), 37 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 36f575f8..d6c5b134 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1380,46 +1380,15 @@ mod tests { #[tracing_test::traced_test] fn test_disabled_sampler_initialization_warning() { let warning = "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue"; - let mut expected_warnings = 0; - for feature in [false, true] { - for tps in [0.0, -1.0, 10.0] { - let mut config = rescue_config(always_keep_sampler_config()); - config.error_rare_sample_tracer_drop = feature; - config.error_sampler.target_tps = tps; - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - trace_processor::ServerlessTraceProcessor::new_error_sampler(&config), - ); - expected_warnings += usize::from(feature && tps <= 0.0); - logs_assert(|lines| { - let count = lines.iter().filter(|line| line.contains(warning)).count(); - if count == expected_warnings { - Ok(()) - } else { - Err(format!( - "expected {expected_warnings} warnings, got {count}" - )) - } - }); - for _ in 0..3 { - run_rescue( - &processor.clone(), - &config, - vec![errored_root_chunk(1, -1)], - RESCUE_NOW, - ); - } - } - } + let mut config = rescue_config(rate_limited_sampler_config(0.0)); + config.error_rare_sample_tracer_drop = true; + let _state = trace_processor::ServerlessTraceProcessor::new_error_sampler(&config); logs_assert(|lines| { - let count = lines - .iter() - .filter(|line| line.contains("DD_APM_ERROR_TPS <= 0 disables error rescue")) - .count(); - if count == 2 { + let count = lines.iter().filter(|line| line.contains(warning)).count(); + if count == 1 { Ok(()) } else { - Err(format!("expected two initialization warnings, got {count}")) + Err(format!("expected one initialization warning, got {count}")) } }); } From 26c78cc0372ee7606bae6edf52c411074147891c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 18:09:43 -0400 Subject: [PATCH 33/49] Remove redundant rescue root bounds checks --- crates/datadog-trace-agent/src/trace_processor.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index d6c5b134..5c503701 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -270,9 +270,9 @@ impl ServerlessTraceProcessor { }; if let SampleDecision::Keep { errors_sr } = state.sampler.sample(now_unix_secs, &trace) - && let Some(root) = chunk.spans.get_mut(root_index) { - root.metrics + chunk.spans[root_index] + .metrics .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); } } @@ -393,9 +393,7 @@ fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Optio } // No identifiable root span: leave the chunk unchanged rather than // guessing an index. - trace_utils::get_root_span_index(&chunk.spans) - .ok() - .filter(|&root_index| root_index < chunk.spans.len()) + trace_utils::get_root_span_index(&chunk.spans).ok() } #[async_trait] From 04254aca9050b87cb67cf6aa631ef041659dd40e Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 18:21:26 -0400 Subject: [PATCH 34/49] Clarify shared error rescue state construction --- crates/datadog-serverless-compat/src/main.rs | 2 +- crates/datadog-trace-agent/src/trace_processor.rs | 4 ++-- crates/datadog-trace-agent/tests/integration_test.rs | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/datadog-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs index 8254a248..2ce7fa96 100644 --- a/crates/datadog-serverless-compat/src/main.rs +++ b/crates/datadog-serverless-compat/src/main.rs @@ -208,7 +208,7 @@ pub async fn main() { let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor::new( stats_concentrator.as_ref().map(|c| c.handle.clone()), - trace_processor::ServerlessTraceProcessor::new_error_sampler(&config), + trace_processor::ServerlessTraceProcessor::new_error_rescue_state(&config), )); let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher { diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 5c503701..89b9e760 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -190,7 +190,7 @@ impl ServerlessTraceProcessor { /// Builds the shared error rescue state from parsed config settings. #[must_use] - pub fn new_error_sampler(config: &Config) -> Arc> { + pub fn new_error_rescue_state(config: &Config) -> Arc> { if config.error_rare_sample_tracer_drop && config.error_sampler.target_tps <= 0.0 { warn!( "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue" @@ -1380,7 +1380,7 @@ mod tests { let warning = "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue"; let mut config = rescue_config(rate_limited_sampler_config(0.0)); config.error_rare_sample_tracer_drop = true; - let _state = trace_processor::ServerlessTraceProcessor::new_error_sampler(&config); + let _state = trace_processor::ServerlessTraceProcessor::new_error_rescue_state(&config); logs_assert(|lines| { let count = lines.iter().filter(|line| line.contains(warning)).count(); if count == 1 { diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 3105962f..255ce3b2 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -101,7 +101,7 @@ pub fn create_mini_agent_with_real_flushers( config: config.clone(), trace_processor: Arc::new(ServerlessTraceProcessor::new( Some(stats_concentrator_handle.clone()), - ServerlessTraceProcessor::new_error_sampler(&config), + ServerlessTraceProcessor::new_error_rescue_state(&config), )), trace_flusher: Arc::new(ServerlessTraceFlusher::new( aggregator.clone(), @@ -298,7 +298,7 @@ async fn test_mini_agent_tcp_handles_requests() { config: config.clone(), trace_processor: Arc::new(ServerlessTraceProcessor::new( None, - ServerlessTraceProcessor::new_error_sampler(&config), + ServerlessTraceProcessor::new_error_rescue_state(&config), )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), @@ -415,7 +415,7 @@ async fn test_mini_agent_named_pipe_handles_requests() { config: config.clone(), trace_processor: Arc::new(ServerlessTraceProcessor::new( None, - ServerlessTraceProcessor::new_error_sampler(&config), + ServerlessTraceProcessor::new_error_rescue_state(&config), )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), @@ -588,7 +588,7 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { config: config.clone(), trace_processor: Arc::new(ServerlessTraceProcessor::new( None, - ServerlessTraceProcessor::new_error_sampler(&config), + ServerlessTraceProcessor::new_error_rescue_state(&config), )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), From 0d2981da22f2c5a4eb5d2c23a7d1080d9435fb83 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 28 Oct 2026 18:56:04 -0400 Subject: [PATCH 35/49] Consolidate error rescue eligibility docs --- .../src/trace_processor.rs | 36 +++++-------------- 1 file changed, 9 insertions(+), 27 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 89b9e760..6175d422 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -207,15 +207,9 @@ impl ServerlessTraceProcessor { .unwrap_or_else(PoisonError::into_inner) } - /// Applies the error rescue pass to a payload collection. - /// - /// For every eligible V07 chunk containing a non-zero error flag, consults - /// the shared sampler. P0 and no-priority (-128) chunks are always eligible; - /// the opt-in feature adds other negative priorities except exact manual drops. - /// On a keep, stamps the sampler's `_dd.errors_sr` on the chunk's root - /// span without priority promotion. Retention is decided by the backend. - /// Chunks are never removed or reordered: unrescued chunks stay in the - /// payload for the backend to handle. + /// Samples V07 chunks selected by [`rescue_root_index`] and stamps + /// `_dd.errors_sr` on kept roots. Priorities and chunk order are unchanged; + /// every chunk is forwarded, with retention decided by the backend. fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { self.apply_error_rescue_with_clock(payload, config, || { SystemTime::now() @@ -361,19 +355,11 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { } } -/// Returns the root span index of a chunk that needs an error sampler decision, -/// or `None` when the chunk is not a rescue candidate. -/// -/// P0 and no-priority (-128) chunks are candidates. No priority means a missing -/// sampling decision, not an explicit drop, matching Go's baseline eligibility. -/// With the feature enabled, other negative priorities also qualify except -/// priority -1 with chunk DM -4. -/// Only promoted chunk tags are consulted, never span metadata. The chunk -/// must contain at least one span with a non-zero error flag; HTTP status or -/// error metadata alone does not qualify. The root is resolved with -/// `get_root_span_index`: empty chunks are left unchanged, and non-empty -/// chunks always resolve a root, falling back to the last span when no span -/// has `parent_id 0` (e.g. a cyclic chunk). +/// Returns the root index of an eligible errored chunk, otherwise `None`. +/// P0 and missing-priority (-128) chunks qualify by default. The feature adds +/// other negative priorities except manual drops (priority -1 with chunk DM -4). +/// Only chunk tags determine DM; any span with `error != 0` supplies the error. +/// Root selection delegates to `get_root_span_index`; empty chunks yield `None`. fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Option { let eligible = chunk.priority == 0 || chunk.priority == i32::from(i8::MIN) @@ -385,14 +371,10 @@ fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Optio if !eligible { return None; } - // An error anywhere in the chunk makes it a rescue candidate, not just an - // error on the root span. Checked before the root-span search because - // non-errored chunks are the common case on this path. + // Skip root resolution for healthy chunks. if !chunk.spans.iter().any(|span| span.error != 0) { return None; } - // No identifiable root span: leave the chunk unchanged rather than - // guessing an index. trace_utils::get_root_span_index(&chunk.spans).ok() } From ee070ea9ca3b5cceb257e7b35995257f13108573 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:18:54 -0400 Subject: [PATCH 36/49] Stop rescuing errored no-priority chunks Chunks without a sampling priority are kept by the backend, so they no longer use the shared error sampling budget. This matches the Go agent, where only chunks dropped by the no-priority sampler reach the error sampler. --- .../datadog-trace-agent/src/trace_processor.rs | 18 +++++++++--------- .../tests/integration_test.rs | 4 ++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 6175d422..2038688a 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -356,15 +356,15 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { } /// Returns the root index of an eligible errored chunk, otherwise `None`. -/// P0 and missing-priority (-128) chunks qualify by default. The feature adds -/// other negative priorities except manual drops (priority -1 with chunk DM -4). +/// P0 chunks qualify by default. The feature adds other negative priorities +/// except no-priority (-128) and manual drops (priority -1 with chunk DM -4). /// Only chunk tags determine DM; any span with `error != 0` supplies the error. /// Root selection delegates to `get_root_span_index`; empty chunks yield `None`. fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Option { let eligible = chunk.priority == 0 - || chunk.priority == i32::from(i8::MIN) || (tracer_drop_enabled && chunk.priority < 0 + && chunk.priority != i32::from(i8::MIN) && !(chunk.priority == -1 && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) == Some(MANUAL_DECISION_MAKER))); @@ -1284,8 +1284,8 @@ mod tests { for (priority, dm, expected) in [ (0, None, [true, true]), (0, Some("-4"), [true, true]), - (-128, None, [true, true]), - (-128, Some("-4"), [true, true]), + (-128, None, [false, false]), + (-128, Some("-4"), [false, false]), (-1, None, [false, true]), (-1, Some(""), [false, true]), (-1, Some("-4"), [false, false]), @@ -1622,9 +1622,9 @@ mod tests { let processor_clone = processor.clone(); // Ten errored chunks with the same signature but distinct IDs: five - // through the original, five through the clone. Mix P0, explicit drops, - // and no-priority chunks. All must draw from the same budget, matching - // a standalone sampler fed the same sequence. + // through the original, five through the clone. Mix P0 and explicit + // drops. All must draw from the same budget, matching a standalone + // sampler fed the same sequence. let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); let mut observed = Vec::new(); let mut standalone = ErrorsSampler::new(sampler_config); @@ -1643,7 +1643,7 @@ mod tests { RESCUE_NOW, ); assert_eq!(rescued_count(&excluded), 0); - let priority = [0, -1, i32::from(i8::MIN)][i % 3]; + let priority = if i % 2 == 0 { 0 } else { -1 }; let chunks = run_rescue( target, &config, diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 255ce3b2..9efa30d9 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1454,9 +1454,9 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { .find(|s| s.parent_id == 0) .and_then(|root| root.metrics.get("_dd.errors_sr").copied()); let rescued = match name { - "rescued_error_p0" | "no_priority" => true, + "rescued_error_p0" => true, "negative_nonmanual" | "other_negative_manual_dm" => feature, - "healthy_p0" | "manual_child" | "priority_1" => false, + "healthy_p0" | "manual_child" | "no_priority" | "priority_1" => false, _ => panic!("missing expected rescue outcome for {name}"), }; assert_eq!( From 2ee0a89aa856c6aac67fa9c31453289340a128b8 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:20:34 -0400 Subject: [PATCH 37/49] Reuse sampler defaults for error TPS config --- crates/datadog-trace-agent/src/config.rs | 25 +++++++++--------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 1b072f01..364f58c2 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -9,7 +9,7 @@ use std::env; use std::str::FromStr; use std::sync::OnceLock; -use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; +use datadog_agent_trace_sampler::ErrorSamplerConfig; use libdd_trace_obfuscation::obfuscation_config; use libdd_trace_utils::config_utils::{ read_cloud_env, trace_intake_url, trace_intake_url_prefixed, trace_stats_url, @@ -19,30 +19,23 @@ use libdd_trace_utils::trace_utils; use tracing::warn; const DEFAULT_APM_RECEIVER_PORT: u16 = 8126; -/// Default error rescue budget: error traces per second, matching the Go agent's `ErrorTPS`. -const DEFAULT_ERROR_SAMPLER_TPS: f64 = 10.0; -/// Parses the rate-limited error rescue budget from an optional raw TPS value. +/// Parses the error rescue budget from an optional raw TPS value. /// -/// `DD_APM_ERROR_TPS` sets the target error traces per second, default 10. Parsed -/// as `f64`; non-finite or unparsable values warn and fall back to the default. -/// A finite value <= 0 disables rescue. -/// -/// The strategy is fixed to `RateLimited`, and `extra_sample_rate` is fixed at -/// 1.0. Startup stays operational when the optional TPS setting is malformed. +/// Starts from `ErrorSamplerConfig::default()` (the Go agent defaults) and +/// overrides only `target_tps` with `DD_APM_ERROR_TPS`, parsed as `f64`. +/// Non-finite or unparsable values warn and keep the default. A finite +/// value <= 0 disables rescue. fn parse_error_sampler_config(raw: Option<&str>) -> ErrorSamplerConfig { - let mut config = ErrorSamplerConfig { - mode: ErrorSamplerMode::RateLimited, - target_tps: DEFAULT_ERROR_SAMPLER_TPS, - extra_sample_rate: 1.0, - }; + let mut config = ErrorSamplerConfig::default(); if let Some(raw) = raw { match raw.trim().parse::() { Ok(tps) if tps.is_finite() => config.target_tps = tps, _ => { warn!( - "Invalid DD_APM_ERROR_TPS {raw:?}; using default {DEFAULT_ERROR_SAMPLER_TPS}" + "Invalid DD_APM_ERROR_TPS {raw:?}; using default {}", + config.target_tps ); } } From 8def48afd01ab2eb030c6a9669458d7970d42c2c Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:20:54 -0400 Subject: [PATCH 38/49] Treat blank DD_APM_ERROR_TPS as unset An empty or whitespace-only value now uses the default without logging an invalid-value warning, as the Go agent does. --- crates/datadog-trace-agent/src/config.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 364f58c2..98fe4b9f 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -24,13 +24,13 @@ const DEFAULT_APM_RECEIVER_PORT: u16 = 8126; /// /// Starts from `ErrorSamplerConfig::default()` (the Go agent defaults) and /// overrides only `target_tps` with `DD_APM_ERROR_TPS`, parsed as `f64`. -/// Non-finite or unparsable values warn and keep the default. A finite -/// value <= 0 disables rescue. +/// Blank values count as unset. Non-finite or unparsable values warn and keep +/// the default. A finite value <= 0 disables rescue. fn parse_error_sampler_config(raw: Option<&str>) -> ErrorSamplerConfig { let mut config = ErrorSamplerConfig::default(); - if let Some(raw) = raw { - match raw.trim().parse::() { + if let Some(raw) = raw.map(str::trim).filter(|raw| !raw.is_empty()) { + match raw.parse::() { Ok(tps) if tps.is_finite() => config.target_tps = tps, _ => { warn!( From eab92693165db37591019c6f02cc95c7167f85de Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:21:17 -0400 Subject: [PATCH 39/49] Remove redundant unreadable DD_APM_FEATURES test --- crates/datadog-trace-agent/src/config.rs | 25 ------------------------ 1 file changed, 25 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 98fe4b9f..1befb503 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -990,31 +990,6 @@ mod tests { } } - #[cfg(unix)] - #[test] - #[serial] - fn test_unreadable_apm_features_is_disabled() { - use std::os::unix::ffi::OsStringExt; - temp_env::with_var( - "DD_APM_FEATURES", - Some(std::ffi::OsString::from_vec(vec![0xff])), - || { - temp_env::with_vars( - [ - ("DD_API_KEY", Some("test")), - ("FUNCTIONS_EXTENSION_VERSION", Some("~4")), - ("FUNCTIONS_WORKER_RUNTIME", Some("dotnet")), - ("WEBSITE_SITE_NAME", Some("test")), - ("DD_APM_ERROR_TPS", None), - ], - || { - assert!(!config::Config::new().unwrap().error_rare_sample_tracer_drop); - }, - ); - }, - ); - } - #[test] #[serial] fn test_error_sampler_env_is_wired() { From 0e987499c2f93a51258cac707b9a2f889d875167 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:24:12 -0400 Subject: [PATCH 40/49] Trim error rescue tests that re-test the sampler Drop the AlwaysKeep disabled case and the lock-ordering test, and assert literal keep decisions instead of replaying a standalone sampler. --- .../src/trace_processor.rs | 186 ++++-------------- 1 file changed, 40 insertions(+), 146 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 2038688a..e6110e63 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -566,9 +566,7 @@ mod tests { TraceProcessor, encoded_size, split_oversized_payloads, }, }; - use datadog_agent_trace_sampler::{ - ErrorSamplerConfig, ErrorSamplerMode, ErrorsSampler, SampleDecision, SpanView, TraceView, - }; + use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; use libdd_common::{Endpoint, http_common}; use libdd_trace_protobuf::pb; use libdd_trace_utils::test_utils::{create_test_gcp_json_span, create_test_gcp_span}; @@ -1187,24 +1185,6 @@ mod tests { span.metrics.get("_dd.errors_sr").copied() } - /// Builds the equivalent standalone `TraceView` for a fixture chunk, used to - /// compare adapter decisions with direct shared-sampler usage. - fn equivalent_trace_view<'a>( - chunk: &pb::TraceChunk, - env: &'a str, - views: &'a [SpanView<'a>], - ) -> TraceView<'a> { - let root_index = trace_utils::get_root_span_index(&chunk.spans).unwrap(); - let root = &chunk.spans[root_index]; - TraceView { - env, - trace_id: root.trace_id, - root_index, - root_global_sample_rate: root.metrics.get("_sample_rate").copied().unwrap_or(1.0), - spans: views, - } - } - fn rescued_count(chunks: &[pb::TraceChunk]) -> usize { chunks .iter() @@ -1521,50 +1501,40 @@ mod tests { } #[test] - fn test_disabled_tps_disables_rescue_in_both_modes() { - for mode in [ErrorSamplerMode::RateLimited, ErrorSamplerMode::AlwaysKeep] { - for tps in [0.0_f64, -3.0] { - let (mut config, processor) = rescue_setup(ErrorSamplerConfig { - mode, - target_tps: tps, - extra_sample_rate: 1.0, - }); - config.error_rare_sample_tracer_drop = true; - - let original = pb::TracerPayload { - chunks: vec![ - errored_root_chunk(0xe000, 0), - errored_root_chunk(0xe001, -1), - errored_root_chunk(0xe002, i32::from(i8::MIN)), - ], - ..Default::default() - }; - let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); - processor.apply_error_rescue_with_clock(&mut payload, &config, || { - panic!("disabled sampler must not read the clock"); - }); + fn test_disabled_tps_skips_rescue() { + for tps in [0.0_f64, -3.0] { + let (mut config, processor) = rescue_setup(rate_limited_sampler_config(tps)); + config.error_rare_sample_tracer_drop = true; - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); - let TracerPayloadCollection::V07(payloads) = payload else { - panic!("disabled rescue must preserve the payload format"); - }; - assert_eq!(payloads, vec![original], "mode={mode:?} tps={tps}"); - } + let original = pb::TracerPayload { + chunks: vec![ + errored_root_chunk(0xe000, 0), + errored_root_chunk(0xe001, -1), + ], + ..Default::default() + }; + let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); + processor.apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("disabled sampler must not read the clock"); + }); + + assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + let TracerPayloadCollection::V07(payloads) = payload else { + panic!("disabled rescue must preserve the payload format"); + }; + assert_eq!(payloads, vec![original], "tps={tps}"); } } #[test] fn test_rate_limited_under_sustained_load_keeps_and_rejects() { - for sample_rate in [None, Some(0.25_f64)] { - let sampler_config = rate_limited_sampler_config(10.0); - let (config, processor) = rescue_setup(sampler_config); + for (sample_rate, expected_keeps) in [(None, 83_usize), (Some(0.25_f64), 9)] { + let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); let context = format!("sample_rate={sample_rate:?}"); - // With the default client rate, 100 chunks with the same signature - // in one 5-second bucket give a default rate of 10 / (100 / 5) = 0.5. - // A finite client rate affects both weighting and applied probability. - // No exact keep count is asserted: this is an adaptive rolling-window - // sampler, not a hard token bucket. + // 100 errored chunks with the same signature in one bucket. The + // client `_sample_rate` changes both weighting and applied + // probability, so the two cases keep different counts. let ids: Vec = (0..100_u64).map(|i| 0x11_0000 + i).collect(); let input: Vec = ids .iter() @@ -1578,56 +1548,30 @@ mod tests { .collect(); let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); - let keeps = rescued_count(&chunks); - let drops = chunks.len() - keeps; - if sample_rate.is_none() { - assert!(keeps > 0, "{context}: expected some keeps, got none"); - } - assert!(drops > 0, "{context}: expected some drops, got none"); assert_eq!( chunks.len(), 100, "{context}: every chunk is still forwarded" ); - - // Compare every decision against direct crate usage. Set the expected - // client rate from the input case, independently of the output metric. - let mut standalone = ErrorsSampler::new(sampler_config); - for (chunk, id) in chunks.iter().zip(&ids) { - assert_eq!( - chunk.priority, 0, - "{context}: priority changed for trace id {id}" - ); - let views: Vec = chunk.spans.iter().map(super::span_view).collect(); - let mut trace = equivalent_trace_view(chunk, &config.env, &views); - trace.root_global_sample_rate = sample_rate.unwrap_or(1.0); - let expected = standalone.sample(RESCUE_NOW, &trace); - let actual = match errors_sr(root(chunk)) { - Some(errors_sr) => SampleDecision::Keep { errors_sr }, - None => SampleDecision::Drop, - }; - assert_eq!( - actual, expected, - "{context}: decision mismatch for trace id {id}" - ); - } + assert!( + chunks.iter().all(|chunk| chunk.priority == 0), + "{context}: priority changed" + ); + assert_eq!(rescued_count(&chunks), expected_keeps, "{context}"); } } #[test] fn test_processor_clones_share_one_budget() { - let sampler_config = rate_limited_sampler_config(1.0); - let (mut config, processor) = rescue_setup(sampler_config); + let (mut config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); config.error_rare_sample_tracer_drop = true; let processor_clone = processor.clone(); // Ten errored chunks with the same signature but distinct IDs: five // through the original, five through the clone. Mix P0 and explicit - // drops. All must draw from the same budget, matching a standalone - // sampler fed the same sequence. + // drops. Ineligible chunks interleaved with them spend no budget. let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); let mut observed = Vec::new(); - let mut standalone = ErrorsSampler::new(sampler_config); for (i, id) in ids.iter().enumerate() { let target = if i < 5 { &processor } else { &processor_clone }; let mut manual = errored_root_chunk(*id, -1); @@ -1650,26 +1594,13 @@ mod tests { vec![errored_root_chunk(*id, priority)], RESCUE_NOW, ); - let actual = errors_sr(root(&chunks[0])).is_some(); - let views: Vec = chunks[0].spans.iter().map(super::span_view).collect(); - let trace = equivalent_trace_view(&chunks[0], &config.env, &views); - let expected = matches!( - standalone.sample(RESCUE_NOW, &trace), - SampleDecision::Keep { .. } - ); - assert_eq!( - actual, expected, - "clone {i} decision diverged from the shared-budget standalone sampler" - ); - observed.push(actual); + observed.push(errors_sr(root(&chunks[0])).is_some()); } - assert!( - observed.iter().any(|kept| *kept), - "expected at least one keep" - ); - assert!( - !observed.iter().all(|kept| *kept), - "expected at least one drop" + assert_eq!( + observed, + [ + true, true, true, true, true, true, false, false, false, false + ] ); } @@ -1701,43 +1632,6 @@ mod tests { } } - #[test] - fn test_error_rescue_reads_clock_under_lock() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - let mut payload = TracerPayloadCollection::V07(vec![ - pb::TracerPayload { - chunks: vec![ - errored_root_chunk(0x10000, 1), - errored_root_chunk(0x10001, 0), - ], - ..Default::default() - }, - pb::TracerPayload { - chunks: vec![errored_root_chunk(0x20000, 0)], - ..Default::default() - }, - ]); - let mut clock_called = false; - - processor.apply_error_rescue_with_clock(&mut payload, &config, || { - clock_called = true; - assert!(matches!( - processor.error_rescue_state.try_lock(), - Err(std::sync::TryLockError::WouldBlock) - )); - RESCUE_NOW - }); - - assert!(clock_called); - assert_eq!(processor.lock_sampler().last_sampler_timestamp, RESCUE_NOW); - let TracerPayloadCollection::V07(payloads) = payload else { - panic!("expected V07 payload"); - }; - assert_eq!(errors_sr(root(&payloads[0].chunks[0])), None); - assert_eq!(errors_sr(root(&payloads[0].chunks[1])), Some(1.0)); - assert_eq!(errors_sr(root(&payloads[1].chunks[0])), Some(1.0)); - } - #[test] fn test_error_rescue_without_candidates_skips_sampler() { let (config, processor) = rescue_setup(always_keep_sampler_config()); From 3986519d5ec7bb3946e5afc6a8253113ace4cec7 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:26:58 -0400 Subject: [PATCH 41/49] Build error rescue state inside the trace processor ServerlessTraceProcessor::new now takes the config and owns its error sampler, instead of exposing the shared lock in its public signature. Clones still share one budget. --- crates/datadog-serverless-compat/src/main.rs | 2 +- .../src/trace_processor.rs | 59 ++++++------------- .../tests/integration_test.rs | 17 ++---- 3 files changed, 24 insertions(+), 54 deletions(-) diff --git a/crates/datadog-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs index 2ce7fa96..5bc20f48 100644 --- a/crates/datadog-serverless-compat/src/main.rs +++ b/crates/datadog-serverless-compat/src/main.rs @@ -208,7 +208,7 @@ pub async fn main() { let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor::new( stats_concentrator.as_ref().map(|c| c.handle.clone()), - trace_processor::ServerlessTraceProcessor::new_error_rescue_state(&config), + &config, )); let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher { diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index e6110e63..723ccde1 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -146,7 +146,7 @@ impl TraceChunkProcessor for ChunkProcessor { const MAX_IN_FLIGHT_ENQUEUES: usize = 10; /// Error rescue sampler and its monotonic timestamp floor, shared under one mutex. -pub struct ErrorRescueState { +struct ErrorRescueState { sampler: ErrorsSampler, last_sampler_timestamp: i64, } @@ -170,33 +170,23 @@ impl ErrorRescueState { pub struct ServerlessTraceProcessor { pub stats_concentrator: Option, enqueue_permits: Arc, - /// Processor clones and independently constructed processors using this - /// `Arc` share one TPS budget and timestamp floor under the same lock. + /// Processor clones share one TPS budget and timestamp floor under the same lock. error_rescue_state: Arc>, } impl ServerlessTraceProcessor { #[allow(clippy::must_use_candidate)] - pub fn new( - stats_concentrator: Option, - error_rescue_state: Arc>, - ) -> Self { - ServerlessTraceProcessor { - stats_concentrator, - enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), - error_rescue_state, - } - } - - /// Builds the shared error rescue state from parsed config settings. - #[must_use] - pub fn new_error_rescue_state(config: &Config) -> Arc> { + pub fn new(stats_concentrator: Option, config: &Config) -> Self { if config.error_rare_sample_tracer_drop && config.error_sampler.target_tps <= 0.0 { warn!( "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue" ); } - Arc::new(Mutex::new(ErrorRescueState::new(config.error_sampler))) + ServerlessTraceProcessor { + stats_concentrator, + enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), + error_rescue_state: Arc::new(Mutex::new(ErrorRescueState::new(config.error_sampler))), + } } /// Locks the shared sampler and timestamp floor, recovering from a poisoned @@ -562,8 +552,8 @@ mod tests { config::{Config, Tags}, peer_tags::peer_tag_keys, trace_processor::{ - self, ErrorRescueState, MAX_IN_FLIGHT_ENQUEUES, TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY, - TraceProcessor, encoded_size, split_oversized_payloads, + self, MAX_IN_FLIGHT_ENQUEUES, TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY, TraceProcessor, + encoded_size, split_oversized_payloads, }, }; use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; @@ -636,10 +626,6 @@ mod tests { } } - fn default_error_sampler() -> Arc> { - sampler_for(&ErrorSamplerConfig::default()) - } - fn create_test_metadata() -> MiniAgentMetadata { MiniAgentMetadata { azure_spring_app_hostname: Default::default(), @@ -756,7 +742,7 @@ mod tests { .unwrap(); let trace_processor = - trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); + trace_processor::ServerlessTraceProcessor::new(None, &create_test_config()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -830,7 +816,7 @@ mod tests { .unwrap(); let trace_processor = - trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); + trace_processor::ServerlessTraceProcessor::new(None, &create_test_config()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -917,7 +903,7 @@ mod tests { .unwrap(); let trace_processor = - trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); + trace_processor::ServerlessTraceProcessor::new(None, &create_test_config()); let res = trace_processor .process_traces( Arc::new(create_test_config()), @@ -976,7 +962,7 @@ mod tests { ) = mpsc::channel(1); let trace_processor = - trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); + trace_processor::ServerlessTraceProcessor::new(None, &create_test_config()); let config = Arc::new(Config { enqueue_permit_timeout_secs: 0, ..create_test_config() @@ -1045,7 +1031,7 @@ mod tests { ) = mpsc::channel(MAX_IN_FLIGHT_ENQUEUES + 1); let trace_processor = - trace_processor::ServerlessTraceProcessor::new(None, default_error_sampler()); + trace_processor::ServerlessTraceProcessor::new(None, &create_test_config()); // Uses the default (non-zero) enqueue_permit_timeout_secs, so the 11th request's permit // acquire has real time to succeed once an earlier request's send completes and // releases its permit, rather than racing a 0-second timeout against task scheduling. @@ -1141,20 +1127,13 @@ mod tests { } } - fn sampler_for(config: &ErrorSamplerConfig) -> Arc> { - Arc::new(std::sync::Mutex::new(ErrorRescueState::new(*config))) - } - /// Builds the standard rescue test fixture: a config with the given /// sampler settings and a processor whose shared sampler uses them. fn rescue_setup( sampler_config: ErrorSamplerConfig, ) -> (Config, trace_processor::ServerlessTraceProcessor) { let config = rescue_config(sampler_config); - let processor = trace_processor::ServerlessTraceProcessor::new( - None, - sampler_for(&config.error_sampler), - ); + let processor = trace_processor::ServerlessTraceProcessor::new(None, &config); (config, processor) } @@ -1342,7 +1321,7 @@ mod tests { let warning = "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue"; let mut config = rescue_config(rate_limited_sampler_config(0.0)); config.error_rare_sample_tracer_drop = true; - let _state = trace_processor::ServerlessTraceProcessor::new_error_rescue_state(&config); + let _processor = trace_processor::ServerlessTraceProcessor::new(None, &config); logs_assert(|lines| { let count = lines.iter().filter(|line| line.contains(warning)).count(); if count == 1 { @@ -1474,7 +1453,7 @@ mod tests { let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); let processor = trace_processor::ServerlessTraceProcessor::new( Some(super::StatsConcentratorHandle::new(stats_tx)), - sampler_for(&config.error_sampler), + &config, ); let (tx, mut rx) = mpsc::channel(1); let res = processor @@ -1686,7 +1665,7 @@ mod tests { let (stats_tx, mut stats_rx) = tokio::sync::mpsc::unbounded_channel(); let processor = trace_processor::ServerlessTraceProcessor::new( Some(super::StatsConcentratorHandle::new(stats_tx)), - sampler_for(&config.error_sampler), + &config, ); for (priority, dm, rescued) in [(0, None, true), (-1, Some("-4"), false)] { let (tx, mut rx) = mpsc::channel(1); diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 9efa30d9..0381a322 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -101,7 +101,7 @@ pub fn create_mini_agent_with_real_flushers( config: config.clone(), trace_processor: Arc::new(ServerlessTraceProcessor::new( Some(stats_concentrator_handle.clone()), - ServerlessTraceProcessor::new_error_rescue_state(&config), + &config, )), trace_flusher: Arc::new(ServerlessTraceFlusher::new( aggregator.clone(), @@ -296,10 +296,7 @@ async fn test_mini_agent_tcp_handles_requests() { let test_port = config.dd_apm_receiver_port; let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new( - None, - ServerlessTraceProcessor::new_error_rescue_state(&config), - )), + trace_processor: Arc::new(ServerlessTraceProcessor::new(None, &config)), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -413,10 +410,7 @@ async fn test_mini_agent_named_pipe_handles_requests() { let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new( - None, - ServerlessTraceProcessor::new_error_rescue_state(&config), - )), + trace_processor: Arc::new(ServerlessTraceProcessor::new(None, &config)), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -586,10 +580,7 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { let mini_agent = MiniAgent { config: config.clone(), - trace_processor: Arc::new(ServerlessTraceProcessor::new( - None, - ServerlessTraceProcessor::new_error_rescue_state(&config), - )), + trace_processor: Arc::new(ServerlessTraceProcessor::new(None, &config)), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), From 5aca291cfd178baae95cbdf2cdd0f4d7602f0501 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:28:48 -0400 Subject: [PATCH 42/49] Hold the error sampler lock only while sampling Disabled rescue no longer takes the shared lock, and span views are built before locking, so large payloads hold it for less time. --- .../src/trace_processor.rs | 109 +++++++++++------- 1 file changed, 66 insertions(+), 43 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 723ccde1..13026d62 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -172,6 +172,8 @@ pub struct ServerlessTraceProcessor { enqueue_permits: Arc, /// Processor clones share one TPS budget and timestamp floor under the same lock. error_rescue_state: Arc>, + /// False when the configured TPS disables the sampler, so rescue never locks. + error_rescue_enabled: bool, } impl ServerlessTraceProcessor { @@ -182,10 +184,12 @@ impl ServerlessTraceProcessor { "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue" ); } + let error_rescue_state = ErrorRescueState::new(config.error_sampler); ServerlessTraceProcessor { stats_concentrator, enqueue_permits: Arc::new(tokio::sync::Semaphore::new(MAX_IN_FLIGHT_ENQUEUES)), - error_rescue_state: Arc::new(Mutex::new(ErrorRescueState::new(config.error_sampler))), + error_rescue_enabled: !error_rescue_state.sampler.is_disabled(), + error_rescue_state: Arc::new(Mutex::new(error_rescue_state)), } } @@ -217,49 +221,68 @@ impl ServerlessTraceProcessor { let TracerPayloadCollection::V07(tracer_payloads) = payload else { return; }; - let mut state = None; - let mut read_clock = Some(read_clock); - for tracer_payload in tracer_payloads.iter_mut() { - let env = resolve_payload_env(&tracer_payload.env, &config.env); - for chunk in &mut tracer_payload.chunks { - let Some(root_index) = - rescue_root_index(chunk, config.error_rare_sample_tracer_drop) - else { - continue; - }; - // Requests without candidates never take the shared lock. - let state = state.get_or_insert_with(|| self.lock_sampler()); - if let Some(read_clock) = read_clock.take() { - if state.sampler.is_disabled() { - return; + if !self.error_rescue_enabled { + return; + } + + // Build the sampler inputs before locking, so the shared lock covers + // only sampling. Requests without candidates never take the lock. + let candidates: Vec<(usize, usize, &str, usize, Vec)> = + tracer_payloads + .iter() + .enumerate() + .flat_map(|(payload_index, tracer_payload)| { + let env = resolve_payload_env(&tracer_payload.env, &config.env); + tracer_payload.chunks.iter().enumerate().filter_map( + move |(chunk_index, chunk)| { + let root_index = + rescue_root_index(chunk, config.error_rare_sample_tracer_drop)?; + let views = chunk.spans.iter().map(span_view).collect(); + Some((payload_index, chunk_index, env, root_index, views)) + }, + ) + }) + .collect(); + if candidates.is_empty() { + return; + } + + let kept: Vec<(usize, usize, usize, f64)> = { + let mut state = self.lock_sampler(); + // Read once under the lock, and clamp backward wall-clock steps so + // the sampler's rolling window cannot move backwards. + let now_unix_secs = state.clamp_timestamp(read_clock()); + candidates + .iter() + .filter_map(|(payload_index, chunk_index, env, root_index, views)| { + let root = + &tracer_payloads[*payload_index].chunks[*chunk_index].spans[*root_index]; + let trace = TraceView { + env, + trace_id: root.trace_id, + root_index: *root_index, + // The shared sampler sanitizes the raw wire rate. + root_global_sample_rate: root + .metrics + .get("_sample_rate") + .copied() + .unwrap_or(1.0), + spans: views, + }; + match state.sampler.sample(now_unix_secs, &trace) { + SampleDecision::Keep { errors_sr } => { + Some((*payload_index, *chunk_index, *root_index, errors_sr)) + } + SampleDecision::Drop => None, } - // Read once under the lock, and clamp backward wall-clock steps - // so the sampler's rolling window cannot move backwards. - state.clamp_timestamp(read_clock()); - } - let now_unix_secs = state.last_sampler_timestamp; - let root = &chunk.spans[root_index]; - let views: Vec = chunk.spans.iter().map(span_view).collect(); - let trace = TraceView { - env, - trace_id: root.trace_id, - root_index, - // The shared sampler sanitizes the raw wire rate. - root_global_sample_rate: root - .metrics - .get("_sample_rate") - .copied() - .unwrap_or(1.0), - spans: &views, - }; - if let SampleDecision::Keep { errors_sr } = - state.sampler.sample(now_unix_secs, &trace) - { - chunk.spans[root_index] - .metrics - .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); - } - } + }) + .collect() + }; + + for (payload_index, chunk_index, root_index, errors_sr) in kept { + tracer_payloads[payload_index].chunks[chunk_index].spans[root_index] + .metrics + .insert(ERRORS_SR_METRIC_KEY.to_string(), errors_sr); } } From d1bd5137c7caf31ec938dbca752c742645c0fa06 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:29:34 -0400 Subject: [PATCH 43/49] Move payload env fallback next to the config --- crates/datadog-trace-agent/src/config.rs | 31 +++++++++++++++++ .../src/stats_concentrator_service.rs | 3 +- .../src/trace_processor.rs | 33 +------------------ 3 files changed, 33 insertions(+), 34 deletions(-) diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 1befb503..d31a5716 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -53,6 +53,18 @@ fn error_rare_sample_tracer_drop_enabled(raw: &str) -> bool { .any(|entry| entry == "error_rare_sample_tracer_drop") } +/// Chooses the env used to key per-signature state: the payload-reported env +/// when nonempty, otherwise the agent's configured env. Shared by stats +/// flushing and error-rescue sampling so the fallback rule can't drift +/// between the two. +pub(crate) fn resolve_payload_env<'a>(payload_env: &'a str, config_env: &'a str) -> &'a str { + if payload_env.is_empty() { + config_env + } else { + payload_env + } +} + const DEFAULT_DOGSTATSD_PORT: u16 = 8125; const DSM_PIPELINE_STATS_ROUTE: &str = "/api/v0.1/pipeline_stats"; @@ -1036,6 +1048,25 @@ mod tests { assert_eq!(error_sampler.extra_sample_rate, 1.0); } } + + #[test] + fn test_resolve_payload_env_prefers_payload_env() { + assert_eq!( + super::resolve_payload_env("tracer-env", "agent-env"), + "tracer-env", + "nonempty payload env wins" + ); + assert_eq!( + super::resolve_payload_env("", "agent-env"), + "agent-env", + "empty payload env falls back to the agent config env" + ); + assert_eq!( + super::resolve_payload_env("", ""), + "", + "both empty stays empty" + ); + } } /// Test helpers for creating Config instances in tests diff --git a/crates/datadog-trace-agent/src/stats_concentrator_service.rs b/crates/datadog-trace-agent/src/stats_concentrator_service.rs index 46897746..1e69ba21 100644 --- a/crates/datadog-trace-agent/src/stats_concentrator_service.rs +++ b/crates/datadog-trace-agent/src/stats_concentrator_service.rs @@ -6,8 +6,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; use tokio::sync::{mpsc, oneshot}; -use crate::config::Config; -use crate::trace_processor::resolve_payload_env; +use crate::config::{Config, resolve_payload_env}; use libdd_library_config::tracer_metadata::TracerMetadata; use libdd_trace_protobuf::pb::{ClientStatsPayload, TraceChunk}; use libdd_trace_stats::span_concentrator::{CardinalityLimitConfig, SpanConcentrator}; diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 13026d62..da636bff 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -25,7 +25,7 @@ use prost::Message; use crate::{ aggregator::MAX_CONTENT_SIZE_BYTES, - config::Config, + config::{Config, resolve_payload_env}, http_utils::{self, log_and_create_http_response, log_and_create_traces_success_http_response}, stats_concentrator_service::StatsConcentratorHandle, }; @@ -343,18 +343,6 @@ impl ServerlessTraceProcessor { } } -/// Chooses the env used to key per-signature state: the payload-reported env -/// when nonempty, otherwise the agent's configured env. Shared by stats -/// flushing and error-rescue sampling so the fallback rule can't drift -/// between the two. -pub(crate) fn resolve_payload_env<'a>(payload_env: &'a str, config_env: &'a str) -> &'a str { - if payload_env.is_empty() { - config_env - } else { - payload_env - } -} - /// Builds the sampler's read-only view of a span from its already enriched and /// obfuscated fields. fn span_view(span: &pb::Span) -> SpanView<'_> { @@ -1661,25 +1649,6 @@ mod tests { assert_eq!(payloads, vec![original]); } - #[test] - fn test_resolve_payload_env_prefers_payload_env() { - assert_eq!( - super::resolve_payload_env("tracer-env", "agent-env"), - "tracer-env", - "nonempty payload env wins" - ); - assert_eq!( - super::resolve_payload_env("", "agent-env"), - "agent-env", - "empty payload env falls back to the agent config env" - ); - assert_eq!( - super::resolve_payload_env("", ""), - "", - "both empty stays empty" - ); - } - #[tokio::test] async fn test_stats_concentrator_observes_pre_rescue_chunks() { let mut config = rescue_config(always_keep_sampler_config()); From 4d623ebbb7d2ec06a0211b990f4cf3b188f8c3b7 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Thu, 29 Oct 2026 17:30:27 -0400 Subject: [PATCH 44/49] Drop the tracing-test dev-dependency --- Cargo.lock | 1 - crates/datadog-trace-agent/Cargo.toml | 1 - .../datadog-trace-agent/src/trace_processor.rs | 17 ----------------- 3 files changed, 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 44f9c5b7..d28ba062 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -674,7 +674,6 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tracing", - "tracing-test", ] [[package]] diff --git a/crates/datadog-trace-agent/Cargo.toml b/crates/datadog-trace-agent/Cargo.toml index 295e07b8..a9b73ded 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -44,7 +44,6 @@ prost = "0.14.1" [dev-dependencies] datadog-mock-intake = { path = "../datadog-mock-intake" } -tracing-test = { version = "0.2.5", default-features = false } rmp-serde = "1.1.1" serial_test = "2.0.0" duplicate = "2.0.1" diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index da636bff..d65251b9 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1326,23 +1326,6 @@ mod tests { assert_eq!(chunks, original); } - #[test] - #[tracing_test::traced_test] - fn test_disabled_sampler_initialization_warning() { - let warning = "error_rare_sample_tracer_drop is enabled but DD_APM_ERROR_TPS <= 0 disables error rescue"; - let mut config = rescue_config(rate_limited_sampler_config(0.0)); - config.error_rare_sample_tracer_drop = true; - let _processor = trace_processor::ServerlessTraceProcessor::new(None, &config); - logs_assert(|lines| { - let count = lines.iter().filter(|line| line.contains(warning)).count(); - if count == 1 { - Ok(()) - } else { - Err(format!("expected one initialization warning, got {count}")) - } - }); - } - #[test] fn test_rescue_stamps_errored_p0_root_and_keeps_priority() { let (config, processor) = rescue_setup(always_keep_sampler_config()); From 503addb2622e8952d3e8cf26b7d158c53a3e3c7a Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 2 Nov 2026 12:38:03 -0500 Subject: [PATCH 45/49] Remove duplicate error-rescue priority test --- .../datadog-trace-agent/src/trace_processor.rs | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index d65251b9..e809905a 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1326,22 +1326,6 @@ mod tests { assert_eq!(chunks, original); } - #[test] - fn test_rescue_stamps_errored_p0_root_and_keeps_priority() { - let (config, processor) = rescue_setup(always_keep_sampler_config()); - - let chunks = run_rescue( - &processor, - &config, - vec![errored_root_chunk(0xdead_beef, 0)], - RESCUE_NOW, - ); - - assert_eq!(chunks.len(), 1, "rescued chunk must still be forwarded"); - assert_eq!(chunks[0].priority, 0, "priority must not be promoted"); - assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - } - #[test] fn test_rescue_uses_resolved_root_not_first_span() { let (config, processor) = rescue_setup(always_keep_sampler_config()); From f45614488caef2dbc62da3d973b29037674f25b8 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 2 Nov 2026 12:49:06 -0500 Subject: [PATCH 46/49] Simplify empty-chunk rescue test to one config --- .../src/trace_processor.rs | 69 +++++++++---------- 1 file changed, 31 insertions(+), 38 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index e809905a..91c7490d 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1381,45 +1381,38 @@ mod tests { #[test] fn test_empty_chunk_and_fallback_root() { - for (feature, priority) in [(false, 0), (true, -1)] { - let (mut config, processor) = rescue_setup(always_keep_sampler_config()); - config.error_rare_sample_tracer_drop = feature; - let context = format!("feature={feature} priority={priority}"); - - // Eligible empty chunks return at the no-error guard before root resolution. - // The cyclic chunk has no parent_id 0, so the last span is the fallback root. - let cyclic = test_chunk( - vec![ - test_span(0xd00d, 0xd01, 0xd02, 1), - test_span(0xd00d, 0xd02, 0xd01, 0), - ], - priority, - ); - let chunks = run_rescue( - &processor, - &config, - vec![test_chunk(vec![], priority), cyclic], - RESCUE_NOW, - ); + let (config, processor) = rescue_setup(always_keep_sampler_config()); - assert_eq!(chunks.len(), 2, "{context}"); - assert!( - chunks[0].spans.is_empty(), - "empty chunk unchanged: {context}" - ); - assert_eq!(chunks[0].priority, priority, "{context}"); - assert_eq!(chunks[1].priority, priority, "{context}"); - assert_eq!( - errors_sr(&chunks[1].spans[0]), - None, - "non-root span not stamped: {context}" - ); - assert_eq!( - errors_sr(&chunks[1].spans[1]), - Some(1.0), - "fallback root (last span) stamped: {context}" - ); - } + // Eligible empty chunks return at the no-error guard before root resolution. + // The cyclic chunk has no parent_id 0, so the last span is the fallback root. + let cyclic = test_chunk( + vec![ + test_span(0xd00d, 0xd01, 0xd02, 1), + test_span(0xd00d, 0xd02, 0xd01, 0), + ], + 0, + ); + let chunks = run_rescue( + &processor, + &config, + vec![test_chunk(vec![], 0), cyclic], + RESCUE_NOW, + ); + + assert_eq!(chunks.len(), 2); + assert!(chunks[0].spans.is_empty(), "empty chunk unchanged"); + assert_eq!(chunks[0].priority, 0); + assert_eq!(chunks[1].priority, 0); + assert_eq!( + errors_sr(&chunks[1].spans[0]), + None, + "non-root span not stamped" + ); + assert_eq!( + errors_sr(&chunks[1].spans[1]), + Some(1.0), + "fallback root (last span) stamped" + ); } #[tokio::test] From 70ed34901167d12b13fe96133a0e80eb12531830 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 2 Nov 2026 13:11:32 -0500 Subject: [PATCH 47/49] Test client sample rate wiring, not keep counts --- .../src/trace_processor.rs | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 91c7490d..d61b012b 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -1477,19 +1477,17 @@ mod tests { } #[test] - fn test_rate_limited_under_sustained_load_keeps_and_rejects() { - for (sample_rate, expected_keeps) in [(None, 83_usize), (Some(0.25_f64), 9)] { + fn test_client_sample_rate_affects_error_rescue() { + let mut rescued_by_rate = Vec::new(); + for sample_rate in [None, Some(0.25_f64)] { let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); - let context = format!("sample_rate={sample_rate:?}"); // 100 errored chunks with the same signature in one bucket. The - // client `_sample_rate` changes both weighting and applied - // probability, so the two cases keep different counts. - let ids: Vec = (0..100_u64).map(|i| 0x11_0000 + i).collect(); - let input: Vec = ids - .iter() - .map(|id| { - let mut chunk = errored_root_chunk(*id, 0); + // client `_sample_rate` weights each chunk, so the two cases keep + // different counts. + let input: Vec = (0..100_u64) + .map(|i| { + let mut chunk = errored_root_chunk(0x11_0000 + i, 0); if let Some(rate) = sample_rate { chunk.spans[0].metrics.insert("_sample_rate".into(), rate); } @@ -1498,17 +1496,18 @@ mod tests { .collect(); let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); - assert_eq!( - chunks.len(), - 100, - "{context}: every chunk is still forwarded" - ); + assert_eq!(chunks.len(), 100, "every chunk is still forwarded"); assert!( chunks.iter().all(|chunk| chunk.priority == 0), - "{context}: priority changed" + "priority changed" ); - assert_eq!(rescued_count(&chunks), expected_keeps, "{context}"); + rescued_by_rate.push(rescued_count(&chunks)); } + + assert!( + rescued_by_rate[1] < rescued_by_rate[0], + "the client's root sample rate must affect error-rescue decisions" + ); } #[test] From 8d484f75608b697c9da2716ed94cbd3ad96cfe44 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 2 Nov 2026 13:31:41 -0500 Subject: [PATCH 48/49] Simplify error rescue rules and verification --- .../src/trace_processor.rs | 27 ++++++++++--------- .../tests/integration_test.rs | 17 ++++++------ 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index d61b012b..7ca28fe0 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -362,13 +362,16 @@ fn span_view(span: &pb::Span) -> SpanView<'_> { /// Only chunk tags determine DM; any span with `error != 0` supplies the error. /// Root selection delegates to `get_root_span_index`; empty chunks yield `None`. fn rescue_root_index(chunk: &pb::TraceChunk, tracer_drop_enabled: bool) -> Option { - let eligible = chunk.priority == 0 - || (tracer_drop_enabled - && chunk.priority < 0 - && chunk.priority != i32::from(i8::MIN) - && !(chunk.priority == -1 + let eligible = match chunk.priority { + 0 => true, + -128 => false, + -1 => { + tracer_drop_enabled && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) - == Some(MANUAL_DECISION_MAKER))); + != Some(MANUAL_DECISION_MAKER) + } + priority => tracer_drop_enabled && priority < 0, + }; if !eligible { return None; } @@ -1519,20 +1522,20 @@ mod tests { // Ten errored chunks with the same signature but distinct IDs: five // through the original, five through the clone. Mix P0 and explicit // drops. Ineligible chunks interleaved with them spend no budget. - let ids: Vec = (0..10_u64).map(|i| 0x33_0000 + i).collect(); let mut observed = Vec::new(); - for (i, id) in ids.iter().enumerate() { + for i in 0..10_u64 { + let id = 0x33_0000 + i; let target = if i < 5 { &processor } else { &processor_clone }; - let mut manual = errored_root_chunk(*id, -1); + let mut manual = errored_root_chunk(id, -1); manual .tags .insert(super::DECISION_MAKER_KEY.into(), "-4".into()); - let mut healthy = errored_root_chunk(*id, -1); + let mut healthy = errored_root_chunk(id, -1); healthy.spans[0].error = 0; let excluded = run_rescue( target, &config, - vec![manual, healthy, errored_root_chunk(*id, 1)], + vec![manual, healthy, errored_root_chunk(id, 1)], RESCUE_NOW, ); assert_eq!(rescued_count(&excluded), 0); @@ -1540,7 +1543,7 @@ mod tests { let chunks = run_rescue( target, &config, - vec![errored_root_chunk(*id, priority)], + vec![errored_root_chunk(id, priority)], RESCUE_NOW, ); observed.push(errors_sr(root(&chunks[0])).is_some()); diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 0381a322..d6ae5811 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1346,8 +1346,7 @@ async fn test_error_rescue_feature_on_outbound_payload_and_stats() { async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { use datadog_agent_trace_sampler::{ErrorSamplerConfig, ErrorSamplerMode}; - use libdd_trace_protobuf::pb::{AgentPayload, TraceChunk}; - use prost::Message as _; + use libdd_trace_protobuf::pb::TraceChunk; use std::collections::HashMap; let mock_intake = MockIntake::start().await; @@ -1411,13 +1410,15 @@ async fn verify_error_rescue_outbound_payload_and_stats(feature: bool) { .expect("bounded agent shutdown") .expect("agent task"); - // Decode the actual outbound protobuf payload(s) and locate each submitted - // chunk by its root span name. - let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); + // Locate each submitted chunk by its root span name in the outbound payloads. + let trace_payloads = mock_intake.trace_payloads(); + assert_eq!( + trace_payloads.len(), + mock_intake.requests_for_path("/api/v0.2/traces").len(), + "every outbound trace request must decode successfully" + ); let mut chunks_by_root_name: HashMap = HashMap::new(); - for req in &trace_reqs { - let agent_payload = AgentPayload::decode(&req.body[..]) - .expect("Failed to decode outbound AgentPayload protobuf"); + for agent_payload in trace_payloads { for tracer_payload in agent_payload.tracer_payloads { for chunk in tracer_payload.chunks { let root_name = chunk From 1bfef7da57b8c4d188585a6a81cb7b966b2066a9 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Mon, 2 Nov 2026 18:05:03 -0500 Subject: [PATCH 49/49] Avoid blocking async workers during error rescue --- .../src/trace_processor.rs | 177 ++++++++++++------ 1 file changed, 122 insertions(+), 55 deletions(-) diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index 7ca28fe0..497c8dfc 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -11,9 +11,8 @@ use http_body_util::BodyExt; use hyper::{StatusCode, http}; use libdd_common::http_common; use libdd_library_config::tracer_metadata::TracerMetadata; -use std::sync::{Mutex, MutexGuard, PoisonError}; use std::time::{SystemTime, UNIX_EPOCH}; -use tokio::sync::mpsc::Sender; +use tokio::sync::{Mutex, mpsc::Sender}; use tracing::{debug, error, warn}; use libdd_trace_obfuscation::obfuscate::obfuscate_pb_span; @@ -193,26 +192,19 @@ impl ServerlessTraceProcessor { } } - /// Locks the shared sampler and timestamp floor, recovering from a poisoned - /// guard (a panic in another thread while holding the lock). - fn lock_sampler(&self) -> MutexGuard<'_, ErrorRescueState> { - self.error_rescue_state - .lock() - .unwrap_or_else(PoisonError::into_inner) - } - /// Samples V07 chunks selected by [`rescue_root_index`] and stamps /// `_dd.errors_sr` on kept roots. Priorities and chunk order are unchanged; /// every chunk is forwarded, with retention decided by the backend. - fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { + async fn apply_error_rescue(&self, payload: &mut TracerPayloadCollection, config: &Config) { self.apply_error_rescue_with_clock(payload, config, || { SystemTime::now() .duration_since(UNIX_EPOCH) .map_or(0, |d| d.as_secs() as i64) - }); + }) + .await; } - fn apply_error_rescue_with_clock i64>( + async fn apply_error_rescue_with_clock i64>( &self, payload: &mut TracerPayloadCollection, config: &Config, @@ -248,7 +240,7 @@ impl ServerlessTraceProcessor { } let kept: Vec<(usize, usize, usize, f64)> = { - let mut state = self.lock_sampler(); + let mut state = self.error_rescue_state.lock().await; // Read once under the lock, and clamp backward wall-clock steps so // the sampler's rolling window cannot move backwards. let now_unix_secs = state.clamp_timestamp(read_clock()); @@ -289,13 +281,14 @@ impl ServerlessTraceProcessor { /// Test-only variant that injects a synthetic clock reading through the /// production clamp path, exercising the rolling window without sleeping. #[cfg(test)] - fn apply_error_rescue_at( + async fn apply_error_rescue_at( &self, payload: &mut TracerPayloadCollection, config: &Config, now_unix_secs: i64, ) { - self.apply_error_rescue_with_clock(payload, config, || now_unix_secs); + self.apply_error_rescue_with_clock(payload, config, || now_unix_secs) + .await; } fn send_to_concentrator( @@ -490,7 +483,7 @@ impl TraceProcessor for ServerlessTraceProcessor { // conditional stats submission so the concentrator observes chunks before // rescue stamping, and before payload splitting so the inserted metric is // included in the recomputed outbound size. - self.apply_error_rescue(&mut payload, &config); + self.apply_error_rescue(&mut payload, &config).await; let pieces: Vec<(TracerPayloadCollection, usize)> = match payload { TracerPayloadCollection::V07(payloads) => { @@ -1151,7 +1144,7 @@ mod tests { (config, processor) } - fn run_rescue( + async fn run_rescue( processor: &trace_processor::ServerlessTraceProcessor, config: &Config, chunks: Vec, @@ -1161,7 +1154,9 @@ mod tests { chunks, ..Default::default() }]); - processor.apply_error_rescue_at(&mut payload, config, now_unix_secs); + processor + .apply_error_rescue_at(&mut payload, config, now_unix_secs) + .await; match payload { TracerPayloadCollection::V07(mut payloads) => payloads.remove(0).chunks, _ => unreachable!(), @@ -1302,8 +1297,8 @@ mod tests { } } - #[test] - fn test_rescue_preserves_chunks_except_root_metric() { + #[tokio::test] + async fn test_rescue_preserves_chunks_except_root_metric() { let (config, processor) = rescue_setup(always_keep_sampler_config()); let mut chunk = test_chunk(vec![test_span(1, 2, 0, 0), test_span(1, 3, 2, 1)], 0); chunk.spans[0].metrics.extend([ @@ -1319,7 +1314,7 @@ mod tests { chunk.tags.insert("origin".into(), "synthetics".into()); let healthy = test_chunk(vec![test_span(2, 4, 0, 0)], 0); let original = vec![chunk, healthy]; - let mut chunks = run_rescue(&processor, &config, original.clone(), RESCUE_NOW); + let mut chunks = run_rescue(&processor, &config, original.clone(), RESCUE_NOW).await; assert_eq!(chunks.len(), 2, "every chunk must still be forwarded"); assert_eq!(errors_sr(&chunks[0].spans[0]), Some(1.0)); assert_eq!(errors_sr(&chunks[0].spans[1]), None); @@ -1329,8 +1324,8 @@ mod tests { assert_eq!(chunks, original); } - #[test] - fn test_rescue_uses_resolved_root_not_first_span() { + #[tokio::test] + async fn test_rescue_uses_resolved_root_not_first_span() { let (config, processor) = rescue_setup(always_keep_sampler_config()); // Root (parent_id 0) appears last; the first span is a child. @@ -1339,7 +1334,7 @@ mod tests { test_span(0xbeef, 0x202, 0x203, 1), test_span(0xbeef, 0x203, 0, 0), ]; - let chunks = run_rescue(&processor, &config, vec![test_chunk(spans, 0)], RESCUE_NOW); + let chunks = run_rescue(&processor, &config, vec![test_chunk(spans, 0)], RESCUE_NOW).await; assert_eq!(chunks.len(), 1); assert_eq!( @@ -1351,8 +1346,8 @@ mod tests { assert_eq!(errors_sr(&chunks[0].spans[1]), None); } - #[test] - fn test_no_error_flags_leaves_chunk_unchanged_and_spends_no_budget() { + #[tokio::test] + async fn test_no_error_flags_leaves_chunk_unchanged_and_spends_no_budget() { // With RateLimited, counting is global across signatures via // `all_sigs_seen`, so if the healthy chunks below were wrongly fed to // the sampler, the final errored chunk's rate would drop below 1.0 and @@ -1370,7 +1365,7 @@ mod tests { .collect(); chunks.push(errored_root_chunk(0xb000, 0)); - let chunks = run_rescue(&processor, &config, chunks, RESCUE_NOW); + let chunks = run_rescue(&processor, &config, chunks, RESCUE_NOW).await; for (i, chunk) in chunks.iter().enumerate() { let is_last = i == chunks.len() - 1; @@ -1382,8 +1377,8 @@ mod tests { } } - #[test] - fn test_empty_chunk_and_fallback_root() { + #[tokio::test] + async fn test_empty_chunk_and_fallback_root() { let (config, processor) = rescue_setup(always_keep_sampler_config()); // Eligible empty chunks return at the no-error guard before root resolution. @@ -1400,7 +1395,8 @@ mod tests { &config, vec![test_chunk(vec![], 0), cyclic], RESCUE_NOW, - ); + ) + .await; assert_eq!(chunks.len(), 2); assert!(chunks[0].spans.is_empty(), "empty chunk unchanged"); @@ -1453,8 +1449,8 @@ mod tests { assert_eq!(errors_sr(&payloads[0].chunks[0].spans[0]), Some(1.0)); } - #[test] - fn test_disabled_tps_skips_rescue() { + #[tokio::test] + async fn test_disabled_tps_skips_rescue() { for tps in [0.0_f64, -3.0] { let (mut config, processor) = rescue_setup(rate_limited_sampler_config(tps)); config.error_rare_sample_tracer_drop = true; @@ -1467,11 +1463,20 @@ mod tests { ..Default::default() }; let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); - processor.apply_error_rescue_with_clock(&mut payload, &config, || { - panic!("disabled sampler must not read the clock"); - }); + processor + .apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("disabled sampler must not read the clock"); + }) + .await; - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + assert_eq!( + processor + .error_rescue_state + .lock() + .await + .last_sampler_timestamp, + i64::MIN + ); let TracerPayloadCollection::V07(payloads) = payload else { panic!("disabled rescue must preserve the payload format"); }; @@ -1479,8 +1484,8 @@ mod tests { } } - #[test] - fn test_client_sample_rate_affects_error_rescue() { + #[tokio::test] + async fn test_client_sample_rate_affects_error_rescue() { let mut rescued_by_rate = Vec::new(); for sample_rate in [None, Some(0.25_f64)] { let (config, processor) = rescue_setup(rate_limited_sampler_config(10.0)); @@ -1497,7 +1502,7 @@ mod tests { chunk }) .collect(); - let chunks = run_rescue(&processor, &config, input, RESCUE_NOW); + let chunks = run_rescue(&processor, &config, input, RESCUE_NOW).await; assert_eq!(chunks.len(), 100, "every chunk is still forwarded"); assert!( @@ -1513,8 +1518,8 @@ mod tests { ); } - #[test] - fn test_processor_clones_share_one_budget() { + #[tokio::test] + async fn test_processor_clones_share_one_budget() { let (mut config, processor) = rescue_setup(rate_limited_sampler_config(1.0)); config.error_rare_sample_tracer_drop = true; let processor_clone = processor.clone(); @@ -1537,7 +1542,8 @@ mod tests { &config, vec![manual, healthy, errored_root_chunk(id, 1)], RESCUE_NOW, - ); + ) + .await; assert_eq!(rescued_count(&excluded), 0); let priority = if i % 2 == 0 { 0 } else { -1 }; let chunks = run_rescue( @@ -1545,7 +1551,8 @@ mod tests { &config, vec![errored_root_chunk(id, priority)], RESCUE_NOW, - ); + ) + .await; observed.push(errors_sr(root(&chunks[0])).is_some()); } assert_eq!( @@ -1556,11 +1563,54 @@ mod tests { ); } - #[test] - fn test_error_rescue_clamps_backward_timestamps_across_clones() { + #[tokio::test] + async fn test_error_rescue_yields_while_sampler_is_locked() { + use std::{ + future::{Future, poll_fn}, + task::Poll, + }; + + let (config, processor) = rescue_setup(always_keep_sampler_config()); + let guard = processor.error_rescue_state.lock().await; + let mut rescue = std::pin::pin!(run_rescue( + &processor, + &config, + vec![errored_root_chunk(0x10000, 0)], + RESCUE_NOW, + )); + + poll_fn(|cx| { + assert!(rescue.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(guard.last_sampler_timestamp, i64::MIN); + drop(guard); + + let chunks = rescue.await; + assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); + assert_eq!( + processor + .error_rescue_state + .lock() + .await + .last_sampler_timestamp, + RESCUE_NOW + ); + } + + #[tokio::test] + async fn test_error_rescue_clamps_backward_timestamps_across_clones() { let (config, processor) = rescue_setup(always_keep_sampler_config()); let clone = processor.clone(); - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + assert_eq!( + processor + .error_rescue_state + .lock() + .await + .last_sampler_timestamp, + i64::MIN + ); for (i, (now, expected)) in [ (RESCUE_NOW, RESCUE_NOW), @@ -1578,14 +1628,22 @@ mod tests { &config, vec![errored_root_chunk(0x10000 + i as u64, 0)], now, - ); + ) + .await; assert_eq!(errors_sr(root(&chunks[0])), Some(1.0)); - assert_eq!(target.lock_sampler().last_sampler_timestamp, expected); + assert_eq!( + target + .error_rescue_state + .lock() + .await + .last_sampler_timestamp, + expected + ); } } - #[test] - fn test_error_rescue_without_candidates_skips_sampler() { + #[tokio::test] + async fn test_error_rescue_without_candidates_skips_sampler() { let (config, processor) = rescue_setup(always_keep_sampler_config()); let original = pb::TracerPayload { chunks: vec![ @@ -1600,11 +1658,20 @@ mod tests { }; let mut payload = TracerPayloadCollection::V07(vec![original.clone()]); - processor.apply_error_rescue_with_clock(&mut payload, &config, || { - panic!("a payload without rescue candidates must not consult the sampler"); - }); + processor + .apply_error_rescue_with_clock(&mut payload, &config, || { + panic!("a payload without rescue candidates must not consult the sampler"); + }) + .await; - assert_eq!(processor.lock_sampler().last_sampler_timestamp, i64::MIN); + assert_eq!( + processor + .error_rescue_state + .lock() + .await + .last_sampler_timestamp, + i64::MIN + ); let TracerPayloadCollection::V07(payloads) = payload else { panic!("rescue must preserve the payload format"); };