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 std::sync::LazyLock;

use agent_data_plane_config::domains;
use async_trait::async_trait;
use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
use saluki_core::{
    components::{transforms::*, BuildContext},
    data_model::event::{
        trace::{AttributeValue, Span, Trace},
        Event, EventType,
    },
    topology::OutputDefinition,
};
use saluki_error::GenericError;
use stringtheory::MetaString;
use tokio::select;
use tracing::debug;

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

use self::probabilistic::PROB_RATE_KEY;
use self::telemetry::{DecisionWindow, WINDOW};
use crate::common::datadog::{
    compute_top_level, get_root_span_index, get_trace_env, sample_by_rate, DECISION_MAKER_MANUAL,
    DECISION_MAKER_PROBABILISTIC, OTEL_TRACE_ID_META_KEY, SAMPLING_PRIORITY_METRIC_KEY, TAG_DECISION_MAKER, TAG_ORIGIN,
};

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

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

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 {
    probabilistic_sampler_enabled: bool,
    probabilistic_hash_seed: u32,
    probabilistic_full_trace_id: bool,
    sampling_percentage: f64,
    error_sampling_enabled: bool,
    error_tracking_standalone: bool,
    errors_per_second: f64,
    target_traces_per_second: f64,
    extra_sample_rate: f64,
    max_catalog_entries: usize,
    default_env: MetaString,
    rare_sampler_enabled: bool,
    rare_sampler_tps: f64,
    rare_sampler_cooldown_secs: f64,
    rare_sampler_cardinality: usize,
    otlp_sampling_rate: f64,
    compute_top_level_by_span_kind: bool,
}

impl TraceSamplerConfiguration {
    /// Creates a new `TraceSamplerConfiguration` from the resolved traces domain.
    ///
    /// The OTLP trace settings live in their own domain, so they arrive as a separate slice rather
    /// than through the traces domain.
    pub fn from_configuration(traces: &domains::traces::Domain, otlp_traces: &domains::otlp::Traces) -> Self {
        let otlp_sampling_rate = normalize_sampling_rate(otlp_traces.probabilistic_sampler_sampling_percentage / 100.0);
        Self {
            probabilistic_sampler_enabled: traces.probabilistic_sampler.enabled,
            probabilistic_hash_seed: traces.probabilistic_sampler.hash_seed,
            probabilistic_full_trace_id: traces
                .features
                .contains(&domains::traces::ApmFeature::ProbabilisticSamplerFullTraceId),
            sampling_percentage: traces.probabilistic_sampler.sampling_percentage,
            error_sampling_enabled: traces.error_sampling_enabled,
            error_tracking_standalone: traces.error_tracking_standalone_enabled,
            errors_per_second: traces.errors_per_second,
            target_traces_per_second: traces.target_traces_per_second,
            extra_sample_rate: traces.extra_sample_rate,
            max_catalog_entries: traces.max_catalog_entries,
            default_env: MetaString::from(traces.default_env.clone()),
            rare_sampler_enabled: traces.enable_rare_sampler,
            rare_sampler_tps: traces.rare_sampler.tps,
            rare_sampler_cooldown_secs: traces.rare_sampler.cooldown,
            rare_sampler_cardinality: traces.rare_sampler.cardinality,
            otlp_sampling_rate,
            compute_top_level_by_span_kind: otlp_traces.enable_compute_top_level_by_span_kind,
        }
    }
}

#[async_trait]
impl TransformBuilder for TraceSamplerConfiguration {
    fn input_event_type(&self) -> EventType {
        EventType::Trace
    }

    fn outputs(&self) -> &[OutputDefinition<EventType>] {
        static OUTPUTS: LazyLock<Vec<OutputDefinition<EventType>>> = LazyLock::new(|| {
            vec![
                OutputDefinition::default_output(EventType::Trace),
                OutputDefinition::named_output("metrics", EventType::Metric),
            ]
        });
        &OUTPUTS
    }

