diff --git a/src/collection.rs b/src/collection.rs index 93a16fa1..1d07d8a2 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")] + // SAFETY: This is an empty string. + process_buffer: unsafe { String::from_utf8_unchecked(Vec::with_capacity(16_384)) } } } diff --git a/src/collection/processes/linux/mod.rs b/src/collection/processes/linux/mod.rs index a64952f7..ac8e84a3 100644 --- a/src/collection/processes/linux/mod.rs +++ b/src/collection/processes/linux/mod.rs @@ -404,14 +404,16 @@ pub(crate) fn linux_process_data( }; // TODO: Maybe pre-allocate these buffers in the future w/ routine cleanup. - let mut buffer = String::new(); + // SAFETY: This is safe, we're converting an empty string. 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(); @@ -455,7 +457,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 b6c745be..71146760 100644 --- a/src/collection/processes/linux/process.rs +++ b/src/collection/processes/linux/process.rs @@ -3,7 +3,7 @@ use std::{ fs::File, - io::{self, BufRead, BufReader, Read}, + io::{self, BufRead, BufReader}, path::PathBuf, sync::OnceLock, }; @@ -65,10 +65,13 @@ impl Stat { /// Get process stats from a file; this assumes the file is located at /// `/proc//stat`. For documentation, see /// [here](https://manpages.ubuntu.com/manpages/noble/man5/proc_pid_stat.5.html) as a reference. - fn from_file(mut f: File, buffer: &mut String) -> anyhow::Result { + fn from_file(fd: OwnedFd, buffer: &mut String) -> anyhow::Result { // Since this is just one line, we can read it all at once. However, since it // (technically) might have non-utf8 characters, we can't just use read_to_string. - f.read_to_end(unsafe { buffer.as_mut_vec() })?; + // + // TODO: Can we read per delim. token to avoid memory? + // SAFETY: We are only going to be reading strings. + rustix::io::read(&fd, unsafe { buffer.as_mut_vec() })?; // TODO: Is this needed? let line = buffer.trim(); @@ -136,13 +139,15 @@ pub(crate) struct Io { impl Io { #[inline] - fn from_file(f: File, buffer: &mut String) -> anyhow::Result { + fn from_file(fd: OwnedFd, buffer: &mut String) -> anyhow::Result { const NUM_FIELDS: u16 = 0; // Make sure to update this if you want more fields! enum Fields { ReadBytes, WriteBytes, } + let f = File::from(fd); + let mut read_fields = 0; let mut reader = BufReader::new(f); @@ -272,7 +277,7 @@ impl Process { // Stat is pretty long, do this first to pre-allocate up-front. let stat = - open_at(&mut root, "stat", &pid_dir).and_then(|file| Stat::from_file(file, buffer))?; + open_at(&mut root, "stat", &pid_dir).and_then(|fd| Stat::from_file(fd, buffer))?; reset(&mut root, buffer); let cmdline = if cmdline(&mut root, &pid_dir, buffer).is_ok() { @@ -306,20 +311,22 @@ impl Process { #[inline] fn cmdline(root: &mut PathBuf, fd: &OwnedFd, buffer: &mut String) -> anyhow::Result<()> { - let _ = open_at(root, "cmdline", fd).map(|mut file| file.read_to_string(buffer))?; + // SAFETY: This is safe, we are only writing strings. + let _ = open_at(root, "cmdline", fd) + .map(|cmdline_fd| rustix::io::read(cmdline_fd, unsafe { buffer.as_mut_vec() }))?; Ok(()) } -/// Opens a path. Note that this function takes in a mutable root - this will +/// Opens a path and return the file. Note that this function takes in a mutable root - this will /// mutate it to avoid allocations. You probably will want to pop the most /// recent child after if you need to use the buffer again. #[inline] -fn open_at(root: &mut PathBuf, child: &str, fd: &OwnedFd) -> anyhow::Result { +fn open_at(root: &mut PathBuf, child: &str, fd: &OwnedFd) -> anyhow::Result { root.push(child); let new_fd = rustix::fs::openat(fd, &*root, OFlags::RDONLY | OFlags::CLOEXEC, Mode::empty())?; - Ok(File::from(new_fd)) + Ok(new_fd) } #[inline]