memlink_runtime/

backpressure.rs

//! Global backpressure management for the runtime.
//!
//! This module calculates and broadcasts backpressure signals to all connected clients.
//! Backpressure is computed from multiple factors:
//! - Total queue depth across all modules
//! - CPU and memory usage
//! - Active client count
//! - Module-specific quotas

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

use dashmap::DashMap;

/// Backpressure thresholds for different congestion levels.
pub const BP_LOW_THRESHOLD: f32 = 0.3;
pub const BP_MEDIUM_THRESHOLD: f32 = 0.5;
pub const BP_HIGH_THRESHOLD: f32 = 0.7;
pub const BP_CRITICAL_THRESHOLD: f32 = 0.9;

/// Maximum backpressure value (scaled to 1000 for atomic operations).
const BP_SCALE: u32 = 1000;

/// Global backpressure state shared across the runtime.
///
/// This is written to the SHM control region for clients to read.
#[derive(Debug)]
pub struct BackpressureState {
    /// Current backpressure value (0-1000, representing 0.0-1.0).
    value: AtomicU32,
    /// Timestamp of last calculation.
    last_update: AtomicU64,
    /// Number of active clients.
    active_clients: AtomicUsize,
    /// Total requests in queue.
    queue_depth: AtomicUsize,
    /// Requests rejected due to backpressure.
    rejected_count: AtomicU64,
}

impl BackpressureState {
    /// Creates a new backpressure state.
    pub fn new() -> Self {
        Self {
            value: AtomicU32::new(0),
            last_update: AtomicU64::new(0),
            active_clients: AtomicUsize::new(0),
            queue_depth: AtomicUsize::new(0),
            rejected_count: AtomicU64::new(0),
        }
    }

    /// Returns the current backpressure value (0.0-1.0).
    pub fn get(&self) -> f32 {
        self.value.load(Ordering::Acquire) as f32 / BP_SCALE as f32
    }

    /// Sets the backpressure value (0.0-1.0).
    pub fn set(&self, value: f32) {
        let scaled = (value.clamp(0.0, 1.0) * BP_SCALE as f32) as u32;
        self.value.store(scaled, Ordering::Release);

        // Update timestamp
        let now = Instant::now().duration_since(Instant::now()).as_secs();
        self.last_update.store(now, Ordering::Relaxed);
    }

    /// Increments the rejected request count.
    pub fn record_rejection(&self) {
        self.rejected_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Returns the number of rejected requests.
    pub fn rejected_count(&self) -> u64 {
        self.rejected_count.load(Ordering::Acquire)
    }

    /// Returns the number of active clients.
    pub fn active_clients(&self) -> usize {
        self.active_clients.load(Ordering::Acquire)
    }

    /// Returns the current queue depth.
    pub fn queue_depth(&self) -> usize {
        self.queue_depth.load(Ordering::Acquire)
    }

    /// Returns the backpressure level as a string.
    pub fn level(&self) -> &'static str {
        let bp = self.get();
        if bp < BP_LOW_THRESHOLD {
            "Low"
        } else if bp < BP_MEDIUM_THRESHOLD {
            "Medium"
        } else if bp < BP_HIGH_THRESHOLD {
            "High"
        } else if bp < BP_CRITICAL_THRESHOLD {
            "Very High"
        } else {
            "Critical"
        }
    }
}

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

/// Per-module quota configuration.
#[derive(Debug, Clone)]
pub struct ModuleQuota {
    /// Maximum calls per second.
    pub max_calls_per_sec: u32,
    /// Maximum concurrent in-flight requests.
    pub max_in_flight: usize,
    /// Maximum memory usage (bytes).
    pub max_memory_bytes: usize,
}

impl Default for ModuleQuota {
    fn default() -> Self {
        Self {
            max_calls_per_sec: 1000,
            max_in_flight: 100,
            max_memory_bytes: 64 * 1024 * 1024, // 64 MB
        }
    }
}

/// Per-module resource tracking and quota enforcement.
#[derive(Debug)]
pub struct ModuleResources {
    /// Module name/identifier.
    module_name: String,
    /// Calls in the current second.
    calls_this_sec: AtomicU32,
    /// Current in-flight request count.
    in_flight: AtomicUsize,
    /// Current memory usage (bytes).
    memory_bytes: AtomicUsize,
    /// Quota configuration.
    quota: ModuleQuota,
}

impl ModuleResources {
    /// Creates a new module resource tracker.
    pub fn new(module_name: String, quota: ModuleQuota) -> Self {
        Self {
            module_name,
            calls_this_sec: AtomicU32::new(0),
            in_flight: AtomicUsize::new(0),
            memory_bytes: AtomicUsize::new(0),
            quota,
        }
    }