    async fn build(&self, _context: BuildContext) -> Result<Box<dyn Transform + Send>, GenericError> {
        // TODO: Need to support remote configuration changing these at runtime
        // See https://github.com/DataDog/saluki/issues/1326
        let telemetry = DecisionWindow::new();

        let sampler = TraceSampler {
            sampling_rate: self.sampling_percentage / 100.0,
            error_sampling_enabled: self.error_sampling_enabled,
            error_tracking_standalone: self.error_tracking_standalone,
            probabilistic_sampler_enabled: self.probabilistic_sampler_enabled,
            probabilistic: probabilistic::ProbabilisticSampler::new(
                self.probabilistic_hash_seed,
                self.probabilistic_full_trace_id,
            ),
            otlp_sampling_rate: self.otlp_sampling_rate,
            error_sampler: errors::ErrorsSampler::new(self.errors_per_second, self.extra_sample_rate),
            priority_sampler: priority_sampler::PrioritySampler::new(
                self.default_env.clone(),
                self.extra_sample_rate,
                self.target_traces_per_second,
                self.max_catalog_entries,
            ),
            no_priority_sampler: score_sampler::NoPrioritySampler::new(
                self.target_traces_per_second,
                self.extra_sample_rate,
            ),
            rare_sampler: rare_sampler::RareSampler::new(
                self.rare_sampler_enabled,
                self.rare_sampler_tps,
                std::time::Duration::from_secs_f64(self.rare_sampler_cooldown_secs),
                self.rare_sampler_cardinality,
                telemetry.counters().clone(),
            ),
            telemetry,
            compute_top_level_by_span_kind: self.compute_top_level_by_span_kind,
        };

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

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

        // The catalog holds a hash map slot plus one LRU slab entry per tracked signature. A
        // configured 0 selects the default cap, matching the catalog's own resolution.
        let catalog_capacity = if self.max_catalog_entries == 0 {
            catalog::MAX_CATALOG_ENTRIES
        } else {
            self.max_catalog_entries
        };
        builder
            .minimum()
            .with_map::<signature::ServiceSignature, u32>("priority sampler catalog map", catalog_capacity)
            .with_map::<signature::ServiceSignature, signature::Signature>(
                "priority sampler catalog entries",
                catalog_capacity,
            );

        // Decision keys span the same (service, env) universe as the catalog, across a handful
        // of samplers. The window resets every ten seconds, so steady-state growth follows active
        // combinations; within a window, growth is capped at the telemetry's decision-key limit,
        // with excess combinations rolled into one bucket per sampler.
        builder
            .minimum()
            .with_map::<telemetry::DecisionKey, telemetry::DecisionCounts>("sampler decision window", catalog_capacity);
    }
}

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

/// The outcome of evaluating a trace against the configured samplers.
struct SamplerOutcome {
    /// Whether the trace should be kept.
    keep: bool,
    /// The final sampling priority.
    priority: i32,
    /// The decision maker tag value for the deciding sampler, if any.
    decision_maker: &'static str,
    /// The index of the root span the samplers evaluated against, when one was found.
    root_span_idx: Option<usize>,
    /// The sampler that decided the trace.
    sampler: telemetry::SamplerName,
}

impl SamplerOutcome {
    const fn new(
        keep: bool, priority: i32, decision_maker: &'static str, root_span_idx: Option<usize>,
        sampler: telemetry::SamplerName,
    ) -> Self {
        Self {
            keep,
            priority,
            decision_maker,
            root_span_idx,
            sampler,
        }
    }
}

impl TraceSampler {
    /// Find the root span index of a trace.
    fn get_root_span_index(&self, trace: &Trace) -> Option<usize> {
        // The shared root finder, so every consumer anchors metadata to identical root spans.
        get_root_span_index(trace.spans())
    }

