diff options
author | Clement Tsang <34804052+ClementTsang@users.noreply.github.com> | 2021-05-15 18:57:02 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2021-05-15 21:57:02 -0400 |
commit | 6847f2ff0ce5827c0779a1dfcb9e6a8a657652b1 (patch) | |
tree | 2b957d6701f06720baf2cab836303400f340e6f2 /src/app/data_harvester/disks | |
parent | 39c5ee991e8a02a72da398433c62f4a1b8c7acbb (diff) |
refactor: split up data collection by OS (#482)
Refactor to split up data collection by OS and/or the backing library. The goal is to make it easier to work with and add new OS support, as opposed to how it was prior where we stored OS-independent implementations all in the same file.
Diffstat (limited to 'src/app/data_harvester/disks')
-rw-r--r-- | src/app/data_harvester/disks/heim/linux.rs | 34 | ||||
-rw-r--r-- | src/app/data_harvester/disks/heim/mod.rs | 154 | ||||
-rw-r--r-- | src/app/data_harvester/disks/heim/windows_macos.rs | 14 | ||||
-rw-r--r-- | src/app/data_harvester/disks/mod.rs | 10 |
4 files changed, 212 insertions, 0 deletions
diff --git a/src/app/data_harvester/disks/heim/linux.rs b/src/app/data_harvester/disks/heim/linux.rs new file mode 100644 index 00000000..cbc99d9f --- /dev/null +++ b/src/app/data_harvester/disks/heim/linux.rs @@ -0,0 +1,34 @@ +//! Linux-specific things for Heim disk data collection. + +use heim::disk::Partition; + +pub fn get_device_name(partition: &Partition) -> String { + if let Some(device) = partition.device() { + // See if this disk is actually mounted elsewhere on Linux... + // This is a workaround to properly map I/O in some cases (i.e. disk encryption), see + // https://github.com/ClementTsang/bottom/issues/419 + if let Ok(path) = std::fs::read_link(device) { + if path.is_absolute() { + path.into_os_string() + } else { + let mut combined_path = std::path::PathBuf::new(); + combined_path.push(device); + combined_path.pop(); // Pop the current file... + combined_path.push(path); + + if let Ok(canon_path) = std::fs::canonicalize(combined_path) { + // Resolve the local path into an absolute one... + canon_path.into_os_string() + } else { + device.to_os_string() + } + } + } else { + device.to_os_string() + } + .into_string() + .unwrap_or_else(|_| "Name Unavailable".to_string()) + } else { + "Name Unavailable".to_string() + } +} diff --git a/src/app/data_harvester/disks/heim/mod.rs b/src/app/data_harvester/disks/heim/mod.rs new file mode 100644 index 00000000..a79d00db --- /dev/null +++ b/src/app/data_harvester/disks/heim/mod.rs @@ -0,0 +1,154 @@ +use crate::app::Filter; + +cfg_if::cfg_if! { + if #[cfg(target_os = "linux")] { + pub mod linux; + pub use linux::*; + } else if #[cfg(any(target_os = "macos", target_os = "windows"))] { + pub mod windows_macos; + pub use windows_macos::*; + } +} + +#[derive(Debug, Clone, Default)] +pub struct DiskHarvest { + pub name: String, + pub mount_point: String, + pub free_space: Option<u64>, + pub used_space: Option<u64>, + pub total_space: Option<u64>, +} + +#[derive(Clone, Debug)] +pub struct IoData { + pub read_bytes: u64, + pub write_bytes: u64, +} + +pub type IoHarvest = std::collections::HashMap<String, Option<IoData>>; + +pub async fn get_io_usage(actually_get: bool) -> crate::utils::error::Result<Option<IoHarvest>> { + if !actually_get { + return Ok(None); + } + + use futures::StreamExt; + + let mut io_hash: std::collections::HashMap<String, Option<IoData>> = + std::collections::HashMap::new(); + + let counter_stream = heim::disk::io_counters().await?; + futures::pin_mut!(counter_stream); + + while let Some(io) = counter_stream.next().await { + if let Ok(io) = io { + let mount_point = io.device_name().to_str().unwrap_or("Name Unavailable"); + + io_hash.insert( + mount_point.to_string(), + Some(IoData { + read_bytes: io.read_bytes().get::<heim::units::information::byte>(), + write_bytes: io.write_bytes().get::<heim::units::information::byte>(), + }), + ); + } + } + + Ok(Some(io_hash)) +} + +pub async fn get_disk_usage( + actually_get: bool, disk_filter: &Option<Filter>, mount_filter: &Option<Filter>, +) -> crate::utils::error::Result<Option<Vec<DiskHarvest>>> { + if !actually_get { + return Ok(None); + } + + use futures::StreamExt; + + let mut vec_disks: Vec<DiskHarvest> = Vec::new(); + let partitions_stream = heim::disk::partitions_physical().await?; + futures::pin_mut!(partitions_stream); + + while let Some(part) = partitions_stream.next().await { + if let Ok(partition) = part { + let name = get_device_name(&partition); + + let mount_point = (partition + .mount_point() + .to_str() + .unwrap_or("Name Unavailable")) + .to_string(); + + // Precedence ordering in the case where name and mount filters disagree, "allow" takes precedence over "deny". + // + // For implementation, we do this as follows: + // 1. Is the entry allowed through any filter? That is, does it match an entry in a filter where `is_list_ignored` is `false`? If so, we always keep this entry. + // 2. Is the entry denied through any filter? That is, does it match an entry in a filter where `is_list_ignored` is `true`? If so, we always deny this entry. + // 3. Anything else is allowed. + + let filter_check_map = [(disk_filter, &name), (mount_filter, &mount_point)]; + + // This represents case 1. That is, if there is a match in an allowing list - if there is, then + // immediately allow it! + let matches_allow_list = filter_check_map.iter().any(|(filter, text)| { + if let Some(filter) = filter { + if !filter.is_list_ignored { + for r in &filter.list { + if r.is_match(text) { + return true; + } + } + } + } + false + }); + + let to_keep = if matches_allow_list { + true + } else { + // If it doesn't match an allow list, then check if it is denied. + // That is, if it matches in a reject filter, then reject. Otherwise, we always keep it. + !filter_check_map.iter().any(|(filter, text)| { + if let Some(filter) = filter { + if filter.is_list_ignored { + for r in &filter.list { + if r.is_match(text) { + return true; + } + } + } + } + false + }) + }; + + if to_keep { + // The usage line can fail in some cases (for example, if you use Void Linux + LUKS, + // see https://github.com/ClementTsang/bottom/issues/419 for details). As such, check + // it like this instead. + if let Ok(usage) = heim::disk::usage(partition.mount_point().to_path_buf()).await { + vec_disks.push(DiskHarvest { + free_space: Some(usage.free().get::<heim::units::information::byte>()), + used_space: Some(usage.used().get::<heim::units::information::byte>()), + total_space: Some(usage.total().get::<heim::units::information::byte>()), + mount_point, + name, + }); + } else { + vec_disks.push(DiskHarvest { + free_space: None, + used_space: None, + total_space: None, + mount_point, + name, + }); + } + } + } + } + + vec_disks.sort_by(|a, b| a.name.cmp(&b.name)); + + Ok(Some(vec_disks)) +} diff --git a/src/app/data_harvester/disks/heim/windows_macos.rs b/src/app/data_harvester/disks/heim/windows_macos.rs new file mode 100644 index 00000000..428733bf --- /dev/null +++ b/src/app/data_harvester/disks/heim/windows_macos.rs @@ -0,0 +1,14 @@ +//! macOS and Windows-specific things for Heim disk data collection. + +use heim::disk::Partition; + +pub fn get_device_name(partition: &Partition) -> String { + if let Some(device) = partition.device() { + device + .to_os_string() + .into_string() + .unwrap_or_else(|_| "Name Unavailable".to_string()) + } else { + "Name Unavailable".to_string() + } +} diff --git a/src/app/data_harvester/disks/mod.rs b/src/app/data_harvester/disks/mod.rs new file mode 100644 index 00000000..e5a52336 --- /dev/null +++ b/src/app/data_harvester/disks/mod.rs @@ -0,0 +1,10 @@ +//! Data collection for disks (IO, usage, space, etc.). +//! +//! For Linux, macOS, and Windows, this is handled by heim. + +cfg_if::cfg_if! { + if #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] { + pub mod heim; + pub use self::heim::*; + } +} |