fix(packets): take render duration into account so as not to lose any packets

3 files changed, 9 insertions, 30 deletions

diff --git a/src/main.rs b/src/main.rs
index 1387490..bbe3554 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -21,6 +21,7 @@ use ::tui::backend::Backend;
 use std::process;
 
 use ::std::io;
+use ::std::time::Instant;
 use ::termion::raw::IntoRawMode;
 use ::tui::backend::TermionBackend;
 
@@ -164,12 +165,10 @@ where
             let ui = ui.clone();
             move || {
                 while running.load(Ordering::Acquire) {
+                    let render_start_time = Instant::now();
+                    let utilization = { network_utilization.lock().unwrap().clone_and_reset() };
                     let connections_to_procs = get_open_sockets();
                     let ip_to_host = { ip_to_host.lock().unwrap().clone() };
-                    let utilization = {
-                        let mut network_utilization = network_utilization.lock().unwrap();
-                        network_utilization.clone_and_reset()
-                    };
                     let mut unresolved_ips = Vec::new();
                     for connection in connections_to_procs.keys() {
                         if !ip_to_host.contains_key(&connection.remote_socket.ip) {
@@ -190,7 +189,8 @@ where
                             ui.draw();
                         }
                     }
-                    park_timeout(time::Duration::from_secs(1));
+                    let render_duration = render_start_time.elapsed();
+                    park_timeout(time::Duration::from_millis(1000) - render_duration);
                 }
                 if !raw_mode {
                     let mut ui = ui.lock().unwrap();
diff --git a/src/network/utilization.rs b/src/network/utilization.rs
index bb3f581..54870dd 100644
--- a/src/network/utilization.rs
+++ b/src/network/utilization.rs
@@ -1,7 +1,6 @@
 use crate::network::{Connection, Direction, Segment};
 
 use ::std::collections::HashMap;
-use ::std::time::SystemTime;
 
 #[derive(Clone)]
 pub struct TotalBandwidth {
@@ -9,27 +8,9 @@ pub struct TotalBandwidth {
     pub total_bytes_uploaded: u128,
 }
 
-impl TotalBandwidth {
-    pub fn increment_bytes_downloaded(&mut self, data_length: u128, reset_time: &SystemTime) {
-        if let Ok(elapsed) = reset_time.elapsed() {
-            if elapsed.as_millis() < 1000 {
-                self.total_bytes_downloaded += data_length;
-            }
-        }
-    }
-    pub fn increment_bytes_uploaded(&mut self, data_length: u128, reset_time: &SystemTime) {
-        if let Ok(elapsed) = reset_time.elapsed() {
-            if elapsed.as_millis() < 1000 {
-                self.total_bytes_uploaded += data_length;
-            }
-        }
-    }
-}
-
 #[derive(Clone)]
 pub struct Utilization {
     pub connections: HashMap<Connection, TotalBandwidth>,
-    reset_time: SystemTime,
 }
 
 impl Utilization {
@@ -37,12 +18,10 @@ impl Utilization {
         let connections = HashMap::new();
         Utilization {
             connections,
-            reset_time: SystemTime::now(),
         }
     }
     pub fn clone_and_reset(&mut self) -> Self {
         let clone = self.clone();
-        self.reset_time = SystemTime::now();
         self.connections.clear();
         clone
     }
@@ -56,10 +35,10 @@ impl Utilization {
                 });
         match seg.direction {
             Direction::Download => {
-                total_bandwidth.increment_bytes_downloaded(seg.data_length, &self.reset_time);
+                total_bandwidth.total_bytes_downloaded += seg.data_length;
             }
             Direction::Upload => {
-                total_bandwidth.increment_bytes_uploaded(seg.data_length, &self.reset_time);
+                total_bandwidth.total_bytes_uploaded += seg.data_length;
             }
         }
     }
diff --git a/src/tests/fakes/fake_input.rs b/src/tests/fakes/fake_input.rs
index 5e238e4..d11d974 100644
--- a/src/tests/fakes/fake_input.rs
+++ b/src/tests/fakes/fake_input.rs
@@ -25,7 +25,7 @@ impl Iterator for KeyboardEvents {
             Some(ev) => match ev {
                 Some(ev) => Some(ev),
                 None => {
-                    thread::sleep(time::Duration::from_secs(1));
+                    thread::sleep(time::Duration::from_millis(900));
                     self.next()
                 }
             },
@@ -147,7 +147,7 @@ pub fn create_fake_lookup_addr(
 pub fn create_fake_on_winch(should_send_winch_event: bool) -> Box<dyn Fn(Box<dyn Fn()>) + Send> {
     Box::new(move |cb| {
         if should_send_winch_event {
-            thread::sleep(time::Duration::from_secs(1));
+            thread::sleep(time::Duration::from_millis(900));
             cb()
         }
     })