nova-shell/stats-daemon/src/gpu.rs

use std::fs;
use std::io::Write;
use std::time::Instant;

pub struct GpuInfo {
    pub usage: u32,
    pub vram_used_gb: f64,
    pub vram_total_gb: f64,
    pub temp_c: i32,
    pub vendor: &'static str,
}

pub enum GpuBackend {
    Amd {
        card_path: String,
        hwmon_path: Option<String>,
    },
    Nvidia,
    Intel {
        card_path: String,
        device_path: String,
        hwmon_path: Option<String>,
        prev_rc6: Option<(u64, Instant)>,
    },
    None,
}

fn read_sysfs(path: &str) -> Option<String> {
    fs::read_to_string(path).ok().map(|s| s.trim().to_string())
}

fn read_sysfs_u64(path: &str) -> Option<u64> {
    read_sysfs(path)?.parse().ok()
}

fn find_hwmon(driver_name: &str) -> Option<String> {
    for i in 0..32 {
        if let Some(name) = read_sysfs(&format!("/sys/class/hwmon/hwmon{i}/name")) {
            if name == driver_name {
                return Some(format!("/sys/class/hwmon/hwmon{i}"));
            }
        }
    }
    None
}

pub fn detect_gpu() -> GpuBackend {
    // AMD: look for gpu_busy_percent exposed by the amdgpu driver
    for i in 0..8 {
        let p = format!("/sys/class/drm/card{i}/device/gpu_busy_percent");
        if fs::read_to_string(&p).is_ok() {
            let card = format!("/sys/class/drm/card{i}/device");
            let hwmon = find_hwmon("amdgpu");
            return GpuBackend::Amd {
                card_path: card,
                hwmon_path: hwmon,
            };
        }
    }
    // NVIDIA: probe nvidia-smi
    let nvidia_ok = std::process::Command::new("nvidia-smi")
        .args(["--query-gpu=name", "--format=csv,noheader"])
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false);
    if nvidia_ok {
        return GpuBackend::Nvidia;
    }
    // Intel: look for i915 or xe driver
    for i in 0..8 {
        let driver_link = format!("/sys/class/drm/card{i}/device/driver");
        if let Some(drv) = fs::read_link(&driver_link)
            .ok()
            .and_then(|p| p.file_name().map(|n| n.to_string_lossy().into_owned()))
        {
            if drv == "i915" || drv == "xe" {
                let card_path = format!("/sys/class/drm/card{i}");
                let device_path = format!("/sys/class/drm/card{i}/device");
                let hwmon = find_hwmon(&drv);
                return GpuBackend::Intel {
                    card_path,
                    device_path,
                    hwmon_path: hwmon,
                    prev_rc6: None,
                };
            }
        }
    }
    GpuBackend::None
}

fn read_amd(card: &str, hwmon: Option<&String>) -> Option<GpuInfo> {
    let usage: u32 = read_sysfs(&format!("{card}/gpu_busy_percent"))?
        .parse()
        .ok()?;
    let vram_used: u64 = read_sysfs(&format!("{card}/mem_info_vram_used"))?
        .parse()
        .ok()?;
    let vram_total: u64 = read_sysfs(&format!("{card}/mem_info_vram_total"))?
        .parse()
        .ok()?;
    let temp_c = hwmon
        .and_then(|h| read_sysfs(&format!("{h}/temp1_input")))
        .and_then(|s| s.parse::<i32>().ok())
        .map_or(0, |mc| mc / 1000);
    Some(GpuInfo {
        usage,
        vram_used_gb: vram_used as f64 / 1_073_741_824.0,
        vram_total_gb: vram_total as f64 / 1_073_741_824.0,
        temp_c,
        vendor: "amd",
    })
}

fn read_nvidia() -> Option<GpuInfo> {
    let out = std::process::Command::new("nvidia-smi")
        .args([
            "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu",
            "--format=csv,noheader,nounits",
        ])
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    let s = String::from_utf8_lossy(&out.stdout);
    let p: Vec<&str> = s.trim().split(',').map(str::trim).collect();
    if p.len() < 4 {
        return None;
    }
    Some(GpuInfo {
        usage: p[0].parse().ok()?,
        vram_used_gb: p[1].parse::<f64>().ok()? / 1024.0,
        vram_total_gb: p[2].parse::<f64>().ok()? / 1024.0,
        temp_c: p[3].parse().ok()?,
        vendor: "nvidia",
    })
}

fn read_intel(
    card: &str,
    device: &str,
    hwmon: Option<&String>,
    prev_rc6: &mut Option<(u64, Instant)>,
) -> GpuInfo {
    let usage = read_intel_usage(card, prev_rc6).unwrap_or(0);

    // VRAM - only present on discrete GPUs (Arc)
    let vram_total = read_sysfs_u64(&format!("{device}/mem_info_vram_total")).unwrap_or(0);
    let vram_used = if vram_total > 0 {
        read_sysfs_u64(&format!("{device}/mem_info_vram_used")).unwrap_or(0)
    } else {
        0
    };

    let temp_c = hwmon
        .and_then(|h| read_sysfs(&format!("{h}/temp1_input")))
        .and_then(|s| s.parse::<i32>().ok())
        .map_or(0, |mc| mc / 1000);

    GpuInfo {
        usage,
        vram_used_gb: vram_used as f64 / 1_073_741_824.0,
        vram_total_gb: vram_total as f64 / 1_073_741_824.0,
        temp_c,
        vendor: "intel",
    }
}

fn read_intel_usage(card: &str, prev_rc6: &mut Option<(u64, Instant)>) -> Option<u32> {
    // RC6 is the GPU idle state - compute usage from residency delta
    let rc6_ms = read_sysfs_u64(&format!("{card}/gt/gt0/rc6_residency_ms"))
        .or_else(|| read_sysfs_u64(&format!("{card}/power/rc6_residency_ms")))?;

    let now = Instant::now();
    let usage = if let Some((prev_ms, prev_time)) = prev_rc6.take() {
        let dt_ms = now.duration_since(prev_time).as_millis() as u64;
        if dt_ms > 0 {
            let idle_ms = rc6_ms.saturating_sub(prev_ms);
            let idle_pct = (idle_ms * 100 / dt_ms).min(100);
            (100 - idle_pct) as u32
        } else {
            0
        }
    } else {
        0 // first reading, no delta yet
    };

    *prev_rc6 = Some((rc6_ms, now));
    Some(usage)
}

pub fn emit_gpu(out: &mut impl Write, backend: &mut GpuBackend) {
    let info = match backend {
        GpuBackend::Amd {
            card_path,
            hwmon_path,
        } => read_amd(card_path, hwmon_path.as_ref()),
        GpuBackend::Nvidia => read_nvidia(),
        GpuBackend::Intel {
            card_path,
            device_path,
            hwmon_path,
            prev_rc6,
        } => Some(read_intel(
            card_path,
            device_path,
            hwmon_path.as_ref(),
            prev_rc6,
        )),
        GpuBackend::None => return,
    };
    if let Some(g) = info {
        let _ = writeln!(
            out,
            "{{\"type\":\"gpu\",\"usage\":{},\"vram_used_gb\":{:.3},\"vram_total_gb\":{:.3},\"temp_c\":{},\"vendor\":\"{}\"}}",
            g.usage, g.vram_used_gb, g.vram_total_gb, g.temp_c, g.vendor
        );
    }
}