net-pool 0.5.1

A set of types for network connection pool.
Documentation
use crate::backend::{Address, BackendState};
use crate::error::Error;
use crate::strategy::Strategy;
use std::sync::Arc;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::{AtomicU64, AtomicUsize};
use std::time::Duration;

pub trait Pool: Send + Sync {
    /// id
    fn get_id(&self) -> &str;

    /// 设置id
    fn set_id(&mut self, id: &str);
    
    /// 设置最大连接数
    fn set_max_conn(&self, max: Option<usize>);

    /// 获取最大连接数
    fn get_max_conn(&self) -> Option<usize>;

    /// 获取当前的连接数
    fn get_cur_conn(&self) -> usize;

    /// 设置空闲连接保留时长
    fn set_keepalive(&self, _: Option<Duration>) {}

    /// 获取空闲连接保留时长
    fn get_keepalive(&self) -> Option<Duration> {
        None
    }

    /// 获取转发策略
    fn get_strategy(&self) -> Arc<dyn Strategy>;

    /// 设置策略
    fn set_strategy(&mut self, strategy: Arc<dyn Strategy>);
    
    /// 添加一个后端地址
    fn add_backend(&self, id: Option<u32>, addr: Address) {
        self.get_strategy().add_backend(id, addr)
    }

    /// 根据策略获取一个后端
    fn get_backend(&self, key: &str) -> Option<BackendState> {
        self.get_strategy().get_backend(key)
    }

    /// 移除一个后端地址
    fn remove_backend(&self, addr: &Address) -> bool {
        self.get_strategy().remove_backend(addr)
    }

    /// 获取所有后端地址切片
    fn get_backends(&self) -> Vec<BackendState> {
        self.get_strategy().get_backends()
    }

    fn get_backend_by_id(&self, id: u32) -> Option<BackendState> {
        self.get_strategy().get_backend_by_id(id)
    }

    fn get_backend_by_code(&self, code: u64) -> Option<BackendState> {
        self.get_strategy().get_backend_by_code(code)
    }

    // 是否使用tls
    fn use_tls(&self, _: bool) {}

    fn tls(&self) -> bool {
        false
    }
}

pub struct BaseState {
    /// id标识
    pub id: String,
    /// usize::MAX表示无设置
    pub max_conn: AtomicUsize, 
    pub cur_conn: AtomicUsize,
    /// 精度只到秒级别, u64::MAX表示无设置
    pub keepalive: AtomicU64, 
    pub lb_strategy: Arc<dyn Strategy>,
}

impl BaseState {
    pub fn new(strategy: Arc<dyn Strategy>) -> Self {
        BaseState {
            id: String::new(),
            max_conn: AtomicUsize::new(usize::MAX),
            cur_conn: AtomicUsize::new(0),
            keepalive: AtomicU64::new(u64::MAX),
            lb_strategy: strategy,
        }
    }
}

pub fn increase_current(max: &AtomicUsize, cur: &AtomicUsize) -> Result<(), Error> {
    let m = max.load(Relaxed);
    if m == usize::MAX {
        cur.fetch_add(1, Relaxed);
        Ok(())
    } else {
        let mut c = cur.load(Relaxed);
        loop {
            if c == m {
                break Err(Error::PoolFull);
            }

            match cur.compare_exchange(c, c + 1, Relaxed, Relaxed) {
                Ok(_) => break Ok(()),
                Err(s) => c = s,
            }
        }
    }
}