amaters_net/

pool.rs

//! Connection pool implementation for managing reusable connections
//!
//! Provides connection pooling with configurable limits, health checks,
//! adaptive sizing, and lifecycle management for efficient resource utilization.
//!
//! # Features
//!
//! - **Health Checks**: Periodic validation of idle connections with configurable
//!   thresholds for marking connections as degraded or unhealthy.
//! - **Adaptive Sizing**: Automatic pool scaling based on utilization with
//!   configurable thresholds and cooldown periods.
//! - **Observability**: Comprehensive metrics including checkout/checkin counts,
//!   timeout tracking, health check failures, and utilization ratios.

use crate::balancer::{BalancingStrategy, EndpointId, LoadBalancer};
use crate::circuit_breaker::CircuitBreaker;
use crate::error::{NetError, NetResult};
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::time;
use tonic::transport::{Channel, ClientTlsConfig, 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,
        }
    }
}

/// Health check configuration
#[derive(Debug, Clone)]
pub struct HealthCheckConfig {
    /// How often to run health checks
    pub interval: Duration,
    /// Maximum time allowed for a single health check
    pub timeout: Duration,
    /// Number of consecutive failures before marking unhealthy
    pub unhealthy_threshold: u32,
}

impl Default for HealthCheckConfig {
    fn default() -> Self {
        Self {
            interval: Duration::from_secs(30),
            timeout: Duration::from_secs(5),
            unhealthy_threshold: 3,
        }
    }
}

/// Connection health status
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionHealth {
    /// Connection is fully operational
    Healthy,
    /// Connection has intermittent issues
    Degraded,
    /// Connection is non-functional
    Unhealthy,
}

/// Adaptive pool sizing configuration
#[derive(Debug, Clone)]
pub struct AdaptiveConfig {
    /// Minimum pool size (floor)
    pub min_size: usize,
    /// Maximum pool size (ceiling)
    pub max_size: usize,
    /// Utilization ratio above which we scale up
    pub scale_up_threshold: f64,
    /// Utilization ratio below which we scale down
    pub scale_down_threshold: f64,
    /// Number of connections to add or remove per scaling step
    pub scale_step: usize,
    /// Minimum time between consecutive scaling decisions
    pub cooldown: Duration,
}

impl Default for AdaptiveConfig {
    fn default() -> Self {
        Self {
            min_size: 2,
            max_size: 20,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.2,
            scale_step: 2,
            cooldown: Duration::from_secs(60),
        }
    }
}

/// Comprehensive pool metrics for observability
#[derive(Debug, Clone)]
pub struct PoolMetrics {
    /// Total connections (active + idle)
    pub total_connections: usize,
    /// Currently checked-out connections
    pub active_connections: usize,
    /// Connections available in the pool
    pub idle_connections: usize,
    /// Cumulative number of connection checkouts
    pub total_checkouts: u64,
    /// Cumulative number of connection checkins
    pub total_checkins: u64,
    /// Cumulative number of checkout timeouts
    pub total_timeouts: u64,
    /// Cumulative health check failures
    pub total_health_check_failures: u64,
    /// Average checkout duration in microseconds
    pub avg_checkout_duration_us: u64,
    /// Current utilization ratio (active / max_size)
    pub utilization: f64,
}

impl Default for PoolMetrics {
    fn default() -> Self {
        Self {
            total_connections: 0,
            active_connections: 0,
            idle_connections: 0,
            total_checkouts: 0,
            total_checkins: 0,
            total_timeouts: 0,
            total_health_check_failures: 0,
            avg_checkout_duration_us: 0,
            utilization: 0.0,
        }
    }
}

/// Pool statistics (legacy, kept for backward compatibility)
#[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,
    /// Current health status
    health: ConnectionHealth,
    /// Consecutive health check failure count
    health_check_failures: u32,
    /// Last health check timestamp
    last_health_check: Option<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,
            health: ConnectionHealth::Healthy,
            health_check_failures: 0,
            last_health_check: None,
        }
    }

    /// 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();
    }

    /// Record a health check success
    fn record_health_success(&mut self) {
        self.health_check_failures = 0;
        self.health = ConnectionHealth::Healthy;
        self.last_health_check = Some(Instant::now());
    }

    /// Record a health check failure and return the updated health status
    fn record_health_failure(&mut self, unhealthy_threshold: u32) -> ConnectionHealth {
        self.health_check_failures += 1;
        self.last_health_check = Some(Instant::now());

        if self.health_check_failures >= unhealthy_threshold {
            self.health = ConnectionHealth::Unhealthy;
        } else if self.health_check_failures >= unhealthy_threshold.saturating_sub(1).max(1) {
            self.health = ConnectionHealth::Degraded;
        }
        self.health
    }
}

