xrpc/

pool.rs

use parking_lot::Mutex;
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;

use crate::error::{Result, RpcError};

pub struct PoolConfig {
    pub min_connections: usize,
    pub max_connections: usize,
    pub acquire_timeout: Duration,
    pub idle_timeout: Duration,
    pub health_check_interval: Duration,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            min_connections: 1,
            max_connections: 10,
            acquire_timeout: Duration::from_secs(30),
            idle_timeout: Duration::from_secs(300),
            health_check_interval: Duration::from_secs(30),
        }
    }
}

impl PoolConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_min_connections(mut self, min: usize) -> Self {
        self.min_connections = min;
        self
    }

    pub fn with_max_connections(mut self, max: usize) -> Self {
        self.max_connections = max;
        self
    }

    pub fn with_acquire_timeout(mut self, timeout: Duration) -> Self {
        self.acquire_timeout = timeout;
        self
    }

    pub fn with_idle_timeout(mut self, timeout: Duration) -> Self {
        self.idle_timeout = timeout;
        self
    }
}

struct PooledItem<T> {
    item: T,
    last_used: Instant,
}

impl<T> PooledItem<T> {
    fn new(item: T) -> Self {
        Self {
            item,
            last_used: Instant::now(),
        }
    }

    fn touch(&mut self) {
        self.last_used = Instant::now();
    }

    #[allow(dead_code)]
    fn is_idle(&self, timeout: Duration) -> bool {
        self.last_used.elapsed() > timeout
    }
}

type BoxedFactory<T> =
    Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<T>> + Send>> + Send + Sync>;

pub struct ConnectionPool<T> {
    config: PoolConfig,
    factory: BoxedFactory<T>,
    available: Arc<Mutex<VecDeque<PooledItem<T>>>>,
    semaphore: Arc<Semaphore>,
    total_created: AtomicUsize,
    shutdown: AtomicBool,
}

impl<T: Send + 'static> ConnectionPool<T> {
    pub async fn new<F, Fut>(config: PoolConfig, factory: F) -> Result<Arc<Self>>
    where
        F: Fn() -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<T>> + Send + 'static,
    {
        let boxed_factory: BoxedFactory<T> = Box::new(move || Box::pin(factory()));

        let pool = Arc::new(Self {
            semaphore: Arc::new(Semaphore::new(config.max_connections)),
            config,
            factory: boxed_factory,
            available: Arc::new(Mutex::new(VecDeque::new())),
            total_created: AtomicUsize::new(0),
            shutdown: AtomicBool::new(false),
        });

        for _ in 0..pool.config.min_connections {
            let item = (pool.factory)().await?;
            pool.available.lock().push_back(PooledItem::new(item));
            pool.total_created.fetch_add(1, Ordering::Relaxed);
        }

        Ok(pool)
    }

    pub async fn get(self: &Arc<Self>) -> Result<PoolGuard<T>> {
        if self.shutdown.load(Ordering::Acquire) {
            return Err(RpcError::ClientError("Pool is shutdown".to_string()));
        }

        let permit = tokio::time::timeout(
            self.config.acquire_timeout,
            self.semaphore.clone().acquire_owned(),
        )
        .await
        .map_err(|_| RpcError::Timeout("Acquire connection timeout".to_string()))?
        .map_err(|_| RpcError::ClientError("Pool is closed".to_string()))?;

        if let Some(mut pooled) = self.available.lock().pop_front() {
            pooled.touch();
            return Ok(PoolGuard {
                item: Some(pooled.item),
                pool: self.clone(),
                _permit: permit,
            });
        }

        let item = (self.factory)().await?;
        self.total_created.fetch_add(1, Ordering::Relaxed);

        Ok(PoolGuard {
            item: Some(item),
            pool: self.clone(),
            _permit: permit,
        })
    }

    fn return_item(&self, item: T) {
        if self.shutdown.load(Ordering::Acquire) {
            return;
        }
        self.available.lock().push_back(PooledItem::new(item));
    }

    pub fn available_count(&self) -> usize {
        self.available.lock().len()
    }

    pub fn total_created(&self) -> usize {
        self.total_created.load(Ordering::Relaxed)
    }

    pub fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
        self.available.lock().clear();
    }
}

pub struct PoolGuard<T: Send + 'static> {
    item: Option<T>,
    pool: Arc<ConnectionPool<T>>,
    _permit: tokio::sync::OwnedSemaphorePermit,
}

impl<T: Send + 'static> PoolGuard<T> {
    pub fn get(&self) -> &T {
        self.item.as_ref().unwrap()
    }

    pub fn get_mut(&mut self) -> &mut T {
        self.item.as_mut().unwrap()
    }
}

impl<T: Send + 'static> Drop for PoolGuard<T> {
    fn drop(&mut self) {
        if let Some(item) = self.item.take() {
            self.pool.return_item(item);
        }
    }
}

impl<T: Send + 'static> std::ops::Deref for PoolGuard<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.get()
    }
}

impl<T: Send + 'static> std::ops::DerefMut for PoolGuard<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.get_mut()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicI32;

    #[tokio::test]
    async fn test_pool_basic() {
        let counter = Arc::new(AtomicI32::new(0));
        let counter_clone = counter.clone();

        let config = PoolConfig::default()
            .with_min_connections(2)
            .with_max_connections(5);

        let pool = ConnectionPool::new(config, move || {
            let c = counter_clone.clone();
            async move {
                let val = c.fetch_add(1, Ordering::Relaxed);
                Ok(val)
            }
        })
        .await
        .unwrap();

        assert_eq!(pool.available_count(), 2);
        assert_eq!(pool.total_created(), 2);

        {
            let conn = pool.get().await.unwrap();
            assert!(*conn == 0 || *conn == 1);
            assert_eq!(pool.available_count(), 1);
        }

        assert_eq!(pool.available_count(), 2);
    }

    #[tokio::test]
    async fn test_pool_max_connections() {
        let config = PoolConfig::default()
            .with_min_connections(0)
            .with_max_connections(2)
            .with_acquire_timeout(Duration::from_millis(100));

        let pool = ConnectionPool::new(config, || async { Ok(()) })
            .await
            .unwrap();

        let _conn1 = pool.get().await.unwrap();
        let _conn2 = pool.get().await.unwrap();

        let result = pool.get().await;
        assert!(result.is_err());
    }
}