rustytorch/crates/production/rtx-streaming/src/connection_manager.rs

//! # Connection Manager
//!
//! Manages streaming connection lifecycle, pooling, and health monitoring
//! for high-performance concurrent connections.

use crate::{StreamingError, StreamingResult};
use crossbeam::channel::{self, Receiver, Sender};
use dashmap::DashMap;
use parking_lot::RwLock as SyncRwLock;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, RwLock};
use uuid::Uuid;

/// Connection manager handling lifecycle and pooling
#[derive(Debug)]
pub struct ConnectionManager {
    /// Pool of available connections
    connection_pool: Arc<ConnectionPool>,

    /// Active connection registry
    active_connections: Arc<DashMap<Uuid, ManagedConnection>>,

    /// Connection health monitor
    health_monitor: Arc<HealthMonitor>,

    /// Configuration
    config: ConnectionConfig,

    /// Manager state
    state: Arc<RwLock<ManagerState>>,
}

/// Connection pool for efficient reuse
#[derive(Debug)]
pub struct ConnectionPool {
    /// Available connections
    available: Arc<Mutex<Vec<PooledConnection>>>,

    /// Pool configuration
    config: PoolConfig,

    /// Pool statistics
    stats: Arc<SyncRwLock<PoolStats>>,
}

/// Individual pooled connection
#[derive(Debug)]
pub struct PooledConnection {
    /// Connection identifier
    pub id: Uuid,

    /// Creation timestamp
    pub created_at: Instant,

    /// Last used timestamp
    pub last_used: Instant,

    /// Usage count
    pub usage_count: u64,

    /// Connection state
    pub state: PoolConnectionState,
}

/// Managed connection with monitoring
#[derive(Debug)]
pub struct ManagedConnection {
    /// Connection ID
    pub id: Uuid,

    /// Client identifier
    pub client_id: String,

    /// Connection establishment time
    pub established_at: Instant,

    /// Last activity timestamp
    pub last_activity: Instant,

    /// Bytes sent
    pub bytes_sent: u64,

    /// Bytes received
    pub bytes_received: u64,

    /// Connection health status
    pub health: ConnectionHealth,

    /// Performance metrics
    pub metrics: ConnectionMetrics,
}

/// Health monitoring for connections
#[derive(Debug)]
pub struct HealthMonitor {
    /// Health check interval
    check_interval: Duration,

    /// Unhealthy connection threshold
    unhealthy_threshold: Duration,

    /// Health check sender
    health_sender: Sender<HealthCheckCommand>,

    /// Health check receiver
    health_receiver: Arc<Mutex<Receiver<HealthCheckCommand>>>,

    /// Monitor state
    state: Arc<RwLock<MonitorState>>,
}

/// Connection configuration
#[derive(Debug, Clone)]
pub struct ConnectionConfig {
    /// Maximum concurrent connections
    pub max_connections: usize,

    /// Connection timeout
    pub connection_timeout: Duration,

    /// Keep-alive interval
    pub keep_alive_interval: Duration,

    /// Maximum idle time
    pub max_idle_time: Duration,

    /// Buffer size per connection
    pub buffer_size: usize,
}

/// Pool configuration
#[derive(Debug, Clone)]
pub struct PoolConfig {
    /// Initial pool size
    pub initial_size: usize,

    /// Maximum pool size
    pub max_size: usize,

    /// Connection lifetime
    pub connection_lifetime: Duration,

    /// Pool cleanup interval
    pub cleanup_interval: Duration,
}

/// Pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Total connections created
    pub total_created: u64,

    /// Total connections destroyed
    pub total_destroyed: u64,

    /// Current pool size
    pub current_size: usize,

    /// Pool hit rate
    pub hit_rate: f64,

    /// Average connection lifetime
    pub avg_lifetime: Duration,
}

/// Connection performance metrics
#[derive(Debug, Clone)]
pub struct ConnectionMetrics {
    /// Average response time
    pub avg_response_time: Duration,

    /// Message throughput
    pub message_throughput: f64,

    /// Error rate
    pub error_rate: f64,

    /// Last error timestamp
    pub last_error: Option<Instant>,
}

/// Connection health status
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionHealth {
    /// Connection is healthy
    Healthy,

    /// Connection is degraded
    Degraded,

    /// Connection is unhealthy
    Unhealthy,

    /// Connection is failed
    Failed,
}

/// Pool connection state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PoolConnectionState {
    /// Available for use
    Available,

    /// Currently in use
    InUse,

    /// Being validated
    Validating,

    /// Marked for cleanup
    Cleanup,
}

/// Manager operational state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ManagerState {
    /// Manager is starting
    Starting,

    /// Manager is active
    Active,

    /// Manager is shutting down
    Shutdown,

    /// Manager is stopped
    Stopped,
}

/// Monitor state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MonitorState {
    /// Monitor is active
    Active,

    /// Monitor is paused
    Paused,

    /// Monitor is stopped
    Stopped,
}

