diff --git a/src/util/cpu.rs b/src/util/cpu.rs new file mode 100644 index 0000000..3da3de9 --- /dev/null +++ b/src/util/cpu.rs @@ -0,0 +1,73 @@ +use std::{fs::File, io::Read}; + +use anyhow::{Context, Error}; + +/// Parse a set of CPU ranges. They can be either a single number of a fully qualified range +/// which is separated to one another with a comma (`,`) and use a dash (`-`) to indicate the +/// start and (inclusive) end of the range. +fn _read_cpu_range(ranges: &str) -> Result, Error> { + let mut cpus = vec![]; + + for cpu_range in ranges.split(',') { + let rangeop_result = cpu_range.find('-'); + match rangeop_result { + None => cpus.push( + cpu_range + .trim_end() + .parse::() + .with_context(|| "Failed to parse lone CPU".to_string())?, + ), + Some(index) => { + let start = cpu_range[..index] + .trim_end() + .parse::() + .with_context(|| "Failed to parse starting CPU".to_string())?; + let end = cpu_range[index + 1..] + .trim_end() + .parse::() + .with_context(|| "Failed to parse ending CPU".to_string())?; + cpus.extend(start..end + 1); + } + } + } + + Ok(cpus) +} + +/// Parses `/sys/devices/system/cpu/online` and returns the online CPUs in the system. +pub fn get_online_cpus() -> Result, Error> { + let mut file = File::open("/sys/devices/system/cpu/online")?; + let mut ranges = String::new(); + file.read_to_string(&mut ranges)?; + + _read_cpu_range(&ranges) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cpu_ranges_to_list() { + let cpus = _read_cpu_range("0").unwrap(); + assert_eq!(cpus, vec![0]); + + let cpus = _read_cpu_range("0-7").unwrap(); + assert_eq!(cpus, (0..=7).collect::>()); + + let cpus = _read_cpu_range("0-7,16-23").unwrap(); + let expected = (0..=7).chain(16..=23).collect::>(); + assert_eq!(cpus, expected); + + let cpus = _read_cpu_range("0-1,3,7-9,48,49").unwrap(); + assert_eq!( + cpus, + (0..=1) + .chain(3..=3) + .chain(7..=9) + .chain(48..=48) + .chain(49..=49) + .collect::>() + ); + } +} diff --git a/src/util.rs b/src/util/lpm.rs similarity index 62% rename from src/util.rs rename to src/util/lpm.rs index 1a09212..ac1d017 100644 --- a/src/util.rs +++ b/src/util/lpm.rs @@ -1,7 +1,3 @@ -use std::{fs::File, io::Read}; - -use anyhow::{Context, Error}; - #[derive(Debug, PartialEq)] pub struct AddressBlockRange { pub addr: u64, @@ -11,16 +7,13 @@ pub struct AddressBlockRange { /// Calculate addresses for longest prefix match. /// /// For a given address range, calculate all the prefix ranges to ensure searching -/// with Longest Prefix Match returns the precise value we want. This is typically -/// used in networking to select the right subnet. +/// with Longest Prefix Match algorithm returns the precise value we want. This is +/// typically used in networking to select the right subnet but we use it to store +/// memory mappings. pub fn summarize_address_range(low: u64, high: u64) -> Vec { let mut res = Vec::new(); let mut curr = low; - if low > high { - panic!("first {:x} > last {:x}", low, high); - } - while curr <= high { let number_of_bits = std::cmp::min( curr.trailing_zeros(), @@ -39,48 +32,9 @@ pub fn summarize_address_range(low: u64, high: u64) -> Vec { res } -fn _read_cpu_range(path: &str) -> Result, Error> { - let mut cpus: Vec<_> = vec![]; - let mut fh = File::open(path)?; - let mut cpu_range_str = String::new(); - fh.read_to_string(&mut cpu_range_str)?; - - for cpu_range in cpu_range_str.split(',') { - let rangeop_result = cpu_range.find('-'); - match rangeop_result { - None => cpus.push( - cpu_range - .trim_end() - .parse::() - .with_context(|| "Failed to parse lone CPU".to_string())?, - ), - Some(index) => { - let start = cpu_range[..index] - .trim_end() - .parse::() - .with_context(|| "Failed to parse starting CPU".to_string())?; - let end = cpu_range[index + 1..] - .trim_end() - .parse::() - .with_context(|| "Failed to parse ending CPU".to_string())?; - cpus.extend(start..end + 1); - } - } - } - - Ok(cpus) -} - -pub fn get_online_cpus() -> Result, Error> { - let cpus: Vec = _read_cpu_range("/sys/devices/system/cpu/online")?; - - Ok(cpus) -} - #[cfg(test)] mod tests { use std::mem::size_of; - use tempdir::TempDir; use libbpf_rs::libbpf_sys; use libbpf_rs::MapCore; @@ -208,50 +162,4 @@ mod tests { assert_eq!(parsed.executable_id, mapping2.executable_id); } } - - #[test] - fn cpu_ranges_to_list() { - use std::io::Seek; - use std::io::Write; - - let tmp_dir = TempDir::new("cpu_devs").unwrap(); - let file_path = tmp_dir.path().join("online"); - let mut tmp_file = File::create(file_path.clone()).unwrap(); - let file_str = file_path.to_str().unwrap(); - - writeln!(tmp_file, "0").unwrap(); - let cpus = _read_cpu_range(file_str).unwrap(); - assert_eq!(cpus, vec![0]); - - tmp_file.rewind().unwrap(); - writeln!(tmp_file, "0-7").unwrap(); - - let cpus = _read_cpu_range(file_str).unwrap(); - assert_eq!(cpus, (0..=7).collect::>()); - - tmp_file.rewind().unwrap(); - writeln!(tmp_file, "0-7,16-23").unwrap(); - - let cpus = _read_cpu_range(file_str).unwrap(); - let expected = (0..=7).chain(16..=23).collect::>(); - - assert_eq!(cpus, expected); - - tmp_file.rewind().unwrap(); - writeln!(tmp_file, "0-1,3,7-9,48,49").unwrap(); - - let cpus = _read_cpu_range(file_str).unwrap(); - assert_eq!( - cpus, - (0..=1) - .chain(3..=3) - .chain(7..=9) - .chain(48..=48) - .chain(49..=49) - .collect::>() - ); - - drop(tmp_file); - tmp_dir.close().expect("tempdir should be destroyed"); - } } diff --git a/src/util/mod.rs b/src/util/mod.rs new file mode 100644 index 0000000..043e42a --- /dev/null +++ b/src/util/mod.rs @@ -0,0 +1,6 @@ +mod cpu; +mod lpm; + +pub use cpu::get_online_cpus; +pub use lpm::summarize_address_range; +pub use lpm::AddressBlockRange;