some file/implementation organization

Turns out sysinfo lacks a lot of data I need. I can still use it for the Windows disk usage implementation, but I'm probably going to manually implement macos/linux usage and all io usage stats.
ClementTsang · ClementTsang · Apr 10, 2023 · Mar 9, 2023 · Mar 11, 2023 · Mar 14, 2023
commit 7754d02fbd5eca689c4a352205320498bc452a42
diff --git a/src/app/data_harvester.rs b/src/app/data_harvester.rs
@@ -229,12 +229,6 @@ impl DataCollector {
             self.sys.refresh_networks();
         }
 
-        // FreeBSD uses a custom implementation.
-        #[cfg(not(target_os = "freebsd"))]
-        if self.widgets_to_harvest.use_disk {
-            self.sys.refresh_disks();
-        }
-
         #[cfg(not(target_os = "linux"))]
         {
             if self.widgets_to_harvest.use_proc {
@@ -268,27 +262,11 @@ impl DataCollector {
         );
         self.update_temps();
         self.update_network_usage(current_instant);
+        self.update_disks();
 
         #[cfg(feature = "battery")]
         self.update_batteries();
 
-        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)
-        );
-
-        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;
@@ -446,6 +424,31 @@ impl DataCollector {
             }
         }
     }
+
+    #[inline]
+    fn update_disks(&mut self) {
+        if self.widgets_to_harvest.use_disk {
+            #[cfg(any(target_os = "freebsd", target_os = "linux", target_os = "macos"))]
+            {
+                self.data.disks = disks::usage::get_disk_usage(
+                    &self.filters.disk_filter,
+                    &self.filters.mount_filter,
+                )
+                .ok();
+            }
+
+            #[cfg(target_os = "windows")]
+            {
+                self.data.disks = Some(disks::usage::get_disk_usage(
+                    &self.sys,
+                    &self.filters.disk_filter,
+                    &self.filters.mount_filter,
+                ));
+            }
+
+            self.data.io = disks::io::get_io_usage().ok();
+        }
+    }
 }
 
 #[cfg(target_os = "freebsd")]

diff --git a/src/app/data_harvester/disks.rs b/src/app/data_harvester/disks.rs
@@ -1,27 +1,19 @@
-//! Data collection for disks (IO, usage, space, etc.).
-//!
-//! For Linux, macOS, and Windows, this is handled by heim. For FreeBSD there is a custom
-//! implementation.
+//! Data collection about disks (e.g. I/O, usage, space).
 
-cfg_if::cfg_if! {
-    if #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] {
-        pub mod heim;
-        pub use self::heim::*;
-    } else if #[cfg(target_os = "freebsd")] {
-        pub mod freebsd;
-        pub use self::freebsd::*;
-    }
-}
+use std::collections::HashMap;
+
+pub mod io;
+pub mod usage;
 
 #[derive(Debug, Clone, Default)]
 pub struct DiskHarvest {
     pub name: String,
     pub mount_point: String,
 
     // TODO: Maybe unify all these?
-    pub free_space: Option<u64>,
-    pub used_space: Option<u64>,
-    pub total_space: Option<u64>,
+    pub free_space: u64,
+    pub used_space: u64,
+    pub total_space: u64,
 }
 
 #[derive(Clone, Debug)]
@@ -30,4 +22,4 @@ pub struct IoData {
     pub write_bytes: u64,
 }
 
