saluki_components/transforms/trace_sampler/

mod.rs

//! Decides which incoming traces to forward before they reach the backend.
//!
//! The transform honors user sampling decisions and can keep traces through probabilistic,
//! priority, error, and rare sampling. It also handles OTLP sampling metadata. When given an
//! `APM_SAMPLING` subscription, it updates priority and error sampler targets and the rare sampler
//! switch before processing each buffer. Without a subscription, it uses local configuration.
//!
//! TODO: Add trace metrics from `datadog-agent/pkg/trace/sampler/metrics.go`.

use std::collections::HashMap;
use std::sync::{Arc, LazyLock};

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

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

use self::event_processor::EventProcessor;
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, SAMPLING_PRIORITY_METRIC_KEY, TAG_DECISION_MAKER, TAG_ORIGIN,
};
use crate::config::AgentEnv;

// 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";

/// 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,
    /// Environment the priority sampler uses for traces without one, and used to select a remote `by_env` entry.
    agent_env: AgentEnv,
    rare_sampler_enabled: bool,
    rare_sampler_tps: f64,
    rare_sampler_cooldown_secs: f64,
    rare_sampler_cardinality: usize,
    compute_top_level_by_span_kind: bool,
    analyzed_spans_by_service: HashMap<String, HashMap<String, f64>>,
    analyzed_rate_by_service: HashMap<String, f64>,
    max_events_per_second: f64,

    /// Remote updates, or an inert subscription when remote sampling is not configured.
    remote_sampling: TraceSamplingSubscription,
}

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. The priority sampler samples traces without an environment
    /// under `agent_env`, and remote sampling selects its `by_env` entry with it.
    pub fn from_configuration(
        traces: &domains::traces::Domain, otlp_traces: &domains::otlp::Traces, agent_env: AgentEnv,
    ) -> Self {
        if !traces.analyzed_rate_by_service.is_empty() {
            warn!("analyzed_rate_by_service is deprecated, please use analyzed_spans instead");
        }
        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,
            agent_env,
            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,
            compute_top_level_by_span_kind: otlp_traces.enable_compute_top_level_by_span_kind,
            analyzed_spans_by_service: traces.analyzed_spans_by_service.clone(),
            analyzed_rate_by_service: traces.analyzed_rate_by_service.clone(),
            max_events_per_second: traces.max_events_per_second,
            remote_sampling: TraceSamplingSubscription::default(),
        }
    }

    /// Connects the trace sampler to `APM_SAMPLING` remote updates.
    ///
    /// Only priority and error sampling targets and rare sampling can change. For each field
    /// left unset by the remote configuration, the sampler uses its local setting.
    pub fn with_remote_sampling(mut self, remote_sampling: TraceSamplingSubscription) -> Self {
        self.remote_sampling = remote_sampling;
        self
    }

    fn build_sampler(&self) -> TraceSampler {
        let telemetry = DecisionWindow::new();
        let event_processor = EventProcessor::new(
            &self.analyzed_spans_by_service,
            &self.analyzed_rate_by_service,
            self.max_events_per_second,
        );
        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,
            ),
            error_sampler: errors::ErrorsSampler::new(self.errors_per_second, self.extra_sample_rate),
            priority_sampler: priority_sampler::PrioritySampler::new(
                self.agent_env.clone().into(),
                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(),
            ),
            event_processor,
            telemetry,
            compute_top_level_by_span_kind: self.compute_top_level_by_span_kind,
            static_settings: SamplerSettings {
                target_traces_per_second: self.target_traces_per_second,
                errors_per_second: self.errors_per_second,
                rare_sampler_enabled: self.rare_sampler_enabled,
            },
            agent_env: self.agent_env.clone().into(),
            remote_sampling: self.remote_sampling.clone(),
            applied_sampling: None,
        }
    }
}

#[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> {
        Ok(Box::new(self.build_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);

        // Event extraction rates are bounded by their configured size, plus one nested operation
        // map per configured service.
        let analyzed_operations: usize = self.analyzed_spans_by_service.values().map(|ops| ops.len()).sum();
        builder
            .minimum()
            .with_map::<String, HashMap<String, f64>>(
                "analyzed spans service rates",
                self.analyzed_spans_by_service.len(),
            )
            .with_map::<String, f64>("analyzed spans operation rates", analyzed_operations)
            .with_map::<String, f64>("legacy analyzed service rates", self.analyzed_rate_by_service.len());
    }
}

pub struct TraceSampler {
    sampling_rate: f64,
    error_tracking_standalone: bool,
    error_sampling_enabled: bool,
    probabilistic_sampler_enabled: bool,
    probabilistic: probabilistic::ProbabilisticSampler,
    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,
    event_processor: EventProcessor,
    telemetry: DecisionWindow,
    /// Local settings used when a remote configuration leaves a field unset.
    static_settings: SamplerSettings,
    /// Agent environment, used to select a remote `by_env` entry.
    agent_env: MetaString,
    remote_sampling: TraceSamplingSubscription,
    /// Last update handled, so each buffer does not reapply the same settings.
    applied_sampling: Option<Arc<RemoteSampling>>,
}

/// 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,
    /// Whether the samplers decided without short-circuiting, so event extraction should run.
    extract_events: bool,
}

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