    /// 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_high: u64, trace_id_low: u64) -> bool {
        self.probabilistic
            .sample(trace_id_high, trace_id_low, 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.
    fn run_samplers(&mut self, trace: &mut Trace) -> SamplerOutcome {
        let outcome = self.run_samplers_inner(trace);

        // Unclaimed decisions carry no sampler identity, so recording them would inflate a
        // series that exists only to expose attribution gaps; empty and rootless traces are the
        // only unclaimed cases.
        if outcome.sampler != telemetry::SamplerName::Unknown {
            let (service, env) = self.decision_dimensions(trace, outcome.root_span_idx);
            self.telemetry
                .record_decision(outcome.keep, outcome.sampler, outcome.priority, service, env);
        }

        outcome
    }

    /// Reads the service and env dimensions a decision is recorded under.
    fn decision_dimensions(&self, trace: &Trace, root_span_idx: Option<usize>) -> (MetaString, MetaString) {
        match root_span_idx {
            Some(root_span_idx) => {
                let service = trace
                    .spans()
                    .get(root_span_idx)
                    .map(|span| MetaString::from(span.service()))
                    .unwrap_or_default();
                let env = get_trace_env(trace, root_span_idx).cloned().unwrap_or_default();
                (service, env)
            }
            None => (MetaString::empty(), MetaString::empty()),
        }
    }

    fn run_samplers_inner(&mut self, trace: &mut Trace) -> SamplerOutcome {
        // The name starts unclaimed; whichever sampler decides the trace claims it, so every
        // decision is recorded under exactly one sampler.
        let mut sampler_name = telemetry::SamplerName::Unknown;

        // 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 SamplerOutcome::new(false, PRIORITY_AUTO_DROP, "", None, sampler_name);
        }

        let Some(root_span_idx) = self.get_root_span_index(trace) else {
            return SamplerOutcome::new(false, PRIORITY_AUTO_DROP, "", None, sampler_name);
        };

        // ETS: only sample traces containing errors (including exception span events); skip all other samplers.
        // logic taken from: https://github.com/DataDog/datadog-agent/blob/be33ac1490c4a34602cbc65a211406b73ad6d00/pkg/trace/agent/agent.go#L1068
        if self.error_tracking_standalone {
            sampler_name = telemetry::SamplerName::Error;
            let otlp_pre_sample = self.otlp_pre_sample(trace, root_span_idx);
            if self.trace_contains_error(trace, true) {
                let now = std::time::SystemTime::now();
                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 SamplerOutcome::new(keep, priority, dm, Some(root_span_idx), sampler_name);
            }
            let (pre_priority, pre_dm) = otlp_pre_sample.unwrap_or((PRIORITY_AUTO_DROP, ""));
            return SamplerOutcome::new(false, pre_priority, pre_dm, Some(root_span_idx), sampler_name);
        }

        // 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 {
            sampler_name = telemetry::SamplerName::Probabilistic;
            let mut prob_keep = false;
            let mut decision_maker = "";

            if rare {
                // Rare sampler wins over probabilistic sampling.
                sampler_name = telemetry::SamplerName::Rare;
                prob_keep = true;
            } else {
                if self.sample_probabilistic(trace.trace_id_high, trace.trace_id_low) {
                    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_static(PROB_RATE_KEY),
                            AttributeValue::Float(self.sampling_rate),
                        );
                    }
                } else if self.error_sampling_enabled && self.trace_contains_error(trace, false) {
                    sampler_name = telemetry::SamplerName::Error;
                    let now = std::time::SystemTime::now();
                    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 SamplerOutcome::new(prob_keep, priority, decision_maker, Some(root_span_idx), sampler_name);
        }

