diff --git a/src/display/components/display_bandwidth.rs b/src/display/components/display_bandwidth.rs index 35fe9efcf..e7bfaab97 100644 --- a/src/display/components/display_bandwidth.rs +++ b/src/display/components/display_bandwidth.rs @@ -1,22 +1,21 @@ use ::std::fmt; -// Should we make a `DisplayMode` enum with `Rate` and `Total`? -// Instead of the ambiguous bool? - -// This might be better as a traditional struct now? -pub struct DisplayBandwidth(pub f64, pub bool); +pub struct DisplayBandwidth { + pub bandwidth: f64, + pub as_rate: bool, +} impl fmt::Display for DisplayBandwidth { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let suffix = if self.1 { "" } else { "ps" }; // Should this be the other way around? - if self.0 > 999_999_999.0 { - write!(f, "{:.2}GB{}", self.0 / 1_000_000_000.0, suffix) - } else if self.0 > 999_999.0 { - write!(f, "{:.2}MB{}", self.0 / 1_000_000.0, suffix) - } else if self.0 > 999.0 { - write!(f, "{:.2}KB{}", self.0 / 1000.0, suffix) + let suffix = if self.as_rate { "" } else { "ps" }; + if self.bandwidth > 999_999_999.0 { + write!(f, "{:.2}GB{}", self.bandwidth / 1_000_000_000.0, suffix) + } else if self.bandwidth > 999_999.0 { + write!(f, "{:.2}MB{}", self.bandwidth / 1_000_000.0, suffix) + } else if self.bandwidth > 999.0 { + write!(f, "{:.2}KB{}", self.bandwidth / 1000.0, suffix) } else { - write!(f, "{}B{}", self.0, suffix) + write!(f, "{}B{}", self.bandwidth, suffix) } } } diff --git a/src/display/components/table.rs b/src/display/components/table.rs index 190b8c457..451c61f87 100644 --- a/src/display/components/table.rs +++ b/src/display/components/table.rs @@ -9,46 +9,22 @@ use ::tui::widgets::{Block, Borders, Row, Widget}; use crate::display::{Bandwidth, DisplayBandwidth, UIState}; use crate::network::{display_connection_string, display_ip_or_host}; -use ::std::iter::FromIterator; use ::std::net::IpAddr; fn display_upload_and_download(bandwidth: &impl Bandwidth, total: bool) -> String { format!( "{} / {}", - DisplayBandwidth(bandwidth.get_total_bytes_uploaded() as f64, total), - DisplayBandwidth(bandwidth.get_total_bytes_downloaded() as f64, total) + DisplayBandwidth{ + bandwidth: bandwidth.get_total_bytes_uploaded() as f64, + as_rate: !total, + }, + DisplayBandwidth{ + bandwidth: bandwidth.get_total_bytes_downloaded() as f64, + as_rate: !total, + }, ) } -// Why is this returning the sorted list when it mutates things? -fn sort_by_bandwidth<'a, T>( - list: &'a mut Vec<(T, &impl Bandwidth)>, -) -> &'a Vec<(T, &'a impl Bandwidth)> { - list.sort_by(|(_, a), (_, b)| { - let a_highest = if a.get_total_bytes_downloaded() > a.get_total_bytes_uploaded() { - a.get_total_bytes_downloaded() - } else { - a.get_total_bytes_uploaded() - }; - let b_highest = if b.get_total_bytes_downloaded() > b.get_total_bytes_uploaded() { - b.get_total_bytes_downloaded() - } else { - b.get_total_bytes_uploaded() - }; - b_highest.cmp(&a_highest) - }); - list -} -// I'd suggest using this: -// fn sort_by_bandwidth(list: &mut Vec<(T, impl Bandwidth)>) { -// list.sort_by_key(|(_, b)| { -// cmp::Reverse(cmp::max( -// b.get_total_bytes_downloaded(), -// b.get_total_bytes_uploaded(), -// )) -// }); -// } - pub enum ColumnCount { Two, Three, @@ -87,9 +63,8 @@ fn truncate_middle(row: &str, max_length: u16) -> String { impl<'a> Table<'a> { pub fn create_connections_table(state: &UIState, ip_to_host: &HashMap) -> Self { - let mut connections_list = Vec::from_iter(&state.connections); - sort_by_bandwidth(&mut connections_list); - let connections_rows = connections_list + let connections_rows = state + .connections .iter() .map(|(connection, connection_data)| { vec![ @@ -99,7 +74,7 @@ impl<'a> Table<'a> { &connection_data.interface_name, ), connection_data.process_name.to_string(), - display_upload_and_download(*connection_data, state.cumulative_mode), + display_upload_and_download(connection_data, state.cumulative_mode), ] }) .collect(); @@ -142,15 +117,14 @@ impl<'a> Table<'a> { } } pub fn create_processes_table(state: &UIState) -> Self { - let mut processes_list = Vec::from_iter(&state.processes); - sort_by_bandwidth(&mut processes_list); - let processes_rows = processes_list + let processes_rows = state + .processes .iter() .map(|(process_name, data_for_process)| { vec![ (*process_name).to_string(), data_for_process.connection_count.to_string(), - display_upload_and_download(*data_for_process, state.cumulative_mode), + display_upload_and_download(data_for_process, state.cumulative_mode), ] }) .collect(); @@ -196,16 +170,15 @@ impl<'a> Table<'a> { state: &UIState, ip_to_host: &HashMap, ) -> Self { - let mut remote_addresses_list = Vec::from_iter(&state.remote_addresses); - sort_by_bandwidth(&mut remote_addresses_list); - let remote_addresses_rows = remote_addresses_list + let remote_addresses_rows = state + .remote_addresses .iter() .map(|(remote_address, data_for_remote_address)| { - let remote_address = display_ip_or_host(**remote_address, &ip_to_host); + let remote_address = display_ip_or_host(*remote_address, &ip_to_host); vec![ remote_address, data_for_remote_address.connection_count.to_string(), - display_upload_and_download(*data_for_remote_address, state.cumulative_mode), + display_upload_and_download(data_for_remote_address, state.cumulative_mode), ] }) .collect(); diff --git a/src/display/components/total_bandwidth.rs b/src/display/components/total_bandwidth.rs index c7c87553e..64c08a5b8 100644 --- a/src/display/components/total_bandwidth.rs +++ b/src/display/components/total_bandwidth.rs @@ -25,8 +25,14 @@ impl<'a> TotalBandwidth<'a> { [Text::styled( format!( " Total Up / Down: {} / {} {}", - DisplayBandwidth(self.state.total_bytes_uploaded as f64, c_mode), - DisplayBandwidth(self.state.total_bytes_downloaded as f64, c_mode), + DisplayBandwidth{ + bandwidth: self.state.total_bytes_uploaded as f64, + as_rate: !c_mode, + }, + DisplayBandwidth{ + bandwidth: self.state.total_bytes_downloaded as f64, + as_rate: !c_mode, + }, paused_str ), Style::default().fg(color).modifier(Modifier::BOLD), diff --git a/src/display/ui.rs b/src/display/ui.rs index deb03f682..f0ed57fe4 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -30,12 +30,11 @@ where let mut terminal = Terminal::new(terminal_backend).unwrap(); terminal.clear().unwrap(); terminal.hide_cursor().unwrap(); + let mut state: UIState = Default::default(); + state.cumulative_mode = opts.total_utilization; Ui { terminal, - state: UIState { - cumulative_mode: opts.total_utilization, - ..Default::default() - }, + state, ip_to_host: Default::default(), opts, } diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index 586b449e0..a65e610ad 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -1,5 +1,6 @@ use ::std::cmp; -use ::std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use ::std::collections::{HashMap, HashSet, VecDeque}; +use ::std::hash::Hash; use ::std::iter::FromIterator; use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; @@ -11,7 +12,8 @@ static MAX_BANDWIDTH_ITEMS: usize = 1000; pub trait Bandwidth { fn get_total_bytes_downloaded(&self) -> u128; fn get_total_bytes_uploaded(&self) -> u128; - fn combine_bandwidth(&mut self, other: &impl Bandwidth); + fn combine_bandwidth(&mut self, other: &Self); + fn divide_by(&mut self, amount: u128); } #[derive(Clone, Default)] @@ -29,44 +31,39 @@ pub struct ConnectionData { pub interface_name: String, } -impl NetworkData { - pub fn divide_by(&mut self, amount: u128) { - self.total_bytes_downloaded /= amount; - self.total_bytes_uploaded /= amount; - } -} - -impl ConnectionData { - pub fn divide_by(&mut self, amount: u128) { - self.total_bytes_downloaded /= amount; - self.total_bytes_uploaded /= amount; - } -} - -impl Bandwidth for ConnectionData { +impl Bandwidth for NetworkData { fn get_total_bytes_downloaded(&self) -> u128 { self.total_bytes_downloaded } fn get_total_bytes_uploaded(&self) -> u128 { self.total_bytes_uploaded } - fn combine_bandwidth(&mut self, other: &impl Bandwidth) { + fn combine_bandwidth(&mut self, other: &NetworkData) { self.total_bytes_downloaded += other.get_total_bytes_downloaded(); self.total_bytes_uploaded += other.get_total_bytes_uploaded(); + self.connection_count = other.connection_count; + } + fn divide_by(&mut self, amount: u128) { + self.total_bytes_downloaded /= amount; + self.total_bytes_uploaded /= amount; } } -impl Bandwidth for NetworkData { +impl Bandwidth for ConnectionData { fn get_total_bytes_downloaded(&self) -> u128 { self.total_bytes_downloaded } fn get_total_bytes_uploaded(&self) -> u128 { self.total_bytes_uploaded } - fn combine_bandwidth(&mut self, other: &impl Bandwidth) { + fn combine_bandwidth(&mut self, other: &ConnectionData) { self.total_bytes_downloaded += other.get_total_bytes_downloaded(); self.total_bytes_uploaded += other.get_total_bytes_uploaded(); } + fn divide_by(&mut self, amount: u128) { + self.total_bytes_downloaded /= amount; + self.total_bytes_uploaded /= amount; + } } pub struct UtilizationData { @@ -76,20 +73,16 @@ pub struct UtilizationData { #[derive(Default)] pub struct UIState { - // We aren't really taking advantage of the self-sorting of the BTreeMap here, - // but it would be rather difficult to, as it sorts by key, not value. Perhaps we - // should just use a HashMap? - // We could also use a HashMap internally for insertions and merging, but cache a - // sorted vector so that doesn't need sorting during the display step. This would - // also fix Bandwhich eating CPU cycles when paused, but is a bit of a non-issue - // with `MAX_BANDWIDTH_ITEMS` in place. - pub processes: BTreeMap, - pub remote_addresses: BTreeMap, - pub connections: BTreeMap, + pub processes: Vec<(String, NetworkData)>, + pub remote_addresses: Vec<(IpAddr, NetworkData)>, + pub connections: Vec<(Connection, ConnectionData)>, pub total_bytes_downloaded: u128, pub total_bytes_uploaded: u128, pub cumulative_mode: bool, - pub utilization_data: VecDeque, + utilization_data: VecDeque, + processes_map: HashMap, + remote_addresses_map: HashMap, + connections_map: HashMap, } impl UIState { @@ -125,9 +118,9 @@ impl UIState { if self.utilization_data.len() > RECALL_LENGTH { self.utilization_data.pop_front(); } - let mut processes: BTreeMap = BTreeMap::new(); - let mut remote_addresses: BTreeMap = BTreeMap::new(); - let mut connections: BTreeMap = BTreeMap::new(); + let mut processes: HashMap = HashMap::new(); + let mut remote_addresses: HashMap = HashMap::new(); + let mut connections: HashMap = HashMap::new(); let mut total_bytes_downloaded: u128 = 0; let mut total_bytes_uploaded: u128 = 0; @@ -190,27 +183,27 @@ impl UIState { } if self.cumulative_mode { - merge_bandwidth(&mut self.processes, processes); - merge_bandwidth(&mut self.remote_addresses, remote_addresses); - merge_bandwidth(&mut self.connections, connections); + merge_bandwidth(&mut self.processes_map, processes); + merge_bandwidth(&mut self.remote_addresses_map, remote_addresses); + merge_bandwidth(&mut self.connections_map, connections); self.total_bytes_downloaded += total_bytes_downloaded / divide_by; self.total_bytes_uploaded += total_bytes_uploaded / divide_by; } else { - self.processes = processes; - self.remote_addresses = remote_addresses; - self.connections = connections; + self.processes_map = processes; + self.remote_addresses_map = remote_addresses; + self.connections_map = connections; self.total_bytes_downloaded = total_bytes_downloaded / divide_by; self.total_bytes_uploaded = total_bytes_uploaded / divide_by; } - prune_map(&mut self.processes); - prune_map(&mut self.remote_addresses); - prune_map(&mut self.connections); + self.processes = sort_and_prune(&mut self.processes_map); + self.remote_addresses = sort_and_prune(&mut self.remote_addresses_map); + self.connections = sort_and_prune(&mut self.connections_map); } } -fn merge_bandwidth(self_map: &mut BTreeMap, other_map: BTreeMap) +fn merge_bandwidth(self_map: &mut HashMap, other_map: HashMap) where - K: Eq + Ord, + K: Eq + Hash, V: Bandwidth, { for (key, b_other) in other_map { @@ -221,22 +214,21 @@ where } } -fn prune_map(map: &mut BTreeMap) { - if map.len() > MAX_BANDWIDTH_ITEMS { - let mut bandwidth_list = Vec::from_iter(map.clone()); - sort_by_bandwidth(&mut bandwidth_list); +fn sort_and_prune(map: &mut HashMap) -> Vec<(K, V)> +where + K: Eq + Hash + Clone, + V: Bandwidth + Clone, +{ + let mut bandwidth_list = Vec::from_iter(map.clone()); + bandwidth_list.sort_by_key(|(_, b)| { + cmp::Reverse(b.get_total_bytes_downloaded() + b.get_total_bytes_uploaded()) + }); + + if bandwidth_list.len() > MAX_BANDWIDTH_ITEMS { for (key, _) in &bandwidth_list[MAX_BANDWIDTH_ITEMS..] { map.remove(key); } } -} -// This is duplicated from table.rs temporarily -fn sort_by_bandwidth(list: &mut Vec<(T, impl Bandwidth)>) { - list.sort_by_key(|(_, b)| { - cmp::Reverse(cmp::max( - b.get_total_bytes_downloaded(), - b.get_total_bytes_uploaded(), - )) - }); + bandwidth_list }