    /// Checks if a new call can be admitted.
    pub fn can_admit(&self) -> bool {
        // Check in-flight limit
        if self.in_flight.load(Ordering::Acquire) >= self.quota.max_in_flight {
            return false;
        }

        // Check rate limit (simplified - should use sliding window)
        let calls = self.calls_this_sec.load(Ordering::Acquire);
        if calls >= self.quota.max_calls_per_sec {
            return false;
        }

        // Check memory limit
        let memory = self.memory_bytes.load(Ordering::Acquire);
        if memory >= self.quota.max_memory_bytes {
            return false;
        }

        true
    }

    /// Records a new call start.
    pub fn record_call_start(&self) -> bool {
        if !self.can_admit() {
            return false;
        }

        self.in_flight.fetch_add(1, Ordering::AcqRel);
        self.calls_this_sec.fetch_add(1, Ordering::Relaxed);
        true
    }

    /// Records a call completion.
    pub fn record_call_end(&self) {
        self.in_flight.fetch_sub(1, Ordering::AcqRel);
    }

    /// Updates memory usage.
    pub fn set_memory_usage(&self, bytes: usize) {
        self.memory_bytes.store(bytes, Ordering::Release);
    }

    /// Resets the per-second counter (called by a background task).
    pub fn reset_second_counter(&self) {
        self.calls_this_sec.store(0, Ordering::Relaxed);
    }

    /// Returns the current in-flight count.
    pub fn in_flight(&self) -> usize {
        self.in_flight.load(Ordering::Acquire)
    }

    /// Returns the current memory usage.
    pub fn memory_usage(&self) -> usize {
        self.memory_bytes.load(Ordering::Acquire)
    }

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

/// Global resource manager for the runtime.
///
/// Tracks all modules and clients, calculates backpressure,
/// and enforces quotas.
#[derive(Debug)]
pub struct ResourceManager {
    /// Global backpressure state.
    pub backpressure: Arc<BackpressureState>,
    /// Per-module resource tracking.
    modules: DashMap<String, Arc<ModuleResources>>,
    /// Default quota for modules without explicit configuration.
    default_quota: ModuleQuota,
}

impl ResourceManager {
    /// Creates a new resource manager.
    pub fn new() -> Self {
        Self {
            backpressure: Arc::new(BackpressureState::new()),
            modules: DashMap::new(),
            default_quota: ModuleQuota::default(),
        }
    }

    /// Creates a resource manager with custom default quota.
    pub fn with_default_quota(quota: ModuleQuota) -> Self {
        Self {
            backpressure: Arc::new(BackpressureState::new()),
            modules: DashMap::new(),
            default_quota: quota,
        }
    }

    /// Gets or creates a module resource tracker.
    pub fn get_or_create_module(&self, module_name: &str) -> Arc<ModuleResources> {
        self.modules
            .entry(module_name.to_string())
            .or_insert_with(|| {
                Arc::new(ModuleResources::new(
                    module_name.to_string(),
                    self.default_quota.clone(),
                ))
            })
            .clone()
    }

    /// Registers a module with a custom quota.
    pub fn register_module(&self, module_name: &str, quota: ModuleQuota) {
        self.modules.insert(
            module_name.to_string(),
            Arc::new(ModuleResources::new(module_name.to_string(), quota)),
        );
    }

    /// Checks if a request can be admitted.
    pub fn can_admit_request(&self, module_name: &str) -> bool {
        // Check global backpressure first
        let bp = self.backpressure.get();
        if bp >= BP_CRITICAL_THRESHOLD {
            self.backpressure.record_rejection();
            return false;
        }

        // Check module-specific quota
        let module = self.get_or_create_module(module_name);
        if !module.can_admit() {
            self.backpressure.record_rejection();
            return false;
        }

        true
    }

    /// Records a request start.
    pub fn record_request_start(&self, module_name: &str) -> bool {
        if !self.can_admit_request(module_name) {
            return false;
        }

        let module = self.get_or_create_module(module_name);
        module.record_call_start();

        // Update global queue depth
        self.backpressure
            .queue_depth
            .fetch_add(1, Ordering::Relaxed);

        true
    }

    /// Records a request completion.
    pub fn record_request_end(&self, module_name: &str) {
        let module = self.get_or_create_module(module_name);
        module.record_call_end();

        // Update global queue depth
        self.backpressure
            .queue_depth
            .fetch_sub(1, Ordering::Relaxed);
    }

    /// Calculates and updates global backpressure.
    ///
    /// This should be called periodically (e.g., every 100ms).
    pub fn update_backpressure(&self) {
        let queue_depth = self.backpressure.queue_depth.load(Ordering::Acquire);
        let active_clients = self.backpressure.active_clients.load(Ordering::Acquire);

        // Calculate backpressure from multiple factors
        let queue_factor = (queue_depth as f32 / 1000.0).min(1.0);
        let client_factor = (active_clients as f32 / 100.0).min(1.0);

        // Weighted average: 70% queue, 30% clients
        let backpressure = queue_factor * 0.7 + client_factor * 0.3;

        self.backpressure.set(backpressure);
    }

