prock/platform/

linux.rs

//! Linux implementation using /proc filesystem.
//!
//! Reads `/proc/[pid]/stat` for CPU times (~10µs per call).
//! Reads `/proc/[pid]/task/[pid]/children` for child discovery (~10µs per call).

#![allow(unsafe_code)]
#![allow(rustdoc::invalid_html_tags)]

use super::CpuTime;
use std::fs;

/// Disk I/O statistics for a process.
#[derive(Debug, Clone, Copy, Default)]
pub struct DiskIo {
    /// Total bytes read from disk
    pub read_bytes: u64,
    /// Total bytes written to disk
    pub write_bytes: u64,
}

/// Combined stats from reading multiple /proc files.
#[derive(Debug, Clone, Default)]
pub struct AllStats {
    pub cpu_time: CpuTime,
    pub memory_rss: u64,
    pub disk_io: DiskIo,
}

/// Clock ticks per second (usually 100 on Linux).
fn clock_ticks_per_sec() -> u64 {
    // SAFETY: sysconf is safe to call
    let ticks = unsafe { libc::sysconf(libc::_SC_CLK_TCK) };
    if ticks <= 0 { 100 } else { ticks as u64 }
}

/// Get CPU time for a process by reading `/proc/[pid]/stat`.
///
/// Returns the process's own CPU time (utime + stime).
/// For cumulative time including exited children, use `get_cpu_time_with_children`.
/// ~10µs per call.
pub fn get_cpu_time(pid: i32) -> Option<CpuTime> {
    get_cpu_time_inner(pid, false)
}

/// Get CPU time for a process INCLUDING time from waited-for (exited) children.
///
/// This is critical for accurate tracking of short-lived child processes.
/// Uses cutime/cstime fields which accumulate CPU from children that have exited.
/// ~10µs per call.
pub fn get_cpu_time_with_children(pid: i32) -> Option<CpuTime> {
    get_cpu_time_inner(pid, true)
}

fn get_cpu_time_inner(pid: i32, include_children: bool) -> Option<CpuTime> {
    let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;

    // Format: pid (comm) state ppid pgrp session tty_nr tpgid flags
    //         minflt cminflt majflt cmajflt utime stime cutime cstime ...
    // Fields (1-indexed): 14=utime, 15=stime, 16=cutime, 17=cstime

    // Find the closing paren of comm (process name can contain spaces/parens)
    let comm_end = stat.rfind(')')?;
    let fields: Vec<&str> = stat[comm_end + 2..].split_whitespace().collect();

    // After (comm), fields are: state(0) ppid(1) ... utime(11) stime(12) cutime(13) cstime(14)
    if fields.len() < 15 {
        return None;
    }

    let utime_ticks: u64 = fields[11].parse().ok()?;
    let stime_ticks: u64 = fields[12].parse().ok()?;

    let (user_ticks, system_ticks) = if include_children {
        let cutime_ticks: i64 = fields[13].parse().ok()?; // Can be negative
        let cstime_ticks: i64 = fields[14].parse().ok()?; // Can be negative
        (
            utime_ticks.saturating_add(cutime_ticks.max(0) as u64),
            stime_ticks.saturating_add(cstime_ticks.max(0) as u64),
        )
    } else {
        (utime_ticks, stime_ticks)
    };

    let ticks_per_sec = clock_ticks_per_sec();
    let ns_per_tick = 1_000_000_000 / ticks_per_sec;

    Some(CpuTime {
        user_ns: user_ticks * ns_per_tick,
        system_ns: system_ticks * ns_per_tick,
    })
}

/// Get memory usage (resident set size) for a process in bytes.
///
/// This is the physical RAM currently used by the process.
/// ~10µs per call.
pub fn get_memory(pid: i32) -> Option<u64> {
    get_statm(pid).map(|(rss, _vsz)| rss)
}

/// Get virtual memory size for a process in bytes.
///
/// This is the total address space mapped by the process.
/// ~10µs per call (same file read as get_memory).
pub fn get_memory_virtual(pid: i32) -> Option<u64> {
    get_statm(pid).map(|(_rss, vsz)| vsz)
}

