some progress on linux

This commit is contained in:
ClementTsang 2024-09-12 04:13:30 -04:00
parent 1b98b967a8
commit 2ba435f22e
No known key found for this signature in database
GPG key ID: DC3B7867D8D97095
29 changed files with 406 additions and 204 deletions

View file

@ -18,6 +18,3 @@ pub struct CpuData {
}
pub type CpuHarvest = Vec<CpuData>;
pub type PastCpuWork = f64;
pub type PastCpuTotal = f64;

View file

@ -19,6 +19,6 @@ pub mod arc;
pub struct MemHarvest {
pub used_bytes: u64,
pub total_bytes: u64,
pub use_percent: Option<f64>, /* TODO: Might be find to just make this an f64, and any
pub use_percent: Option<f64>, /* TODO: Might be fine to just make this an f64, and any
* consumer checks NaN. */
}

View file

@ -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<Vec<TemperatureData>>,
// pub process_data: Option<Vec<ProcessHarvest>>,
// pub disk_data: Option<DiskHarvest>,
// }
/// 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<Vec<TemperatureData>>;
@ -23,4 +38,14 @@ pub(crate) trait DataCollector {
/// Return disk data.
fn get_disk_data(&mut self) -> CollectionResult<DiskHarvest>;
/// Return CPU data.
fn get_cpu_data(&mut self) -> CollectionResult<CpuHarvest>;
/// Return memory data.
fn get_memory_data(&mut self) -> CollectionResult<MemHarvest>;
#[cfg(feature = "battery")]
/// Return battery data.
fn get_battery_data(&mut self) -> CollectionResult<Vec<BatteryHarvest>>;
}

View file

@ -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<Filter>,
}
impl DataCollector for FreeBsdDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
let mut results = get_temperature_data(&self.temp_type, &self.temp_filters);
for entry in sysctl_temp_iter(&self.temp_type, &self.temp_filters) {
results.push(entry);
}
Ok(Some(results))
}
}

View file

@ -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<Manager>,
#[cfg(feature = "battery")]
battery_list: Option<Vec<Battery>>,
batteries: Option<(Manager, Vec<Battery>)>,
#[cfg(feature = "gpu")]
nvml: nvml_wrapper::Nvml,
#[cfg(feature = "gpu")]
gpus_total_mem: Option<u64>,
}
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<Vec<TemperatureData>> {
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<DiskHarvest> {
todo!()
}
fn get_cpu_data(&mut self) -> CollectionResult<CpuHarvest> {
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<MemHarvest> {
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<Vec<BatteryHarvest>> {
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),
}
}
}

View file

@ -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<Filter>,
}
impl DataCollector for MacOsDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
Ok(Some(get_temperature_data(
&self.temp_type,
&self.temp_filters,
)))
}
}

View file

@ -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<Filter>,
}
impl DataCollector for WindowsDataCollector {
fn refresh_data(&mut self) -> CollectionResult<()> {
Ok(())
}
fn get_temperature_data(&self) -> CollectionResult<Option<Vec<TemperatureData>>> {
Ok(Some(get_temperature_data(
&self.temp_type,
&self.temp_filters,
)))
}
}

View file

@ -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,

View file

@ -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;

View file

@ -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<i64>,
pub secs_until_empty: Option<i64>,
pub power_consumption_rate_watts: f64,
pub health_percent: f64,
pub state: State,
}

View file

@ -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<CpuData>,
pub load_average: [f32; 3],
}

View file

@ -26,6 +26,7 @@ pub struct IoData {
pub write_bytes: u64,
}
#[derive(Clone, Debug)]
pub struct DiskHarvest {
/// Disk entries.
pub entries: Vec<DiskEntry>,

View file

@ -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)>,
}

View file

@ -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;

View file

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

View file

@ -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<Filter>,
) -> impl Iterator<Item = TemperatureData> {
const KEY: &str = "hw.temperature";
if let Ok(root) = sysctl::Ctl::new(KEY) {
sysctl::CtlIter::below(root).flatten().filter_map(|ctl| {
if let (Ok(name), Ok(temp)) = (ctl.name(), ctl.value()) {
if let Some(temp) = temp.as_temperature() {
if Filter::optional_should_keep(filter, &name) {
return Some(TemperatureData {
name,
temperature: Some(match temp_type {
TemperatureType::Celsius => temp.celsius(),
TemperatureType::Kelvin => temp.kelvin(),
TemperatureType::Fahrenheit => temp.fahrenheit(),
}),
});
}
}
}
None
})
} else {
std::iter::empty()
}
}

View file

@ -0,0 +1,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::<u64>() {
*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,
}
}

View file

@ -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::*;
}
}

View file

@ -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 {

View file

@ -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;
// }
}

View file

@ -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()
}

View file

@ -0,0 +1,2 @@
mod memory;
pub(crate) use memory::*;

View file

@ -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<i64>,
pub secs_until_empty: Option<i64>,
pub power_consumption_rate_watts: f64,
pub health_percent: f64,
pub state: State,
use super::battery::{BatteryHarvest, State};
impl From<starship_battery::State> 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<BatteryHarvest> {
@ -40,7 +44,7 @@ pub fn refresh_batteries(manager: &Manager, batteries: &mut [Battery]) -> Vec<Ba
charge_percent: f64::from(battery.state_of_charge().get::<percent>()),
power_consumption_rate_watts: f64::from(battery.energy_rate().get::<watt>()),
health_percent: f64::from(battery.state_of_health().get::<percent>()),
state: battery.state(),
state: battery.state().into(),
})
} else {
None

View file

@ -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<CpuData> {
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]
}

View file

@ -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,
}
}

View file

@ -1 +1,4 @@
pub mod temperature;
pub mod cpu;
pub mod disk;
pub mod memory;

View file

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

View file

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