qrusty 0.20.6

// src/memory_monitor.rs
// Implements: SYS-0022

//! # Memory Monitor
//!
//! Monitors container/process memory usage via cgroup v2 (with v1 fallback)
//! or an explicit `QRUSTY_MEMORY_LIMIT_MB` override.  Reports usage ratio
//! and pressure state so the server can shed load (e.g. shrink hot tiers,
//! trigger payload compaction) before hitting OOM.

use std::path::{Path, PathBuf};

/// Graduated pressure level so the server can respond proportionally
/// instead of going from "fine" to "shed everything" in one step.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum PressureLevel {
    /// Below the enter-pressure threshold — no action needed.
    None,
    /// Moderate pressure (above enter threshold but below critical).
    /// Recommended actions: flush RocksDB, shrink block cache.
    Warning,
    /// High pressure (above critical threshold).
    /// Recommended actions: shed publishes, shrink hot tiers, compact.
    Critical,
}

/// Snapshot of current memory state.
#[derive(Debug, Clone)]
pub struct MemoryState {
    /// Current RSS / working-set usage in bytes.
    pub usage_bytes: u64,
    /// Configured or detected memory limit in bytes (0 = unknown).
    pub limit_bytes: u64,
    /// True when usage exceeds the pressure threshold (kept for
    /// backwards compat with the atomic flag in main.rs).
    #[allow(dead_code)]
    pub under_pressure: bool,
    /// Graduated pressure level for proportional response.
    pub pressure_level: PressureLevel,
    /// Current usage as a fraction of the limit (0.0–1.0), or 0.0 if
    /// no limit is configured.
    pub usage_ratio: f64,
}

/// File paths used for reading cgroup/proc memory info.
/// Extracted to a struct so tests can point them at temp files.
struct MemPaths {
    cgroup_v2_max: PathBuf,
    cgroup_v2_current: PathBuf,
    cgroup_v1_limit: PathBuf,
    cgroup_v1_usage: PathBuf,
    proc_statm: PathBuf,
}

impl Default for MemPaths {
    fn default() -> Self {
        Self {
            cgroup_v2_max: PathBuf::from("/sys/fs/cgroup/memory.max"),
            cgroup_v2_current: PathBuf::from("/sys/fs/cgroup/memory.current"),
            cgroup_v1_limit: PathBuf::from("/sys/fs/cgroup/memory/memory.limit_in_bytes"),
            cgroup_v1_usage: PathBuf::from("/sys/fs/cgroup/memory/memory.usage_in_bytes"),
            proc_statm: PathBuf::from("/proc/self/statm"),
        }
    }
}

/// Reads memory usage and limits from cgroup or env override.
pub struct MemoryMonitor {
    limit_bytes: u64,
    /// Fraction of limit at which pressure is entered (default 0.80).
    pressure_threshold: f64,
    /// Fraction of limit at which pressure is exited (default 0.70).
    /// Creates hysteresis so the system doesn't flap at the boundary.
    pressure_exit_threshold: f64,
    /// Fraction of limit at which critical pressure is declared (default 0.90).
    critical_threshold: f64,
    /// Paths for reading memory info (mockable in tests).
    paths: MemPaths,
    /// Sticky pressure state — once entered, stays true until usage drops
    /// below `pressure_exit_threshold`.
    was_under_pressure: std::cell::Cell<bool>,
}

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

impl MemoryMonitor {
    /// Creates a new monitor.  Tries, in order:
    /// 1. `QRUSTY_MEMORY_LIMIT_MB` environment variable (explicit override)
    /// 2. cgroup v2 `memory.max`
    /// 3. cgroup v1 `memory.limit_in_bytes`
    /// 4. Falls back to 0 (unknown — pressure detection disabled)
    pub fn new() -> Self {
        Self::with_paths(MemPaths::default())
    }

