diff options
Diffstat (limited to 'src/display')
-rw-r--r-- | src/display/ui.rs | 4 | ||||
-rw-r--r-- | src/display/ui_state.rs | 119 |
2 files changed, 72 insertions, 51 deletions
diff --git a/src/display/ui.rs b/src/display/ui.rs index 62dfb2a..dd8ecfb 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -98,8 +98,8 @@ where utilization: Utilization, ip_to_host: HashMap<Ipv4Addr, String>, ) { - self.state = UIState::new(connections_to_procs, utilization, &self.state); - self.ip_to_host = ip_to_host; + self.state.update(connections_to_procs, utilization); + self.ip_to_host.extend(ip_to_host); } pub fn end(&mut self) { self.terminal.clear().unwrap(); diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index a5bdfcb..4b42039 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -1,9 +1,10 @@ -use ::std::collections::{BTreeMap, HashMap}; +use ::std::collections::{BTreeMap, HashMap, VecDeque}; use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use crate::network::{Connection, LocalSocket, Utilization}; static BANDWIDTH_DECAY_FACTOR: f32 = 0.5; +static RECALL_LENGTH: usize = 5; pub trait Bandwidth { fn get_total_bytes_downloaded(&self) -> u128; @@ -32,6 +33,20 @@ pub struct ConnectionData { pub interface_name: String, } +impl NetworkData { + pub fn divide_by(&mut self, amount: usize) { + self.total_bytes_downloaded = self.total_bytes_downloaded / amount as u128; + self.total_bytes_uploaded = self.total_bytes_uploaded / amount as u128; + } +} + +impl ConnectionData { + pub fn divide_by(&mut self, amount: usize) { + self.total_bytes_downloaded = self.total_bytes_downloaded / amount as u128; + self.total_bytes_uploaded = self.total_bytes_uploaded / amount as u128; + } +} + fn calc_avg_bandwidth(prev_bandwidth: u128, curr_bandwidth: u128) -> u128 { if prev_bandwidth == 0 { curr_bandwidth @@ -77,6 +92,11 @@ impl Bandwidth for NetworkData { } } +pub struct UtilizationData { + connections_to_procs: HashMap<LocalSocket, String>, + network_utilization: Utilization, +} + #[derive(Default)] pub struct UIState { pub processes: BTreeMap<String, NetworkData>, @@ -84,7 +104,8 @@ pub struct UIState { pub connections: BTreeMap<Connection, ConnectionData>, pub total_bytes_downloaded: u128, pub total_bytes_uploaded: u128, -} + utilization_data: VecDeque<UtilizationData>, + } impl UIState { fn get_proc_name<'a>( @@ -107,64 +128,64 @@ impl UIState { }) } } - pub fn new( + pub fn update( + &mut self, connections_to_procs: HashMap<LocalSocket, String>, network_utilization: Utilization, - old_state: &UIState, - ) -> Self { + ) { + self.utilization_data.push_back(UtilizationData { connections_to_procs, network_utilization }); + if self.utilization_data.len() > RECALL_LENGTH { + self.utilization_data.pop_front(); + } let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new(); let mut remote_addresses: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new(); let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new(); let mut total_bytes_downloaded: u128 = 0; let mut total_bytes_uploaded: u128 = 0; - for (connection, connection_info) in network_utilization.connections { - let connection_data = connections.entry(connection).or_default(); - let data_for_remote_address = remote_addresses - .entry(connection.remote_socket.ip) - .or_default(); - connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded; - connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded; - connection_data.interface_name = connection_info.interface_name; - data_for_remote_address.total_bytes_downloaded += - connection_info.total_bytes_downloaded; - data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded; - data_for_remote_address.connection_count += 1; - total_bytes_downloaded += connection_info.total_bytes_downloaded; - total_bytes_uploaded += connection_info.total_bytes_uploaded; + for state in self.utilization_data.iter().rev() { + let connections_to_procs = &state.connections_to_procs; + let network_utilization = &state.network_utilization; + for (connection, connection_info) in &network_utilization.connections { + let connection_data = connections.entry(connection.clone()).or_default(); + let data_for_remote_address = remote_addresses + .entry(connection.remote_socket.ip) + .or_default(); + connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded; + connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded; + connection_data.interface_name = connection_info.interface_name.clone(); + data_for_remote_address.total_bytes_downloaded += + connection_info.total_bytes_downloaded; + data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded; + 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; + 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>"); } - } else { - connection_data.process_name = String::from("<UNKNOWN>"); } } - UIState { - processes, - remote_addresses, - connections, - total_bytes_downloaded, - total_bytes_uploaded, + for (_, network_data) in processes.iter_mut() { + network_data.divide_by(self.utilization_data.len()) + } + for (_, network_data) in remote_addresses.iter_mut() { + network_data.divide_by(self.utilization_data.len()) + } + for (_, connection_data) in connections.iter_mut() { + connection_data.divide_by(self.utilization_data.len()) } + self.processes = processes; + self.remote_addresses = remote_addresses; + self.connections = connections; + self.total_bytes_downloaded = total_bytes_downloaded; + self.total_bytes_uploaded = total_bytes_uploaded; } } |