Skip to main content

cdrs/load_balancing/
random.rs

1use rand;
2
3use super::LoadBalancingStrategy;
4
5pub struct Random<N> {
6    pub cluster: Vec<N>,
7}
8
9impl<N> Random<N> {
10    pub fn new(cluster: Vec<N>) -> Self {
11        Random { cluster }
12    }
13
14    /// Returns a random number from a range
15    fn rnd_idx(bounds: (usize, usize)) -> usize {
16        let min = bounds.0;
17        let max = bounds.1;
18        let rnd = rand::random::<usize>();
19        rnd % (max - min) + min
20    }
21}
22
23impl<N> From<Vec<N>> for Random<N> {
24    fn from(cluster: Vec<N>) -> Random<N> {
25        Random { cluster }
26    }
27}
28
29impl<N> LoadBalancingStrategy<N> for Random<N> {
30    fn init(&mut self, cluster: Vec<N>) {
31        self.cluster = cluster;
32    }
33
34    /// Returns next random node from a cluster
35    fn next(&self) -> Option<&N> {
36        let len = self.cluster.len();
37        if len == 0 {
38            return None;
39        }
40        self.cluster.get(Self::rnd_idx((0, len)))
41    }
42
43    fn remove_node<F>(&mut self, filter: F)
44    where
45        F: FnMut(&N) -> bool,
46    {
47        if let Some(i) = self.cluster.iter().position(filter) {
48            self.cluster.remove(i);
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn next_random() {
59        let nodes = vec!["a", "b", "c", "d", "e", "f", "g"];
60        let load_balancer = Random::from(nodes);
61        for _ in 0..100 {
62            let s = load_balancer.next();
63            assert!(s.is_some());
64        }
65    }
66
67    #[test]
68    fn remove_from_random() {
69        let nodes = vec!["a"];
70        let mut load_balancer = Random::from(nodes);
71
72        let s = load_balancer.next();
73        assert!(s.is_some());
74
75        load_balancer.remove_node(|n| n == &"a");
76        let s = load_balancer.next();
77        assert!(s.is_none());
78    }
79}