[Traces] Add error sampling with opt-in tracer-drop rescue

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
 

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_sampler(&config),
     ));
 
     let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher {

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"

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::<f64>() {
+            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<usize>,
     /// 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,227 @@ 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 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]
+    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);
+                    },
+                );
+            },
+        );
+    }
+
+    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);
+                },
+            );
+        }
+    }
+
+    #[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 [
+            (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() {
+        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
@@ -930,6 +1204,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,
         }
     }
 }