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
33/// Returns [`PickResult::Keep`] if the current item should be sampled.
34///
35/// The passed `rate` is expected to be `0 <= rate <= 1`.
36pub fn sample(rate: f32) -> PickResult {
37    match (rate >= 1.0) || (rate > 0.0 && rand::random::<f32>() < rate) {
38        true => PickResult::Keep,
39        false => PickResult::Discard,
40    }
41}
42
43#[cfg(test)]
44mod test {
45    use super::*;
46
47    #[test]
48    fn test_rollout() {
49        assert_eq!(is_rolled_out(1, 0.0), PickResult::Discard);
50        assert_eq!(is_rolled_out(1, 0.0001), PickResult::Keep); // Id 1 should always be rolled out
51        assert_eq!(is_rolled_out(1, 0.1), PickResult::Keep);
52        assert_eq!(is_rolled_out(1, 0.9), PickResult::Keep);
53        assert_eq!(is_rolled_out(1, 1.0), PickResult::Keep);
54
55        assert_eq!(is_rolled_out(0, 0.0), PickResult::Discard);
56        assert_eq!(is_rolled_out(0, 1.0), PickResult::Keep);
57        assert_eq!(is_rolled_out(100000, 0.0), PickResult::Discard);
58        assert_eq!(is_rolled_out(100000, 1.0), PickResult::Keep);
59        assert_eq!(is_rolled_out(100001, 0.0), PickResult::Discard);
60        assert_eq!(is_rolled_out(100001, 1.0), PickResult::Keep);
61
62        assert_eq!(is_rolled_out(42, -100.0), PickResult::Discard);
63        assert_eq!(is_rolled_out(42, 100.0), PickResult::Keep);
64    }
65
66    #[cfg(feature = "processing")]
67    #[test]
68    fn test_sample() {
69        assert_eq!(sample(1.0), PickResult::Keep);
70        assert_eq!(sample(0.0), PickResult::Discard);
71
72        let mut r: i64 = 0;
73        for _ in 0..10000 {
74            if sample(0.5).is_keep() {
75                r += 1;
76            } else {
77                r -= 1;
78            }
79        }
80        assert!(r.abs() < 500);
81    }
82}
relay_server/utils/pick.rs

relay_server/utils/
pick.rs
