druid_io/
connection.rs

1use std::sync::atomic::AtomicUsize;
2use std::sync::atomic::Ordering;
3use std::sync::Arc;
4
5pub trait BrokersPool {
6    fn broker(&self) -> &String;
7}
8
9pub struct StaticPool {
10    brokers: Vec<String>,
11    selection: SelectionStategy,
12}
13
14impl StaticPool {
15    pub fn new(brokers: Vec<String>, stategy: SelectionStategy) -> Self {
16        StaticPool {
17            brokers: brokers,
18            selection: stategy,
19        }
20    }
21}
22
23impl BrokersPool for StaticPool {
24    fn broker(&self) -> &String {
25        self.selection.select(&self.brokers)
26    }
27}
28
29pub enum SelectionStategy {
30    Constant,
31    // Random,
32    RoundRobin(Arc<AtomicUsize>),
33}
34
35fn get_and_increment(index: &AtomicUsize, max_val: usize) -> usize {
36    loop {
37        let val = index.load(Ordering::SeqCst);
38        let new_val = if val >= max_val { 1 } else { val + 1 };
39        if index.compare_and_swap(val, new_val, Ordering::SeqCst) == val {
40            return if val >= max_val { 0 as usize } else { val };
41        }
42    }
43}
44
45impl SelectionStategy {
46    pub fn select<'a, T>(&self, list: &'a Vec<T>) -> &'a T {
47        if let SelectionStategy::RoundRobin(ref index) = *self {
48            let idx = get_and_increment(index, list.len());
49            list.get(idx).unwrap()
50        } else {
51            list.get(0).unwrap()
52        }
53    }
54
55    pub fn round_robin() -> SelectionStategy {
56        SelectionStategy::RoundRobin(Arc::new(AtomicUsize::new(0)))
57    }
58    pub fn constant() -> SelectionStategy {
59        SelectionStategy::Constant
60    }
61
62    pub fn default_for<T>(list: &Vec<T>) -> SelectionStategy {
63        if list.len() == 1 {
64            SelectionStategy::constant()
65        } else {
66            SelectionStategy::round_robin()
67        }
68    }
69}
70
71#[cfg(test)]
72mod test {
73    use super::*;
74
75    #[test]
76    fn test_round_robin_selection() {
77        let nodes = vec![
78            "localhost:8081".to_string(),
79            "localhost:80".to_string(),
80            "example.com:8080".to_string(),
81        ];
82
83        let strategy = SelectionStategy::default_for(&nodes);
84        assert_eq!(strategy.select(&nodes), "localhost:8081");
85        assert_eq!(strategy.select(&nodes), "localhost:80");
86        assert_eq!(strategy.select(&nodes), "example.com:8080");
87        assert_eq!(strategy.select(&nodes), "localhost:8081");
88        assert_eq!(strategy.select(&nodes), "localhost:80");
89        assert_eq!(strategy.select(&nodes), "example.com:8080");
90    }
91    #[test]
92    fn test_constaint() {
93        let nodes = vec![
94            "localhost:8081".to_string(),
95            "localhost:80".to_string(),
96            "example.com:8080".to_string(),
97        ];
98
99        let strategy = SelectionStategy::constant();
100        assert_eq!(strategy.select(&nodes), "localhost:8081");
101        assert_eq!(strategy.select(&nodes), "localhost:8081");
102        assert_eq!(strategy.select(&nodes), "localhost:8081");
103    }
104}
druid_io/connection.rs

druid_io/
connection.rs