-pub type IoHarvest = std::collections::HashMap<String, Option<IoData>>;
+pub type IoHarvest = HashMap<String, Option<IoData>>;
diff --git a/src/app/data_harvester/disks/heim/windows_macos.rs b/src/app/data_harvester/disks/heim/windows_macos.rs
diff --git a/src/app/data_harvester/disks/io.rs b/src/app/data_harvester/disks/io.rs
@@ -0,0 +1,15 @@
+cfg_if::cfg_if! {
+    if #[cfg(target_os = "freebsd")] {
+        pub mod freebsd;
+        pub use self::freebsd::*;
+    } else if #[cfg(target_os = "windows")] {
+        pub mod windows;
+        pub use self::windows::*;
+    } else if #[cfg(target_os = "linux")] {
+        pub mod linux;
+        pub use self::linux::*;
+    } else if #[cfg(target_os = "macos")] {
+        pub mod macos;
+        pub use self::macos::*;
+    }
+}
diff --git a/src/app/data_harvester/disks/io/freebsd.rs b/src/app/data_harvester/disks/io/freebsd.rs
@@ -0,0 +1,13 @@
+use crate::{app::data_harvester::disks::IoHarvest, utils::error};
+
+pub fn get_io_usage() -> error::Result<IoHarvest> {
+    let io_harvest = get_disk_info().map(|storage_system_information| {
+        storage_system_information
+            .filesystem
+            .into_iter()
+            .map(|disk| (disk.name, None))
+            .collect()
+    })?;
+
+    Ok(io_harvest)
+}
diff --git a/src/app/data_harvester/disks/io/linux.rs b/src/app/data_harvester/disks/io/linux.rs
@@ -0,0 +1,5 @@
+use crate::{app::data_harvester::disks::IoHarvest, utils::error};
+
+pub fn get_io_usage() -> error::Result<IoHarvest> {
+    Ok(IoHarvest::default())
+}
diff --git a/src/app/data_harvester/disks/io/macos.rs b/src/app/data_harvester/disks/io/macos.rs
@@ -0,0 +1,5 @@
+use crate::{app::data_harvester::disks::IoHarvest, utils::error};
+
+pub fn get_io_usage() -> error::Result<IoHarvest> {
+    Ok(IoHarvest::default())
+}
diff --git a/src/app/data_harvester/disks/io/windows.rs b/src/app/data_harvester/disks/io/windows.rs
@@ -0,0 +1,5 @@
+use crate::{app::data_harvester::disks::IoHarvest, utils::error};
+
+pub fn get_io_usage() -> error::Result<IoHarvest> {
+    Ok(IoHarvest::default())
+}
diff --git a/src/app/data_harvester/disks/sysinfo.rs b/src/app/data_harvester/disks/sysinfo.rs
diff --git a/src/app/data_harvester/disks/usage.rs b/src/app/data_harvester/disks/usage.rs
@@ -0,0 +1,12 @@
+cfg_if::cfg_if! {
+    if #[cfg(target_os = "freebsd")] {
+        mod freebsd;
+        pub(crate) use self::freebsd::*;
+    } else if #[cfg(target_os = "windows")] {
+        mod windows;
+        pub(crate) use self::windows::*;
+    } else if #[cfg(target_family = "unix")] {
+        mod unix;
+        pub(crate) use self::unix::*;
+    }
+}
diff --git a/src/app/data_harvester/disks/freebsd.rs → ...app/data_harvester/disks/usage/freebsd.rs b/src/app/data_harvester/disks/freebsd.rs → ...app/data_harvester/disks/usage/freebsd.rs
@@ -4,9 +4,10 @@ use std::io;
 
 use serde::Deserialize;
 
-use super::{DiskHarvest, IoHarvest};
 use crate::app::Filter;
 use crate::data_harvester::deserialize_xo;
+use crate::data_harvester::disks::DiskHarvest;
+use crate::utils::error;
 
 #[derive(Deserialize, Debug, Default)]
 #[serde(rename_all = "kebab-case")]
@@ -24,21 +25,9 @@ struct FileSystem {
     mounted_on: String,
 }
 
