Display weighted average of bandwidth

3 files changed, 76 insertions, 15 deletions

diff --git a/src/display/components/table.rs b/src/display/components/table.rs
index 2b823d7..ae29d3c 100644
--- a/src/display/components/table.rs
+++ b/src/display/components/table.rs
@@ -23,8 +23,8 @@ const THIRD_COLUMN_WIDTHS: [u16; 4] = [20, 20, 20, 20];
 fn display_upload_and_download(bandwidth: &impl Bandwidth) -> String {
     format!(
         "{} / {}",
-        DisplayBandwidth(bandwidth.get_total_bytes_uploaded() as f64),
-        DisplayBandwidth(bandwidth.get_total_bytes_downloaded() as f64)
+        DisplayBandwidth(bandwidth.get_avg_bytes_uploaded() as f64),
+        DisplayBandwidth(bandwidth.get_avg_bytes_downloaded() as f64)
     )
 }
 
diff --git a/src/display/ui.rs b/src/display/ui.rs
index e61f5e7..62dfb2a 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -98,7 +98,7 @@ where
         utilization: Utilization,
         ip_to_host: HashMap<Ipv4Addr, String>,
     ) {
-        self.state = UIState::new(connections_to_procs, utilization);
+        self.state = UIState::new(connections_to_procs, utilization, &self.state);
         self.ip_to_host = ip_to_host;
     }
     pub fn end(&mut self) {
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index acbd04f..290f972 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -3,15 +3,22 @@ use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
 use crate::network::{Connection, LocalSocket, Utilization};
 
+static BANDWIDTH_DECAY_FACTOR: f32 = 0.5;
+
 pub trait Bandwidth {
     fn get_total_bytes_downloaded(&self) -> u128;
     fn get_total_bytes_uploaded(&self) -> u128;
+
+    fn get_avg_bytes_downloaded(&self) -> u128;
+    fn get_avg_bytes_uploaded(&self) -> u128;
 }
 
 #[derive(Default)]
 pub struct NetworkData {
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
+    pub prev_total_bytes_downloaded: u128,
+    pub prev_total_bytes_uploaded: u128,
     pub connection_count: u128,
 }
 
@@ -19,6 +26,8 @@ pub struct NetworkData {
 pub struct ConnectionData {
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
+    pub prev_total_bytes_downloaded: u128,
+    pub prev_total_bytes_uploaded: u128,
     pub process_name: String,
     pub interface_name: String,
 }
@@ -30,6 +39,24 @@ impl Bandwidth for ConnectionData {
     fn get_total_bytes_downloaded(&self) -> u128 {
         self.total_bytes_downloaded
     }
+    fn get_avg_bytes_uploaded(&self) -> u128 {
+        if self.prev_total_bytes_uploaded == 0 {
+            self.total_bytes_uploaded
+        } else {
+            (self.prev_total_bytes_uploaded as f32 * BANDWIDTH_DECAY_FACTOR
+                + (1.0 - BANDWIDTH_DECAY_FACTOR) * self.total_bytes_uploaded as f32)
+                as u128
+        }
+    }
+    fn get_avg_bytes_downloaded(&self) -> u128 {
+        if self.prev_total_bytes_downloaded == 0 {
+            self.total_bytes_downloaded
+        } else {
+            (self.prev_total_bytes_downloaded as f32 * BANDWIDTH_DECAY_FACTOR
+                + (1.0 - BANDWIDTH_DECAY_FACTOR) * self.total_bytes_downloaded as f32)
+                as u128
+        }
+    }
 }
 
 impl Bandwidth for NetworkData {
@@ -39,6 +66,24 @@ impl Bandwidth for NetworkData {
     fn get_total_bytes_downloaded(&self) -> u128 {
         self.total_bytes_downloaded
     }
+    fn get_avg_bytes_uploaded(&self) -> u128 {
+        if self.prev_total_bytes_uploaded == 0 {
+            self.total_bytes_uploaded
+        } else {
+            (self.prev_total_bytes_uploaded as f32 * BANDWIDTH_DECAY_FACTOR
+                + (1.0 - BANDWIDTH_DECAY_FACTOR) * self.total_bytes_uploaded as f32)
+                as u128
+        }
+    }
+    fn get_avg_bytes_downloaded(&self) -> u128 {
+        if self.prev_total_bytes_downloaded == 0 {
+            self.total_bytes_downloaded
+        } else {
+            (self.prev_total_bytes_downloaded as f32 * BANDWIDTH_DECAY_FACTOR
+                + (1.0 - BANDWIDTH_DECAY_FACTOR) * self.total_bytes_downloaded as f32)
+                as u128
+        }
+    }
 }
 
 #[derive(Default)]
@@ -74,6 +119,7 @@ impl UIState {
     pub fn new(
         connections_to_procs: HashMap<LocalSocket, String>,
         network_utilization: Utilization,
+        old_state: &UIState,
     ) -> Self {
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
         let mut remote_addresses: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
@@ -82,18 +128,6 @@ impl UIState {
         let mut total_bytes_uploaded: u128 = 0;
         for (connection, connection_info) in network_utilization.connections {
             let connection_data = connections.entry(connection).or_default();
-
-            if let Some(process_name) =
-                UIState::get_proc_name(&connections_to_procs, &connection.local_socket)
-            {
-                let data_for_process = processes.entry(process_name.clone()).or_default();
-                data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded;
-                data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded;
-                data_for_process.connection_count += 1;
-                connection_data.process_name = process_name.clone();
-            } else {
-                connection_data.process_name = String::from("<UNKNOWN>");
-            }
             let data_for_remote_address = remote_addresses
                 .entry(connection.remote_socket.ip)
                 .or_default();
@@ -106,6 +140,33 @@ impl UIState {
             data_for_remote_address.connection_count += 1;
             total_bytes_downloaded += connection_info.total_bytes_downloaded;
             total_bytes_uploaded += connection_info.total_bytes_uploaded;
+
+            if let Some(process_name) =
+                UIState::get_proc_name(&connections_to_procs, &connection.local_socket)
+            {
+                let data_for_process = processes.entry(process_name.clone()).or_default();
+                data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded;
+                data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+                data_for_process.connection_count += 1;
+                connection_data.process_name = process_name.clone();
+                // Record bandwidth data of last iteration
+                if let Some(prev_connection_info) = old_state.connections.get(&connection) {
+                    // Using previous round's weighted average. Exponential decay
+                    let prev_bytes_downloaded = prev_connection_info.get_avg_bytes_downloaded();
+                    let prev_bytes_uploaded = prev_connection_info.get_avg_bytes_uploaded();
+    
+                    connection_data.prev_total_bytes_downloaded += prev_bytes_downloaded;
+                    connection_data.prev_total_bytes_uploaded += prev_bytes_uploaded;
+    
+                    data_for_process.prev_total_bytes_downloaded += prev_bytes_downloaded;
+                    data_for_process.prev_total_bytes_uploaded += prev_bytes_uploaded;
+    
+                    data_for_remote_address.prev_total_bytes_downloaded += prev_bytes_downloaded;
+                    data_for_remote_address.prev_total_bytes_uploaded += prev_bytes_uploaded;
+                }
+            } else {
+                connection_data.process_name = String::from("<UNKNOWN>");
+            }
         }
         UIState {
             processes,