rustytorch/crates/production/rtx-serving-api/src/load_balancer.rs

//! Load balancer and auto-scaling implementation

use crate::error::{ApiError, ApiResult};
use crate::grpc::ScalingAction;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

/// Load balancing strategy
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum BalancingStrategy {
    RoundRobin,
    LeastConnections,
    WeightedRoundRobin,
    Random,
    IpHash,
}

/// Load balancer configuration
#[derive(Debug, Clone)]
pub struct LoadBalancerConfig {
    pub strategy: BalancingStrategy,
    pub health_check_interval: Option<Duration>,
    pub connection_timeout: Duration,
    pub request_timeout: Duration,
}

impl Default for LoadBalancerConfig {
    fn default() -> Self {
        Self {
            strategy: BalancingStrategy::RoundRobin,
            health_check_interval: Some(Duration::from_secs(30)),
            connection_timeout: Duration::from_secs(5),
            request_timeout: Duration::from_secs(30),
        }
    }
}

impl LoadBalancerConfig {
    #[must_use]
    pub fn with_strategy(strategy: BalancingStrategy) -> Self {
        Self {
            strategy,
            ..Default::default()
        }
    }

    pub fn enable_health_checks(&mut self, interval: Duration) {
        self.health_check_interval = Some(interval);
    }
}

/// Backend server health status
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum BackendHealth {
    Healthy,
    Unhealthy,
    Unknown,
}

/// Backend server representation
#[derive(Debug, Clone)]
pub struct Backend {
    name: String,
    url: String,
    weight: usize,
    active_connections: usize,
    health: BackendHealth,
    last_health_check: Option<Instant>,
}

impl Backend {
    #[must_use]
    pub fn new(name: &str, url: &str) -> Self {
        Self {
            name: name.to_string(),
            url: url.to_string(),
            weight: 1,
            active_connections: 0,
            health: BackendHealth::Unknown,
            last_health_check: None,
        }
    }

    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn set_weight(&mut self, weight: usize) {
        self.weight = weight;
    }

    pub fn set_active_connections(&mut self, count: usize) {
        self.active_connections = count;
    }

    pub fn mark_unhealthy(&mut self) {
        self.health = BackendHealth::Unhealthy;
    }

    #[must_use]
    pub fn is_healthy(&self) -> bool {
        self.health == BackendHealth::Healthy
    }
}

/// Load balancer for distributing requests
#[derive(Debug)]
pub struct LoadBalancer {
    config: LoadBalancerConfig,
    backends: Arc<RwLock<Vec<Backend>>>,
    current_index: Arc<RwLock<usize>>,
}

impl LoadBalancer {
    #[must_use]
    pub fn new(config: LoadBalancerConfig) -> Self {
        Self {
            config,
            backends: Arc::new(RwLock::new(Vec::new())),
            current_index: Arc::new(RwLock::new(0)),
        }
    }

    #[must_use]
    pub fn strategy(&self) -> BalancingStrategy {
        self.config.strategy
    }

    #[must_use]
    pub fn num_backends(&self) -> usize {
        futures::executor::block_on(async { self.backends.read().await.len() })
    }

    #[must_use]
    pub fn is_ready(&self) -> bool {
        self.num_backends() > 0
    }

    pub async fn add_backend(&mut self, backend: Backend) -> ApiResult<()> {
        let mut backends = self.backends.write().await;
        backends.push(backend);
        Ok(())
    }

    #[must_use]
    pub fn has_backend(&self, name: &str) -> bool {
        futures::executor::block_on(async {
            self.backends.read().await.iter().any(|b| b.name == name)
        })
    }

    #[must_use]
    pub fn has_healthy_backend(&self, name: &str) -> bool {
        futures::executor::block_on(async {
            self.backends
                .read()
                .await
                .iter()
                .any(|b| b.name == name && b.is_healthy())
        })
    }

