From 9d7046732733f4d098557b55048af5df49a9d7ca Mon Sep 17 00:00:00 2001 From: bin liu Date: Tue, 25 Aug 2020 10:51:49 +0800 Subject: [PATCH] add hugetlb functions Signed-off-by: bin liu --- Cargo.toml | 1 + src/error.rs | 3 + src/hugetlb.rs | 144 +++++++++++++++++++++++++++++++++++++++++++++-- tests/hugetlb.rs | 34 +++++++++++ 4 files changed, 177 insertions(+), 5 deletions(-) create mode 100644 tests/hugetlb.rs diff --git a/Cargo.toml b/Cargo.toml index 524e006..1556a40 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ edition = "2018" [dependencies] log = "0.4" +regex = "1.1" [dev-dependencies] nix = "0.11.0" diff --git a/src/error.rs b/src/error.rs index 70eb4ec..f14a15c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -27,6 +27,8 @@ pub enum ErrorKind { /// This crate checks against this and operations will fail with this error. InvalidPath, + InvalidBytesSize, + /// An unknown error has occured. Other, } @@ -45,6 +47,7 @@ impl fmt::Display for Error { ErrorKind::ParseError => "unable to parse control group file", ErrorKind::InvalidOperation => "the requested operation is invalid", ErrorKind::InvalidPath => "the given path is invalid", + ErrorKind::InvalidBytesSize => "invalid bytes size", ErrorKind::Other => "an unknown error", }; diff --git a/src/hugetlb.rs b/src/hugetlb.rs index 5c3b1b7..3a1560f 100644 --- a/src/hugetlb.rs +++ b/src/hugetlb.rs @@ -20,8 +20,9 @@ use crate::{ /// the control group. #[derive(Debug, Clone)] pub struct HugeTlbController { - base: PathBuf, - path: PathBuf, + base: PathBuf, + path: PathBuf, + sizes: Vec, } impl ControllerInternal for HugeTlbController { @@ -87,16 +88,26 @@ impl HugeTlbController { pub fn new(oroot: PathBuf) -> Self { let mut root = oroot; root.push(Self::controller_type().to_string()); + let sizes = get_hugepage_sizes().unwrap(); Self { base: root.clone(), path: root, + sizes: sizes, } } /// Whether the system supports `hugetlb_size` hugepages. - pub fn size_supported(&self, _hugetlb_size: &str) -> bool { - // TODO - true + pub fn size_supported(&self, hugetlb_size: &str) -> bool { + for s in &self.sizes { + if s == hugetlb_size { + return true + } + } + false + } + + pub fn get_sizes(&self) -> Vec { + self.sizes.clone() } /// Check how many times has the limit of `hugetlb_size` hugepages been hit. @@ -138,3 +149,126 @@ impl HugeTlbController { }) } } + + +pub const HUGEPAGESIZE_DIR: &'static str = "/sys/kernel/mm/hugepages"; +use regex::Regex; +use std::collections::HashMap; +use std::fs; + +fn get_hugepage_sizes() -> Result> { + let mut m = Vec::new(); + let dirs = fs::read_dir(HUGEPAGESIZE_DIR); + if dirs.is_err() { + return Ok(m); + } + + for e in dirs.unwrap() { + let entry = e.unwrap(); + let name = entry.file_name().into_string().unwrap(); + let parts: Vec<&str> = name.split('-').collect(); + if parts.len() != 2 { + continue; + } + let bmap= get_binary_size_map(); + let size = parse_size(parts[1], &bmap)?; + let dabbrs = get_decimal_abbrs(); + m.push(custom_size(size as f64, 1024.0, &dabbrs)); + } + + Ok(m) +} + + +pub const KB: u128 = 1000; +pub const MB: u128 = 1000 * KB; +pub const GB: u128 = 1000 * MB; +pub const TB: u128 = 1000 * GB; +pub const PB: u128 = 1000 * TB; + +pub const KiB: u128 = 1024; +pub const MiB: u128 = 1024 * KiB; +pub const GiB: u128 = 1024 * MiB; +pub const TiB: u128 = 1024 * GiB; +pub const PiB: u128 = 1024 * TiB; + + +pub fn get_binary_size_map() -> HashMap { + let mut m = HashMap::new(); + m.insert("k".to_string(), KiB); + m.insert("m".to_string(), MiB); + m.insert("g".to_string(), GiB); + m.insert("t".to_string(), TiB); + m.insert("p".to_string(), PiB); + m +} + +pub fn get_decimal_size_map() -> HashMap { + let mut m = HashMap::new(); + m.insert("k".to_string(), KB); + m.insert("m".to_string(), MB); + m.insert("g".to_string(), GB); + m.insert("t".to_string(), TB); + m.insert("p".to_string(), PB); + m +} + +pub fn get_decimal_abbrs() -> Vec { + let m = vec![ + "B".to_string(), + "KB".to_string(), + "MB".to_string(), + "GB".to_string(), + "TB".to_string(), + "PB".to_string(), + "EB".to_string(), + "ZB".to_string(), + "YB".to_string(), + ]; + m +} + +fn parse_size(s: &str, m: &HashMap) -> Result { + let re = Regex::new(r"(?P\d+)(?P[kKmMgGtTpP]?)[bB]?$"); + + if re.is_err() { + return Err(Error::new(InvalidBytesSize)); + } + let caps = re.unwrap().captures(s).unwrap(); + + let num = caps.name("num"); + let size: u128 = if num.is_some() { + let n = num.unwrap().as_str().trim().parse::(); + if n.is_err(){ + return Err(Error::new(InvalidBytesSize)); + } + n.unwrap() + } else { + return Err(Error::new(InvalidBytesSize)); + }; + + let q = caps.name("mul"); + let mul: u128 = if q.is_some() { + let t = m.get(q.unwrap().as_str()); + if t.is_some() { + *t.unwrap() + } else { + return Err(Error::new(InvalidBytesSize)); + } + } else { + return Err(Error::new(InvalidBytesSize)); + }; + + Ok(size * mul) +} + +fn custom_size(mut size: f64, base: f64, m: &Vec) -> String { + let mut i = 0; + while size >= base && i < m.len() - 1 { + size /= base; + i += 1; + } + + format!("{}{}", size, m[i].as_str()) +} + diff --git a/tests/hugetlb.rs b/tests/hugetlb.rs new file mode 100644 index 0000000..8ceeecf --- /dev/null +++ b/tests/hugetlb.rs @@ -0,0 +1,34 @@ +//! Integration tests about the hugetlb subsystem +use cgroups::hugetlb::{HugeTlbController}; +use cgroups::Controller; +use cgroups::Cgroup; + +use cgroups::error::*; +use cgroups::error::ErrorKind::*; + +#[test] +fn test_hugetlb_sizes() { + let hier = cgroups::hierarchies::V1::new(); + let cg = Cgroup::new(&hier, String::from("test_hugetlb_sizes")); + { + let hugetlb_controller: &HugeTlbController = cg.controller_of().unwrap(); + let sizes = hugetlb_controller.get_sizes(); + + let size = "2MB"; + assert_eq!(sizes, vec![size.to_string()]); + + let supported = hugetlb_controller.size_supported(size); + assert_eq!(supported, true); + + assert_no_error( hugetlb_controller.failcnt(size)); + assert_no_error( hugetlb_controller.limit_in_bytes(size)); + assert_no_error( hugetlb_controller.usage_in_bytes(size)); + assert_no_error( hugetlb_controller.max_usage_in_bytes(size)); + } + cg.delete(); +} + + +fn assert_no_error(r: Result ) { + assert_eq!(!r.is_err() , true) +} \ No newline at end of file