From fe7932047098ca6c2c2ba334052532b829cb7a3e Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Sun, 11 Aug 2024 17:38:03 -0400 Subject: [PATCH 1/7] refactor: move data_collection.rs to mod.rs --- src/{data_collection.rs => data_collection/mod.rs} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/{data_collection.rs => data_collection/mod.rs} (100%) diff --git a/src/data_collection.rs b/src/data_collection/mod.rs similarity index 100% rename from src/data_collection.rs rename to src/data_collection/mod.rs From ac4d1a0cb1e4ae4de8f864542e9a24406cb584fd Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Sun, 11 Aug 2024 18:22:41 -0400 Subject: [PATCH 2/7] add readme for data collection --- src/data_collection/README.md | 44 +++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 src/data_collection/README.md diff --git a/src/data_collection/README.md b/src/data_collection/README.md new file mode 100644 index 000000000..1d8a955a8 --- /dev/null +++ b/src/data_collection/README.md @@ -0,0 +1,44 @@ +# Data Collection + +Data collection in bottom has two main components: **sources** and **collectors**. + +**Sources** are either libraries or system APIs that actually extract the data. +These may map to multiple different operating systems. Examples are `sysinfo`, +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 +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): + +```mermaid +flowchart TB + subgraph sources + direction TB + linux + windows + macos + unix + sysinfo + freebsd + end + subgraph collectors + direction TB + Linux + Windows + macOS + FreeBSD + end + linux -..-> Linux + unix -..-> Linux + sysinfo -..-> Linux + windows -..-> Windows + sysinfo -..-> Windows + macos -..-> macOS + unix -..-> macOS + sysinfo -..-> macOS + freebsd -..-> FreeBSD + sysinfo -..-> FreeBSD +``` From 9652a262cdd142d221975bd317511e2055f6443e Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Sun, 11 Aug 2024 19:44:04 -0400 Subject: [PATCH 3/7] create new data collection directory --- src/main.rs | 1 + src/new_data_collection/collectors/freebsd.rs | 0 src/new_data_collection/collectors/linux.rs | 0 src/new_data_collection/collectors/macos.rs | 0 src/new_data_collection/collectors/mod.rs | 0 src/new_data_collection/collectors/windows.rs | 0 src/new_data_collection/mod.rs | 0 .../sources/common/temperature.rs | 0 .../sources/linux/temperature.rs | 491 ++++++++++++++++++ src/new_data_collection/sources/mod.rs | 14 + .../sources/starship_battery.rs | 50 ++ .../sources/sysinfo/temperature.rs | 51 ++ 12 files changed, 607 insertions(+) create mode 100644 src/new_data_collection/collectors/freebsd.rs create mode 100644 src/new_data_collection/collectors/linux.rs create mode 100644 src/new_data_collection/collectors/macos.rs create mode 100644 src/new_data_collection/collectors/mod.rs create mode 100644 src/new_data_collection/collectors/windows.rs create mode 100644 src/new_data_collection/mod.rs create mode 100644 src/new_data_collection/sources/common/temperature.rs create mode 100644 src/new_data_collection/sources/linux/temperature.rs create mode 100644 src/new_data_collection/sources/mod.rs create mode 100644 src/new_data_collection/sources/starship_battery.rs create mode 100644 src/new_data_collection/sources/sysinfo/temperature.rs diff --git a/src/main.rs b/src/main.rs index be0799766..16561bb6f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,6 +21,7 @@ pub mod constants; pub mod data_collection; pub mod data_conversion; pub mod event; +pub mod new_data_collection; pub mod options; pub mod widgets; diff --git a/src/new_data_collection/collectors/freebsd.rs b/src/new_data_collection/collectors/freebsd.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/collectors/linux.rs b/src/new_data_collection/collectors/linux.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/collectors/macos.rs b/src/new_data_collection/collectors/macos.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/collectors/mod.rs b/src/new_data_collection/collectors/mod.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/collectors/windows.rs b/src/new_data_collection/collectors/windows.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/mod.rs b/src/new_data_collection/mod.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/sources/common/temperature.rs b/src/new_data_collection/sources/common/temperature.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/sources/linux/temperature.rs b/src/new_data_collection/sources/linux/temperature.rs new file mode 100644 index 000000000..4aaf1f52b --- /dev/null +++ b/src/new_data_collection/sources/linux/temperature.rs @@ -0,0 +1,491 @@ +//! Gets temperature sensor data for Linux platforms. + +use std::{ + fs, + path::{Path, PathBuf}, +}; + +use anyhow::Result; +use hashbrown::{HashMap, HashSet}; + +use super::{TempHarvest, TemperatureType}; +use crate::app::filter::Filter; + +const EMPTY_NAME: &str = "Unknown"; + +/// Returned results from grabbing hwmon/coretemp temperature sensor +/// values/names. +struct HwmonResults { + temperatures: Vec, + num_hwmon: usize, +} + +/// Parses and reads temperatures that were in millidegree Celsius, and if +/// successful, returns a temperature in Celsius. +fn parse_temp(path: &Path) -> Result { + Ok(fs::read_to_string(path)?.trim_end().parse::()? / 1_000.0) +} + +/// Get all candidates from hwmon and coretemp. It will also return the number +/// of entries from hwmon. +fn get_hwmon_candidates() -> (HashSet, usize) { + let mut dirs = HashSet::default(); + + if let Ok(read_dir) = Path::new("/sys/class/hwmon").read_dir() { + for entry in read_dir.flatten() { + let mut path = entry.path(); + + // hwmon includes many sensors, we only want ones with at least one temperature + // sensor Reading this file will wake the device, but we're only + // checking existence, so it should be fine. + if !path.join("temp1_input").exists() { + // Note we also check for a `device` subdirectory (e.g. + // `/sys/class/hwmon/hwmon*/device/`). This is needed for + // CentOS, which adds this extra `/device` directory. See: + // - https://github.com/nicolargo/glances/issues/1060 + // - https://github.com/giampaolo/psutil/issues/971 + // - https://github.com/giampaolo/psutil/blob/642438375e685403b4cd60b0c0e25b80dd5a813d/psutil/_pslinux.py#L1316 + // + // If it does match, then add the `device/` directory to the path. + if path.join("device/temp1_input").exists() { + path.push("device"); + } + } + + dirs.insert(path); + } + } + + let num_hwmon = dirs.len(); + + if let Ok(read_dir) = Path::new("/sys/devices/platform").read_dir() { + for entry in read_dir.flatten() { + if entry.file_name().to_string_lossy().starts_with("coretemp.") { + if let Ok(read_dir) = entry.path().join("hwmon").read_dir() { + for entry in read_dir.flatten() { + let path = entry.path(); + + if path.join("temp1_input").exists() { + // It's possible that there are dupes (represented by symlinks) - the + // easy way is to just substitute the parent + // directory and check if the hwmon + // variant exists already in a set. + // + // For more info, see https://github.com/giampaolo/psutil/pull/1822/files + if let Some(child) = path.file_name() { + let to_check_path = Path::new("/sys/class/hwmon").join(child); + + if !dirs.contains(&to_check_path) { + dirs.insert(path); + } + } + } + } + } + } + } + } + + (dirs, num_hwmon) +} + +#[inline] +fn read_to_string_lossy>(path: P) -> Option { + fs::read(path) + .map(|v| String::from_utf8_lossy(&v).trim().to_string()) + .ok() +} + +#[inline] +fn humanize_name(name: String, sensor_name: Option<&String>) -> String { + match sensor_name { + Some(ty) => format!("{name} ({ty})"), + None => name, + } +} + +#[inline] +fn counted_name(seen_names: &mut HashMap, name: String) -> String { + if let Some(count) = seen_names.get_mut(&name) { + *count += 1; + format!("{name} ({count})") + } else { + seen_names.insert(name.clone(), 0); + name + } +} + +#[inline] +fn finalize_name( + hwmon_name: Option, sensor_label: Option, + fallback_sensor_name: &Option, seen_names: &mut HashMap, +) -> String { + let candidate_name = match (hwmon_name, sensor_label) { + (Some(name), Some(label)) => match (name.is_empty(), label.is_empty()) { + (false, false) => { + format!("{name}: {label}") + } + (true, false) => match fallback_sensor_name { + Some(fallback) if !fallback.is_empty() => { + if label.is_empty() { + fallback.to_owned() + } else { + format!("{fallback}: {label}") + } + } + _ => { + if label.is_empty() { + EMPTY_NAME.to_string() + } else { + label + } + } + }, + (false, true) => name.to_owned(), + (true, true) => EMPTY_NAME.to_string(), + }, + (None, Some(label)) => match fallback_sensor_name { + Some(fallback) if !fallback.is_empty() => { + if label.is_empty() { + fallback.to_owned() + } else { + format!("{fallback}: {label}") + } + } + _ => { + if label.is_empty() { + EMPTY_NAME.to_string() + } else { + label + } + } + }, + (Some(name), None) => { + if name.is_empty() { + EMPTY_NAME.to_string() + } else { + name + } + } + (None, None) => match fallback_sensor_name { + Some(sensor_name) if !sensor_name.is_empty() => sensor_name.to_owned(), + _ => EMPTY_NAME.to_string(), + }, + }; + + counted_name(seen_names, candidate_name) +} + +/// 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 +/// `device/power_state`. +#[inline] +fn is_device_awake(path: &Path) -> bool { + // Whether the temperature should *actually* be read during enumeration. + // Set to false if the device is in ACPI D3cold. + // Documented at https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-devices-power_state + let device = path.join("device"); + let power_state = device.join("power_state"); + if power_state.exists() { + if let Ok(state) = fs::read_to_string(power_state) { + let state = state.trim(); + // The zenpower3 kernel module (incorrectly?) reports "unknown", causing this + // check to fail and temperatures to appear as zero instead of + // having the file not exist. + // + // Their self-hosted git instance has disabled sign up, so this bug cant be + // reported either. + state == "D0" || state == "unknown" + } else { + true + } + } else { + true + } +} + +/// Get temperature sensors from the linux sysfs interface `/sys/class/hwmon` +/// and `/sys/devices/platform/coretemp.*`. It returns all found temperature +/// sensors, and the number of checked hwmon directories (not coretemp +/// directories). +/// +/// For more details, see the relevant Linux kernel documentation: +/// - [`/sys/class/hwmon`](https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-class-hwmon) +/// - [`/sys/devices/platform/coretemp.*`](https://www.kernel.org/doc/html/v5.14/hwmon/coretemp.html) +/// +/// This method will return `0` as the temperature for devices, such as GPUs, +/// that support power management features that have powered themselves off. +/// Specifically, in laptops with iGPUs and dGPUs, if the dGPU is capable of +/// entering ACPI D3cold, reading the temperature sensors will wake it, +/// and keep it awake, wasting power. +/// +/// For such devices, this method will only query the sensors *only* if +/// the device is already in ACPI D0. This has the notable issue that +/// 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) -> HwmonResults { + let mut temperatures: Vec = vec![]; + let mut seen_names: HashMap = HashMap::new(); + + let (dirs, num_hwmon) = get_hwmon_candidates(); + + // Note that none of this is async if we ever go back to it, but sysfs is in + // memory, so in theory none of this should block if we're slightly careful. + // Of note is that reading the temperature sensors of a device that has + // `/sys/class/hwmon/hwmon*/device/power_state` == `D3cold` will + // wake the device up, and will block until it initializes. + // + // Reading the `hwmon*/device/power_state` or `hwmon*/temp*_label` properties + // will not wake the device, and thus not block, + // and meaning no sensors have to be hidden depending on `power_state` + // + // It would probably be more ideal to use a proper async runtime; this would + // also allow easy cancellation/timeouts. + for file_path in dirs { + let sensor_name = read_to_string_lossy(file_path.join("name")); + + if !is_device_awake(&file_path) { + let name = finalize_name(None, None, &sensor_name, &mut seen_names); + temperatures.push(TempHarvest { + name, + temperature: None, + }); + + continue; + } + + if let Ok(dir_entries) = file_path.read_dir() { + // Enumerate the devices temperature sensors + for file in dir_entries.flatten() { + let name = file.file_name(); + let name = name.to_string_lossy(); + + // We only want temperature sensors, skip others early + if !(name.starts_with("temp") && name.ends_with("input")) { + continue; + } + + let temp_path = file.path(); + let sensor_label_path = file_path.join(name.replace("input", "label")); + let sensor_label = read_to_string_lossy(sensor_label_path); + + // Do some messing around to get a more sensible name for sensors: + // - For GPUs, this will use the kernel device name, ex `card0` + // - For nvme drives, this will also use the kernel name, ex `nvme0`. This is + // found differently than for GPUs + // - For whatever acpitz is, on my machine this is now `thermal_zone0`. + // - For k10temp, this will still be k10temp, but it has to be handled special. + let hwmon_name = { + let device = file_path.join("device"); + + // This will exist for GPUs but not others, this is how we find their kernel + // name. + let drm = device.join("drm"); + if drm.exists() { + // This should never actually be empty. If it is though, we'll fall back to + // the sensor name later on. + let mut gpu = None; + + if let Ok(cards) = drm.read_dir() { + for card in cards.flatten() { + if let Some(name) = card.file_name().to_str() { + if name.starts_with("card") { + gpu = Some(humanize_name( + name.trim().to_string(), + sensor_name.as_ref(), + )); + break; + } + } + } + } + + gpu + } else { + // This little mess is to account for stuff like k10temp. This is needed + // because the `device` symlink points to `nvme*` + // for nvme drives, but to PCI buses for anything + // else. If the first character is alphabetic, it's an actual name like + // k10temp or nvme0, not a PCI bus. + fs::read_link(device).ok().and_then(|link| { + let link = link + .file_name() + .and_then(|f| f.to_str()) + .map(|s| s.trim().to_owned()); + + match link { + Some(link) if link.as_bytes()[0].is_ascii_alphabetic() => { + Some(humanize_name(link, sensor_name.as_ref())) + } + _ => None, + } + }) + } + }; + + let name = finalize_name(hwmon_name, sensor_label, &sensor_name, &mut seen_names); + + // TODO: It's possible we may want to move the filter check further up to avoid + // probing hwmon if not needed? + if Filter::optional_should_keep(filter, &name) { + if let Ok(temp_celsius) = parse_temp(&temp_path) { + temperatures.push(TempHarvest { + name, + temperature: Some(temp_type.convert_temp_unit(temp_celsius)), + }); + } + } + } + } + } + + HwmonResults { + temperatures, + num_hwmon, + } +} + +/// Gets data from `/sys/class/thermal/thermal_zone*`. This should only be used +/// if [`hwmon_temperatures`] doesn't return anything to avoid duplicate sensor +/// results. +/// +/// 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, temp_type: &TemperatureType, filter: &Option, +) { + let path = Path::new("/sys/class/thermal"); + let Ok(read_dir) = path.read_dir() else { + return; + }; + + let mut seen_names: HashMap = HashMap::new(); + + for entry in read_dir.flatten() { + if entry + .file_name() + .to_string_lossy() + .starts_with("thermal_zone") + { + let file_path = entry.path(); + let name_path = file_path.join("type"); + + if let Some(name) = read_to_string_lossy(name_path) { + let name = if name.is_empty() { + EMPTY_NAME.to_string() + } else { + name + }; + + if Filter::optional_should_keep(filter, &name) { + let temp_path = file_path.join("temp"); + if let Ok(temp_celsius) = parse_temp(&temp_path) { + let name = counted_name(&mut seen_names, name); + + temperatures.push(TempHarvest { + name, + temperature: Some(temp_type.convert_temp_unit(temp_celsius)), + }); + } + } + } + } + } +} + +/// Gets temperature sensors and data. +pub fn get_temperature_data( + temp_type: &TemperatureType, filter: &Option, +) -> Result>> { + 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)) +} + +#[cfg(test)] +mod tests { + use hashbrown::HashMap; + + use super::finalize_name; + + #[test] + fn test_finalize_name() { + let mut seen_names = HashMap::new(); + + assert_eq!( + finalize_name( + Some("hwmon".to_string()), + Some("sensor".to_string()), + &Some("test".to_string()), + &mut seen_names + ), + "hwmon: sensor" + ); + + assert_eq!( + finalize_name( + Some("hwmon".to_string()), + None, + &Some("test".to_string()), + &mut seen_names + ), + "hwmon" + ); + + assert_eq!( + finalize_name( + None, + Some("sensor".to_string()), + &Some("test".to_string()), + &mut seen_names + ), + "test: sensor" + ); + + assert_eq!( + finalize_name( + Some("hwmon".to_string()), + Some("sensor".to_string()), + &Some("test".to_string()), + &mut seen_names + ), + "hwmon: sensor (1)" + ); + + assert_eq!( + finalize_name(None, None, &Some("test".to_string()), &mut seen_names), + "test" + ); + + assert_eq!(finalize_name(None, None, &None, &mut seen_names), "Unknown"); + + assert_eq!( + finalize_name(None, None, &Some("test".to_string()), &mut seen_names), + "test (1)" + ); + + assert_eq!( + finalize_name(None, None, &None, &mut seen_names), + "Unknown (1)" + ); + + assert_eq!( + finalize_name(Some(String::default()), None, &None, &mut seen_names), + "Unknown (2)" + ); + + assert_eq!( + finalize_name(None, Some(String::default()), &None, &mut seen_names), + "Unknown (3)" + ); + + assert_eq!( + finalize_name(None, None, &Some(String::default()), &mut seen_names), + "Unknown (4)" + ); + } +} diff --git a/src/new_data_collection/sources/mod.rs b/src/new_data_collection/sources/mod.rs new file mode 100644 index 000000000..5f3015811 --- /dev/null +++ b/src/new_data_collection/sources/mod.rs @@ -0,0 +1,14 @@ +//! Re-exports all of the sources. + +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; + } +} diff --git a/src/new_data_collection/sources/starship_battery.rs b/src/new_data_collection/sources/starship_battery.rs new file mode 100644 index 000000000..1235e9384 --- /dev/null +++ b/src/new_data_collection/sources/starship_battery.rs @@ -0,0 +1,50 @@ +//! Covers battery usage for: +//! - Linux 2.6.39+ +//! - MacOS 10.10+ +//! - iOS +//! - Windows 7+ +//! - FreeBSD +//! - DragonFlyBSD +//! +//! For more information, refer to the [starship_battery](https://github.com/starship/rust-battery) repo/docs. + +use starship_battery::{ + units::{power::watt, ratio::percent, time::second}, + Battery, Manager, State, +}; + +#[derive(Debug, Clone)] +pub struct BatteryHarvest { + pub charge_percent: f64, + pub secs_until_full: Option, + pub secs_until_empty: Option, + pub power_consumption_rate_watts: f64, + pub health_percent: f64, + pub state: State, +} + +pub fn refresh_batteries(manager: &Manager, batteries: &mut [Battery]) -> Vec { + batteries + .iter_mut() + .filter_map(|battery| { + if manager.refresh(battery).is_ok() { + Some(BatteryHarvest { + secs_until_full: { + let optional_time = battery.time_to_full(); + optional_time.map(|time| f64::from(time.get::()) as i64) + }, + secs_until_empty: { + let optional_time = battery.time_to_empty(); + optional_time.map(|time| f64::from(time.get::()) as i64) + }, + charge_percent: f64::from(battery.state_of_charge().get::()), + power_consumption_rate_watts: f64::from(battery.energy_rate().get::()), + health_percent: f64::from(battery.state_of_health().get::()), + state: battery.state(), + }) + } else { + None + } + }) + .collect::>() +} diff --git a/src/new_data_collection/sources/sysinfo/temperature.rs b/src/new_data_collection/sources/sysinfo/temperature.rs new file mode 100644 index 000000000..d1a9e634b --- /dev/null +++ b/src/new_data_collection/sources/sysinfo/temperature.rs @@ -0,0 +1,51 @@ +//! Gets temperature data via sysinfo. + +use anyhow::Result; + +use super::{TempHarvest, TemperatureType}; +use crate::app::filter::Filter; + +pub fn get_temperature_data( + components: &sysinfo::Components, temp_type: &TemperatureType, filter: &Option, +) -> Result>> { + let mut temperature_vec: Vec = Vec::new(); + + for component in components { + let name = component.label().to_string(); + + if Filter::optional_should_keep(filter, &name) { + temperature_vec.push(TempHarvest { + 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)) +} From 890d0aeedcf3f3497e4869a33facccf17abf267a Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Wed, 14 Aug 2024 22:36:49 -0400 Subject: [PATCH 4/7] initial side-by-side impl for linux temmp and process as a poc --- src/data_collection/batteries.rs | 9 +- .../README.md | 13 +- src/new_data_collection/collectors/common.rs | 27 ++ .../collectors/fallback.rs | 7 + src/new_data_collection/collectors/freebsd.rs | 36 ++ src/new_data_collection/collectors/linux.rs | 71 +++ src/new_data_collection/collectors/macos.rs | 33 ++ src/new_data_collection/collectors/windows.rs | 33 ++ src/new_data_collection/error.rs | 42 ++ src/new_data_collection/mod.rs | 28 ++ src/new_data_collection/sources/common/mod.rs | 2 + .../sources/common/processes.rs | 87 ++++ .../sources/common/temperature.rs | 70 +++ .../sources/freebsd/mod.rs | 1 + .../sources/freebsd/temperature.rs | 35 ++ src/new_data_collection/sources/linux/mod.rs | 2 + .../sources/linux/processes/mod.rs | 450 ++++++++++++++++++ .../sources/linux/processes/process.rs | 308 ++++++++++++ .../sources/linux/temperature.rs | 22 +- .../{collectors => sources/macos}/mod.rs | 0 src/new_data_collection/sources/mod.rs | 22 +- src/new_data_collection/sources/nvidia/mod.rs | 0 .../sources/sysinfo/mod.rs | 1 + .../sources/sysinfo/temperature.rs | 41 +- src/new_data_collection/sources/unix/mod.rs | 1 + .../sources/unix/processes/mod.rs | 5 + .../sources/unix/processes/user_table.rs | 33 ++ .../sources/windows/mod.rs | 1 + .../sources/windows/processes.rs | 3 + 29 files changed, 1328 insertions(+), 55 deletions(-) rename src/{data_collection => new_data_collection}/README.md (68%) create mode 100644 src/new_data_collection/collectors/common.rs create mode 100644 src/new_data_collection/collectors/fallback.rs create mode 100644 src/new_data_collection/error.rs create mode 100644 src/new_data_collection/sources/common/mod.rs create mode 100644 src/new_data_collection/sources/common/processes.rs create mode 100644 src/new_data_collection/sources/freebsd/mod.rs create mode 100644 src/new_data_collection/sources/freebsd/temperature.rs create mode 100644 src/new_data_collection/sources/linux/mod.rs create mode 100644 src/new_data_collection/sources/linux/processes/mod.rs create mode 100644 src/new_data_collection/sources/linux/processes/process.rs rename src/new_data_collection/{collectors => sources/macos}/mod.rs (100%) create mode 100644 src/new_data_collection/sources/nvidia/mod.rs create mode 100644 src/new_data_collection/sources/sysinfo/mod.rs create mode 100644 src/new_data_collection/sources/unix/mod.rs create mode 100644 src/new_data_collection/sources/unix/processes/mod.rs create mode 100644 src/new_data_collection/sources/unix/processes/user_table.rs create mode 100644 src/new_data_collection/sources/windows/mod.rs create mode 100644 src/new_data_collection/sources/windows/processes.rs diff --git a/src/data_collection/batteries.rs b/src/data_collection/batteries.rs index a155ad2d2..0bdab1c6d 100644 --- a/src/data_collection/batteries.rs +++ b/src/data_collection/batteries.rs @@ -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::*; } diff --git a/src/data_collection/README.md b/src/new_data_collection/README.md similarity index 68% rename from src/data_collection/README.md rename to src/new_data_collection/README.md index 1d8a955a8..f06eaa70b 100644 --- a/src/data_collection/README.md +++ b/src/new_data_collection/README.md @@ -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. diff --git a/src/new_data_collection/collectors/common.rs b/src/new_data_collection/collectors/common.rs new file mode 100644 index 000000000..6d7e507a9 --- /dev/null +++ b/src/new_data_collection/collectors/common.rs @@ -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> { + Ok(vec![]) + } + + /// Return process data. + fn get_process_data(&mut self) -> CollectionResult> { + Ok(vec![]) + } +} diff --git a/src/new_data_collection/collectors/fallback.rs b/src/new_data_collection/collectors/fallback.rs new file mode 100644 index 000000000..a445a8c00 --- /dev/null +++ b/src/new_data_collection/collectors/fallback.rs @@ -0,0 +1,7 @@ +use super::common::DataCollector; + +/// A fallback [`DataCollector`] for unsupported systems +/// that does nothing. +pub struct FallbackDataCollector {} + +impl DataCollector for FallbackDataCollector {} diff --git a/src/new_data_collection/collectors/freebsd.rs b/src/new_data_collection/collectors/freebsd.rs index e69de29bb..ccf5aad9b 100644 --- a/src/new_data_collection/collectors/freebsd.rs +++ b/src/new_data_collection/collectors/freebsd.rs @@ -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, +} + +impl DataCollector for FreeBsdDataCollector { + fn refresh_data(&mut self) -> CollectionResult<()> { + Ok(()) + } + + fn get_temperature_data(&self) -> CollectionResult>> { + 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)) + } +} diff --git a/src/new_data_collection/collectors/linux.rs b/src/new_data_collection/collectors/linux.rs index e69de29bb..4a7191a4e 100644 --- a/src/new_data_collection/collectors/linux.rs +++ b/src/new_data_collection/collectors/linux.rs @@ -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, + + proc_collector: ProcessCollector, + + system: sysinfo::System, + network: sysinfo::Networks, + + #[cfg(feature = "battery")] + battery_manager: Option, + #[cfg(feature = "battery")] + battery_list: Option>, + + #[cfg(feature = "gpu")] + gpus_total_mem: Option, +} + +impl DataCollector for LinuxDataCollector { + fn refresh_data(&mut self) -> CollectionResult<()> { + Ok(()) + } + + fn get_temperature_data(&mut self) -> CollectionResult> { + Ok(get_temperature_data(&self.temp_type, &self.temp_filters)) + } + + fn get_process_data(&mut self) -> CollectionResult> { + 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, + ) + } +} diff --git a/src/new_data_collection/collectors/macos.rs b/src/new_data_collection/collectors/macos.rs index e69de29bb..5a4efb745 100644 --- a/src/new_data_collection/collectors/macos.rs +++ b/src/new_data_collection/collectors/macos.rs @@ -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, +} + +impl DataCollector for MacOsDataCollector { + fn refresh_data(&mut self) -> CollectionResult<()> { + Ok(()) + } + + fn get_temperature_data(&self) -> CollectionResult>> { + Ok(Some(get_temperature_data( + &self.temp_type, + &self.temp_filters, + ))) + } +} diff --git a/src/new_data_collection/collectors/windows.rs b/src/new_data_collection/collectors/windows.rs index e69de29bb..607d94663 100644 --- a/src/new_data_collection/collectors/windows.rs +++ b/src/new_data_collection/collectors/windows.rs @@ -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, +} + +impl DataCollector for WindowsDataCollector { + fn refresh_data(&mut self) -> CollectionResult<()> { + Ok(()) + } + + fn get_temperature_data(&self) -> CollectionResult>> { + Ok(Some(get_temperature_data( + &self.temp_type, + &self.temp_filters, + ))) + } +} diff --git a/src/new_data_collection/error.rs b/src/new_data_collection/error.rs new file mode 100644 index 000000000..dc5dd7bd3 --- /dev/null +++ b/src/new_data_collection/error.rs @@ -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 = Result; + +impl From 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)) + } +} diff --git a/src/new_data_collection/mod.rs b/src/new_data_collection/mod.rs index e69de29bb..8f9a2cef1 100644 --- a/src/new_data_collection/mod.rs +++ b/src/new_data_collection/mod.rs @@ -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; diff --git a/src/new_data_collection/sources/common/mod.rs b/src/new_data_collection/sources/common/mod.rs new file mode 100644 index 000000000..d918321f6 --- /dev/null +++ b/src/new_data_collection/sources/common/mod.rs @@ -0,0 +1,2 @@ +pub mod processes; +pub mod temperature; diff --git a/src/new_data_collection/sources/common/processes.rs b/src/new_data_collection/sources/common/processes.rs new file mode 100644 index 000000000..c556fcc97 --- /dev/null +++ b/src/new_data_collection/sources/common/processes.rs @@ -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, + + /// 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, + + /// 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; + } + } +} diff --git a/src/new_data_collection/sources/common/temperature.rs b/src/new_data_collection/sources/common/temperature.rs index e69de29bb..abae51ac6 100644 --- a/src/new_data_collection/sources/common/temperature.rs +++ b/src/new_data_collection/sources/common/temperature.rs @@ -0,0 +1,70 @@ +use std::str::FromStr; + +#[derive(Default, Debug, Clone)] +pub struct TemperatureData { + pub name: String, + pub temperature: Option, +} + +#[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 { + 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); + } +} diff --git a/src/new_data_collection/sources/freebsd/mod.rs b/src/new_data_collection/sources/freebsd/mod.rs new file mode 100644 index 000000000..d7f95f24f --- /dev/null +++ b/src/new_data_collection/sources/freebsd/mod.rs @@ -0,0 +1 @@ +mod temperature; diff --git a/src/new_data_collection/sources/freebsd/temperature.rs b/src/new_data_collection/sources/freebsd/temperature.rs new file mode 100644 index 000000000..c9058f15d --- /dev/null +++ b/src/new_data_collection/sources/freebsd/temperature.rs @@ -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, +) -> impl Iterator { + 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() + } +} diff --git a/src/new_data_collection/sources/linux/mod.rs b/src/new_data_collection/sources/linux/mod.rs new file mode 100644 index 000000000..d918321f6 --- /dev/null +++ b/src/new_data_collection/sources/linux/mod.rs @@ -0,0 +1,2 @@ +pub mod processes; +pub mod temperature; diff --git a/src/new_data_collection/sources/linux/processes/mod.rs b/src/new_data_collection/sources/linux/processes/mod.rs new file mode 100644 index 000000000..ef978a8e0 --- /dev/null +++ b/src/new_data_collection/sources/linux/processes/mod.rs @@ -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//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::().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 { + 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, + pub user_table: UserTable, + + #[cfg(feature = "gpu")] + pub gpu_pids: Option>>, +} + +pub(crate) fn linux_process_data( + system: &sysinfo::System, time_difference_in_secs: u64, collector: &mut ProcessCollector, + #[cfg(feature = "gpu")] gpus_total_mem: Option, +) -> CollectionResult> { + 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 = 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 = 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" + ); + } +} diff --git a/src/new_data_collection/sources/linux/processes/process.rs b/src/new_data_collection/sources/linux/processes/process.rs new file mode 100644 index 000000000..45ed80a67 --- /dev/null +++ b/src/new_data_collection/sources/linux/processes/process.rs @@ -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 = OnceLock::new(); + +#[inline] +fn next_part<'a>(iter: &mut impl Iterator) -> Result<&'a str, io::Error> { + iter.next() + .ok_or_else(|| io::Error::from(io::ErrorKind::InvalidData)) +} + +/// A wrapper around the data in `/proc//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 { + // 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//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 { + 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::()?; + 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/`. +/// +/// 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, + pub stat: Stat, + pub io: anyhow::Result, + pub cmdline: anyhow::Result>, +} + +#[inline] +fn reset(root: &mut PathBuf, buffer: &mut String) { + root.pop(); + buffer.clear(); +} + +impl Process { + /// Creates a new [`Process`] given a `/proc/` 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 { + // 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::().ok()) + .or_else(|| { + rustix::fs::readlinkat(rustix::fs::CWD, &pid_path, vec![]) + .ok() + .and_then(|s| s.to_string_lossy().parse::().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> { + 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::>() + }) + .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 { + root.push(child); + let new_fd = rustix::fs::openat(fd, &*root, OFlags::RDONLY | OFlags::CLOEXEC, Mode::empty())?; + + Ok(File::from(new_fd)) +} diff --git a/src/new_data_collection/sources/linux/temperature.rs b/src/new_data_collection/sources/linux/temperature.rs index 4aaf1f52b..55438bfdd 100644 --- a/src/new_data_collection/sources/linux/temperature.rs +++ b/src/new_data_collection/sources/linux/temperature.rs @@ -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, + temperatures: Vec, 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) -> HwmonResults { - let mut temperatures: Vec = vec![]; + let mut temperatures: Vec = vec![]; let mut seen_names: HashMap = HashMap::new(); let (dirs, num_hwmon) = get_hwmon_candidates(); @@ -246,7 +248,7 @@ fn hwmon_temperatures(temp_type: &TemperatureType, filter: &Option) -> 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) -> 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) -> 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, temp_type: &TemperatureType, filter: &Option, + temperatures: &mut Vec, temp_type: &TemperatureType, filter: &Option, ) { 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, -) -> Result>> { +) -> Vec { 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)] diff --git a/src/new_data_collection/collectors/mod.rs b/src/new_data_collection/sources/macos/mod.rs similarity index 100% rename from src/new_data_collection/collectors/mod.rs rename to src/new_data_collection/sources/macos/mod.rs diff --git a/src/new_data_collection/sources/mod.rs b/src/new_data_collection/sources/mod.rs index 5f3015811..e5436906f 100644 --- a/src/new_data_collection/sources/mod.rs +++ b/src/new_data_collection/sources/mod.rs @@ -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; } } diff --git a/src/new_data_collection/sources/nvidia/mod.rs b/src/new_data_collection/sources/nvidia/mod.rs new file mode 100644 index 000000000..e69de29bb diff --git a/src/new_data_collection/sources/sysinfo/mod.rs b/src/new_data_collection/sources/sysinfo/mod.rs new file mode 100644 index 000000000..a0749b8fb --- /dev/null +++ b/src/new_data_collection/sources/sysinfo/mod.rs @@ -0,0 +1 @@ +pub mod temperature; diff --git a/src/new_data_collection/sources/sysinfo/temperature.rs b/src/new_data_collection/sources/sysinfo/temperature.rs index d1a9e634b..a6a60099a 100644 --- a/src/new_data_collection/sources/sysinfo/temperature.rs +++ b/src/new_data_collection/sources/sysinfo/temperature.rs @@ -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, -) -> Result>> { - let mut temperature_vec: Vec = Vec::new(); +) -> Vec { + let mut temperature_vec: Vec = 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 } diff --git a/src/new_data_collection/sources/unix/mod.rs b/src/new_data_collection/sources/unix/mod.rs new file mode 100644 index 000000000..6f5902e45 --- /dev/null +++ b/src/new_data_collection/sources/unix/mod.rs @@ -0,0 +1 @@ +pub mod processes; diff --git a/src/new_data_collection/sources/unix/processes/mod.rs b/src/new_data_collection/sources/unix/processes/mod.rs new file mode 100644 index 000000000..d11af02e7 --- /dev/null +++ b/src/new_data_collection/sources/unix/processes/mod.rs @@ -0,0 +1,5 @@ +pub mod user_table; + +/// A UNIX process ID. +#[cfg(target_family = "unix")] +pub type Pid = libc::pid_t; diff --git a/src/new_data_collection/sources/unix/processes/user_table.rs b/src/new_data_collection/sources/unix/processes/user_table.rs new file mode 100644 index 000000000..dc8e0ab45 --- /dev/null +++ b/src/new_data_collection/sources/unix/processes/user_table.rs @@ -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, +} + +impl UserTable { + pub fn get_uid_to_username_mapping(&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 + 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) + } + } + } +} diff --git a/src/new_data_collection/sources/windows/mod.rs b/src/new_data_collection/sources/windows/mod.rs new file mode 100644 index 000000000..6f5902e45 --- /dev/null +++ b/src/new_data_collection/sources/windows/mod.rs @@ -0,0 +1 @@ +pub mod processes; diff --git a/src/new_data_collection/sources/windows/processes.rs b/src/new_data_collection/sources/windows/processes.rs new file mode 100644 index 000000000..4cd96315c --- /dev/null +++ b/src/new_data_collection/sources/windows/processes.rs @@ -0,0 +1,3 @@ +/// A Windows process ID. +#[cfg(target_family = "windows")] +pub type Pid = usize; From 1b98b967a80906a93548cad3aab863646165d533 Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Wed, 11 Sep 2024 22:35:10 -0400 Subject: [PATCH 5/7] temp commit --- src/new_data_collection/collectors/common.rs | 19 ++- src/new_data_collection/collectors/linux.rs | 5 + src/new_data_collection/error.rs | 2 +- .../sources/common/disk.rs | 127 ++++++++++++++++++ src/new_data_collection/sources/common/mod.rs | 1 + .../sources/sysinfo/disk.rs | 0 6 files changed, 143 insertions(+), 11 deletions(-) create mode 100644 src/new_data_collection/sources/common/disk.rs create mode 100644 src/new_data_collection/sources/sysinfo/disk.rs diff --git a/src/new_data_collection/collectors/common.rs b/src/new_data_collection/collectors/common.rs index 6d7e507a9..d16956648 100644 --- a/src/new_data_collection/collectors/common.rs +++ b/src/new_data_collection/collectors/common.rs @@ -2,7 +2,9 @@ use crate::new_data_collection::{ error::CollectionResult, - sources::common::{processes::ProcessHarvest, temperature::TemperatureData}, + sources::common::{ + disk::DiskHarvest, processes::ProcessHarvest, temperature::TemperatureData, + }, }; /// The trait representing what a per-platform data collector should implement. @@ -11,17 +13,14 @@ pub(crate) trait DataCollector { /// /// Note that depending on the implementation, this may /// not actually need to do anything. - fn refresh_data(&mut self) -> CollectionResult<()> { - Ok(()) - } + fn refresh_data(&mut self) -> CollectionResult<()>; /// Return temperature data. - fn get_temperature_data(&mut self) -> CollectionResult> { - Ok(vec![]) - } + fn get_temperature_data(&mut self) -> CollectionResult>; /// Return process data. - fn get_process_data(&mut self) -> CollectionResult> { - Ok(vec![]) - } + fn get_process_data(&mut self) -> CollectionResult>; + + /// Return disk data. + fn get_disk_data(&mut self) -> CollectionResult; } diff --git a/src/new_data_collection/collectors/linux.rs b/src/new_data_collection/collectors/linux.rs index 4a7191a4e..8786762be 100644 --- a/src/new_data_collection/collectors/linux.rs +++ b/src/new_data_collection/collectors/linux.rs @@ -10,6 +10,7 @@ use crate::{ error::CollectionResult, sources::{ common::{ + disk::DiskHarvest, processes::ProcessHarvest, temperature::{TemperatureData, TemperatureType}, }, @@ -68,4 +69,8 @@ impl DataCollector for LinuxDataCollector { self.gpus_total_mem, ) } + + fn get_disk_data(&mut self) -> CollectionResult { + todo!() + } } diff --git a/src/new_data_collection/error.rs b/src/new_data_collection/error.rs index dc5dd7bd3..7652ac6fe 100644 --- a/src/new_data_collection/error.rs +++ b/src/new_data_collection/error.rs @@ -6,7 +6,7 @@ pub enum CollectionError { /// A general error to propagate back up. A wrapper around [`anyhow::Error`]. General(anyhow::Error), - /// The collection is unsupported. + /// Collection is unsupported. Unsupported, } diff --git a/src/new_data_collection/sources/common/disk.rs b/src/new_data_collection/sources/common/disk.rs new file mode 100644 index 000000000..63e1a6e27 --- /dev/null +++ b/src/new_data_collection/sources/common/disk.rs @@ -0,0 +1,127 @@ +use hashbrown::HashMap; + +use crate::app::filter::Filter; + +#[derive(Clone, Debug, Default)] +pub struct DiskEntry { + pub name: String, + pub mount_point: String, + + /// Windows also contains an additional volume name field. + #[cfg(target_os = "windows")] + pub volume_name: Option, + + // TODO: Maybe unify all these? + pub free_space: Option, + pub used_space: Option, + pub total_space: Option, +} + +#[derive(Clone, Debug)] +pub struct IoData { + /// How many bytes are read. + pub read_bytes: u64, + + /// How many bytes are written. + pub write_bytes: u64, +} + +pub struct DiskHarvest { + /// Disk entries. + pub entries: Vec, + /// I/O stats, mapped to device names. + pub device_io_stats: HashMap>, +} + +/// Whether to keep the current disk entry given the filters, disk name, and +/// disk mount. Precedence ordering in the case where name and mount filters +/// disagree, "allow" takes precedence over "deny". +/// +/// For implementation, we do this as follows: +/// +/// 1. Is the entry allowed through any filter? That is, does it match an entry +/// in a filter where `is_list_ignored` is `false`? If so, we always keep +/// this entry. +/// 2. Is the entry denied through any filter? That is, does it match an entry +/// in a filter where `is_list_ignored` is `true`? If so, we always deny this +/// entry. +/// 3. Anything else is allowed. +pub fn keep_disk_entry( + disk_name: &str, mount_point: &str, disk_filter: &Option, mount_filter: &Option, +) -> bool { + match (disk_filter, mount_filter) { + (Some(d), Some(m)) => match (d.ignore_matches(), m.ignore_matches()) { + (true, true) => !(d.has_match(disk_name) || m.has_match(mount_point)), + (true, false) => { + if m.has_match(mount_point) { + true + } else { + d.should_keep(disk_name) + } + } + (false, true) => { + if d.has_match(disk_name) { + true + } else { + m.should_keep(mount_point) + } + } + (false, false) => d.has_match(disk_name) || m.has_match(mount_point), + }, + (Some(d), None) => d.should_keep(disk_name), + (None, Some(m)) => m.should_keep(mount_point), + (None, None) => true, + } +} + +#[cfg(test)] +mod test { + use regex::Regex; + + use super::keep_disk_entry; + use crate::app::filter::Filter; + + fn run_filter(disk_filter: &Option, mount_filter: &Option) -> Vec { + let targets = [ + ("/dev/nvme0n1p1", "/boot"), + ("/dev/nvme0n1p2", "/"), + ("/dev/nvme0n1p3", "/home"), + ("/dev/sda1", "/mnt/test"), + ("/dev/sda2", "/mnt/boot"), + ]; + + targets + .into_iter() + .enumerate() + .filter_map(|(itx, (name, mount))| { + if keep_disk_entry(name, mount, disk_filter, mount_filter) { + Some(itx) + } else { + None + } + }) + .collect() + } + + #[test] + fn test_keeping_disk_entry() { + let disk_ignore = Some(Filter::new(true, vec![Regex::new("nvme").unwrap()])); + let disk_keep = Some(Filter::new(false, vec![Regex::new("nvme").unwrap()])); + let mount_ignore = Some(Filter::new(true, vec![Regex::new("boot").unwrap()])); + let mount_keep = Some(Filter::new(false, vec![Regex::new("boot").unwrap()])); + + assert_eq!(run_filter(&None, &None), vec![0, 1, 2, 3, 4]); + + assert_eq!(run_filter(&disk_ignore, &None), vec![3, 4]); + assert_eq!(run_filter(&disk_keep, &None), vec![0, 1, 2]); + + assert_eq!(run_filter(&None, &mount_ignore), vec![1, 2, 3]); + assert_eq!(run_filter(&None, &mount_keep), vec![0, 4]); + + assert_eq!(run_filter(&disk_ignore, &mount_ignore), vec![3]); + assert_eq!(run_filter(&disk_keep, &mount_ignore), vec![0, 1, 2, 3]); + + assert_eq!(run_filter(&disk_ignore, &mount_keep), vec![0, 3, 4]); + assert_eq!(run_filter(&disk_keep, &mount_keep), vec![0, 1, 2, 4]); + } +} diff --git a/src/new_data_collection/sources/common/mod.rs b/src/new_data_collection/sources/common/mod.rs index d918321f6..9a68e2856 100644 --- a/src/new_data_collection/sources/common/mod.rs +++ b/src/new_data_collection/sources/common/mod.rs @@ -1,2 +1,3 @@ +pub mod disk; pub mod processes; pub mod temperature; diff --git a/src/new_data_collection/sources/sysinfo/disk.rs b/src/new_data_collection/sources/sysinfo/disk.rs new file mode 100644 index 000000000..e69de29bb From 2ba435f22ed63e39ab0fd20e2351d3d263aee91b Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Thu, 12 Sep 2024 04:13:30 -0400 Subject: [PATCH 6/7] some progress on linux --- src/data_collection/cpu.rs | 3 - src/data_collection/memory.rs | 2 +- src/new_data_collection/collectors/common.rs | 43 +++++++-- src/new_data_collection/collectors/freebsd.rs | 36 -------- src/new_data_collection/collectors/linux.rs | 87 ++++++++++++++++--- src/new_data_collection/collectors/macos.rs | 33 ------- src/new_data_collection/collectors/windows.rs | 33 ------- src/new_data_collection/error.rs | 6 ++ src/new_data_collection/mod.rs | 18 ++-- .../sources/common/battery.rs | 20 +++++ src/new_data_collection/sources/common/cpu.rs | 21 +++++ .../sources/common/disk.rs | 1 + .../sources/common/memory.rs | 22 +++++ src/new_data_collection/sources/common/mod.rs | 4 + .../sources/freebsd/mod.rs | 1 - .../sources/freebsd/temperature.rs | 35 -------- .../sources/linux/memory.rs | 46 ++++++++++ src/new_data_collection/sources/linux/mod.rs | 15 +++- .../sources/linux/processes/mod.rs | 18 ++-- src/new_data_collection/sources/macos/mod.rs | 0 src/new_data_collection/sources/mod.rs | 14 +-- .../sources/nvidia/memory.rs | 32 +++++++ src/new_data_collection/sources/nvidia/mod.rs | 2 + .../{starship_battery.rs => starship.rs} | 24 ++--- .../sources/sysinfo/cpu.rs | 40 +++++++++ .../sources/sysinfo/memory.rs | 47 ++++++++++ .../sources/sysinfo/mod.rs | 3 + .../sources/windows/mod.rs | 1 - .../sources/windows/processes.rs | 3 - 29 files changed, 406 insertions(+), 204 deletions(-) delete mode 100644 src/new_data_collection/collectors/freebsd.rs delete mode 100644 src/new_data_collection/collectors/macos.rs delete mode 100644 src/new_data_collection/collectors/windows.rs create mode 100644 src/new_data_collection/sources/common/battery.rs create mode 100644 src/new_data_collection/sources/common/cpu.rs create mode 100644 src/new_data_collection/sources/common/memory.rs delete mode 100644 src/new_data_collection/sources/freebsd/mod.rs delete mode 100644 src/new_data_collection/sources/freebsd/temperature.rs create mode 100644 src/new_data_collection/sources/linux/memory.rs delete mode 100644 src/new_data_collection/sources/macos/mod.rs create mode 100644 src/new_data_collection/sources/nvidia/memory.rs rename src/new_data_collection/sources/{starship_battery.rs => starship.rs} (71%) create mode 100644 src/new_data_collection/sources/sysinfo/cpu.rs create mode 100644 src/new_data_collection/sources/sysinfo/memory.rs delete mode 100644 src/new_data_collection/sources/windows/mod.rs delete mode 100644 src/new_data_collection/sources/windows/processes.rs diff --git a/src/data_collection/cpu.rs b/src/data_collection/cpu.rs index 843df161c..5aae7e460 100644 --- a/src/data_collection/cpu.rs +++ b/src/data_collection/cpu.rs @@ -18,6 +18,3 @@ pub struct CpuData { } pub type CpuHarvest = Vec; - -pub type PastCpuWork = f64; -pub type PastCpuTotal = f64; diff --git a/src/data_collection/memory.rs b/src/data_collection/memory.rs index b953e2bfe..64fc7caf0 100644 --- a/src/data_collection/memory.rs +++ b/src/data_collection/memory.rs @@ -19,6 +19,6 @@ pub mod arc; pub struct MemHarvest { pub used_bytes: u64, pub total_bytes: u64, - pub use_percent: Option, /* TODO: Might be find to just make this an f64, and any + pub use_percent: Option, /* TODO: Might be fine to just make this an f64, and any * consumer checks NaN. */ } diff --git a/src/new_data_collection/collectors/common.rs b/src/new_data_collection/collectors/common.rs index d16956648..e418e0bc3 100644 --- a/src/new_data_collection/collectors/common.rs +++ b/src/new_data_collection/collectors/common.rs @@ -1,19 +1,34 @@ //! Common code amongst all data collectors. -use crate::new_data_collection::{ - error::CollectionResult, - sources::common::{ - disk::DiskHarvest, processes::ProcessHarvest, temperature::TemperatureData, +use crate::{ + data_collection::Data, + new_data_collection::{ + error::CollectionResult, + sources::{ + cpu::CpuHarvest, disk::DiskHarvest, memory::MemHarvest, processes::ProcessHarvest, + temperature::TemperatureData, + }, }, }; +#[cfg(feature = "battery")] +use crate::new_data_collection::sources::battery::BatteryHarvest; + +// /// Represents data collected at an instance. +// #[derive(Debug)] +// pub struct Data { +// pub collection_time: Instant, +// pub temperature_data: Option>, +// pub process_data: Option>, +// pub disk_data: Option, +// } + /// 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. +pub trait DataCollector { + /// Return data. /// - /// Note that depending on the implementation, this may - /// not actually need to do anything. - fn refresh_data(&mut self) -> CollectionResult<()>; + /// For now, this returns the old data type for cross-compatibility as we migrate. + fn get_data(&mut self) -> Data; /// Return temperature data. fn get_temperature_data(&mut self) -> CollectionResult>; @@ -23,4 +38,14 @@ pub(crate) trait DataCollector { /// Return disk data. fn get_disk_data(&mut self) -> CollectionResult; + + /// Return CPU data. + fn get_cpu_data(&mut self) -> CollectionResult; + + /// Return memory data. + fn get_memory_data(&mut self) -> CollectionResult; + + #[cfg(feature = "battery")] + /// Return battery data. + fn get_battery_data(&mut self) -> CollectionResult>; } diff --git a/src/new_data_collection/collectors/freebsd.rs b/src/new_data_collection/collectors/freebsd.rs deleted file mode 100644 index ccf5aad9b..000000000 --- a/src/new_data_collection/collectors/freebsd.rs +++ /dev/null @@ -1,36 +0,0 @@ -//! 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, -} - -impl DataCollector for FreeBsdDataCollector { - fn refresh_data(&mut self) -> CollectionResult<()> { - Ok(()) - } - - fn get_temperature_data(&self) -> CollectionResult>> { - 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)) - } -} diff --git a/src/new_data_collection/collectors/linux.rs b/src/new_data_collection/collectors/linux.rs index 8786762be..88771bd22 100644 --- a/src/new_data_collection/collectors/linux.rs +++ b/src/new_data_collection/collectors/linux.rs @@ -2,28 +2,36 @@ use std::time::Instant; -use starship_battery::{Battery, Manager}; - use crate::{ app::filter::Filter, + data_collection::Data, new_data_collection::{ error::CollectionResult, sources::{ - common::{ - disk::DiskHarvest, - processes::ProcessHarvest, - temperature::{TemperatureData, TemperatureType}, - }, - linux::{ - processes::{linux_process_data, ProcessCollector}, - temperature::get_temperature_data, + cpu::CpuHarvest, + disk::DiskHarvest, + linux::{get_temperature_data, linux_process_data, ProcessCollector}, + memory::MemHarvest, + processes::ProcessHarvest, + sysinfo::{ + cpu::{get_cpu_data_list, get_load_avg}, + memory::{get_cache_usage, get_ram_usage, get_swap_usage}, }, + temperature::{TemperatureData, TemperatureType}, }, }, }; use super::common::DataCollector; +cfg_if::cfg_if! { + if #[cfg(feature = "battery")] { + use starship_battery::{Battery, Manager}; + use crate::new_data_collection::sources::battery::BatteryHarvest; + } + +} + /// The [`DataCollector`] for Linux. pub struct LinuxDataCollector { current_collection_time: Instant, @@ -37,19 +45,30 @@ pub struct LinuxDataCollector { system: sysinfo::System, network: sysinfo::Networks, + show_average_cpu: bool, + #[cfg(feature = "battery")] - battery_manager: Option, - #[cfg(feature = "battery")] - battery_list: Option>, + batteries: Option<(Manager, Vec)>, + + #[cfg(feature = "gpu")] + nvml: nvml_wrapper::Nvml, #[cfg(feature = "gpu")] gpus_total_mem: Option, } -impl DataCollector for LinuxDataCollector { +impl LinuxDataCollector { fn refresh_data(&mut self) -> CollectionResult<()> { Ok(()) } +} + +impl DataCollector for LinuxDataCollector { + fn get_data(&mut self) -> Data { + let collection_time = Instant::now(); + + todo!() + } fn get_temperature_data(&mut self) -> CollectionResult> { Ok(get_temperature_data(&self.temp_type, &self.temp_filters)) @@ -73,4 +92,44 @@ impl DataCollector for LinuxDataCollector { fn get_disk_data(&mut self) -> CollectionResult { todo!() } + + fn get_cpu_data(&mut self) -> CollectionResult { + let usages = get_cpu_data_list(&self.system, self.show_average_cpu); + let load_average = get_load_avg(); + + CollectionResult::Ok(CpuHarvest { + usages, + load_average, + }) + } + + fn get_memory_data(&mut self) -> CollectionResult { + let memory = get_ram_usage(&self.system); + let swap = get_swap_usage(&self.system); + let cache = get_cache_usage(&self.system); + + CollectionResult::Ok(MemHarvest { + memory, + swap, + cache, + #[cfg(feature = "zfs")] + arc: crate::new_data_collection::sources::linux::get_arc_usage(), + #[cfg(feature = "gpu")] + gpu: crate::new_data_collection::sources::nvidia::get_gpu_memory_usage(&self.nvml), + }) + } + + #[cfg(feature = "battery")] + fn get_battery_data(&mut self) -> CollectionResult> { + use crate::new_data_collection::{ + error::CollectionError, sources::starship::refresh_batteries, + }; + + match &mut self.batteries { + Some((battery_manager, battery_list)) => { + CollectionResult::Ok(refresh_batteries(battery_manager, battery_list)) + } + None => CollectionResult::Err(CollectionError::NoData), + } + } } diff --git a/src/new_data_collection/collectors/macos.rs b/src/new_data_collection/collectors/macos.rs deleted file mode 100644 index 5a4efb745..000000000 --- a/src/new_data_collection/collectors/macos.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! 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, -} - -impl DataCollector for MacOsDataCollector { - fn refresh_data(&mut self) -> CollectionResult<()> { - Ok(()) - } - - fn get_temperature_data(&self) -> CollectionResult>> { - Ok(Some(get_temperature_data( - &self.temp_type, - &self.temp_filters, - ))) - } -} diff --git a/src/new_data_collection/collectors/windows.rs b/src/new_data_collection/collectors/windows.rs deleted file mode 100644 index 607d94663..000000000 --- a/src/new_data_collection/collectors/windows.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! 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, -} - -impl DataCollector for WindowsDataCollector { - fn refresh_data(&mut self) -> CollectionResult<()> { - Ok(()) - } - - fn get_temperature_data(&self) -> CollectionResult>> { - Ok(Some(get_temperature_data( - &self.temp_type, - &self.temp_filters, - ))) - } -} diff --git a/src/new_data_collection/error.rs b/src/new_data_collection/error.rs index 7652ac6fe..5a8b3347d 100644 --- a/src/new_data_collection/error.rs +++ b/src/new_data_collection/error.rs @@ -6,6 +6,9 @@ pub enum CollectionError { /// A general error to propagate back up. A wrapper around [`anyhow::Error`]. General(anyhow::Error), + /// No data. + NoData, + /// Collection is unsupported. Unsupported, } @@ -14,6 +17,9 @@ 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::NoData => { + write!(f, "no data found") + } CollectionError::Unsupported => { write!( f, diff --git a/src/new_data_collection/mod.rs b/src/new_data_collection/mod.rs index 8f9a2cef1..cace97014 100644 --- a/src/new_data_collection/mod.rs +++ b/src/new_data_collection/mod.rs @@ -9,15 +9,15 @@ mod collectors { 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 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; diff --git a/src/new_data_collection/sources/common/battery.rs b/src/new_data_collection/sources/common/battery.rs new file mode 100644 index 000000000..e8b9b5af7 --- /dev/null +++ b/src/new_data_collection/sources/common/battery.rs @@ -0,0 +1,20 @@ +//! Common code for retrieving battery data. + +#[derive(Debug, Clone)] +pub enum State { + Unknown, + Charging, + Discharging, + Empty, + Full, +} + +#[derive(Debug, Clone)] +pub struct BatteryHarvest { + pub charge_percent: f64, + pub secs_until_full: Option, + pub secs_until_empty: Option, + pub power_consumption_rate_watts: f64, + pub health_percent: f64, + pub state: State, +} diff --git a/src/new_data_collection/sources/common/cpu.rs b/src/new_data_collection/sources/common/cpu.rs new file mode 100644 index 000000000..83c463e44 --- /dev/null +++ b/src/new_data_collection/sources/common/cpu.rs @@ -0,0 +1,21 @@ +//! Common code for retrieving CPU data. + +#[derive(Debug, Clone, Copy)] +pub(crate) enum CpuDataType { + Avg, + Cpu(usize), +} + +/// Represents a single core/thread and its usage. +#[derive(Debug, Clone)] +pub(crate) struct CpuData { + pub entry_type: CpuDataType, + pub usage: f64, +} + +/// Collected CPU data at an instance. +#[derive(Debug, Clone)] +pub(crate) struct CpuHarvest { + pub usages: Vec, + pub load_average: [f32; 3], +} diff --git a/src/new_data_collection/sources/common/disk.rs b/src/new_data_collection/sources/common/disk.rs index 63e1a6e27..07dd1748b 100644 --- a/src/new_data_collection/sources/common/disk.rs +++ b/src/new_data_collection/sources/common/disk.rs @@ -26,6 +26,7 @@ pub struct IoData { pub write_bytes: u64, } +#[derive(Clone, Debug)] pub struct DiskHarvest { /// Disk entries. pub entries: Vec, diff --git a/src/new_data_collection/sources/common/memory.rs b/src/new_data_collection/sources/common/memory.rs new file mode 100644 index 000000000..bdbe1cdf5 --- /dev/null +++ b/src/new_data_collection/sources/common/memory.rs @@ -0,0 +1,22 @@ +//! Code pertaining to memory data retrieval. + +#[derive(Debug)] +pub(crate) struct MemData { + pub used_bytes: u64, + pub total_bytes: u64, +} + +#[derive(Debug)] +pub(crate) struct MemHarvest { + pub memory: MemData, + pub swap: MemData, + + #[cfg(not(target_os = "windows"))] + pub cache: MemData, + + #[cfg(feature = "zfs")] + pub arc: MemData, + + #[cfg(feature = "gpu")] + pub gpu: Vec<(String, MemData)>, +} diff --git a/src/new_data_collection/sources/common/mod.rs b/src/new_data_collection/sources/common/mod.rs index 9a68e2856..dc3e32418 100644 --- a/src/new_data_collection/sources/common/mod.rs +++ b/src/new_data_collection/sources/common/mod.rs @@ -1,3 +1,7 @@ +#[cfg(feature = "battery")] +pub mod battery; +pub mod cpu; pub mod disk; +pub mod memory; pub mod processes; pub mod temperature; diff --git a/src/new_data_collection/sources/freebsd/mod.rs b/src/new_data_collection/sources/freebsd/mod.rs deleted file mode 100644 index d7f95f24f..000000000 --- a/src/new_data_collection/sources/freebsd/mod.rs +++ /dev/null @@ -1 +0,0 @@ -mod temperature; diff --git a/src/new_data_collection/sources/freebsd/temperature.rs b/src/new_data_collection/sources/freebsd/temperature.rs deleted file mode 100644 index c9058f15d..000000000 --- a/src/new_data_collection/sources/freebsd/temperature.rs +++ /dev/null @@ -1,35 +0,0 @@ -//! 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, -) -> impl Iterator { - 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() - } -} diff --git a/src/new_data_collection/sources/linux/memory.rs b/src/new_data_collection/sources/linux/memory.rs new file mode 100644 index 000000000..b33ebec80 --- /dev/null +++ b/src/new_data_collection/sources/linux/memory.rs @@ -0,0 +1,46 @@ +use crate::new_data_collection::sources::memory::MemData; + +pub(crate) fn get_arc_usage() -> MemData { + // TODO: [OPT] is this efficient? + use std::fs::read_to_string; + + let (total_bytes, used_bytes) = + if let Ok(arc_stats) = read_to_string("/proc/spl/kstat/zfs/arcstats") { + let mut mem_arc = 0; + let mut mem_total = 0; + let mut zfs_keys_read: u8 = 0; + const ZFS_KEYS_NEEDED: u8 = 2; + + for line in arc_stats.lines() { + if let Some((label, value)) = line.split_once(' ') { + let to_write = match label { + "size" => &mut mem_arc, + "c_max" => &mut mem_total, + _ => { + continue; + } + }; + + if let Some((_type, number)) = value.trim_start().rsplit_once(' ') { + // Parse the value, remember it's in bytes! + if let Ok(number) = number.parse::() { + *to_write = number; + // We only need a few keys, so we can bail early. + zfs_keys_read += 1; + if zfs_keys_read == ZFS_KEYS_NEEDED { + break; + } + } + } + } + } + (mem_total, mem_arc) + } else { + (0, 0) + }; + + MemData { + used_bytes, + total_bytes, + } +} diff --git a/src/new_data_collection/sources/linux/mod.rs b/src/new_data_collection/sources/linux/mod.rs index d918321f6..a57a9884d 100644 --- a/src/new_data_collection/sources/linux/mod.rs +++ b/src/new_data_collection/sources/linux/mod.rs @@ -1,2 +1,13 @@ -pub mod processes; -pub mod temperature; +mod processes; +mod temperature; + +pub(crate) use processes::*; +pub(crate) use temperature::*; + +// For now we only use a Linux-specific implementation for zfs ARC usage. +cfg_if::cfg_if! { + if #[cfg(feature = "zfs")] { + mod memory; + pub(crate) use memory::*; + } +} diff --git a/src/new_data_collection/sources/linux/processes/mod.rs b/src/new_data_collection/sources/linux/processes/mod.rs index ef978a8e0..93de68dd5 100644 --- a/src/new_data_collection/sources/linux/processes/mod.rs +++ b/src/new_data_collection/sources/linux/processes/mod.rs @@ -272,13 +272,13 @@ fn read_proc( )) } -pub(crate) struct PrevProc { +pub struct PrevProc { pub prev_idle: f64, pub prev_non_idle: f64, } #[derive(Clone, Copy)] -pub(crate) struct ProcHarvestOptions { +pub struct ProcHarvestOptions { pub use_current_cpu_total: bool, pub unnormalized_cpu: bool, } @@ -289,13 +289,13 @@ fn is_str_numeric(s: &str) -> bool { /// 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 ReadProcArgs { + pub use_current_cpu_total: bool, + pub cpu_usage: f64, + pub cpu_fraction: f64, + pub total_memory: u64, + pub time_difference_in_secs: u64, + pub uptime: u64, } pub struct ProcessCollector { diff --git a/src/new_data_collection/sources/macos/mod.rs b/src/new_data_collection/sources/macos/mod.rs deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/new_data_collection/sources/mod.rs b/src/new_data_collection/sources/mod.rs index e5436906f..3ba1016a8 100644 --- a/src/new_data_collection/sources/mod.rs +++ b/src/new_data_collection/sources/mod.rs @@ -1,16 +1,20 @@ pub mod common; pub mod linux; -pub mod macos; +// pub mod macos; #[cfg(feature = "gpu")] pub mod nvidia; +pub mod starship; pub mod sysinfo; pub mod unix; -pub mod windows; +// pub mod windows; + +pub use common::*; cfg_if::cfg_if! { - if #[cfg(target_family = "windows")] { - pub use windows::processes::Pid as Pid; - } else if #[cfg(target_family = "unix")] { + if #[cfg(target_family = "unix")] { pub use unix::processes::Pid as Pid; } + // else if #[cfg(target_family = "windows")] { + // pub use windows::processes::Pid as Pid; + // } } diff --git a/src/new_data_collection/sources/nvidia/memory.rs b/src/new_data_collection/sources/nvidia/memory.rs new file mode 100644 index 000000000..a9e7fce4c --- /dev/null +++ b/src/new_data_collection/sources/nvidia/memory.rs @@ -0,0 +1,32 @@ +use nvml_wrapper::Nvml; + +use crate::new_data_collection::sources::memory::MemData; + +/// Returns GPU memory usage per device name. +pub(crate) fn get_gpu_memory_usage(nvml: &Nvml) -> Vec<(String, MemData)> { + let Ok(num_gpu) = nvml.device_count() else { + return vec![]; + }; + + (0..num_gpu) + .filter_map(|i| nvml.device_by_index(i).ok()) + .filter_map(|device| match device.name() { + Ok(name) => { + match device.memory_info() { + Ok(mem_info) => Some(( + name, + MemData { + used_bytes: mem_info.used, + total_bytes: mem_info.total, + }, + )), + Err(_) => { + // TODO: Maybe we should still return something here if it errors out. + None + } + } + } + Err(_) => None, + }) + .collect() +} diff --git a/src/new_data_collection/sources/nvidia/mod.rs b/src/new_data_collection/sources/nvidia/mod.rs index e69de29bb..118c99e35 100644 --- a/src/new_data_collection/sources/nvidia/mod.rs +++ b/src/new_data_collection/sources/nvidia/mod.rs @@ -0,0 +1,2 @@ +mod memory; +pub(crate) use memory::*; diff --git a/src/new_data_collection/sources/starship_battery.rs b/src/new_data_collection/sources/starship.rs similarity index 71% rename from src/new_data_collection/sources/starship_battery.rs rename to src/new_data_collection/sources/starship.rs index 1235e9384..9086c0501 100644 --- a/src/new_data_collection/sources/starship_battery.rs +++ b/src/new_data_collection/sources/starship.rs @@ -10,17 +10,21 @@ use starship_battery::{ units::{power::watt, ratio::percent, time::second}, - Battery, Manager, State, + Battery, Manager, }; -#[derive(Debug, Clone)] -pub struct BatteryHarvest { - pub charge_percent: f64, - pub secs_until_full: Option, - pub secs_until_empty: Option, - pub power_consumption_rate_watts: f64, - pub health_percent: f64, - pub state: State, +use super::battery::{BatteryHarvest, State}; + +impl From for State { + fn from(value: starship_battery::State) -> Self { + match value { + starship_battery::State::Unknown => State::Unknown, + starship_battery::State::Charging => State::Charging, + starship_battery::State::Discharging => State::Discharging, + starship_battery::State::Empty => State::Empty, + starship_battery::State::Full => State::Full, + } + } } pub fn refresh_batteries(manager: &Manager, batteries: &mut [Battery]) -> Vec { @@ -40,7 +44,7 @@ pub fn refresh_batteries(manager: &Manager, batteries: &mut [Battery]) -> Vec()), power_consumption_rate_watts: f64::from(battery.energy_rate().get::()), health_percent: f64::from(battery.state_of_health().get::()), - state: battery.state(), + state: battery.state().into(), }) } else { None diff --git a/src/new_data_collection/sources/sysinfo/cpu.rs b/src/new_data_collection/sources/sysinfo/cpu.rs new file mode 100644 index 000000000..a9069bfdb --- /dev/null +++ b/src/new_data_collection/sources/sysinfo/cpu.rs @@ -0,0 +1,40 @@ +use std::collections::VecDeque; + +use sysinfo::{LoadAvg, System}; + +use crate::{ + data_collection::cpu::LoadAvgHarvest, + new_data_collection::sources::cpu::{CpuData, CpuDataType}, +}; + +pub(crate) fn get_cpu_data_list(sys: &System, show_average_cpu: bool) -> Vec { + let mut cpu_deque: VecDeque<_> = sys + .cpus() + .iter() + .enumerate() + .map(|(i, cpu)| CpuData { + entry_type: CpuDataType::Cpu(i), + usage: cpu.cpu_usage() as f64, + }) + .collect(); + + if show_average_cpu { + let cpu = sys.global_cpu_info(); + + cpu_deque.push_front(CpuData { + entry_type: CpuDataType::Avg, + usage: cpu.cpu_usage() as f64, + }) + } + + Vec::from(cpu_deque) +} + +#[cfg(not(target_os = "windows"))] +pub(crate) fn get_load_avg() -> LoadAvgHarvest { + // The API for sysinfo apparently wants you to call it like this, rather than + // using a &System. + let LoadAvg { one, five, fifteen } = sysinfo::System::load_average(); + + [one as f32, five as f32, fifteen as f32] +} diff --git a/src/new_data_collection/sources/sysinfo/memory.rs b/src/new_data_collection/sources/sysinfo/memory.rs new file mode 100644 index 000000000..f25b64521 --- /dev/null +++ b/src/new_data_collection/sources/sysinfo/memory.rs @@ -0,0 +1,47 @@ +//! Collecting memory data using sysinfo. + +use sysinfo::System; + +use crate::new_data_collection::sources::memory::MemData; + +/// Returns RAM usage. +pub(crate) fn get_ram_usage(sys: &System) -> MemData { + let mem_used = sys.used_memory(); + let mem_total = sys.total_memory(); + + MemData { + used_bytes: mem_used, + total_bytes: mem_total, + } +} + +/// Returns SWAP usage. +pub(crate) fn get_swap_usage(sys: &System) -> MemData { + let mem_used = sys.used_swap(); + let mem_total = sys.total_swap(); + + MemData { + used_bytes: mem_used, + total_bytes: mem_total, + } +} + +/// Returns cache usage. sysinfo has no way to do this directly but it should +/// equal the difference between the available and free memory. Free memory is +/// defined as memory not containing any data, which means cache and buffer +/// memory are not "free". Available memory is defined as memory able +/// to be allocated by processes, which includes cache and buffer memory. On +/// Windows, this will always be 0 - as such, we do not use this on Windows. +/// +/// For more information, see [docs](https://docs.rs/sysinfo/latest/sysinfo/struct.System.html#method.available_memory) +/// and [memory explanation](https://askubuntu.com/questions/867068/what-is-available-memory-while-using-free-command) +#[cfg(not(target_os = "windows"))] +pub(crate) fn get_cache_usage(sys: &System) -> MemData { + let mem_used = sys.available_memory().saturating_sub(sys.free_memory()); + let mem_total = sys.total_memory(); + + MemData { + total_bytes: mem_total, + used_bytes: mem_used, + } +} diff --git a/src/new_data_collection/sources/sysinfo/mod.rs b/src/new_data_collection/sources/sysinfo/mod.rs index a0749b8fb..aa022b05b 100644 --- a/src/new_data_collection/sources/sysinfo/mod.rs +++ b/src/new_data_collection/sources/sysinfo/mod.rs @@ -1 +1,4 @@ pub mod temperature; +pub mod cpu; +pub mod disk; +pub mod memory; \ No newline at end of file diff --git a/src/new_data_collection/sources/windows/mod.rs b/src/new_data_collection/sources/windows/mod.rs deleted file mode 100644 index 6f5902e45..000000000 --- a/src/new_data_collection/sources/windows/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod processes; diff --git a/src/new_data_collection/sources/windows/processes.rs b/src/new_data_collection/sources/windows/processes.rs deleted file mode 100644 index 4cd96315c..000000000 --- a/src/new_data_collection/sources/windows/processes.rs +++ /dev/null @@ -1,3 +0,0 @@ -/// A Windows process ID. -#[cfg(target_family = "windows")] -pub type Pid = usize; From 30bb37bf5de18cfae05e82e6b3c81d47f4357969 Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Sun, 22 Sep 2024 02:11:24 -0400 Subject: [PATCH 7/7] some notes --- src/data_collection/disks/unix/macos/counters.rs | 2 +- src/new_data_collection/sources/sysinfo/disk.rs | 1 + src/new_data_collection/sources/sysinfo/mod.rs | 4 ++-- src/utils/logging.rs | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/data_collection/disks/unix/macos/counters.rs b/src/data_collection/disks/unix/macos/counters.rs index 7d3d67686..55ce6d52f 100644 --- a/src/data_collection/disks/unix/macos/counters.rs +++ b/src/data_collection/disks/unix/macos/counters.rs @@ -6,7 +6,7 @@ use crate::data_collection::disks::IoCounters; fn get_device_io(device: io_kit::IoObject) -> anyhow::Result { let parent = device.service_parent()?; - // XXX: Re: Conform check being disabled. + // NB: Regarding the conform check being disabled. // // Okay, so this is weird. // diff --git a/src/new_data_collection/sources/sysinfo/disk.rs b/src/new_data_collection/sources/sysinfo/disk.rs index e69de29bb..8b1378917 100644 --- a/src/new_data_collection/sources/sysinfo/disk.rs +++ b/src/new_data_collection/sources/sysinfo/disk.rs @@ -0,0 +1 @@ + diff --git a/src/new_data_collection/sources/sysinfo/mod.rs b/src/new_data_collection/sources/sysinfo/mod.rs index aa022b05b..ea41d9566 100644 --- a/src/new_data_collection/sources/sysinfo/mod.rs +++ b/src/new_data_collection/sources/sysinfo/mod.rs @@ -1,4 +1,4 @@ -pub mod temperature; pub mod cpu; pub mod disk; -pub mod memory; \ No newline at end of file +pub mod memory; +pub mod temperature; diff --git a/src/utils/logging.rs b/src/utils/logging.rs index 85230ccb6..975122c5a 100644 --- a/src/utils/logging.rs +++ b/src/utils/logging.rs @@ -18,7 +18,7 @@ pub fn init_logger( // logging at all in a debug context which is generally fine, // release builds should have this logging disabled entirely for now. unsafe { - // XXX: If we ever DO add general logging as a release feature, evaluate this + // NB: If we ever DO add general logging as a release feature, evaluate this // again and whether this is something we want enabled in // release builds! What might be safe is falling back to the non-set-soundness // mode when specifically using certain feature flags (e.g. dev-logging feature