    pub async fn next_backend(&self) -> ApiResult<Backend> {
        let backends = self.backends.read().await;

        if backends.is_empty() {
            return Err(ApiError::LoadBalancer("No backends available".to_string()));
        }

        match self.config.strategy {
            BalancingStrategy::RoundRobin => {
                let mut index = self.current_index.write().await;
                let backend = backends[*index % backends.len()].clone();
                *index += 1;
                Ok(backend)
            }
            BalancingStrategy::LeastConnections => {
                let backend = backends
                    .iter()
                    .filter(|b| b.is_healthy())
                    .min_by_key(|b| b.active_connections)
                    .ok_or_else(|| ApiError::LoadBalancer("No healthy backends".to_string()))?
                    .clone();
                Ok(backend)
            }
            BalancingStrategy::WeightedRoundRobin => {
                let mut index = self.current_index.write().await;

                // Build weighted list
                let mut weighted_backends = Vec::new();
                for backend in backends.iter() {
                    for _ in 0..backend.weight {
                        weighted_backends.push(backend.clone());
                    }
                }

                if weighted_backends.is_empty() {
                    return Err(ApiError::LoadBalancer("No weighted backends".to_string()));
                }

                let backend = weighted_backends[*index % weighted_backends.len()].clone();
                *index += 1;
                Ok(backend)
            }
            _ => {
                // Default to round-robin for other strategies
                let mut index = self.current_index.write().await;
                let backend = backends[*index % backends.len()].clone();
                *index += 1;
                Ok(backend)
            }
        }
    }

    pub async fn check_backend_health(&self, name: &str) -> ApiResult<BackendHealth> {
        let backends = self.backends.read().await;
        let backend = backends
            .iter()
            .find(|b| b.name == name)
            .ok_or_else(|| ApiError::LoadBalancer(format!("Backend {name} not found")))?;

        // Simulate health check
        Ok(backend.health)
    }

    pub async fn perform_health_checks(&mut self) {
        let mut backends = self.backends.write().await;

        for backend in backends.iter_mut() {
            // Simulate health check
            if backend.health == BackendHealth::Unhealthy {
                // Keep unhealthy backends marked as such
                backend.last_health_check = Some(Instant::now());
            } else {
                backend.health = BackendHealth::Healthy;
                backend.last_health_check = Some(Instant::now());
            }
        }

        // Remove persistently unhealthy backends
        backends.retain(|b| b.health != BackendHealth::Unhealthy);
    }
}

/// Auto-scaling policy
#[derive(Debug, Clone)]
pub enum ScalingPolicy {
    CpuBased { target: f64 },
    RequestRateBased { target_rps: f64 },
    Predictive,
    Custom(String),
}

/// Auto-scaler configuration
#[derive(Debug, Clone)]
pub struct AutoScalerConfig {
    pub min_instances: usize,
    pub max_instances: usize,
    pub scaling_policy: ScalingPolicy,
    pub cooldown_period: Duration,
    pub scale_up_threshold: f64,
    pub scale_down_threshold: f64,
}

impl Default for AutoScalerConfig {
    fn default() -> Self {
        Self {
            min_instances: 1,
            max_instances: 10,
            scaling_policy: ScalingPolicy::CpuBased { target: 70.0 },
            cooldown_period: Duration::from_secs(300),
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.3,
        }
    }
}

impl AutoScalerConfig {
    pub fn set_min_instances(&mut self, min: usize) {
        self.min_instances = min;
    }

    pub fn set_max_instances(&mut self, max: usize) {
        self.max_instances = max;
    }

    pub fn set_scaling_policy(&mut self, policy: ScalingPolicy) {
        self.scaling_policy = policy;
    }

    pub fn set_cooldown_period(&mut self, period: Duration) {
        self.cooldown_period = period;
    }
}

/// Metrics collector for auto-scaling decisions
#[derive(Debug)]
pub struct MetricsCollector {
    cpu_usage: Arc<RwLock<Vec<f64>>>,
    request_rate: Arc<RwLock<Vec<f64>>>,
    memory_usage: Arc<RwLock<Vec<f64>>>,
}

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

