initial side-by-side impl for linux temmp and process as a poc

This commit is contained in:
ClementTsang 2024-08-14 22:36:49 -04:00
parent e268c48c25
commit c758bb241f
No known key found for this signature in database
GPG Key ID: DC3B7867D8D97095
29 changed files with 1328 additions and 55 deletions

View File

@ -4,7 +4,14 @@
//! the battery crate.
cfg_if::cfg_if! {
if #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux", target_os = "freebsd", target_os = "dragonfly", target_os = "ios"))] {
if #[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux",
target_os = "freebsd",
target_os = "dragonfly",
target_os = "ios",
))] {
pub mod battery;
pub use self::battery::*;
}

View File

@ -1,5 +1,8 @@
# Data Collection
**Note:** This information is really only useful to _developers_ of bottom,
and can be ignored by users.
Data collection in bottom has two main components: **sources** and **collectors**.
**Sources** are either libraries or system APIs that actually extract the data.
@ -8,7 +11,7 @@ or `libc` bindings, or Linux-specific code.
**Collectors** are _platform-specific_ (typically OS-specific), and can pull from
different sources to get all the data needed, with some glue code in between. As
such, sources should be written to be per-"job", and be divisible such that
such, sources should be written to be per-"job", and be divided such that
collectors can import specific code as needed.
We can kinda visualize this with a quick-and-dirty diagram (note this is not accurate or up-to-date):
@ -42,3 +45,11 @@ flowchart TB
freebsd -..-> FreeBSD
sysinfo -..-> FreeBSD
```
## Sources
As mentioned above, sources should be written in a way where collectors can easily pick the necessary code required.
## Collectors
Each platform should implement the `DataCollector` trait in `collectors/common.rs`. The trait has default implementations where essentially no work is done, which is used as fallback behaviour.

View File

@ -0,0 +1,27 @@
//! Common code amongst all data collectors.
use crate::new_data_collection::{
error::CollectionResult,
sources::common::{processes::ProcessHarvest, temperature::TemperatureData},
};
/// The trait representing what a per-platform data collector should implement.
pub(crate) trait DataCollector {
/// Refresh inner data sources to prepare them for gathering data.
///
/// Note that depending on the implementation, this may
/// not actually need to do anything.
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
/// Return temperature data.
fn get_temperature_data(&mut self) -> CollectionResult<Vec<TemperatureData>> {
Ok(vec![])
}
/// Return process data.
fn get_process_data(&mut self) -> CollectionResult<Vec<ProcessHarvest>> {
Ok(vec![])
}
}

View File

@ -0,0 +1,7 @@
use super::common::DataCollector;
/// A fallback [`DataCollector`] for unsupported systems
/// that does nothing.
pub struct FallbackDataCollector {}
impl DataCollector for FallbackDataCollector {}

View File

@ -0,0 +1,36 @@
//! The data collector for FreeBSD.
use crate::{
app::filter::Filter,
new_data_collection::{
error::CollectionResult,
sources::{
common::temperature::{TemperatureData, TemperatureType},
sysinfo::temperature::get_temperature_data,
},
},
};
use super::common::DataCollector;
/// The [`DataCollector`] for FreeBSD.
pub struct FreeBsdDataCollector {
temp_type: TemperatureType,
temp_filters: Option<Filter>,
}
impl DataCollector for FreeBsdDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
let mut results = get_temperature_data(&self.temp_type, &self.temp_filters);
for entry in sysctl_temp_iter(&self.temp_type, &self.temp_filters) {
results.push(entry);
}
Ok(Some(results))
}
}

View File

@ -0,0 +1,71 @@
//! The data collector for Linux.
use std::time::Instant;
use starship_battery::{Battery, Manager};
use crate::{
app::filter::Filter,
new_data_collection::{
error::CollectionResult,
sources::{
common::{
processes::ProcessHarvest,
temperature::{TemperatureData, TemperatureType},
},
linux::{
processes::{linux_process_data, ProcessCollector},
temperature::get_temperature_data,
},
},
},
};
use super::common::DataCollector;
/// The [`DataCollector`] for Linux.
pub struct LinuxDataCollector {
current_collection_time: Instant,
last_collection_time: Instant,
temp_type: TemperatureType,
temp_filters: Option<Filter>,
proc_collector: ProcessCollector,
system: sysinfo::System,
network: sysinfo::Networks,
#[cfg(feature = "battery")]
battery_manager: Option<Manager>,
#[cfg(feature = "battery")]
battery_list: Option<Vec<Battery>>,
#[cfg(feature = "gpu")]
gpus_total_mem: Option<u64>,
}
impl DataCollector for LinuxDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&mut self) -> CollectionResult<Vec<TemperatureData>> {
Ok(get_temperature_data(&self.temp_type, &self.temp_filters))
}
fn get_process_data(&mut self) -> CollectionResult<Vec<ProcessHarvest>> {
let time_diff = self
.current_collection_time
.duration_since(self.last_collection_time)
.as_secs();
linux_process_data(
&self.system,
time_diff,
&mut self.proc_collector,
#[cfg(feature = "gpu")]
self.gpus_total_mem,
)
}
}

View File

@ -0,0 +1,33 @@
//! The data collector for macOS.
use crate::{
app::filter::Filter,
new_data_collection::{
error::CollectionResult,
sources::{
common::temperature::{TemperatureData, TemperatureType},
sysinfo::temperature::get_temperature_data,
},
},
};
use super::common::DataCollector;
/// The [`DataCollector`] for macOS.
pub struct MacOsDataCollector {
temp_type: TemperatureType,
temp_filters: Option<Filter>,
}
impl DataCollector for MacOsDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
Ok(Some(get_temperature_data(
&self.temp_type,
&self.temp_filters,
)))
}
}

View File

@ -0,0 +1,33 @@
//! The data collector for Windows.
use crate::{
app::filter::Filter,
new_data_collection::{
error::CollectionResult,
sources::{
common::temperature::{TemperatureData, TemperatureType},
sysinfo::temperature::get_temperature_data,
},
},
};
use super::common::DataCollector;
/// The [`DataCollector`] for Windows.
pub struct WindowsDataCollector {
temp_type: TemperatureType,
temp_filters: Option<Filter>,
}
impl DataCollector for WindowsDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
Ok(Some(get_temperature_data(
&self.temp_type,
&self.temp_filters,
)))
}
}

View File

@ -0,0 +1,42 @@
use anyhow::anyhow;
/// An error to do with data collection.
#[derive(Debug)]
pub enum CollectionError {
/// A general error to propagate back up. A wrapper around [`anyhow::Error`].
General(anyhow::Error),
/// The collection is unsupported.
Unsupported,
}
impl std::fmt::Display for CollectionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CollectionError::General(err) => err.fmt(f),
CollectionError::Unsupported => {
write!(
f,
"bottom does not support this type of data collection for this platform."
)
}
}
}
}
impl std::error::Error for CollectionError {}
/// A [`Result`] with the error type being a [`DataCollectionError`].
pub(crate) type CollectionResult<T> = Result<T, CollectionError>;
impl From<std::io::Error> for CollectionError {
fn from(err: std::io::Error) -> Self {
Self::General(err.into())
}
}
impl From<&'static str> for CollectionError {
fn from(msg: &'static str) -> Self {
Self::General(anyhow!(msg))
}
}

View File

@ -0,0 +1,28 @@
//! Module that just re-exports the right data collector for a given platform.
pub mod error;
mod collectors {
pub mod common;
cfg_if::cfg_if! {
if #[cfg(target_os = "linux")] {
pub mod linux;
pub use linux::LinuxDataCollector as DataCollectorImpl;
} else if #[cfg(target_os = "macos")] {
pub mod macos;
pub use macos::MacOsDataCollector as DataCollectorImpl;
} else if #[cfg(target_os = "windows")] {
pub mod windows;
pub use windows::WindowsDataCollector as DataCollectorImpl;
} else if #[cfg(target_os = "freebsd")] {
pub mod freebsd;
pub use freebsd::FreeBsdDataCollector as DataCollectorImpl;
} else {
pub mod fallback;
pub use fallback::FallbackDataCollector as DataCollectorImpl;
}
}
}
pub mod sources;

View File

@ -0,0 +1,2 @@
pub mod processes;
pub mod temperature;

View File

@ -0,0 +1,87 @@
use std::{borrow::Cow, time::Duration};
use crate::new_data_collection::sources::Pid;
#[derive(Debug, Clone, Default)]
pub struct ProcessHarvest {
/// The pid of the process.
pub pid: Pid,
/// The parent PID of the process. A `parent_pid` of 0 is usually the root.
pub parent_pid: Option<Pid>,
/// CPU usage as a percentage.
pub cpu_usage_percent: f32,
/// Memory usage as a percentage.
pub mem_usage_percent: f32,
/// Memory usage as bytes.
pub mem_usage_bytes: u64,
/// The name of the process.
pub name: String,
/// The exact command for the process.
pub command: String,
/// Bytes read per second.
pub read_bytes_per_sec: u64,
/// Bytes written per second.
pub write_bytes_per_sec: u64,
/// The total number of bytes read by the process.
pub total_read_bytes: u64,
/// The total number of bytes written by the process.
pub total_write_bytes: u64,
/// The current state of the process (e.g. zombie, asleep).
pub process_state: (String, char),
/// Cumulative process uptime.
pub time: Duration,
/// This is the *effective* user ID of the process. This is only used on
/// Unix platforms.
#[cfg(target_family = "unix")]
pub uid: Option<libc::uid_t>,
/// This is the process' user.
pub user: Cow<'static, str>,
/// GPU memory usage as bytes.
#[cfg(feature = "gpu")]
pub gpu_mem: u64,
/// GPU memory usage as percentage.
#[cfg(feature = "gpu")]
pub gpu_mem_percent: f32,
/// GPU utilization as a percentage.
#[cfg(feature = "gpu")]
pub gpu_util: u32,
// TODO: Additional fields
// pub rss_kb: u64,
// pub virt_kb: u64,
}
impl ProcessHarvest {
pub(crate) fn add(&mut self, rhs: &ProcessHarvest) {
self.cpu_usage_percent += rhs.cpu_usage_percent;
self.mem_usage_bytes += rhs.mem_usage_bytes;
self.mem_usage_percent += rhs.mem_usage_percent;
self.read_bytes_per_sec += rhs.read_bytes_per_sec;
self.write_bytes_per_sec += rhs.write_bytes_per_sec;
self.total_read_bytes += rhs.total_read_bytes;
self.total_write_bytes += rhs.total_write_bytes;
self.time = self.time.max(rhs.time);
#[cfg(feature = "gpu")]
{
self.gpu_mem += rhs.gpu_mem;
self.gpu_util += rhs.gpu_util;
self.gpu_mem_percent += rhs.gpu_mem_percent;
}
}
}

View File

@ -0,0 +1,70 @@
use std::str::FromStr;
#[derive(Default, Debug, Clone)]
pub struct TemperatureData {
pub name: String,
pub temperature: Option<f32>,
}
#[derive(Clone, Debug, Copy, PartialEq, Eq, Default)]
pub enum TemperatureType {
#[default]
Celsius,
Kelvin,
Fahrenheit,
}
impl FromStr for TemperatureType {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"fahrenheit" | "f" => Ok(TemperatureType::Fahrenheit),
"kelvin" | "k" => Ok(TemperatureType::Kelvin),
"celsius" | "c" => Ok(TemperatureType::Celsius),
_ => Err(format!(
"'{s}' is an invalid temperature type, use one of: [kelvin, k, celsius, c, fahrenheit, f]."
)),
}
}
}
impl TemperatureType {
/// Given a temperature in Celsius, covert it if necessary for a different
/// unit.
pub fn convert_temp_unit(&self, temp_celsius: f32) -> f32 {
fn convert_celsius_to_kelvin(celsius: f32) -> f32 {
celsius + 273.15
}
fn convert_celsius_to_fahrenheit(celsius: f32) -> f32 {
(celsius * (9.0 / 5.0)) + 32.0
}
match self {
TemperatureType::Celsius => temp_celsius,
TemperatureType::Kelvin => convert_celsius_to_kelvin(temp_celsius),
TemperatureType::Fahrenheit => convert_celsius_to_fahrenheit(temp_celsius),
}
}
}
#[cfg(test)]
mod test {
use crate::new_data_collection::sources::common::temperature::TemperatureType;
#[test]
fn temp_conversions() {
const TEMP: f32 = 100.0;
assert_eq!(
TemperatureType::Celsius.convert_temp_unit(TEMP),
TEMP,
"celsius to celsius is the same"
);
assert_eq!(TemperatureType::Kelvin.convert_temp_unit(TEMP), 373.15);
assert_eq!(TemperatureType::Fahrenheit.convert_temp_unit(TEMP), 212.0);
}
}

View File

@ -0,0 +1 @@
mod temperature;

View File

@ -0,0 +1,35 @@
//! FreeBSD-specific temperature extraction code.
// For RockPro64 boards on FreeBSD, they apparently use "hw.temperature" for
// sensors.
use sysctl::Sysctl;
/// Return an iterator of temperature data pulled from sysctl.
pub(crate) fn sysctl_temp_iter(
temp_type: &TemperatureType, filter: &Option<Filter>,
) -> impl Iterator<Item = TemperatureData> {
const KEY: &str = "hw.temperature";
if let Ok(root) = sysctl::Ctl::new(KEY) {
sysctl::CtlIter::below(root).flatten().filter_map(|ctl| {
if let (Ok(name), Ok(temp)) = (ctl.name(), ctl.value()) {
if let Some(temp) = temp.as_temperature() {
if Filter::optional_should_keep(filter, &name) {
return Some(TemperatureData {
name,
temperature: Some(match temp_type {
TemperatureType::Celsius => temp.celsius(),
TemperatureType::Kelvin => temp.kelvin(),
TemperatureType::Fahrenheit => temp.fahrenheit(),
}),
});
}
}
}
None
})
} else {
std::iter::empty()
}
}

View File

@ -0,0 +1,2 @@
pub mod processes;
pub mod temperature;

View File

@ -0,0 +1,450 @@
//! Process data collection for Linux.
mod process;
use std::{
fs::{self, File},
io::{BufRead, BufReader},
time::Duration,
};
use hashbrown::{HashMap, HashSet};
use process::*;
use sysinfo::ProcessStatus;
use crate::new_data_collection::{
error::CollectionResult,
sources::{common::processes::ProcessHarvest, unix::processes::user_table::UserTable, Pid},
};
/// Maximum character length of a `/proc/<PID>/stat`` process name.
/// If it's equal or greater, then we instead refer to the command for the name.
const MAX_STAT_NAME_LEN: usize = 15;
#[derive(Debug, Clone, Default)]
pub struct PrevProcDetails {
total_read_bytes: u64,
total_write_bytes: u64,
cpu_time: u64,
}
/// Given `/proc/stat` file contents, determine the idle and non-idle values of
/// the CPU used to calculate CPU usage.
fn fetch_cpu_usage(line: &str) -> (f64, f64) {
/// Converts a `Option<&str>` value to an f64. If it fails to parse or is
/// `None`, it will return `0_f64`.
fn str_to_f64(val: Option<&str>) -> f64 {
val.and_then(|v| v.parse::<f64>().ok()).unwrap_or(0_f64)
}
let mut val = line.split_whitespace();
let user = str_to_f64(val.next());
let nice: f64 = str_to_f64(val.next());
let system: f64 = str_to_f64(val.next());
let idle: f64 = str_to_f64(val.next());
let iowait: f64 = str_to_f64(val.next());
let irq: f64 = str_to_f64(val.next());
let softirq: f64 = str_to_f64(val.next());
let steal: f64 = str_to_f64(val.next());
// Note we do not get guest/guest_nice, as they are calculated as part of
// user/nice respectively See https://github.com/htop-dev/htop/blob/main/linux/LinuxProcessList.c
let idle = idle + iowait;
let non_idle = user + nice + system + irq + softirq + steal;
(idle, non_idle)
}
struct CpuUsage {
/// Difference between the total delta and the idle delta.
cpu_usage: f64,
/// Overall CPU usage as a fraction.
cpu_fraction: f64,
}
fn cpu_usage_calculation(
prev_idle: &mut f64, prev_non_idle: &mut f64,
) -> CollectionResult<CpuUsage> {
let (idle, non_idle) = {
// From SO answer: https://stackoverflow.com/a/23376195
let first_line = {
// We just need a single line from this file. Read it and return it.
let mut reader = BufReader::new(File::open("/proc/stat")?);
let mut buffer = String::new();
reader.read_line(&mut buffer)?;
buffer
};
fetch_cpu_usage(&first_line)
};
let total = idle + non_idle;
let prev_total = *prev_idle + *prev_non_idle;
let total_delta = total - prev_total;
let idle_delta = idle - *prev_idle;
*prev_idle = idle;
*prev_non_idle = non_idle;
// TODO: Should these return errors instead?
let cpu_usage = if total_delta - idle_delta != 0.0 {
total_delta - idle_delta
} else {
1.0
};
let cpu_fraction = if total_delta != 0.0 {
cpu_usage / total_delta
} else {
0.0
};
Ok(CpuUsage {
cpu_usage,
cpu_fraction,
})
}
/// Returns the usage and a new set of process times.
///
/// NB: cpu_fraction should be represented WITHOUT the x100 factor!
fn get_linux_cpu_usage(
stat: &Stat, cpu_usage: f64, cpu_fraction: f64, prev_proc_times: u64,
use_current_cpu_total: bool,
) -> (f32, u64) {
// Based heavily on https://stackoverflow.com/a/23376195 and https://stackoverflow.com/a/1424556
let new_proc_times = stat.utime + stat.stime;
let diff = (new_proc_times - prev_proc_times) as f64; // No try_from for u64 -> f64... oh well.
if cpu_usage == 0.0 {
(0.0, new_proc_times)
} else if use_current_cpu_total {
(((diff / cpu_usage) * 100.0) as f32, new_proc_times)
} else {
(
((diff / cpu_usage) * 100.0 * cpu_fraction) as f32,
new_proc_times,
)
}
}
fn read_proc(
prev_proc: &PrevProcDetails, process: Process, args: ReadProcArgs, user_table: &mut UserTable,
) -> CollectionResult<(ProcessHarvest, u64)> {
let Process {
pid: _,
uid,
stat,
io,
cmdline,
} = process;
let ReadProcArgs {
use_current_cpu_total,
cpu_usage,
cpu_fraction,
total_memory,
time_difference_in_secs,
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_state_char,
);
let (cpu_usage_percent, new_process_times) = get_linux_cpu_usage(
&stat,
cpu_usage,
cpu_fraction,
prev_proc.cpu_time,
use_current_cpu_total,
);
let parent_pid = Some(stat.ppid);
let mem_usage_bytes = stat.rss_bytes();
let mem_usage_percent = (mem_usage_bytes as f64 / total_memory as f64 * 100.0) as f32;
// 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 {
let total_read_bytes = io.read_bytes;
let total_write_bytes = io.write_bytes;
let prev_total_read_bytes = prev_proc.total_read_bytes;
let prev_total_write_bytes = prev_proc.total_write_bytes;
let read_bytes_per_sec = total_read_bytes
.saturating_sub(prev_total_read_bytes)
.checked_div(time_difference_in_secs)
.unwrap_or(0);
let write_bytes_per_sec = total_write_bytes
.saturating_sub(prev_total_write_bytes)
.checked_div(time_difference_in_secs)
.unwrap_or(0);
(
total_read_bytes,
total_write_bytes,
read_bytes_per_sec,
write_bytes_per_sec,
)
} else {
(0, 0, 0, 0)
};
let user = uid
.and_then(|uid| {
user_table
.get_uid_to_username_mapping(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()) {
if ticks_per_sec == 0 {
Duration::ZERO
} else {
Duration::from_secs(uptime.saturating_sub(stat.start_time / ticks_per_sec as u64))
}
} else {
Duration::ZERO
};
Ok((
ProcessHarvest {
pid: process.pid,
parent_pid,
cpu_usage_percent,
mem_usage_percent,
mem_usage_bytes,
name,
command,
read_bytes_per_sec,
write_bytes_per_sec,
total_read_bytes,
total_write_bytes,
process_state,
uid,
user,
time,
#[cfg(feature = "gpu")]
gpu_mem: 0,
#[cfg(feature = "gpu")]
gpu_mem_percent: 0.0,
#[cfg(feature = "gpu")]
gpu_util: 0,
},
new_process_times,
))
}
pub(crate) struct PrevProc {
pub prev_idle: f64,
pub prev_non_idle: f64,
}
#[derive(Clone, Copy)]
pub(crate) struct ProcHarvestOptions {
pub use_current_cpu_total: bool,
pub unnormalized_cpu: bool,
}
fn is_str_numeric(s: &str) -> bool {
s.chars().all(|c| c.is_ascii_digit())
}
/// General args to keep around for reading proc data.
#[derive(Copy, Clone)]
pub(crate) struct ReadProcArgs {
pub(crate) use_current_cpu_total: bool,
pub(crate) cpu_usage: f64,
pub(crate) cpu_fraction: f64,
pub(crate) total_memory: u64,
pub(crate) time_difference_in_secs: u64,
pub(crate) uptime: u64,
}
pub struct ProcessCollector {
pub options: ProcHarvestOptions,
pub prev_proc: PrevProc,
pub pid_mapping: HashMap<Pid, PrevProcDetails>,
pub user_table: UserTable,
#[cfg(feature = "gpu")]
pub gpu_pids: Option<Vec<HashMap<u32, (u64, u32)>>>,
}
pub(crate) fn linux_process_data(
system: &sysinfo::System, time_difference_in_secs: u64, collector: &mut ProcessCollector,
#[cfg(feature = "gpu")] gpus_total_mem: Option<u64>,
) -> CollectionResult<Vec<ProcessHarvest>> {
let total_memory = system.total_memory();
let ProcHarvestOptions {
use_current_cpu_total,
unnormalized_cpu,
} = collector.options;
let PrevProc {
prev_idle,
prev_non_idle,
} = &mut collector.prev_proc;
// TODO: [PROC THREADS] Add threads
let CpuUsage {
mut cpu_usage,
cpu_fraction,
} = cpu_usage_calculation(prev_idle, prev_non_idle)?;
if unnormalized_cpu {
let num_processors = system.cpus().len() as f64;
// Note we *divide* here because the later calculation divides `cpu_usage` - in
// effect, multiplying over the number of cores.
cpu_usage /= num_processors;
}
let mut pids_to_clear: HashSet<Pid> = collector.pid_mapping.keys().cloned().collect();
let pids = fs::read_dir("/proc")?.flatten().filter_map(|dir| {
if is_str_numeric(dir.file_name().to_string_lossy().trim()) {
Some(dir.path())
} else {
None
}
});
let args = ReadProcArgs {
use_current_cpu_total,
cpu_usage,
cpu_fraction,
total_memory,
time_difference_in_secs,
uptime: sysinfo::System::uptime(),
};
let process_vector: Vec<ProcessHarvest> = pids
.filter_map(|pid_path| {
if let Ok(process) = Process::from_path(pid_path) {
let pid = process.pid;
let prev_proc_details = collector.pid_mapping.entry(pid).or_default();
#[allow(unused_mut)]
if let Ok((mut process_harvest, new_process_times)) =
read_proc(prev_proc_details, process, args, &mut collector.user_table)
{
#[cfg(feature = "gpu")]
if let Some(gpus) = &collector.gpu_pids {
gpus.iter().for_each(|gpu| {
// add mem/util for all gpus to pid
if let Some((mem, util)) = gpu.get(&(pid as u32)) {
process_harvest.gpu_mem += mem;
process_harvest.gpu_util += util;
}
});
if let Some(gpu_total_mem) = gpus_total_mem {
process_harvest.gpu_mem_percent =
(process_harvest.gpu_mem as f64 / gpu_total_mem as f64 * 100.0)
as f32;
}
}
prev_proc_details.cpu_time = new_process_times;
prev_proc_details.total_read_bytes = process_harvest.total_read_bytes;
prev_proc_details.total_write_bytes = process_harvest.total_write_bytes;
pids_to_clear.remove(&pid);
return Some(process_harvest);
}
}
None
})
.collect();
pids_to_clear.iter().for_each(|pid| {
collector.pid_mapping.remove(pid);
});
Ok(process_vector)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_proc_cpu_parse() {
assert_eq!(
(100_f64, 200_f64),
fetch_cpu_usage("100 0 100 100"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 4 values"
);
assert_eq!(
(120_f64, 200_f64),
fetch_cpu_usage("100 0 100 100 20"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 5 values"
);
assert_eq!(
(120_f64, 230_f64),
fetch_cpu_usage("100 0 100 100 20 30"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 6 values"
);
assert_eq!(
(120_f64, 270_f64),
fetch_cpu_usage("100 0 100 100 20 30 40"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 7 values"
);
assert_eq!(
(120_f64, 320_f64),
fetch_cpu_usage("100 0 100 100 20 30 40 50"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 8 values"
);
assert_eq!(
(120_f64, 320_f64),
fetch_cpu_usage("100 0 100 100 20 30 40 50 100"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 9 values"
);
assert_eq!(
(120_f64, 320_f64),
fetch_cpu_usage("100 0 100 100 20 30 40 50 100 200"),
"Failed to properly calculate idle/non-idle for /proc/stat CPU with 10 values"
);
}
}

View File

@ -0,0 +1,308 @@
//! Linux process code for getting process data via `/proc/`.
//! Based on the [procfs](https://github.com/eminence/procfs) crate.
use std::{
fs::File,
io::{self, BufRead, BufReader, Read},
path::PathBuf,
sync::OnceLock,
};
use anyhow::anyhow;
use libc::uid_t;
use rustix::{
fd::OwnedFd,
fs::{Mode, OFlags},
path::Arg,
};
use crate::data_collection::processes::Pid;
static PAGESIZE: OnceLock<u64> = OnceLock::new();
#[inline]
fn next_part<'a>(iter: &mut impl Iterator<Item = &'a str>) -> Result<&'a str, io::Error> {
iter.next()
.ok_or_else(|| io::Error::from(io::ErrorKind::InvalidData))
}
/// A wrapper around the data in `/proc/<PID>/stat`. For documentation, see
/// [here](https://man7.org/linux/man-pages/man5/proc.5.html).
///
/// Note this does not necessarily get all fields, only the ones we use in
/// bottom.
pub(crate) struct Stat {
/// The filename of the executable without parentheses.
pub comm: String,
/// The current process state, represented by a char.
pub state: char,
/// The parent process PID.
pub ppid: Pid,
/// The amount of time this process has been scheduled in user mode in clock
/// ticks.
pub utime: u64,
/// The amount of time this process has been scheduled in kernel mode in
/// clock ticks.
pub stime: u64,
/// The resident set size, or the number of pages the process has in real
/// memory.
pub rss: u64,
/// The start time of the process, represented in clock ticks.
pub start_time: u64,
}
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.
f.read_to_end(unsafe { buffer.as_mut_vec() })?;
let line = buffer.to_string_lossy();
let line = line.trim();
let (comm, rest) = {
let start_paren = line
.find('(')
.ok_or_else(|| anyhow!("start paren missing"))?;
let end_paren = line.find(')').ok_or_else(|| anyhow!("end paren missing"))?;
(
line[start_paren + 1..end_paren].to_string(),
&line[end_paren + 2..],
)
};
let mut rest = rest.split(' ');
let state = next_part(&mut rest)?
.chars()
.next()
.ok_or_else(|| anyhow!("missing state"))?;
let ppid: Pid = next_part(&mut rest)?.parse()?;
// Skip 9 fields until utime (pgrp, session, tty_nr, tpgid, flags, minflt,
// cminflt, majflt, cmajflt).
let mut rest = rest.skip(9);
let utime: u64 = next_part(&mut rest)?.parse()?;
let stime: u64 = next_part(&mut rest)?.parse()?;
// Skip 6 fields until starttime (cutime, cstime, priority, nice, num_threads,
// itrealvalue).
let mut rest = rest.skip(6);
let start_time: u64 = next_part(&mut rest)?.parse()?;
// Skip one field until rss (vsize)
let mut rest = rest.skip(1);
let rss: u64 = next_part(&mut rest)?.parse()?;
Ok(Stat {
comm,
state,
ppid,
utime,
stime,
rss,
start_time,
})
}
/// Returns the Resident Set Size in bytes.
#[inline]
pub fn rss_bytes(&self) -> u64 {
self.rss * PAGESIZE.get_or_init(|| rustix::param::page_size() as u64)
}
}
/// A wrapper around the data in `/proc/<PID>/io`.
///
/// Note this does not necessarily get all fields, only the ones we use in
/// bottom.
pub(crate) struct Io {
pub read_bytes: u64,
pub write_bytes: u64,
}
impl Io {
#[inline]
fn from_file(f: File, buffer: &mut String) -> anyhow::Result<Io> {
const NUM_FIELDS: u16 = 0; // Make sure to update this if you want more fields!
enum Fields {
ReadBytes,
WriteBytes,
}
let mut read_fields = 0;
let mut reader = BufReader::new(f);
let mut read_bytes = 0;
let mut write_bytes = 0;
// This saves us from doing a string allocation on each iteration compared to
// `lines()`.
while let Ok(bytes) = reader.read_line(buffer) {
if bytes > 0 {
if buffer.is_empty() {
// Empty, no need to clear.
continue;
}
let mut parts = buffer.split_whitespace();
if let Some(field) = parts.next() {
let curr_field = match field {
"read_bytes:" => Fields::ReadBytes,
"write_bytes:" => Fields::WriteBytes,
_ => {
buffer.clear();
continue;
}
};
if let Some(value) = parts.next() {
let value = value.parse::<u64>()?;
match curr_field {
Fields::ReadBytes => {
read_bytes = value;
read_fields += 1;
}
Fields::WriteBytes => {
write_bytes = value;
read_fields += 1;
}
}
}
}
// Quick short circuit if we have already read all the required fields.
if read_fields == NUM_FIELDS {
break;
}
buffer.clear();
} else {
break;
}
}
Ok(Io {
read_bytes,
write_bytes,
})
}
}
/// A wrapper around a Linux process operations in `/proc/<PID>`.
///
/// Core documentation based on [proc's manpages](https://man7.org/linux/man-pages/man5/proc.5.html).
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>>,
}
#[inline]
fn reset(root: &mut PathBuf, buffer: &mut String) {
root.pop();
buffer.clear();
}
impl Process {
/// Creates a new [`Process`] given a `/proc/<PID>` path. This may fail if
/// the process no longer exists or there are permissions issues.
///
/// Note that this pre-allocates fields on **creation**! As such, some data
/// might end up "outdated" depending on when you call some of the
/// 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?
let fd = rustix::fs::openat(
rustix::fs::CWD,
&pid_path,
OFlags::PATH | OFlags::DIRECTORY | OFlags::CLOEXEC,
Mode::empty(),
)?;
let pid = pid_path
.as_path()
.components()
.last()
.and_then(|s| s.to_string_lossy().parse::<Pid>().ok())
.or_else(|| {
rustix::fs::readlinkat(rustix::fs::CWD, &pid_path, vec![])
.ok()
.and_then(|s| s.to_string_lossy().parse::<Pid>().ok())
})
.ok_or_else(|| anyhow!("PID for {pid_path:?} was not found"))?;
let uid = {
let metadata = rustix::fs::fstat(&fd);
match metadata {
Ok(md) => Some(md.st_uid),
Err(_) => None,
}
};
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!
let stat =
open_at(&mut root, "stat", &fd).and_then(|file| Stat::from_file(file, &mut buffer))?;
reset(&mut root, &mut buffer);
let cmdline = cmdline(&mut root, &fd, &mut buffer);
reset(&mut root, &mut buffer);
let io = open_at(&mut root, "io", &fd).and_then(|file| Io::from_file(file, &mut buffer));
Ok(Process {
pid,
uid,
stat,
io,
cmdline,
})
}
}
#[inline]
fn cmdline(root: &mut PathBuf, fd: &OwnedFd, buffer: &mut String) -> anyhow::Result<Vec<String>> {
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)
}
/// Opens a path. 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<File> {
root.push(child);
let new_fd = rustix::fs::openat(fd, &*root, OFlags::RDONLY | OFlags::CLOEXEC, Mode::empty())?;
Ok(File::from(new_fd))
}