impl TraceSampler {
    /// Applies the latest accepted remote sampling configuration, if it has changed.
    ///
    /// Unset fields use local settings, not previous remote values. When the assignment becomes
    /// empty, the subscription keeps reporting the last configuration, as the Datadog Agent does. A newly published
    /// configuration can have the same values after a worker restart; applying it again preserves
    /// learned sampling rates and the rare sampler's record of seen span signatures.
    fn apply_remote_sampling(&mut self) {
        let Some(current) = self.remote_sampling.current() else {
            return;
        };
        if self
            .applied_sampling
            .as_ref()
            .is_some_and(|applied| Arc::ptr_eq(applied, &current))
        {
            return;
        }
        if let RemoteSampling::Remote { id, config } = &*current {
            let settings = config.resolve(&self.agent_env, &self.static_settings);
            debug!(config_id = %id, ?settings, "Applying remote sampling configuration.");
            // The Datadog Agent does not update the no-priority sampler from this product.
            self.priority_sampler
                .update_target_tps(settings.target_traces_per_second);
            self.error_sampler.update_target_tps(settings.errors_per_second);
            self.rare_sampler.set_enabled(settings.rare_sampler_enabled);
        }
        self.applied_sampling = Some(current);
    }

    /// Applies pending sampling settings, then samples every trace in the buffer with those settings.
    fn process_buffer(&mut self, events: &mut EventsBuffer) {
        self.apply_remote_sampling();
        events.remove_if(|event| match event {
            Event::Trace(trace) => !self.process_trace(trace),
            _ => false,
        });
    }

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

