refactor: remove heim network usage (#833)

4 files changed, 93 insertions, 187 deletions

diff --git a/src/app/data_harvester.rs b/src/app/data_harvester.rs
index f289f41f..090fb06f 100644
--- a/src/app/data_harvester.rs
+++ b/src/app/data_harvester.rs
@@ -5,7 +5,6 @@ use std::time::Instant;
 #[cfg(target_os = "linux")]
 use fxhash::FxHashMap;
 
-#[cfg(not(target_os = "linux"))]
 use sysinfo::{System, SystemExt};
 
 #[cfg(feature = "battery")]
@@ -100,7 +99,6 @@ impl Data {
 #[derive(Debug)]
 pub struct DataCollector {
     pub data: Data,
-    #[cfg(not(target_os = "linux"))]
     sys: System,
     previous_cpu_times: Vec<(cpu::PastCpuWork, cpu::PastCpuTotal)>,
     previous_average_cpu_time: Option<(cpu::PastCpuWork, cpu::PastCpuTotal)>,
@@ -132,7 +130,6 @@ impl DataCollector {
     pub fn new(filters: DataFilters) -> Self {
         DataCollector {
             data: Data::default(),
-            #[cfg(not(target_os = "linux"))]
             sys: System::new_with_specifics(sysinfo::RefreshKind::new()),
             previous_cpu_times: vec![],
             previous_average_cpu_time: None,
@@ -161,35 +158,8 @@ impl DataCollector {
     }
 
     pub fn init(&mut self) {
-        #[cfg(target_os = "linux")]
-        {
-            futures::executor::block_on(self.initialize_memory_size());
-        }
-        #[cfg(not(target_os = "linux"))]
-        {
-            self.sys.refresh_memory();
-            self.mem_total_kb = self.sys.total_memory();
-
-            // TODO: Would be good to get this and network list running on a timer instead...?
-            // Refresh components list once...
-            if self.widgets_to_harvest.use_temp {
-                self.sys.refresh_components_list();
-            }
-
-            // Refresh network list once...
-            if cfg!(target_os = "windows") && self.widgets_to_harvest.use_net {
-                self.sys.refresh_networks_list();
-            }
-
-            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_cpu {
-                self.sys.refresh_cpu();
-            }
-
-            // Refresh disk list once...
-            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_disk {
-                self.sys.refresh_disks_list();
-            }
-        }
+        self.refresh_sysinfo();
+        self.mem_total_kb = self.sys.total_memory();
 
         #[cfg(feature = "battery")]
         {
@@ -211,17 +181,6 @@ impl DataCollector {
         std::thread::sleep(std::time::Duration::from_millis(250));
 
         self.data.cleanup();
-
-        // trace!("Enabled widgets to harvest: {:#?}", self.widgets_to_harvest);
-    }
-
-    #[cfg(target_os = "linux")]
-    async fn initialize_memory_size(&mut self) {
-        self.mem_total_kb = if let Ok(mem) = heim::memory::memory().await {
-            mem.total().get::<heim::units::information::kilobyte>()
-        } else {
-            1
-        };
     }
 
     pub fn set_data_collection(&mut self, used_widgets: UsedWidgets) {
@@ -240,28 +199,44 @@ impl DataCollector {
         self.show_average_cpu = show_average_cpu;
     }
 
-    pub async fn update_data(&mut self) {
+    fn refresh_sysinfo(&mut self) {
         #[cfg(not(target_os = "linux"))]
         {
             if self.widgets_to_harvest.use_proc || self.widgets_to_harvest.use_cpu {
                 self.sys.refresh_cpu();
             }
-            if self.widgets_to_harvest.use_proc {
-                self.sys.refresh_processes();
-            }
+
             if self.widgets_to_harvest.use_temp {
                 self.sys.refresh_components();
             }
-            if cfg!(target_os = "windows") && self.widgets_to_harvest.use_net {
-                self.sys.refresh_networks();
+        }
+
+        #[cfg(target_os = "freebsd")]
+        {
+            if self.widgets_to_harvest.use_mem {
+                self.sys.refresh_memory();
             }
-            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_disk {
+
+            if self.widgets_to_harvest.use_disk {
                 self.sys.refresh_disks();
             }
-            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_mem {
-                self.sys.refresh_memory();
+        }
+
+        if self.widgets_to_harvest.use_net {
+            // TODO: Would be good to get this and network list running on a timer instead...?
+            self.sys.refresh_networks();
+        }
+
+        #[cfg(not(target_os = "linux"))]
+        {
+            if self.widgets_to_harvest.use_proc {
+                self.sys.refresh_processes();
             }
         }
+    }
+
+    pub async fn update_data(&mut self) {
+        self.refresh_sysinfo();
 
         let current_instant = std::time::Instant::now();
 
@@ -377,31 +352,21 @@ impl DataCollector {
             }
         }
 
-        let network_data_fut = {
-            #[cfg(any(target_os = "windows", target_os = "freebsd"))]
-            {
-                network::get_network_data(
-                    &self.sys,
-                    self.last_collection_time,
-                    &mut self.total_rx,
-                    &mut self.total_tx,
-                    current_instant,
-                    self.widgets_to_harvest.use_net,
-                    &self.filters.net_filter,
-                )
-            }
-            #[cfg(not(any(target_os = "windows", target_os = "freebsd")))]
-            {
-                network::get_network_data(
-                    self.last_collection_time,
-                    &mut self.total_rx,
-                    &mut self.total_tx,
-                    current_instant,
-                    self.widgets_to_harvest.use_net,
-                    &self.filters.net_filter,
-                )
+        if self.widgets_to_harvest.use_net {
+            if let Ok(net_data) = network::get_network_data(
+                &self.sys,
+                self.last_collection_time,
+                &mut self.total_rx,
+                &mut self.total_tx,
+                current_instant,
+                &self.filters.net_filter,
+            ) {
+                self.total_rx = net_data.total_rx;
+                self.total_tx = net_data.total_tx;
+                self.data.network = Some(net_data);
             }
-        };
+        }
+
         let mem_data_fut = {
             #[cfg(not(target_os = "freebsd"))]
             {
@@ -426,20 +391,7 @@ impl DataCollector {
         );
         let disk_io_usage_fut = disks::get_io_usage(self.widgets_to_harvest.use_disk);
 
-        let (net_data, mem_res, disk_res, io_res) = join!(
-            network_data_fut,
-            mem_data_fut,
-            disk_data_fut,
-            disk_io_usage_fut,
-        );
-
-        if let Ok(net_data) = net_data {
-            if let Some(net_data) = &net_data {
-                self.total_rx = net_data.total_rx;
-                self.total_tx = net_data.total_tx;
-            }
-            self.data.network = net_data;
-        }
+        let (mem_res, disk_res, io_res) = join!(mem_data_fut, disk_data_fut, disk_io_usage_fut,);
 
         if let Ok(memory) = mem_res.ram {
             self.data.memory = memory;
diff --git a/src/app/data_harvester/network.rs b/src/app/data_harvester/network.rs
index 3497ceca..62cd7b32 100644
--- a/src/app/data_harvester/network.rs
+++ b/src/app/data_harvester/network.rs
@@ -1,24 +1,21 @@
-//! Data collection for network usage/IO.
-//!
-//! For Linux and macOS, this is handled by Heim.
-//! For Windows, this is handled by sysinfo.
-
-cfg_if::cfg_if! {
-    if #[cfg(any(target_os = "linux", target_os = "macos"))] {
-        pub mod heim;
-        pub use self::heim::*;
-    } else if #[cfg(any(target_os = "freebsd", target_os = "windows"))] {
-        pub mod sysinfo;
-        pub use self::sysinfo::*;
-    }
-}
+//! Data collection for network usage/IO. This is handled by sysinfo.
+
+pub mod sysinfo;
+pub use self::sysinfo::*;
 
 #[derive(Default, Clone, Debug)]
-/// All units in bits.
+/// Harvested network data. Note that all units in bits.
 pub struct NetworkHarvest {
+    /// Current incoming bits/s.
     pub rx: u64,
+
+    /// Current outgoing bits/s.
     pub tx: u64,
+
+    /// Total number of incoming bits.
     pub total_rx: u64,
+
+    /// Total number of outgoing bits.
     pub total_tx: u64,
 }
 
@@ -28,3 +25,34 @@ impl NetworkHarvest {
         self.tx = 0;
     }
 }
+
+#[cfg(test)]
+mod test {
+    use std::{
+        thread::sleep,
+        time::{Duration, Instant},
+    };
+
+    use super::*;
+    use ::sysinfo::{System, SystemExt};
+
+    #[test]
+    fn test_getting_network() {
+        let sys = System::new_all();
+        let prev = Instant::now();
+
+        sleep(Duration::from_secs(2));
+        let mut prev_rx = 0;
+        let mut prev_tx = 0;
+
+        get_network_data(
+            &sys,
+            prev,
+            &mut prev_rx,
+            &mut prev_tx,
+            Instant::now(),
+            &None,
+        )
+        .unwrap();
+    }
+}
diff --git a/src/app/data_harvester/network/heim.rs b/src/app/data_harvester/network/heim.rs
deleted file mode 100644
index b7980ddb..00000000
--- a/src/app/data_harvester/network/heim.rs
+++ /dev/null
@@ -1,71 +0,0 @@
-//! Gets network data via heim.
-
-use super::NetworkHarvest;
-use std::time::Instant;
-
-// TODO: Eventually make it so that this thing also takes individual usage into account, so we can show per-interface!
-pub async fn get_network_data(
-    prev_net_access_time: Instant, prev_net_rx: &mut u64, prev_net_tx: &mut u64,
-    curr_time: Instant, actually_get: bool, filter: &Option<crate::app::Filter>,
-) -> crate::utils::error::Result<Option<NetworkHarvest>> {
-    use futures::StreamExt;
-
-    if !actually_get {
-        return Ok(None);
-    }
-
-    let io_data = heim::net::io_counters().await?;
-    futures::pin_mut!(io_data);
-    let mut total_rx: u64 = 0;
-    let mut total_tx: u64 = 0;
-
-    while let Some(io) = io_data.next().await {
-        if let Ok(io) = io {
-            let to_keep = if let Some(filter) = filter {
-                if filter.is_list_ignored {
-                    let mut ret = true;
-                    for r in &filter.list {
-                        if r.is_match(io.interface()) {
-                            ret = false;
-                            break;
-                        }
-                    }
-                    ret
-                } else {
-                    true
-                }
-            } else {
-                true
-            };
-
-            if to_keep {
-                // TODO: Use bytes as the default instead, perhaps?
-                // Since you might have to do a double conversion (bytes -> bits -> bytes) in some cases;
-                // but if you stick to bytes, then in the bytes, case, you do no conversion, and in the bits case,
-                // you only do one conversion...
-                total_rx += io.bytes_recv().get::<heim::units::information::bit>();
-                total_tx += io.bytes_sent().get::<heim::units::information::bit>();
-            }
-        }
-    }
-
-    let elapsed_time = curr_time.duration_since(prev_net_access_time).as_secs_f64();
-
-    let (rx, tx) = if elapsed_time == 0.0 {
-        (0, 0)
-    } else {
-        (
-            ((total_rx.saturating_sub(*prev_net_rx)) as f64 / elapsed_time) as u64,
-            ((total_tx.saturating_sub(*prev_net_tx)) as f64 / elapsed_time) as u64,
-        )
-    };
-
-    *prev_net_rx = total_rx;
-    *prev_net_tx = total_tx;
-    Ok(Some(NetworkHarvest {
-        rx,
-        tx,
-        total_rx,
-        total_tx,
-    }))
-}
diff --git a/src/app/data_harvester/network/sysinfo.rs b/src/app/data_harvester/network/sysinfo.rs
index 634cf779..9fb4b38f 100644
--- a/src/app/data_harvester/network/sysinfo.rs
+++ b/src/app/data_harvester/network/sysinfo.rs
@@ -3,17 +3,12 @@
 use super::NetworkHarvest;
 use std::time::Instant;
 
-pub async fn get_network_data(
+pub fn get_network_data(
     sys: &sysinfo::System, prev_net_access_time: Instant, prev_net_rx: &mut u64,
-    prev_net_tx: &mut u64, curr_time: Instant, actually_get: bool,
-    filter: &Option<crate::app::Filter>,
-) -> crate::utils::error::Result<Option<NetworkHarvest>> {
+    prev_net_tx: &mut u64, current_time: Instant, filter: &Option<crate::app::Filter>,
+) -> crate::utils::error::Result<NetworkHarvest> {
     use sysinfo::{NetworkExt, SystemExt};
 
-    if !actually_get {
-        return Ok(None);
-    }
-
     let mut total_rx: u64 = 0;
     let mut total_tx: u64 = 0;
 
@@ -38,7 +33,9 @@ pub async fn get_network_data(
         }
     }
 
-    let elapsed_time = curr_time.duration_since(prev_net_access_time).as_secs_f64();
+    let elapsed_time = current_time
+        .duration_since(prev_net_access_time)
+        .as_secs_f64();
 
     let (rx, tx) = if elapsed_time == 0.0 {
         (0, 0)
@@ -51,10 +48,10 @@ pub async fn get_network_data(
 
     *prev_net_rx = total_rx;
     *prev_net_tx = total_tx;
-    Ok(Some(NetworkHarvest {
+    Ok(NetworkHarvest {
         rx,
         tx,
         total_rx,
         total_tx,
-    }))
+    })
 }