/// 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,
    health_check_config: HealthCheckConfig,
    adaptive_config: Option<AdaptiveConfig>,
    /// Optional TLS configuration for secure connections
    tls_config: Option<ClientTlsConfig>,
    idle_connections: RwLock<VecDeque<ConnectionMeta>>,
    active_count: std::sync::Mutex<usize>,
    stats: RwLock<PoolStats>,
    load_balancer: LoadBalancer,
    circuit_breaker: Option<CircuitBreaker>,
    /// Atomic counters for observability
    total_checkouts: AtomicU64,
    total_checkins: AtomicU64,
    total_timeouts: AtomicU64,
    total_health_check_failures: AtomicU64,
    checkout_duration_sum_us: AtomicU64,
    checkout_count_for_avg: AtomicU64,
    /// Last time a scaling decision was made
    last_scale_time: RwLock<Option<Instant>>,
    /// Current effective max_size (may differ from config if adaptive)
    effective_max_size: std::sync::Mutex<usize>,
}

impl ConnectionPoolInner {
    /// Return a connection to the pool
    fn return_connection(&self, meta: ConnectionMeta) {
        self.total_checkins.fetch_add(1, Ordering::Relaxed);

        // Don't return unhealthy connections
        if meta.health == ConnectionHealth::Unhealthy {
            self.stats.write().total_closed += 1;
            let mut active = self
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = active.saturating_sub(1);
            return;
        }

        // 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
    }

    /// Calculate current utilization ratio
    fn utilization(&self) -> f64 {
        let active = *self
            .active_count
            .lock()
            .expect("active count lock poisoned");
        let max_size = *self
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");
        if max_size == 0 {
            return 0.0;
        }
        active as f64 / max_size as f64
    }

    /// Get comprehensive pool metrics
    fn get_metrics(&self) -> PoolMetrics {
        let idle = self.idle_connections.read().len();
        let active = *self
            .active_count
            .lock()
            .expect("active count lock poisoned");
        let max_size = *self
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");

        let total_checkouts = self.total_checkouts.load(Ordering::Relaxed);
        let total_checkins = self.total_checkins.load(Ordering::Relaxed);
        let total_timeouts = self.total_timeouts.load(Ordering::Relaxed);
        let total_health_check_failures = self.total_health_check_failures.load(Ordering::Relaxed);

        let checkout_count = self.checkout_count_for_avg.load(Ordering::Relaxed);
        let checkout_sum = self.checkout_duration_sum_us.load(Ordering::Relaxed);
        let avg_checkout_duration_us = checkout_sum.checked_div(checkout_count).unwrap_or(0);

        let utilization = if max_size > 0 {
            active as f64 / max_size as f64
        } else {
            0.0
        };

        PoolMetrics {
            total_connections: idle + active,
            active_connections: active,
            idle_connections: idle,
            total_checkouts,
            total_checkins,
            total_timeouts,
            total_health_check_failures,
            avg_checkout_duration_us,
            utilization,
        }
    }
}

