amaters_net/

pool.rs

//! Connection pool implementation for managing reusable connections
//!
//! Provides connection pooling with configurable limits, health checks,
//! and lifecycle management for efficient resource utilization.

use crate::balancer::{BalancingStrategy, EndpointId, LoadBalancer};
use crate::circuit_breaker::CircuitBreaker;
use crate::error::{NetError, NetResult};
use async_trait::async_trait;
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::time;
use tonic::transport::{Channel, Endpoint};

/// Configuration for connection pool
#[derive(Debug, Clone)]
pub struct PoolConfig {
    /// Minimum number of connections to maintain
    pub min_size: usize,
    /// Maximum number of connections allowed
    pub max_size: usize,
    /// Connection idle timeout (connections idle longer are closed)
    pub idle_timeout: Duration,
    /// Connection maximum lifetime (connections older are closed)
    pub max_lifetime: Duration,
    /// Connection timeout for establishing new connections
    pub connect_timeout: Duration,
    /// Health check interval
    pub health_check_interval: Duration,
    /// Load balancing strategy
    pub balancing_strategy: BalancingStrategy,
    /// Enable circuit breaker
    pub enable_circuit_breaker: bool,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            min_size: 2,
            max_size: 10,
            idle_timeout: Duration::from_secs(300), // 5 minutes
            max_lifetime: Duration::from_secs(1800), // 30 minutes
            connect_timeout: Duration::from_secs(10),
            health_check_interval: Duration::from_secs(30),
            balancing_strategy: BalancingStrategy::LeastConnections,
            enable_circuit_breaker: true,
        }
    }
}

/// Pool statistics
#[derive(Debug, Clone, Default)]
pub struct PoolStats {
    /// Total number of connections (active + idle)
    pub total_connections: usize,
    /// Number of active (in-use) connections
    pub active_connections: usize,
    /// Number of idle (available) connections
    pub idle_connections: usize,
    /// Number of failed connection attempts
    pub failed_connections: u64,
    /// Total connections created
    pub total_created: u64,
    /// Total connections closed
    pub total_closed: u64,
    /// Number of times pool was exhausted (max size reached)
    pub pool_exhausted_count: u64,
    /// Average connection wait time in milliseconds
    pub avg_wait_time_ms: u64,
}

/// Connection metadata
#[derive(Debug)]
struct ConnectionMeta {
    /// gRPC channel
    channel: Channel,
    /// Endpoint ID
    endpoint_id: EndpointId,
    /// Time when connection was created
    created_at: Instant,
    /// Time when connection was last used
    last_used: Instant,
}

impl ConnectionMeta {
    /// Create new connection metadata
    fn new(channel: Channel, endpoint_id: EndpointId) -> Self {
        let now = Instant::now();
        Self {
            channel,
            endpoint_id,
            created_at: now,
            last_used: now,
        }
    }

    /// Check if connection is expired based on idle timeout
    fn is_idle_expired(&self, idle_timeout: Duration) -> bool {
        self.last_used.elapsed() > idle_timeout
    }

    /// Check if connection exceeded max lifetime
    fn is_lifetime_expired(&self, max_lifetime: Duration) -> bool {
        self.created_at.elapsed() > max_lifetime
    }

    /// Update last used timestamp
    fn touch(&mut self) {
        self.last_used = Instant::now();
    }
}

/// Pooled connection wrapper
pub struct PooledConnection {
    meta: Option<ConnectionMeta>,
    pool: Arc<ConnectionPoolInner>,
}

impl PooledConnection {
    /// Get the underlying gRPC channel
    pub fn channel(&self) -> &Channel {
        &self.meta.as_ref().expect("connection should exist").channel
    }

    /// Get endpoint ID
    pub fn endpoint_id(&self) -> &str {
        &self
            .meta
            .as_ref()
            .expect("connection should exist")
            .endpoint_id
    }
}

impl Drop for PooledConnection {
    fn drop(&mut self) {
        if let Some(mut meta) = self.meta.take() {
            meta.touch();
            self.pool.return_connection(meta);
        }
    }
}

/// Internal connection pool state
struct ConnectionPoolInner {
    config: PoolConfig,
    idle_connections: RwLock<VecDeque<ConnectionMeta>>,
    active_count: std::sync::Mutex<usize>,
    stats: RwLock<PoolStats>,
    load_balancer: LoadBalancer,
    circuit_breaker: Option<CircuitBreaker>,
}

