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

relay_server/utils/
pick.rs
