momentry_core/src/worker/resources.rs

use tracing::{info, warn};

pub struct SystemResources {
    pub cpu_idle_percent: f64,
    pub cpu_used_percent: f64,
    pub memory_available_mb: u64,
    pub memory_total_mb: u64,
    pub memory_used_percent: f64,
    pub gpu_available: bool,
    pub gpu_utilization: Option<f64>,
    pub gpu_memory_used_pct: Option<f64>,
}

impl SystemResources {
    pub fn check() -> Self {
        let cpu_idle = Self::get_cpu_idle();
        let (mem_available, mem_total) = Self::get_memory_info();
        let mem_used_pct = if mem_total > 0 && mem_available <= mem_total {
            ((mem_total - mem_available) as f64 / mem_total as f64) * 100.0
        } else if mem_total > 0 {
            100.0
        } else {
            0.0
        };
        let (gpu_avail, gpu_util, gpu_mem) = Self::get_gpu_info();

        Self {
            cpu_idle_percent: cpu_idle,
            cpu_used_percent: 100.0 - cpu_idle,
            memory_available_mb: mem_available,
            memory_total_mb: mem_total,
            memory_used_percent: mem_used_pct,
            gpu_available: gpu_avail,
            gpu_utilization: gpu_util,
            gpu_memory_used_pct: gpu_mem,
        }
    }

    pub fn safe_max_concurrent(&self, config_max: usize) -> usize {
        let mut limit = config_max;

        // 記憶體限制
        if self.memory_available_mb < 1000 {
            limit = limit.min(1);
            warn!(
                "Low memory ({}MB available), reducing concurrency to 1",
                self.memory_available_mb
            );
        } else if self.memory_available_mb < 2000 {
            limit = limit.min(2);
            info!(
                "Moderate memory ({}MB available), limiting concurrency to 2",
                self.memory_available_mb
            );
        } else if self.memory_available_mb < 4000 {
            limit = limit.min(3);
            info!(
                "Adequate memory ({}MB available), limiting concurrency to 3",
                self.memory_available_mb
            );
        } else if self.memory_available_mb < 8000 {
            limit = limit.min(4);
            info!(
                "Good memory ({}MB available), limiting concurrency to 4",
                self.memory_available_mb
            );
        }

        // CPU 限制
        if self.cpu_idle_percent < 15.0 {
            limit = limit.min(1);
            warn!(
                "High CPU load (idle={:.1}%), reducing concurrency to 1",
                self.cpu_idle_percent
            );
        } else if self.cpu_idle_percent < 30.0 {
            limit = limit.min(2);
            info!(
                "Moderate CPU load (idle={:.1}%), limiting concurrency to 2",
                self.cpu_idle_percent
            );
        }

        // GPU 限制：利用率 > 80% 或記憶體 > 90% 時降並發
        if let Some(util) = self.gpu_utilization {
            if util > 80.0 {
                limit = limit.min(1);
                warn!(
                    "High GPU utilization ({:.1}%), reducing concurrency to 1",
                    util
                );
            } else if util > 60.0 {
                limit = limit.min(2);
                info!(
                    "Moderate GPU utilization ({:.1}%), limiting concurrency to 2",
                    util
                );
            }
        }
        if let Some(mem) = self.gpu_memory_used_pct {
            if mem > 90.0 {
                limit = limit.min(1);
                warn!(
                    "High GPU memory usage ({:.1}%), reducing concurrency to 1",
                    mem
                );
            }
        }

        limit.max(1)
    }

    fn get_cpu_idle() -> f64 {
        use std::process::Command;
        let output = Command::new("top").args(["-l", "1", "-n", "1"]).output();
        match output {
            Ok(o) => {
                let s = String::from_utf8_lossy(&o.stdout);
                if let Some(line) = s.lines().find(|l| l.contains("idle")) {
                    if let Some(pct) = line
                        .split_whitespace()
                        .find_map(|s| s.strip_suffix("%idle"))
                    {
                        pct.trim().parse().ok().unwrap_or(50.0)
                    } else {
                        50.0
                    }
                } else {
                    50.0
                }
            }
            Err(_) => 50.0,
        }
    }