    fn with_paths(paths: MemPaths) -> Self {
        let limit_bytes = Self::detect_limit(&paths);
        let pressure_threshold: f64 = std::env::var("QRUSTY_MEMORY_PRESSURE_THRESHOLD")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.80);
        let pressure_exit_threshold: f64 = std::env::var("QRUSTY_MEMORY_PRESSURE_EXIT_THRESHOLD")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.70);
        let critical_threshold: f64 = std::env::var("QRUSTY_MEMORY_CRITICAL_THRESHOLD")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.90);

        if limit_bytes > 0 {
            tracing::info!(
                "Memory monitor: limit={} MB, pressure_threshold={:.0}%, exit={:.0}%, critical={:.0}%",
                limit_bytes / (1024 * 1024),
                pressure_threshold * 100.0,
                pressure_exit_threshold * 100.0,
                critical_threshold * 100.0
            );
        } else {
            tracing::info!(
                "Memory monitor: no cgroup limit detected, pressure detection disabled. \
                 Set QRUSTY_MEMORY_LIMIT_MB to enable."
            );
        }

        Self {
            limit_bytes,
            pressure_threshold,
            pressure_exit_threshold,
            critical_threshold,
            paths,
            was_under_pressure: std::cell::Cell::new(false),
        }
    }

    fn detect_limit(paths: &MemPaths) -> u64 {
        Self::detect_limit_with_override(paths, std::env::var("QRUSTY_MEMORY_LIMIT_MB").ok())
    }

    /// Inner implementation that accepts an optional explicit limit string
    /// so unit tests can exercise the override path without touching env vars.
    fn detect_limit_with_override(paths: &MemPaths, env_override_mb: Option<String>) -> u64 {
        // 1. Explicit override
        if let Some(mb) = env_override_mb.and_then(|v| v.parse::<u64>().ok()) {
            return mb * 1024 * 1024;
        }

        // 2. cgroup v2
        if let Some(bytes) = read_cgroup_v2_limit(&paths.cgroup_v2_max) {
            return bytes;
        }

        // 3. cgroup v1
        if let Some(bytes) = read_cgroup_v1_limit(&paths.cgroup_v1_limit) {
            return bytes;
        }

        0
    }

    /// Reads current memory usage.
    fn read_usage(&self) -> u64 {
        // cgroup v2
        if let Some(bytes) = read_bytes_from_file(&self.paths.cgroup_v2_current) {
            return bytes;
        }

        // cgroup v1
        if let Some(bytes) = read_bytes_from_file(&self.paths.cgroup_v1_usage) {
            return bytes;
        }

        // Fallback: /proc/self/statm (page count * page size)
        if let Ok(content) = std::fs::read_to_string(&self.paths.proc_statm) {
            if let Some(rss_pages) = content.split_whitespace().nth(1) {
                if let Ok(pages) = rss_pages.parse::<u64>() {
                    return pages * PAGE_SIZE;
                }
            }
        }

        0
    }

    /// Returns current memory state with hysteresis.
    ///
    /// Pressure is entered when usage exceeds `pressure_threshold` and only
    /// exited when it drops below `pressure_exit_threshold`.  This prevents
    /// the system from flapping between "normal" and "pressure" every poll
    /// when usage hovers near the threshold.
    pub fn state(&self) -> MemoryState {
        let usage_bytes = self.read_usage();
        let usage_ratio = if self.limit_bytes > 0 {
            usage_bytes as f64 / self.limit_bytes as f64
        } else {
            0.0
        };

        // Hysteresis: once under pressure, stay there until we drop below the exit threshold.
        let under_pressure = if self.limit_bytes == 0 {
            false
        } else if self.was_under_pressure.get() {
            // Already under pressure — only exit when below the exit threshold.
            usage_ratio > self.pressure_exit_threshold
        } else {
            // Not under pressure — only enter when above the enter threshold.
            usage_ratio > self.pressure_threshold
        };
        self.was_under_pressure.set(under_pressure);

        let pressure_level = if self.limit_bytes == 0 || !under_pressure {
            PressureLevel::None
        } else if usage_ratio > self.critical_threshold {
            PressureLevel::Critical
        } else {
            PressureLevel::Warning
        };

        MemoryState {
            usage_bytes,
            limit_bytes: self.limit_bytes,
            under_pressure,
            pressure_level,
            usage_ratio,
        }
    }

    /// Returns the configured memory limit in bytes.
    #[allow(dead_code)]
    pub fn limit_bytes(&self) -> u64 {
        self.limit_bytes
    }

    /// Lightweight snapshot of current memory usage and limit, suitable for
    /// the `/stats` endpoint.  Does NOT create a full monitor, does NOT log,
    /// and has no hysteresis state.
    pub fn usage_snapshot() -> (u64, u64) {
        let paths = MemPaths::default();
        let limit = Self::detect_limit(&paths);

        let usage = read_bytes_from_file(&paths.cgroup_v2_current)
            .or_else(|| read_bytes_from_file(&paths.cgroup_v1_usage))
            .unwrap_or_else(|| {
                std::fs::read_to_string(&paths.proc_statm)
                    .ok()
                    .and_then(|c| c.split_whitespace().nth(1)?.parse::<u64>().ok())
                    .map(|pages| pages * PAGE_SIZE)
                    .unwrap_or(0)
            });

        (usage, limit)
    }
}

const PAGE_SIZE: u64 = 4096;

/// Reads a single u64 value from a file (trimmed).
fn read_bytes_from_file(path: &Path) -> Option<u64> {
    std::fs::read_to_string(path)
        .ok()?
        .trim()
        .parse::<u64>()
        .ok()
}

