cf_rustracing/
sampler.rs

1//! `Sampler` trait and its built-in implementations.
2use crate::span::CandidateSpan;
3use crate::{ErrorKind, Result};
4use rand::{self, Rng};
5use std::fmt::Debug;
6
7/// `Sampler` decides whether a new trace should be sampled or not.
8pub trait Sampler<T>: Debug {
9    /// This method decides whether a trace with given `span` should be sampled.
10    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool;
11
12    /// Returns the sampler that samples a trace if `self` or `other` decides to sample it.
13    fn or<U>(self, other: U) -> OrSampler<Self, U>
14    where
15        Self: Sized,
16        U: Sampler<T>,
17    {
18        OrSampler(self, other)
19    }
20
21    /// Returns the sampler that samples a trace if both of `self` and `other` decides to sample it.
22    fn and<U>(self, other: U) -> AndSampler<Self, U>
23    where
24        Self: Sized,
25        U: Sampler<T>,
26    {
27        AndSampler(self, other)
28    }
29
30    /// Converts into `BoxSampler`.
31    fn boxed(self) -> BoxSampler<T>
32    where
33        Self: Sized + Send + Sync + 'static,
34    {
35        Box::new(self)
36    }
37}
38impl<T> Sampler<T> for BoxSampler<T> {
39    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
40        (**self).is_sampled(span)
41    }
42    fn boxed(self) -> BoxSampler<T>
43    where
44        Self: Sized + Send + 'static,
45    {
46        self
47    }
48}
49
50/// Boxed version of `Sampler`.
51pub type BoxSampler<T> = Box<dyn Sampler<T> + Send + Sync + 'static>;
52
53/// This samples a certain percentage of traces.
54#[derive(Debug, Clone)]
55pub struct ProbabilisticSampler {
56    sampling_rate: f64,
57}
58impl ProbabilisticSampler {
59    /// Makes a new `ProbabilisticSampler` instance.
60    ///
61    /// # Errors
62    ///
63    /// If `sampling_rate` is not in the range `0.0...1.0`,
64    /// it will return an error with the kind `ErrorKind::InvalidInput`.
65    pub fn new(sampling_rate: f64) -> Result<Self> {
66        track_assert!(0.0 <= sampling_rate, ErrorKind::InvalidInput);
67        track_assert!(sampling_rate <= 1.0, ErrorKind::InvalidInput);
68        Ok(ProbabilisticSampler { sampling_rate })
69    }
70}
71impl<T> Sampler<T> for ProbabilisticSampler {
72    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
73        rand::thread_rng().gen_range(0.0..1.0) < self.sampling_rate
74    }
75}
76
77/// This samples traces which have one or more references.
78#[derive(Debug, Clone)]
79pub struct PassiveSampler;
80impl<T> Sampler<T> for PassiveSampler {
81    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
82        !span.references().is_empty()
83    }
84}
85
86/// This samples no traces.
87#[derive(Debug, Clone)]
88pub struct NullSampler;
89impl<T> Sampler<T> for NullSampler {
90    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
91        false
92    }
93}
94
95/// This samples all traces.
96#[derive(Debug, Clone)]
97pub struct AllSampler;
98impl<T> Sampler<T> for AllSampler {
99    fn is_sampled(&self, _span: &CandidateSpan<T>) -> bool {
100        true
101    }
102}
103
104/// `or` combinator.
105#[derive(Debug, Clone)]
106pub struct OrSampler<A, B>(A, B);
107impl<A, B, T> Sampler<T> for OrSampler<A, B>
108where
109    A: Sampler<T>,
110    B: Sampler<T>,
111{
112    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
113        self.0.is_sampled(span) || self.1.is_sampled(span)
114    }
115}
116
117/// `and` combinator.
118#[derive(Debug, Clone)]
119pub struct AndSampler<A, B>(A, B);
120impl<A, B, T> Sampler<T> for AndSampler<A, B>
121where
122    A: Sampler<T>,
123    B: Sampler<T>,
124{
125    fn is_sampled(&self, span: &CandidateSpan<T>) -> bool {
126        self.0.is_sampled(span) && self.1.is_sampled(span)
127    }
128}