relay_server/utils/
pick.rs1#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
3pub enum PickResult {
4 #[default]
6 Keep,
7 Discard,
9}
10
11pub 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 pub fn is_keep(self) -> bool {
25 matches!(self, PickResult::Keep)
26 }
27
28 #[cfg(feature = "processing")]
30 pub fn is_discard(self) -> bool {
31 !self.is_keep()
32 }
33}
34
35pub 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); 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}