Revert "refactor: remove heim network usage (#833)"

This reverts commit 913562e7e6867e3eb32cdb034792e01cb6e98e81.

4 files changed, 187 insertions, 93 deletions

diff --git a/src/app/data_harvester.rs b/src/app/data_harvester.rs
index 090fb06f..f289f41f 100644
--- a/src/app/data_harvester.rs
+++ b/src/app/data_harvester.rs
@@ -5,6 +5,7 @@ use std::time::Instant;
 #[cfg(target_os = "linux")]
 use fxhash::FxHashMap;
 
+#[cfg(not(target_os = "linux"))]
 use sysinfo::{System, SystemExt};
 
 #[cfg(feature = "battery")]
@@ -99,6 +100,7 @@ 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)>,
@@ -130,6 +132,7 @@ 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,
@@ -158,8 +161,35 @@ impl DataCollector {
     }
 
     pub fn init(&mut self) {
-        self.refresh_sysinfo();
-        self.mem_total_kb = self.sys.total_memory();
+        #[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();
+            }
+        }
 
         #[cfg(feature = "battery")]
         {
@@ -181,6 +211,17 @@ 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) {
@@ -199,44 +240,28 @@ impl DataCollector {
         self.show_average_cpu = show_average_cpu;
     }
 
-    fn refresh_sysinfo(&mut self) {
+    pub async fn update_data(&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();
             }
-        }
-
-        #[cfg(target_os = "freebsd")]
-        {
-            if self.widgets_to_harvest.use_mem {
-                self.sys.refresh_memory();
+            if cfg!(target_os = "windows") && self.widgets_to_harvest.use_net {
+                self.sys.refresh_networks();
             }
-
-            if self.widgets_to_harvest.use_disk {
+            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_disk {
                 self.sys.refresh_disks();
             }
-        }
-
-        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();
+            if cfg!(target_os = "freebsd") && self.widgets_to_harvest.use_mem {
+                self.sys.refresh_memory();
             }
         }
-    }
-
-    pub async fn update_data(&mut self) {
-        self.refresh_sysinfo();
 
         let current_instant = std::time::Instant::now();
 
@@ -352,21 +377,31 @@ impl DataCollector {
             }
         }
 
-        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 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,
+                )
+            }
+        };
         let mem_data_fut = {
             #[cfg(not(target_os = "freebsd"))]
             {
@@ -391,7 +426,20 @@ impl DataCollector {
         );
         let disk_io_usage_fut = disks::get_io_usage(self.widgets_to_harvest.use_disk);
 
-        let (mem_res, disk_res, io_res) = join!(mem_data_fut, disk_data_fut, disk_io_usage_fut,);
+        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;
+        }
 
         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 62cd7b32..3497ceca 100644
--- a/src/app/data_harvester/network.rs
+++ b/src/app/data_harvester/network.rs
@@ -1,21 +1,24 @@
-//! Data collection for network usage/IO. This is handled by sysinfo.
-
-pub mod sysinfo;
-pub use self::sysinfo::*;
+//! 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::*;
+    }
+}
 
 #[derive(Default, Clone, Debug)]
-/// Harvested network data. Note that all units in bits.
+/// 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,
 }
 
@@ -25,34 +28,3 @@ 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
new file mode 100644
index 00000000..b7980ddb
--- /dev/null
+++ b/src/app/data_harvester/network/heim.rs
@@ -0,0 +1,71 @@
+//! 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 9fb4b38f..634cf779 100644
--- a/src/app/data_harvester/network/sysinfo.rs
+++ b/src/app/data_harvester/network/sysinfo.rs
@@ -3,12 +3,17 @@
 use super::NetworkHarvest;
 use std::time::Instant;
 
-pub fn get_network_data(
+pub async fn get_network_data(
     sys: &sysinfo::System, prev_net_access_time: Instant, prev_net_rx: &mut u64,
-    prev_net_tx: &mut u64, current_time: Instant, filter: &Option<crate::app::Filter>,
-) -> crate::utils::error::Result<NetworkHarvest> {
+    prev_net_tx: &mut u64, curr_time: Instant, actually_get: bool,
+    filter: &Option<crate::app::Filter>,
+) -> crate::utils::error::Result<Option<NetworkHarvest>> {
     use sysinfo::{NetworkExt, SystemExt};
 
+    if !actually_get {
+        return Ok(None);
+    }
+
     let mut total_rx: u64 = 0;
     let mut total_tx: u64 = 0;
 
@@ -33,9 +38,7 @@ pub fn get_network_data(
         }
     }
 
-    let elapsed_time = current_time
-        .duration_since(prev_net_access_time)
-        .as_secs_f64();
+    let elapsed_time = curr_time.duration_since(prev_net_access_time).as_secs_f64();
 
     let (rx, tx) = if elapsed_time == 0.0 {
         (0, 0)
@@ -48,10 +51,10 @@ pub fn get_network_data(
 
     *prev_net_rx = total_rx;
     *prev_net_tx = total_tx;
-    Ok(NetworkHarvest {
+    Ok(Some(NetworkHarvest {
         rx,
         tx,
         total_rx,
         total_tx,
-    })
+    }))
 }