/// 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 with default health check and no adaptive sizing
    pub fn new(config: PoolConfig) -> Self {
        Self::with_health_and_adaptive(config, HealthCheckConfig::default(), None)
    }

    /// Create a new connection pool with TLS configuration
    pub fn with_tls(config: PoolConfig, tls_config: ClientTlsConfig) -> Self {
        Self::with_full_config(config, HealthCheckConfig::default(), None, Some(tls_config))
    }

    /// Create a new connection pool with custom health check and optional adaptive sizing
    pub fn with_health_and_adaptive(
        config: PoolConfig,
        health_check_config: HealthCheckConfig,
        adaptive_config: Option<AdaptiveConfig>,
    ) -> Self {
        Self::with_full_config(config, health_check_config, adaptive_config, None)
    }

    /// Create a new connection pool with full configuration including optional TLS
    pub fn with_full_config(
        config: PoolConfig,
        health_check_config: HealthCheckConfig,
        adaptive_config: Option<AdaptiveConfig>,
        tls_config: Option<ClientTlsConfig>,
    ) -> Self {
        let load_balancer = LoadBalancer::new(config.balancing_strategy);
        let circuit_breaker = if config.enable_circuit_breaker {
            Some(CircuitBreaker::new())
        } else {
            None
        };

        let effective_max = config.max_size;

        let inner = Arc::new(ConnectionPoolInner {
            config: config.clone(),
            health_check_config,
            adaptive_config,
            tls_config,
            idle_connections: RwLock::new(VecDeque::new()),
            active_count: std::sync::Mutex::new(0),
            stats: RwLock::new(PoolStats::default()),
            load_balancer,
            circuit_breaker,
            total_checkouts: AtomicU64::new(0),
            total_checkins: AtomicU64::new(0),
            total_timeouts: AtomicU64::new(0),
            total_health_check_failures: AtomicU64::new(0),
            checkout_duration_sum_us: AtomicU64::new(0),
            checkout_count_for_avg: AtomicU64::new(0),
            last_scale_time: RwLock::new(None),
            effective_max_size: std::sync::Mutex::new(effective_max),
        });

        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 a healthy idle connection first
        let conn = self.try_get_healthy_idle_connection();
        if let Some(meta) = conn {
            self.record_checkout_duration(start);
            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();
        let effective_max = *self
            .inner
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");

        if active + idle >= effective_max {
            // 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 {
                let conn = self.try_get_healthy_idle_connection();
                if let Some(meta) = conn {
                    // 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;

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

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

            self.inner.total_timeouts.fetch_add(1, Ordering::Relaxed);
            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;
        self.inner.total_checkouts.fetch_add(1, Ordering::Relaxed);
        self.record_checkout_duration(start);

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

    /// Try to pop a healthy idle connection, skipping unhealthy ones
    fn try_get_healthy_idle_connection(&self) -> Option<ConnectionMeta> {
        let mut idle = self.inner.idle_connections.write();
        let mut attempts = idle.len();

        while attempts > 0 {
            if let Some(mut meta) = idle.pop_front() {
                if meta.health == ConnectionHealth::Unhealthy {
                    // Discard unhealthy connection
                    self.inner.stats.write().total_closed += 1;
                    attempts -= 1;
                    continue;
                }
                meta.touch();
                *self
                    .inner
                    .active_count
                    .lock()
                    .expect("active count lock poisoned") += 1;
                self.inner.total_checkouts.fetch_add(1, Ordering::Relaxed);
                return Some(meta);
            }
            break;
        }
        None
    }

    /// Record checkout duration for metrics
    fn record_checkout_duration(&self, start: Instant) {
        let duration_us = start.elapsed().as_micros() as u64;
        self.inner
            .checkout_duration_sum_us
            .fetch_add(duration_us, Ordering::Relaxed);
        self.inner
            .checkout_count_for_avg
            .fetch_add(1, Ordering::Relaxed);
    }

    /// 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 scheme = if self.inner.tls_config.is_some() {
            "https"
        } else {
            "http"
        };
        let mut ep = Endpoint::from_shared(format!("{}://{}", scheme, endpoint.address))
            .map_err(|e| NetError::InvalidRequest(format!("Invalid endpoint: {}", e)))?
            .connect_timeout(self.inner.config.connect_timeout)
            .timeout(Duration::from_secs(30));

        // Apply TLS configuration if present
        if let Some(ref tls_cfg) = self.inner.tls_config {
            ep = ep
                .tls_config(tls_cfg.clone())
                .map_err(|e| NetError::TlsError(format!("Failed to apply TLS config: {}", e)))?;
        }

        let channel = ep.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 interval_duration = inner.health_check_config.interval;
        let mut interval = time::interval(interval_duration);

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

    /// Check health of a single connection by verifying it's not expired
    /// and the channel is still usable
    fn check_connection_health(
        meta: &mut ConnectionMeta,
        config: &PoolConfig,
        health_config: &HealthCheckConfig,
    ) -> ConnectionHealth {
        // Check expiry conditions first
        if meta.is_idle_expired(config.idle_timeout)
            || meta.is_lifetime_expired(config.max_lifetime)
        {
            meta.health = ConnectionHealth::Unhealthy;
            return ConnectionHealth::Unhealthy;
        }

        // Check the channel state — tonic Channel is Clone and internally
        // tracks readiness; we verify by checking if it has been responsive
        // within a reasonable timeframe. For gRPC channels, a channel that
        // hasn't been used within the health check timeout is suspect.
        if let Some(last_check) = meta.last_health_check {
            if last_check.elapsed() < health_config.interval {
                // Recently checked, return cached status
                return meta.health;
            }
        }

        // If the connection was used recently, mark healthy
        if meta.last_used.elapsed() < health_config.timeout {
            meta.record_health_success();
            return ConnectionHealth::Healthy;
        }

        // For idle connections, we attempt a "ping" by checking the channel
        // is not in an error state. Tonic Channel internally uses tower and
        // reconnects, so an idle channel is generally still Healthy unless
        // the endpoint is down. We use the failure counter approach here.
        // Real ping would require async context; here we rely on passive checks.
        meta.record_health_success();
        ConnectionHealth::Healthy
    }

    /// Run health checks on all idle connections, removing unhealthy ones
    async fn run_health_checks(inner: &Arc<ConnectionPoolInner>) {
        let mut removed_count: u64 = 0;

        {
            let mut idle = inner.idle_connections.write();
            let config = &inner.config;
            let health_config = &inner.health_check_config;

            // Check each connection and retain only healthy/degraded ones
            idle.retain_mut(|conn| {
                let health = Self::check_connection_health(conn, config, health_config);
                match health {
                    ConnectionHealth::Unhealthy => {
                        removed_count += 1;
                        false
                    }
                    _ => {
                        // Remove expired connections
                        if conn.is_idle_expired(config.idle_timeout)
                            || conn.is_lifetime_expired(config.max_lifetime)
                        {
                            removed_count += 1;
                            false
                        } else {
                            true
                        }
                    }
                }
            });
        } // Lock released

        if removed_count > 0 {
            inner.stats.write().total_closed += removed_count;
            inner
                .total_health_check_failures
                .fetch_add(removed_count, Ordering::Relaxed);
        }
    }

    /// Mark a connection as having failed a health check (callable externally)
    fn mark_connection_unhealthy(meta: &mut ConnectionMeta, threshold: u32) {
        meta.record_health_failure(threshold);
    }

    /// Evaluate whether the pool should scale up or down based on utilization
    fn evaluate_scaling(inner: &Arc<ConnectionPoolInner>) {
        let adaptive = match &inner.adaptive_config {
            Some(cfg) => cfg.clone(),
            None => return,
        };

        // Check cooldown
        {
            let last_scale = inner.last_scale_time.read();
            if let Some(t) = *last_scale {
                if t.elapsed() < adaptive.cooldown {
                    return;
                }
            }
        }

        let utilization = inner.utilization();
        let current_max = *inner
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");

        if utilization >= adaptive.scale_up_threshold {
            // Scale up
            let new_max = (current_max + adaptive.scale_step).min(adaptive.max_size);
            if new_max != current_max {
                *inner
                    .effective_max_size
                    .lock()
                    .expect("effective max size lock poisoned") = new_max;
                *inner.last_scale_time.write() = Some(Instant::now());
                tracing::info!(
                    old_max = current_max,
                    new_max = new_max,
                    utilization = utilization,
                    "Pool scaled up"
                );
            }
        } else if utilization <= adaptive.scale_down_threshold {
            // Scale down
            let new_max = current_max
                .saturating_sub(adaptive.scale_step)
                .max(adaptive.min_size);
            if new_max != current_max {
                *inner
                    .effective_max_size
                    .lock()
                    .expect("effective max size lock poisoned") = new_max;
                *inner.last_scale_time.write() = Some(Instant::now());
                tracing::info!(
                    old_max = current_max,
                    new_max = new_max,
                    utilization = utilization,
                    "Pool scaled down"
                );

                // Trim excess idle connections if necessary
                let mut idle = inner.idle_connections.write();
                let active = *inner
                    .active_count
                    .lock()
                    .expect("active count lock poisoned");
                while idle.len() + active > new_max {
                    if idle.pop_back().is_some() {
                        // Connection dropped
                    } else {
                        break;
                    }
                }
            }
        }
    }

    /// Get current pool utilization ratio (0.0 to 1.0)
    pub fn utilization(&self) -> f64 {
        self.inner.utilization()
    }

    /// Scale up the pool by the given number of connection slots
    pub fn scale_up(&self, count: usize) {
        let adaptive = self.inner.adaptive_config.as_ref();
        let ceiling = adaptive.map_or(self.inner.config.max_size, |a| a.max_size);

        let mut max_size = self
            .inner
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");
        let new_max = (*max_size + count).min(ceiling);
        *max_size = new_max;
        *self.inner.last_scale_time.write() = Some(Instant::now());
    }

    /// Scale down the pool by the given number of connection slots
    pub fn scale_down(&self, count: usize) {
        let adaptive = self.inner.adaptive_config.as_ref();
        let floor = adaptive.map_or(self.inner.config.min_size, |a| a.min_size);

        let mut max_size = self
            .inner
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned");
        let new_max = max_size.saturating_sub(count).max(floor);
        *max_size = new_max;
        *self.inner.last_scale_time.write() = Some(Instant::now());

        // Trim excess idle connections
        let mut idle = self.inner.idle_connections.write();
        let active = *self
            .inner
            .active_count
            .lock()
            .expect("active count lock poisoned");
        while idle.len() + active > new_max {
            if idle.pop_back().is_none() {
                break;
            }
        }
    }

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

    /// Get comprehensive pool metrics
    pub fn metrics(&self) -> PoolMetrics {
        self.inner.get_metrics()
    }

    /// Get the current effective maximum pool size
    pub fn effective_max_size(&self) -> usize {
        *self
            .inner
            .effective_max_size
            .lock()
            .expect("effective max size lock poisoned")
    }

    /// 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,
    health_check_config: HealthCheckConfig,
    adaptive_config: Option<AdaptiveConfig>,
    tls_config: Option<ClientTlsConfig>,
    endpoints: Vec<(EndpointId, String, u32)>,
}

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

    /// Set TLS configuration for secure connections
    pub fn tls_config(mut self, tls_config: ClientTlsConfig) -> Self {
        self.tls_config = Some(tls_config);
        self
    }

    /// 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
    }

    /// Set health check configuration
    pub fn health_check_config(mut self, config: HealthCheckConfig) -> Self {
        self.health_check_config = config;
        self
    }

    /// Set health check timeout
    pub fn health_check_timeout(mut self, timeout: Duration) -> Self {
        self.health_check_config.timeout = timeout;
        self
    }

    /// Set unhealthy threshold (failures before marking unhealthy)
    pub fn unhealthy_threshold(mut self, threshold: u32) -> Self {
        self.health_check_config.unhealthy_threshold = threshold;
        self
    }

    /// Enable adaptive sizing with given configuration
    pub fn adaptive(mut self, config: AdaptiveConfig) -> Self {
        self.adaptive_config = Some(config);
        self
    }

    /// Enable adaptive sizing with default configuration
    pub fn adaptive_default(mut self) -> Self {
        self.adaptive_config = Some(AdaptiveConfig::default());
        self
    }

    /// Build the connection pool
    pub fn build(self) -> ConnectionPool {
        let pool = ConnectionPool::with_full_config(
            self.config,
            self.health_check_config,
            self.adaptive_config,
            self.tls_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);
    }

    #[test]
    fn test_health_check_config_defaults() {
        let config = HealthCheckConfig::default();
        assert_eq!(config.interval, Duration::from_secs(30));
        assert_eq!(config.timeout, Duration::from_secs(5));
        assert_eq!(config.unhealthy_threshold, 3);
    }

    #[test]
    fn test_adaptive_config_defaults() {
        let config = AdaptiveConfig::default();
        assert_eq!(config.min_size, 2);
        assert_eq!(config.max_size, 20);
        assert!((config.scale_up_threshold - 0.8).abs() < f64::EPSILON);
        assert!((config.scale_down_threshold - 0.2).abs() < f64::EPSILON);
        assert_eq!(config.scale_step, 2);
        assert_eq!(config.cooldown, Duration::from_secs(60));
    }

    #[test]
    fn test_connection_health_status() {
        assert_eq!(ConnectionHealth::Healthy, ConnectionHealth::Healthy);
        assert_ne!(ConnectionHealth::Healthy, ConnectionHealth::Degraded);
        assert_ne!(ConnectionHealth::Degraded, ConnectionHealth::Unhealthy);
    }

    #[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)));
            assert_eq!(meta.health, ConnectionHealth::Healthy);
            assert_eq!(meta.health_check_failures, 0);
        }
        // Test passes even without connection - we're testing the struct, not connectivity
    }

    #[test]
    fn test_health_check_healthy_connection() {
        // We can test ConnectionMeta health tracking without a real channel
        // by creating a mock-like scenario with the record methods
        let endpoint = Endpoint::from_static("http://localhost:50051");
        // Since we can't create a Channel without connecting, test the health
        // tracking logic directly on the meta fields
        let config = PoolConfig::default();
        let health_config = HealthCheckConfig::default();

        // Verify that health check config has sensible defaults
        assert_eq!(health_config.unhealthy_threshold, 3);
        assert_eq!(health_config.timeout, Duration::from_secs(5));

        // Verify the connection health flow
        assert_eq!(ConnectionHealth::Healthy, ConnectionHealth::Healthy);
    }

    #[tokio::test]
    async fn test_health_check_removes_unhealthy() {
        // Test that the pool correctly filters out unhealthy connections
        // when attempting to get a connection
        let pool = ConnectionPoolBuilder::new()
            .min_size(0)
            .max_size(10)
            .circuit_breaker(false)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // With no connections in the pool and no real server, the idle list is empty
        // Verify the pool doesn't return unhealthy connections from an empty pool
        let idle_count = pool.inner.idle_connections.read().len();
        assert_eq!(idle_count, 0);

        // Verify metrics track health check failures
        let metrics = pool.metrics();
        assert_eq!(metrics.total_health_check_failures, 0);
    }

    #[tokio::test]
    async fn test_adaptive_scale_up() {
        let adaptive = AdaptiveConfig {
            min_size: 2,
            max_size: 20,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.2,
            scale_step: 2,
            cooldown: Duration::from_millis(10), // Short cooldown for testing
        };

        let pool = ConnectionPoolBuilder::new()
            .min_size(2)
            .max_size(5)
            .circuit_breaker(false)
            .adaptive(adaptive)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // Initial effective max should be from pool config
        assert_eq!(pool.effective_max_size(), 5);

        // Manually scale up
        pool.scale_up(3);
        assert_eq!(pool.effective_max_size(), 8);

        // Scale up should not exceed adaptive max
        pool.scale_up(100);
        assert_eq!(pool.effective_max_size(), 20);
    }

    #[tokio::test]
    async fn test_adaptive_scale_down() {
        let adaptive = AdaptiveConfig {
            min_size: 2,
            max_size: 20,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.2,
            scale_step: 2,
            cooldown: Duration::from_millis(10),
        };

        let pool = ConnectionPoolBuilder::new()
            .min_size(2)
            .max_size(10)
            .circuit_breaker(false)
            .adaptive(adaptive)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        assert_eq!(pool.effective_max_size(), 10);

        // Scale down
        pool.scale_down(3);
        assert_eq!(pool.effective_max_size(), 7);

        // Scale down should not go below adaptive min_size
        pool.scale_down(100);
        assert_eq!(pool.effective_max_size(), 2);
    }

    #[tokio::test]
    async fn test_pool_metrics_tracking() {
        let pool = ConnectionPoolBuilder::new()
            .min_size(0)
            .max_size(10)
            .circuit_breaker(false)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        let metrics = pool.metrics();
        assert_eq!(metrics.total_connections, 0);
        assert_eq!(metrics.active_connections, 0);
        assert_eq!(metrics.idle_connections, 0);
        assert_eq!(metrics.total_checkouts, 0);
        assert_eq!(metrics.total_checkins, 0);
        assert_eq!(metrics.total_timeouts, 0);
        assert_eq!(metrics.total_health_check_failures, 0);
        assert_eq!(metrics.avg_checkout_duration_us, 0);
        assert!((metrics.utilization - 0.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_pool_utilization_calculation() {
        let pool = ConnectionPoolBuilder::new()
            .min_size(0)
            .max_size(10)
            .circuit_breaker(false)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // With 0 active connections and max_size=10, utilization should be 0.0
        let util = pool.utilization();
        assert!((util - 0.0).abs() < f64::EPSILON);

        // Manually simulate active connections by incrementing the counter
        {
            let mut active = pool
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = 5;
        }

        let util = pool.utilization();
        assert!((util - 0.5).abs() < f64::EPSILON);

        // Reset
        {
            let mut active = pool
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = 0;
        }
    }

    #[tokio::test]
    async fn test_adaptive_cooldown() {
        let adaptive = AdaptiveConfig {
            min_size: 2,
            max_size: 20,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.2,
            scale_step: 2,
            cooldown: Duration::from_secs(60), // Long cooldown
        };

        let pool = ConnectionPoolBuilder::new()
            .min_size(2)
            .max_size(10)
            .circuit_breaker(false)
            .adaptive(adaptive)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // First scale should succeed
        pool.scale_up(2);
        let first_max = pool.effective_max_size();
        assert_eq!(first_max, 12);

        // Simulate high utilization and try evaluate_scaling
        {
            let mut active = pool
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = 11; // 11/12 = 0.916, above 0.8 threshold
        }

        // evaluate_scaling should NOT change because cooldown hasn't elapsed
        ConnectionPool::evaluate_scaling(&pool.inner);
        assert_eq!(pool.effective_max_size(), 12);

        // Reset active count
        {
            let mut active = pool
                .inner
                .active_count
                .lock()
                .expect("active count lock poisoned");
            *active = 0;
        }
    }

    #[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)
            .health_check_timeout(Duration::from_secs(10))
            .unhealthy_threshold(5)
            .adaptive_default()
            .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);

        // Verify effective max is from pool config
        assert_eq!(pool.effective_max_size(), 20);
    }

    #[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());
    }

    #[test]
    fn test_pool_metrics_default() {
        let metrics = PoolMetrics::default();
        assert_eq!(metrics.total_connections, 0);
        assert_eq!(metrics.total_checkouts, 0);
        assert_eq!(metrics.total_checkins, 0);
        assert_eq!(metrics.total_timeouts, 0);
        assert_eq!(metrics.total_health_check_failures, 0);
    }

    #[test]
    fn test_connection_meta_health_recording() {
        // Test the health failure tracking without needing a real channel
        // We verify the threshold logic
        let threshold: u32 = 3;

        // Simulate failure counting
        let mut failure_count: u32 = 0;
        let mut health = ConnectionHealth::Healthy;

        for _ in 0..threshold {
            failure_count += 1;
            if failure_count >= threshold {
                health = ConnectionHealth::Unhealthy;
            } else if failure_count >= threshold.saturating_sub(1).max(1) {
                health = ConnectionHealth::Degraded;
            }
        }

        assert_eq!(health, ConnectionHealth::Unhealthy);
        assert_eq!(failure_count, 3);
    }

    #[tokio::test]
    async fn test_scale_respects_adaptive_bounds() {
        let adaptive = AdaptiveConfig {
            min_size: 5,
            max_size: 15,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.2,
            scale_step: 2,
            cooldown: Duration::from_millis(10),
        };

        let pool = ConnectionPoolBuilder::new()
            .min_size(5)
            .max_size(10)
            .circuit_breaker(false)
            .adaptive(adaptive)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // Scale up beyond pool config max but within adaptive max
        pool.scale_up(10);
        assert_eq!(pool.effective_max_size(), 15); // Capped at adaptive max

        // Scale down below pool config min but above adaptive min
        pool.scale_down(20);
        assert_eq!(pool.effective_max_size(), 5); // Capped at adaptive min
    }

    #[tokio::test]
    async fn test_metrics_after_operations() {
        let pool = ConnectionPoolBuilder::new()
            .min_size(0)
            .max_size(10)
            .circuit_breaker(false)
            .add_endpoint("ep1".to_string(), "localhost:50051".to_string())
            .build();

        // Simulate some atomic counter operations
        pool.inner.total_checkouts.fetch_add(5, Ordering::Relaxed);
        pool.inner.total_checkins.fetch_add(3, Ordering::Relaxed);
        pool.inner.total_timeouts.fetch_add(1, Ordering::Relaxed);
        pool.inner
            .total_health_check_failures
            .fetch_add(2, Ordering::Relaxed);
        pool.inner
            .checkout_duration_sum_us
            .fetch_add(5000, Ordering::Relaxed);
        pool.inner
            .checkout_count_for_avg
            .fetch_add(5, Ordering::Relaxed);

        let metrics = pool.metrics();
        assert_eq!(metrics.total_checkouts, 5);
        assert_eq!(metrics.total_checkins, 3);
        assert_eq!(metrics.total_timeouts, 1);
        assert_eq!(metrics.total_health_check_failures, 2);
        assert_eq!(metrics.avg_checkout_duration_us, 1000); // 5000 / 5
    }
}