Skip to content

refactor: some optimizations for linux data collection #1807

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ClementTsang wants to merge 15 commits into
base: main
Choose a base branch
from
Open

refactor: some optimizations for linux data collection #1807

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
8015f0c
driveby refactor
ClementTsang Aug 31, 2025
bb741cf
reduce some allocations in amd step
ClementTsang Aug 31, 2025
f5833c5
some todos
ClementTsang Aug 31, 2025
85eee5d
hmmm
ClementTsang Aug 31, 2025
d510c74
try custom read_link
ClementTsang Sep 1, 2025
80d2ef2
fix name
ClementTsang Sep 1, 2025
7b88dbe
fmt
ClementTsang Sep 1, 2025
a226a10
what
ClementTsang Sep 1, 2025
e3398d8
driveby fix for dash cmdline bin name
ClementTsang Sep 1, 2025
462cf76
fmt
ClementTsang Sep 1, 2025
8343fb7
probably not needed to check for symlink
ClementTsang Sep 1, 2025
2900ed0
why was this collecting
ClementTsang Sep 1, 2025
5856521
okay so temp commit before I rewrite this whole thing
ClementTsang Sep 1, 2025
5a74cf9
Merge branch 'main' of https://github.com/ClementTsang/bottom into so…
ClementTsang Sep 2, 2025
73edcd9
hmmmm
ClementTsang Sep 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/collection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Manager>,
Expand Down Expand Up @@ -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()
}
}

Expand Down
66 changes: 40 additions & 26 deletions src/collection/amd.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec> for example is a bit redundant.
pub struct AmdGpuData {
Expand Down Expand Up @@ -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<PathBuf>) -> Option<Vec<u32>> {
fn get_amdgpu_pid_fds(pid: u32, device_path: &[String], buffer: &mut Vec<u8>) -> Option<Vec<u32>> {
let Ok(fd_list) = fs::read_dir(format!("/proc/{pid}/fd/")) else {
return None;
};

let valid_fds: Vec<u32> = 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::<u32>().ok()
// TODO: Should we bother parsing here?
path.file_name()?.to_str()?.parse::<u32>().ok()
} else {
None
}
Expand All @@ -188,7 +193,7 @@ fn get_amdgpu_pid_fds(pid: u32, device_path: Vec<PathBuf>) -> Option<Vec<u32>> {
}
}

fn get_amdgpu_drm(device_path: &Path) -> Option<Vec<PathBuf>> {
fn get_amdgpu_drm(device_path: &Path) -> Option<Vec<String>> {
let mut drm_devices = Vec::new();
let drm_root = device_path.join("drm");

Expand All @@ -212,7 +217,7 @@ fn get_amdgpu_drm(device_path: &Path) -> Option<Vec<PathBuf>> {
continue;
}

drm_devices.push(PathBuf::from(format!("/dev/dri/{drm_name}")));
drm_devices.push(format!("/dev/dri/{drm_name}"));
}

if drm_devices.is_empty() {
Expand All @@ -222,44 +227,51 @@ fn get_amdgpu_drm(device_path: &Path) -> Option<Vec<PathBuf>> {
}
}

// 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<HashMap<u32, AmdGpuProc>> {
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)?;

let Ok(proc_dir) = fs::read_dir("/proc") else {
return None;
};

let pids: Vec<u32> = 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::<u32>().ok()?;
let pid = dir_entry.file_name().to_str()?.parse::<u32>().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<usize> = HashSet::new();

for fd in fds {
Expand All @@ -271,6 +283,7 @@ fn get_amd_fdinfo(device_path: &Path) -> Option<HashMap<u32, AmdGpuProc>> {
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()
Expand Down Expand Up @@ -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());
Expand Down
38 changes: 37 additions & 1 deletion src/collection/linux/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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<u8>) -> std::io::Result<Cow<'a, str>> {
// 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]))
}
44 changes: 44 additions & 0 deletions src/collection/processes/linux/gpu.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<u32>) {
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;
};
}
}
30 changes: 15 additions & 15 deletions src/collection/processes/linux/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Process data collection for Linux.

mod gpu;
mod process;

use std::{
Expand Down Expand Up @@ -137,7 +138,7 @@ fn read_proc(
thread_parent: Option<Pid>,
) -> CollectionResult<(ProcessHarvest, u64)> {
let Process {
pid: _pid,
pid: _,
uid,
stat,
io,
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -351,6 +347,10 @@ pub(crate) struct ReadProcArgs {
pub(crate) fn linux_process_data(
collector: &mut DataCollector, time_difference_in_secs: u64,
) -> CollectionResult<Vec<ProcessHarvest>> {
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,
Expand All @@ -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,
Expand Down Expand Up @@ -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<ProcessHarvest> = 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();

Expand Down Expand Up @@ -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();

Expand Down
19 changes: 16 additions & 3 deletions src/collection/processes/linux/process.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<u64> = OnceLock::new();

Expand Down Expand Up @@ -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('(')
Expand Down Expand Up @@ -250,9 +256,16 @@ impl Process {
.next_back()
.and_then(|s| s.to_string_lossy().parse::<Pid>().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::<Pid>().ok())
.and_then(|s| s.parse::<Pid>().ok());

buffer.clear();

out
})
.ok_or_else(|| anyhow!("PID for {pid_path:?} was not found"))?;

Expand Down
7 changes: 1 addition & 6 deletions src/collection/processes/unix/process_ext.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading