memlink_runtime/

pool.rs

//! Module instance pooling and load balancing.
//!
//! Provides multiple instances of the same module for parallel execution
//! with load-aware request routing.

use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use dashmap::DashMap;

use crate::instance::ModuleInstance;

/// Pool configuration.
#[derive(Debug, Clone)]
pub struct PoolConfig {
    /// Minimum number of instances.
    pub min_instances: usize,
    /// Maximum number of instances.
    pub max_instances: usize,
    /// Target load factor per instance (0.0-1.0).
    pub target_load: f32,
    /// Scale up threshold (load factor).
    pub scale_up_threshold: f32,
    /// Scale down threshold (load factor).
    pub scale_down_threshold: f32,
    /// Cooldown between scaling operations.
    pub scale_cooldown: Duration,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            min_instances: 1,
            max_instances: 4,
            target_load: 0.6,
            scale_up_threshold: 0.8,
            scale_down_threshold: 0.3,
            scale_cooldown: Duration::from_secs(30),
        }
    }
}

/// Statistics for a pooled instance.
#[derive(Debug, Clone)]
pub struct InstanceStats {
    pub instance_id: usize,
    pub total_requests: u64,
    pub pending_requests: usize,
    pub avg_latency_us: f64,
    pub last_request: Option<Instant>,
    pub load_factor: f32,
}

/// A pooled module instance with load tracking.
#[derive(Debug)]
pub struct PooledInstance {
    /// Instance identifier.
    id: usize,
    /// The module instance.
    instance: ModuleInstance,
    /// Total requests handled.
    total_requests: AtomicU64,
    /// Current pending requests.
    pending_requests: AtomicUsize,
    /// Total latency for averaging (microseconds).
    total_latency_us: AtomicU64,
    /// Latency count for averaging.
    latency_count: AtomicU64,
    /// Last request timestamp.
    last_request: AtomicU64,
}

impl PooledInstance {
    /// Creates a new pooled instance.
    pub fn new(id: usize, instance: ModuleInstance) -> Self {
        let now = Instant::now().duration_since(Instant::now()).as_secs();
        Self {
            id,
            instance,
            total_requests: AtomicU64::new(0),
            pending_requests: AtomicUsize::new(0),
            total_latency_us: AtomicU64::new(0),
            latency_count: AtomicU64::new(0),
            last_request: AtomicU64::new(now),
        }
    }

    /// Records a request start.
    pub fn record_request_start(&self) {
        self.pending_requests.fetch_add(1, Ordering::AcqRel);
        self.total_requests.fetch_add(1, Ordering::Relaxed);

        let now = Instant::now().duration_since(Instant::now()).as_secs();
        self.last_request.store(now, Ordering::Relaxed);
    }

    /// Records a request completion.
    pub fn record_request_end(&self, latency_us: u64) {
        self.pending_requests.fetch_sub(1, Ordering::AcqRel);
        self.total_latency_us.fetch_add(latency_us, Ordering::Relaxed);
        self.latency_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Returns the instance ID.
    pub fn id(&self) -> usize {
        self.id
    }

    /// Returns the module instance.
    pub fn instance(&self) -> &ModuleInstance {
        &self.instance
    }

    /// Returns the load factor (0.0-1.0).
    pub fn load_factor(&self) -> f32 {
        let pending = self.pending_requests.load(Ordering::Acquire) as f32;
        (pending / 100.0).min(1.0)
    }

    /// Returns statistics.
    pub fn stats(&self) -> InstanceStats {
        let total_latency = self.total_latency_us.load(Ordering::Acquire);
        let count = self.latency_count.load(Ordering::Acquire);

        InstanceStats {
            instance_id: self.id,
            total_requests: self.total_requests.load(Ordering::Acquire),
            pending_requests: self.pending_requests.load(Ordering::Acquire),
            avg_latency_us: if count > 0 { total_latency as f64 / count as f64 } else { 0.0 },
            last_request: None, // Simplified
            load_factor: self.load_factor(),
        }
    }

    /// Returns the module instance (mutable).
    pub fn instance_mut(&mut self) -> &mut ModuleInstance {
        &mut self.instance
    }
}

/// Load balancing strategy.
#[derive(Debug, Clone, Copy, Default)]
pub enum LoadBalanceStrategy {
    /// Route to instance with lowest load.
    #[default]
    LeastLoaded,
    /// Route to instances in round-robin order.
    RoundRobin,
    /// Route randomly.
    Random,
}

/// Module instance pool with load balancing.
#[derive(Debug)]
pub struct ModulePool {
    /// Module name.
    module_name: String,
    /// Pooled instances.
    instances: DashMap<usize, Arc<PooledInstance>>,
    /// Pool configuration.
    config: PoolConfig,
    /// Load balancing strategy.
    strategy: LoadBalanceStrategy,
    /// Round-robin counter.
    rr_counter: AtomicUsize,
}

impl ModulePool {
    /// Creates a new module pool.
    pub fn new(module_name: String, config: PoolConfig) -> Self {
        Self {
            module_name,
            instances: DashMap::new(),
            config,
            strategy: LoadBalanceStrategy::default(),
            rr_counter: AtomicUsize::new(0),
        }
    }

    /// Creates a pool with a single instance.
    pub fn with_instance(module_name: String, instance: ModuleInstance) -> Self {
        let pool = Self::new(module_name, PoolConfig::default());
        pool.add_instance(instance);
        pool
    }

