http-pool 0.7.4

use crate::http1::{SendRequest, Sender};
use hyper::client::conn::http1;
use net_pool::backend::{Address, BackendState};
use net_pool::{Error, Strategy, debug, instrument_current_span, tokio_spawn};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::{Arc, Mutex};

#[derive(Clone)]
struct SrData(Arc<SendRequest>);

impl PartialEq for SrData {
    fn eq(&self, other: &Self) -> bool {
        self.0.as_ref() as *const _ == other.0.as_ref() as *const _
    }
}

impl Eq for SrData {}

impl Hash for SrData {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write_usize(self.0.as_ref() as *const _ as usize)
    }
}

struct Inner {
    /// 空闲的sr
    free: HashSet<SrData>,
    /// 在用的sr
    work: HashSet<SrData>,
}

impl Inner {
    fn new() -> Self {
        Inner {
            free: HashSet::new(),
            work: HashSet::new(),
        }
    }

    fn add_work_sr(&mut self, sr: Arc<SendRequest>) {
        self.work.insert(SrData(sr));
    }

    fn remove_sr(&mut self, sr: Arc<SendRequest>) -> bool {
        let d = SrData(sr);
        // 不应该同时处于work和free中
        if self.work.remove(&d) {
            true
        } else if self.free.remove(&d) {
            true
        } else {
            false
        }
    }

    /// 返回值是被删除的元素个数, 及查找到的sr
    fn get_sr(&mut self) -> (usize, Option<Arc<SendRequest>>) {
        // 删除个数
        let mut dels = 0;

        // 空了的话从work里取, 这里效率不高
        if self.free.is_empty() {
            // 挪动
            self.work.retain(|d| {
                // 引用计数为2时,代表外界不存在引用, 队列里一份, run那里一份
                if Arc::strong_count(&d.0) == 2 {
                    if !d.0.is_closed() {
                        self.free.insert(d.clone());
                    } else {
                        dels += 1;
                    }
                    false
                } else {
                    true
                }
            });
        }

        // 这里虽然删除了一些closed状态的, 但算法却有延迟性,它无法及时删除所有的closed
        loop {
            let d = match self.free.iter().next() {
                None => return (dels, None),
                Some(d) => d.clone(),
            };

            self.free.remove(&d);
            if d.0.is_closed() {
                dels += 1;
            } else {
                let sr = d.0.clone();
                self.work.insert(d);
                return (dels, Some(sr));
            }
        }
    }
}

/// http1连接池
/// 连接池中的连接受max_conn数据限制
/// 默认的keepalive是一分钟
/// 有三种行为会导致连接池中的连接被释放
///     1.对方主动关闭连接
///     2.remove_backend被调用清除后端地址, 这会导致该地址的所有连接失效
///     3.在外界没有引用sender达到keepalive时间后, 注意如果没有设置keepalive则此条无效
pub struct Pool {
    state: net_pool::pool::BaseState,
    use_tls: AtomicBool,
    free_conn_map: Mutex<HashMap<u64, Inner>>,
}

impl Pool {
    pub fn new(strategy: Arc<dyn Strategy>) -> Self {
        let p = Pool {
            state: net_pool::pool::BaseState::new(strategy),
            use_tls: AtomicBool::new(false),
            free_conn_map: Mutex::new(HashMap::new()),
        };
        <Pool as net_pool::pool::Pool>::set_keepalive(
            &p,
            Some(std::time::Duration::from_secs(60 * 1)),
        );
        p
    }
}

impl Default for Pool {
    fn default() -> Self {
        Pool::new(Arc::new(net_pool::strategy::CHStrategy::default()))
    }
}

impl<L: Strategy + 'static> From<L> for Pool {
    fn from(value: L) -> Self {
        Self::new(Arc::new(value))
    }
}

impl net_pool::pool::Pool for Pool {
    net_pool::macros::base_pool_impl! {state}

    fn remove_backend(&self, addr: &Address) -> bool {
        if self.state.lb_strategy.remove_backend(addr) {
            // 清除缓存
            self.clear_bs_sr(addr.hash_code());
            true
        } else {
            false
        }
    }

    fn use_tls(&self, tls: bool) {
        self.use_tls.store(tls, Relaxed);
    }

    fn tls(&self) -> bool {
        self.use_tls.load(Relaxed)
    }
}

impl Pool {
    fn clear_bs_sr(&self, hash_code: u64) {
        let mut guard = self.free_conn_map.lock().unwrap();
        if let Some(inner) = guard.remove(&hash_code) {
            assert!(
                self.state
                    .cur_conn
                    .fetch_sub(inner.free.len() + inner.work.len(), Relaxed)
                    > 0
            );
            debug!(
                "[desc] current connection count: {}",
                self.state.cur_conn.load(Relaxed)
            );
        }
    }

    fn get_sender(&self, bs: &BackendState) -> Option<Sender> {
        let mut guard = self.free_conn_map.lock().unwrap();
        let inner = guard.get_mut(&bs.hash_code())?;

        let sr = {
            let (del_cnt, sr) = inner.get_sr();
            if del_cnt > 0 {
                assert!(self.state.cur_conn.fetch_sub(del_cnt, Relaxed) > 0);
                debug!(
                    "[desc] current connection count: {}",
                    self.state.cur_conn.load(Relaxed)
                );
            }
            sr
        };

        if let Some(sr) = sr {
            let tls = <Pool as net_pool::pool::Pool>::tls(self);
            Some(Sender::new(
                sr,
                crate::utils::base_url(tls, bs.get_address()),
            ))
        } else {
            None
        }
    }

