net_pool/strategy/
ch_strategy.rs

1use crate::backend::{Address, BackendState};
2use crate::utils::bytes_to_hash_code;
3use std::collections::{BTreeMap, HashSet};
4use std::sync::RwLock;
5
6/// 一致性hash负载策略
7
8pub struct CHStrategy {
9    /// usize:虚拟的后端个数
10    /// BTreeMap<u64, BackendState>: 排序后用来表示环形
11    /// HashSet<BackendState>: 真实的节点
12    inner: RwLock<(usize, BTreeMap<u64, BackendState>, HashSet<BackendState>)>,
13}
14
15impl CHStrategy {
16    pub fn new(virtual_backend_count: usize) -> Self {
17        if virtual_backend_count == 0 {
18            return CHStrategy::default();
19        }
20        CHStrategy {
21            inner: RwLock::new((virtual_backend_count, BTreeMap::new(), HashSet::new())),
22        }
23    }
24
25    fn sub_address(id: usize, addr: &Address) -> String {
26        let mut flag = match addr {
27            Address::Ori(ori) => ori.clone(),
28            Address::Addr(addr) => crate::utils::socketaddr_to_ip_string(addr),
29        };
30
31        flag.push('#');
32        flag.push_str(&id.to_string());
33        flag
34    }
35}
36
37impl Default for CHStrategy {
38    fn default() -> Self {
39        CHStrategy {
40            inner: RwLock::new((64, BTreeMap::new(), HashSet::new())),
41        }
42    }
43}
44
45impl super::LbStrategy for CHStrategy {
46    fn strategy(&self) -> super::Strategy {
47        super::Strategy::ConsistentHash
48    }
49
50    fn contain(&self, addr: &Address) -> bool {
51        let hash_code = addr.hash_code();
52        let g = self.inner.read().unwrap();
53        g.2.iter().any(|inner| inner.hash_code() == hash_code)
54    }
55
56    fn add_backend(&self, id: Option<u32>, addr: Address) {
57        let mut g = self.inner.write().unwrap();
58        let bs = BackendState::new(id, addr.clone());
59
60        for id in 0..g.0 {
61            let address = Self::sub_address(id, &addr);
62            let code = bytes_to_hash_code(address.as_bytes());
63            g.1.insert(code, bs.clone());
64        }
65
66        g.2.insert(bs);
67    }
68
69    fn remove_backend(&self, addr: &Address) -> bool {
70        let mut g = self.inner.write().unwrap();
71
72        for id in 0..g.0 {
73            let address = Self::sub_address(id, addr);
74            let code = bytes_to_hash_code(address.as_bytes());
75            if let None = g.1.remove(&code) {
76                return false;
77            }
78        }
79
80        g.2.remove(&BackendState::new(None, addr.clone()))
81    }
82
83    fn get_backend(&self, key: &str) -> Option<BackendState> {
84        let g = self.inner.read().unwrap();
85        if g.1.is_empty() {
86            return None;
87        }
88
89        let code = bytes_to_hash_code(key.as_bytes());
90
91        let o =
92            g.1.iter()
93                .find(|(key, _)| if code <= **key { true } else { false });
94
95        if o.is_none() {
96            g.1.iter().min().map(|(_, bs)| bs.clone())
97        } else {
98            o.map(|(_, bs)| bs.clone())
99        }
100    }
101
102    fn get_backends(&self) -> Vec<BackendState> {
103        let g = self.inner.read().unwrap();
104        let bss: Vec<BackendState> = g.2.iter().map(|kv| kv.clone()).collect();
105        bss
106    }
107}
net_pool/strategy/ch_strategy.rs

net_pool/strategy/
ch_strategy.rs