        // Read once here, where the samplers below consume it; every path above returns without
        // needing it.
        let now = std::time::SystemTime::now();
        let user_priority = self.get_user_priority(trace, root_span_idx);
        if let Some(priority) = user_priority {
            sampler_name = telemetry::SamplerName::Priority;
            if Self::trace_chunk_contains_probability_sampling(trace) {
                // Traces already sampled probabilistically upstream count under the probabilistic
                // sampler; the priority sampler still decides them.
                sampler_name = telemetry::SamplerName::Probabilistic;
            }
            if priority < PRIORITY_AUTO_DROP {
                // Manual drop: short-circuit and skip other samplers.
                return SamplerOutcome::new(false, priority, "", Some(root_span_idx), sampler_name);
            }

            if rare {
                sampler_name = telemetry::SamplerName::Rare;
                return SamplerOutcome::new(true, priority, "", Some(root_span_idx), sampler_name);
            }

            if self.priority_sampler.sample(now, trace, root_span_idx, priority, 0.0) {
                return SamplerOutcome::new(true, priority, "", Some(root_span_idx), sampler_name);
            }
        } 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 {
                sampler_name = telemetry::SamplerName::Rare;
                return SamplerOutcome::new(true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx), sampler_name);
            }

            // The OTLP rate decision is probabilistic: its keeps and its drops both record
            // under the probabilistic sampler, never the no-priority bucket.
            sampler_name = telemetry::SamplerName::Probabilistic;
            // 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 SamplerOutcome::new(
                    true,
                    PRIORITY_AUTO_KEEP,
                    DECISION_MAKER_PROBABILISTIC,
                    Some(root_span_idx),
                    sampler_name,
                );
            }
        } else {
            sampler_name = telemetry::SamplerName::NoPriority;
            if rare {
                sampler_name = telemetry::SamplerName::Rare;
                return SamplerOutcome::new(true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx), sampler_name);
            }
            if self.no_priority_sampler.sample(now, trace, root_span_idx) {
                return SamplerOutcome::new(true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx), sampler_name);
            }
        }

        if self.error_sampling_enabled && self.trace_contains_error(trace, false) {
            sampler_name = telemetry::SamplerName::Error;
            let keep = self.error_sampler.sample_error(now, trace, root_span_idx);
            if keep {
                return SamplerOutcome::new(true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx), sampler_name);
            }
        }

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

    /// Returns whether the trace was already sampled probabilistically upstream.
    ///
    /// Tracer head sampling and receiver-side sampling both stamp the decision maker, so the
    /// first span carrying it identifies the trace.
    fn trace_chunk_contains_probability_sampling(trace: &Trace) -> bool {
        trace
            .spans()
            .iter()
            .find_map(|span| {
                span.attributes
                    .get(TAG_DECISION_MAKER)
                    .and_then(AttributeValue::as_string)
            })
            .is_some_and(|dm| dm == DECISION_MAKER_PROBABILISTIC)
    }

    /// Fills in trace-level metadata from the spans.
    ///
    /// Origin comes from the root span's `_dd.origin` tag, and the decision maker is the first
    /// span in chunk order carrying `_dd.p.dm` - both only when not already set. This runs here,
    /// downstream of span filtering, so the metadata anchors to the final span population rather
    /// than the one that arrived at decode time.
    fn backfill_trace_metadata(&self, trace: &mut Trace, root_span_idx: usize) {
        if trace.origin.is_empty() {
            if let Some(origin) = trace.spans()[root_span_idx]
                .attributes
                .get(TAG_ORIGIN)
                .and_then(AttributeValue::as_string)
            {
                trace.origin = origin.clone();
            }
        }

        if trace.decision_maker.is_none() {
            let first_decision_maker = trace.spans().iter().find_map(|span| {
                span.attributes
                    .get(TAG_DECISION_MAKER)
                    .and_then(AttributeValue::as_string)
            });
            if let Some(dm) = first_decision_maker {
                trace.decision_maker = Some(dm.clone());
            }
        }

        if !self.compute_top_level_by_span_kind {
            compute_top_level(trace.spans_mut());
        }
    }

    /// 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: &'static str, root_span_idx: usize,
    ) {
        let is_otlp = self.is_otlp_trace(trace, root_span_idx);
        // Chunk-level decision maker: the sampler's value when it decided, otherwise the first
        // carrier promoted during backfill (falling back to the root's own tag).
        let promoted_decision_maker = if decision_maker.is_empty() {
            trace.decision_maker.clone()
        } else {
            None
        };
        let root_span_value = match trace.spans_mut().get_mut(root_span_idx) {
            Some(span) => span,
            None => return,
        };

        let decision_maker_meta = if decision_maker.is_empty() {
            promoted_decision_maker.or_else(|| {
                root_span_value
                    .attributes
                    .get(TAG_DECISION_MAKER)
                    .and_then(AttributeValue::as_string)
                    .cloned()
            })
        } else {
            Some(MetaString::from_static(decision_maker))
        };

        // The span-level decision maker is stamped only when a sampler made the call: promotion
        // targets trace chunk metadata, and each span's own tag is left in place because
        // downstream systems rely on the span keeping its original value.
        //
        // With the APM-level probabilistic sampler on OTLP traces, `_dd.p.dm` is written to
        // trace chunk tags only (not span meta); on the legacy sampling path it is written to
        // both. The span meta write is skipped only when both conditions hold; the DM value
        // still flows through trace fields to the encoder.
        if priority > 0 && !decision_maker.is_empty() && !(is_otlp && self.probabilistic_sampler_enabled) {
            root_span_value.attributes.insert(
                MetaString::from_static(TAG_DECISION_MAKER),
                AttributeValue::String(MetaString::from_static(decision_maker)),
            );
        }

        // 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 SamplerOutcome {
            keep,
            priority,
            decision_maker,
            root_span_idx,
            ..
        } = self.run_samplers(trace);

        // Backfill trace-derived metadata before any forwarding path (keep, ETS, or
        // single-span/analytics sampling), so every forwarded trace carries metadata derived
        // from its final span population.
        if let Some(root_idx) = root_span_idx {
            self.backfill_trace_metadata(trace, root_idx);
        }

        // 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
    }
}