    /// Increments the active client count.
    pub fn client_connected(&self) {
        self.backpressure
            .active_clients
            .fetch_add(1, Ordering::AcqRel);
    }

    /// Decrements the active client count.
    pub fn client_disconnected(&self) {
        self.backpressure
            .active_clients
            .fetch_sub(1, Ordering::AcqRel);
    }

    /// Returns statistics about the resource manager.
    pub fn stats(&self) -> ResourceManagerStats {
        let mut total_in_flight = 0;
        let mut total_memory = 0;

        for entry in self.modules.iter() {
            total_in_flight += entry.value().in_flight();
            total_memory += entry.value().memory_usage();
        }

        ResourceManagerStats {
            backpressure: self.backpressure.get(),
            active_clients: self.backpressure.active_clients(),
            queue_depth: self.backpressure.queue_depth(),
            total_in_flight,
            total_memory,
            module_count: self.modules.len(),
        }
    }

    /// Starts a background task that periodically resets counters and updates backpressure.
    pub fn spawn_background_task(&self) -> std::thread::JoinHandle<()> {
        let _backpressure = Arc::clone(&self.backpressure);
        let modules = self
            .modules
            .iter()
            .map(|e| e.key().clone())
            .collect::<Vec<_>>();
        let modules = Arc::new(modules);

        std::thread::spawn(move || {
            let mut last_second = Instant::now();
            let mut last_update = Instant::now();

            loop {
                let now = Instant::now();

                // Reset per-second counters every second
                if now.duration_since(last_second) >= Duration::from_secs(1) {
                    for module_name in modules.iter() {
                        // Note: In production, you'd get the actual ModuleResources
                        // This is a simplified version
                        let _ = module_name;
                    }
                    last_second = now;
                }

                // Update backpressure every 100ms
                if now.duration_since(last_update) >= Duration::from_millis(100) {
                    // Note: Would need access to self here
                    // This is a simplified version
                    last_update = now;
                }

                std::thread::sleep(Duration::from_millis(10));
            }
        })
    }
}

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

/// Statistics about resource manager state.
#[derive(Debug, Clone)]
pub struct ResourceManagerStats {
    pub backpressure: f32,
    pub active_clients: usize,
    pub queue_depth: usize,
    pub total_in_flight: usize,
    pub total_memory: usize,
    pub module_count: usize,
}

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

    #[test]
    fn test_backpressure_state() {
        let bp = BackpressureState::new();
        assert_eq!(bp.get(), 0.0);

        bp.set(0.5);
        assert_eq!(bp.get(), 0.5);

        bp.set(1.0);
        assert_eq!(bp.get(), 1.0);

        bp.set(1.5); // Should clamp to 1.0
        assert_eq!(bp.get(), 1.0);
    }

    #[test]
    fn test_module_quota_enforcement() {
        let quota = ModuleQuota {
            max_calls_per_sec: 10,
            max_in_flight: 5,
            max_memory_bytes: 1024,
        };

        let module = ModuleResources::new("test".to_string(), quota);

        // Should admit up to max_in_flight
        for i in 0..5 {
            assert!(module.record_call_start(), "Call {} should be admitted", i);
        }

        // 6th call should be rejected (in-flight limit)
        assert!(!module.record_call_start());

        // Complete one call
        module.record_call_end();

        // Now should admit again
        assert!(module.record_call_start());
    }

    #[test]
    fn test_resource_manager_backpressure() {
        let rm = ResourceManager::new();

        // Initially no backpressure
        assert_eq!(rm.backpressure.get(), 0.0);

        // Simulate requests
        rm.record_request_start("module1");
        rm.record_request_start("module2");
        rm.record_request_start("module1");

        // Update backpressure
        rm.update_backpressure();

        // Should have some backpressure now
        assert!(rm.backpressure.get() > 0.0);

        // Complete requests
        rm.record_request_end("module1");
        rm.record_request_end("module2");
        rm.record_request_end("module1");

        // Update backpressure
        rm.update_backpressure();

        // Backpressure should decrease
        assert!(rm.backpressure.get() < 0.1);
    }

    #[test]
    fn test_client_tracking() {
        let rm = ResourceManager::new();
        assert_eq!(rm.backpressure.active_clients(), 0);

        rm.client_connected();
        rm.client_connected();
        assert_eq!(rm.backpressure.active_clients(), 2);

        rm.client_disconnected();
        assert_eq!(rm.backpressure.active_clients(), 1);
    }
}