1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

//! `Sampler` trait and its built-in implementations.
use crate::span::CandidateSpan;
use crate::{ErrorKind, Result};
use rand::{self, Rng};

/// `Sampler` decides whether a new trace should be sampled or not.
pub trait Sampler<T> {
    /// This method decides whether a trace with given `span` should be sampled.
    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool;

    /// Returns the sampler that samples a trace if `self` or `other` decides to sample it.
    fn or<U>(self, other: U) -> OrSampler<Self, U>
    where
        Self: Sized,
        U: Sampler<T>,
    {
        OrSampler(self, other)
    }

    /// Returns the sampler that samples a trace if both of `self` and `other` decides to sample it.
    fn and<U>(self, other: U) -> AndSampler<Self, U>
    where
        Self: Sized,
        U: Sampler<T>,
    {
        AndSampler(self, other)
    }

    /// Converts into `BoxSampler`.
    fn boxed(self) -> BoxSampler<T>
    where
        Self: Sized + Send + Sync + 'static,
    {
        Box::new(self)
    }
}
impl<T> Sampler<T> for BoxSampler<T> {
    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
        (**self).is_sampled(span)
    }
    fn boxed(self) -> BoxSampler<T>
    where
        Self: Sized + Send + 'static,
    {
        self
    }
}

/// Boxed version of `Sampler`.
pub type BoxSampler<T> = Box<Sampler<T> + Send + Sync + 'static>;

/// This samples a certain percentage of traces.
#[derive(Debug, Clone)]
pub struct ProbabilisticSampler {
    sampling_rate: f64,
}
impl ProbabilisticSampler {
    /// Makes a new `ProbabilisticSampler` instance.
    ///
    /// # Errors
    ///
    /// If `sampling_rate` is not in the range `0.0...1.0`,
    /// it will return an error with the kind `ErrorKind::InvalidInput`.
    pub fn new(sampling_rate: f64) -> Result<Self> {
        track_assert!(0.0 <= sampling_rate, ErrorKind::InvalidInput);
        track_assert!(sampling_rate <= 1.0, ErrorKind::InvalidInput);
        Ok(ProbabilisticSampler { sampling_rate })
    }
}
impl<T> Sampler<T> for ProbabilisticSampler {
    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
        rand::thread_rng().gen_range(0.0, 1.0) < self.sampling_rate
    }
}

/// This samples traces which have one or more references.
#[derive(Debug, Clone)]
pub struct PassiveSampler;
impl<T> Sampler<T> for PassiveSampler {
    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
        !span.references().is_empty()
    }
}

/// This samples no traces.
#[derive(Debug, Clone)]
pub struct NullSampler;
impl<T> Sampler<T> for NullSampler {
    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
        false
    }
}

/// This samples all traces.
#[derive(Debug, Clone)]
pub struct AllSampler;
impl<T> Sampler<T> for AllSampler {
    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
        true
    }
}

/// `or` combinator.
#[derive(Debug, Clone)]
pub struct OrSampler<A, B>(A, B);
impl<A, B, T> Sampler<T> for OrSampler<A, B>
where
    A: Sampler<T>,
    B: Sampler<T>,
{
    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
        self.0.is_sampled(span) || self.1.is_sampled(span)
    }
}

/// `and` combinator.
#[derive(Debug, Clone)]
pub struct AndSampler<A, B>(A, B);
impl<A, B, T> Sampler<T> for AndSampler<A, B>
where
    A: Sampler<T>,
    B: Sampler<T>,
{
    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
        self.0.is_sampled(span) && self.1.is_sampled(span)
    }
}