    fn add_sr(&self, hash_code: u64, sr: Arc<SendRequest>) {
        let mut guard = self.free_conn_map.lock().unwrap();
        match guard.entry(hash_code) {
            Entry::Occupied(mut o) => {
                o.get_mut().add_work_sr(sr);
            }
            Entry::Vacant(e) => {
                let mut inner = Inner::new();
                inner.add_work_sr(sr);
                e.insert(inner);
            }
        }
    }

    fn remove_sr(&self, hash_code: u64, sr: Arc<SendRequest>) {
        let mut guard = self.free_conn_map.lock().unwrap();
        if let Some(inner) = guard.get_mut(&hash_code) {
            if inner.remove_sr(sr) {
                assert!(self.state.cur_conn.fetch_sub(1, Relaxed) > 0);
                debug!(
                    "[desc] current connection count: {}",
                    self.state.cur_conn.load(Relaxed)
                );
            }
        }
    }

    fn run_conn<C: Future<Output = Result<(), hyper::Error>> + Send + 'static>(
        pool: Arc<Self>,
        c: C,
        sr: Arc<SendRequest>,
        bs: &BackendState,
    ) {
        let hash_code = bs.hash_code();

        // 加入work,且必须在下面执行super::run_conn之前加入
        // 添加sr进去缓存时, 可能此时bs已经被删除了
        // 由于缓存中持有sr, 所以有可能会导致在connection没有设置超时的情况下一直在运行,从而一直占用最大连接数
        // 由于缓存和bs是通过不同的锁持有的, 所以暂时无法对它们进行一致性限制
        pool.add_sr(hash_code, sr.clone());
        let ka = <Pool as net_pool::pool::Pool>::get_keepalive(&pool);

        tokio_spawn! {
            instrument_current_span! {
                async move {
                    let _r = crate::utils::run_conn(sr.clone(), c, ka).await;
                    debug!("connection closed: {:?}", _r);
                    pool.remove_sr(hash_code, sr);
                }
            }
        };
    }

    async fn create_tls_sender(self: Arc<Self>, bs: &BackendState) -> Result<Sender, Error> {
        let tls = <Pool as net_pool::pool::Pool>::tls(&self);
        let addr = bs.get_address();

        // 获取连接
        let tcp = crate::utils::create_https_stream(addr).await?;
        let tls_tcp =
            crate::utils::create_tls_tcp(tcp, addr, crate::utils::HTTP1_TLS_CLIENT_CFG.clone())
                .await?;

        // 握手建立 sender 和 connection
        let io = hyper_util::rt::TokioIo::new(tls_tcp);
        let pair = http1::handshake(io)
            .await
            .map_err(|e| Error::from_other(e))?;

        // 启动连接驱动
        let sr = Arc::new(pair.0);
        Pool::run_conn(self, pair.1, sr.clone(), &bs);

        // 返回一份
        Ok(Sender::new(sr, crate::utils::base_url(tls, addr)))
    }

    async fn create_non_tls_sender(self: Arc<Self>, bs: &BackendState) -> Result<Sender, Error> {
        let addr = bs.get_address();
        let tls = <Pool as net_pool::pool::Pool>::tls(&self);

        // 获取连接
        let tcp = crate::utils::create_http_stream(addr).await?;

        // 握手建立 sender 和 connection
        let io = hyper_util::rt::TokioIo::new(tcp);
        let pair = http1::handshake(io)
            .await
            .map_err(|e| Error::from_other(e))?;

        // 启动连接驱动
        let sr = Arc::new(pair.0);
        Pool::run_conn(self, pair.1, sr.clone(), &bs);

        // 返回一份
        Ok(Sender::new(sr, crate::utils::base_url(tls, addr)))
    }

    async fn create_sender(self: Arc<Self>, bs: &BackendState) -> Result<Sender, Error> {
        if <Pool as net_pool::pool::Pool>::tls(&self) {
            self.create_tls_sender(bs).await
        } else {
            self.create_non_tls_sender(bs).await
        }
    }

    async fn real(self: Arc<Self>, bs: BackendState) -> Result<Sender, Error> {
        if let Some(sender) = self.clone().get_sender(&bs) {
            return Ok(sender);
        }

        let sender = {
            // 预分配数量
            net_pool::pool::increase_current(&self.state.max_conn, &self.state.cur_conn)?;
            self.clone().create_sender(&bs).await.map(|s| {
                debug!(
                    "[incr] current connection count: {}",
                    self.state.cur_conn.load(Relaxed)
                );
                s
            })
        };

        if sender.is_err() {
            assert!(self.state.cur_conn.fetch_sub(1, Relaxed) > 0);
        }

        sender
    }
}

pub trait HttpPool {
    /// 根据key做某种策略获取
    fn get(self: Arc<Self>, key: &str) -> impl Future<Output = Result<Sender, Error>> + Send;

    /// 根据地址获取
    fn target(self: Arc<Self>, _: &Address) -> impl Future<Output = Result<Sender, Error>> + Send;
}

impl HttpPool for Pool {
    async fn get(self: Arc<Self>, key: &str) -> Result<Sender, Error> {
        let bs = self
            .state
            .lb_strategy
            .get_backend(key)
            .ok_or(Error::NoBackend)?;

        self.real(bs).await
    }

    async fn target(self: Arc<Self>, addr: &Address) -> Result<Sender, Error> {
        if !self.state.lb_strategy.contain(addr) {
            return Err(Error::NoBackend);
        }

        self.real(BackendState::new(None, addr.clone())).await
    }
}

pub async fn get<P: HttpPool>(pool: Arc<P>, key: &str) -> Result<Sender, Error> {
    HttpPool::get(pool, key).await
}