refactor: clean up some file structure, process code, and terminal cleanup (#1676)

* move widgets

* reduce allocations needed

* ah

* more possible optimizations around reducing allocs

* some fixes

* I forgot to clear the buffer oops

* missing

* only run terminal cleanup after certain point
This commit is contained in:
Clement Tsang 2025-02-15 02:32:09 -05:00 committed by GitHub
parent 2b5441ca8b
commit d63ca07cae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 206 additions and 92 deletions

View File

@ -25,7 +25,9 @@ pub type Values = ChunkedData<f64>;
#[derive(Clone, Debug, Default)]
pub struct TimeSeriesData {
/// Time values.
pub time: Vec<Instant>, // TODO: (points_rework_v1) should we not store instant, and just store the millisecond component? Write a compatible wrapper!
///
/// TODO: (points_rework_v1) Either store millisecond-level only or offsets only.
pub time: Vec<Instant>,
/// Network RX data.
pub rx: Values,

View File

@ -1,7 +1,11 @@
use bottom::{reset_stdout, start_bottom};
fn main() -> anyhow::Result<()> {
start_bottom().inspect_err(|_| {
reset_stdout();
let mut run_error_hook = false;
start_bottom(&mut run_error_hook).inspect_err(|_| {
if run_error_hook {
reset_stdout();
}
})
}

View File

@ -4,6 +4,7 @@
//! For Windows, macOS, FreeBSD, Android, and Linux, this is handled by sysinfo.
use cfg_if::cfg_if;
use sysinfo::ProcessStatus;
cfg_if! {
if #[cfg(target_os = "linux")] {
@ -81,7 +82,7 @@ pub struct ProcessHarvest {
pub total_write_bytes: u64,
/// The current state of the process (e.g. zombie, asleep).
pub process_state: (String, char),
pub process_state: (&'static str, char),
/// Cumulative process uptime.
pub time: Duration,
@ -150,3 +151,52 @@ impl DataCollector {
}
}
}
/// Pulled from [`ProcessStatus::to_string`] to avoid an alloc.
pub(super) fn process_status_str(status: ProcessStatus) -> &'static str {
cfg_if::cfg_if! {
if #[cfg(target_os = "linux")] {
match status {
ProcessStatus::Idle => "Idle",
ProcessStatus::Run => "Runnable",
ProcessStatus::Sleep => "Sleeping",
ProcessStatus::Stop => "Stopped",
ProcessStatus::Zombie => "Zombie",
ProcessStatus::Tracing => "Tracing",
ProcessStatus::Dead => "Dead",
ProcessStatus::Wakekill => "Wakekill",
ProcessStatus::Waking => "Waking",
ProcessStatus::Parked => "Parked",
ProcessStatus::UninterruptibleDiskSleep => "UninterruptibleDiskSleep",
_ => "Unknown",
}
} else if #[cfg(target_os = "windows")] {
match status {
ProcessStatus::Run => "Runnable",
_ => "Unknown",
}
} else if #[cfg(target_os = "macos")] {
match status {
ProcessStatus::Idle => "Idle",
ProcessStatus::Run => "Runnable",
ProcessStatus::Sleep => "Sleeping",
ProcessStatus::Stop => "Stopped",
ProcessStatus::Zombie => "Zombie",
_ => "Unknown",
}
} else if #[cfg(target_os = "freebsd")] {
match status {
ProcessStatus::Idle => "Idle",
ProcessStatus::Run => "Runnable",
ProcessStatus::Sleep => "Sleeping",
ProcessStatus::Stop => "Stopped",
ProcessStatus::Zombie => "Zombie",
ProcessStatus::Dead => "Dead",
ProcessStatus::LockBlocked => "LockBlocked",
_ => "Unknown",
}
} else {
"Unknown"
}
}
}

View File