/// Health check commands
#[derive(Debug)]
pub enum HealthCheckCommand {
    /// Check specific connection
    CheckConnection(Uuid),

    /// Check all connections
    CheckAll,

    /// Stop monitoring
    Stop,
}

impl ConnectionManager {
    /// Create a new connection manager
    pub async fn new(pool_size: usize) -> StreamingResult<Self> {
        let config = ConnectionConfig {
            max_connections: 1000,
            connection_timeout: Duration::from_secs(30),
            keep_alive_interval: Duration::from_secs(10),
            max_idle_time: Duration::from_secs(300),
            buffer_size: 8_192,
        };

        let pool_config = PoolConfig {
            initial_size: pool_size / 2,
            max_size: pool_size,
            connection_lifetime: Duration::from_secs(3600),
            cleanup_interval: Duration::from_secs(60),
        };

        let connection_pool = Arc::new(ConnectionPool::new(pool_config).await?);
        let health_monitor = Arc::new(HealthMonitor::new().await?);

        Ok(Self {
            connection_pool,
            active_connections: Arc::new(DashMap::new()),
            health_monitor,
            config,
            state: Arc::new(RwLock::new(ManagerState::Starting)),
        })
    }

    /// Start the connection manager
    pub async fn start(&self) -> StreamingResult<()> {
        *self.state.write().await = ManagerState::Active;

        // Initialize connection pool
        self.connection_pool.initialize().await?;

        // Start health monitoring
        self.health_monitor.start().await?;

        Ok(())
    }

    /// Create a new connection
    pub async fn create_connection(&self, client_id: &str) -> StreamingResult<Uuid> {
        // Check if we're at capacity
        if self.active_connections.len() >= self.config.max_connections {
            return Err(StreamingError::Connection(
                "Connection limit reached".to_string(),
            ));
        }

        // Try to get connection from pool
        let pooled_connection = self.connection_pool.acquire().await?;

        let connection_id = pooled_connection.id;

        // Create managed connection
        let managed_connection = ManagedConnection {
            id: connection_id,
            client_id: client_id.to_string(),
            established_at: Instant::now(),
            last_activity: Instant::now(),
            bytes_sent: 0,
            bytes_received: 0,
            health: ConnectionHealth::Healthy,
            metrics: ConnectionMetrics {
                avg_response_time: Duration::ZERO,
                message_throughput: 0.0,
                error_rate: 0.0,
                last_error: None,
            },
        };

        // Register connection
        self.active_connections
            .insert(connection_id, managed_connection);

        // Start monitoring this connection
        self.health_monitor
            .monitor_connection(connection_id)
            .await?;

        Ok(connection_id)
    }

    /// Close a connection
    pub async fn close_connection(&self, connection_id: Uuid) -> StreamingResult<()> {
        if let Some((_, connection)) = self.active_connections.remove(&connection_id) {
            // Return to pool if still healthy
            if connection.health == ConnectionHealth::Healthy {
                self.connection_pool.release(connection_id).await?;
            }

            // Stop monitoring
            self.health_monitor.stop_monitoring(connection_id).await?;
        }

        Ok(())
    }

    /// Get connection statistics
    pub async fn get_connection_stats(&self) -> StreamingResult<ConnectionStats> {
        let pool_stats = self.connection_pool.get_stats().await?;

        Ok(ConnectionStats {
            active_connections: self.active_connections.len(),
            pool_size: pool_stats.current_size,
            pool_hit_rate: pool_stats.hit_rate,
            total_created: pool_stats.total_created,
            total_destroyed: pool_stats.total_destroyed,
        })
    }

    /// Shutdown connection manager
    pub async fn shutdown(&self) -> StreamingResult<()> {
        *self.state.write().await = ManagerState::Shutdown;

        // Close all active connections
        let connection_ids: Vec<Uuid> = self
            .active_connections
            .iter()
            .map(|entry| *entry.key())
            .collect();

        for connection_id in connection_ids {
            self.close_connection(connection_id).await?;
        }

        // Stop health monitoring
        self.health_monitor.stop().await?;

        // Cleanup connection pool
        self.connection_pool.cleanup().await?;

        *self.state.write().await = ManagerState::Stopped;

        Ok(())
    }
}

impl ConnectionPool {
    /// Create a new connection pool
    pub async fn new(config: PoolConfig) -> StreamingResult<Self> {
        let pool = Self {
            available: Arc::new(Mutex::new(Vec::with_capacity(config.max_size))),
            config,
            stats: Arc::new(SyncRwLock::new(PoolStats {
                total_created: 0,
                total_destroyed: 0,
                current_size: 0,
                hit_rate: 0.0,
                avg_lifetime: Duration::ZERO,
            })),
        };

        Ok(pool)
    }

