zus_common/

pool.rs

//! Connection Pool for ZUS RPC Client
//!
//! Provides connection pooling for high-throughput scenarios like gateways.
//! Default configuration (pool_size=1) maintains backward compatibility.

use {
  parking_lot::RwLock,
  std::{
    sync::{
      Arc,
      atomic::{AtomicU64, AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
  },
  tokio::sync::Semaphore,
  tracing::{debug, info, warn},
};

use crate::{ZusError, endpoint::RpcEndpoint, error::Result};

/// Connection pool configuration
#[derive(Debug, Clone)]
pub struct ConnectionPoolConfig {
  /// Minimum connections to maintain (default: 1)
  pub min_connections: usize,

  /// Maximum connections allowed (default: 1 for backward compatibility)
  pub max_connections: usize,

  /// Idle timeout before closing excess connections (default: 60s)
  pub idle_timeout: Duration,

  /// Connection timeout (default: 5s)
  pub connect_timeout: Duration,

  /// Health check interval (default: 10s)
  pub health_check_interval: Duration,
}

impl Default for ConnectionPoolConfig {
  fn default() -> Self {
    Self {
      min_connections: 1,
      max_connections: 1, // Backward compatible
      idle_timeout: Duration::from_secs(60),
      connect_timeout: Duration::from_secs(5),
      health_check_interval: Duration::from_secs(10),
    }
  }
}

impl ConnectionPoolConfig {
  /// Configuration optimized for gateway scenarios
  pub fn for_gateway() -> Self {
    Self {
      min_connections: 2,
      max_connections: 10,
      idle_timeout: Duration::from_secs(300),
      connect_timeout: Duration::from_secs(5),
      health_check_interval: Duration::from_secs(10),
    }
  }

  /// Create configuration with specific pool size
  pub fn with_pool_size(min: usize, max: usize) -> Self {
    Self {
      min_connections: min,
      max_connections: max,
      ..Default::default()
    }
  }
}

/// Strategy for selecting connections from the pool
#[derive(Debug, Clone, Copy, Default)]
pub enum ConnectionSelectionStrategy {
  /// Round-robin selection (default)
  #[default]
  RoundRobin,

  /// Select connection with least pending requests
  LeastPending,
}

/// Health status of a pooled connection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionHealth {
  Healthy,
  Unhealthy,
}

/// A connection wrapper with health and usage tracking
pub struct PooledConnection {
  endpoint: Arc<RpcEndpoint>,
  health: RwLock<ConnectionHealth>,
  pending_requests: AtomicU64,
  last_used: RwLock<Instant>,
  created_at: Instant,
}

impl PooledConnection {
  fn new(endpoint: RpcEndpoint) -> Self {
    Self {
      endpoint: Arc::new(endpoint),
      health: RwLock::new(ConnectionHealth::Healthy),
      pending_requests: AtomicU64::new(0),
      last_used: RwLock::new(Instant::now()),
      created_at: Instant::now(),
    }
  }

  /// Get the underlying endpoint
  pub fn endpoint(&self) -> Arc<RpcEndpoint> {
    *self.last_used.write() = Instant::now();
    self.endpoint.clone()
  }

  /// Increment pending request count
  pub fn inc_pending(&self) {
    self.pending_requests.fetch_add(1, Ordering::SeqCst);
  }

  /// Decrement pending request count
  pub fn dec_pending(&self) {
    self.pending_requests.fetch_sub(1, Ordering::SeqCst);
  }

  /// Get pending request count
  pub fn pending_count(&self) -> u64 {
    self.pending_requests.load(Ordering::SeqCst)
  }

  /// Check if connection is healthy
  pub fn is_healthy(&self) -> bool {
    *self.health.read() == ConnectionHealth::Healthy
  }

  /// Mark connection as unhealthy
  pub fn mark_unhealthy(&self) {
    *self.health.write() = ConnectionHealth::Unhealthy;
  }

  /// Mark connection as healthy
  pub fn mark_healthy(&self) {
    *self.health.write() = ConnectionHealth::Healthy;
  }

  /// Get time since last use
  pub fn idle_duration(&self) -> Duration {
    self.last_used.read().elapsed()
  }

  /// Get connection age
  pub fn age(&self) -> Duration {
    self.created_at.elapsed()
  }

  /// Get the address of this connection
  pub fn address(&self) -> String {
    self.endpoint.address()
  }
}

/// Connection pool for a single server
///
/// Manages multiple connections to the same server for high throughput.
pub struct ConnectionPool {
  host: String,
  port: u16,
  config: ConnectionPoolConfig,
  connections: RwLock<Vec<Arc<PooledConnection>>>,
  round_robin_index: AtomicUsize,
  selection_strategy: ConnectionSelectionStrategy,
  #[allow(dead_code)] // Reserved for future back-pressure support
  connection_semaphore: Semaphore,
}

impl ConnectionPool {
  /// Create a new connection pool
  pub async fn new(host: String, port: u16, config: ConnectionPoolConfig) -> Result<Arc<Self>> {
    let max_connections = config.max_connections;

    let pool = Arc::new(Self {
      host: host.clone(),
      port,
      config: config.clone(),
      connections: RwLock::new(Vec::with_capacity(config.max_connections)),
      round_robin_index: AtomicUsize::new(0),
      selection_strategy: ConnectionSelectionStrategy::RoundRobin,
      connection_semaphore: Semaphore::new(max_connections),
    });

    // Create minimum connections
    pool.ensure_min_connections().await?;

    // Start background maintenance task
    let pool_clone = pool.clone();
    tokio::spawn(async move {
      pool_clone.maintenance_loop().await;
    });

    info!(
      "Connection pool created for {}:{} (min={}, max={})",
      host, port, config.min_connections, config.max_connections
    );

    Ok(pool)
  }

  /// Create pool with default config (single connection, backward compatible)
  pub async fn new_default(host: String, port: u16) -> Result<Arc<Self>> {
    Self::new(host, port, ConnectionPoolConfig::default()).await
  }

  /// Create pool with gateway config
  pub async fn new_for_gateway(host: String, port: u16) -> Result<Arc<Self>> {
    Self::new(host, port, ConnectionPoolConfig::for_gateway()).await
  }

  /// Set connection selection strategy
  pub fn with_selection_strategy(self: Arc<Self>, strategy: ConnectionSelectionStrategy) -> Arc<Self> {
    // Note: This creates a new pool with the strategy set
    // In practice, you'd set this before sharing the pool
    Arc::new(Self {
      host: self.host.clone(),
      port: self.port,
      config: self.config.clone(),
      connections: RwLock::new(self.connections.read().clone()),
      round_robin_index: AtomicUsize::new(0),
      selection_strategy: strategy,
      connection_semaphore: Semaphore::new(self.config.max_connections),
    })
  }

  /// Ensure minimum connections are established
  async fn ensure_min_connections(&self) -> Result<()> {
    let current_count = self.connections.read().len();
    let needed = self.config.min_connections.saturating_sub(current_count);

    for _ in 0..needed {
      if let Err(e) = self.create_connection().await {
        warn!(
          "Failed to create initial connection to {}:{}: {:?}",
          self.host, self.port, e
        );
        // Continue trying to create other connections
      }
    }

    let final_count = self.connections.read().len();
    if final_count == 0 {
      return Err(ZusError::Connection(format!(
        "Failed to create any connections to {}:{}",
        self.host, self.port
      )));
    }

    Ok(())
  }

  /// Create a new connection and add to pool
  async fn create_connection(&self) -> Result<Arc<PooledConnection>> {
    // Check if we can create more connections
    let current_count = self.connections.read().len();
    if current_count >= self.config.max_connections {
      return Err(ZusError::Connection("Pool at maximum capacity".to_string()));
    }

    // Create new connection
    let endpoint = tokio::time::timeout(
      self.config.connect_timeout,
      RpcEndpoint::connect(self.host.clone(), self.port),
    )
    .await
    .map_err(|_| ZusError::Timeout)??;

    let pooled = Arc::new(PooledConnection::new(endpoint));

    // Add to pool
    self.connections.write().push(pooled.clone());

    debug!(
      "Created new connection to {}:{} (pool size: {})",
      self.host,
      self.port,
      self.connections.read().len()
    );

    Ok(pooled)
  }

  /// Get a connection from the pool
  pub async fn get_connection(&self) -> Result<Arc<PooledConnection>> {
    // Try to get existing healthy connection
    if let Some(conn) = self.select_connection() {
      return Ok(conn);
    }

    // No healthy connections, try to create new one
    if self.connections.read().len() < self.config.max_connections
      && let Ok(conn) = self.create_connection().await
    {
      return Ok(conn);
    }

    // Return any connection as last resort
    let connections = self.connections.read();
    if connections.is_empty() {
      return Err(ZusError::Connection("No connections available".to_string()));
    }

    let index = self.round_robin_index.fetch_add(1, Ordering::SeqCst);
    Ok(connections[index % connections.len()].clone())
  }

  /// Select a connection based on strategy
  fn select_connection(&self) -> Option<Arc<PooledConnection>> {
    let connections = self.connections.read();
    if connections.is_empty() {
      return None;
    }

    // Filter healthy connections
    let healthy: Vec<_> = connections.iter().filter(|c| c.is_healthy()).cloned().collect();

    let pool = if healthy.is_empty() {
      &connections[..]
    } else {
      &healthy[..]
    };

    match self.selection_strategy {
      | ConnectionSelectionStrategy::RoundRobin => {
        let index = self.round_robin_index.fetch_add(1, Ordering::SeqCst);
        Some(pool[index % pool.len()].clone())
      }
      | ConnectionSelectionStrategy::LeastPending => pool.iter().min_by_key(|c| c.pending_count()).cloned(),
    }
  }

  /// Get pool statistics
  pub fn stats(&self) -> PoolStats {
    let connections = self.connections.read();
    let total = connections.len();
    let healthy = connections.iter().filter(|c| c.is_healthy()).count();
    let total_pending: u64 = connections.iter().map(|c| c.pending_count()).sum();

    PoolStats {
      total_connections: total,
      healthy_connections: healthy,
      unhealthy_connections: total - healthy,
      total_pending_requests: total_pending,
      max_connections: self.config.max_connections,
      min_connections: self.config.min_connections,
    }
  }

  /// Get the server address
  pub fn address(&self) -> String {
    format!("{}:{}", self.host, self.port)
  }

  /// Background maintenance loop
  async fn maintenance_loop(self: Arc<Self>) {
    let mut interval = tokio::time::interval(self.config.health_check_interval);

    loop {
      interval.tick().await;

      // Remove unhealthy connections that have been unhealthy for too long
      self.cleanup_unhealthy();

      // Remove excess idle connections
      self.cleanup_idle();

      // Ensure minimum connections
      if let Err(e) = self.ensure_min_connections().await {
        warn!("Failed to maintain minimum connections: {:?}", e);
      }
    }
  }

  /// Remove connections that have been unhealthy
  fn cleanup_unhealthy(&self) {
    let mut connections = self.connections.write();
    let before = connections.len();
    let min_connections = self.config.min_connections;

    // Count how many we can remove while staying above minimum
    let unhealthy_count = connections.iter().filter(|c| !c.is_healthy()).count();
    let can_remove = before.saturating_sub(min_connections).min(unhealthy_count);

    if can_remove > 0 {
      let mut removed = 0;
      connections.retain(|c| {
        if !c.is_healthy() && removed < can_remove {
          removed += 1;
          false
        } else {
          true
        }
      });

      debug!(
        "Removed {} unhealthy connections from pool {}:{}",
        removed, self.host, self.port
      );
    }
  }

  /// Remove excess idle connections
  fn cleanup_idle(&self) {
    let mut connections = self.connections.write();
    let before = connections.len();
    let min_connections = self.config.min_connections;
    let idle_timeout = self.config.idle_timeout;

    // Only cleanup if we're above minimum
    if before <= min_connections {
      return;
    }

    // Count how many idle connections we can remove
    let idle_count = connections
      .iter()
      .filter(|c| c.idle_duration() >= idle_timeout && c.pending_count() == 0)
      .count();
    let can_remove = before.saturating_sub(min_connections).min(idle_count);

    if can_remove > 0 {
      let mut removed = 0;
      connections.retain(|c| {
        if c.idle_duration() >= idle_timeout && c.pending_count() == 0 && removed < can_remove {
          removed += 1;
          false
        } else {
          true
        }
      });

      debug!(
        "Removed {} idle connections from pool {}:{}",
        removed, self.host, self.port
      );
    }
  }
}

/// Pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
  pub total_connections: usize,
  pub healthy_connections: usize,
  pub unhealthy_connections: usize,
  pub total_pending_requests: u64,
  pub max_connections: usize,
  pub min_connections: usize,
}

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

  #[test]
  fn test_default_config() {
    let config = ConnectionPoolConfig::default();
    assert_eq!(config.min_connections, 1);
    assert_eq!(config.max_connections, 1);
  }

  #[test]
  fn test_gateway_config() {
    let config = ConnectionPoolConfig::for_gateway();
    assert_eq!(config.min_connections, 2);
    assert_eq!(config.max_connections, 10);
  }

  #[test]
  fn test_pool_config_with_size() {
    let config = ConnectionPoolConfig::with_pool_size(4, 20);
    assert_eq!(config.min_connections, 4);
    assert_eq!(config.max_connections, 20);
  }

  #[test]
  fn test_pooled_connection_pending() {
    // We can't test without a real connection, but we can test the atomic operations
    let pending = AtomicU64::new(0);
    pending.fetch_add(1, Ordering::SeqCst);
    assert_eq!(pending.load(Ordering::SeqCst), 1);
    pending.fetch_sub(1, Ordering::SeqCst);
    assert_eq!(pending.load(Ordering::SeqCst), 0);
  }

  #[test]
  fn test_pool_stats() {
    let stats = PoolStats {
      total_connections: 5,
      healthy_connections: 4,
      unhealthy_connections: 1,
      total_pending_requests: 10,
      max_connections: 10,
      min_connections: 2,
    };

    assert_eq!(stats.total_connections, 5);
    assert_eq!(stats.healthy_connections, 4);
  }
}