net-pool 0.3.0

A set of types for network connection pool.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use crate::backend::{Address, BackendState};
use crate::utils::bytes_to_hash_code;
use std::collections::{BTreeMap, HashSet};
use std::sync::RwLock;

/// 一致性hash负载策略

pub struct CHStrategy {
    /// usize:虚拟的后端个数
    /// BTreeMap<u64, BackendState>: 排序后用来表示环形
    /// HashSet<BackendState>: 真实的节点
    inner: RwLock<(usize, BTreeMap<u64, BackendState>, HashSet<BackendState>)>,
}

impl CHStrategy {
    pub fn new(virtual_backend_count: usize) -> Self {
        if virtual_backend_count == 0 {
            return CHStrategy::default();
        }
        CHStrategy {
            inner: RwLock::new((virtual_backend_count, BTreeMap::new(), HashSet::new())),
        }
    }

    fn sub_address(id: usize, addr: &Address) -> String {
        let mut flag = match addr {
            Address::Ori(ori) => ori.clone(),
            Address::Addr(addr) => crate::utils::socketaddr_to_ip_string(addr),
        };

        flag.push('#');
        flag.push_str(&id.to_string());
        flag
    }
}

impl Default for CHStrategy {
    fn default() -> Self {
        CHStrategy {
            inner: RwLock::new((64, BTreeMap::new(), HashSet::new())),
        }
    }
}

impl super::LbStrategy for CHStrategy {
    fn strategy(&self) -> super::Strategy {
        super::Strategy::ConsistentHash
    }

    fn contain(&self, addr: &Address) -> bool {
        let hash_code = addr.hash_code();
        let g = self.inner.read().unwrap();
        g.2.iter().any(|inner| inner.hash_code() == hash_code)
    }

    fn add_backend(&self, id: Option<u32>, addr: Address) {
        let mut g = self.inner.write().unwrap();
        let bs = BackendState::new(id, addr.clone());

        for id in 0..g.0 {
            let address = Self::sub_address(id, &addr);
            let code = bytes_to_hash_code(address.as_bytes());
            g.1.insert(code, bs.clone());
        }

        g.2.insert(bs);
    }

    fn remove_backend(&self, addr: &Address) -> bool {
        let mut g = self.inner.write().unwrap();

        for id in 0..g.0 {
            let address = Self::sub_address(id, addr);
            let code = bytes_to_hash_code(address.as_bytes());
            if let None = g.1.remove(&code) {
                return false;
            }
        }

        g.2.remove(&BackendState::new(None, addr.clone()))
    }

    fn get_backend(&self, key: &str) -> Option<BackendState> {
        let g = self.inner.read().unwrap();
        if g.1.is_empty() {
            return None;
        }

        let code = bytes_to_hash_code(key.as_bytes());

        let o =
            g.1.iter()
                .find(|(key, _)| if code <= **key { true } else { false });

        if o.is_none() {
            g.1.iter().min().map(|(_, bs)| bs.clone())
        } else {
            o.map(|(_, bs)| bs.clone())
        }
    }

    fn get_backends(&self) -> Vec<BackendState> {
        let g = self.inner.read().unwrap();
        let bss: Vec<BackendState> = g.2.iter().map(|kv| kv.clone()).collect();
        bss
    }
}