impl MetricsCollector {
    #[must_use]
    pub fn new() -> Self {
        Self {
            cpu_usage: Arc::new(RwLock::new(Vec::new())),
            request_rate: Arc::new(RwLock::new(Vec::new())),
            memory_usage: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub async fn record_cpu_usage(&self, usage: f64) {
        let mut cpu = self.cpu_usage.write().await;
        cpu.push(usage);

        // Keep only last 100 samples
        if cpu.len() > 100 {
            cpu.remove(0);
        }
    }

    pub async fn record_request_rate(&self, rate: f64) {
        let mut rps = self.request_rate.write().await;
        rps.push(rate);

        if rps.len() > 100 {
            rps.remove(0);
        }
    }

    pub async fn get_average_cpu(&self) -> f64 {
        let cpu = self.cpu_usage.read().await;
        if cpu.is_empty() {
            return 0.0;
        }
        cpu.iter().sum::<f64>() / cpu.len() as f64
    }

    pub async fn get_average_request_rate(&self) -> f64 {
        let rps = self.request_rate.read().await;
        if rps.is_empty() {
            return 0.0;
        }
        rps.iter().sum::<f64>() / rps.len() as f64
    }
}

/// Auto-scaler for dynamic instance management
#[derive(Debug)]
pub struct AutoScaler {
    config: AutoScalerConfig,
    last_scaling_time: Arc<RwLock<Option<Instant>>>,
}

impl AutoScaler {
    #[must_use]
    pub fn new(config: AutoScalerConfig) -> Self {
        Self {
            config,
            last_scaling_time: Arc::new(RwLock::new(None)),
        }
    }

    #[must_use]
    pub fn min_instances(&self) -> usize {
        self.config.min_instances
    }

    #[must_use]
    pub fn max_instances(&self) -> usize {
        self.config.max_instances
    }

    #[must_use]
    pub fn target_cpu_utilization(&self) -> f64 {
        match &self.config.scaling_policy {
            ScalingPolicy::CpuBased { target } => *target,
            _ => 70.0,
        }
    }

    pub async fn evaluate(
        &mut self,
        lb: &LoadBalancer,
        metrics: &MetricsCollector,
    ) -> ApiResult<ScalingAction> {
        // Check cooldown period
        if !self.can_scale().await {
            return Ok(ScalingAction::None);
        }

        let current_instances = lb.num_backends();

        match &self.config.scaling_policy {
            ScalingPolicy::CpuBased { target } => {
                let avg_cpu = metrics.get_average_cpu().await;

                if avg_cpu > target * self.config.scale_up_threshold {
                    if current_instances < self.config.max_instances {
                        return Ok(ScalingAction::ScaleUp(1));
                    }
                } else if avg_cpu < target * self.config.scale_down_threshold
                    && current_instances > self.config.min_instances
                {
                    return Ok(ScalingAction::ScaleDown(1));
                }
            }
            ScalingPolicy::RequestRateBased { target_rps } => {
                let avg_rps = metrics.get_average_request_rate().await;
                let rps_per_instance = avg_rps / current_instances as f64;

                if rps_per_instance > *target_rps {
                    if current_instances < self.config.max_instances {
                        return Ok(ScalingAction::ScaleUp(1));
                    }
                } else if rps_per_instance < target_rps * 0.5
                    && current_instances > self.config.min_instances
                {
                    return Ok(ScalingAction::ScaleDown(1));
                }
            }
            ScalingPolicy::Predictive => {
                // Simplified predictive scaling
                let cpu_history = metrics.cpu_usage.read().await;
                if cpu_history.len() >= 5 {
                    let recent = &cpu_history[cpu_history.len() - 5..];
                    let trend: f64 = recent.windows(2).map(|w| w[1] - w[0]).sum::<f64>() / 4.0;

                    if trend > 5.0 && current_instances < self.config.max_instances {
                        return Ok(ScalingAction::ScaleUp(1));
                    }
                }
            }
            _ => {}
        }

        Ok(ScalingAction::None)
    }

    pub async fn apply_scaling(&mut self, action: ScalingAction) -> ApiResult<()> {
        if action != ScalingAction::None {
            let mut last_time = self.last_scaling_time.write().await;
            *last_time = Some(Instant::now());
        }
        Ok(())
    }

    async fn can_scale(&self) -> bool {
        let last_time = self.last_scaling_time.read().await;

        match *last_time {
            Some(time) => time.elapsed() >= self.config.cooldown_period,
            None => true,
        }
    }
}