/// Internal: get both RSS and VSZ from `/proc/[pid]/statm`.
fn get_statm(pid: i32) -> Option<(u64, u64)> {
    let statm = fs::read_to_string(format!("/proc/{pid}/statm")).ok()?;
    let fields: Vec<&str> = statm.split_whitespace().collect();

    // Field 0 is size (VSZ) in pages, field 1 is resident (RSS) in pages
    if fields.len() < 2 {
        return None;
    }

    let vsz_pages: u64 = fields[0].parse().ok()?;
    let rss_pages: u64 = fields[1].parse().ok()?;

    // SAFETY: sysconf is always safe to call. Returns -1 on error.
    let page_size_raw = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
    // Default to 4KB if sysconf fails (very unlikely on Linux)
    let page_size = if page_size_raw <= 0 {
        4096
    } else {
        page_size_raw as u64
    };

    Some((rss_pages * page_size, vsz_pages * page_size))
}

/// Get disk I/O statistics for a process.
///
/// Reads `/proc/[pid]/io` for cumulative disk bytes read/written.
/// ~10µs per call.
pub fn get_disk_io(pid: i32) -> Option<DiskIo> {
    let io = fs::read_to_string(format!("/proc/{pid}/io")).ok()?;

    let mut read_bytes = 0u64;
    let mut write_bytes = 0u64;

    // Format is "key: value" lines
    // We want read_bytes and write_bytes (not rchar/wchar which include page cache)
    for line in io.lines() {
        if let Some(value) = line.strip_prefix("read_bytes: ") {
            read_bytes = value.parse().unwrap_or(0);
        } else if let Some(value) = line.strip_prefix("write_bytes: ") {
            write_bytes = value.parse().unwrap_or(0);
        }
    }

    Some(DiskIo {
        read_bytes,
        write_bytes,
    })
}

/// Get all stats (CPU, memory, disk I/O) by reading multiple /proc files.
///
/// On Linux this reads 3 files: `/proc/[pid]/stat`, statm, and io.
/// ~30µs per call (vs ~30µs for 3 separate calls - no savings on Linux).
pub fn get_all_stats(pid: i32) -> Option<AllStats> {
    let cpu_time = get_cpu_time(pid)?;
    let memory_rss = get_memory(pid)?;
    let disk_io = get_disk_io(pid).unwrap_or_default();

    Some(AllStats {
        cpu_time,
        memory_rss,
        disk_io,
    })
}

/// Check if a process exists.
pub fn process_exists(pid: i32) -> bool {
    std::path::Path::new(&format!("/proc/{pid}")).exists()
}

/// Get direct child PIDs of a process.
///
/// Reads `/proc/[pid]/task/[pid]/children` which contains space-separated child PIDs.
/// ~10µs per call.
pub fn get_children(pid: i32) -> Vec<i32> {
    fs::read_to_string(format!("/proc/{pid}/task/{pid}/children"))
        .unwrap_or_default()
        .split_whitespace()
        .filter_map(|s| s.parse().ok())
        .collect()
}

/// Get parent PID of a process.
///
/// Reads `/proc/[pid]/stat` and extracts ppid field.
/// ~10µs per call.
pub fn get_ppid(pid: i32) -> Option<i32> {
    let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;

    // Find the closing paren of comm (process name can contain spaces/parens)
    let comm_end = stat.rfind(')')?;
    let fields: Vec<&str> = stat[comm_end + 2..].split_whitespace().collect();

    // After (comm), fields are: state(0) ppid(1) ...
    if fields.len() < 2 {
        return None;
    }

    fields[1].parse().ok()
}

/// Get process start time as seconds since Unix epoch.
///
/// Reads `/proc/[pid]/stat` field 22 (starttime in clock ticks since boot),
/// then converts to absolute time using system boot time.
/// ~10µs per call.
pub fn get_start_time(pid: i32) -> Option<i64> {
    let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;

    // Find the closing paren of comm (process name can contain spaces/parens)
    let comm_end = stat.rfind(')')?;
    let fields: Vec<&str> = stat[comm_end + 2..].split_whitespace().collect();

    // After (comm), fields are: state(0) ppid(1) ... starttime(19)
    // Field 22 in 1-indexed stat file = field 19 in 0-indexed after comm
    if fields.len() < 20 {
        return None;
    }

    let starttime_ticks: u64 = fields[19].parse().ok()?;
    let ticks_per_sec = clock_ticks_per_sec();

    // Get boot time from /proc/stat
    let boot_time = get_boot_time()?;

    // Convert starttime (ticks since boot) to seconds since epoch
    let starttime_secs = starttime_ticks / ticks_per_sec;
    Some(boot_time + starttime_secs as i64)
}