impl ConnectionPoolInner {
    /// Return a connection to the pool
    fn return_connection(&self, meta: ConnectionMeta) {
        // Check if connection is expired
        if meta.is_idle_expired(self.config.idle_timeout)
            || meta.is_lifetime_expired(self.config.max_lifetime)
        {
            // Connection expired, don't return to pool
            self.stats.write().total_closed += 1;
            let mut active = self
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = active.saturating_sub(1);
            return;
        }

        // Return to pool
        self.idle_connections.write().push_back(meta);
        let mut active = self
            .active_count
            .lock()
            .expect("active count lock poisoned");
        *active = active.saturating_sub(1);
    }

    /// Get pool statistics
    fn get_stats(&self) -> PoolStats {
        let mut stats = self.stats.read().clone();
        let idle = self.idle_connections.read().len();
        let active = *self
            .active_count
            .lock()
            .expect("active count lock poisoned");
        stats.total_connections = idle + active;
        stats.active_connections = active;
        stats.idle_connections = idle;
        stats
    }
}

/// Connection pool for managing gRPC connections
pub struct ConnectionPool {
    inner: Arc<ConnectionPoolInner>,
    shutdown_tx: tokio::sync::watch::Sender<bool>,
}

impl ConnectionPool {
    /// Create a new connection pool
    pub fn new(config: PoolConfig) -> Self {
        let load_balancer = LoadBalancer::new(config.balancing_strategy);
        let circuit_breaker = if config.enable_circuit_breaker {
            Some(CircuitBreaker::new())
        } else {
            None
        };

        let inner = Arc::new(ConnectionPoolInner {
            config: config.clone(),
            idle_connections: RwLock::new(VecDeque::new()),
            active_count: std::sync::Mutex::new(0),
            stats: RwLock::new(PoolStats::default()),
            load_balancer,
            circuit_breaker,
        });

        let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);

        // Spawn health check task
        let health_check_inner = Arc::clone(&inner);
        tokio::spawn(async move {
            Self::health_check_loop(health_check_inner, shutdown_rx).await;
        });