View File

@ -8,15 +8,17 @@ use std::{
use anyhow::Result;
use hashbrown::{HashMap, HashSet};
use super::{TempHarvest, TemperatureType};
use crate::app::filter::Filter;
use crate::{
app::filter::Filter,
new_data_collection::sources::common::temperature::{TemperatureData, TemperatureType},
};
const EMPTY_NAME: &str = "Unknown";
/// Returned results from grabbing hwmon/coretemp temperature sensor
/// values/names.
struct HwmonResults {
temperatures: Vec<TempHarvest>,
temperatures: Vec<TemperatureData>,
num_hwmon: usize,
}
@ -224,7 +226,7 @@ fn is_device_awake(path: &Path) -> bool {
/// once this happens, the device will be *kept* on through the sensor
/// reading, and not be able to re-enter ACPI D3cold.
fn hwmon_temperatures(temp_type: &TemperatureType, filter: &Option<Filter>) -> HwmonResults {
let mut temperatures: Vec<TempHarvest> = vec![];
let mut temperatures: Vec<TemperatureData> = vec![];
let mut seen_names: HashMap<String, u32> = HashMap::new();
let (dirs, num_hwmon) = get_hwmon_candidates();
@ -246,7 +248,7 @@ fn hwmon_temperatures(temp_type: &TemperatureType, filter: &Option<Filter>) -> H
if !is_device_awake(&file_path) {
let name = finalize_name(None, None, &sensor_name, &mut seen_names);
temperatures.push(TempHarvest {
temperatures.push(TemperatureData {
name,
temperature: None,
});
@ -329,7 +331,7 @@ fn hwmon_temperatures(temp_type: &TemperatureType, filter: &Option<Filter>) -> H
// probing hwmon if not needed?
if Filter::optional_should_keep(filter, &name) {
if let Ok(temp_celsius) = parse_temp(&temp_path) {
temperatures.push(TempHarvest {
temperatures.push(TemperatureData {
name,
temperature: Some(temp_type.convert_temp_unit(temp_celsius)),
});
@ -352,7 +354,7 @@ fn hwmon_temperatures(temp_type: &TemperatureType, filter: &Option<Filter>) -> H
/// See [the Linux kernel documentation](https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-class-thermal)
/// for more details.
fn add_thermal_zone_temperatures(
temperatures: &mut Vec<TempHarvest>, temp_type: &TemperatureType, filter: &Option<Filter>,
temperatures: &mut Vec<TemperatureData>, temp_type: &TemperatureType, filter: &Option<Filter>,
) {
let path = Path::new("/sys/class/thermal");
let Ok(read_dir) = path.read_dir() else {
@ -382,7 +384,7 @@ fn add_thermal_zone_temperatures(
if let Ok(temp_celsius) = parse_temp(&temp_path) {
let name = counted_name(&mut seen_names, name);
temperatures.push(TempHarvest {
temperatures.push(TemperatureData {
name,
temperature: Some(temp_type.convert_temp_unit(temp_celsius)),
});
@ -396,14 +398,14 @@ fn add_thermal_zone_temperatures(
/// Gets temperature sensors and data.
pub fn get_temperature_data(
temp_type: &TemperatureType, filter: &Option<Filter>,
) -> Result<Option<Vec<TempHarvest>>> {
) -> Vec<TemperatureData> {
let mut results = hwmon_temperatures(temp_type, filter);
if results.num_hwmon == 0 {
add_thermal_zone_temperatures(&mut results.temperatures, temp_type, filter);
}
Ok(Some(results.temperatures))
results.temperatures
}
#[cfg(test)]

View File

@ -1,14 +1,16 @@
//! Re-exports all of the sources.
pub mod common;
pub mod linux;
pub mod macos;
#[cfg(feature = "gpu")]
pub mod nvidia;
pub mod sysinfo;
pub mod unix;
pub mod windows;
cfg_if::cfg_if! {
if #[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux",
target_os = "freebsd",
target_os = "dragonfly",
target_os = "ios",
))] {
pub mod starship_battery;
if #[cfg(target_family = "windows")] {
pub use windows::processes::Pid as Pid;
} else if #[cfg(target_family = "unix")] {
pub use unix::processes::Pid as Pid;
}
}

View File

@ -0,0 +1 @@
pub mod temperature;

View File

@ -1,51 +1,26 @@
//! Gets temperature data via sysinfo.
use anyhow::Result;
use super::{TempHarvest, TemperatureType};
use crate::app::filter::Filter;
use crate::{
app::filter::Filter,
new_data_collection::sources::common::temperature::{TemperatureData, TemperatureType},
};
pub fn get_temperature_data(
components: &sysinfo::Components, temp_type: &TemperatureType, filter: &Option<Filter>,
) -> Result<Option<Vec<TempHarvest>>> {
let mut temperature_vec: Vec<TempHarvest> = Vec::new();
) -> Vec<TemperatureData> {
let mut temperature_vec: Vec<TemperatureData> = Vec::new();
for component in components {
let name = component.label().to_string();
if Filter::optional_should_keep(filter, &name) {
temperature_vec.push(TempHarvest {
temperature_vec.push(TemperatureData {
name,
temperature: Some(temp_type.convert_temp_unit(component.temperature())),
});
}
}
// For RockPro64 boards on FreeBSD, they apparently use "hw.temperature" for
// sensors.
#[cfg(target_os = "freebsd")]
{
use sysctl::Sysctl;
const KEY: &str = "hw.temperature";
if let Ok(root) = sysctl::Ctl::new(KEY) {
for ctl in sysctl::CtlIter::below(root).flatten() {
if let (Ok(name), Ok(temp)) = (ctl.name(), ctl.value()) {
if let Some(temp) = temp.as_temperature() {
temperature_vec.push(TempHarvest {
name,
temperature: Some(match temp_type {
TemperatureType::Celsius => temp.celsius(),
TemperatureType::Kelvin => temp.kelvin(),
TemperatureType::Fahrenheit => temp.fahrenheit(),
}),
});
}
}
}
}
}
// TODO: Should we instead use a hashmap -> vec to skip dupes?
Ok(Some(temperature_vec))
temperature_vec
}

View File

@ -0,0 +1 @@
pub mod processes;

View File

@ -0,0 +1,5 @@
pub mod user_table;
/// A UNIX process ID.
#[cfg(target_family = "unix")]
pub type Pid = libc::pid_t;

View File

@ -0,0 +1,33 @@
use hashbrown::HashMap;
use crate::data_collection::error::{CollectionError, CollectionResult};
#[derive(Debug, Default)]
pub struct UserTable {
pub uid_user_mapping: HashMap<libc::uid_t, String>,
}
impl UserTable {
pub fn get_uid_to_username_mapping(&mut self, uid: libc::uid_t) -> CollectionResult<String> {
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
let passwd = unsafe { libc::getpwuid(uid) };
if passwd.is_null() {
Err("passwd is inaccessible".into())
} else {
// SAFETY: We return early if passwd is null.
let username = unsafe { std::ffi::CStr::from_ptr((*passwd).pw_name) }
.to_str()
.map_err(|err| CollectionError::General(err.into()))?
.to_string();
self.uid_user_mapping.insert(uid, username.clone());
Ok(username)
}
}
}
}

View File

@ -0,0 +1 @@
pub mod processes;

View File

@ -0,0 +1,3 @@
/// A Windows process ID.
#[cfg(target_family = "windows")]
pub type Pid = usize;