/// Get system boot time as seconds since Unix epoch.
fn get_boot_time() -> Option<i64> {
    let stat = fs::read_to_string("/proc/stat").ok()?;
    for line in stat.lines() {
        if let Some(value) = line.strip_prefix("btime ") {
            return value.trim().parse().ok();
        }
    }
    None
}

/// Get the full executable path for a process.
///
/// Reads `/proc/[pid]/exe` symlink to get the absolute path.
/// ~10µs per call.
pub fn get_process_path(pid: i32) -> Option<String> {
    fs::read_link(format!("/proc/{pid}/exe"))
        .ok()
        .and_then(|p| p.to_str().map(|s| s.to_string()))
}

/// Get the kernel's internal command name for a process.
///
/// Returns the executable basename from `/proc/[pid]/comm`, truncated to 16 characters
/// (kernel's TASK_COMM_LEN - 1). Examples: "zsh", "node", "tmux", "python3.13"
///
/// **Important**: This is NOT equivalent to `ps -o comm=`, which may return:
/// - Full executable path
/// - Or modified argv[0] (e.g., "claude" for renamed node, "-bash" for login shells)
///
/// Use `get_process_path()` if you need the full executable path.
///
/// ~10µs per call (vs ~2ms for spawning ps).
pub fn get_process_comm(pid: i32) -> Option<String> {
    fs::read_to_string(format!("/proc/{pid}/comm"))
        .ok()
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
}

/// Get the controlling TTY device name for a process (like `ps -o tty=`).
///
/// Returns the TTY device name (e.g., "pts/0", "tty1") or None if the
/// process has no controlling terminal.
/// ~10µs per call.
///
/// **Note**: The TTY major number mappings (136-143 for pts, 4 for tty) are
/// based on standard Linux conventions but may vary on some distributions.
/// The fallback to `/proc/[pid]/fd/0` has a potential race condition if the
/// process's file descriptors change between the stat read and the fd/0 read,
/// though this is unlikely in practice.
pub fn get_tty(pid: i32) -> Option<String> {
    let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;

    // Find the closing paren of comm (process name can contain spaces/parens)
    let comm_end = stat.rfind(')')?;
    let fields: Vec<&str> = stat[comm_end + 2..].split_whitespace().collect();

    // After (comm), fields are: state(0) ppid(1) pgrp(2) session(3) tty_nr(4) ...
    if fields.len() < 5 {
        return None;
    }

    let tty_nr: i32 = fields[4].parse().ok()?;

    // 0 means no controlling terminal
    if tty_nr == 0 {
        return None;
    }

    // Convert tty_nr to device name
    // tty_nr encodes major and minor device numbers:
    // major = (tty_nr >> 8) & 0xff
    // minor = (tty_nr & 0xff) | ((tty_nr >> 12) & 0xfff00)
    let major = ((tty_nr >> 8) & 0xff) as u32;
    let minor = ((tty_nr & 0xff) | ((tty_nr >> 12) & 0xfff00)) as u32;

    // Common TTY device mappings:
    // Major 136+ = pts (pseudo-terminal slave)
    // Major 4 = ttyN (virtual console)
    // Major 5 = tty, console, ptmx
    match major {
        136..=143 => Some(format!("pts/{minor}")),
        4 => Some(format!("tty{minor}")),
        _ => {
            // Try to read the link from /proc/[pid]/fd/0 as a fallback
            // This is less reliable but works for many cases
            fs::read_link(format!("/proc/{pid}/fd/0"))
                .ok()
                .and_then(|p| {
                    let path = p.to_string_lossy();
                    if path.starts_with("/dev/") {
                        Some(path.strip_prefix("/dev/")?.to_string())
                    } else {
                        None
                    }
                })
        }
    }
}

/// List all PIDs on the system.
///
/// Reads /proc directory for numeric entries.
/// ~100-200µs for ~500-1000 processes.
pub fn list_all_pids() -> Vec<i32> {
    let Ok(entries) = fs::read_dir("/proc") else {
        return Vec::new();
    };

    entries
        .filter_map(|e| e.ok())
        .filter_map(|e| e.file_name().to_str()?.parse::<i32>().ok())
        .filter(|&pid| pid > 0)
        .collect()
}

/// Info about a single process from a full system scan.
#[derive(Debug, Clone)]
pub struct ProcessInfo {
    pub pid: i32,
    pub ppid: i32,
    pub cpu_time: super::CpuTime,
    pub memory_bytes: u64,
}

