From 1b98b967a80906a93548cad3aab863646165d533 Mon Sep 17 00:00:00 2001 From: ClementTsang <34804052+ClementTsang@users.noreply.github.com> Date: Wed, 11 Sep 2024 22:35:10 -0400 Subject: [PATCH] temp commit --- src/new_data_collection/collectors/common.rs | 19 ++- src/new_data_collection/collectors/linux.rs | 5 + src/new_data_collection/error.rs | 2 +- .../sources/common/disk.rs | 127 ++++++++++++++++++ src/new_data_collection/sources/common/mod.rs | 1 + .../sources/sysinfo/disk.rs | 0 6 files changed, 143 insertions(+), 11 deletions(-) create mode 100644 src/new_data_collection/sources/common/disk.rs create mode 100644 src/new_data_collection/sources/sysinfo/disk.rs diff --git a/src/new_data_collection/collectors/common.rs b/src/new_data_collection/collectors/common.rs index 6d7e507a..d1695664 100644 --- a/src/new_data_collection/collectors/common.rs +++ b/src/new_data_collection/collectors/common.rs @@ -2,7 +2,9 @@ use crate::new_data_collection::{ error::CollectionResult, - sources::common::{processes::ProcessHarvest, temperature::TemperatureData}, + sources::common::{ + disk::DiskHarvest, processes::ProcessHarvest, temperature::TemperatureData, + }, }; /// The trait representing what a per-platform data collector should implement. @@ -11,17 +13,14 @@ pub(crate) trait DataCollector { /// /// Note that depending on the implementation, this may /// not actually need to do anything. - fn refresh_data(&mut self) -> CollectionResult<()> { - Ok(()) - } + fn refresh_data(&mut self) -> CollectionResult<()>; /// Return temperature data. - fn get_temperature_data(&mut self) -> CollectionResult> { - Ok(vec![]) - } + fn get_temperature_data(&mut self) -> CollectionResult>; /// Return process data. - fn get_process_data(&mut self) -> CollectionResult> { - Ok(vec![]) - } + fn get_process_data(&mut self) -> CollectionResult>; + + /// Return disk data. + fn get_disk_data(&mut self) -> CollectionResult; } diff --git a/src/new_data_collection/collectors/linux.rs b/src/new_data_collection/collectors/linux.rs index 4a7191a4..8786762b 100644 --- a/src/new_data_collection/collectors/linux.rs +++ b/src/new_data_collection/collectors/linux.rs @@ -10,6 +10,7 @@ use crate::{ error::CollectionResult, sources::{ common::{ + disk::DiskHarvest, processes::ProcessHarvest, temperature::{TemperatureData, TemperatureType}, }, @@ -68,4 +69,8 @@ impl DataCollector for LinuxDataCollector { self.gpus_total_mem, ) } + + fn get_disk_data(&mut self) -> CollectionResult { + todo!() + } } diff --git a/src/new_data_collection/error.rs b/src/new_data_collection/error.rs index dc5dd7bd..7652ac6f 100644 --- a/src/new_data_collection/error.rs +++ b/src/new_data_collection/error.rs @@ -6,7 +6,7 @@ pub enum CollectionError { /// A general error to propagate back up. A wrapper around [`anyhow::Error`]. General(anyhow::Error), - /// The collection is unsupported. + /// Collection is unsupported. Unsupported, } diff --git a/src/new_data_collection/sources/common/disk.rs b/src/new_data_collection/sources/common/disk.rs new file mode 100644 index 00000000..63e1a6e2 --- /dev/null +++ b/src/new_data_collection/sources/common/disk.rs @@ -0,0 +1,127 @@ +use hashbrown::HashMap; + +use crate::app::filter::Filter; + +#[derive(Clone, Debug, Default)] +pub struct DiskEntry { + pub name: String, + pub mount_point: String, + + /// Windows also contains an additional volume name field. + #[cfg(target_os = "windows")] + pub volume_name: Option, + + // TODO: Maybe unify all these? + pub free_space: Option, + pub used_space: Option, + pub total_space: Option, +} + +#[derive(Clone, Debug)] +pub struct IoData { + /// How many bytes are read. + pub read_bytes: u64, + + /// How many bytes are written. + pub write_bytes: u64, +} + +pub struct DiskHarvest { + /// Disk entries. + pub entries: Vec, + /// I/O stats, mapped to device names. + pub device_io_stats: HashMap>, +} + +/// Whether to keep the current disk entry given the filters, disk name, and +/// disk mount. 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. +pub fn keep_disk_entry( + disk_name: &str, mount_point: &str, disk_filter: &Option, mount_filter: &Option, +) -> bool { + match (disk_filter, mount_filter) { + (Some(d), Some(m)) => match (d.ignore_matches(), m.ignore_matches()) { + (true, true) => !(d.has_match(disk_name) || m.has_match(mount_point)), + (true, false) => { + if m.has_match(mount_point) { + true + } else { + d.should_keep(disk_name) + } + } + (false, true) => { + if d.has_match(disk_name) { + true + } else { + m.should_keep(mount_point) + } + } + (false, false) => d.has_match(disk_name) || m.has_match(mount_point), + }, + (Some(d), None) => d.should_keep(disk_name), + (None, Some(m)) => m.should_keep(mount_point), + (None, None) => true, + } +} + +#[cfg(test)] +mod test { + use regex::Regex; + + use super::keep_disk_entry; + use crate::app::filter::Filter; + + fn run_filter(disk_filter: &Option, mount_filter: &Option) -> Vec { + let targets = [ + ("/dev/nvme0n1p1", "/boot"), + ("/dev/nvme0n1p2", "/"), + ("/dev/nvme0n1p3", "/home"), + ("/dev/sda1", "/mnt/test"), + ("/dev/sda2", "/mnt/boot"), + ]; + + targets + .into_iter() + .enumerate() + .filter_map(|(itx, (name, mount))| { + if keep_disk_entry(name, mount, disk_filter, mount_filter) { + Some(itx) + } else { + None + } + }) + .collect() + } + + #[test] + fn test_keeping_disk_entry() { + let disk_ignore = Some(Filter::new(true, vec![Regex::new("nvme").unwrap()])); + let disk_keep = Some(Filter::new(false, vec![Regex::new("nvme").unwrap()])); + let mount_ignore = Some(Filter::new(true, vec![Regex::new("boot").unwrap()])); + let mount_keep = Some(Filter::new(false, vec![Regex::new("boot").unwrap()])); + + assert_eq!(run_filter(&None, &None), vec![0, 1, 2, 3, 4]); + + assert_eq!(run_filter(&disk_ignore, &None), vec![3, 4]); + assert_eq!(run_filter(&disk_keep, &None), vec![0, 1, 2]); + + assert_eq!(run_filter(&None, &mount_ignore), vec![1, 2, 3]); + assert_eq!(run_filter(&None, &mount_keep), vec![0, 4]); + + assert_eq!(run_filter(&disk_ignore, &mount_ignore), vec![3]); + assert_eq!(run_filter(&disk_keep, &mount_ignore), vec![0, 1, 2, 3]); + + assert_eq!(run_filter(&disk_ignore, &mount_keep), vec![0, 3, 4]); + assert_eq!(run_filter(&disk_keep, &mount_keep), vec![0, 1, 2, 4]); + } +} diff --git a/src/new_data_collection/sources/common/mod.rs b/src/new_data_collection/sources/common/mod.rs index d918321f..9a68e285 100644 --- a/src/new_data_collection/sources/common/mod.rs +++ b/src/new_data_collection/sources/common/mod.rs @@ -1,2 +1,3 @@ +pub mod disk; pub mod processes; pub mod temperature; diff --git a/src/new_data_collection/sources/sysinfo/disk.rs b/src/new_data_collection/sources/sysinfo/disk.rs new file mode 100644 index 00000000..e69de29b