grafos_observe/

sampling.rs

//! Sampling strategies for distributed tracing.
//!
//! Sampling controls which traces are recorded and exported. The decision is
//! made at the trace root (head-based sampling) and propagated to all child
//! spans via the [`FLAG_SAMPLED`](crate::trace::FLAG_SAMPLED) bit in
//! [`TraceContext::flags`](crate::trace::TraceContext).

use crate::trace::TraceContext;

/// Sampling strategy enumeration.
#[derive(Debug, Clone, PartialEq)]
pub enum SamplingStrategy {
    /// Sample all traces (default for development).
    AlwaysOn,
    /// Sample no traces (useful for benchmarks).
    AlwaysOff,
    /// Probabilistic head-based sampling. The `rate` field in
    /// [`SamplingConfig`] controls the probability (0.0 to 1.0).
    HeadBased,
    /// Rate-limited sampling: at most N traces per second.
    RateBased {
        /// Maximum traces per second.
        traces_per_sec: u32,
    },
}

/// Sampling configuration.
#[derive(Debug, Clone)]
pub struct SamplingConfig {
    /// Sample rate (0.0 to 1.0). Only used with [`SamplingStrategy::HeadBased`].
    pub rate: f64,
    /// Strategy to use.
    pub strategy: SamplingStrategy,
}

impl Default for SamplingConfig {
    /// Default: AlwaysOn (sample everything, suitable for dev).
    fn default() -> Self {
        SamplingConfig {
            rate: 1.0,
            strategy: SamplingStrategy::AlwaysOn,
        }
    }
}

/// A sampler that decides whether a trace should be recorded.
pub struct Sampler {
    config: SamplingConfig,
    /// For RateBased: count of traces started in the current second window.
    rate_count: u32,
    /// For RateBased: the second (Unix timestamp) of the current window.
    rate_window_sec: u64,
}

impl Sampler {
    /// Create a new sampler with the given configuration.
    pub fn new(config: SamplingConfig) -> Self {
        Sampler {
            config,
            rate_count: 0,
            rate_window_sec: 0,
        }
    }

    /// Decide whether a new root trace should be sampled.
    ///
    /// For child spans, the decision is inherited from the parent trace
    /// context (if the parent is sampled, the child is sampled).
    ///
    /// `trace_id_low` is the low 64 bits of the trace ID, used as a
    /// deterministic hash for HeadBased sampling. `now_unix_sec` is the
    /// current Unix timestamp in seconds, used for RateBased windows.
    pub fn should_sample(&mut self, trace_id_low: u64, now_unix_sec: u64) -> bool {
        match &self.config.strategy {
            SamplingStrategy::AlwaysOn => true,
            SamplingStrategy::AlwaysOff => false,
            SamplingStrategy::HeadBased => {
                // Use the low bits of trace_id as a deterministic hash.
                // Scale to 0.0..1.0 range and compare with rate.
                let hash_f = (trace_id_low as f64) / (u64::MAX as f64);
                hash_f < self.config.rate
            }
            SamplingStrategy::RateBased { traces_per_sec } => {
                if now_unix_sec != self.rate_window_sec {
                    self.rate_window_sec = now_unix_sec;
                    self.rate_count = 0;
                }
                if self.rate_count < *traces_per_sec {
                    self.rate_count += 1;
                    true
                } else {
                    false
                }
            }
        }
    }

    /// Check if a child span should be sampled, given its parent context.
    ///
    /// If the parent is sampled, the child is always sampled (propagation).
    pub fn should_sample_child(parent: &TraceContext) -> bool {
        parent.is_sampled()
    }

    /// Get the current configuration.
    pub fn config(&self) -> &SamplingConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn always_on_samples_everything() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 1.0,
            strategy: SamplingStrategy::AlwaysOn,
        });
        for i in 0..100 {
            assert!(sampler.should_sample(i, 1000));
        }
    }

    #[test]
    fn always_off_samples_nothing() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 0.0,
            strategy: SamplingStrategy::AlwaysOff,
        });
        for i in 0..100 {
            assert!(!sampler.should_sample(i, 1000));
        }
    }

    #[test]
    fn head_based_rate_zero_samples_none() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 0.0,
            strategy: SamplingStrategy::HeadBased,
        });
        // rate=0.0 means nothing should be sampled
        let mut sampled = 0;
        for i in 0..1000u64 {
            if sampler.should_sample(i, 1000) {
                sampled += 1;
            }
        }
        assert_eq!(sampled, 0);
    }

    #[test]
    fn head_based_rate_one_samples_all() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 1.0,
            strategy: SamplingStrategy::HeadBased,
        });
        let mut sampled = 0;
        for i in 0..1000u64 {
            if sampler.should_sample(i, 1000) {
                sampled += 1;
            }
        }
        assert_eq!(sampled, 1000);
    }

    #[test]
    fn head_based_approximately_half() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 0.5,
            strategy: SamplingStrategy::HeadBased,
        });
        let mut sampled = 0u64;
        // Use well-distributed trace IDs
        for i in 0..10_000u64 {
            let trace_id_low = i
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            if sampler.should_sample(trace_id_low, 1000) {
                sampled += 1;
            }
        }
        // Should be approximately 5000, allow 15% tolerance
        assert!(
            sampled > 4000 && sampled < 6000,
            "expected ~5000, got {sampled}"
        );
    }

    #[test]
    fn rate_based_caps_per_second() {
        let mut sampler = Sampler::new(SamplingConfig {
            rate: 1.0,
            strategy: SamplingStrategy::RateBased { traces_per_sec: 5 },
        });

        // First 5 in second 100 should be sampled
        for i in 0..5 {
            assert!(sampler.should_sample(i, 100), "trace {i} should be sampled");
        }
        // 6th should be rejected
        assert!(!sampler.should_sample(5, 100));

        // New second resets the window
        assert!(sampler.should_sample(6, 101));
    }

    #[test]
    fn child_inherits_parent_sampling() {
        let mut ctx = TraceContext::default();
        ctx.set_sampled(true);
        assert!(Sampler::should_sample_child(&ctx));

        ctx.set_sampled(false);
        assert!(!Sampler::should_sample_child(&ctx));
    }

    #[test]
    fn default_config_is_always_on() {
        let config = SamplingConfig::default();
        assert_eq!(config.strategy, SamplingStrategy::AlwaysOn);
        assert!((config.rate - 1.0).abs() < f64::EPSILON);
    }
}