diff --git a/src/collection.rs b/src/collection.rs index 93a16fa10..447af3b63 100644 --- a/src/collection.rs +++ b/src/collection.rs @@ -171,6 +171,8 @@ pub struct DataCollector { prev_idle: f64, #[cfg(target_os = "linux")] prev_non_idle: f64, + #[cfg(target_os = "linux")] + process_buffer: String, #[cfg(feature = "battery")] battery_manager: Option, @@ -224,6 +226,9 @@ impl DataCollector { gpus_total_mem: None, last_list_collection_time: last_collection_time, should_run_less_routine_tasks: true, + #[cfg(target_os = "linux")] + // TODO: Maybe pre-allocate this? I've tried this before with 16_384 bytes and it was ok? + process_buffer: String::new() } } diff --git a/src/collection/amd.rs b/src/collection/amd.rs index 18fe5d739..dd7799ca7 100644 --- a/src/collection/amd.rs +++ b/src/collection/amd.rs @@ -11,7 +11,10 @@ use std::{ use hashbrown::{HashMap, HashSet}; use super::linux::utils::is_device_awake; -use crate::{app::layout_manager::UsedWidgets, collection::memory::MemData}; +use crate::{ + app::layout_manager::UsedWidgets, + collection::{linux::utils::read_link, memory::MemData}, +}; // TODO: May be able to clean up some of these, Option for example is a bit redundant. pub struct AmdGpuData { @@ -162,19 +165,21 @@ fn diff_usage(pre: u64, cur: u64, interval: &Duration) -> u64 { } // from amdgpu_top: https://github.com/Umio-Yasuno/amdgpu_top/blob/c961cf6625c4b6d63fda7f03348323048563c584/crates/libamdgpu_top/src/stat/fdinfo/proc_info.rs#L13-L27 -fn get_amdgpu_pid_fds(pid: u32, device_path: Vec) -> Option> { +fn get_amdgpu_pid_fds(pid: u32, device_path: &[String], buffer: &mut Vec) -> Option> { let Ok(fd_list) = fs::read_dir(format!("/proc/{pid}/fd/")) else { return None; }; let valid_fds: Vec = fd_list - .filter_map(|fd_link| { - let dir_entry = fd_link.map(|fd_link| fd_link.path()).ok()?; - let link = fs::read_link(&dir_entry).ok()?; + .flatten() + .filter_map(|entry| { + let path = entry.path(); + let link = read_link(path.as_path(), buffer).ok()?; // e.g. "/dev/dri/renderD128" or "/dev/dri/card0" if device_path.iter().any(|path| link.starts_with(path)) { - dir_entry.file_name()?.to_str()?.parse::().ok() + // TODO: Should we bother parsing here? + path.file_name()?.to_str()?.parse::().ok() } else { None } @@ -188,7 +193,7 @@ fn get_amdgpu_pid_fds(pid: u32, device_path: Vec) -> Option> { } } -fn get_amdgpu_drm(device_path: &Path) -> Option> { +fn get_amdgpu_drm(device_path: &Path) -> Option> { let mut drm_devices = Vec::new(); let drm_root = device_path.join("drm"); @@ -212,7 +217,7 @@ fn get_amdgpu_drm(device_path: &Path) -> Option> { continue; } - drm_devices.push(PathBuf::from(format!("/dev/dri/{drm_name}"))); + drm_devices.push(format!("/dev/dri/{drm_name}")); } if drm_devices.is_empty() { @@ -222,8 +227,11 @@ fn get_amdgpu_drm(device_path: &Path) -> Option> { } } +// Based on https://github.com/Umio-Yasuno/amdgpu_top/blob/c961cf6625c4b6d63fda7f03348323048563c584/crates/libamdgpu_top/src/stat/fdinfo/proc_info.rs#L13-L27. fn get_amd_fdinfo(device_path: &Path) -> Option> { let mut fdinfo = HashMap::new(); + let mut buffer = Vec::new(); + const SYSTEMD_TO_SKIP: &[&str] = &["/lib/systemd", "/usr/lib/systemd"]; let drm_paths = get_amdgpu_drm(device_path)?; @@ -231,35 +239,39 @@ fn get_amd_fdinfo(device_path: &Path) -> Option> { return None; }; - let pids: Vec = proc_dir - .filter_map(|dir_entry| { - // check if pid is valid - let dir_entry = dir_entry.ok()?; - let metadata = dir_entry.metadata().ok()?; + let pids = proc_dir.filter_map(|dir_entry| { + // check if pid is valid + let dir_entry = dir_entry.ok()?; + let metadata = dir_entry.metadata().ok()?; - if !metadata.is_dir() { - return None; - } + if !metadata.is_dir() { + return None; + } - let pid = dir_entry.file_name().to_str()?.parse::().ok()?; + let pid = dir_entry.file_name().to_str()?.parse::().ok()?; - // skip init process - if pid == 1 { - return None; - } + // skip init process/systemd + if pid == 1 { + return None; + } - Some(pid) - }) - .collect(); + // TODO: We can instead refer to our already-obtained processes? I think we could maybe just do + // this with processes at the same time... + let cmdline = fs::read_to_string(format!("/proc/{pid}/cmdline")).ok()?; + if SYSTEMD_TO_SKIP.iter().any(|path| cmdline.starts_with(path)) { + return None; + } + + Some(pid) + }); for pid in pids { // collect file descriptors that point to our device renderers - let Some(fds) = get_amdgpu_pid_fds(pid, drm_paths.clone()) else { + let Some(fds) = get_amdgpu_pid_fds(pid, &drm_paths, &mut buffer) else { continue; }; let mut usage: AmdGpuProc = Default::default(); - let mut observed_ids: HashSet = HashSet::new(); for fd in fds { @@ -271,6 +283,7 @@ fn get_amd_fdinfo(device_path: &Path) -> Option> { let mut fdinfo_lines = fdinfo_data .lines() .skip_while(|l| !l.starts_with("drm-client-id")); + if let Some(id) = fdinfo_lines.next().and_then(|fdinfo_line| { const LEN: usize = "drm-client-id:\t".len(); fdinfo_line.get(LEN..)?.parse().ok() @@ -331,6 +344,7 @@ pub fn get_amd_vecs(widgets_to_harvest: &UsedWidgets, prev_time: Instant) -> Opt let mut proc_vec = Vec::with_capacity(num_gpu); let mut total_mem = 0; + // TODO: We can optimize this to do this all in one pass, rather than for loop + for loop. This reduces syscalls. for device_path in device_path_list { let device_name = get_amd_name(&device_path) .unwrap_or(amd_gpu_marketing::AMDGPU_DEFAULT_NAME.to_string()); diff --git a/src/collection/linux/utils.rs b/src/collection/linux/utils.rs index ff148a6ff..b8d017ba9 100644 --- a/src/collection/linux/utils.rs +++ b/src/collection/linux/utils.rs @@ -1,4 +1,6 @@ -use std::{fs, path::Path}; +use std::{borrow::Cow, fs, os::unix::ffi::OsStrExt, path::Path}; + +use libc::PATH_MAX; /// Whether the temperature should *actually* be read during enumeration. /// Will return false if the state is not D0/unknown, or if it does not support @@ -28,3 +30,37 @@ pub fn is_device_awake(device: &Path) -> bool { true } } + +/// A custom implementation to read a symlink while allowing for buffer reuse. If the path is +/// not a symlink, this will also return an error. +/// +/// If successful, then a [`Cow`] will be returned referencing the contents of `buffer`. +pub(crate) fn read_link<'a>(path: &Path, buffer: &'a mut Vec) -> std::io::Result> { + // if !path.is_symlink() { + // return Err(std::io::Error::new( + // std::io::ErrorKind::InvalidInput, + // "path is not a symlink", + // )); + // } + + let c_path = std::ffi::CString::new(path.as_os_str().as_bytes())?; + + buffer.clear(); + if buffer.len() < PATH_MAX as usize { + buffer.resize(PATH_MAX as usize, 0); + } + + // SAFETY: this is a libc API; we must check the length which we do below. + let len = unsafe { + libc::readlink( + c_path.as_ptr(), + buffer.as_mut_ptr() as *mut libc::c_char, + buffer.len(), + ) + }; + + if len < 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(String::from_utf8_lossy(&buffer[..len as usize])) +} diff --git a/src/collection/processes/linux/gpu.rs b/src/collection/processes/linux/gpu.rs new file mode 100644 index 000000000..d2f8ec0c9 --- /dev/null +++ b/src/collection/processes/linux/gpu.rs @@ -0,0 +1,44 @@ +//! Extract GPU process information on Linux. + +use std::{os::fd::BorrowedFd, path::Path}; + +use hashbrown::HashSet; +use rustix::fs::{Mode, OFlags}; + +use crate::collection::processes::Pid; + +fn is_drm_fd(fd: &BorrowedFd<'_>) -> bool { + true +} + +/// Get fdinfo for a process given the PID. +/// +/// Based on the method from nvtop [here](https://github.com/Syllo/nvtop/blob/339ee0b10a64ec51f43d27357b0068a40f16e9e4/src/extract_processinfo_fdinfo.c#L101). +pub(crate) fn get_fdinfo(pid: Pid, seen_fds: &mut HashSet) { + let fdinfo_path = format!("/proc/{pid}/fdinfo"); + let fdinfo_path = Path::new(&fdinfo_path); + + let Ok(fd_entries) = std::fs::read_dir(fdinfo_path) else { + return; + }; + + for fd_entry in fd_entries.flatten() { + let path = fd_entry.path(); + + if !path.is_file() { + continue; + } + + if !(path.to_string_lossy().chars().all(|c| c.is_ascii_digit())) { + continue; + } + + let Ok(fd) = rustix::fs::openat( + pid_path.as_path(), + OFlags::PATH | OFlags::DIRECTORY | OFlags::CLOEXEC, + Mode::empty(), + ) else { + continue; + }; + } +} diff --git a/src/collection/processes/linux/mod.rs b/src/collection/processes/linux/mod.rs index b49fb27d1..d59302314 100644 --- a/src/collection/processes/linux/mod.rs +++ b/src/collection/processes/linux/mod.rs @@ -1,5 +1,6 @@ //! Process data collection for Linux. +mod gpu; mod process; use std::{ @@ -137,7 +138,7 @@ fn read_proc( thread_parent: Option, ) -> CollectionResult<(ProcessHarvest, u64)> { let Process { - pid: _pid, + pid: _, uid, stat, io, @@ -200,12 +201,7 @@ fn read_proc( }; let user = uid - .and_then(|uid| { - user_table - .get_uid_to_username_mapping(uid) - .map(Into::into) - .ok() - }) + .and_then(|uid| user_table.uid_to_username(uid).map(Into::into).ok()) .unwrap_or_else(|| "N/A".into()); let time = if let Ok(ticks_per_sec) = u32::try_from(rustix::param::clock_ticks_per_second()) { @@ -351,6 +347,10 @@ pub(crate) struct ReadProcArgs { pub(crate) fn linux_process_data( collector: &mut DataCollector, time_difference_in_secs: u64, ) -> CollectionResult> { + if collector.should_run_less_routine_tasks { + collector.process_buffer = String::new(); + } + let total_memory = collector.total_memory(); let prev_proc = PrevProc { prev_idle: &mut collector.prev_idle, @@ -375,8 +375,6 @@ pub(crate) fn linux_process_data( prev_non_idle, } = prev_proc; - // TODO: [PROC THREADS] Add threads - let CpuUsage { mut cpu_usage, cpu_fraction, @@ -414,15 +412,15 @@ pub(crate) fn linux_process_data( get_process_threads: get_threads, }; - // TODO: Maybe pre-allocate these buffers in the future w/ routine cleanup. - let mut buffer = String::new(); let mut process_threads_to_check = HashMap::new(); let mut process_vector: Vec = pids .filter_map(|pid_path| { - if let Ok((process, threads)) = - Process::from_path(pid_path, &mut buffer, args.get_process_threads) - { + if let Ok((process, threads)) = Process::from_path( + pid_path, + &mut collector.process_buffer, + args.get_process_threads, + ) { let pid = process.pid; let prev_proc_details = prev_process_details.entry(pid).or_default(); @@ -466,7 +464,9 @@ pub(crate) fn linux_process_data( // Get thread data. for (pid, tid_paths) in process_threads_to_check { for tid_path in tid_paths { - if let Ok((process, _)) = Process::from_path(tid_path, &mut buffer, false) { + if let Ok((process, _)) = + Process::from_path(tid_path, &mut collector.process_buffer, false) + { let tid = process.pid; let prev_proc_details = prev_process_details.entry(tid).or_default(); diff --git a/src/collection/processes/linux/process.rs b/src/collection/processes/linux/process.rs index 6d6f5b326..a55736fab 100644 --- a/src/collection/processes/linux/process.rs +++ b/src/collection/processes/linux/process.rs @@ -16,7 +16,10 @@ use rustix::{ path::Arg, }; -use crate::collection::processes::{Pid, linux::is_str_numeric}; +use crate::collection::{ + linux::utils::read_link, + processes::{Pid, linux::is_str_numeric}, +}; static PAGESIZE: OnceLock = OnceLock::new(); @@ -73,6 +76,9 @@ impl Stat { // TODO: Is this needed? let line = buffer.trim(); + // TODO: comm is max 16, so we could in theory pre-allocate this. Also get it from /proc/pid/comm instead? + // They slightly differ though, see https://unix.stackexchange.com/questions/769962/thread-name-is-proc-pid-comm-always-identical-to-the-name-line-of-proc-pid-s + let (comm, rest) = { let start_paren = line .find('(') @@ -250,9 +256,16 @@ impl Process { .next_back() .and_then(|s| s.to_string_lossy().parse::().ok()) .or_else(|| { - rustix::fs::readlinkat(rustix::fs::CWD, pid_path.as_path(), vec![]) + // SAFETY: We can do this safely, we plan to only put a valid string in here. + let buffer = unsafe { buffer.as_mut_vec() }; + + let out = read_link(pid_path.as_path(), buffer) .ok() - .and_then(|s| s.to_string_lossy().parse::().ok()) + .and_then(|s| s.parse::().ok()); + + buffer.clear(); + + out }) .ok_or_else(|| anyhow!("PID for {pid_path:?} was not found"))?; diff --git a/src/collection/processes/unix/process_ext.rs b/src/collection/processes/unix/process_ext.rs index d303323b3..9498f0e85 100644 --- a/src/collection/processes/unix/process_ext.rs +++ b/src/collection/processes/unix/process_ext.rs @@ -89,12 +89,7 @@ pub(crate) trait UnixProcessExt { process_state, uid, user: uid - .and_then(|uid| { - user_table - .get_uid_to_username_mapping(uid) - .map(Into::into) - .ok() - }) + .and_then(|uid| user_table.uid_to_username(uid).map(Into::into).ok()) .unwrap_or_else(|| "N/A".into()), time: if process_val.start_time() == 0 { // Workaround for sysinfo occasionally returning a start time equal to UNIX diff --git a/src/collection/processes/unix/user_table.rs b/src/collection/processes/unix/user_table.rs index 036fe60e3..b2836a51c 100644 --- a/src/collection/processes/unix/user_table.rs +++ b/src/collection/processes/unix/user_table.rs @@ -8,11 +8,13 @@ pub struct UserTable { } impl UserTable { - pub fn get_uid_to_username_mapping(&mut self, uid: libc::uid_t) -> CollectionResult { + /// Get the username associated with a UID. On first access of a name, it will + /// be cached for future accesses. + pub fn uid_to_username(&mut self, uid: libc::uid_t) -> CollectionResult { if let Some(user) = self.uid_user_mapping.get(&uid) { Ok(user.clone()) } else { - // SAFETY: getpwuid returns a null pointer if no passwd entry is found for the uid. + // SAFETY: getpwuid returns a null pointer if no passwd entry is found for the uid which we check. let passwd = unsafe { libc::getpwuid(uid) }; if passwd.is_null() {