relay_server/utils/
pick.rs

1/// Result of a sampling operation, whether to keep or discard something.
2#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
3pub enum PickResult {
4    /// The item should be kept.
5    #[default]
6    Keep,
7    /// The item should be dropped.
8    Discard,
9}
10
11/// Returns [`PickResult::Keep`] if `id` is rolled out with a rollout rate `rate`.
12///
13/// Deterministically makes a rollout decision for an id, usually organization id,
14/// and rate.
15#[allow(
16    dead_code,
17    reason = "A utility which is frequently used and removed again"
18)]
19pub fn is_rolled_out(id: u64, rate: f32) -> PickResult {
20    match ((id % 100000) as f32 / 100000.0f32) < rate {
21        true => PickResult::Keep,
22        false => PickResult::Discard,
23    }
24}
25
26impl PickResult {
27    /// Returns `true` if the sampling result is [`PickResult::Keep`].
28    pub fn is_keep(self) -> bool {
29        matches!(self, PickResult::Keep)
30    }
31
32    /// Returns `true` if the sampling result is [`PickResult::Discard`].
33    #[cfg(feature = "processing")]
34    pub fn is_discard(self) -> bool {
35        !self.is_keep()
36    }
37}
38
39/// Returns [`PickResult::Keep`] if the current item should be sampled.
40///
41/// The passed `rate` is expected to be `0 <= rate <= 1`.
42pub fn sample(rate: f32) -> PickResult {
43    match (rate >= 1.0) || (rate > 0.0 && rand::random::<f32>() < rate) {
44        true => PickResult::Keep,
45        false => PickResult::Discard,
46    }
47}
48
49#[cfg(test)]
50mod test {
51    use super::*;
52
53    #[test]
54    fn test_rollout() {
55        assert_eq!(is_rolled_out(1, 0.0), PickResult::Discard);
56        assert_eq!(is_rolled_out(1, 0.0001), PickResult::Keep); // Id 1 should always be rolled out
57        assert_eq!(is_rolled_out(1, 0.1), PickResult::Keep);
58        assert_eq!(is_rolled_out(1, 0.9), PickResult::Keep);
59        assert_eq!(is_rolled_out(1, 1.0), PickResult::Keep);
60
61        assert_eq!(is_rolled_out(0, 0.0), PickResult::Discard);
62        assert_eq!(is_rolled_out(0, 1.0), PickResult::Keep);
63        assert_eq!(is_rolled_out(100000, 0.0), PickResult::Discard);
64        assert_eq!(is_rolled_out(100000, 1.0), PickResult::Keep);
65        assert_eq!(is_rolled_out(100001, 0.0), PickResult::Discard);
66        assert_eq!(is_rolled_out(100001, 1.0), PickResult::Keep);
67
68        assert_eq!(is_rolled_out(42, -100.0), PickResult::Discard);
69        assert_eq!(is_rolled_out(42, 100.0), PickResult::Keep);
70    }
71
72    #[test]
73    fn test_sample() {
74        assert_eq!(sample(1.0), PickResult::Keep);
75        assert_eq!(sample(0.0), PickResult::Discard);
76
77        let mut r: i64 = 0;
78        for _ in 0..10000 {
79            if sample(0.5).is_keep() {
80                r += 1;
81            } else {
82                r -= 1;
83            }
84        }
85        assert!(r.abs() < 500);
86    }
87}
relay_server/utils/pick.rs

relay_server/utils/
pick.rs
