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    fn valid_backend(&self, addr: &Address, valid: bool) -> bool {
19        let mut g = self.bss.write().unwrap();
20        let bss = g.get_backends_mut();
21
22        if bss.is_empty() {
23            return false;
24        }
25
26        let code = addr.hash_code();
27
28        for bs in bss.iter_mut() {
29            if bs.hash_code() == code {
30                bs.set_valid(valid);
31                return true;
32            }
33        }
34
35        return false;
36    }
37}
38
39impl Default for RRStrategy {
40    fn default() -> Self {
41        RRStrategy {
42            index: AtomicUsize::new(0),
43            bss: RwLock::new(BackendStates::new()),
44        }
45    }
46}
47
48impl super::LbStrategy for RRStrategy {
49    fn strategy(&self) -> super::Strategy {
50        super::Strategy::RoundRobin
51    }
52
53    fn contain(&self, hash_code: u64) -> bool {
54        let g = self.bss.read().unwrap();
55        g.contain(hash_code)
56    }
57
58    fn add_backend(&self, addr: Address) {
59        let mut g = self.bss.write().unwrap();
60        g.add_backend(addr);
61    }
62
63    fn remove_backend(&self, addr: &Address) -> bool {
64        let mut g = self.bss.write().unwrap();
65        g.remove_backend(&addr)
66    }
67
68    fn get_backend(&self, _: &str) -> Option<BackendState> {
69        let g = self.bss.read().unwrap();
70        let bss = g.get_backends();
71        if bss.is_empty() {
72            return None;
73        }
74
75        let index = self.index.fetch_add(1, Ordering::Relaxed);
76        let idx = index % bss.len();
77        Some(bss[idx].clone())
78    }
79
80    fn get_backends(&self) -> Vec<BackendState> {
81        let g = self.bss.read().unwrap();
82        let bss = g.get_backends();
83        bss.iter().map(|bs| bs.clone()).collect()
84    }
85
86    fn get_origin_backends(&self) -> Vec<BackendState> {
87        self.get_backends()
88    }
89
90    fn disable_backend(&self, addr: &Address) -> bool {
91        self.valid_backend(addr, false)
92    }
93
94    fn enable_backend(&self, addr: &Address) {
95        self.valid_backend(addr, true);
96    }
97}
net_pool/strategy/rr_strategy.rs

net_pool/strategy/
rr_strategy.rs
