rocketmq_observability/
sampling.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::atomic::AtomicU64;
16use std::sync::atomic::Ordering;
17
18const SAMPLE_SCALE: u64 = 1_000_000;
19
20#[derive(Debug)]
21pub struct SamplingGate {
22    sampled_per_million: u64,
23    accumulator: AtomicU64,
24}
25
26impl SamplingGate {
27    pub fn new(sample_ratio: f64) -> Self {
28        Self {
29            sampled_per_million: sampled_per_million(sample_ratio),
30            accumulator: AtomicU64::new(0),
31        }
32    }
33
34    #[inline]
35    pub fn should_sample(&self) -> bool {
36        match self.sampled_per_million {
37            0 => false,
38            SAMPLE_SCALE => true,
39            sampled_per_million => {
40                let previous = self.accumulator.fetch_add(sampled_per_million, Ordering::Relaxed);
41                previous % SAMPLE_SCALE + sampled_per_million >= SAMPLE_SCALE
42            }
43        }
44    }
45
46    pub fn sampled_per_million(&self) -> u64 {
47        self.sampled_per_million
48    }
49}
50
51impl Default for SamplingGate {
52    fn default() -> Self {
53        Self::new(1.0)
54    }
55}
56
57fn sampled_per_million(sample_ratio: f64) -> u64 {
58    if sample_ratio.is_nan() || sample_ratio <= 0.0 {
59        return 0;
60    }
61    if sample_ratio >= 1.0 {
62        return SAMPLE_SCALE;
63    }
64    (sample_ratio * SAMPLE_SCALE as f64).ceil() as u64
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn full_sample_ratio_always_samples() {
73        let gate = SamplingGate::new(1.0);
74
75        assert_eq!(gate.sampled_per_million(), SAMPLE_SCALE);
76        assert!((0..8).all(|_| gate.should_sample()));
77    }
78
79    #[test]
80    fn non_positive_or_invalid_ratio_never_samples() {
81        for ratio in [0.0, -1.0, f64::NAN, f64::NEG_INFINITY] {
82            let gate = SamplingGate::new(ratio);
83
84            assert_eq!(gate.sampled_per_million(), 0);
85            assert!((0..8).all(|_| !gate.should_sample()));
86        }
87    }
88
89    #[test]
90    fn fractional_ratio_is_evenly_spread() {
91        let one_quarter = SamplingGate::new(0.25);
92        let three_quarters = SamplingGate::new(0.75);
93
94        assert_eq!((0..8).filter(|_| one_quarter.should_sample()).count(), 2);
95        assert_eq!((0..8).filter(|_| three_quarters.should_sample()).count(), 6);
96    }
97}
rocketmq_observability/sampling.rs

rocketmq_observability/
sampling.rs
