net-pool 0.1.3

use crate::backend::{Address, BackendState};
use crate::error::Error;
use crate::strategy::LbStrategy;
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 id(&self) -> &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 LbStrategy>;

    /// 添加一个后端地址
    fn add_backend(&self, addr: Address) {
        self.get_strategy().add_backend(addr)
    }

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

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

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

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

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

impl BaseState {
    pub fn new(strategy: Arc<dyn LbStrategy>) -> Self {
        BaseState {
            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,
            }
        }
    }
}