net_pool/
pool.rs

1use crate::backend::{Address, BackendState};
2use crate::error::Error;
3use crate::strategy::LbStrategy;
4use std::sync::Arc;
5use std::sync::atomic::Ordering::Relaxed;
6use std::sync::atomic::{AtomicU64, AtomicUsize};
7use std::time::Duration;
8
9pub trait Pool: Send + Sync {
10    /// id
11    fn id(&self) -> &str {
12        ""
13    }
14
15    /// 设置最大连接数
16    fn set_max_conn(&self, max: Option<usize>);
17
18    /// 获取最大连接数
19    fn get_max_conn(&self) -> Option<usize>;
20
21    /// 获取当前的连接数
22    fn get_cur_conn(&self) -> usize;
23
24    /// 设置空闲连接保留时长
25    fn set_keepalive(&self, _: Option<Duration>) {}
26
27    /// 获取空闲连接保留时长
28    fn get_keepalive(&self) -> Option<Duration> {
29        None
30    }
31
32    /// 获取转发策略
33    fn get_strategy(&self) -> Arc<dyn LbStrategy>;
34
35    /// 添加一个后端地址
36    fn add_backend(&self, addr: Address) {
37        self.get_strategy().add_backend(addr)
38    }
39
40    /// 移除一个后端地址
41    fn remove_backend(&self, addr: &Address) -> bool {
42        self.get_strategy().remove_backend(addr)
43    }
44
45    /// 获取所有后端地址切片
46    fn get_backends(&self) -> Vec<BackendState> {
47        self.get_strategy().get_backends()
48    }
49
50    // 是否使用tls
51    fn use_tls(&self, _: bool) {}
52
53    fn tls(&self) -> bool {
54        false
55    }
56}
57
58pub struct BaseState {
59    pub max_conn: AtomicUsize, // usize::MAX表示无设置
60    pub cur_conn: AtomicUsize,
61    pub keepalive: AtomicU64, // 精度只到秒级别, u64::MAX表示无设置
62    pub lb_strategy: Arc<dyn LbStrategy>,
63}
64
65impl BaseState {
66    pub fn new(strategy: Arc<dyn LbStrategy>) -> Self {
67        BaseState {
68            max_conn: AtomicUsize::new(usize::MAX),
69            cur_conn: AtomicUsize::new(0),
70            keepalive: AtomicU64::new(u64::MAX),
71            lb_strategy: strategy,
72        }
73    }
74}
75
76pub fn increase_current(max: &AtomicUsize, cur: &AtomicUsize) -> Result<(), Error> {
77    let m = max.load(Relaxed);
78    if m == usize::MAX {
79        cur.fetch_add(1, Relaxed);
80        Ok(())
81    } else {
82        let mut c = cur.load(Relaxed);
83        loop {
84            if c == m {
85                break Err(Error::PoolFull);
86            }
87
88            match cur.compare_exchange(c, c + 1, Relaxed, Relaxed) {
89                Ok(_) => break Ok(()),
90                Err(s) => c = s,
91            }
92        }
93    }
94}
net_pool/pool.rs

net_pool/
pool.rs