@ -8,11 +8,12 @@ use std::{
time::Duration,
};
use concat_string::concat_string;
use hashbrown::HashSet;
use process::*;
use sysinfo::ProcessStatus;
use super::{Pid, ProcessHarvest, UserTable};
use super::{process_status_str, Pid, ProcessHarvest, UserTable};
use crate::collection::{error::CollectionResult, DataCollector};
/// Maximum character length of a `/proc/<PID>/stat`` process name.
@ -149,39 +150,9 @@ fn read_proc(
uptime,
} = args;
let (command, name) = {
let truncated_name = stat.comm.as_str();
if let Ok(cmdline) = cmdline {
if cmdline.is_empty() {
(format!("[{truncated_name}]"), truncated_name.to_string())
} else {
(
cmdline.join(" "),
if truncated_name.len() >= MAX_STAT_NAME_LEN {
if let Some(first_part) = cmdline.first() {
// We're only interested in the executable part... not the file path.
// That's for command.
first_part
.rsplit_once('/')
.map(|(_prefix, suffix)| suffix)
.unwrap_or(truncated_name)
.to_string()
} else {
truncated_name.to_string()
}
} else {
truncated_name.to_string()
},
)
}
} else {
(truncated_name.to_string(), truncated_name.to_string())
}
};
let process_state_char = stat.state;
let process_state = (
ProcessStatus::from(process_state_char).to_string(),
process_status_str(ProcessStatus::from(process_state_char)),
process_state_char,
);
let (cpu_usage_percent, new_process_times) = get_linux_cpu_usage(
@ -197,7 +168,7 @@ fn read_proc(
// This can fail if permission is denied!
let (total_read_bytes, total_write_bytes, read_bytes_per_sec, write_bytes_per_sec) =
if let Ok(io) = io {
if let Some(io) = io {
let total_read_bytes = io.read_bytes;
let total_write_bytes = io.write_bytes;
let prev_total_read_bytes = prev_proc.total_read_bytes;
@ -242,6 +213,37 @@ fn read_proc(
Duration::ZERO
};
let (command, name) = {
let truncated_name = stat.comm;
if let Some(cmdline) = cmdline {
if cmdline.is_empty() {
(concat_string!("[", truncated_name, "]"), truncated_name)
} else {
let name = if truncated_name.len() >= MAX_STAT_NAME_LEN {
let first_part = match cmdline.split_once(' ') {
Some((first, _)) => first,
None => &cmdline,
};
// We're only interested in the executable part, not the file path (part of command),
// so strip everything but the command name if needed.
let last_part = match first_part.rsplit_once('/') {
Some((_, last)) => last,
None => first_part,
};
last_part.to_string()
} else {
truncated_name
};
(cmdline, name)
}
} else {
(truncated_name.clone(), truncated_name)
}
};
Ok((
ProcessHarvest {
pid: process.pid,
@ -354,9 +356,11 @@ pub(crate) fn linux_process_data(
uptime: sysinfo::System::uptime(),
};
let mut buffer = String::new();
let process_vector: Vec<ProcessHarvest> = pids
.filter_map(|pid_path| {
if let Ok(process) = Process::from_path(pid_path) {
if let Ok(process) = Process::from_path(pid_path, &mut buffer) {
let pid = process.pid;
let prev_proc_details = pid_mapping.entry(pid).or_default();

View File

@ -58,14 +58,13 @@ pub(crate) struct Stat {
}
impl Stat {
#[inline]
fn from_file(mut f: File, buffer: &mut String) -> anyhow::Result<Stat> {
// Since this is just one line, we can read it all at once. However, since it
// might have non-utf8 characters, we can't just use read_to_string.
// (technically) might have non-utf8 characters, we can't just use read_to_string.
f.read_to_end(unsafe { buffer.as_mut_vec() })?;
let line = buffer.to_string_lossy();
let line = line.trim();
// TODO: Is this needed?
let line = buffer.trim();
let (comm, rest) = {
let start_paren = line
@ -204,8 +203,8 @@ pub(crate) struct Process {
pub pid: Pid,
pub uid: Option<uid_t>,
pub stat: Stat,
pub io: anyhow::Result<Io>,
pub cmdline: anyhow::Result<Vec<String>>,
pub io: Option<Io>,
pub cmdline: Option<String>,
}
#[inline]
@ -223,8 +222,10 @@ impl Process {
/// methods. Therefore, this struct is only useful for either fields
/// that are unlikely to change, or are short-lived and
/// will be discarded quickly.
pub(crate) fn from_path(pid_path: PathBuf) -> anyhow::Result<Process> {
// TODO: Pass in a buffer vec/string to share?
///
/// This takes in a buffer to avoid allocs; this function will clear the buffer.
pub(crate) fn from_path(pid_path: PathBuf, buffer: &mut String) -> anyhow::Result<Process> {
buffer.clear();
let fd = rustix::fs::openat(
rustix::fs::CWD,
@ -254,18 +255,26 @@ impl Process {
};
let mut root = pid_path;
let mut buffer = String::new();
// NB: Whenever you add a new stat, make sure to pop the root and clear the
// buffer!
// Stat is pretty long, do this first to pre-allocate up-front.
let stat =
open_at(&mut root, "stat", &fd).and_then(|file| Stat::from_file(file, &mut buffer))?;
reset(&mut root, &mut buffer);
open_at(&mut root, "stat", &fd).and_then(|file| Stat::from_file(file, buffer))?;
reset(&mut root, buffer);
let cmdline = cmdline(&mut root, &fd, &mut buffer);
reset(&mut root, &mut buffer);
let cmdline = if cmdline(&mut root, &fd, buffer).is_ok() {
// The clone will give a string with the capacity of the length of buffer, don't worry.
Some(buffer.clone())
} else {
None
};
reset(&mut root, buffer);
let io = open_at(&mut root, "io", &fd).and_then(|file| Io::from_file(file, &mut buffer));
let io = open_at(&mut root, "io", &fd)
.and_then(|file| Io::from_file(file, buffer))
.ok();
Ok(Process {
pid,
@ -278,22 +287,22 @@ impl Process {
}
#[inline]
fn cmdline(root: &mut PathBuf, fd: &OwnedFd, buffer: &mut String) -> anyhow::Result<Vec<String>> {
open_at(root, "cmdline", fd)
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))
.map(|_| {
buffer
.split('\0')
.filter_map(|s| {
if !s.is_empty() {
Some(s.to_string())
} else {
None
}
})
.collect::<Vec<_>>()
})
.map_err(Into::into)
.inspect(|_| {
// SAFETY: We are only replacing a single char (NUL) with another single char (space).
let buf_mut = unsafe { buffer.as_mut_vec() };
for byte in buf_mut {
if *byte == 0 {
const SPACE: u8 = ' '.to_ascii_lowercase() as u8;
*byte = SPACE;
}
}
})?;
Ok(())
}
/// Opens a path. Note that this function takes in a mutable root - this will

View File

@ -5,7 +5,7 @@ use std::{io, time::Duration};
use hashbrown::HashMap;
use sysinfo::{ProcessStatus, System};
use super::ProcessHarvest;
use super::{process_status_str, ProcessHarvest};
use crate::collection::{error::CollectionResult, processes::UserTable, Pid};
pub(crate) trait UnixProcessExt {
@ -60,7 +60,7 @@ pub(crate) trait UnixProcessExt {
let disk_usage = process_val.disk_usage();
let process_state = {
let ps = process_val.status();
(ps.to_string(), convert_process_status_to_char(ps))
(process_status_str(ps), convert_process_status_to_char(ps))
};
let uid = process_val.user_id().map(|u| **u);
let pid = process_val.pid().as_u32() as Pid;
@ -146,11 +146,57 @@ pub(crate) trait UnixProcessExt {
}
fn convert_process_status_to_char(status: ProcessStatus) -> char {
match status {
ProcessStatus::Run => 'R',
ProcessStatus::Sleep => 'S',
ProcessStatus::Idle => 'D',
ProcessStatus::Zombie => 'Z',
_ => '?',
// TODO: Based on https://github.com/GuillaumeGomez/sysinfo/blob/baa46efb46d82f21b773088603720262f4a34646/src/unix/freebsd/process.rs#L13?
cfg_if::cfg_if! {
if #[cfg(target_os = "macos")] {
// SAFETY: These are all const and should be valid characters.
const SIDL: char = unsafe { char::from_u32_unchecked(libc::SIDL) };
// SAFETY: These are all const and should be valid characters.
const SRUN: char = unsafe { char::from_u32_unchecked(libc::SRUN) };
// SAFETY: These are all const and should be valid characters.
const SSLEEP: char = unsafe { char::from_u32_unchecked(libc::SSLEEP) };
// SAFETY: These are all const and should be valid characters.
const SSTOP: char = unsafe { char::from_u32_unchecked(libc::SSTOP) };
// SAFETY: These are all const and should be valid characters.
const SZOMB: char = unsafe { char::from_u32_unchecked(libc::SZOMB) };
match status {
ProcessStatus::Idle => SIDL,
ProcessStatus::Run => SRUN,
ProcessStatus::Sleep => SSLEEP,
ProcessStatus::Stop => SSTOP,
ProcessStatus::Zombie => SZOMB,
_ => '?'
}
} else if #[cfg(target_os = "freebsd")] {
const fn assert_u8(val: i8) -> u8 {
if val < 0 { panic!("there was an invalid i8 constant that is supposed to be a char") } else { val as u8 }
}
const SIDL: u8 = assert_u8(libc::SIDL);
const SRUN: u8 = assert_u8(libc::SRUN);
const SSLEEP: u8 = assert_u8(libc::SSLEEP);
const SSTOP: u8 = assert_u8(libc::SSTOP);
const SZOMB: u8 = assert_u8(libc::SZOMB);
const SWAIT: u8 = assert_u8(libc::SWAIT);
const SLOCK: u8 = assert_u8(libc::SLOCK);
match status {
ProcessStatus::Idle => SIDL as char,
ProcessStatus::Run => SRUN as char,
ProcessStatus::Sleep => SSLEEP as char,
ProcessStatus::Stop => SSTOP as char,
ProcessStatus::Zombie => SZOMB as char,
ProcessStatus::Dead => SWAIT as char,
ProcessStatus::LockBlocked => SLOCK as char,
_ => '?'
}
} else {
'?'
}
}
}

View File

@ -12,8 +12,7 @@ impl UserTable {
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.
let passwd = unsafe { libc::getpwuid(uid) };
if passwd.is_null() {

View File

@ -2,7 +2,7 @@
use std::time::Duration;
use super::ProcessHarvest;
use super::{process_status_str, ProcessHarvest};
use crate::collection::{error::CollectionResult, DataCollector};
// TODO: There's a lot of shared code with this and the unix impl.
@ -60,7 +60,7 @@ pub fn sysinfo_process_data(
} as f32;
let disk_usage = process_val.disk_usage();
let process_state = (process_val.status().to_string(), 'R');
let process_state = (process_status_str(process_val.status()), 'R');
#[cfg(feature = "gpu")]
let (gpu_mem, gpu_util, gpu_mem_percent) = {

View File

@ -25,7 +25,7 @@ pub mod widgets;
use std::{
boxed::Box,
io::{stderr, stdout, Stdout, Write},
io::{stderr, stdout, Write},
panic::{self, PanicHookInfo},
sync::{
mpsc::{self, Receiver, Sender},
@ -78,8 +78,8 @@ fn cleanup_terminal(
disable_raw_mode()?;
execute!(
terminal.backend_mut(),
DisableBracketedPaste,
DisableMouseCapture,
DisableBracketedPaste,
LeaveAlternateScreen,
Show,
)?;
@ -103,18 +103,16 @@ fn check_if_terminal() {
}
/// This manually resets stdout back to normal state.
pub fn reset_stdout() -> Stdout {
pub fn reset_stdout() {
let mut stdout = stdout();
let _ = disable_raw_mode();
let _ = execute!(
stdout,
DisableBracketedPaste,
DisableMouseCapture,
DisableBracketedPaste,
LeaveAlternateScreen,
Show,
);
stdout
}
/// A panic hook to properly restore the terminal in the case of a panic.
@ -275,7 +273,7 @@ fn create_collection_thread(
/// Main code to call.
#[inline]
pub fn start_bottom() -> anyhow::Result<()> {
pub fn start_bottom(enable_error_hook: &mut bool) -> anyhow::Result<()> {
// let _profiler = dhat::Profiler::new_heap();
let args = args::get_args();
@ -337,13 +335,15 @@ pub fn start_bottom() -> anyhow::Result<()> {
};
// Set up tui and crossterm
*enable_error_hook = true;
let mut stdout_val = stdout();
execute!(
stdout_val,
Hide,
EnterAlternateScreen,
EnableBracketedPaste,
EnableMouseCapture,
EnableBracketedPaste
)?;
enable_raw_mode()?;

View File

@ -1040,7 +1040,7 @@ mod test {
wps: 0,
total_read: 0,
total_write: 0,
process_state: "N/A".to_string(),
process_state: "N/A",
process_char: '?',
#[cfg(target_family = "unix")]
user: "root".to_string(),

View File

@ -203,7 +203,7 @@ pub struct ProcWidgetData {
pub wps: u64,
pub total_read: u64,
pub total_write: u64,
pub process_state: String,
pub process_state: &'static str,
pub process_char: char,
pub user: String,
pub num_similar: u64,
@ -242,7 +242,7 @@ impl ProcWidgetData {
wps: process.write_bytes_per_sec,
total_read: process.total_read_bytes,
total_write: process.total_write_bytes,
process_state: process.process_state.0.clone(),
process_state: process.process_state.0,
process_char: process.process_state.1,
user: process.user.to_string(),
num_similar: 1,
@ -348,7 +348,7 @@ impl DataToCell<ProcColumn> for ProcWidgetData {
if calculated_width < 8 {
self.process_char.to_string().into()
} else {
self.process_state.clone().into()
self.process_state.into()
}
}
ProcColumn::User => self.user.clone().into(),

View File

@ -814,7 +814,7 @@ impl Prefix {
process.name.as_str()
}),
PrefixType::Pid => r.is_match(process.pid.to_string().as_str()),
PrefixType::State => r.is_match(process.process_state.0.as_str()),
PrefixType::State => r.is_match(process.process_state.0),
PrefixType::User => r.is_match(process.user.as_ref()),
_ => true,
}