Add support for rsvd hugetlb cgroup

Signed-off-by: omprakaash <omsuseela@gmail.com>
Signed-off-by: om <omsuseela@gmail.com>
Signed-off-by: Om Prakaash <omsuseela@gmail.com>

Author: omprakaash

crates/libcgroups/src/v1/hugetlb.rs

@@ -1,10 +1,10 @@
-use std::{collections::HashMap, num::ParseIntError, path::Path};
-
 use crate::{
     common::{self, ControllerOpt, EitherError, MustBePowerOfTwo, WrappedIoError},
     stats::{supported_page_sizes, HugeTlbStats, StatsProvider, SupportedPageSizesError},
 };
+use std::{collections::HashMap, num::ParseIntError, path::Path};
 
+use crate::common::read_cgroup_file;
 use oci_spec::runtime::LinuxHugepageLimit;
 
 use super::controller::Controller;
@@ -109,6 +109,15 @@ impl HugeTlb {
             root_path.join(format!("hugetlb.{}.limit_in_bytes", hugetlb.page_size())),
             hugetlb.limit(),
         )?;
+
+        let rsvd_file_path = root_path.join(format!(
+            "hugetlb.{}.rsvd.limit_in_bytes",
+            hugetlb.page_size()
+        ));
+        if rsvd_file_path.exists() {
+            common::write_cgroup_file(rsvd_file_path, hugetlb.limit())?;
+        }
+
         Ok(())
     }
 
@@ -121,16 +130,20 @@ impl HugeTlb {
         page_size: &str,
     ) -> Result<HugeTlbStats, V1HugeTlbStatsError> {
         let mut stats = HugeTlbStats::default();
-
-        let usage_file = format!("hugetlb.{page_size}.usage_in_bytes");
-        let usage_content = common::read_cgroup_file(cgroup_path.join(usage_file))?;
+        let mut file_prefix = format!("hugetlb.{page_size}.rsvd");
+        let mut usage_file = format!("{file_prefix}.usage_in_bytes");
+        let usage_content = read_cgroup_file(cgroup_path.join(&usage_file)).or_else(|_| {
+            file_prefix = format!("hugetlb.{page_size}");
+            usage_file = format!("{file_prefix}.usage_in_bytes");
+            read_cgroup_file(cgroup_path.join(&usage_file))
+        })?;
         stats.usage = usage_content.trim().parse()?;
 
-        let max_file = format!("hugetlb.{page_size}.max_usage_in_bytes");
+        let max_file = format!("{file_prefix}.max_usage_in_bytes");
         let max_content = common::read_cgroup_file(cgroup_path.join(max_file))?;
         stats.max_usage = max_content.trim().parse()?;
 
-        let failcnt_file = format!("hugetlb.{page_size}.failcnt");
+        let failcnt_file = format!("{file_prefix}.failcnt");
         let failcnt_content = common::read_cgroup_file(cgroup_path.join(failcnt_file))?;
         stats.fail_count = failcnt_content.trim().parse()?;
 
@@ -163,6 +176,32 @@ mod tests {
         assert_eq!(hugetlb.limit().to_string(), content);
     }
 
+    #[test]
+    fn test_set_rsvd_hugetlb() {
+        let page_file_name = "hugetlb.2MB.limit_in_bytes";
+        let rsvd_page_file_name = "hugetlb.2MB.rsvd.limit_in_bytes";
+        let tmp = tempfile::tempdir().unwrap();
+        set_fixture(tmp.path(), page_file_name, "0").expect("Set fixture for 2 MB page size");
+        set_fixture(tmp.path(), rsvd_page_file_name, "0")
+            .expect("Set fixture for 2 MB rsvd page size");
+
+        let hugetlb = LinuxHugepageLimitBuilder::default()
+            .page_size("2MB")
+            .limit(16384)
+            .build()
+            .unwrap();
+
+        HugeTlb::apply(tmp.path(), &hugetlb).expect("apply hugetlb");
+        let content =
+            read_to_string(tmp.path().join(page_file_name)).expect("Read hugetlb file content");
+        let rsvd_content = read_to_string(tmp.path().join(rsvd_page_file_name))
+            .expect("Read rsvd hugetlb file content");
+
+        // Both files should have been written to
+        assert_eq!(hugetlb.limit().to_string(), content);
+        assert_eq!(hugetlb.limit().to_string(), rsvd_content);
+    }
+
     #[test]
     fn test_set_hugetlb_with_invalid_page_size() {
         let tmp = tempfile::tempdir().unwrap();
@@ -222,4 +261,30 @@ mod tests {
         };
         assert_eq!(actual, expected);
     }
+
+    #[test]
+    fn test_stat_rsvd_hugetlb() {
+        let tmp = tempfile::tempdir().unwrap();
+
+        set_fixture(tmp.path(), "hugetlb.2MB.rsvd.usage_in_bytes", "1024\n")
+            .expect("set hugetlb usage");
+        set_fixture(tmp.path(), "hugetlb.2MB.rsvd.max_usage_in_bytes", "4096\n")
+            .expect("set hugetlb max usage");
+        set_fixture(tmp.path(), "hugetlb.2MB.rsvd.failcnt", "5").expect("set hugetlb fail count");
+
+        set_fixture(tmp.path(), "hugetlb.2MB.usage_in_bytes", "2048\n").expect("set hugetlb usage");
+        set_fixture(tmp.path(), "hugetlb.2MB.max_usage_in_bytes", "8192\n")
+            .expect("set hugetlb max usage");
+        set_fixture(tmp.path(), "hugetlb.2MB.failcnt", "10").expect("set hugetlb fail count");
+
+        let actual = HugeTlb::stats_for_page_size(tmp.path(), "2MB").expect("get cgroup stats");
+
+        // Should prefer rsvd stats over non-rsvd stats
+        let expected = HugeTlbStats {
+            usage: 1024,
+            max_usage: 4096,
+            fail_count: 5,
+        };
+        assert_eq!(actual, expected);
+    }
 }

