saluki_components/transforms/trace_sampler/

mod.rs

//! Trace sampling transform.
//!
//! This transform implements agent-side head sampling for traces, supporting:
//! - Probabilistic sampling based on trace ID
//! - User-set priority preservation
//! - Error-based sampling as a safety net
//! - OTLP trace ingestion with proper sampling decision handling
//!
//! TODO:
//!
//! - add trace metrics: datadog-agent/pkg/trace/sampler/metrics.go
//! - adding missing samplers (priority, nopriority)
//! - add error tracking standalone mode

use async_trait::async_trait;
use saluki_common::collections::FastHashMap;
use saluki_config::GenericConfiguration;
use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
use saluki_core::{
    components::{transforms::*, ComponentContext},
    data_model::event::{
        trace::{AttributeValue, Span, Trace},
        Event,
    },
    topology::EventsBuffer,
};
use saluki_error::GenericError;
use stringtheory::MetaString;
use tracing::debug;

mod catalog;
mod core_sampler;
mod errors;
mod priority_sampler;
mod probabilistic;
mod rare_sampler;
mod score_sampler;
mod signature;

use self::probabilistic::PROB_RATE_KEY;
use crate::common::datadog::{
    apm::ApmConfig, sample_by_rate, DECISION_MAKER_MANUAL, DECISION_MAKER_PROBABILISTIC, OTEL_TRACE_ID_META_KEY,
    SAMPLING_PRIORITY_METRIC_KEY, TAG_DECISION_MAKER,
};
use crate::common::otlp::config::TracesConfig;

// Sampling priority constants (matching datadog-agent)
const PRIORITY_AUTO_DROP: i32 = 0;
const PRIORITY_AUTO_KEEP: i32 = 1;
const PRIORITY_USER_KEEP: i32 = 2;

const ERROR_SAMPLE_RATE: f64 = 1.0; // Default extra sample rate (matches agent's ExtraSampleRate)

// Single Span Sampling and Analytics Events keys
const KEY_SPAN_SAMPLING_MECHANISM: &str = "_dd.span_sampling.mechanism";
const KEY_ANALYZED_SPANS: &str = "_dd.analyzed";

// Decision maker values for `_dd.p.dm` (matching datadog-agent).

fn normalize_sampling_rate(rate: f64) -> f64 {
    if rate <= 0.0 || rate >= 1.0 {
        1.0
    } else {
        rate
    }
}

/// Configuration for the trace sampler transform.
#[derive(Debug)]
pub struct TraceSamplerConfiguration {
    apm_config: ApmConfig,
    otlp_sampling_rate: f64,
}

impl TraceSamplerConfiguration {
    /// Creates a new `TraceSamplerConfiguration` from the given configuration.
    pub fn from_configuration(config: &GenericConfiguration) -> Result<Self, GenericError> {
        let apm_config = ApmConfig::from_configuration(config)?;
        let otlp_traces: TracesConfig = config.try_get_typed("otlp_config.traces")?.unwrap_or_default();
        let otlp_sampling_rate = normalize_sampling_rate(otlp_traces.probabilistic_sampler.sampling_percentage / 100.0);
        Ok(Self {
            apm_config,
            otlp_sampling_rate,
        })
    }
}

#[async_trait]
impl SynchronousTransformBuilder for TraceSamplerConfiguration {
    async fn build(&self, _context: ComponentContext) -> Result<Box<dyn SynchronousTransform + Send>, GenericError> {
        // TODO: Need to support remote configuration changing these at runtime
        // See https://github.com/DataDog/saluki/issues/1326
        let sampler = TraceSampler {
            sampling_rate: self.apm_config.probabilistic_sampler_sampling_percentage() / 100.0,
            error_sampling_enabled: self.apm_config.error_sampling_enabled(),
            error_tracking_standalone: self.apm_config.error_tracking_standalone_enabled(),
            probabilistic_sampler_enabled: self.apm_config.probabilistic_sampler_enabled(),
            otlp_sampling_rate: self.otlp_sampling_rate,
            error_sampler: errors::ErrorsSampler::new(self.apm_config.errors_per_second(), ERROR_SAMPLE_RATE),
            priority_sampler: priority_sampler::PrioritySampler::new(
                self.apm_config.default_env().clone(),
                ERROR_SAMPLE_RATE,
                self.apm_config.target_traces_per_second(),
            ),
            no_priority_sampler: score_sampler::NoPrioritySampler::new(
                self.apm_config.target_traces_per_second(),
                ERROR_SAMPLE_RATE,
            ),
            rare_sampler: rare_sampler::RareSampler::new(
                self.apm_config.rare_sampler_enabled(),
                self.apm_config.rare_sampler_tps(),
                std::time::Duration::from_secs_f64(self.apm_config.rare_sampler_cooldown_period_secs()),
                self.apm_config.rare_sampler_cardinality(),
            ),
        };

        Ok(Box::new(sampler))
    }
}

impl MemoryBounds for TraceSamplerConfiguration {
    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
        builder.minimum().with_single_value::<TraceSampler>("component struct");
    }
}

pub struct TraceSampler {
    sampling_rate: f64,
    error_tracking_standalone: bool,
    error_sampling_enabled: bool,
    probabilistic_sampler_enabled: bool,
    otlp_sampling_rate: f64,
    error_sampler: errors::ErrorsSampler,
    priority_sampler: priority_sampler::PrioritySampler,
    no_priority_sampler: score_sampler::NoPrioritySampler,
    rare_sampler: rare_sampler::RareSampler,
}

impl TraceSampler {
    // TODO: merge this with the other duplicate "find root span of trace" functions
    /// Find the root span index of a trace.
    fn get_root_span_index(&self, trace: &Trace) -> Option<usize> {
        // logic taken from here: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/traceutil/trace.go#L36
        let spans = trace.spans();
        if spans.is_empty() {
            return None;
        }
        let length = spans.len();
        // General case: go over all spans and check for one without a matching parent.
        // This intentionally mirrors `datadog-agent/pkg/trace/traceutil/trace.go:GetRoot`:
        // - Fast-path: return the last span with `parent_id == 0` (some clients report the root last)
        // - Otherwise: build a map of `parent_id -> child_span_index`, delete entries whose parent
        //   exists in the trace, and pick any remaining "orphan" child span.
        let mut parent_id_to_child: FastHashMap<u64, usize> = FastHashMap::default();

        for i in 0..length {
            // Common case optimization: check for span with parent_id == 0, starting from the end,
            // since some clients report the root last.
            let j = length - 1 - i;
            if spans[j].parent_id() == 0 {
                return Some(j);
            }
            parent_id_to_child.insert(spans[j].parent_id(), j);
        }

        for span in spans.iter() {
            parent_id_to_child.remove(&span.span_id());
        }

        // Here, if the trace is valid, we should have `len(parent_id_to_child) == 1`.
        if parent_id_to_child.len() != 1 {
            debug!(
                "Didn't reliably find the root span for traceID:{:016x}{:016x}",
                trace.trace_id_high, trace.trace_id_low,
            );
        }

        // Have a safe behavior if that's not the case.
        // Pick a random span without its parent.
        if let Some((_, child_idx)) = parent_id_to_child.iter().next() {
            return Some(*child_idx);
        }

        // Gracefully fail with the last span of the trace.
        Some(length - 1)
    }