/// Scan all processes on the system and return their info.
///
/// This is the "full scan" approach - get all processes in one pass,
/// then filter to the ones you need. Faster than targeted discovery
/// when you have many sessions (crossover at ~20-30 sessions).
///
/// ~300-500µs for ~500-1000 processes (vs ~50µs for 5 targeted sessions).
pub fn scan_all_processes() -> Vec<ProcessInfo> {
    let pids = list_all_pids();
    let mut result = Vec::with_capacity(pids.len());

    for pid in pids {
        // Get ppid
        let ppid = match get_ppid(pid) {
            Some(p) => p,
            None => continue, // Process died
        };

        // Get CPU time
        let cpu_time = match get_cpu_time(pid) {
            Some(c) => c,
            None => continue, // Process died
        };

        // Get memory
        let memory_bytes = get_memory(pid).unwrap_or(0);

        result.push(ProcessInfo {
            pid,
            ppid,
            cpu_time,
            memory_bytes,
        });
    }

    result
}

/// Build a HashMap of pid -> parent pid for ALL processes on the system.
///
/// This is significantly faster than spawning `ps -eo pid=,ppid=`:
/// - Direct /proc reads: ~5-10ms for 500 processes
/// - ps subprocess: ~100ms
///
/// Iterates `/proc` and reads `/proc/[pid]/stat` for each PID.
/// Returns HashMap<pid, ppid> for all running processes.
pub fn build_parent_map() -> std::collections::HashMap<i32, i32> {
    let pids = list_all_pids();
    let mut map = std::collections::HashMap::with_capacity(pids.len());

    for pid in pids {
        if let Some(ppid) = get_ppid(pid) {
            map.insert(pid, ppid);
        }
    }

    map
}

/// Get environment variables for a process.
///
/// Reads `/proc/[pid]/environ` which contains null-separated KEY=VALUE pairs.
/// Returns None if process doesn't exist or environ can't be read.
/// ~10-50µs depending on environment size.
pub fn get_process_environ(pid: i32) -> Option<std::collections::HashMap<String, String>> {
    let environ = fs::read(format!("/proc/{pid}/environ")).ok()?;

    let mut map = std::collections::HashMap::new();

    // Environment is null-separated KEY=VALUE pairs
    for entry in environ.split(|&b| b == 0) {
        if entry.is_empty() {
            continue;
        }
        if let Ok(s) = std::str::from_utf8(entry)
            && let Some((key, value)) = s.split_once('=')
        {
            map.insert(key.to_string(), value.to_string());
        }
    }

    Some(map)
}

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

    #[test]
    fn test_build_parent_map() {
        let map = build_parent_map();
        // Should have many processes
        assert!(map.len() > 10);
        // Our process should be in it
        let our_pid = std::process::id() as i32;
        assert!(map.contains_key(&our_pid));
        // Our parent should exist
        let ppid = map.get(&our_pid).unwrap();
        assert!(map.contains_key(ppid) || *ppid == 1);
    }

    #[test]
    fn test_get_cpu_time_self() {
        let pid = std::process::id() as i32;
        let cpu = get_cpu_time(pid);
        assert!(cpu.is_some());
    }

    #[test]
    fn test_get_memory_self() {
        let pid = std::process::id() as i32;
        let mem = get_memory(pid);
        assert!(mem.is_some());
        assert!(mem.unwrap() > 0);
    }

    #[test]
    fn test_get_cpu_time_invalid() {
        let cpu = get_cpu_time(999_999_999);
        assert!(cpu.is_none());
    }

    #[test]
    fn test_process_exists() {
        let pid = std::process::id() as i32;
        assert!(process_exists(pid));
        assert!(!process_exists(999_999_999));
    }

    #[test]
    fn test_get_disk_io_self() {
        let pid = std::process::id() as i32;
        let io = get_disk_io(pid);
        assert!(io.is_some());
        // Disk I/O may be 0 if binary is cached, just verify we got a result
        let _io = io.unwrap();
    }

    #[test]
    fn test_get_all_stats_self() {
        let pid = std::process::id() as i32;
        let stats = get_all_stats(pid);
        assert!(stats.is_some());
        let stats = stats.unwrap();

        // Memory should be non-zero
        assert!(stats.memory_rss > 0, "Memory should be non-zero");

        // Disk I/O may be 0 if binary is cached - just verify struct is populated
    }

    #[test]
    fn test_get_all_stats_invalid_pid() {
        let stats = get_all_stats(999_999_999);
        assert!(stats.is_none());
    }
}