net_pool/strategy/
rr_strategy.rs

1use crate::backend::{Address, BackendState, BackendStates};
2use std::sync::RwLock;
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5/// 轮询策略
6pub struct RRStrategy {
7    /// 记录轮询索引
8    index: AtomicUsize,
9
10    bss: RwLock<BackendStates>,
11}
12
13impl RRStrategy {
14    pub fn new() -> Self {
15        RRStrategy::default()
16    }
17}
18
19impl Default for RRStrategy {
20    fn default() -> Self {
21        RRStrategy {
22            index: AtomicUsize::new(0),
23            bss: RwLock::new(BackendStates::new()),
24        }
25    }
26}
27
28impl super::LbStrategy for RRStrategy {
29    fn strategy(&self) -> super::Strategy {
30        super::Strategy::RoundRobin
31    }
32
33    fn contain(&self, addr: &Address) -> bool {
34        let g = self.bss.read().unwrap();
35        g.get_backend(addr).is_some()
36    }
37
38    fn add_backend(&self, addr: Address) {
39        let mut g = self.bss.write().unwrap();
40        g.add_backend(addr);
41    }
42
43    fn remove_backend(&self, addr: &Address) -> bool {
44        let mut g = self.bss.write().unwrap();
45        g.remove_backend(&addr)
46    }
47
48    fn get_backend(&self, _: &str) -> Option<BackendState> {
49        let g = self.bss.read().unwrap();
50        let bss = g.get_backends();
51        if bss.is_empty() {
52            return None;
53        }
54
55        let index = self.index.fetch_add(1, Ordering::Relaxed);
56        let idx = index % bss.len();
57        Some(bss[idx].clone())
58    }
59
60    fn get_backends(&self) -> Vec<BackendState> {
61        let g = self.bss.read().unwrap();
62        let bss = g.get_backends();
63        bss.iter().map(|bs| bs.clone()).collect()
64    }
65}