    /// Adds an instance to the pool.
    pub fn add_instance(&self, instance: ModuleInstance) -> usize {
        let id = self.instances.len();
        let pooled = Arc::new(PooledInstance::new(id, instance));
        self.instances.insert(id, pooled);
        id
    }

    /// Removes an instance from the pool.
    pub fn remove_instance(&self, id: usize) -> Option<Arc<PooledInstance>> {
        if self.instances.len() <= self.config.min_instances {
            return None; // Can't go below minimum
        }
        self.instances.remove(&id).map(|(_, v)| v)
    }

    /// Selects the best instance for a new request.
    pub fn select_instance(&self) -> Option<Arc<PooledInstance>> {
        if self.instances.is_empty() {
            return None;
        }

        match self.strategy {
            LoadBalanceStrategy::LeastLoaded => self.select_least_loaded(),
            LoadBalanceStrategy::RoundRobin => self.select_round_robin(),
            LoadBalanceStrategy::Random => self.select_random(),
        }
    }

    /// Selects the instance with lowest load.
    fn select_least_loaded(&self) -> Option<Arc<PooledInstance>> {
        let mut best: Option<(usize, f32)> = None;

        for entry in self.instances.iter() {
            let load = entry.value().load_factor();
            if best.is_none() || load < best.unwrap().1 {
                best = Some((*entry.key(), load));
            }
        }

        best.and_then(|(id, _)| self.instances.get(&id).map(|e| e.clone()))
    }

    /// Selects instance in round-robin order.
    fn select_round_robin(&self) -> Option<Arc<PooledInstance>> {
        let count = self.instances.len();
        if count == 0 {
            return None;
        }

        let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed) % count;
        self.instances.get(&idx).map(|e| e.clone())
    }

    /// Selects a random instance.
    fn select_random(&self) -> Option<Arc<PooledInstance>> {
        let count = self.instances.len();
        if count == 0 {
            return None;
        }

        // Simple hash-based selection
        use std::collections::hash_map::RandomState;
        use std::hash::{BuildHasher, Hasher};

        let hasher = RandomState::new().build_hasher();
        let idx = hasher.finish() as usize % count;
        self.instances.get(&idx).map(|e| e.clone())
    }

    /// Returns the average load factor across all instances.
    pub fn avg_load(&self) -> f32 {
        let mut total = 0.0;
        let count = self.instances.len();

        if count == 0 {
            return 0.0;
        }

        for entry in self.instances.iter() {
            total += entry.value().load_factor();
        }

        total / count as f32
    }

    /// Checks if the pool should scale up.
    pub fn should_scale_up(&self) -> bool {
        if self.instances.len() >= self.config.max_instances {
            return false;
        }

        let avg_load = self.avg_load();
        avg_load >= self.config.scale_up_threshold
    }

    /// Checks if the pool should scale down.
    pub fn should_scale_down(&self) -> bool {
        if self.instances.len() <= self.config.min_instances {
            return false;
        }

        let avg_load = self.avg_load();
        avg_load <= self.config.scale_down_threshold
    }

    /// Returns pool statistics.
    pub fn stats(&self) -> PoolStats {
        let mut total_requests = 0;
        let mut total_pending = 0;
        let mut instance_stats = Vec::new();

        for entry in self.instances.iter() {
            let stats = entry.value().stats();
            total_requests += stats.total_requests;
            total_pending += stats.pending_requests;
            instance_stats.push(stats);
        }

        PoolStats {
            module_name: self.module_name.clone(),
            instance_count: self.instances.len(),
            total_requests,
            total_pending,
            avg_load: self.avg_load(),
            instances: instance_stats,
        }
    }

    /// Returns the module name.
    pub fn module_name(&self) -> &str {
        &self.module_name
    }

    /// Returns the number of instances.
    pub fn instance_count(&self) -> usize {
        self.instances.len()
    }

    /// Returns all instances.
    pub fn all_instances(&self) -> Vec<Arc<PooledInstance>> {
        self.instances.iter().map(|e| e.value().clone()).collect()
    }
}

/// Pool statistics.
#[derive(Debug, Clone)]
pub struct PoolStats {
    pub module_name: String,
    pub instance_count: usize,
    pub total_requests: u64,
    pub total_pending: usize,
    pub avg_load: f32,
    pub instances: Vec<InstanceStats>,
}

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

    #[test]
    fn test_pool_creation() {
        let pool = ModulePool::new("test".to_string(), PoolConfig::default());
        assert_eq!(pool.instance_count(), 0);
        assert_eq!(pool.module_name(), "test");
    }

    #[test]
    fn test_pool_add_instance() {
        // Note: Can't create real ModuleInstance without loading a module
        // This test verifies the pool structure
        let pool = ModulePool::new("test".to_string(), PoolConfig::default());
        // Would add instance here if we had a mock
        assert_eq!(pool.instance_count(), 0);
    }

    #[test]
    fn test_pool_config() {
        let config = PoolConfig {
            min_instances: 2,
            max_instances: 8,
            target_load: 0.5,
            scale_up_threshold: 0.7,
            scale_down_threshold: 0.2,
            scale_cooldown: Duration::from_secs(60),
        };

        assert_eq!(config.min_instances, 2);
        assert_eq!(config.max_instances, 8);
        assert_eq!(config.target_load, 0.5);
    }
}