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.
15pub fn is_rolled_out(id: u64, rate: f32) -> PickResult {
16    match ((id % 100000) as f32 / 100000.0f32) < rate {
17        true => PickResult::Keep,
18        false => PickResult::Discard,
19    }
20}
21
22impl PickResult {
23    /// Returns `true` if the sampling result is [`PickResult::Keep`].
24    pub fn is_keep(self) -> bool {
25        matches!(self, PickResult::Keep)
26    }
27
28    /// Returns `true` if the sampling result is [`PickResult::Discard`].
29    #[cfg(feature = "processing")]
30    pub fn is_discard(self) -> bool {
31        !self.is_keep()
32    }
33}
34
35/// Returns [`PickResult::Keep`] if the current item should be sampled.
36///
37/// The passed `rate` is expected to be `0 <= rate <= 1`.
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    #[test]
69    fn test_sample() {
70        assert_eq!(sample(1.0), PickResult::Keep);
71        assert_eq!(sample(0.0), PickResult::Discard);
72
73        let mut r: i64 = 0;
74        for _ in 0..10000 {
75            if sample(0.5).is_keep() {
76                r += 1;
77            } else {
78                r -= 1;
79            }
80        }
81        assert!(r.abs() < 500);
82    }
83}
relay_server/utils/pick.rs

relay_server/utils/
pick.rs
