udp_pool/

pool.rs

use crate::id::ID;
use crate::{Sender, UdpTx, connection};
use net_pool::backend::{Address, BackendState};
use net_pool::error::Error;
use net_pool::strategy::LbStrategy;
use net_pool::{debug, instrument_current_span, tokio_spawn, trace};
use std::collections::HashMap;
use std::io;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use tokio::net::UdpSocket;

/// udp连接池
/// 连接池中的连接受max_conn数据限制
/// 默认的keepalive是5分钟
/// 有两种行为会导致连接池中的连接被释放
///     1. remove_backend被调用清除后端地址, 这会导致该地址的所有连接失效
///     2. 连接keepalive时间后没有数据发送
pub struct Pool {
    state: net_pool::pool::BaseState,
    free_conn_map: Mutex<HashMap<ID, UdpTx>>,
}

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

impl Default for Pool {
    fn default() -> Self {
        // robin round
        Pool::new(Arc::new(net_pool::strategy::RRStrategy::default()))
    }
}

impl<L: LbStrategy + '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_tx(addr);
            true
        } else {
            false
        }
    }
}

impl Pool {
    fn get_tx(&self, a: &SocketAddr) -> Option<UdpTx> {
        let id = a.into();
        let mut guard = self.free_conn_map.lock().unwrap();

        // 这里虽然删除了一些closed状态的, 但算法却有延迟性,它无法及时删除所有的closed
        let tx = guard.get(&id)?;
        if tx.is_closed() {
            guard.remove(&id);
            assert!(
                self.state
                    .cur_conn
                    .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                    > 0
            );
            trace!(
                "[udp pool] [desc] current socket count: {}",
                self.state
                    .cur_conn
                    .load(std::sync::atomic::Ordering::Relaxed)
            );
            None
        } else {
            Some(tx.clone())
        }
    }

    /// 清除所有的后端, 因为清除了连接缓存, 但可能因为外界仍持有sender所以会导致conn还无法及时退出.但cur_conn数值已发生了变化
    fn clear_bs_tx(&self, a: &Address) {
        let mut guard = self.free_conn_map.lock().unwrap();
        guard.retain(|k, _| {
            if k.get_bid() != a.hash_code() {
                true
            } else {
                assert!(
                    self.state
                        .cur_conn
                        .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                        > 0
                );
                false
            }
        });

        trace!(
            "[udp pool] [desc] current socket count: {}",
            self.state
                .cur_conn
                .load(std::sync::atomic::Ordering::Relaxed)
        );
    }

    fn add_tx(&self, id: ID, tx: UdpTx) {
        let mut guard = self.free_conn_map.lock().unwrap();
        guard.insert(id, tx);
    }

    fn remove_tx(&self, id: ID) {
        let mut guard = self.free_conn_map.lock().unwrap();
        if guard.remove(&id).is_some() {
            assert!(
                self.state
                    .cur_conn
                    .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                    > 0
            );
            trace!(
                "[udp pool] [desc] current socket count: {}",
                self.state
                    .cur_conn
                    .load(std::sync::atomic::Ordering::Relaxed)
            );
        }
    }

    async fn create_tx(
        self: Arc<Self>,
        a: SocketAddr,
        client: Option<Arc<UdpSocket>>,
    ) -> Result<UdpTx, Error> {
        let mut id = ID::new(&a);
        let bs = self
            .state
            .lb_strategy
            .get_backend(&id.to_string())
            .ok_or(Error::NoBackend)?;
        id.set_bid(bs.hash_code());

        let proxy = create_udp_socket(&bs).await?;
        let ka = <Pool as net_pool::pool::Pool>::get_keepalive(&self);

        let (tx, conn) = io(a, client, proxy, ka);

        // 驱动
        let pool = self.clone();
        tokio_spawn! {
            instrument_current_span! {
                async move {
                    let _res = conn.await;
                    debug!("[udp pool] udp socket close: {:?}", _res);
                    // conn可能会因为超时而退出, 但外界仍持有sender
                    pool.remove_tx(id);
                }
            }
        };

        self.add_tx(id, tx.clone());
        Ok(tx)
    }
}

pub trait UdpPool {
    fn get(
        self: Arc<Self>,
        a: SocketAddr,
        send: Option<Arc<UdpSocket>>,
    ) -> impl Future<Output = Result<Sender, Error>> + Send;
}

impl UdpPool for Pool {
    async fn get(
        self: Arc<Self>,
        a: SocketAddr,
        send: Option<Arc<UdpSocket>>,
    ) -> Result<Sender, Error> {
        let tx = match self.get_tx(&a) {
            Some(s) => Ok(s),
            None => {
                // 预分配数量
                net_pool::pool::increase_current(&self.state.max_conn, &self.state.cur_conn)?;
                self.clone().create_tx(a, send).await.map(|s| {
                    trace!(
                        "[udp pool] [incr] current socket count: {}",
                        self.state
                            .cur_conn
                            .load(std::sync::atomic::Ordering::Relaxed)
                    );
                    s
                })
            }
        };

        if tx.is_err() {
            assert!(
                self.state
                    .cur_conn
                    .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
                    > 0
            );
        }

        tx.map(|tx| Sender::new(tx))
    }
}

pub async fn get<P: UdpPool>(
    pool: Arc<P>,
    a: SocketAddr,
    send: Option<Arc<UdpSocket>>,
) -> Result<Sender, Error> {
    UdpPool::get(pool, a, send).await
}

fn io(
    client_addr: SocketAddr,
    client: Option<Arc<UdpSocket>>,
    proxy: UdpSocket,
    keepalive: Option<std::time::Duration>,
) -> (UdpTx, connection::Connection) {
    let (tx, rx) = tokio::sync::mpsc::channel(20);
    (
        tx,
        connection::Connection::new(client_addr, client, proxy, rx, keepalive),
    )
}

async fn create_udp_socket(bs: &BackendState) -> Result<UdpSocket, Error> {
    let recv = UdpSocket::bind("0.0.0.0:0").await?;

    let a = to_socket_addrs(bs.get_address()).await?;
    recv.connect(&a[..]).await?;

    Ok(recv)
}

async fn to_socket_addrs(addr: &Address) -> io::Result<Vec<SocketAddr>> {
    async fn inner<T: tokio::net::ToSocketAddrs>(host: &T) -> io::Result<Vec<SocketAddr>> {
        tokio::net::lookup_host(host)
            .await
            .map(|a| a.into_iter().collect())
    }

    match addr {
        Address::Ori(ori) => inner(ori).await,
        Address::Addr(addr) => inner(addr).await,
    }
}