#[async_trait]
impl Transform for TraceSampler {
    async fn run(mut self: Box<Self>, mut context: TransformContext) -> Result<(), GenericError> {
        let mut health = context.take_health_handle();
        health.mark_ready();
        debug!("Trace sampler transform started.");

        let mut window = tokio::time::interval_at(tokio::time::Instant::now() + WINDOW, WINDOW);

        loop {
            select! {
                _ = health.live() => continue,
                _ = window.tick() => {
                    // Report the sampling telemetry through the dedicated metrics output, so the
                    // metrics flow into the metrics pipeline and on to the backend. The signature
                    // counts are read here, at report time, so a quiet window cannot republish
                    // stale gauges.
                    let events = self.telemetry.take_window_events(
                        self.priority_sampler.tracked_signature_count(),
                        self.no_priority_sampler.tracked_signature_count(),
                        self.error_sampler.tracked_signature_count(),
                    );
                    context.dispatcher().buffered_named("metrics")?.send_all(events).await?;
                }
                maybe_events = context.events().next() => match maybe_events {
                    Some(mut events) => {
                        events.remove_if(|event| match event {
                            Event::Trace(trace) => !self.process_trace(trace),
                            _ => false,
                        });

                        context.dispatcher().buffered()?.send_all(events).await?;
                    }
                    None => {
                        // The input stream has ended, so report the final window before stopping.
                        let events = self.telemetry.take_window_events(
                            self.priority_sampler.tracked_signature_count(),
                            self.no_priority_sampler.tracked_signature_count(),
                            self.error_sampler.tracked_signature_count(),
                        );
                        context.dispatcher().buffered_named("metrics")?.send_all(events).await?;
                        break;
                    }
                }
            }
        }

        debug!("Trace sampler transform stopped.");
        Ok(())
    }
}

