cdrs/load_balancing/
round_robin.rs

1use std::cell::RefCell;
2
3use super::LoadBalancingStrategy;
4
5pub struct RoundRobin<N> {
6    cluster: Vec<N>,
7    prev_idx: RefCell<usize>,
8}
9
10impl<N> RoundRobin<N> {
11    pub fn new() -> Self {
12        RoundRobin {
13            prev_idx: RefCell::new(0),
14            cluster: vec![],
15        }
16    }
17}
18
19impl<N> From<Vec<N>> for RoundRobin<N> {
20    fn from(cluster: Vec<N>) -> RoundRobin<N> {
21        RoundRobin {
22            prev_idx: RefCell::new(0),
23            cluster: cluster,
24        }
25    }
26}
27
28impl<N> LoadBalancingStrategy<N> for RoundRobin<N> {
29    fn init(&mut self, cluster: Vec<N>) {
30        self.cluster = cluster;
31    }
32
33    /// Returns next node from a cluster
34    fn next(&self) -> Option<&N> {
35        let prev_idx = *self.prev_idx.borrow();
36        let next_idx = (prev_idx + 1) % self.cluster.len();
37        self.prev_idx.replace(next_idx);
38        self.cluster.get(next_idx)
39    }
40
41    fn remove_node<F>(&mut self, filter: F)
42    where
43        F: FnMut(&N) -> bool,
44    {
45        if let Some(i) = self.cluster.iter().position(filter) {
46            self.cluster.remove(i);
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn round_robin() {
57        let nodes = vec!["a", "b", "c"];
58        let nodes_c = nodes.clone();
59        let load_balancer = RoundRobin::from(nodes);
60        for i in 0..10 {
61            assert_eq!(&nodes_c[(i + 1) % 3], load_balancer.next().unwrap());
62        }
63    }
64
65    #[test]
66    fn remove_from_round_robin() {
67        let nodes = vec!["a", "b"];
68        let mut load_balancer = RoundRobin::from(nodes);
69        assert_eq!(&"b", load_balancer.next().unwrap());
70
71        load_balancer.remove_node(|n| n == &"a");
72        assert_eq!(&"b", load_balancer.next().unwrap());
73    }
74}
cdrs/load_balancing/round_robin.rs

cdrs/load_balancing/
round_robin.rs
