agent-manager/src/ps.rs

//! ps: Nushell-style process table.
//!
//! Windows uses the Toolhelp32 snapshot API (pid, ppid, name, threads, cpu
//! time, working set); Unix reads /proc.

use crate::tables::{Cell, DataTable};

/// Columns: pid, ppid, name, cpu, mem, threads.
fn columns() -> Vec<String> {
    vec![
        "pid".into(),
        "ppid".into(),
        "name".into(),
        "cpu".into(),
        "mem".into(),
        "threads".into(),
    ]
}

/// Build the process table.
pub fn ps_table() -> DataTable {
    #[cfg(windows)]
    let rows = windows_rows();
    #[cfg(unix)]
    let rows = unix_rows();
    #[cfg(not(any(windows, unix)))]
    let rows = Vec::new();

    let mut rows = rows;
    rows.sort_by(|a, b| {
        a[0]
            .num
            .unwrap_or(0.0)
            .partial_cmp(&b[0].num.unwrap_or(0.0))
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    DataTable {
        columns: columns(),
        rows,
    }
}

// ---------------------------------------------------------------------------
// Windows: Toolhelp32 snapshot
// ---------------------------------------------------------------------------

#[cfg(windows)]
fn windows_rows() -> Vec<Vec<Cell>> {
    use crate::nav::format_size;
    use std::mem::size_of;
    use windows_sys::Win32::Foundation::{CloseHandle, INVALID_HANDLE_VALUE};
    use windows_sys::Win32::System::Diagnostics::ToolHelp::{
        CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W,
        TH32CS_SNAPPROCESS,
    };
    use windows_sys::Win32::System::ProcessStatus::{
        K32GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS,
    };
    use windows_sys::Win32::System::Threading::{
        GetProcessTimes, OpenProcess, PROCESS_QUERY_INFORMATION,
    };

    let mut rows: Vec<Vec<Cell>> = Vec::new();
    let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
    if snapshot == INVALID_HANDLE_VALUE {
        return rows;
    }
    let mut entry: PROCESSENTRY32W = unsafe { std::mem::zeroed() };
    entry.dwSize = size_of::<PROCESSENTRY32W>() as u32;
    let mut ok = unsafe { Process32FirstW(snapshot, &mut entry) };
    while ok != 0 {
        let pid = entry.th32ProcessID;
        let ppid = entry.th32ParentProcessID;
        let threads = entry.cntThreads;
        let mut name = String::new();
        for &unit in entry.szExeFile.iter() {
            if unit == 0 {
                break;
            }
            if let Some(c) = char::from_u32(unit as u32) {
                name.push(c);
            }
        }

        // CPU time and memory need a process handle; some system processes
        // refuse — those cells stay "—".
        let mut cpu = "—".to_string();
        let mut mem = "—".to_string();
        let handle = unsafe { OpenProcess(PROCESS_QUERY_INFORMATION, 0, pid) };
        if handle != 0 {
            let mut counters: PROCESS_MEMORY_COUNTERS = unsafe { std::mem::zeroed() };
            if unsafe {
                K32GetProcessMemoryInfo(
                    handle,
                    &mut counters,
                    size_of::<PROCESS_MEMORY_COUNTERS>() as u32,
                )
            } != 0
            {
                mem = format_size(counters.WorkingSetSize as u64);
            }
            let (mut created, mut exited, mut kernel, mut user): (
                windows_sys::Win32::Foundation::FILETIME,
                windows_sys::Win32::Foundation::FILETIME,
                windows_sys::Win32::Foundation::FILETIME,
                windows_sys::Win32::Foundation::FILETIME,
            ) = unsafe {
                (
                    std::mem::zeroed(),
                    std::mem::zeroed(),
                    std::mem::zeroed(),
                    std::mem::zeroed(),
                )
            };
            if unsafe { GetProcessTimes(handle, &mut created, &mut exited, &mut kernel, &mut user) } != 0 {
                let ticks = (filetime_u64(&kernel) + filetime_u64(&user)) as f64 / 10_000_000.0;
                cpu = format!("{ticks:.0}s");
            }
            unsafe { CloseHandle(handle) };
        }

        rows.push(vec![
            Cell::num(pid as f64, pid.to_string()),
            Cell::num(ppid as f64, ppid.to_string()),
            Cell::str(name),
            Cell::str(cpu),
            Cell::str(mem),
            Cell::num(threads as f64, threads.to_string()),
        ]);
        ok = unsafe { Process32NextW(snapshot, &mut entry) };
    }
    unsafe { CloseHandle(snapshot) };
    rows
}

#[cfg(windows)]
fn filetime_u64(ft: &windows_sys::Win32::Foundation::FILETIME) -> u64 {
    ((ft.dwHighDateTime as u64) << 32) | (ft.dwLowDateTime as u64)
}

// ---------------------------------------------------------------------------
// Unix: /proc
// ---------------------------------------------------------------------------

#[cfg(unix)]
fn unix_rows() -> Vec<Vec<Cell>> {
    use crate::nav::format_size;
    let page_size = page_size();
    let mut rows: Vec<Vec<Cell>> = Vec::new();
    let Ok(entries) = std::fs::read_dir("/proc") else {
        return rows;
    };
    for entry in entries.flatten() {
        let name = entry.file_name().to_string_lossy().to_string();
        let Ok(pid) = name.parse::<f64>() else {
            continue;
        };
        let base = entry.path();
        let Ok(stat) = std::fs::read_to_string(base.join("stat")) else {
            continue;
        };
        // The comm field is in parentheses and may contain spaces.
        let Some(close) = stat.rfind(')') else {
            continue;
        };
        let comm = stat[stat.find('(').unwrap_or(0) + 1..close].to_string();
        let rest: Vec<&str> = stat[close + 2..].split_whitespace().collect();
        if rest.len() < 22 {
            continue;
        }
        // rest[0]=state, rest[1]=ppid, rest[11]=utime, rest[12]=stime,
        // rest[17]=num_threads, rest[21]=rss (pages).
        let ppid = rest[1].parse::<f64>().unwrap_or(0.0);
        let utime = rest[11].parse::<f64>().unwrap_or(0.0);
        let stime = rest[12].parse::<f64>().unwrap_or(0.0);
        let threads = rest[17].parse::<f64>().unwrap_or(0.0);
        let rss_pages = rest[21].parse::<u64>().unwrap_or(0);
        let cpu = (utime + stime) / 100.0;
        rows.push(vec![
            Cell::num(pid, pid.to_string()),
            Cell::num(ppid, ppid.to_string()),
            Cell::str(comm),
            Cell::str(format!("{cpu:.0}s")),
            Cell::str(format_size(rss_pages * page_size)),
            Cell::num(threads, threads.to_string()),
        ]);
    }
    rows
}

#[cfg(unix)]
fn page_size() -> u64 {
    // Standard page size; RSS in pages * 4096 is close enough for display.
    4096
}

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

    #[test]
    fn ps_produces_a_table() {
        let table = ps_table();
        assert_eq!(
            table.columns,
            vec!["pid", "ppid", "name", "cpu", "mem", "threads"]
        );
        assert!(!table.rows.is_empty(), "process table should not be empty");
        // Our own pid is in the table.
        let me = std::process::id().to_string();
        assert!(
            table.rows.iter().any(|r| r[0].text == me),
            "our pid should be listed"
        );
        // Rows have as many cells as columns.
        assert!(table.rows.iter().all(|r| r.len() == 6));
    }
}