    fn get_memory_info() -> (u64, u64) {
        use std::process::Command;

        // 總記憶體
        let total = Command::new("sysctl")
            .args(["hw.memsize"])
            .output()
            .ok()
            .and_then(|o| {
                let s = String::from_utf8_lossy(&o.stdout);
                s.split_whitespace()
                    .nth(1)
                    .and_then(|v| v.parse::<u64>().ok())
            })
            .unwrap_or(0)
            / 1024
            / 1024;

        // 用 memory_pressure 取得真實可用記憶體
        // macOS 的可用記憶體 = free + inactive + speculative
        let available = Command::new("memory_pressure")
            .output()
            .ok()
            .and_then(|o| {
                let s = String::from_utf8_lossy(&o.stdout);
                // 從 "System-wide memory free percentage: XX%" 這行解析
                for line in s.lines() {
                    if line.contains("free percentage") {
                        if let Some(pct_str) = line.split(':').nth(1) {
                            let pct = pct_str.trim().trim_end_matches('%').parse::<f64>().ok()?;
                            return Some((total as f64 * pct / 100.0) as u64);
                        }
                    }
                }
                None
            })
            .unwrap_or_else(|| {
                // fallback: vm_stat
                Command::new("vm_stat")
                    .output()
                    .ok()
                    .and_then(|v| {
                        let vs = String::from_utf8_lossy(&v.stdout);
                        let mut free_pages: u64 = 0;
                        let mut inactive_pages: u64 = 0;
                        let mut speculative_pages: u64 = 0;
                        for line in vs.lines() {
                            if line.contains("Pages free:") {
                                free_pages = line
                                    .split_whitespace()
                                    .last()
                                    .and_then(|v| v.trim_end_matches('.').parse().ok())
                                    .unwrap_or(0);
                            } else if line.contains("Pages inactive:") {
                                inactive_pages = line
                                    .split_whitespace()
                                    .last()
                                    .and_then(|v| v.trim_end_matches('.').parse().ok())
                                    .unwrap_or(0);
                            } else if line.contains("Pages speculative:") {
                                speculative_pages = line
                                    .split_whitespace()
                                    .last()
                                    .and_then(|v| v.trim_end_matches('.').parse().ok())
                                    .unwrap_or(0);
                            }
                        }
                        Some(
                            (free_pages + inactive_pages + speculative_pages) * 16384 / 1024 / 1024,
                        )
                    })
                    .unwrap_or(total / 4)
            });

        (available, total)
    }

    fn get_gpu_info() -> (bool, Option<f64>, Option<f64>) {
        use std::process::Command;

        // Apple Silicon (MPS) — 用 ioreg 取 GPU 利用率
        let ioreg = Command::new("ioreg")
            .args(["-r", "-c", "AppleM2ScalerCSCDriver"])
            .output();
        if let Ok(o) = ioreg {
            let s = String::from_utf8_lossy(&o.stdout);
            if s.contains("PerformanceStatistics") {
                let util = s.lines().find_map(|l| {
                    if l.contains("GPU Utilization") {
                        l.split('=').nth(1).and_then(|v| {
                            v.trim()
                                .trim_matches('"')
                                .trim_end_matches('}')
                                .parse::<f64>()
                                .ok()
                        })
                    } else {
                        None
                    }
                });
                let mem = s.lines().find_map(|l| {
                    if l.contains("GPU Memory Utilization") {
                        l.split('=').nth(1).and_then(|v| {
                            v.trim()
                                .trim_matches('"')
                                .trim_end_matches('}')
                                .parse::<f64>()
                                .ok()
                        })
                    } else {
                        None
                    }
                });
                return (true, util, mem);
            }
        }

        // NVIDIA GPU via nvidia-smi
        let nvidia = Command::new("nvidia-smi")
            .args([
                "--query-gpu=utilization.gpu,memory.used,memory.total",
                "--format=csv,noheader,nounits",
            ])
            .output();
        if let Ok(o) = nvidia {
            if o.status.success() {
                let s = String::from_utf8_lossy(&o.stdout);
                let parts: Vec<&str> = s.trim().split(',').collect();
                if parts.len() >= 3 {
                    let util = parts[0].trim().parse::<f64>().ok();
                    let mem_used = parts[1].trim().parse::<f64>().ok();
                    let mem_total = parts[2].trim().parse::<f64>().ok();
                    let mem_pct = match (mem_used, mem_total) {
                        (Some(u), Some(t)) if t > 0.0 => Some(u / t * 100.0),
                        _ => None,
                    };
                    return (true, util, mem_pct);
                }
            }
        }

        (false, None, None)
    }
}