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, id: Option<u32>, addr: Address) {
37        self.get_strategy().add_backend(id, addr)
38    }
39    
40    /// 根据策略获取一个后端
41    fn get_backend(&self, key: &str) -> Option<BackendState> {
42        self.get_strategy().get_backend(key)
43    }
44
45    /// 移除一个后端地址
46    fn remove_backend(&self, addr: &Address) -> bool {
47        self.get_strategy().remove_backend(addr)
48    }
49
50    /// 获取所有后端地址切片
51    fn get_backends(&self) -> Vec<BackendState> {
52        self.get_strategy().get_backends()
53    }
54
55    // 是否使用tls
56    fn use_tls(&self, _: bool) {}
57
58    fn tls(&self) -> bool {
59        false
60    }
61}
62
63pub struct BaseState {
64    pub max_conn: AtomicUsize, // usize::MAX表示无设置
65    pub cur_conn: AtomicUsize,
66    pub keepalive: AtomicU64, // 精度只到秒级别, u64::MAX表示无设置
67    pub lb_strategy: Arc<dyn LbStrategy>,
68}
69
70impl BaseState {
71    pub fn new(strategy: Arc<dyn LbStrategy>) -> Self {
72        BaseState {
73            max_conn: AtomicUsize::new(usize::MAX),
74            cur_conn: AtomicUsize::new(0),
75            keepalive: AtomicU64::new(u64::MAX),
76            lb_strategy: strategy,
77        }
78    }
79}
80
81pub fn increase_current(max: &AtomicUsize, cur: &AtomicUsize) -> Result<(), Error> {
82    let m = max.load(Relaxed);
83    if m == usize::MAX {
84        cur.fetch_add(1, Relaxed);
85        Ok(())
86    } else {
87        let mut c = cur.load(Relaxed);
88        loop {
89            if c == m {
90                break Err(Error::PoolFull);
91            }
92
93            match cur.compare_exchange(c, c + 1, Relaxed, Relaxed) {
94                Ok(_) => break Ok(()),
95                Err(s) => c = s,
96            }
97        }
98    }
99}
net_pool/pool.rs

net_pool/
pool.rs