    /// Check for user-set sampling priority in trace
    fn get_user_priority(&self, trace: &Trace, root_span_idx: usize) -> Option<i32> {
        // First check trace-level sampling priority (last-seen priority from OTLP ingest)
        if let Some(priority) = trace.priority {
            return Some(priority);
        }

        if trace.spans().is_empty() {
            return None;
        }

        // Fall back to checking spans (for compatibility with non-OTLP traces)
        // Prefer the root span (common case), but fall back to scanning all spans to be robust to ordering.
        if let Some(root) = trace.spans().get(root_span_idx) {
            if let Some(p) = root
                .attributes
                .get(SAMPLING_PRIORITY_METRIC_KEY)
                .and_then(AttributeValue::as_num)
            {
                return Some(p as i32);
            }
        }
        let spans = trace.spans();
        spans.iter().find_map(|span| {
            span.attributes
                .get(SAMPLING_PRIORITY_METRIC_KEY)
                .and_then(AttributeValue::as_num)
                .map(|p| p as i32)
        })
    }

    /// Returns `true` if the given trace ID should be probabilistically sampled.
    fn sample_probabilistic(&self, trace_id: u64) -> bool {
        probabilistic::ProbabilisticSampler::sample(trace_id, self.sampling_rate)
    }

    fn is_otlp_trace(&self, trace: &Trace, root_span_idx: usize) -> bool {
        trace
            .spans()
            .get(root_span_idx)
            .map(|span| {
                span.attributes
                    .contains_key(&MetaString::from_static(OTEL_TRACE_ID_META_KEY))
            })
            .unwrap_or(false)
    }

    /// Returns `true` if the trace contains a span with an error.
    fn trace_contains_error(&self, trace: &Trace, consider_exception_span_events: bool) -> bool {
        trace.spans().iter().any(|span| {
            span.error() != 0 || (consider_exception_span_events && self.span_contains_exception_span_event(span))
        })
    }

    /// Returns `true` if the span has exception span events.
    ///
    /// This checks for the `_dd.span_events.has_exception` meta field set to `"true"`.
    fn span_contains_exception_span_event(&self, span: &Span) -> bool {
        if let Some(has_exception) = span
            .attributes
            .get("_dd.span_events.has_exception")
            .and_then(AttributeValue::as_string)
        {
            return has_exception == "true";
        }
        false
    }

    /// Computes the OTLP pre-sampling priority and decision maker for a trace, mirroring
    /// `OTLPReceiver.createChunks` in DDA which runs before `runSamplersV1`.
    ///
    /// Returns `Some((priority, dm))` for OTLP traces when the probabilistic sampler is disabled,
    /// or `None` if pre-sampling doesn't apply.
    ///
    /// See: https://github.com/DataDog/datadog-agent/blob/be33ac1490c4a34602cbc65a211406b73ad6d00b/pkg/trace/api/otlp.go#L561-L585
    fn otlp_pre_sample(&mut self, trace: &mut Trace, root_span_idx: usize) -> Option<(i32, &'static str)> {
        if self.probabilistic_sampler_enabled || !self.is_otlp_trace(trace, root_span_idx) {
            return None;
        }
        let (priority, dm) = if let Some(user_priority) = self.get_user_priority(trace, root_span_idx) {
            (user_priority, DECISION_MAKER_MANUAL)
        } else {
            let root_trace_id = trace.trace_id_low;
            if sample_by_rate(root_trace_id, self.otlp_sampling_rate) {
                (PRIORITY_AUTO_KEEP, DECISION_MAKER_PROBABILISTIC)
            } else {
                (PRIORITY_AUTO_DROP, DECISION_MAKER_PROBABILISTIC)
            }
        };
        if priority == PRIORITY_AUTO_KEEP {
            if let Some(root_span) = trace.spans_mut().get_mut(root_span_idx) {
                root_span.attributes.remove(PROB_RATE_KEY);
            }
        }
        Some((priority, dm))
    }

    /// Apply analyzed span sampling to the trace.
    ///
    /// Returns `true` if the trace was modified.
    fn analyzed_span_sampling(&self, trace: &mut Trace) -> bool {
        let retained = trace.retain_spans(|_, span| span.attributes.contains_key(KEY_ANALYZED_SPANS));
        if retained > 0 {
            trace.dropped_trace = false;
            trace.priority = Some(PRIORITY_USER_KEEP);
            trace.otlp_sampling_rate = Some(self.sampling_rate);
            true
        } else {
            false
        }
    }

    /// Returns `true` if the given trace has any analyzed spans.
    fn has_analyzed_spans(&self, trace: &Trace) -> bool {
        trace
            .spans()
            .iter()
            .any(|span| span.attributes.contains_key(KEY_ANALYZED_SPANS))
    }

    /// Apply Single Span Sampling to the trace
    /// Returns true if the trace was modified
    fn single_span_sampling(&self, trace: &mut Trace) -> bool {
        let retained = trace.retain_spans(|_, span| span.attributes.contains_key(KEY_SPAN_SAMPLING_MECHANISM));
        if retained > 0 {
            trace.dropped_trace = false;
            trace.priority = Some(PRIORITY_USER_KEEP);
            trace.otlp_sampling_rate = Some(self.sampling_rate);
            true
        } else {
            false
        }
    }