    /// Initialize the pool with initial connections
    pub async fn initialize(&self) -> StreamingResult<()> {
        let mut available = self.available.lock().await;

        for _ in 0..self.config.initial_size {
            let connection = PooledConnection {
                id: Uuid::new_v4(),
                created_at: Instant::now(),
                last_used: Instant::now(),
                usage_count: 0,
                state: PoolConnectionState::Available,
            };

            available.push(connection);
        }

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.total_created += self.config.initial_size as u64;
            stats.current_size = self.config.initial_size;
        }

        Ok(())
    }

    /// Acquire a connection from the pool
    pub async fn acquire(&self) -> StreamingResult<PooledConnection> {
        let mut available = self.available.lock().await;

        if let Some(mut connection) = available.pop() {
            connection.last_used = Instant::now();
            connection.usage_count += 1;
            connection.state = PoolConnectionState::InUse;

            // Update hit rate
            {
                let mut stats = self.stats.write();
                stats.hit_rate = (stats.hit_rate * 0.9) + (1.0 * 0.1); // Exponential moving average
            }

            Ok(connection)
        } else {
            // No idle connection cached for reuse: create a fresh one on demand.
            //
            // `max_size` bounds how many *idle* connections this pool caches for
            // reuse (a soft sizing knob for reuse efficiency), not the total
            // number of connections the system may ever create — overall
            // concurrency is enforced by `ConnectionManager::config.max_connections`
            // in `create_connection`. Refusing to create a connection here just
            // because the idle cache is full would turn a reuse-cache limit into
            // a spurious hard connection cap and drop otherwise-valid clients.
            let connection = PooledConnection {
                id: Uuid::new_v4(),
                created_at: Instant::now(),
                last_used: Instant::now(),
                usage_count: 1,
                state: PoolConnectionState::InUse,
            };

            // Update stats
            {
                let mut stats = self.stats.write();
                stats.total_created += 1;
                if stats.current_size < self.config.max_size {
                    stats.current_size += 1;
                }
                stats.hit_rate = (stats.hit_rate * 0.9) + (0.0 * 0.1); // Miss
            }

            Ok(connection)
        }
    }

    /// Release a connection back to the pool
    pub async fn release(&self, connection_id: Uuid) -> StreamingResult<()> {
        // In a real implementation, we would find the connection and return it to the pool
        // For now, just simulate the release
        let mut available = self.available.lock().await;

        if available.len() < self.config.max_size {
            let connection = PooledConnection {
                id: connection_id,
                created_at: Instant::now(),
                last_used: Instant::now(),
                usage_count: 0,
                state: PoolConnectionState::Available,
            };

            available.push(connection);
        }

        Ok(())
    }

    /// Get pool statistics
    pub async fn get_stats(&self) -> StreamingResult<PoolStats> {
        Ok(self.stats.read().clone())
    }

    /// Cleanup expired connections
    pub async fn cleanup(&self) -> StreamingResult<()> {
        let mut available = self.available.lock().await;
        let now = Instant::now();

        available
            .retain(|conn| now.duration_since(conn.created_at) < self.config.connection_lifetime);

        Ok(())
    }
}

impl HealthMonitor {
    /// Create a new health monitor
    pub async fn new() -> StreamingResult<Self> {
        let (health_sender, health_receiver) = channel::unbounded();

        Ok(Self {
            check_interval: Duration::from_secs(30),
            unhealthy_threshold: Duration::from_secs(60),
            health_sender,
            health_receiver: Arc::new(Mutex::new(health_receiver)),
            state: Arc::new(RwLock::new(MonitorState::Stopped)),
        })
    }

    /// Start health monitoring
    pub async fn start(&self) -> StreamingResult<()> {
        *self.state.write().await = MonitorState::Active;
        Ok(())
    }

    /// Monitor a specific connection
    pub async fn monitor_connection(&self, _connection_id: Uuid) -> StreamingResult<()> {
        // In real implementation, this would start monitoring the connection
        Ok(())
    }

    /// Stop monitoring a connection
    pub async fn stop_monitoring(&self, _connection_id: Uuid) -> StreamingResult<()> {
        // In real implementation, this would stop monitoring the connection
        Ok(())
    }

    /// Stop the health monitor
    pub async fn stop(&self) -> StreamingResult<()> {
        *self.state.write().await = MonitorState::Stopped;
        self.health_sender
            .send(HealthCheckCommand::Stop)
            .map_err(|e| StreamingError::Connection(format!("Failed to stop monitor: {e}")))?;
        Ok(())
    }
}

/// Connection statistics
#[derive(Debug, Clone)]
pub struct ConnectionStats {
    /// Number of active connections
    pub active_connections: usize,

    /// Current pool size
    pub pool_size: usize,

    /// Pool hit rate
    pub pool_hit_rate: f64,

    /// Total connections created
    pub total_created: u64,

    /// Total connections destroyed
    pub total_destroyed: u64,
}

impl Default for ConnectionConfig {
    fn default() -> Self {
        Self {
            max_connections: 1000,
            connection_timeout: Duration::from_secs(30),
            keep_alive_interval: Duration::from_secs(10),
            max_idle_time: Duration::from_secs(300),
            buffer_size: 8_192,
        }
    }
}