refactor: more data collection cleanup (#1047)

* refactor: more memory collection cleanup

* more cleanup

* clean up data_harvester, remove heim sensor flag

Separate out most individual components to separate functions. Also
remove Linux's usage of heim's sensors feature, since I wasn't using it
apparently.

* clean up GPU section

* fix cond

* fix feature flags

* more cleanup

* even more cleanup
This commit is contained in:
Clement Tsang 2023-03-07 00:18:24 -05:00 committed by GitHub
parent 8489c4bc10
commit b745684156
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 226 additions and 245 deletions

12
Cargo.lock generated
View File

@ -821,7 +821,6 @@ dependencies = [
"heim-common",
"heim-disk",
"heim-runtime",
"heim-sensors",
]
[[package]]
@ -872,17 +871,6 @@ dependencies = [
"smol",
]
[[package]]
name = "heim-sensors"
version = "0.1.0-rc.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "82de7f0784d3b0c53f2e8875c63f430bf6718b03ec8ddce905c12887031158f5"
dependencies = [
"cfg-if",
"heim-common",
"heim-runtime",
]
[[package]]
name = "hermit-abi"
version = "0.1.19"

View File

@ -109,7 +109,7 @@ unicode-width = "0.1.10"
libc = "0.2.124"
[target.'cfg(target_os = "linux")'.dependencies]
heim = { version = "0.1.0-rc.1", features = ["disk", "sensors"] }
heim = { version = "0.1.0-rc.1", features = ["disk"] }
procfs = { version = "0.15.1", default-features = false }
[target.'cfg(target_os = "macos")'.dependencies]

View File

@ -147,6 +147,7 @@ pub struct App {
pub filters: DataFilters,
}
// TODO: Should probably set a fallback max signal/not supported for this.
#[cfg(target_os = "windows")]
const MAX_SIGNAL: usize = 1;
#[cfg(target_os = "linux")]

View File

@ -214,17 +214,13 @@ impl DataCollection {
}
#[cfg(feature = "zfs")]
{
if let Some(arc) = harvested_data.arc {
self.eat_arc(arc, &mut new_entry);
}
if let Some(arc) = harvested_data.arc {
self.eat_arc(arc, &mut new_entry);
}
#[cfg(feature = "gpu")]
{
if let Some(gpu) = harvested_data.gpu {
self.eat_gpu(gpu, &mut new_entry);
}
if let Some(gpu) = harvested_data.gpu {
self.eat_gpu(gpu, &mut new_entry);
}
// CPU
@ -417,7 +413,7 @@ impl DataCollection {
#[cfg(feature = "gpu")]
fn eat_gpu(&mut self, gpu: Vec<(String, memory::MemHarvest)>, new_entry: &mut TimedData) {
// Note this only pre-calculates the data points - the names will be
// within the local copy of gpu_harvest. Since it's all sequential
// within the local copy of gpu_harvest. Since it's all sequential
// it probably doesn't matter anyways.
gpu.iter().for_each(|data| {
new_entry.gpu_data.push(data.1.use_percent);

View File

@ -2,8 +2,6 @@
use std::time::Instant;
use futures::join;
#[cfg(target_os = "linux")]
use fxhash::FxHashMap;
@ -12,7 +10,7 @@ use starship_battery::{Battery, Manager};
use sysinfo::{System, SystemExt};
use self::memory::MemCollect;
use self::temperature::TemperatureType;
use super::DataFilters;
use crate::app::layout_manager::UsedWidgets;
@ -22,6 +20,7 @@ pub mod nvidia;
#[cfg(feature = "battery")]
pub mod batteries;
pub mod cpu;
pub mod disks;
pub mod memory;
@ -108,7 +107,7 @@ pub struct DataCollector {
#[cfg(target_os = "linux")]
prev_non_idle: f64,
mem_total_kb: u64,
temperature_type: temperature::TemperatureType,
temperature_type: TemperatureType,
use_current_cpu_total: bool,
unnormalized_cpu: bool,
last_collection_time: Instant,
@ -138,7 +137,7 @@ impl DataCollector {
#[cfg(target_os = "linux")]
prev_non_idle: 0_f64,
mem_total_kb: 0,
temperature_type: temperature::TemperatureType::Celsius,
temperature_type: TemperatureType::Celsius,
use_current_cpu_total: false,
unnormalized_cpu: false,
last_collection_time: Instant::now(),
@ -161,9 +160,8 @@ impl DataCollector {
self.mem_total_kb = self.sys.total_memory();
// Refresh network list once at the start.
// TODO: may be worth refreshing every once in a while (maybe on a separate timer).
if self.widgets_to_harvest.use_net {
self.sys.refresh_networks_list();
self.sys.refresh_networks_list(); // TODO: refresh on a timer?
}
if self.widgets_to_harvest.use_proc || self.widgets_to_harvest.use_cpu {
@ -172,11 +170,9 @@ impl DataCollector {
#[cfg(not(target_os = "linux"))]
{
// TODO: Would be good to get this and network list running on a timer instead...?
// Refresh components list once...
// Refresh components list once.
if self.widgets_to_harvest.use_temp {
self.sys.refresh_components_list();
self.sys.refresh_components_list(); // TODO: refresh on a timer?
}
if cfg!(target_os = "windows") && self.widgets_to_harvest.use_proc {
@ -207,7 +203,6 @@ impl DataCollector {
futures::executor::block_on(self.update_data());
std::thread::sleep(std::time::Duration::from_millis(250));
self.data.cleanup();
}
@ -215,7 +210,7 @@ impl DataCollector {
self.widgets_to_harvest = used_widgets;
}
pub fn set_temperature_type(&mut self, temperature_type: temperature::TemperatureType) {
pub fn set_temperature_type(&mut self, temperature_type: TemperatureType) {
self.temperature_type = temperature_type;
}
@ -249,42 +244,72 @@ impl DataCollector {
if self.widgets_to_harvest.use_proc {
self.sys.refresh_processes();
}
if self.widgets_to_harvest.use_temp {
self.sys.refresh_components();
}
}
#[cfg(target_os = "freebsd")]
{
if self.widgets_to_harvest.use_disk {
self.sys.refresh_disks();
}
#[cfg(target_os = "freebsd")]
if self.widgets_to_harvest.use_disk {
self.sys.refresh_disks();
}
let current_instant = Instant::now();
self.update_cpu_usage();
self.update_processes(
#[cfg(target_os = "linux")]
current_instant,
);
self.update_temps();
self.update_memory_usage();
self.update_network_usage(current_instant);
#[cfg(feature = "battery")]
if let Some(battery_manager) = &self.battery_manager {
if let Some(battery_list) = &mut self.battery_list {
self.data.list_of_batteries =
Some(batteries::refresh_batteries(battery_manager, battery_list));
}
}
let current_instant = std::time::Instant::now();
let (disk_res, io_res) = futures::join!(
disks::get_disk_usage(
self.widgets_to_harvest.use_disk,
&self.filters.disk_filter,
&self.filters.mount_filter,
),
disks::get_io_usage(self.widgets_to_harvest.use_disk)
);
// CPU
if let Ok(disks) = disk_res {
self.data.disks = disks;
}
if let Ok(io) = io_res {
self.data.io = io;
}
// Update times for future reference.
self.last_collection_time = current_instant;
self.data.last_collection_time = current_instant;
}
#[inline]
fn update_cpu_usage(&mut self) {
if self.widgets_to_harvest.use_cpu {
self.data.cpu = cpu::get_cpu_data_list(&self.sys, self.show_average_cpu).ok();
#[cfg(target_family = "unix")]
{
// Load Average
self.data.load_avg = cpu::get_load_avg().ok();
}
}
}
// Batteries
#[cfg(feature = "battery")]
{
if let Some(battery_manager) = &self.battery_manager {
if let Some(battery_list) = &mut self.battery_list {
self.data.list_of_batteries =
Some(batteries::refresh_batteries(battery_manager, battery_list));
}
}
}
#[inline]
fn update_processes(&mut self, #[cfg(target_os = "linux")] current_instant: Instant) {
if self.widgets_to_harvest.use_proc {
if let Ok(mut process_list) = {
#[cfg(target_os = "linux")]
@ -301,14 +326,16 @@ impl DataCollector {
unnormalized_cpu: self.unnormalized_cpu,
};
let time_diff = current_instant
.duration_since(self.last_collection_time)
.as_secs();
processes::get_process_data(
&self.sys,
prev_proc,
&mut self.pid_mapping,
proc_harvest_options,
current_instant
.duration_since(self.last_collection_time)
.as_secs(),
time_diff,
self.mem_total_kb,
&mut self.user_table,
)
@ -343,54 +370,49 @@ impl DataCollector {
self.data.list_of_processes = Some(process_list);
}
}
}
#[inline]
fn update_temps(&mut self) {
if self.widgets_to_harvest.use_temp {
#[cfg(not(target_os = "linux"))]
{
if let Ok(data) = temperature::get_temperature_data(
&self.sys,
&self.temperature_type,
&self.filters.temp_filter,
) {
self.data.temperature_sensors = data;
}
if let Ok(data) = temperature::get_temperature_data(
&self.sys,
&self.temperature_type,
&self.filters.temp_filter,
) {
self.data.temperature_sensors = data;
}
#[cfg(target_os = "linux")]
if let Ok(data) =
temperature::get_temperature_data(&self.temperature_type, &self.filters.temp_filter)
{
if let Ok(data) = temperature::get_temperature_data(
&self.temperature_type,
&self.filters.temp_filter,
) {
self.data.temperature_sensors = data;
}
self.data.temperature_sensors = data;
}
}
}
#[inline]
fn update_memory_usage(&mut self) {
if self.widgets_to_harvest.use_mem {
let MemCollect {
ram,
swap,
#[cfg(feature = "gpu")]
gpus,
#[cfg(feature = "zfs")]
arc,
} = memory::get_mem_data(&self.sys, self.widgets_to_harvest.use_gpu);
self.data.memory = ram;
self.data.swap = swap;
self.data.memory = memory::get_ram_usage(&self.sys);
self.data.swap = memory::get_swap_usage(&self.sys);
#[cfg(feature = "zfs")]
{
self.data.arc = arc;
self.data.arc = memory::arc::get_arc_usage();
}
#[cfg(feature = "gpu")]
{
self.data.gpu = gpus;
if self.widgets_to_harvest.use_gpu {
self.data.gpu = memory::gpu::get_gpu_mem_usage();
}
}
}
#[inline]
fn update_network_usage(&mut self, current_instant: Instant) {
if self.widgets_to_harvest.use_net {
let net_data = network::get_network_data(
&self.sys,
@ -405,27 +427,6 @@ impl DataCollector {
self.total_tx = net_data.total_tx;
self.data.network = Some(net_data);
}
let disk_data_fut = disks::get_disk_usage(
self.widgets_to_harvest.use_disk,
&self.filters.disk_filter,
&self.filters.mount_filter,
);
let disk_io_usage_fut = disks::get_io_usage(self.widgets_to_harvest.use_disk);
let (disk_res, io_res) = join!(disk_data_fut, disk_io_usage_fut,);
if let Ok(disks) = disk_res {
self.data.disks = disks;
}
if let Ok(io) = io_res {
self.data.io = io;
}
// Update time
self.data.last_collection_time = current_instant;
self.last_collection_time = current_instant;
}
}

View File

@ -1,7 +1,13 @@
//! Memory data collection.
pub mod sysinfo;
pub(crate) use self::sysinfo::*;
pub(crate) use self::sysinfo::{get_ram_usage, get_swap_usage};
#[cfg(feature = "gpu")]
pub mod gpu;
#[cfg(feature = "zfs")]
pub mod arc;
#[derive(Debug, Clone, Default)]
pub struct MemHarvest {
@ -9,13 +15,3 @@ pub struct MemHarvest {
pub used_kib: u64,
pub use_percent: Option<f64>,
}
#[derive(Debug)]
pub struct MemCollect {
pub ram: Option<MemHarvest>,
pub swap: Option<MemHarvest>,
#[cfg(feature = "zfs")]
pub arc: Option<MemHarvest>,
#[cfg(feature = "gpu")]
pub gpus: Option<Vec<(String, MemHarvest)>>,
}

View File

@ -0,0 +1,75 @@
use super::MemHarvest;
/// Return ARC usage.
#[cfg(feature = "zfs")]
pub(crate) fn get_arc_usage() -> Option<MemHarvest> {
let (mem_total_in_kib, mem_used_in_kib) = {
cfg_if::cfg_if! {
if #[cfg(target_os = "linux")] {
// TODO: [OPT] is this efficient?
use std::fs::read_to_string;
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,
"memory_all_bytes" => &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 / 1024, mem_arc / 1024)
} else {
(0, 0)
}
} else if #[cfg(target_os = "freebsd")] {
use sysctl::Sysctl;
if let (Ok(mem_arc_value), Ok(mem_sys_value)) = (
sysctl::Ctl::new("kstat.zfs.misc.arcstats.size"),
sysctl::Ctl::new("hw.physmem"),
) {
if let (Ok(sysctl::CtlValue::U64(arc)), Ok(sysctl::CtlValue::Ulong(mem))) =
(mem_arc_value.value(), mem_sys_value.value())
{
(mem / 1024, arc / 1024)
} else {
(0, 0)
}
} else {
(0, 0)
}
} else {
(0, 0)
}
}
};
Some(MemHarvest {
total_kib: mem_total_in_kib,
used_kib: mem_used_in_kib,
use_percent: if mem_total_in_kib == 0 {
None
} else {
Some(mem_used_in_kib as f64 / mem_total_in_kib as f64 * 100.0)
},
})
}

View File

@ -0,0 +1,49 @@
use super::MemHarvest;
/// Return GPU memory usage.
#[cfg(feature = "gpu")]
pub(crate) fn get_gpu_mem_usage() -> Option<Vec<(String, MemHarvest)>> {
// As we add more support, expand on this.
#[cfg(feature = "nvidia")]
get_nvidia_mem_usage()
}
/// Returns the memory usage of NVIDIA cards.
#[inline]
#[cfg(feature = "nvidia")]
fn get_nvidia_mem_usage() -> Option<Vec<(String, MemHarvest)>> {
use crate::data_harvester::nvidia::NVML_DATA;
if let Ok(nvml) = &*NVML_DATA {
if let Ok(num_gpu) = nvml.device_count() {
let mut results = Vec::with_capacity(num_gpu as usize);
for i in 0..num_gpu {
if let Ok(device) = nvml.device_by_index(i) {
if let (Ok(name), Ok(mem)) = (device.name(), device.memory_info()) {
// add device memory in bytes
let mem_total_in_kib = mem.total / 1024;
let mem_used_in_kib = mem.used / 1024;
results.push((
name,
MemHarvest {
total_kib: mem_total_in_kib,
used_kib: mem_used_in_kib,
use_percent: if mem_total_in_kib == 0 {
None
} else {
Some(mem_used_in_kib as f64 / mem_total_in_kib as f64 * 100.0)
},
},
));
}
}
}
Some(results)
} else {
None
}
} else {
None
}
}

View File

@ -2,22 +2,10 @@
use sysinfo::{System, SystemExt};
use crate::data_harvester::memory::{MemCollect, MemHarvest};
/// Returns all memory data.
pub(crate) fn get_mem_data(sys: &System, _get_gpu: bool) -> MemCollect {
MemCollect {
ram: get_ram_data(sys),
swap: get_swap_data(sys),
#[cfg(feature = "zfs")]
arc: get_arc_data(),
#[cfg(feature = "gpu")]
gpus: if _get_gpu { get_gpu_data() } else { None },
}
}
use crate::data_harvester::memory::MemHarvest;
/// Returns RAM usage.
pub(crate) fn get_ram_data(sys: &System) -> Option<MemHarvest> {
pub(crate) fn get_ram_usage(sys: &System) -> Option<MemHarvest> {
let mem_used_in_kib = sys.used_memory() / 1024;
let mem_total_in_kib = sys.total_memory() / 1024;
@ -33,7 +21,7 @@ pub(crate) fn get_ram_data(sys: &System) -> Option<MemHarvest> {
}
/// Returns SWAP usage.
pub(crate) fn get_swap_data(sys: &System) -> Option<MemHarvest> {
pub(crate) fn get_swap_usage(sys: &System) -> Option<MemHarvest> {
let mem_used_in_kib = sys.used_swap() / 1024;
let mem_total_in_kib = sys.total_swap() / 1024;
@ -47,116 +35,3 @@ pub(crate) fn get_swap_data(sys: &System) -> Option<MemHarvest> {
},
})
}
/// Return ARC usage.
#[cfg(feature = "zfs")]
pub(crate) fn get_arc_data() -> Option<MemHarvest> {
let (mem_total_in_kib, mem_used_in_kib) = {
cfg_if::cfg_if! {
if #[cfg(target_os = "linux")]
{
// TODO: [OPT] is this efficient?
use std::fs::read_to_string;
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,
"memory_all_bytes" => &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 / 1024, mem_arc / 1024)
} else {
(0, 0)
}
} else if #[cfg(target_os = "freebsd")] {
use sysctl::Sysctl;
if let (Ok(mem_arc_value), Ok(mem_sys_value)) = (
sysctl::Ctl::new("kstat.zfs.misc.arcstats.size"),
sysctl::Ctl::new("hw.physmem"),
) {
if let (Ok(sysctl::CtlValue::U64(arc)), Ok(sysctl::CtlValue::Ulong(mem))) =
(mem_arc_value.value(), mem_sys_value.value())
{
(mem / 1024, arc / 1024)
} else {
(0, 0)
}
} else {
(0, 0)
}
} else {
(0, 0)
}
}
};
Some(MemHarvest {
total_kib: mem_total_in_kib,
used_kib: mem_used_in_kib,
use_percent: if mem_total_in_kib == 0 {
None
} else {
Some(mem_used_in_kib as f64 / mem_total_in_kib as f64 * 100.0)
},
})
}
/// Return GPU data. Currently only supports NVIDIA cards.
#[cfg(feature = "nvidia")]
pub(crate) fn get_gpu_data() -> Option<Vec<(String, MemHarvest)>> {
use crate::data_harvester::nvidia::NVML_DATA;
if let Ok(nvml) = &*NVML_DATA {
if let Ok(ngpu) = nvml.device_count() {
let mut results = Vec::with_capacity(ngpu as usize);
for i in 0..ngpu {
if let Ok(device) = nvml.device_by_index(i) {
if let (Ok(name), Ok(mem)) = (device.name(), device.memory_info()) {
// add device memory in bytes
let mem_total_in_kib = mem.total / 1024;
let mem_used_in_kib = mem.used / 1024;
results.push((
name,
MemHarvest {
total_kib: mem_total_in_kib,
used_kib: mem_used_in_kib,
use_percent: if mem_total_in_kib == 0 {
None
} else {
Some(mem_used_in_kib as f64 / mem_total_in_kib as f64 * 100.0)
},
},
));
}
}
}
Some(results)
} else {
None
}
} else {
None
}
}