crates/libcgroups/src/v2/hugetlb.rs

@@ -13,6 +13,7 @@ use crate::{
     },
 };
 
+use crate::common::read_cgroup_file;
 use oci_spec::runtime::LinuxHugepageLimit;
 
 #[derive(thiserror::Error, Debug)]
@@ -104,6 +105,12 @@ impl HugeTlb {
             root_path.join(format!("hugetlb.{}.max", hugetlb.page_size())),
             hugetlb.limit(),
         )?;
+        
+        let rsvd_file_path = root_path.join(format!("hugetlb.{}.rsvd.max", hugetlb.page_size()));
+        if rsvd_file_path.exists() {
+            common::write_cgroup_file(rsvd_file_path, hugetlb.limit())?;
+        }
+
         Ok(())
     }
 
@@ -115,9 +122,14 @@ impl HugeTlb {
         cgroup_path: &Path,
         page_size: &str,
     ) -> Result<HugeTlbStats, V2HugeTlbStatsError> {
-        let events_file = format!("hugetlb.{page_size}.events");
-        let path = cgroup_path.join(events_file);
-        let events = common::read_cgroup_file(&path)?;
+        let mut file_prefix = format!("hugetlb.{page_size}.rsvd");
+        let mut path = cgroup_path.join(format!("{file_prefix}.events"));
+        let events = read_cgroup_file(&path).or_else(|_| {
+            file_prefix = format!("hugetlb.{page_size}");
+            path = cgroup_path.join(format!("{file_prefix}.events"));
+            read_cgroup_file(&path)
+        })?;
+
         let fail_count: u64 = events
             .lines()
             .find(|l| l.starts_with("max"))
@@ -130,7 +142,7 @@ impl HugeTlb {
             .unwrap_or_default();
 
         Ok(HugeTlbStats {
-            usage: parse_single_value(&cgroup_path.join(format!("hugetlb.{page_size}.current")))?,
+            usage: parse_single_value(&cgroup_path.join(format!("{file_prefix}.current")))?,
             fail_count,
             ..Default::default()
         })
@@ -178,6 +190,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_set_rsvd_hugetlb() {
+        let page_file_name = "hugetlb.2MB.max";
+        let rsvd_page_file_name = "hugetlb.2MB.rsvd.max";
+        let tmp = tempfile::tempdir().unwrap();
+        set_fixture(tmp.path(), page_file_name, "0").expect("Set fixture for 2 MB page size");
+        set_fixture(tmp.path(), rsvd_page_file_name, "0")
+            .expect("Set fixture for 2 MB rsvd page size");
+
+        let hugetlb = LinuxHugepageLimitBuilder::default()
+            .page_size("2MB")
+            .limit(16384)
+            .build()
+            .unwrap();
+        HugeTlb::apply(tmp.path(), &hugetlb).expect("apply hugetlb");
+
+        let content =
+            read_to_string(tmp.path().join(page_file_name)).expect("Read hugetlb file content");
+        let rsvd_content = read_to_string(tmp.path().join(rsvd_page_file_name))
+            .expect("Read hugetlb file content");
+
+        assert_eq!(hugetlb.limit().to_string(), content);
+        assert_eq!(hugetlb.limit().to_string(), rsvd_content);
+    }
+
     quickcheck! {
         fn property_test_set_hugetlb(hugetlb: LinuxHugepageLimit) -> bool {
             let page_file_name = format!("hugetlb.{:?}.max", hugetlb.page_size());
@@ -217,4 +254,25 @@ mod tests {
         };
         assert_eq!(actual, expected);
     }
+
+    #[test]
+    fn test_stat_rsvd_hugetbl() {
+        let tmp = tempfile::tempdir().unwrap();
+        set_fixture(tmp.path(), "hugetlb.2MB.current", "2048\n").expect("set hugetlb current");
+        set_fixture(tmp.path(), "hugetlb.2MB.events", "max 5\n").expect("set hugetlb events");
+        set_fixture(tmp.path(), "hugetlb.2MB.rsvd.current", "1024\n")
+            .expect("set hugetlb rsvd current");
+        set_fixture(tmp.path(), "hugetlb.2MB.rsvd.events", "max 5\n")
+            .expect("set hugetlb rsvd events");
+
+        let actual = HugeTlb::stats_for_page_size(tmp.path(), "2MB").expect("get cgroup stats");
+
+        // Should prefer rsvd stats over non-rsvd stats if available
+        let expected = HugeTlbStats {
+            usage: 1024,
+            max_usage: 0,
+            fail_count: 5,
+        };
+        assert_eq!(actual, expected);
+    }
 }

tests/contest/contest/src/tests/tlb/tlb_test.rs

@@ -11,6 +11,20 @@ fn check_hugetlb() -> bool {
     PathBuf::from("/sys/fs/cgroup/hugetlb").exists()
 }
 
+fn check_hugetlb_rsvd() -> bool {
+    let sizes = get_tlb_sizes();
+    for size in sizes.iter() {
+        let rsvd_path = format!(
+            "/sys/fs/cgroup/hugetlb/hugetlb.{}.rsvd.limit_in_bytes",
+            size
+        );
+        if !PathBuf::from(rsvd_path).exists() {
+            return false;
+        }
+    }
+    true
+}
+
 fn make_hugetlb_spec(page_size: &str, limit: i64) -> Spec {
     SpecBuilder::default()
         .linux(
@@ -110,6 +124,23 @@ fn validate_tlb(id: &str, size: &str, limit: i64) -> TestResult {
     }
 }
 
+fn validate_rsvd_tlb(id: &str, size: &str, limit: i64) -> TestResult {
+    let root = "/sys/fs/cgroup/hugetlb";
+    let path = format!("{root}/{id}/hugetlb.{size}.rsvd.limit_in_bytes");
+    let val_str = std::fs::read_to_string(path).unwrap();
+    let val: i64 = val_str.trim().parse().unwrap();
+    if val == limit {
+        TestResult::Passed
+    } else {
+        TestResult::Failed(anyhow!(
+            "page limit not set correctly : for size {}, expected {}, got {}",
+            size,
+            limit,
+            val
+        ))
+    }
+}
+
 fn test_valid_tlb() -> TestResult {
     // When setting the limit just for checking if writing works, the amount of memory
     // requested does not matter, as all insigned integers will be accepted.
@@ -134,6 +165,30 @@ fn test_valid_tlb() -> TestResult {
     TestResult::Passed
 }
 
+fn test_valid_rsvd_tlb() -> TestResult {
+    let limit: i64 = 1 << 30;
+    let tlb_sizes = get_tlb_sizes();
+    for size in tlb_sizes.iter() {
+        let spec = make_hugetlb_spec(size, limit);
+        let res = test_outside_container(spec, &|data| {
+            test_result!(check_container_created(&data));
+            // Currentle, we write the same value to both limit_in_bytes and rsvd.limit_in_bytes
+            let non_rsvd = validate_tlb(&data.id, size, limit);
+            let rsvd = validate_rsvd_tlb(&data.id, size, limit);
+            if matches!(non_rsvd, TestResult::Failed(_)) {
+                return non_rsvd;
+            } else if matches!(rsvd, TestResult::Failed(_)) {
+                return rsvd;
+            }
+            TestResult::Passed
+        });
+        if matches!(res, TestResult::Failed(_)) {
+            return res;
+        }
+    }
+    TestResult::Passed
+}
+
 pub fn get_tlb_test() -> TestGroup {
     let wrong_tlb = ConditionalTest::new(
         "invalid_tlb",
@@ -145,7 +200,16 @@ pub fn get_tlb_test() -> TestGroup {
         Box::new(check_hugetlb),
         Box::new(test_valid_tlb),
     );
+    let valid_rsvd_tlb = ConditionalTest::new(
+        "valid_rsvd_tlb",
+        Box::new(check_hugetlb_rsvd),
+        Box::new(test_valid_rsvd_tlb),
+    );
     let mut tg = TestGroup::new("huge_tlb");
-    tg.add(vec![Box::new(wrong_tlb), Box::new(valid_tlb)]);
+    tg.add(vec![
+        Box::new(wrong_tlb),
+        Box::new(valid_tlb),
+        Box::new(valid_rsvd_tlb),
+    ]);
     tg
 }