        Self { inner, shutdown_tx }
    }

    /// Add an endpoint to the connection pool
    pub fn add_endpoint(&self, id: EndpointId, address: String) {
        self.add_endpoint_with_weight(id, address, 1);
    }

    /// Add an endpoint with weight
    pub fn add_endpoint_with_weight(&self, id: EndpointId, address: String, weight: u32) {
        let endpoint = crate::balancer::Endpoint::with_weight(id, address, weight);
        self.inner.load_balancer.add_endpoint(endpoint);
    }

    /// Remove an endpoint from the connection pool
    pub fn remove_endpoint(&self, endpoint_id: &str) -> bool {
        // Remove from load balancer
        let removed = self.inner.load_balancer.remove_endpoint(endpoint_id);

        // Close connections for this endpoint
        if removed {
            let mut idle = self.inner.idle_connections.write();
            idle.retain(|conn| conn.endpoint_id != endpoint_id);
        }

        removed
    }

    /// Get a connection from the pool
    pub async fn get_connection(&self) -> NetResult<PooledConnection> {
        let start = Instant::now();

        // Check circuit breaker
        if let Some(ref cb) = self.inner.circuit_breaker {
            cb.is_request_allowed()?;
        }

        // Try to get idle connection first
        if let Some(mut meta) = self.inner.idle_connections.write().pop_front() {
            meta.touch();
            *self
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned") += 1;

            return Ok(PooledConnection {
                meta: Some(meta),
                pool: Arc::clone(&self.inner),
            });
        }

        // No idle connections, check if we can create a new one
        let active = *self
            .inner
            .active_count
            .lock()
            .expect("active count lock poisoned");
        let idle = self.inner.idle_connections.read().len();

        if active + idle >= self.inner.config.max_size {
            // Pool exhausted, wait for available connection
            self.inner.stats.write().pool_exhausted_count += 1;

            // Wait with timeout
            let timeout = Duration::from_secs(30);
            let deadline = Instant::now() + timeout;

            while Instant::now() < deadline {
                if let Some(mut meta) = self.inner.idle_connections.write().pop_front() {
                    meta.touch();
                    *self
                        .inner
                        .active_count
                        .lock()
                        .expect("active count lock poisoned") += 1;

                    // Update wait time stats
                    let wait_time = start.elapsed().as_millis() as u64;
                    let mut stats = self.inner.stats.write();
                    stats.avg_wait_time_ms = (stats.avg_wait_time_ms + wait_time) / 2;

                    return Ok(PooledConnection {
                        meta: Some(meta),
                        pool: Arc::clone(&self.inner),
                    });
                }

                // Wait a bit before retrying
                time::sleep(Duration::from_millis(100)).await;
            }

            return Err(NetError::ServerOverloaded(
                "Connection pool exhausted".to_string(),
            ));
        }

        // Create new connection
        let meta = self.create_connection().await?;
        *self
            .inner
            .active_count
            .lock()
            .expect("active count lock poisoned") += 1;

        Ok(PooledConnection {
            meta: Some(meta),
            pool: Arc::clone(&self.inner),
        })
    }

    /// Create a new connection
    async fn create_connection(&self) -> NetResult<ConnectionMeta> {
        // Select endpoint using load balancer
        let endpoint = self.inner.load_balancer.select_endpoint()?;

        // Create gRPC channel
        let channel = Endpoint::from_shared(format!("http://{}", endpoint.address))
            .map_err(|e| NetError::InvalidRequest(format!("Invalid endpoint: {}", e)))?
            .connect_timeout(self.inner.config.connect_timeout)
            .timeout(Duration::from_secs(30))
            .connect()
            .await
            .map_err(|e| {
                self.inner.stats.write().failed_connections += 1;
                if let Some(ref cb) = self.inner.circuit_breaker {
                    cb.record_failure();
                }
                NetError::ConnectionRefused(format!("Failed to connect: {}", e))
            })?;

        // Record success
        if let Some(ref cb) = self.inner.circuit_breaker {
            cb.record_success();
        }

        self.inner.stats.write().total_created += 1;

        Ok(ConnectionMeta::new(channel, endpoint.id.clone()))
    }

    /// Health check loop
    async fn health_check_loop(
        inner: Arc<ConnectionPoolInner>,
        mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
    ) {
        let mut interval = time::interval(inner.config.health_check_interval);

        loop {
            tokio::select! {
                _ = interval.tick() => {
                    Self::perform_health_check(&inner).await;
                }
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        break;
                    }
                }
            }
        }
    }

    /// Perform health check on idle connections
    async fn perform_health_check(inner: &Arc<ConnectionPoolInner>) {
        let needed = {
            // Scope the lock to ensure it's dropped before any await
            let mut idle = inner.idle_connections.write();
            let config = &inner.config;

            // Remove expired connections
            idle.retain(|conn| {
                !conn.is_idle_expired(config.idle_timeout)
                    && !conn.is_lifetime_expired(config.max_lifetime)
            });

            // Ensure minimum pool size
            let current_size = idle.len()
                + *inner
                    .active_count
                    .lock()
                    .expect("active count lock poisoned");
            config.min_size.saturating_sub(current_size)
        }; // Lock is dropped here

        // Create needed connections (async operation)
        for _ in 0..needed {
            // This is best effort - we don't wait for results
            // Real implementation would handle this more carefully
            let _ = async {
                // Would create connection here
            }
            .await;
        }
    }

    /// Get pool statistics
    pub fn stats(&self) -> PoolStats {
        self.inner.get_stats()
    }

    /// Get circuit breaker statistics
    pub fn circuit_breaker_stats(&self) -> Option<crate::circuit_breaker::CircuitBreakerStats> {
        self.inner.circuit_breaker.as_ref().map(|cb| cb.stats())
    }

    /// Shutdown the connection pool gracefully
    pub async fn shutdown(self) -> NetResult<()> {
        // Signal shutdown to background tasks
        self.shutdown_tx
            .send(true)
            .map_err(|_| NetError::ServerInternal("Failed to signal shutdown".to_string()))?;

        // Wait for a short period to allow tasks to complete
        time::sleep(Duration::from_millis(500)).await;

        // Close all idle connections
        let mut idle = self.inner.idle_connections.write();
        let count = idle.len();
        idle.clear();

        self.inner.stats.write().total_closed += count as u64;

        Ok(())
    }

    /// Drain the pool (prepare for graceful shutdown)
    pub async fn drain(&self) -> NetResult<()> {
        // Wait for active connections to complete
        let timeout = Duration::from_secs(30);
        let deadline = Instant::now() + timeout;

        while Instant::now() < deadline {
            let active = *self
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned");
            if active == 0 {
                break;
            }
            time::sleep(Duration::from_millis(100)).await;
        }

        let active = *self
            .inner
            .active_count
            .lock()
            .expect("active count lock poisoned");
        if active > 0 {
            return Err(NetError::Timeout(format!(
                "Drain timeout: {} active connections remaining",
                active
            )));
        }

        Ok(())
    }
}

/// Connection pool builder for fluent configuration
pub struct ConnectionPoolBuilder {
    config: PoolConfig,
    endpoints: Vec<(EndpointId, String, u32)>,
}