    /// Evaluates the given trace against all configured samplers.
    ///
    /// Return a tuple containing whether or not the trace should be kept, the decision maker tag (which sampler is responsible),
    /// and the index of the root span used for evaluation.
    fn run_samplers(&mut self, trace: &mut Trace) -> (bool, i32, &'static str, Option<usize>) {
        // logic taken from: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/agent.go#L1066
        // Empty trace check
        if trace.spans().is_empty() {
            return (false, PRIORITY_AUTO_DROP, "", None);
        }

        let now = std::time::SystemTime::now();
        let Some(root_span_idx) = self.get_root_span_index(trace) else {
            return (false, PRIORITY_AUTO_DROP, "", None);
        };

        // ETS: only sample traces containing errors (including exception span events); skip all other samplers.
        // logic taken from: https://github.com/DataDog/datadog-agent/blob/be33ac1490c4a34602cbc65a211406b73ad6d00b/pkg/trace/agent/agent.go#L1068
        if self.error_tracking_standalone {
            let otlp_pre_sample = self.otlp_pre_sample(trace, root_span_idx);
            if self.trace_contains_error(trace, true) {
                let keep = self.error_sampler.sample_error(now, trace, root_span_idx);
                let default_priority = if keep { PRIORITY_AUTO_KEEP } else { PRIORITY_AUTO_DROP };
                let (priority, dm) = otlp_pre_sample.unwrap_or((default_priority, ""));
                return (keep, priority, dm, Some(root_span_idx));
            }
            let (pre_priority, pre_dm) = otlp_pre_sample.unwrap_or((PRIORITY_AUTO_DROP, ""));
            return (false, pre_priority, pre_dm, Some(root_span_idx));
        }

        let contains_error = self.trace_contains_error(trace, false);

        // Run the rare sampler early, before all other samplers. This mirrors the Go agent behavior
        // where the rare sampler runs first to catch traces that would otherwise be dropped entirely.
        // logic taken from: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/agent.go#L1078
        let rare = self.rare_sampler.sample(trace, root_span_idx);

        // Modern path: ProbabilisticSamplerEnabled = true
        if self.probabilistic_sampler_enabled {
            let mut prob_keep = false;
            let mut decision_maker = "";

            if rare {
                // Rare sampler wins over probabilistic sampling.
                prob_keep = true;
            } else {
                // Run probabilistic sampler - use trace ID
                let root_trace_id = trace.trace_id_low;
                if self.sample_probabilistic(root_trace_id) {
                    decision_maker = DECISION_MAKER_PROBABILISTIC;
                    prob_keep = true;

                    if let Some(root_span) = trace.spans_mut().get_mut(root_span_idx) {
                        root_span.attributes.insert(
                            MetaString::from(PROB_RATE_KEY),
                            AttributeValue::Float(self.sampling_rate),
                        );
                    }
                } else if self.error_sampling_enabled && contains_error {
                    prob_keep = self.error_sampler.sample_error(now, trace, root_span_idx);
                }
            }

            let priority = if prob_keep {
                PRIORITY_AUTO_KEEP
            } else {
                PRIORITY_AUTO_DROP
            };

            return (prob_keep, priority, decision_maker, Some(root_span_idx));
        }

        let user_priority = self.get_user_priority(trace, root_span_idx);
        if let Some(priority) = user_priority {
            if priority < PRIORITY_AUTO_DROP {
                // Manual drop: short-circuit and skip other samplers.
                return (false, priority, "", Some(root_span_idx));
            }

            if rare {
                return (true, priority, "", Some(root_span_idx));
            }

            if self.priority_sampler.sample(now, trace, root_span_idx, priority, 0.0) {
                return (true, priority, "", Some(root_span_idx));
            }
        } else if self.is_otlp_trace(trace, root_span_idx) {
            // Rare check mirrors agent behavior: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/agent.go#L1129-L1140
            if rare {
                return (true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx));
            }

            // some sampling happens upstream in the otlp receiver in the agent: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/api/otlp.go#L572
            let root_trace_id = trace.trace_id_low;
            if sample_by_rate(root_trace_id, self.otlp_sampling_rate) {
                if let Some(root_span) = trace.spans_mut().get_mut(root_span_idx) {
                    root_span.attributes.remove(PROB_RATE_KEY);
                }
                return (
                    true,
                    PRIORITY_AUTO_KEEP,
                    DECISION_MAKER_PROBABILISTIC,
                    Some(root_span_idx),
                );
            }
        } else {
            if rare {
                return (true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx));
            }
            if self.no_priority_sampler.sample(now, trace, root_span_idx) {
                return (true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx));
            }
        }

        if self.error_sampling_enabled && contains_error {
            let keep = self.error_sampler.sample_error(now, trace, root_span_idx);
            if keep {
                return (true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx));
            }
        }

        // Default: drop the trace
        (false, PRIORITY_AUTO_DROP, "", Some(root_span_idx))
    }

    /// Apply sampling metadata to the trace in-place.
    ///
    /// The `root_span_id` parameter identifies which span should receive the sampling metadata.
    /// This avoids recalculating the root span since it was already found in `run_samplers`.
    fn apply_sampling_metadata(
        &self, trace: &mut Trace, keep: bool, priority: i32, decision_maker: &str, root_span_idx: usize,
    ) {
        let is_otlp = self.is_otlp_trace(trace, root_span_idx);
        let root_span_value = match trace.spans_mut().get_mut(root_span_idx) {
            Some(span) => span,
            None => return,
        };

        // Add tag for the decision maker
        let existing_decision_maker = if decision_maker.is_empty() {
            root_span_value
                .attributes
                .get(TAG_DECISION_MAKER)
                .and_then(AttributeValue::as_string)
                .cloned()
        } else {
            None
        };
        let decision_maker_meta = if decision_maker.is_empty() {
            existing_decision_maker
        } else {
            Some(MetaString::from(decision_maker))
        };

        // When the APM-level probabilistic sampler is used with OTLP traces, the DD Agent writes
        // _dd.p.dm to trace chunk tags only (not span meta). For the legacy OTLP sampling path,
        // it is written to both. We match that behavior by skipping the span meta write only when
        // both conditions hold; the DM value still flows through trace fields to the encoder.
        if priority > 0 && !(is_otlp && self.probabilistic_sampler_enabled) {
            if let Some(dm) = decision_maker_meta.as_ref() {
                root_span_value
                    .attributes
                    .insert(MetaString::from(TAG_DECISION_MAKER), AttributeValue::String(dm.clone()));
            }
        }

        // Now set sampling metadata directly on the trace.
        trace.dropped_trace = !keep;
        trace.priority = Some(priority);
        trace.decision_maker = if priority > 0 { decision_maker_meta } else { None };
        trace.otlp_sampling_rate = Some(if is_otlp {
            self.otlp_sampling_rate
        } else {
            self.sampling_rate
        });
    }

    fn process_trace(&mut self, trace: &mut Trace) -> bool {
        // keep is a boolean that indicates if the trace should be kept or dropped
        // priority is the sampling priority
        // decision_maker is the tag that indicates the decision maker (probabilistic, error, etc.)
        // root_span_idx is the index of the root span of the trace
        let (keep, priority, decision_maker, root_span_idx) = self.run_samplers(trace);

        // Apply sampling metadata and forward if kept, or if ETS (dropped non-error traces are
        // forwarded with DroppedTrace=true, suppressing SSS/analytics).
        if keep || self.error_tracking_standalone {
            if let Some(root_idx) = root_span_idx {
                self.apply_sampling_metadata(trace, keep, priority, decision_maker, root_idx);
            }
            return true;
        }

        // logic taken from here: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/agent.go#L980-L990
        // try single span sampling (keeps spans marked for sampling when trace would be dropped)
        let modified = self.single_span_sampling(trace);
        if !modified {
            // Fall back to analytics events if no SSS spans
            if self.analyzed_span_sampling(trace) {
                return true;
            }
        } else if self.has_analyzed_spans(trace) {
            // Warn about both SSS and analytics events
            debug!(
                "Detected both analytics events AND single span sampling in the same trace. Single span sampling wins because App Analytics is deprecated."
            );
            return true;
        }

        // If we modified the trace with SSS, send it
        if modified {
            return true;
        }

        // Neither SSS nor analytics events found, drop the trace
        debug!("Dropping trace with priority {}", priority);
        false
    }
}

