diff --git a/Cargo.lock b/Cargo.lock index f889c50f..37f96bd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -712,6 +712,7 @@ dependencies = [ "anyhow", "async-trait", "bytes", + "datadog-agent-trace-sampler", "datadog-fips", "datadog-mock-intake", "duplicate", @@ -738,6 +739,7 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tracing", + "tracing-test", ] [[package]] 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-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs index cdcdcabb..b68b1e59 100644 --- a/crates/datadog-serverless-compat/src/main.rs +++ b/crates/datadog-serverless-compat/src/main.rs @@ -186,6 +186,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), )); 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 4e868937..5ff113ce 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -33,6 +33,7 @@ libdd-trace-stats = { version = "7.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", @@ -43,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 186ce0ae..759a7828 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -9,14 +9,57 @@ 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 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. +/// 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. +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 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}" + ); + } + } + } + + 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"; @@ -128,6 +171,12 @@ pub struct Config { pub additional_metric_tags_cardinality_limit: Option, /// Whether the agent should compute trace stats pub agent_stats_computation_enabled: bool, + /// 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. + pub error_rare_sample_tracer_drop: bool, } impl Config { @@ -304,6 +353,9 @@ 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(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)), }) } } @@ -315,6 +367,7 @@ mod tests { use std::collections::HashMap; use crate::config; + use datadog_agent_trace_sampler::ErrorSamplerMode; #[test] #[serial] @@ -888,6 +941,106 @@ 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:?}" + ); + } + } + + #[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() { + 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] + 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 helpers for creating Config instances in tests @@ -930,6 +1083,8 @@ pub mod test_helpers { additional_metric_tags: vec![], 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_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/stats_processor.rs b/crates/datadog-trace-agent/src/stats_processor.rs index 990d7067..42c82be1 100644 --- a/crates/datadog-trace-agent/src/stats_processor.rs +++ b/crates/datadog-trace-agent/src/stats_processor.rs @@ -164,6 +164,8 @@ mod tests { additional_metric_tags: vec![], 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 e26465c9..559eaa43 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -4,10 +4,15 @@ use std::sync::Arc; use async_trait::async_trait; +use datadog_agent_trace_sampler::{ + ErrorSamplerConfig, 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 +32,13 @@ use crate::{ const TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY: &str = "_dd.tags.function"; +/// 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 const V07_ENVELOPE_OVERHEAD_BYTES: usize = 64; @@ -94,6 +106,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(); @@ -120,19 +145,144 @@ 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, + /// 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) -> Self { + 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> { + 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))) + } + + /// 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) { + 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 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; + } + // 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); + } + } + } + } + + /// 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, + payload: &mut TracerPayloadCollection, + config: &Config, + now_unix_secs: i64, + ) { + self.apply_error_rescue_with_clock(payload, config, || now_unix_secs); } fn send_to_concentrator( @@ -179,6 +329,54 @@ 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<'_> { + 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), + } +} + +/// 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) + || (tracer_drop_enabled + && chunk.priority < 0 + && !(chunk.priority == -1 + && chunk.tags.get(DECISION_MAKER_KEY).map(String::as_str) + == Some(MANUAL_DECISION_MAKER))); + if !eligible { + return None; + } + // Skip root resolution for healthy chunks. + if !chunk.spans.iter().any(|span| span.error != 0) { + return None; + } + trace_utils::get_root_span_index(&chunk.spans).ok() +} + #[async_trait] impl TraceProcessor for ServerlessTraceProcessor { async fn process_traces( @@ -282,6 +480,13 @@ impl TraceProcessor for ServerlessTraceProcessor { Self::send_to_concentrator(concentrator, &payload); } + // 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) => { let split_budget = @@ -356,10 +561,13 @@ 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::{ + 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}; @@ -424,9 +632,15 @@ mod tests { additional_metric_tags: vec![], additional_metric_tags_cardinality_limit: None, agent_stats_computation_enabled: false, + error_sampler: ErrorSamplerConfig::default(), + error_rare_sample_tracer_drop: false, } } + fn default_error_sampler() -> Arc> { + sampler_for(&ErrorSamplerConfig::default()) + } + fn create_test_metadata() -> MiniAgentMetadata { MiniAgentMetadata { azure_spring_app_hostname: Default::default(), @@ -542,7 +756,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()), @@ -615,7 +830,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()), @@ -701,7 +917,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()), @@ -759,7 +976,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() @@ -827,7 +1045,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. @@ -868,4 +1087,753 @@ 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(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), + ); + (config, processor) + } + + 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_at(&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, + } + } + + 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 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!(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") + .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_decision_maker_promotion_precedence() { + use libdd_trace_utils::tracer_payload::TraceChunkProcessor as _; + 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 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 + ); + } + } + + #[test] + fn test_tracer_drop_rescue_eligibility_matrix() { + // 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() { + 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); + 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()); + + 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()); + + // 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, processor) = rescue_setup(rate_limited_sampler_config(1.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); + + 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_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, "{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] + async fn test_rescue_stamps_eligible_chunks_when_agent_stats_disabled() { + 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] + 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"); + }); + + 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}"); + } + } + } + + #[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); + 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" + ); + + // 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_processor_clones_share_one_budget() { + let sampler_config = rate_limited_sampler_config(1.0); + 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 + // 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(); + 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); + 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, 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, 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); + } + 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_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); + + 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] + 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()); + 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_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()); + config.error_rare_sample_tracer_drop = true; + let config = Arc::new(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), (-1, Some("-4"), 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!( + 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 57de50ef..255ce3b2 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_rescue_state(&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_rescue_state(&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_rescue_state(&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_rescue_state(&config), + )), trace_flusher: Arc::new(MockTraceFlusher), stats_processor: Arc::new(MockStatsProcessor), stats_flusher: Arc::new(MockStatsFlusher), @@ -1278,3 +1289,234 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { let _ = agent_handle.await; verify_stats_request(&mock_intake).await; } + +/// 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_nonmanual", -1, 1, Some("-3"), 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), + ] +} + +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<_> = 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") +} + +/// 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_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_intake = MockIntake::start().await; + + // 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, + 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_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_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[..]) + .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 cases = error_rescue_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: {:?}", + 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()); + 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, + rescued.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 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 is_child { + assert!(!span.metrics.contains_key("_dd.errors_sr")); + } + } + } + + // 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 + .into_iter() + .flat_map(|csp| csp.stats.into_iter()) + .flat_map(|bucket| bucket.stats.into_iter()) + }) + .collect(); + + for name in cases + .into_iter() + .map(|case| case.0) + .chain(["manual_child_span"]) + { + 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.hits > 0, + "expected span {name} to have a positive hit count in stats" + ); + } +}