#[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,
            probabilistic: probabilistic::ProbabilisticSampler::new(0, false),
            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, 5000),
            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,
                telemetry::SamplerCounters::new(),
            ),
            telemetry: telemetry::DecisionWindow::new(),
            compute_top_level_by_span_kind: false,
        }
    }

    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)
    }

    fn create_span_with_parent_and_service(
        span_id: u64, parent_id: u64, service: &str, meta: &[(&str, &str)],
    ) -> DdSpan {
        let mut attrs = saluki_common::collections::FastHashMap::default();
        for (k, v) in meta {
            attrs.insert(MetaString::from(*k), AttributeValue::String(MetaString::from(*v)));
        }
        DdSpan::new(
            MetaString::from(service),
            MetaString::from("operation"),
            MetaString::from("resource"),
            MetaString::from("type"),
            span_id,
            parent_id,
            0,
            1000,
            0,
        )
        .with_attributes(attrs)
    }

    #[test]
    fn apply_sampling_metadata_preserves_backfilled_decision_maker() {
        let sampler = create_test_sampler();
        let root = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![root]);
        trace.decision_maker = Some(MetaString::from("-8"));

        sampler.apply_sampling_metadata(&mut trace, true, PRIORITY_AUTO_KEEP, "", 0);

        assert_eq!(
            trace.decision_maker.as_deref(),
            Some("-8"),
            "the chunk-level decision maker keeps the promoted value"
        );
        assert!(
            !trace.spans()[0].attributes.contains_key(TAG_DECISION_MAKER),
            "a promoted decision maker is never stamped onto the root span; span-level _dd.p.dm is written only when a sampler decides"
        );
    }

    #[test]
    fn apply_sampling_metadata_keeps_root_decision_maker_without_sampler_decision() {
        let sampler = create_test_sampler();
        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(
            MetaString::from(TAG_DECISION_MAKER),
            AttributeValue::String(MetaString::from("-9")),
        );
        let root = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![root]);
        // Promoted from an earlier child span during backfill.
        trace.decision_maker = Some(MetaString::from("-8"));

        sampler.apply_sampling_metadata(&mut trace, true, PRIORITY_AUTO_KEEP, "", 0);

        assert_eq!(
            trace.decision_maker.as_deref(),
            Some("-8"),
            "the chunk-level decision maker keeps the promoted value"
        );
        assert_eq!(
            trace.spans()[0]
                .attributes
                .get(TAG_DECISION_MAKER)
                .and_then(AttributeValue::as_string)
                .map(|s| s.as_ref()),
            Some("-9"),
            "the root span's own decision maker is never overwritten by a promoted value"
        );
    }

    #[test]
    fn backfill_promotes_origin_from_root_span() {
        let sampler = create_test_sampler();
        // The root is reported last here.
        let child = create_span_with_parent_and_service(2, 1, "svc", &[]);
        let root = create_span_with_parent_and_service(1, 0, "svc", &[("_dd.origin", "lambda")]);
        let mut trace = create_test_trace(vec![child, root]);

        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        sampler.backfill_trace_metadata(&mut trace, root_idx);

        assert_eq!(trace.origin.as_ref(), "lambda");
    }

    #[test]
    fn backfill_ignores_origin_from_non_root_span() {
        let sampler = create_test_sampler();
        let root = create_span_with_parent_and_service(1, 0, "svc", &[]);
        let child = create_span_with_parent_and_service(2, 1, "svc", &[("_dd.origin", "rum")]);
        let mut trace = create_test_trace(vec![root, child]);

        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        sampler.backfill_trace_metadata(&mut trace, root_idx);

        assert!(trace.origin.as_ref().is_empty());
    }

    #[test]
    fn backfill_promotes_first_decision_maker() {
        let sampler = create_test_sampler();
        let root = create_span_with_parent_and_service(1, 0, "svc", &[]);
        let early = create_span_with_parent_and_service(2, 1, "svc", &[("_dd.p.dm", "-8")]);
        let late = create_span_with_parent_and_service(3, 1, "svc", &[("_dd.p.dm", "-9")]);
        let mut trace = create_test_trace(vec![root, early, late]);

        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        sampler.backfill_trace_metadata(&mut trace, root_idx);

        assert_eq!(
            trace.decision_maker.as_deref(),
            Some("-8"),
            "first span carrying the tag wins - not a root-based rule"
        );
    }

    #[test]
    fn backfill_marks_top_level_fallback_when_span_kind_computation_disabled() {
        let sampler = create_test_sampler();
        // Roots, orphans, and service boundaries get marked; interior same-service spans don't.
        let root = create_span_with_parent_and_service(1, 0, "svc-a", &[]);
        let child = create_span_with_parent_and_service(2, 1, "svc-a", &[]);
        let orphan = create_span_with_parent_and_service(3, 0xEE, "svc-a", &[]);
        let other_service = create_span_with_parent_and_service(5, 1, "svc-b", &[]);
        let local_entry = create_span_with_parent_and_service(6, 5, "svc-a", &[]);
        let mut trace = create_test_trace(vec![root, child, orphan, other_service, local_entry]);

        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        sampler.backfill_trace_metadata(&mut trace, root_idx);

        let spans = trace.spans();
        let top_level = |sid: u64| {
            spans
                .iter()
                .find(|s| s.span_id() == sid)
                .unwrap()
                .attributes
                .get(crate::common::datadog::TOP_LEVEL_KEY)
                .and_then(AttributeValue::as_num)
        };
        assert_eq!(top_level(1), Some(1.0), "root spans are marked");
        assert_eq!(
            top_level(3),
            Some(1.0),
            "orphans whose parent is missing from the chunk are marked"
        );
        assert_eq!(top_level(5), Some(1.0), "spans entering a different service are marked");
        assert_eq!(
            top_level(6),
            Some(1.0),
            "spans crossing back into the resource service are marked (local root)"
        );
        assert_eq!(
            top_level(2),
            None,
            "same-service children with their parent present are not marked"
        );
    }

    #[test]
    fn backfill_skips_top_level_when_span_kind_computation_enabled() {
        let sampler = TraceSampler {
            compute_top_level_by_span_kind: true,
            ..create_test_sampler()
        };
        let root = create_span_with_parent_and_service(1, 0, "svc", &[]);
        let child = create_span_with_parent_and_service(2, 1, "svc", &[]);
        let orphan = create_span_with_parent_and_service(3, 0xEE, "svc", &[]);
        let mut trace = create_test_trace(vec![root, child, orphan]);

        let root_idx = sampler.get_root_span_index(&trace).unwrap();
        sampler.backfill_trace_metadata(&mut trace, root_idx);

        for span in trace.spans() {
            assert!(
                !span.attributes.contains_key(crate::common::datadog::TOP_LEVEL_KEY),
                "no span should carry a top-level mark with the fallback disabled"
            );
        }
    }

    #[test]
    fn process_trace_backfills_metadata_for_forwarded_traces() {
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 1.0; // keeps the trace, so it is forwarded
        sampler.probabilistic_sampler_enabled = true;

        let child = create_span_with_parent_and_service(2, 1, "svc", &[]);
        let root = create_span_with_parent_and_service(1, 0, "svc", &[("_dd.origin", "lambda")]);
        let mut trace = create_test_trace(vec![child, root]);

        let forwarded = sampler.process_trace(&mut trace);
        assert!(forwarded);
        assert_eq!(
            trace.origin.as_ref(),
            "lambda",
            "origin is backfilled for forwarded traces"
        );
        assert!(
            trace
                .spans()
                .iter()
                .any(|s| s.attributes.contains_key(crate::common::datadog::TOP_LEVEL_KEY)),
            "top-level marks are set for forwarded traces"
        );
    }

    #[test]
    fn apply_sampling_metadata_sampler_decision_overrides_backfill() {
        let sampler = create_test_sampler();
        let root = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![root]);
        trace.decision_maker = Some(MetaString::from("-8"));

        sampler.apply_sampling_metadata(&mut trace, true, PRIORITY_AUTO_KEEP, DECISION_MAKER_PROBABILISTIC, 0);

        assert_eq!(trace.decision_maker.as_deref(), Some(DECISION_MAKER_PROBABILISTIC));
        assert_eq!(
            trace.spans()[0]
                .attributes
                .get(TAG_DECISION_MAKER)
                .and_then(AttributeValue::as_string)
                .map(|s| s.as_ref()),
            Some(DECISION_MAKER_PROBABILISTIC)
        );
    }

    #[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 SamplerOutcome {
            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 SamplerOutcome { 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 SamplerOutcome {
            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`.
        // 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(0, 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(0, trace_id),
            sampler.sample_probabilistic(0, trace_id)
        );
    }

    #[test]
    fn probabilistic_sampling_hash_seed_changes_decisions() {
        // A seed re-buckets every trace ID: 0x1234567890ABCDEF moves from bucket 1764 with seed 0 to
        // bucket 7338 with seed 22, so rates that straddle only one of the buckets flip decisions.
        let mut seeded = create_test_sampler();
        seeded.probabilistic = probabilistic::ProbabilisticSampler::new(22, false);

        // Bucket 7338 (7338/16384 = 0.4479): dropped below 0.45, kept above it.
        seeded.sampling_rate = 0.40;
        assert!(!seeded.sample_probabilistic(0, 0x1234567890ABCDEF));
        seeded.sampling_rate = 0.50;
        assert!(seeded.sample_probabilistic(0, 0x1234567890ABCDEF));

        // At rate 0.20 the zero-seed sampler keeps this trace (bucket 1764, pinned in
        // `probabilistic_sampling_known_decisions`) while the seeded sampler drops it.
        seeded.sampling_rate = 0.20;
        assert!(!seeded.sample_probabilistic(0, 0x1234567890ABCDEF));
        let mut unseeded = create_test_sampler();
        unseeded.sampling_rate = 0.20;
        assert!(unseeded.sample_probabilistic(0, 0x1234567890ABCDEF));
    }

    #[test]
    fn probabilistic_sampling_full_trace_id_mode_hashes_both_halves() {
        let high = 0xAABBCCDD00112233_u64;
        let low = 0x1234567890ABCDEF_u64;

        // Legacy mode buckets on the low half alone, so the high half cannot move the decision.
        let mut legacy = create_test_sampler();
        legacy.probabilistic = probabilistic::ProbabilisticSampler::new(22, false);
        legacy.sampling_rate = 0.45;
        assert_eq!(
            legacy.sample_probabilistic(0, low),
            legacy.sample_probabilistic(high, low)
        );

        // Full mode hashes the 16-byte big-endian ID, high half first: (high, low) lands in bucket
        // 454 (kept at 0.05, dropped at 0.02), moving the decision relative to legacy mode.
        let mut full = create_test_sampler();
        full.probabilistic = probabilistic::ProbabilisticSampler::new(22, true);
        full.sampling_rate = 0.02;
        assert!(!full.sample_probabilistic(high, low));
        full.sampling_rate = 0.05;
        assert!(full.sample_probabilistic(high, low));

        // A zero high half is still a different input than legacy mode: the halves swap byte
        // ranges, bucket 14666 versus 7338 (kept at 0.90, dropped at 0.89).
        full.sampling_rate = 0.89;
        assert!(!full.sample_probabilistic(0, low));
        full.sampling_rate = 0.90;
        assert!(full.sample_probabilistic(0, low));
    }

    #[test]
    fn from_configuration_reads_the_full_trace_id_feature() {
        let traces = domains::traces::Domain {
            features: vec![domains::traces::ApmFeature::ProbabilisticSamplerFullTraceId],
            ..Default::default()
        };
        let config = TraceSamplerConfiguration::from_configuration(&traces, &domains::otlp::Traces::default());
        assert!(config.probabilistic_full_trace_id);

        // Unrecognized flags are carried, not rejected, and enable nothing.
        let traces = domains::traces::Domain {
            features: vec![domains::traces::ApmFeature::Other("table_names".to_owned())],
            ..Default::default()
        };
        let config = TraceSamplerConfiguration::from_configuration(&traces, &domains::otlp::Traces::default());
        assert!(!config.probabilistic_full_trace_id);
    }

    #[test]
    fn every_trace_is_counted_under_exactly_one_sampler() {
        // With the probabilistic sampler enabled at rate 1.0, every trace is decided (and kept) by
        // it: the sum of seen across all decision keys must equal the number of traces, with no
        // double counting across samplers.
        let mut sampler = create_test_sampler();

        for i in 0..10 {
            let span = create_test_span(i, 0);
            let mut trace = create_test_trace(vec![span]);
            sampler.run_samplers(&mut trace);
        }

        let decisions = sampler.telemetry.snapshot_decisions();
        let total_seen: u64 = decisions.values().map(|counts| counts.seen).sum();
        let total_kept: u64 = decisions.values().map(|counts| counts.kept).sum();
        assert_eq!(total_seen, 10);
        assert_eq!(total_kept, 10);
        assert!(decisions
            .keys()
            .all(|key| key.sampler == telemetry::SamplerName::Probabilistic));
        assert!(decisions.keys().all(|key| key.service == "test-service"));
    }

    #[test]
    fn priority_decisions_are_counted_under_the_priority_sampler_with_priority_tag() {
        // A trace with a user-set sampling priority is decided by the priority sampler, and its
        // decision key carries that priority; no other sampler claims it.
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false;

        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]);
        sampler.run_samplers(&mut trace);

        let decisions = sampler.telemetry.snapshot_decisions();
        assert_eq!(decisions.len(), 1);
        let (key, counts) = decisions.iter().next().unwrap();
        assert_eq!(key.sampler, telemetry::SamplerName::Priority);
        assert_eq!(key.priority, Some(telemetry::Priority::ManualKeep));
        assert_eq!(counts.seen, 1);
        assert_eq!(counts.kept, 1);
    }

    #[test]
    fn upstream_probability_sampling_reattributes_to_the_probabilistic_sampler() {
        // A trace already sampled probabilistically upstream carries the probabilistic decision
        // maker; its decision counts under the probabilistic sampler, without the priority
        // sampler's sampling_priority or target_env tags.
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false;

        let mut metrics = HashMap::new();
        metrics.insert(SAMPLING_PRIORITY_METRIC_KEY.to_string(), 2.0);
        let mut span = create_test_span_with_metrics(1, metrics);
        span.attributes.insert(
            MetaString::from_static(TAG_DECISION_MAKER),
            AttributeValue::String(MetaString::from_static(DECISION_MAKER_PROBABILISTIC)),
        );
        let mut trace = create_test_trace(vec![span]);
        sampler.run_samplers(&mut trace);

        let decisions = sampler.telemetry.snapshot_decisions();
        assert_eq!(decisions.len(), 1);
        let (key, _) = decisions.iter().next().unwrap();
        assert_eq!(key.sampler, telemetry::SamplerName::Probabilistic);
        assert_eq!(key.priority, None);
    }

    #[test]
    fn empty_traces_record_no_decision() {
        // Unclaimed decisions carry no sampler identity, so they never inflate the unknown series;
        // a real attribution gap remains the only thing that could produce one.
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false;

        let mut trace = create_test_trace(vec![]);
        sampler.run_samplers(&mut trace);
        assert!(sampler.telemetry.snapshot_decisions().is_empty());
    }

    #[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 SamplerOutcome {
            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 SamplerOutcome {
            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 SamplerOutcome { 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 SamplerOutcome {
            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,
                telemetry::SamplerCounters::new(),
            ),
            ..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 SamplerOutcome {
            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 SamplerOutcome {
            keep,
            root_span_idx: 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 SamplerOutcome { keep: 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 SamplerOutcome {
            keep: keep2,
            priority: 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 SamplerOutcome { 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 SamplerOutcome {
            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 SamplerOutcome { 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 SamplerOutcome {
            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 SamplerOutcome { 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 otlp_rate_decisions_record_under_the_probabilistic_sampler() {
        // Keeps and drops from the OTLP rate decision both belong to the probabilistic sampler;
        // the drop must not fall through to the no-priority bucket.
        let mut sampler = create_test_sampler();
        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 SamplerOutcome { keep, priority, .. } = sampler.run_samplers(&mut trace);
        assert!(!keep);
        assert_eq!(priority, PRIORITY_AUTO_DROP);

        let decisions = sampler.telemetry.snapshot_decisions();
        let (key, counts) = decisions.iter().next().unwrap();
        assert_eq!(key.sampler, telemetry::SamplerName::Probabilistic);
        assert_eq!(counts.seen, 1);
        assert_eq!(counts.kept, 0);
    }

    #[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 SamplerOutcome {
            keep,
            priority,
            decision_maker,
            root_span_idx: 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 SamplerOutcome {
            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 SamplerOutcome {
            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 SamplerOutcome { 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 SamplerOutcome {
            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 SamplerOutcome { 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 SamplerOutcome { 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 SamplerOutcome {
            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 SamplerOutcome {
            keep,
            priority,
            decision_maker: 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 SamplerOutcome {
            keep,
            priority,
            decision_maker: 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 SamplerOutcome {
            keep,
            priority,
            decision_maker: 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 SamplerOutcome {
            keep,
            priority,
            decision_maker: 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 SamplerOutcome {
            keep,
            priority,
            decision_maker: 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");
    }
}