Skip to main content

vector_common/
sampling.rs

1/// Stateful sampler that retains events at a configured ratio.
2///
3/// Each call to [`Self::sample`] advances the sampler. Ratios must be between zero and one,
4/// inclusive.
5#[derive(Clone, Debug)]
6pub struct RatioSampler {
7    ratio: f64,
8    value: f64,
9}
10
11impl RatioSampler {
12    /// Creates a sampler that retains events at `ratio`.
13    #[must_use]
14    pub fn new(ratio: f64) -> Self {
15        debug_assert!((0.0..=1.0).contains(&ratio));
16        Self {
17            ratio,
18            value: 1.0 - ratio,
19        }
20    }
21
22    /// Advances the sampler and returns whether the current event is retained.
23    pub fn sample(&mut self) -> bool {
24        let increment = self.value + self.ratio;
25        self.value = if increment >= 1.0 {
26            increment - 1.0
27        } else {
28            increment
29        };
30        increment >= 1.0
31    }
32}
33
34#[cfg(test)]
35mod tests {
36    use super::RatioSampler;
37
38    #[test]
39    fn retains_events_at_half_ratio() {
40        let mut sampler = RatioSampler::new(0.5);
41        let decisions = (0..6).map(|_| sampler.sample()).collect::<Vec<_>>();
42        assert_eq!(decisions, [true, false, true, false, true, false]);
43    }
44
45    #[test]
46    fn retains_every_event_at_full_ratio() {
47        let mut sampler = RatioSampler::new(1.0);
48        assert!((0..6).all(|_| sampler.sample()));
49    }
50
51    #[test]
52    fn sampler_instances_advance_independently() {
53        let mut first = RatioSampler::new(0.5);
54        let mut second = RatioSampler::new(0.5);
55
56        assert!(first.sample());
57        assert!(!first.sample());
58        assert!(second.sample());
59    }
60}