impl SynchronousTransform for TraceSampler {
    fn transform_buffer(&mut self, buffer: &mut EventsBuffer) {
        buffer.remove_if(|event| match event {
            Event::Trace(trace) => !self.process_trace(trace),
            _ => false,
        });
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use saluki_core::data_model::event::trace::{AttributeValue, Span as DdSpan, Trace};
    const PRIORITY_USER_DROP: i32 = -1;

    use super::*;
    fn create_test_sampler() -> TraceSampler {
        TraceSampler {
            sampling_rate: 1.0,
            error_sampling_enabled: true,
            error_tracking_standalone: false,
            probabilistic_sampler_enabled: true,
            otlp_sampling_rate: 1.0,
            error_sampler: errors::ErrorsSampler::new(10.0, 1.0),
            priority_sampler: priority_sampler::PrioritySampler::new(MetaString::from("agent-env"), 1.0, 10.0),
            no_priority_sampler: score_sampler::NoPrioritySampler::new(10.0, 1.0),
            rare_sampler: rare_sampler::RareSampler::new(false, 5.0, std::time::Duration::from_secs(300), 200),
        }
    }

    fn create_test_span(span_id: u64, error: i32) -> DdSpan {
        DdSpan::new(
            MetaString::from("test-service"),
            MetaString::from("test-operation"),
            MetaString::from("test-resource"),
            MetaString::from("test-type"),
            span_id,
            0,    // parent_id
            0,    // start
            1000, // duration
            error,
        )
    }

    fn create_test_span_with_metrics(span_id: u64, metrics: HashMap<String, f64>) -> DdSpan {
        let attrs: saluki_common::collections::FastHashMap<MetaString, AttributeValue> = metrics
            .into_iter()
            .map(|(k, v)| (MetaString::from(k), AttributeValue::Float(v)))
            .collect();
        create_test_span(span_id, 0).with_attributes(attrs)
    }

    #[allow(dead_code)]
    fn create_test_span_with_meta(span_id: u64, meta: HashMap<String, String>) -> DdSpan {
        let attrs: saluki_common::collections::FastHashMap<MetaString, AttributeValue> = meta
            .into_iter()
            .map(|(k, v)| (MetaString::from(k), AttributeValue::String(MetaString::from(v))))
            .collect();
        create_test_span(span_id, 0).with_attributes(attrs)
    }

    fn create_test_trace(spans: Vec<DdSpan>) -> Trace {
        Trace::new(spans)
    }

    #[test]
    fn user_priority_detection() {
        let sampler = create_test_sampler();

        // Test trace with user-set priority = 2 (UserKeep)
        let mut metrics = HashMap::new();
        metrics.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), 2.0);
        let span = create_test_span_with_metrics(1, metrics);
        let trace = create_test_trace(vec![span]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();

        assert_eq!(sampler.get_user_priority(&trace, root_idx), Some(2));

        // Test trace with user-set priority = -1 (UserDrop)
        let mut metrics = HashMap::new();
        metrics.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), -1.0);
        let span = create_test_span_with_metrics(1, metrics);
        let trace = create_test_trace(vec![span]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();

        assert_eq!(sampler.get_user_priority(&trace, root_idx), Some(-1));

        // Test trace without user priority
        let span = create_test_span(1, 0);
        let trace = create_test_trace(vec![span]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();

        assert_eq!(sampler.get_user_priority(&trace, root_idx), None);
    }

    #[test]
    fn trace_level_priority_takes_precedence() {
        let sampler = create_test_sampler();

        // Test trace-level priority overrides span priorities (last-seen priority)
        // Create spans with different priorities - root has 0, later span has 2
        let mut metrics_root = HashMap::new();
        metrics_root.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), 0.0);
        let root_span = create_test_span_with_metrics(1, metrics_root);

        let mut metrics_later = HashMap::new();
        metrics_later.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), 1.0);
        let later_span = create_test_span_with_metrics(2, metrics_later).with_parent_id(1);

        let mut trace = create_test_trace(vec![root_span, later_span]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();

        // Without trace-level priority, should get priority from root (0)
        assert_eq!(sampler.get_user_priority(&trace, root_idx), Some(0));

        // Now set trace-level priority to 2 (simulating last-seen priority from OTLP translator)
        trace.priority = Some(2);

        // Trace-level priority should take precedence
        assert_eq!(sampler.get_user_priority(&trace, root_idx), Some(2));

        // Test that trace-level priority is used even when no span has priority
        let span_no_priority = create_test_span(3, 0);
        let mut trace_only_trace_level = create_test_trace(vec![span_no_priority]);
        trace_only_trace_level.priority = Some(1);
        let root_idx = sampler.get_root_span_index(&trace_only_trace_level).unwrap();

        assert_eq!(sampler.get_user_priority(&trace_only_trace_level, root_idx), Some(1));
    }

    #[test]
    fn manual_keep_with_trace_level_priority() {
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false; // Use legacy path that checks user priority

        // Test that manual keep (priority = 2) works via trace-level priority
        let span = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![span]);
        trace.priority = Some(PRIORITY_USER_KEEP);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, PRIORITY_USER_KEEP);
        assert_eq!(decision_maker, "");

        // Test manual drop (priority = -1) via trace-level priority
        let span = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![span]);
        trace.priority = Some(PRIORITY_USER_DROP);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep); // Should not keep when user drops
        assert_eq!(priority, PRIORITY_USER_DROP);

        // Test that priority = 1 (auto keep) via trace-level is also respected
        let span = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![span]);
        trace.priority = Some(PRIORITY_AUTO_KEEP);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "");
    }

    #[test]
    fn probabilistic_sampling_known_decisions() {
        // The bucketed probabilistic sampler is fully deterministic: it hashes the trace ID into one of 0x4000
        // buckets and keeps the trace when `bucket < (rate * 0x4000)`. These cases pin the exact keep/drop decision
        // for known trace IDs at known rates, so a regression in the hash, the bucket mask, or the comparison is
        // caught (a determinism-only check would not catch any of those).
        //
        // Expected values were computed directly from the FNV-1a bucket math in `ProbabilisticSampler::sample`
        // (mirrors datadog-agent/pkg/trace/sampler/probabilistic.go). For reference, the trace IDs below hash to
        // these buckets (out of 0x4000 = 16384): 0x1234567890ABCDEF -> 1764, 0x0 -> 9301, u64::MAX -> 12365.
        struct Case {
            trace_id: u64,
            rate: f64,
            expected_keep: bool,
        }

        let cases = [
            // rate 1.0 keeps every trace (the maximum bucket, 16383, is always below 16384).
            Case {
                trace_id: 0x1234567890ABCDEF,
                rate: 1.0,
                expected_keep: true,
            },
            // rate 0.0 drops every trace (no bucket is below 0).
            Case {
                trace_id: 0x1234567890ABCDEF,
                rate: 0.0,
                expected_keep: false,
            },
            // Same trace ID (bucket 1764) flips from drop to keep as the rate crosses its bucket ratio (~0.108).
            Case {
                trace_id: 0x1234567890ABCDEF,
                rate: 0.10,
                expected_keep: false,
            },
            Case {
                trace_id: 0x1234567890ABCDEF,
                rate: 0.20,
                expected_keep: true,
            },
            // Trace ID 0 (bucket 9301) straddles rate 0.5 (scaled bucket 8192) vs 0.6 (scaled bucket 9830).
            Case {
                trace_id: 0,
                rate: 0.50,
                expected_keep: false,
            },
            Case {
                trace_id: 0,
                rate: 0.60,
                expected_keep: true,
            },
            // u64::MAX (bucket 12365) straddles rate 0.5 vs 0.8.
            Case {
                trace_id: u64::MAX,
                rate: 0.50,
                expected_keep: false,
            },
            Case {
                trace_id: u64::MAX,
                rate: 0.80,
                expected_keep: true,
            },
        ];

        for case in cases {
            let mut sampler = create_test_sampler();
            sampler.sampling_rate = case.rate;
            assert_eq!(
                sampler.sample_probabilistic(case.trace_id),
                case.expected_keep,
                "trace_id={:#018x} rate={}",
                case.trace_id,
                case.rate
            );
        }
    }

    #[test]
    fn probabilistic_sampling_is_deterministic() {
        // Determinism is a documented property of `ProbabilisticSampler::sample` (same trace ID + rate always yields
        // the same decision). This is intentionally a determinism-only check; correctness is covered by
        // `test_probabilistic_sampling_known_decisions`.
        let sampler = create_test_sampler();
        let trace_id = 0x1234567890ABCDEF_u64;
        assert_eq!(
            sampler.sample_probabilistic(trace_id),
            sampler.sample_probabilistic(trace_id)
        );
    }

    #[test]
    fn error_detection() {
        let sampler = create_test_sampler();

        // Test trace with error field set
        let span_with_error = create_test_span(1, 1);
        let trace = create_test_trace(vec![span_with_error]);
        assert!(sampler.trace_contains_error(&trace, false));

        // Test trace without error
        let span_without_error = create_test_span(1, 0);
        let trace = create_test_trace(vec![span_without_error]);
        assert!(!sampler.trace_contains_error(&trace, false));
    }

    #[test]
    fn sampling_priority_order() {
        // Test modern path: error sampler overrides probabilistic drop
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 0.5; // 50% sampling rate
        sampler.probabilistic_sampler_enabled = true;

        // Create trace with error that would be dropped by probabilistic
        // Using a trace ID that we know will be dropped at 50% rate
        let span_with_error = create_test_span(1, 1);
        let mut trace = create_test_trace(vec![span_with_error]);
        trace.trace_id_low = u64::MAX - 1;

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, ""); // Error sampler doesn't set decision_maker

        // Test legacy path: user priority is respected
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false; // Use legacy path

        let mut metrics = HashMap::new();
        metrics.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), 2.0);
        let span = create_test_span_with_metrics(1, metrics);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, 2); // UserKeep
        assert_eq!(decision_maker, "");
    }

    #[test]
    fn empty_trace_handling() {
        let mut sampler = create_test_sampler();
        let mut trace = create_test_trace(vec![]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep);
        assert_eq!(priority, PRIORITY_AUTO_DROP);
    }

    #[test]
    fn root_span_detection() {
        let sampler = create_test_sampler();

        // Test 1: Root span with parent_id = 0 (common case)
        let root_span = DdSpan::new(
            MetaString::from("service"),
            MetaString::from("operation"),
            MetaString::from("resource"),
            MetaString::from("type"),
            1,
            0, // parent_id = 0 indicates root
            0,
            1000,
            0,
        );
        let child_span = DdSpan::new(
            MetaString::from("service"),
            MetaString::from("child_op"),
            MetaString::from("resource"),
            MetaString::from("type"),
            2,
            1, // parent_id = 1 (points to root)
            100,
            500,
            0,
        );
        // Put root span second to test that we find it even when not first
        let trace = create_test_trace(vec![child_span.clone(), root_span.clone()]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        assert_eq!(trace.spans()[root_idx].span_id(), 1);

        // Test 2: Orphaned span (parent not in trace)
        let orphan_span = DdSpan::new(
            MetaString::from("service"),
            MetaString::from("orphan"),
            MetaString::from("resource"),
            MetaString::from("type"),
            3,
            999, // parent_id = 999 (doesn't exist in trace)
            200,
            300,
            0,
        );
        let trace = create_test_trace(vec![orphan_span]);
        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        assert_eq!(trace.spans()[root_idx].span_id(), 3);

        // Test 3: Multiple root candidates: should return the last one found (index 1)
        let span1 = create_test_span(1, 0);
        let span2 = create_test_span(2, 0);
        let trace = create_test_trace(vec![span1, span2]);
        // Both have parent_id = 0, should return the last one found (span_id = 2)
        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        assert_eq!(trace.spans()[root_idx].span_id(), 2);
    }

    #[test]
    fn single_span_sampling() {
        let mut sampler = create_test_sampler();

        // Test 1: Trace with SSS tags should be kept even when probabilistic would drop it
        sampler.sampling_rate = 0.0; // 0% sampling rate - should drop everything
        sampler.probabilistic_sampler_enabled = true;

        // Create span with SSS metric
        let mut attrs_map = saluki_common::collections::FastHashMap::default();
        attrs_map.insert(
            MetaString::from(KEY_SPAN_SAMPLING_MECHANISM),
            AttributeValue::Float(8.0),
        );
        let sss_span = create_test_span(1, 0).with_attributes(attrs_map.clone());

        // Create regular span without SSS
        let regular_span = create_test_span(2, 0);

        let mut trace = create_test_trace(vec![sss_span.clone(), regular_span]);

        // Apply SSS
        let modified = sampler.single_span_sampling(&mut trace);
        assert!(modified);
        assert_eq!(trace.spans().len(), 1); // Only SSS span kept
        assert_eq!(trace.spans()[0].span_id(), 1); // It's the SSS span

        // Check that trace has been marked as kept with high priority
        assert_eq!(trace.priority, Some(PRIORITY_USER_KEEP));

        // Test 2: Trace without SSS tags should not be modified
        let trace_without_sss = create_test_trace(vec![create_test_span(3, 0)]);
        let mut trace_copy = trace_without_sss.clone();
        let modified = sampler.single_span_sampling(&mut trace_copy);
        assert!(!modified);
        assert_eq!(trace_copy.spans().len(), trace_without_sss.spans().len());
    }

    #[test]
    fn analytics_events() {
        let sampler = create_test_sampler();

        // Test 1: Trace with analyzed spans
        let mut attrs_map = saluki_common::collections::FastHashMap::default();
        attrs_map.insert(MetaString::from(KEY_ANALYZED_SPANS), AttributeValue::Float(1.0));
        let analyzed_span = create_test_span(1, 0).with_attributes(attrs_map.clone());
        let regular_span = create_test_span(2, 0);

        let mut trace = create_test_trace(vec![analyzed_span.clone(), regular_span]);

        let analyzed_span_ids: Vec<u64> = trace
            .spans()
            .iter()
            .filter(|span| span.attributes.contains_key(KEY_ANALYZED_SPANS))
            .map(|span| span.span_id())
            .collect();
        assert_eq!(analyzed_span_ids, vec![1]);

        assert!(sampler.has_analyzed_spans(&trace));
        let modified = sampler.analyzed_span_sampling(&mut trace);
        assert!(modified);
        assert_eq!(trace.spans().len(), 1);
        assert_eq!(trace.spans()[0].span_id(), 1);
        assert_eq!(trace.priority, Some(PRIORITY_USER_KEEP));

        // Test 2: Trace without analyzed spans
        let trace_no_analytics = create_test_trace(vec![create_test_span(3, 0)]);
        let mut trace_no_analytics_copy = trace_no_analytics.clone();
        let analyzed_span_ids: Vec<u64> = trace_no_analytics
            .spans()
            .iter()
            .filter(|span| span.attributes.contains_key(KEY_ANALYZED_SPANS))
            .map(|span| span.span_id())
            .collect();
        assert!(analyzed_span_ids.is_empty());
        assert!(!sampler.has_analyzed_spans(&trace_no_analytics));
        let modified = sampler.analyzed_span_sampling(&mut trace_no_analytics_copy);
        assert!(!modified);
        assert_eq!(trace_no_analytics_copy.spans().len(), trace_no_analytics.spans().len());
    }

    #[test]
    fn probabilistic_sampling_with_prob_rate_key() {
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 0.75; // 75% sampling rate
        sampler.probabilistic_sampler_enabled = true;

        // Use a trace ID that we know will be sampled
        let trace_id = 12345_u64;
        let root_span = DdSpan::new(
            MetaString::from("service"),
            MetaString::from("operation"),
            MetaString::from("resource"),
            MetaString::from("type"),
            1,
            0, // parent_id = 0 indicates root
            0,
            1000,
            0,
        );
        let mut trace = create_test_trace(vec![root_span]);
        trace.trace_id_low = trace_id;

        let (keep, priority, decision_maker, root_span_idx) = sampler.run_samplers(&mut trace);

        if keep && decision_maker == DECISION_MAKER_PROBABILISTIC {
            // If sampled probabilistically, check that probRateKey was already added
            assert_eq!(priority, PRIORITY_AUTO_KEEP);
            assert_eq!(decision_maker, DECISION_MAKER_PROBABILISTIC); // probabilistic sampling marker

            // Check that the root span already has the probRateKey (it should have been added in run_samplers)
            let root_idx = root_span_idx.unwrap_or(0);
            let root_span = &trace.spans()[root_idx];
            assert!(root_span.attributes.contains_key(PROB_RATE_KEY));
            assert_eq!(
                root_span
                    .attributes
                    .get(PROB_RATE_KEY)
                    .and_then(AttributeValue::as_float),
                Some(0.75)
            );

            // Test that apply_sampling_metadata still works correctly for other metadata
            let mut trace_with_metadata = trace.clone();
            sampler.apply_sampling_metadata(&mut trace_with_metadata, keep, priority, decision_maker, root_idx);

            // Check that decision maker tag was added
            let modified_root = &trace_with_metadata.spans()[root_idx];
            assert!(modified_root.attributes.contains_key(TAG_DECISION_MAKER));
            assert_eq!(
                modified_root
                    .attributes
                    .get(TAG_DECISION_MAKER)
                    .and_then(AttributeValue::as_string),
                Some(&MetaString::from(DECISION_MAKER_PROBABILISTIC))
            );
        }
    }

    // ── Rare-sampler interaction tests ──────────────────────────────────────────
    // Adapted from datadog-agent/pkg/trace/agent/agent_test.go TestSampling cases:
    // "rare-sampler-catch-unsampled", "rare-sampler-catch-sampled",
    // "rare-sampler-disabled", and related probabilistic path interactions.

    /// Create a top-level span eligible for rare sampling.
    ///
    /// The rare sampler only considers spans that have `_top_level=1` or `_dd.measured=1`.
    /// This helper sets `_top_level=1` so that the rare sampler can consider the span.
    fn create_top_level_span(span_id: u64) -> DdSpan {
        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
        create_test_span(span_id, 0).with_attributes(attrs)
    }

    /// Create a `TraceSampler` with the rare sampler enabled and a very high TPS limit so it
    /// freely samples first occurrences, plus a long TTL so second occurrences stay within TTL.
    fn create_sampler_with_rare_enabled() -> TraceSampler {
        TraceSampler {
            rare_sampler: rare_sampler::RareSampler::new(true, 1000.0, std::time::Duration::from_secs(300), 200),
            ..create_test_sampler()
        }
    }

    /// Adapted from Go "rare-sampler-catch-unsampled":
    ///
    /// Rare is enabled + probabilistic would drop → rare catches it (first occurrence).
    #[test]
    fn rare_sampler_catches_unsampled_trace() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.sampling_rate = 0.0; // probabilistic drops everything
        sampler.probabilistic_sampler_enabled = true;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "rare sampler should catch first occurrence");
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "", "rare sampler does not set _dd.p.dm");
    }

    /// Adapted from Go "rare-sampler-catch-sampled" (first trace):
    ///
    /// Rare is enabled, first occurrence—trace is kept and `_dd.rare` is set on the span.
    #[test]
    fn rare_sampler_sets_rare_metric_on_first_occurrence() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.sampling_rate = 0.0;
        sampler.probabilistic_sampler_enabled = true;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, _, _, root_idx) = sampler.run_samplers(&mut trace);
        assert!(keep);
        let root = &trace.spans()[root_idx.unwrap()];
        assert_eq!(
            root.attributes
                .get(rare_sampler::RARE_KEY)
                .and_then(AttributeValue::as_float),
            Some(1.0),
            "_dd.rare should be 1 on first occurrence"
        );
    }

    /// Adapted from Go "rare-sampler-catch-sampled" (second trace same signature):
    ///
    /// Within the TTL, the same signature is no longer "rare" and rare doesn't re-sample it.
    /// With probabilistic at 0%, the trace should be dropped.
    #[test]
    fn rare_sampler_does_not_resample_within_ttl() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.sampling_rate = 0.0;
        sampler.probabilistic_sampler_enabled = true;

        // First trace: rare catches it.
        let span1 = create_top_level_span(1);
        let mut trace1 = create_test_trace(vec![span1]);
        let (keep1, _, _, _) = sampler.run_samplers(&mut trace1);
        assert!(keep1, "first occurrence should be kept by rare sampler");

        // Second trace: same signature (same service/operation/resource on the top-level span),
        // still within TTL → rare won't catch it; probabilistic at 0% drops it.
        let span2 = create_top_level_span(2);
        let mut trace2 = create_test_trace(vec![span2]);
        let (keep2, priority2, _, _) = sampler.run_samplers(&mut trace2);
        assert!(!keep2, "second occurrence within TTL should be dropped");
        assert_eq!(priority2, PRIORITY_AUTO_DROP);
    }

    /// Adapted from Go "rare-sampler-disabled":
    ///
    /// Rare is disabled + probabilistic at 0% → trace is dropped.
    #[test]
    fn rare_sampler_disabled_does_not_catch_unsampled() {
        let mut sampler = create_test_sampler(); // rare disabled by default
        sampler.sampling_rate = 0.0;
        sampler.probabilistic_sampler_enabled = true;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "rare disabled should not catch the trace");
        assert_eq!(priority, PRIORITY_AUTO_DROP);
    }

    /// Rare + non-probabilistic path (priority path): rare catches `PriorityAutoDrop` on first
    /// occurrence, preserving the tracer-set priority rather than upgrading to AutoKeep.
    #[test]
    fn rare_sampler_catches_priority_auto_drop_in_legacy_path() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.probabilistic_sampler_enabled = false;

        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
        attrs.insert(
            MetaString::from(SAMPLING_PRIORITY_METRIC_KEY),
            AttributeValue::Float(PRIORITY_AUTO_DROP as f64),
        );
        let span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "rare sampler should catch PriorityAutoDrop on first occurrence");
        assert_eq!(priority, PRIORITY_AUTO_DROP, "tracer-set priority should be preserved");
        assert_eq!(decision_maker, "");
    }

    /// Rare + non-probabilistic path (priority path): UserKeep priority is preserved, not
    /// downgraded to AutoKeep. Mirrors Go agent behavior at agent.go#L1129-1131.
    #[test]
    fn rare_sampler_preserves_user_keep_priority_in_legacy_path() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.probabilistic_sampler_enabled = false;

        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
        attrs.insert(
            MetaString::from(SAMPLING_PRIORITY_METRIC_KEY),
            AttributeValue::Float(2.0),
        ); // UserKeep
        let span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, 2, "UserKeep priority must not be downgraded to AutoKeep");
    }

    /// Probabilistic path with 100% rate and rare disabled: keep with `_dd.p.dm = "-9"`.
    #[test]
    fn probabilistic_100_percent_keeps_trace_with_decision_maker() {
        let mut sampler = create_test_sampler(); // rare disabled
        sampler.sampling_rate = 1.0;
        sampler.probabilistic_sampler_enabled = true;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, DECISION_MAKER_PROBABILISTIC);
    }

    /// Probabilistic path with 0% rate and rare disabled: drop.
    #[test]
    fn probabilistic_0_percent_drops_trace() {
        let mut sampler = create_test_sampler(); // rare disabled
        sampler.sampling_rate = 0.0;
        sampler.probabilistic_sampler_enabled = true;
        sampler.error_sampling_enabled = false;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep);
        assert_eq!(priority, PRIORITY_AUTO_DROP);
    }

    /// Rare sampler should catch OTLP traces without a sampling priority on their first occurrence,
    /// matching the Go agent behavior: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/agent.go#L1129-L1140
    #[test]
    fn rare_sampler_catches_otlp_no_priority_trace() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.probabilistic_sampler_enabled = false;
        sampler.error_sampling_enabled = false;
        sampler.otlp_sampling_rate = 0.0;

        let mut span = create_top_level_span(1);
        span.attributes.insert(
            MetaString::from_static(OTEL_TRACE_ID_META_KEY),
            AttributeValue::String(MetaString::from("00000000000000000000000000000001")),
        );
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, root_idx) = sampler.run_samplers(&mut trace);
        assert!(
            keep,
            "rare sampler should keep OTLP trace with no priority on first occurrence"
        );
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "");
        assert_eq!(
            trace.spans()[root_idx.unwrap()]
                .attributes
                .get(rare_sampler::RARE_KEY)
                .and_then(AttributeValue::as_float),
            Some(1.0),
            "_dd.rare should be set to 1 on first occurrence"
        );
    }

    /// Adapted from Go "probabilistic-rare-100":
    ///
    /// Rare fires before probabilistic is consulted, so even at 100% sampling rate the decision
    /// maker tag isn't set—the trace is attributed to rare, not probabilistic.
    #[test]
    fn rare_wins_over_probabilistic_no_decision_maker_tag() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.sampling_rate = 1.0;
        sampler.probabilistic_sampler_enabled = true;

        let span = create_top_level_span(1);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep);
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "", "rare takes precedence—_dd.p.dm must not be set");
    }

    /// Adapted from Go "error-sampled-prio-unsampled":
    ///
    /// Rare fires before the error sampler is reached. A no-priority error trace on its first
    /// occurrence is kept by rare, not by the error sampler.
    #[test]
    fn rare_catches_error_trace_before_error_sampler() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.probabilistic_sampler_enabled = false;
        sampler.error_sampling_enabled = true;

        let span = create_top_level_span(1);
        let error_span = create_test_span(2, 1); // error=1
        let mut trace = create_test_trace(vec![span, error_span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "rare should catch the trace before the error sampler");
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "");
    }

    /// Adapted from Go manual-drop short-circuit behavior:
    ///
    /// UserDrop (-1) priority is checked before rare runs in the priority path. A UserDrop trace
    /// must be dropped even when rare is enabled and would otherwise match.
    #[test]
    fn manual_drop_short_circuits_before_rare() {
        let mut sampler = create_sampler_with_rare_enabled();
        sampler.probabilistic_sampler_enabled = false;

        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
        attrs.insert(
            MetaString::from(SAMPLING_PRIORITY_METRIC_KEY),
            AttributeValue::Float(-1.0),
        ); // UserDrop
        let span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "UserDrop must be dropped even when rare would match");
        assert_eq!(priority, -1);
    }

    // ── Error Tracking Standalone tests ─────────────────────────────────────────
    // Adapted from datadog-agent/pkg/trace/agent/agent.go runSamplers ETS block.

    fn create_sampler_with_ets() -> TraceSampler {
        TraceSampler {
            error_tracking_standalone: true,
            ..create_test_sampler()
        }
    }

    /// ETS enabled + trace with error → kept by error sampler.
    #[test]
    fn ets_keeps_trace_with_error() {
        let mut sampler = create_sampler_with_ets();

        let span = create_test_span(1, 1); // error=1
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "ETS should keep traces with errors");
        assert_eq!(priority, PRIORITY_AUTO_KEEP);
        assert_eq!(decision_maker, "", "ETS does not set a decision maker");
    }

    /// ETS enabled + trace without error → dropped; rare/probabilistic/priority not consulted.
    #[test]
    fn ets_drops_trace_without_error() {
        let mut sampler = create_sampler_with_ets();

        let span = create_test_span(1, 0); // error=0
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, _, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "ETS should drop traces without errors");
        assert_eq!(priority, PRIORITY_AUTO_DROP);
    }

    /// ETS enabled + non-error trace → forwarded with DroppedTrace=true; SSS/analytics suppressed.
    #[test]
    fn ets_forwards_dropped_trace_with_dropped_flag() {
        let mut sampler = create_sampler_with_ets();

        // Span with SSS metric — would trigger single span sampling in non-ETS mode.
        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(
            MetaString::from(KEY_SPAN_SAMPLING_MECHANISM),
            AttributeValue::Float(8.0),
        );
        let span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![span]);

        let forwarded = sampler.process_trace(&mut trace);
        assert!(forwarded, "ETS should forward non-error traces to intake");
        assert!(trace.dropped_trace, "non-error ETS trace should have DroppedTrace=true");
    }

    /// ETS enabled + trace with exception span event → kept (exception events count as errors in ETS).
    #[test]
    fn ets_keeps_trace_with_exception_span_event() {
        let mut sampler = create_sampler_with_ets();

        // Span with error=0 but exception span event metadata.
        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(
            MetaString::from("_dd.span_events.has_exception"),
            AttributeValue::String(MetaString::from("true")),
        );
        let span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![span]);

        let (keep, _, _, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "ETS should treat exception span events as errors");
    }

    /// ETS disabled → normal sampling path (probabilistic) is used.
    #[test]
    fn ets_disabled_uses_normal_sampling() {
        let mut sampler = create_test_sampler(); // ETS disabled
        sampler.sampling_rate = 1.0;
        sampler.probabilistic_sampler_enabled = true;

        let span = create_test_span(1, 0); // no error
        let mut trace = create_test_trace(vec![span]);

        let (keep, _, decision_maker, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "normal probabilistic sampling should keep the trace");
        assert_eq!(decision_maker, DECISION_MAKER_PROBABILISTIC);
    }

    // ── ETS + OTLP pre-sampling tests ────────────────────────────────────────────
    // These mirror DDA's OTLPReceiver.createChunks behavior which pre-assigns
    // priority/dm before runSamplersV1, so ETS sees those values even when it
    // short-circuits. See: pkg/trace/api/otlp.go#L561-L585.

    fn create_otlp_test_span(span_id: u64, error: i32) -> DdSpan {
        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(
            MetaString::from_static(OTEL_TRACE_ID_META_KEY),
            AttributeValue::String(MetaString::from("0000000000000000deadbeefcafebabe")),
        );
        create_test_span(span_id, error).with_attributes(attrs)
    }

    fn create_sampler_with_ets_legacy() -> TraceSampler {
        TraceSampler {
            error_tracking_standalone: true,
            probabilistic_sampler_enabled: false,
            otlp_sampling_rate: 1.0,
            ..create_test_sampler()
        }
    }

    /// ETS + OTLP non-error trace (legacy sampler path): pre-sampling sets priority=AutoKeep and dm=-9.
    /// Mirrors DDA `OTLPReceiver` assigning priority=1 + dm=-9 before ETS returns early.
    #[test]
    fn ets_otlp_non_error_gets_presample_priority_and_dm() {
        let mut sampler = create_sampler_with_ets_legacy();

        let span = create_otlp_test_span(1, 0); // no error
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, dm, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "ETS should drop non-error OTLP traces");
        assert_eq!(
            priority, PRIORITY_AUTO_KEEP,
            "OTLP pre-sampling sets priority=AutoKeep even for ETS-dropped traces"
        );
        assert_eq!(dm, DECISION_MAKER_PROBABILISTIC, "OTLP pre-sampling sets dm=-9");
    }

    /// ETS + OTLP error trace (legacy sampler path): pre-sampling sets priority=AutoKeep and dm=-9.
    #[test]
    fn ets_otlp_error_gets_presample_priority_and_dm() {
        let mut sampler = create_sampler_with_ets_legacy();

        let span = create_otlp_test_span(1, 1); // error=1
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, dm, _) = sampler.run_samplers(&mut trace);
        assert!(keep, "ETS should keep error OTLP traces");
        assert_eq!(priority, PRIORITY_AUTO_KEEP, "OTLP pre-sampling sets priority=AutoKeep");
        assert_eq!(dm, DECISION_MAKER_PROBABILISTIC, "OTLP pre-sampling sets dm=-9");
    }

    /// ETS + OTLP + probabilistic_sampler_enabled=true: `OTLPReceiver` defers, no pre-sampling.
    /// DDA's `OTLPReceiver` sets PriorityNone and skips when ProbabilisticSamplerEnabled.
    #[test]
    fn ets_otlp_probabilistic_path_skips_presample() {
        let mut sampler = create_sampler_with_ets_legacy();
        sampler.probabilistic_sampler_enabled = true; // override to prob path

        let span = create_otlp_test_span(1, 0); // no error
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, dm, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "ETS should drop non-error traces");
        assert_eq!(
            priority, PRIORITY_AUTO_DROP,
            "no pre-sampling when probabilistic path active"
        );
        assert_eq!(dm, "", "no dm when probabilistic path active");
    }

    /// ETS + non-OTLP trace (legacy sampler path): behavior unchanged—no pre-sampling.
    #[test]
    fn ets_non_otlp_unaffected_by_presample() {
        let mut sampler = create_sampler_with_ets_legacy();

        let span = create_test_span(1, 0); // no error, no OTLP meta
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, dm, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "ETS should drop non-error non-OTLP traces");
        assert_eq!(priority, PRIORITY_AUTO_DROP, "non-OTLP traces use default ETS priority");
        assert_eq!(dm, "", "non-OTLP traces get no dm");
    }

    /// ETS + OTLP trace with user-set priority: dm="-4" (manual sampling), matching DDA.
    #[test]
    fn ets_otlp_user_priority_gets_manual_dm() {
        let mut sampler = create_sampler_with_ets_legacy();

        let mut span = create_otlp_test_span(1, 0);
        span.attributes.insert(
            MetaString::from(SAMPLING_PRIORITY_METRIC_KEY),
            AttributeValue::Float(2.0),
        ); // UserKeep
        let mut trace = create_test_trace(vec![span]);

        let (keep, priority, dm, _) = sampler.run_samplers(&mut trace);
        assert!(!keep, "ETS drops non-error traces regardless of user priority");
        assert_eq!(priority, PRIORITY_USER_KEEP, "user priority is preserved");
        assert_eq!(dm, DECISION_MAKER_MANUAL, "user-set priority gets dm=-4");
    }
}