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

Cargo.lock

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

Cargo.toml

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

src/app.rs

@@ -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")]

src/app/data_farmer.rs

@@ -214,17 +214,13 @@ impl DataCollection {
         }
 
         #[cfg(feature = "zfs")]
-        {
+        if let Some(arc) = harvested_data.arc {
-            if let Some(arc) = harvested_data.arc {
+            self.eat_arc(arc, &mut new_entry);
-                self.eat_arc(arc, &mut new_entry);
-            }
         }
 
         #[cfg(feature = "gpu")]
-        {
+        if let Some(gpu) = harvested_data.gpu {
-            if let Some(gpu) = harvested_data.gpu {
+            self.eat_gpu(gpu, &mut new_entry);
-                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);

src/app/data_harvester.rs

@@ -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")]
+        #[cfg(target_os = "freebsd")]
-            {
+        if self.widgets_to_harvest.use_disk {
-                if self.widgets_to_harvest.use_disk {
+            self.sys.refresh_disks();
-                    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
+    #[inline]
-        #[cfg(feature = "battery")]
+    fn update_processes(&mut self, #[cfg(target_os = "linux")] current_instant: Instant) {
-        {
-            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));
-                }
-            }
-        }
-
         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
+                        time_diff,
-                            .duration_since(self.last_collection_time)
-                            .as_secs(),
                         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(
-                if let Ok(data) = temperature::get_temperature_data(
+                &self.sys,
-                    &self.sys,
+                &self.temperature_type,
-                    &self.temperature_type,
+                &self.filters.temp_filter,
-                    &self.filters.temp_filter,
+            ) {
-                ) {
+                self.data.temperature_sensors = data;
-                    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.data.temperature_sensors = data;
-                    &self.temperature_type,
-                    &self.filters.temp_filter,
-                ) {
-                    self.data.temperature_sensors = data;
-                }
             }
         }
+    }
 
+    #[inline]
+    fn update_memory_usage(&mut self) {
         if self.widgets_to_harvest.use_mem {
-            let MemCollect {
+            self.data.memory = memory::get_ram_usage(&self.sys);
-                ram,
+            self.data.swap = memory::get_swap_usage(&self.sys);
-                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;
 
             #[cfg(feature = "zfs")]
             {
-                self.data.arc = arc;
+                self.data.arc = memory::arc::get_arc_usage();
             }
 
             #[cfg(feature = "gpu")]
-            {
+            if self.widgets_to_harvest.use_gpu {
-                self.data.gpu = gpus;
+                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;
     }
 }
 

src/app/data_harvester/memory.rs

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

src/app/data_harvester/memory/arc.rs

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

src/app/data_harvester/memory/gpu.rs

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

src/app/data_harvester/memory/sysinfo.rs

@@ -2,22 +2,10 @@
 
 use sysinfo::{System, SystemExt};
 
-use crate::data_harvester::memory::{MemCollect, MemHarvest};
+use crate::data_harvester::memory::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 },
-    }
-}
 
 /// 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
-    }
-}