-pub fn get_io_usage() -> crate::utils::error::Result<IoHarvest> {
-    let io_harvest = get_disk_info().map(|storage_system_information| {
-        storage_system_information
-            .filesystem
-            .into_iter()
-            .map(|disk| (disk.name, None))
-            .collect()
-    })?;
-
-    Ok(io_harvest)
-}
-
 pub fn get_disk_usage(
     disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
-) -> crate::utils::error::Result<Vec<DiskHarvest>> {
+) -> error::Result<Vec<DiskHarvest>> {
     let vec_disks: Vec<DiskHarvest> = get_disk_info().map(|storage_system_information| {
         storage_system_information
             .filesystem
@@ -60,9 +49,9 @@ pub fn get_disk_usage(
                     || !matches_ignore_list(filter_check_map.as_slice())
                 {
                     Some(DiskHarvest {
-                        free_space: Some(disk.available_blocks * 1024),
-                        used_space: Some(disk.used_blocks * 1024),
-                        total_space: Some(disk.total_blocks * 1024),
+                        free_space: disk.available_blocks * 1024,
+                        used_space: disk.used_blocks * 1024,
+                        total_space: disk.total_blocks * 1024,
                         mount_point: disk.mounted_on,
                         name: disk.name,
                     })

diff --git a/src/app/data_harvester/disks/heim/linux.rs → src/app/data_harvester/disks/usage/unix.rs b/src/app/data_harvester/disks/heim/linux.rs → src/app/data_harvester/disks/usage/unix.rs
@@ -1,8 +1,21 @@
-//! Linux-specific things for Heim disk data collection.
-
 use heim::disk::Partition;
 
-pub fn get_device_name(partition: &Partition) -> String {
+use crate::{
+    app::{data_harvester::disks::DiskHarvest, filter::Filter},
+    utils::error,
+};
+
+#[allow(dead_code)]
+#[allow(unused_variables)]
+pub fn get_disk_usage(
+    disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
+) -> error::Result<Vec<DiskHarvest>> {
+    Ok(vec![])
+}
+
+#[allow(dead_code)]
+#[cfg(target_os = "linux")]
+fn get_device_name(partition: &Partition) -> String {
     if let Some(device) = partition.device() {
         // See if this disk is actually mounted elsewhere on Linux...
         // This is a workaround to properly map I/O in some cases (i.e. disk encryption), see
@@ -32,3 +45,16 @@ pub fn get_device_name(partition: &Partition) -> String {
         "Name Unavailable".to_string()
     }
 }
+
+#[allow(dead_code)]
+#[cfg(not(target_os = "linux"))]
+fn get_device_name(partition: &Partition) -> String {
+    if let Some(device) = partition.device() {
+        device
+            .to_os_string()
+            .into_string()
+            .unwrap_or_else(|_| "Name Unavailable".to_string())
+    } else {
+        "Name Unavailable".to_string()
+    }
+}
diff --git a/src/app/data_harvester/disks/usage/windows.rs b/src/app/data_harvester/disks/usage/windows.rs
@@ -0,0 +1,94 @@
+//! Disk stats via sysinfo.
+
+use sysinfo::{DiskExt, System, SystemExt};
+
+use crate::app::filter::Filter;
+use crate::data_harvester::disks::DiskHarvest;
+
+pub(crate) fn get_disk_usage(
+    sys: &System, disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
+) -> Vec<DiskHarvest> {
+    let disks = sys.disks();
+    disks
+        .iter()
+        .filter_map(|disk| {
+            let name = disk
+                .name()
+                .to_os_string()
+                .into_string()
+                .unwrap_or_else(|_| "Name Unavailable".to_string());
+
+            let mount_point = disk
+                .mount_point()
+                .as_os_str()
+                .to_os_string()
+                .into_string()
+                .unwrap_or_else(|_| "Name Unavailable".to_string());
+
+            if keep_entry(&name, &mount_point, disk_filter, mount_filter) {
+                let free_space = disk.available_space();
+                let total_space = disk.total_space();
+                let used_space = total_space - free_space;
+
+                Some(DiskHarvest {
+                    name,
+                    mount_point,
+                    free_space,
+                    used_space,
+                    total_space,
+                })
+            } else {
+                None
+            }
+        })
+        .collect()
+}
+
+fn keep_entry(
+    name: &str, mount_point: &str, disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
+) -> bool {
+    // 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.
+
+    let filter_check_map = [(disk_filter, &name), (mount_filter, &mount_point)];
+
+    // This represents case 1. That is, if there is a match in an allowing list - if there is, then
+    // immediately allow it!
+    let matches_allow_list = filter_check_map.iter().any(|(filter, text)| {
+        if let Some(filter) = filter {
+            if !filter.is_list_ignored {
+                for r in &filter.list {
+                    if r.is_match(text) {
+                        return true;
+                    }
+                }
+            }
+        }
+        false
+    });
+
+    let to_keep = if matches_allow_list {
+        true
+    } else {
+        // If it doesn't match an allow list, then check if it is denied.
+        // That is, if it matches in a reject filter, then reject.  Otherwise, we always keep it.
+        !filter_check_map.iter().any(|(filter, text)| {
+            if let Some(filter) = filter {
+                if filter.is_list_ignored {
+                    for r in &filter.list {
+                        if r.is_match(text) {
+                            return true;
+                        }
+                    }
+                }
+            }
+            false
+        })
+    };
+
+    to_keep
+}