/// Reads cgroup v2 memory.max, returning None for "max" (unlimited).
fn read_cgroup_v2_limit(path: &Path) -> Option<u64> {
    let content = std::fs::read_to_string(path).ok()?;
    let trimmed = content.trim();
    if trimmed == "max" {
        return None;
    }
    trimmed.parse::<u64>().ok()
}

/// Reads cgroup v1 memory.limit_in_bytes, returning None for very large
/// values (which cgroup v1 reports when unlimited).
fn read_cgroup_v1_limit(path: &Path) -> Option<u64> {
    let bytes = read_bytes_from_file(path)?;
    if bytes >= u64::MAX / 2 {
        return None; // cgroup v1 "unlimited" sentinel
    }
    Some(bytes)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    fn mock_paths(dir: &Path) -> MemPaths {
        MemPaths {
            cgroup_v2_max: dir.join("memory.max"),
            cgroup_v2_current: dir.join("memory.current"),
            cgroup_v1_limit: dir.join("memory.limit_in_bytes"),
            cgroup_v1_usage: dir.join("memory.usage_in_bytes"),
            proc_statm: dir.join("statm"),
        }
    }

    #[test]
    fn test_memory_monitor_new_does_not_panic() {
        let monitor = MemoryMonitor::new();
        let state = monitor.state();
        assert!(state.usage_bytes > 0 || state.limit_bytes == 0);
    }

    /// Verifies that when usage exceeds a tiny limit, pressure is reported.
    /// Uses mock files instead of env vars to avoid flakiness from parallel
    /// tests and real cgroup readings.
    #[test]
    fn test_pressure_with_explicit_limit() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // 1 MB limit, 2 MB usage → well above default 85% threshold.
        fs::write(&paths.cgroup_v2_max, "1048576\n").unwrap(); // 1 MB
        fs::write(&paths.cgroup_v2_current, "2097152\n").unwrap(); // 2 MB

        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert!(state.under_pressure);
        assert_eq!(state.limit_bytes, 1_048_576);
        assert_eq!(state.usage_bytes, 2_097_152);
    }

    #[test]
    fn test_cgroup_v2_limit_detection() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        fs::write(&paths.cgroup_v2_max, "1073741824\n").unwrap();
        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let limit = MemoryMonitor::detect_limit(&paths);
        assert_eq!(limit, 1073741824); // 1 GB
    }

    #[test]
    fn test_cgroup_v2_unlimited_returns_zero() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        fs::write(&paths.cgroup_v2_max, "max\n").unwrap();
        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let limit = MemoryMonitor::detect_limit(&paths);
        assert_eq!(limit, 0);
    }

    #[test]
    fn test_cgroup_v1_limit_detection() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // No v2 file, so falls through to v1.
        fs::write(&paths.cgroup_v1_limit, "2147483648\n").unwrap();
        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let limit = MemoryMonitor::detect_limit(&paths);
        assert_eq!(limit, 2147483648); // 2 GB
    }

    #[test]
    fn test_cgroup_v1_unlimited_returns_zero() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // cgroup v1 reports a huge number when unlimited.
        fs::write(&paths.cgroup_v1_limit, "9223372036854775808\n").unwrap();
        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let limit = MemoryMonitor::detect_limit(&paths);
        assert_eq!(limit, 0);
    }

    #[test]
    fn test_cgroup_v2_usage_reading() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        fs::write(&paths.cgroup_v2_max, "4294967296\n").unwrap(); // 4 GB limit
        fs::write(&paths.cgroup_v2_current, "1073741824\n").unwrap(); // 1 GB used

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert_eq!(state.usage_bytes, 1073741824);
        assert_eq!(state.limit_bytes, 4294967296);
        assert!(!state.under_pressure); // 25% < 80%
    }

    #[test]
    fn test_pressure_detected_at_threshold() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();
        fs::write(&paths.cgroup_v2_current, "900\n").unwrap(); // 90% > 80%

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert!(state.under_pressure);
    }

    #[test]
    fn test_no_pressure_below_threshold() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();
        fs::write(&paths.cgroup_v2_current, "500\n").unwrap(); // 50% < 80%

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert!(!state.under_pressure);
    }

    #[test]
    fn test_cgroup_v1_usage_fallback() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // No v2 files — only v1.
        fs::write(&paths.cgroup_v1_limit, "2000000000\n").unwrap();
        fs::write(&paths.cgroup_v1_usage, "500000000\n").unwrap();

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert_eq!(state.usage_bytes, 500000000);
        assert_eq!(state.limit_bytes, 2000000000);
        assert!(!state.under_pressure);
    }

    #[test]
    fn test_proc_statm_fallback() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // No cgroup files — only statm.
        // Format: size resident shared text lib data dt (in pages)
        fs::write(&paths.proc_statm, "1000 500 100 50 0 300 0\n").unwrap();

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert_eq!(state.usage_bytes, 500 * PAGE_SIZE); // RSS = 500 pages
        assert_eq!(state.limit_bytes, 0); // no limit detected
        assert!(!state.under_pressure); // no limit → no pressure
    }

    #[test]
    fn test_no_files_returns_zero_usage() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());
        // No files written — all reads fail.

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();

        assert_eq!(state.usage_bytes, 0);
        assert_eq!(state.limit_bytes, 0);
        assert!(!state.under_pressure);
    }

    #[test]
    fn test_env_override_takes_precedence() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());

        // Both cgroup files exist, but explicit override should win.
        fs::write(&paths.cgroup_v2_max, "999999\n").unwrap();
        let limit = MemoryMonitor::detect_limit_with_override(&paths, Some("42".to_string()));
        assert_eq!(limit, 42 * 1024 * 1024);
    }

    #[test]
    fn test_limit_bytes_accessor() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());
        fs::write(&paths.cgroup_v2_max, "5000\n").unwrap();

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        let monitor = MemoryMonitor::with_paths(paths);
        assert_eq!(monitor.limit_bytes(), 5000);
    }

    /// Verify the default thresholds are 80%/70%/90% (SYS-0022).
    #[test]
    fn test_default_thresholds() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());
        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();
        fs::write(&paths.cgroup_v2_current, "1\n").unwrap();

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_THRESHOLD");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_EXIT_THRESHOLD");
        std::env::remove_var("QRUSTY_MEMORY_CRITICAL_THRESHOLD");
        let monitor = MemoryMonitor::with_paths(paths);

        assert_eq!(
            monitor.pressure_threshold, 0.80,
            "enter threshold should default to 0.80"
        );
        assert_eq!(
            monitor.pressure_exit_threshold, 0.70,
            "exit threshold should default to 0.70"
        );
        assert_eq!(
            monitor.critical_threshold, 0.90,
            "critical threshold should default to 0.90"
        );
    }

    /// Verify that pressure exits at the new 70% exit threshold, not the
    /// old 75%.  At 72% (between old and new exit) pressure should persist
    /// if entered.
    // Verifies: SYS-0022
    #[test]
    fn test_hysteresis_exits_at_70_percent() {
        let dir = TempDir::new().unwrap();
        let p = dir.path();
        let paths = mock_paths(p);
        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_THRESHOLD");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_EXIT_THRESHOLD");
        std::env::remove_var("QRUSTY_MEMORY_CRITICAL_THRESHOLD");

        // Enter pressure at 85%.
        fs::write(p.join("memory.current"), "850\n").unwrap();
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();
        assert!(state.under_pressure, "should enter pressure at 85%");

        // Drop to 72% — between old exit (75%) and new exit (70%).
        // Pressure should PERSIST because we haven't dropped below 70%.
        fs::write(p.join("memory.current"), "720\n").unwrap();
        let state = monitor.state();
        assert!(
            state.under_pressure,
            "should still be under pressure at 72% (exit is 70%)"
        );

        // Drop to 69% — below the new 70% exit threshold.
        fs::write(p.join("memory.current"), "690\n").unwrap();
        let state = monitor.state();
        assert!(
            !state.under_pressure,
            "should exit pressure at 69% (below 70% exit)"
        );
    }

    /// Verify that at 79% (between old 75% enter and new 80% enter),
    /// pressure is NOT entered with the new default threshold.
    // Verifies: SYS-0022
    #[test]
    fn test_no_pressure_at_79_percent() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());
        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();
        fs::write(&paths.cgroup_v2_current, "790\n").unwrap(); // 79%

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_THRESHOLD");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();
        assert!(
            !state.under_pressure,
            "79% should not trigger pressure with 80% threshold"
        );
    }

    /// Verify that at 81%, pressure IS entered with the new 80% threshold.
    // Verifies: SYS-0022
    #[test]
    fn test_pressure_at_81_percent() {
        let dir = TempDir::new().unwrap();
        let paths = mock_paths(dir.path());
        fs::write(&paths.cgroup_v2_max, "1000\n").unwrap();
        fs::write(&paths.cgroup_v2_current, "810\n").unwrap(); // 81%

        std::env::remove_var("QRUSTY_MEMORY_LIMIT_MB");
        std::env::remove_var("QRUSTY_MEMORY_PRESSURE_THRESHOLD");
        let monitor = MemoryMonitor::with_paths(paths);
        let state = monitor.state();
        assert!(
            state.under_pressure,
            "81% should trigger pressure with 80% threshold"
        );
    }
}