    /// 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)> {
        let TraceIngestInfo::Otlp { receiver_sampling_rate } = trace.ingest_info() else {
            return None;
        };
        if self.probabilistic_sampler_enabled {
            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, receiver_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))
    }

    /// Retains only the analyzed-marker spans of a dropped trace and forwards it as events.
    fn forward_analyzed_spans(&self, trace: &mut Trace, priority: i32) -> bool {
        let retained = trace.retain_spans(|_, span| span.attributes.contains_key(KEY_ANALYZED_SPANS));
        if retained > 0 {
            trace.dropped_trace = true;
            trace.priority = Some(priority);
            true
        } else {
            false
        }
    }

    /// 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);
            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, false);
        }

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

        // 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, false);
            }
            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, false);
        }

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

        // Modern path: ProbabilisticSamplerEnabled = true
        if self.probabilistic_sampler_enabled {
            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,
                true,
            );
        }

        // 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, false);
            }

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

            if self.priority_sampler.sample(now, trace, root_span_idx, priority, 0.0) {
                return SamplerOutcome::new(true, priority, "", Some(root_span_idx), sampler_name, true);
            }
        } else if let TraceIngestInfo::Otlp { receiver_sampling_rate } = trace.ingest_info() {
            // 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, true);
            }

            // 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, receiver_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,
                    true,
                );
            }
        } 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, true);
            }
            if self.no_priority_sampler.sample(now, trace, root_span_idx) {
                return SamplerOutcome::new(true, PRIORITY_AUTO_KEEP, "", Some(root_span_idx), sampler_name, true);
            }
        }

        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, true);
            }
        }

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

    /// 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 = matches!(trace.ingest_info(), TraceIngestInfo::Otlp { .. });
        // 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 };
    }

    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,
            extract_events,
            ..
        } = 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);
        }

        // Extract events from every trace the samplers did not short-circuit; surviving spans
        // carry the analyzed marker and their rate metadata whether the trace is kept or dropped,
        // and are forwarded as standalone events only when the trace is dropped.
        let events = if extract_events {
            self.event_processor.process(trace, root_span_idx, priority)
        } else {
            event_processor::EventCounts::default()
        };

        // 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 extracted events if no SSS spans
            if self.forward_analyzed_spans(trace, priority) {
                return true;
            }
        } else if events.sampled > 0 {
            // 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 mut 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(),
                    );
                    events.extend(self.event_processor.take_window_events());
                    context.dispatcher().buffered_named("metrics")?.send_all(events).await?;
                }
                maybe_events = context.events().next() => match maybe_events {
                    Some(mut events) => {
                        self.process_buffer(&mut events);
                        context.dispatcher().buffered()?.send_all(events).await?;
                    }
                    None => {
                        // The input stream has ended, so report the final window before stopping.
                        let mut 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(),
                        );
                        events.extend(self.event_processor.take_window_events());
                        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 datadog_agent_apm::sampling::test_util::TestSamplingPublisher;
    use saluki_core::components::test_util::TestComponentDriver;
    use saluki_core::data_model::event::trace::{AttributeValue, Span as DdSpan, Trace, TraceIngestInfo};

    use super::*;
    use crate::common::datadog::OTEL_TRACE_ID_META_KEY;

    const PRIORITY_USER_DROP: i32 = -1;

    // Deterministic receiver sampling decisions at a rate of 0.5, from the Go trace agent's `TestCreateChunks`:
    // https://github.com/DataDog/datadog-agent/blob/c6f7e0bc2d5744601a10f231751948ac7a4ee781/pkg/trace/api/otlp_test.go#L740-L810
    const OTLP_TEST_RECEIVER_SAMPLING_RATE: f64 = 0.5;
    const OTLP_RECEIVER_KEPT_TRACE_ID: u64 = 123;
    const OTLP_RECEIVER_DROPPED_TRACE_ID: u64 = 1237892138897;

    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),
            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(),
            ),
            event_processor: EventProcessor::for_tests(&HashMap::new(), &HashMap::new(), f64::INFINITY),
            telemetry: telemetry::DecisionWindow::new(),
            compute_top_level_by_span_kind: false,
            static_settings: SamplerSettings {
                target_traces_per_second: 10.0,
                errors_per_second: 10.0,
                rare_sampler_enabled: false,
            },
            agent_env: MetaString::from("agent-env"),
            remote_sampling: TraceSamplingSubscription::default(),
            applied_sampling: None,
        }
    }

    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, TraceIngestInfo::Datadog)
    }

    fn create_otlp_test_trace(spans: Vec<DdSpan>, receiver_sampling_rate: f64, trace_id_low: u64) -> Trace {
        let mut trace = Trace::new(spans, TraceIngestInfo::Otlp { receiver_sampling_rate });
        trace.trace_id_low = trace_id_low;
        trace
    }

    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(), prod_env());
        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(), prod_env());
        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 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;

        let mut trace = create_otlp_test_trace(
            vec![create_top_level_span(1)],
            OTLP_TEST_RECEIVER_SAMPLING_RATE,
            OTLP_RECEIVER_DROPPED_TRACE_ID,
        );

        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;

        let mut trace = create_otlp_test_trace(
            vec![create_top_level_span(1)],
            OTLP_TEST_RECEIVER_SAMPLING_RATE,
            OTLP_RECEIVER_DROPPED_TRACE_ID,
        );

        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_receiver_kept_otlp_test_trace(spans: Vec<DdSpan>) -> Trace {
        create_otlp_test_trace(spans, OTLP_TEST_RECEIVER_SAMPLING_RATE, OTLP_RECEIVER_KEPT_TRACE_ID)
    }

    fn create_sampler_with_ets_legacy() -> TraceSampler {
        TraceSampler {
            error_tracking_standalone: true,
            probabilistic_sampler_enabled: false,
            ..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_test_span(1, 0); // no error
        let mut trace = create_receiver_kept_otlp_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_test_span(1, 1); // error=1
        let mut trace = create_receiver_kept_otlp_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 trace dropped by the receiver rate (legacy sampler path): a drop is not stamped with a decision maker.
    #[test]
    fn ets_otlp_receiver_drop_gets_drop_priority_and_no_dm() {
        let mut sampler = create_sampler_with_ets_legacy();

        let mut trace = create_otlp_test_trace(
            vec![create_test_span(1, 0)],
            OTLP_TEST_RECEIVER_SAMPLING_RATE,
            OTLP_RECEIVER_DROPPED_TRACE_ID,
        );

        assert!(sampler.process_trace(&mut trace), "ETS forwards non-error traces");
        assert!(trace.dropped_trace);
        assert_eq!(trace.priority, Some(PRIORITY_AUTO_DROP));
        assert_eq!(trace.decision_maker, None);
    }

    /// 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_test_span(1, 0); // no error
        let mut trace = create_receiver_kept_otlp_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_test_span(1, 0);
        span.attributes.insert(
            MetaString::from(SAMPLING_PRIORITY_METRIC_KEY),
            AttributeValue::Float(2.0),
        ); // UserKeep
        let mut trace = create_receiver_kept_otlp_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");
    }

    // ── Ingest information tests ────────────────────────────────────────────────
    // Regression tests: the sampler once classified a trace as OTLP by the `otel.trace_id` attribute on its root
    // span. Removing that attribute, or filtering out the original root, sent OTLP traces down the native path, and
    // a native trace carrying the attribute was sampled at the OTLP receiver rate. Classification now uses the
    // ingest information fixed when the trace is constructed.

    #[test]
    fn otlp_receiver_sampling_does_not_depend_on_span_attributes_or_root() {
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false;

        let span_sets = [
            ("rooted", vec![create_test_span(1, 0)]),
            (
                "original root removed",
                vec![
                    create_test_span(2, 0).with_parent_id(1),
                    create_test_span(3, 0).with_parent_id(2),
                ],
            ),
        ];
        let cases = [
            (
                OTLP_RECEIVER_KEPT_TRACE_ID,
                true,
                PRIORITY_AUTO_KEEP,
                DECISION_MAKER_PROBABILISTIC,
            ),
            (OTLP_RECEIVER_DROPPED_TRACE_ID, false, PRIORITY_AUTO_DROP, ""),
        ];
        for (name, spans) in &span_sets {
            for (trace_id, want_keep, want_priority, want_dm) in cases {
                let mut trace = create_otlp_test_trace(spans.clone(), OTLP_TEST_RECEIVER_SAMPLING_RATE, trace_id);

                let SamplerOutcome {
                    keep,
                    priority,
                    decision_maker,
                    ..
                } = sampler.run_samplers(&mut trace);
                assert_eq!(keep, want_keep, "{name}, trace ID {trace_id}");
                assert_eq!(priority, want_priority, "{name}, trace ID {trace_id}");
                assert_eq!(decision_maker, want_dm, "{name}, trace ID {trace_id}");
            }
        }
    }

    #[test]
    fn native_trace_with_otel_trace_id_attribute_is_not_sampled_as_otlp() {
        let spoofed_trace = || {
            let mut span = create_top_level_span(1);
            span.attributes.insert(
                MetaString::from_static(OTEL_TRACE_ID_META_KEY),
                AttributeValue::String(MetaString::from("0000000000000000000000000000007b")),
            );
            let mut trace = create_test_trace(vec![span]);
            trace.trace_id_low = OTLP_RECEIVER_KEPT_TRACE_ID;
            trace
        };

        // Legacy path: the no-priority sampler decides, not the OTLP receiver rate.
        let mut sampler = create_test_sampler();
        sampler.probabilistic_sampler_enabled = false;
        let SamplerOutcome { decision_maker, .. } = sampler.run_samplers(&mut spoofed_trace());
        assert_eq!(decision_maker, "");
        let decisions = sampler.telemetry.snapshot_decisions();
        let (key, _) = decisions.iter().next().unwrap();
        assert_eq!(key.sampler, telemetry::SamplerName::NoPriority);

        // Error tracking standalone: no OTLP pre-sampling priority or decision maker.
        let mut sampler = create_sampler_with_ets_legacy();
        let SamplerOutcome {
            keep,
            priority,
            decision_maker,
            ..
        } = sampler.run_samplers(&mut spoofed_trace());
        assert!(!keep);
        assert_eq!(priority, PRIORITY_AUTO_DROP);
        assert_eq!(decision_maker, "");
    }

    #[test]
    fn probabilistic_decision_maker_placement_follows_ingest_info() {
        // The APM probabilistic sampler is enabled at a rate of 1.0, so both traces are kept by it. The OTLP trace
        // carries no `otel.trace_id` attribute but still gets its decision maker on the chunk only.
        let mut sampler = create_test_sampler();

        let mut otlp = create_receiver_kept_otlp_test_trace(vec![create_test_span(1, 0)]);
        assert!(sampler.process_trace(&mut otlp));
        assert_eq!(otlp.decision_maker.as_deref(), Some(DECISION_MAKER_PROBABILISTIC));
        assert!(!otlp.spans()[0].attributes.contains_key(TAG_DECISION_MAKER));

        let mut native = create_test_trace(vec![create_test_span(1, 0)]);
        assert!(sampler.process_trace(&mut native));
        assert_eq!(native.decision_maker.as_deref(), Some(DECISION_MAKER_PROBABILISTIC));
    }

    #[test]
    fn analytics_events_forward_dropped_traces_with_only_event_spans() {
        let rates = HashMap::from([("test-service".to_string(), 1.0)]);
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 0.0; // probabilistic drops everything
        sampler.probabilistic_sampler_enabled = true;
        sampler.event_processor = EventProcessor::for_tests(&HashMap::new(), &rates, f64::INFINITY);

        let root = create_test_span(1, 0);
        let child = create_span_with_parent_and_service(2, 1, "other-service", &[]);
        let mut trace = create_test_trace(vec![root, child]);

        assert!(sampler.process_trace(&mut trace));
        // Only the extracted event spans are forwarded; the trace stays marked dropped so the
        // backend indexes them as standalone events, carrying the sampling priority.
        assert_eq!(trace.spans().len(), 1);
        assert!(trace.dropped_trace);
        assert_eq!(trace.priority, Some(PRIORITY_AUTO_DROP));
        let event_span = &trace.spans()[0];
        assert!(event_span.attributes.contains_key(KEY_ANALYZED_SPANS));
        // Extraction rates of 1.0 are not written; the backend assumes 1.0.
        assert!(!event_span.attributes.contains_key("_dd1.sr.eausr"));
    }

    #[test]
    fn analytics_events_dropped_trace_without_survivors_is_not_forwarded() {
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 0.0; // probabilistic drops everything
        sampler.probabilistic_sampler_enabled = true;

        let mut trace = create_test_trace(vec![create_test_span(1, 0)]);
        // No extractor nominates the span, so the trace is dropped without events.
        assert!(!sampler.process_trace(&mut trace));
    }

    #[test]
    fn analytics_events_marks_ride_along_on_kept_traces() {
        let rates = HashMap::from([("test-service".to_string(), 1.0)]);
        let mut sampler = create_test_sampler();
        sampler.event_processor = EventProcessor::for_tests(&HashMap::new(), &rates, f64::INFINITY);

        let root = create_test_span(1, 0);
        let child = create_span_with_parent_and_service(2, 1, "other-service", &[]);
        let mut trace = create_test_trace(vec![root, child]);

        assert!(sampler.process_trace(&mut trace));
        // The full span payload is forwarded, with the analyzed mark riding along on the root.
        assert_eq!(trace.spans().len(), 2);
        assert!(!trace.dropped_trace);
        assert!(trace.spans()[0].attributes.contains_key(KEY_ANALYZED_SPANS));
        assert!(!trace.spans()[1].attributes.contains_key(KEY_ANALYZED_SPANS));
    }

    #[test]
    fn single_span_sampling_wins_over_analytics_events() {
        let rates = HashMap::from([("test-service".to_string(), 1.0)]);
        let mut sampler = create_test_sampler();
        sampler.sampling_rate = 0.0; // probabilistic drops everything
        sampler.probabilistic_sampler_enabled = true;
        sampler.event_processor = EventProcessor::for_tests(&HashMap::new(), &rates, f64::INFINITY);

        let mut attrs = saluki_common::collections::FastHashMap::default();
        attrs.insert(
            MetaString::from(KEY_SPAN_SAMPLING_MECHANISM),
            AttributeValue::Float(8.0),
        );
        let sss_span = create_test_span(1, 0).with_attributes(attrs);
        let mut trace = create_test_trace(vec![sss_span]);

        assert!(sampler.process_trace(&mut trace));
        // The trace is forwarded as kept by single span sampling, not as dropped events.
        assert_eq!(trace.spans().len(), 1);
        assert!(!trace.dropped_trace);
        assert_eq!(trace.priority, Some(PRIORITY_USER_KEEP));
    }

    #[test]
    fn manual_drop_skips_event_extraction() {
        let rates = HashMap::from([("test-service".to_string(), 1.0)]);
        let mut sampler = create_test_sampler();
        // The legacy path is the one that consults user priority.
        sampler.probabilistic_sampler_enabled = false;
        sampler.event_processor = EventProcessor::for_tests(&HashMap::new(), &rates, f64::INFINITY);

        // Manual drop short-circuits the samplers, and the trace is not promoted to events.
        let root = create_test_span(1, 0);
        let mut trace = create_test_trace(vec![root]);
        trace.priority = Some(PRIORITY_USER_DROP);

        assert!(!sampler.process_trace(&mut trace));
        assert!(!trace.spans()[0].attributes.contains_key(KEY_ANALYZED_SPANS));
    }

    // Remote sampling tests based on the Datadog Agent's remote configuration tests:
    // https://github.com/DataDog/datadog-agent/blob/17ecddf4e3e/pkg/trace/remoteconfighandler/remote_config_handler_test.go

    // With probabilistic sampling set to 0, only the rare sampler can keep test traces.
    fn create_remote_sampler(errors_per_second: f64) -> (TestSamplingPublisher, TraceSampler) {
        let (publisher, remote_sampling) = TestSamplingPublisher::new();
        let sampler = TraceSampler {
            sampling_rate: 0.0,
            error_sampler: errors::ErrorsSampler::new(errors_per_second, 1.0),
            rare_sampler: rare_sampler::RareSampler::new(
                false,
                1000.0,
                std::time::Duration::from_secs(300),
                200,
                telemetry::SamplerCounters::new(),
            ),
            static_settings: SamplerSettings {
                target_traces_per_second: 10.0,
                errors_per_second,
                rare_sampler_enabled: false,
            },
            remote_sampling,
            ..create_test_sampler()
        };
        (publisher, sampler)
    }

    /// The agent environment resolved from `env: Prod`.
    fn prod_env() -> AgentEnv {
        AgentEnv::resolve(None, Some("Prod"), &[], &[])
    }

    /// Traces settings where probabilistic sampling drops every trace, so only the rare sampler keeps one.
    fn rare_only_traces() -> domains::traces::Domain {
        domains::traces::Domain {
            target_traces_per_second: 7.0,
            errors_per_second: 9.0,
            enable_rare_sampler: true,
            probabilistic_sampler: domains::traces::ProbabilisticSampler {
                enabled: true,
                sampling_percentage: 0.0,
                ..Default::default()
            },
            ..Default::default()
        }
    }

    fn transform_empty_buffer(sampler: &mut TraceSampler) {
        sampler.process_buffer(&mut EventsBuffer::default());
    }

    fn live_tps(sampler: &TraceSampler) -> (f64, f64) {
        (
            sampler.priority_sampler.get_target_tps(),
            sampler.error_sampler.get_target_tps(),
        )
    }

    fn rare_keeps_new_trace(sampler: &mut TraceSampler, span_id: u64) -> bool {
        let mut trace = create_test_trace(vec![
            create_top_level_span(span_id).with_service(MetaString::from(format!("rare-probe-{span_id}")))
        ]);
        sampler.run_samplers(&mut trace).keep
    }

    /// Covers the Datadog Agent's priority, error, and rare sampler update cases.
    #[test]
    fn remote_configuration_changes_live_sampler_settings() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":41,"errors_sampler_target_TPS":42,
                "rare_sampler_enabled":true}}"#,
        )]);
        assert_eq!(live_tps(&sampler), (10.0, 10.0), "nothing applies before a buffer");
        assert!(!rare_keeps_new_trace(&mut sampler, 1));

        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 42.0));
        assert!(rare_keeps_new_trace(&mut sampler, 2));
        // Remote settings do not change the no-priority sampler.
        assert_eq!(sampler.no_priority_sampler.test_target_tps(), 10.0);
    }

    #[test]
    fn more_than_one_remote_configuration_changes_nothing() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        publisher.assign([
            (
                "a",
                r#"{"all_envs":{"priority_sampler_target_TPS":41,"rare_sampler_enabled":true}}"#,
            ),
            ("b", r#"{"all_envs":{"errors_sampler_target_TPS":42}}"#),
        ]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (10.0, 10.0));
        assert!(!rare_keeps_new_trace(&mut sampler, 1));
    }

    #[test]
    fn more_than_one_remote_configuration_after_one_keeps_remote_settings() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        publisher.assign([(
            "a",
            r#"{"all_envs":{"priority_sampler_target_TPS":41,"errors_sampler_target_TPS":42,
                "rare_sampler_enabled":true}}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 42.0));

        publisher.assign([
            ("a", r#"{"all_envs":{"priority_sampler_target_TPS":7}}"#),
            ("b", r#"{"all_envs":{"errors_sampler_target_TPS":8}}"#),
        ]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(
            live_tps(&sampler),
            (41.0, 42.0),
            "the rejected assignment changes nothing"
        );
        assert!(rare_keeps_new_trace(&mut sampler, 1));
    }

    /// Remote sampling uses the resolved agent environment to select a `by_env` entry.
    #[test]
    fn remote_sampling_subscription_reaches_the_built_sampler() {
        let (publisher, subscription) = TestSamplingPublisher::new();
        let config = TraceSamplerConfiguration::from_configuration(
            &rare_only_traces(),
            &domains::otlp::Traces::default(),
            prod_env(),
        )
        .with_remote_sampling(subscription);
        let mut sampler = config.build_sampler();
        assert_eq!(live_tps(&sampler), (7.0, 9.0));
        assert!(rare_keeps_new_trace(&mut sampler, 1));

        // The Datadog Agent's tests send unset fields as `null`.
        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":null,"errors_sampler_target_TPS":null,
                "rare_sampler_enabled":null},
                "by_env":[{"env":"prod","config":{"priority_sampler_target_TPS":41,"errors_sampler_target_TPS":42,
                "rare_sampler_enabled":false}}]}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 42.0), "the normalized env selects `prod`");
        assert!(!rare_keeps_new_trace(&mut sampler, 2));

        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":null,"errors_sampler_target_TPS":null,
                "rare_sampler_enabled":null},"by_env":null}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(
            live_tps(&sampler),
            (7.0, 9.0),
            "each setting reverts to its static value"
        );
        assert!(rare_keeps_new_trace(&mut sampler, 3));
    }

    /// Remote selection matches the resolved agent environment exactly.
    #[test]
    fn remote_selection_matches_the_resolved_agent_env() {
        let payload = r#"{"all_envs":{"errors_sampler_target_TPS":30},
            "by_env":[{"env":"prod","config":{"priority_sampler_target_TPS":42}},
                      {"env":"none","config":{"priority_sampler_target_TPS":41}}]}"#;
        let cases = [
            // Nothing configured resolves to `none`.
            (AgentEnv::resolve(None, None, &[], &[]), (41.0, 30.0)),
            (prod_env(), (42.0, 30.0)),
            // `!!!` normalizes to empty, which matches no entry: `none` is not applied after normalization.
            (AgentEnv::resolve(Some("!!!"), None, &[], &[]), (7.0, 30.0)),
        ];
        for (agent_env, expected) in cases {
            let (publisher, subscription) = TestSamplingPublisher::new();
            let env = agent_env.as_str().to_owned();
            let mut sampler = TraceSamplerConfiguration::from_configuration(
                &rare_only_traces(),
                &domains::otlp::Traces::default(),
                agent_env,
            )
            .with_remote_sampling(subscription)
            .build_sampler();
            publisher.assign([("sampling", payload)]);
            transform_empty_buffer(&mut sampler);
            assert_eq!(live_tps(&sampler), expected, "agent env {env:?}");
        }
    }

    /// Remote selection and the priority sampler's fallback for traces without an env use the same agent environment.
    ///
    /// Regression: remote selection followed `env` while the priority sampler used a separate default, so the two
    /// could disagree.
    #[test]
    fn remote_selection_and_priority_sampling_share_agent_env() {
        let (publisher, subscription) = TestSamplingPublisher::new();
        let agent_env = AgentEnv::resolve(Some("Staging"), Some("prod"), &[], &[]);
        let mut sampler = TraceSamplerConfiguration::from_configuration(
            &rare_only_traces(),
            &domains::otlp::Traces::default(),
            agent_env,
        )
        .with_remote_sampling(subscription)
        .build_sampler();
        publisher.assign([(
            "sampling",
            r#"{"by_env":[{"env":"staging","config":{"priority_sampler_target_TPS":43}},
                          {"env":"prod","config":{"priority_sampler_target_TPS":42}}]}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler).0, 43.0, "remote selection uses the agent env");

        // Traces without an env learn their rate under the agent env. Counts in one bucket become rates when the next
        // begins.
        let start = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_000_000);
        for now in [start, start + std::time::Duration::from_secs(10)] {
            for trace_id in 0..100 {
                let mut trace = create_test_trace(vec![create_test_span(trace_id, 0)]);
                sampler
                    .priority_sampler
                    .sample(now, &mut trace, 0, PRIORITY_AUTO_KEEP, 0.0);
            }
        }
        let learned: Vec<_> = sampler
            .priority_sampler
            .test_signature_sample_rates()
            .0
            .into_keys()
            .collect();
        assert_eq!(
            learned,
            vec![signature::ServiceSignature::new("test-service", "staging").hash()],
            "priority sampling uses the agent env"
        );
    }

    /// The run loop applies a published update before sampling the next buffer.
    #[tokio::test]
    async fn run_loop_samples_with_the_published_remote_settings() {
        let (publisher, subscription) = TestSamplingPublisher::new();
        let config = TraceSamplerConfiguration::from_configuration(
            &rare_only_traces(),
            &domains::otlp::Traces::default(),
            prod_env(),
        )
        .with_remote_sampling(subscription);
        let (control, mut outputs) = TestComponentDriver::transform(config)
            .await
            .expect("the sampler should start");
        let new_service_trace = |span_id: u64| {
            Event::Trace(create_test_trace(vec![
                create_top_level_span(span_id).with_service(MetaString::from(format!("rare-probe-{span_id}")))
            ]))
        };

        // The rare sampler keeps the first trace. Waiting for it makes sure the update is published after it.
        control.send_event(new_service_trace(1)).await;
        outputs
            .default_output()
            .next()
            .await
            .expect("the rare sampler should keep the first trace");

        publisher.assign([("sampling", r#"{"all_envs":{"rare_sampler_enabled":false}}"#)]);
        control.send_event(new_service_trace(2)).await;
        control.shutdown().await.expect("the sampler should stop cleanly");

        let kept = outputs.default_output().collect_events().await;
        assert!(
            kept.is_empty(),
            "the second trace should be dropped now that remote settings disable the rare sampler: {kept:?}"
        );
    }

    #[test]
    fn remote_field_removed_later_falls_back_to_static() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":41,"rare_sampler_enabled":true}}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 10.0));

        publisher.assign([("sampling", r#"{"all_envs":{"errors_sampler_target_TPS":42}}"#)]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(
            live_tps(&sampler),
            (10.0, 42.0),
            "the priority target reverts to static, not to 41"
        );
        assert!(
            !rare_keeps_new_trace(&mut sampler, 1),
            "the rare sampler reverts to static, not to enabled"
        );
    }

    /// Starting with a zero target disables error sampling; a remote update cannot re-enable it.
    #[test]
    fn static_zero_errors_tps_stays_disabled_under_remote_configuration() {
        let (publisher, mut sampler) = create_remote_sampler(0.0);
        publisher.assign([("sampling", r#"{"all_envs":{"errors_sampler_target_TPS":5}}"#)]);
        transform_empty_buffer(&mut sampler);
        assert_eq!(sampler.error_sampler.get_target_tps(), 5.0);

        let mut trace = create_test_trace(vec![create_test_span(1, 1)]);
        assert!(!sampler
            .error_sampler
            .sample_error(std::time::SystemTime::now(), &mut trace, 0));
    }

    #[test]
    fn inert_remote_sampling_changes_nothing() {
        let mut sampler = TraceSampler {
            remote_sampling: TraceSamplingSubscription::default(),
            ..create_remote_sampler(10.0).1
        };
        // An inert subscription must not reset these overrides.
        sampler.priority_sampler.update_target_tps(41.0);
        sampler.error_sampler.update_target_tps(42.0);
        sampler.rare_sampler.set_enabled(true);

        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 42.0));
        assert!(rare_keeps_new_trace(&mut sampler, 1));
    }

    /// Removing the remote assignment does not undo previously applied settings.
    #[test]
    fn unassigned_after_remote_configuration_keeps_remote_settings() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":41,"errors_sampler_target_TPS":42,
                "rare_sampler_enabled":true}}"#,
        )]);
        transform_empty_buffer(&mut sampler);
        let applied = sampler.remote_sampling.current().expect("a configuration is accepted");

        publisher.assign(std::iter::empty::<(&str, &str)>());
        assert!(Arc::ptr_eq(
            &applied,
            &sampler.remote_sampling.current().expect("the configuration is kept")
        ));
        transform_empty_buffer(&mut sampler);
        assert_eq!(live_tps(&sampler), (41.0, 42.0));
        assert!(rare_keeps_new_trace(&mut sampler, 1));
    }

    // The subscription keeps only the latest state, and the sampler reads it once per buffer. Publications the sampler
    // has not yet read must not erase the newest accepted configuration, since the Datadog Agent applies each update
    // as it arrives, ignores an empty one, and rejects arrays where its JSON decoder expects an object:
    // https://github.com/DataDog/datadog-agent/blob/17ecddf4e3e83ccbb0e68aeb99461d1a1d902927/pkg/trace/remoteconfighandler/remote_config_handler.go#L234-L295
    #[test]
    fn unread_publications_keep_the_newest_accepted_configuration() {
        const A: &str = r#"{"all_envs":{"priority_sampler_target_TPS":41,"errors_sampler_target_TPS":42,
            "rare_sampler_enabled":true}}"#;
        const B: &str = r#"{"all_envs":{"priority_sampler_target_TPS":43,"errors_sampler_target_TPS":44}}"#;
        // Each case is a name, the payload applied before the publications under test, those publications (`None`
        // for an empty assignment), and the expected targets and rare sampler state.
        type Case<'a> = (&'a str, Option<&'a str>, &'a [Option<&'a str>], (f64, f64), bool);
        let cases: [Case; 7] = [
            ("initial empty assignment", None, &[None], (10.0, 10.0), false),
            (
                "A then empty before the first buffer",
                None,
                &[Some(A), None],
                (41.0, 42.0),
                true,
            ),
            (
                "B then empty after A is applied",
                Some(A),
                &[Some(B), None],
                (43.0, 44.0),
                false,
            ),
            (
                "malformed after A is applied",
                Some(A),
                &[Some("not json")],
                (41.0, 42.0),
                true,
            ),
            (
                "array payload after A is applied",
                Some(A),
                &[Some("[]")],
                (41.0, 42.0),
                true,
            ),
            (
                "array settings after A is applied",
                Some(A),
                &[Some(r#"{"all_envs":[0,0,false]}"#)],
                (41.0, 42.0),
                true,
            ),
            (
                "accepted {} after A is applied",
                Some(A),
                &[Some("{}")],
                (10.0, 10.0),
                false,
            ),
        ];

        for (name, applied, pending, tps, rare) in cases {
            let (publisher, mut sampler) = create_remote_sampler(10.0);
            let publish = |payload: &Option<&str>| match payload {
                Some(payload) => publisher.assign([("sampling", *payload)]),
                None => publisher.assign(std::iter::empty::<(&str, &str)>()),
            };
            if let Some(applied) = applied {
                publish(&Some(applied));
                transform_empty_buffer(&mut sampler);
            }
            pending.iter().for_each(publish);
            transform_empty_buffer(&mut sampler);
            assert_eq!(live_tps(&sampler), tps, "{name}");
            assert_eq!(rare_keeps_new_trace(&mut sampler, 1), rare, "{name}");
        }
    }

    /// The run loop applies a configuration whose withdrawal was published before the sampler's first buffer.
    #[tokio::test]
    async fn run_loop_applies_a_configuration_withdrawn_before_the_first_buffer() {
        let (publisher, subscription) = TestSamplingPublisher::new();
        let config = TraceSamplerConfiguration::from_configuration(
            &rare_only_traces(),
            &domains::otlp::Traces::default(),
            prod_env(),
        )
        .with_remote_sampling(subscription);
        publisher.assign([("sampling", r#"{"all_envs":{"rare_sampler_enabled":false}}"#)]);
        publisher.assign(std::iter::empty::<(&str, &str)>());

        let (control, mut outputs) = TestComponentDriver::transform(config)
            .await
            .expect("the sampler should start");
        control
            .send_event(Event::Trace(create_test_trace(vec![
                create_top_level_span(1).with_service(MetaString::from("rare-probe-1"))
            ])))
            .await;
        control.shutdown().await.expect("the sampler should stop cleanly");

        let kept = outputs.default_output().collect_events().await;
        assert!(
            kept.is_empty(),
            "the remote configuration disables the rare sampler, which would otherwise keep the trace: {kept:?}"
        );
    }

    /// Changing the target scales rates already learned for trace signatures.
    #[test]
    fn remote_configuration_scales_learned_rates() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);

        // 100 traces in one 5-second bucket are 20 traces per second, so a target of 10 learns a rate of 0.5 once the
        // next bucket begins.
        let start = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_000_000);
        for now in [start, start + std::time::Duration::from_secs(10)] {
            for trace_id in 0..100 {
                let mut trace = create_test_trace(vec![create_test_span(trace_id, 1)]);
                sampler
                    .priority_sampler
                    .sample(now, &mut trace, 0, PRIORITY_AUTO_KEEP, 0.0);
                sampler.error_sampler.sample_error(now, &mut trace, 0);
            }
        }
        let learned_rates = |sampler: &TraceSampler| {
            let priority: Vec<f64> = sampler
                .priority_sampler
                .test_signature_sample_rates()
                .0
                .into_values()
                .collect();
            let errors: Vec<f64> = sampler
                .error_sampler
                .test_signature_sample_rates()
                .0
                .into_values()
                .collect();
            (priority, errors)
        };
        assert_eq!(learned_rates(&sampler), (vec![0.5], vec![0.5]));

        publisher.assign([(
            "sampling",
            r#"{"all_envs":{"priority_sampler_target_TPS":5,"errors_sampler_target_TPS":2}}"#,
        )]);
        transform_empty_buffer(&mut sampler);

        assert_eq!(live_tps(&sampler), (5.0, 2.0));
        assert_eq!(
            learned_rates(&sampler),
            (vec![0.5 * (5.0 / 10.0)], vec![0.5 * (2.0 / 10.0)]),
            "each learned rate scales by the ratio of the new target to the old one"
        );
    }

    /// Republishing the same settings must not clear learned rates or rare-trace history.
    #[test]
    fn reapplying_equal_remote_settings_keeps_sampler_state() {
        let (publisher, mut sampler) = create_remote_sampler(10.0);
        let payload = r#"{"all_envs":{"priority_sampler_target_TPS":5,"errors_sampler_target_TPS":4,
            "rare_sampler_enabled":true}}"#;
        publisher.assign([("sampling", payload)]);
        transform_empty_buffer(&mut sampler);

        // Learn a rate for one signature in each sampler: counts in one bucket become rates when the next begins.
        let start = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_000_000);
        for now in [start, start + std::time::Duration::from_secs(10)] {
            for trace_id in 0..100 {
                let mut trace = create_test_trace(vec![create_test_span(trace_id, 1)]);
                sampler
                    .priority_sampler
                    .sample(now, &mut trace, 0, PRIORITY_AUTO_KEEP, 0.0);
                sampler.error_sampler.sample_error(now, &mut trace, 0);
            }
        }
        assert!(rare_keeps_new_trace(&mut sampler, 1));

        #[derive(Debug, PartialEq)]
        struct LearnedState {
            priority_rates: (saluki_common::collections::FastHashMap<signature::Signature, f64>, f64),
            priority_size: i64,
            errors_rates: (saluki_common::collections::FastHashMap<signature::Signature, f64>, f64),
            errors_size: i64,
            tps: (f64, f64),
        }
        let state = |sampler: &TraceSampler| LearnedState {
            priority_rates: sampler.priority_sampler.test_signature_sample_rates(),
            priority_size: sampler.priority_sampler.test_size(),
            errors_rates: sampler.error_sampler.test_signature_sample_rates(),
            errors_size: sampler.error_sampler.test_size(),
            tps: live_tps(sampler),
        };
        let before = state(&sampler);
        assert_eq!(
            before.priority_rates.0.len(),
            1,
            "the priority sampler has learned a rate"
        );
        assert_eq!(before.errors_rates.0.len(), 1, "the errors sampler has learned a rate");

        let applied = sampler.remote_sampling.current().expect("a configuration is accepted");
        publisher.assign([("sampling", payload)]);
        let republished = sampler.remote_sampling.current().expect("a configuration is accepted");
        assert!(
            !Arc::ptr_eq(&applied, &republished),
            "the equal snapshot is a new allocation"
        );
        transform_empty_buffer(&mut sampler);

        assert_eq!(state(&sampler), before);
        let mut seen = create_test_trace(vec![
            create_top_level_span(1).with_service(MetaString::from("rare-probe-1"))
        ]);
        let SamplerOutcome { keep, .. } = sampler.run_samplers(&mut seen);
        assert!(!keep, "the rare sampler still remembers the signature it kept");
    }
}