impl ConnectionPoolBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self {
            config: PoolConfig::default(),
            endpoints: Vec::new(),
        }
    }

    /// Set minimum pool size
    pub fn min_size(mut self, size: usize) -> Self {
        self.config.min_size = size;
        self
    }

    /// Set maximum pool size
    pub fn max_size(mut self, size: usize) -> Self {
        self.config.max_size = size;
        self
    }

    /// Set idle timeout
    pub fn idle_timeout(mut self, timeout: Duration) -> Self {
        self.config.idle_timeout = timeout;
        self
    }

    /// Set max lifetime
    pub fn max_lifetime(mut self, lifetime: Duration) -> Self {
        self.config.max_lifetime = lifetime;
        self
    }

    /// Set connect timeout
    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
        self.config.connect_timeout = timeout;
        self
    }

    /// Set health check interval
    pub fn health_check_interval(mut self, interval: Duration) -> Self {
        self.config.health_check_interval = interval;
        self
    }

    /// Set balancing strategy
    pub fn balancing_strategy(mut self, strategy: BalancingStrategy) -> Self {
        self.config.balancing_strategy = strategy;
        self
    }

    /// Enable or disable circuit breaker
    pub fn circuit_breaker(mut self, enabled: bool) -> Self {
        self.config.enable_circuit_breaker = enabled;
        self
    }

    /// Add an endpoint
    pub fn add_endpoint(mut self, id: EndpointId, address: String) -> Self {
        self.endpoints.push((id, address, 1));
        self
    }

    /// Add an endpoint with weight
    pub fn add_endpoint_with_weight(
        mut self,
        id: EndpointId,
        address: String,
        weight: u32,
    ) -> Self {
        self.endpoints.push((id, address, weight));
        self
    }

    /// Build the connection pool
    pub fn build(self) -> ConnectionPool {
        let pool = ConnectionPool::new(self.config);

        for (id, address, weight) in self.endpoints {
            pool.add_endpoint_with_weight(id, address, weight);
        }

        pool
    }
}

impl Default for ConnectionPoolBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pool_config_default() {
        let config = PoolConfig::default();
        assert_eq!(config.min_size, 2);
        assert_eq!(config.max_size, 10);
        assert!(config.enable_circuit_breaker);
    }

    #[tokio::test]
    async fn test_connection_meta_expiry() {
        // Skip if we can't connect (localhost not available)
        let endpoint = Endpoint::from_static("http://localhost:50051");
        if let Ok(channel) = endpoint.connect().await {
            let meta = ConnectionMeta::new(channel, "ep1".to_string());

            assert!(!meta.is_idle_expired(Duration::from_secs(10)));
            assert!(!meta.is_lifetime_expired(Duration::from_secs(10)));
        }
        // Test passes even without connection - we're testing the struct, not connectivity
    }

    #[tokio::test]
    async fn test_pool_builder() {
        let pool = ConnectionPoolBuilder::new()
            .min_size(5)
            .max_size(20)
            .idle_timeout(Duration::from_secs(600))
            .balancing_strategy(BalancingStrategy::RoundRobin)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .add_endpoint("ep2".to_string(), "localhost:50052".to_string())
            .build();

        let stats = pool.stats();
        assert_eq!(stats.active_connections, 0);
        assert_eq!(stats.idle_connections, 0);
    }

    #[tokio::test]
    async fn test_pool_add_remove_endpoint() {
        let pool = ConnectionPool::new(PoolConfig::default());

        pool.add_endpoint("ep1".to_string(), "localhost:50051".to_string());
        pool.add_endpoint("ep2".to_string(), "localhost:50052".to_string());

        assert!(pool.remove_endpoint("ep1"));
        assert!(!pool.remove_endpoint("ep3"));
    }

    #[tokio::test]
    async fn test_pool_stats() {
        let pool = ConnectionPool::new(PoolConfig::default());
        pool.add_endpoint("ep1".to_string(), "localhost:50051".to_string());

        let stats = pool.stats();
        assert_eq!(stats.total_connections, 0);
        assert_eq!(stats.active_connections, 0);
        assert_eq!(stats.idle_connections, 0);
    }

    #[tokio::test]
    async fn test_pool_shutdown() {
        let pool = ConnectionPool::new(PoolConfig::default());
        pool.add_endpoint("ep1".to_string(), "localhost:50051".to_string());

        // Shutdown should complete successfully
        let result = pool.shutdown().await;
        assert!(result.is_ok());
    }
}