tcp-pool 0.4.1

use crate::TcpStream;
use net_pool::backend::Address;
use net_pool::{Error, Strategy, debug};
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;

/// tcp连接池
/// 连接池中的连接受max_conn数据限制
/// keepalive数据没有任何效果
/// 连接池中的连接不复用, get出来后tcp stream不被引用则会导致连接断开
pub struct Pool {
    state: net_pool::pool::BaseState,
}

impl Pool {
    pub fn new(strategy: Arc<dyn Strategy>) -> Self {
        Pool {
            state: net_pool::pool::BaseState::new(strategy),
        }
    }
}

impl Default for Pool {
    fn default() -> Self {
        Pool::new(Arc::new(net_pool::strategy::HashStrategy::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}
}

pub trait TcpPool {
    fn get(self: Arc<Self>, key: &str) -> impl Future<Output = Result<TcpStream, Error>> + Send;
}

impl TcpPool for Pool {
    async fn get(self: Arc<Self>, key: &str) -> Result<TcpStream, Error> {
        // 预分配数量
        net_pool::pool::increase_current(&self.state.max_conn, &self.state.cur_conn)?;

        let tcp = {
            match self
                .state
                .lb_strategy
                .get_backend(key)
                .ok_or(Error::NoBackend)
            {
                Err(e) => Err(e),
                Ok(bs) => create_tcp_stream(bs.get_address()).await,
            }
        };

        if tcp.is_err() {
            assert!(
                self.state
                    .cur_conn
                    .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                    > 0
            );
        } else {
            debug!(
                "[tcp pool] [incr] current connection count: {}",
                self.state
                    .cur_conn
                    .load(std::sync::atomic::Ordering::Relaxed)
            );
        }

        let pool = self.clone();
        tcp.map(|t| {
            TcpStream::new(
                move || {
                    assert!(
                        pool.state
                            .cur_conn
                            .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                            > 0
                    );
                    debug!(
                        "[tcp pool] [desc] current connection count: {}",
                        pool.state
                            .cur_conn
                            .load(std::sync::atomic::Ordering::Relaxed)
                    );
                },
                t,
            )
        })
    }
}

async fn create_tcp_stream(addrs: &Address) -> Result<tokio::net::TcpStream, Error> {
    let a = to_socket_addrs(addrs).await?;
    tokio::net::TcpStream::connect(&a[..])
        .await
        .map_err(|e| Error::from_other(e))
}

async fn to_socket_addrs(addr: &Address) -> io::Result<Vec<SocketAddr>> {
    match addr {
        Address::Ori(ori) => tokio::net::lookup_host(ori)
            .await
            .map(|a| a.into_iter().collect()),
        Address::Addr(addr) => Ok(vec![addr.clone()]),
    }
}

pub async fn get<P: TcpPool + Send>(pool: Arc<P>, key: &str) -> Result<TcpStream, Error> {
    TcpPool::get(pool.clone(), key).await
}