Repository: firecracker-microvm/firecracker
File: src/jailer/src/cgroup.rs (Rust)

Code before:

    fn attach_pid(&self) -> Result<(), JailerError> {
        let pid = process::id();
        let location = &self.0.location.join("cgroup.procs");

        writeln_special(location, pid)?;

        Ok(())
    }
}

#[cfg(test)]
pub mod test_util {
    use std::fmt::Debug;
    use std::fs::{self, File, OpenOptions};
    use std::io::Write;
    use std::path::{Path, PathBuf};

    use utils::rand;

    pub fn get_mock_proc_mounts() -> String {
        format!(
            "/tmp/firecracker/test/{}/jailer/proc/mounts",
            rand::rand_alphanumerics(4).into_string().unwrap()
        )
    }

    #[derive(Debug)]
    pub struct MockCgroupFs<'a> {
        mounts_file: File,
        pub proc_mount_path: &'a str,
    }

    // Helper object that simulates the layout of the cgroup file system
    // This can be used for testing regardless of the availability of a particular
    // version of cgroups on the system
    impl<'a> MockCgroupFs<'a> {
        pub fn create_file_with_contents<P: AsRef<Path> + Debug>(
            filename: P,
            contents: &str,
        ) -> std::result::Result<(), std::io::Error> {
            let mut file = OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .truncate(true)
                .open(&filename)?;

            writeln!(file, "{}", contents)?;
            Ok(())
        }

Review comment (bug):
`get_mock_proc_mounts()` is always called just before `MockCgroupFs::new()`. How about putting `get_mock_proc_monts()` inside the `MockCgroupFs::new()`? By doing this and owning the string, the lifetime parameter `<'a>` of `MockCgroupFs` can be removed as follows:
```
-    pub fn get_mock_proc_mounts() -> String {
-        format!(
-            "/tmp/firecracker/test/{}/jailer/proc/mounts",
-            rand::rand_alphanumerics(4).into_string().unwrap()
-        )
-    }
-
     #[derive(Debug)]
-    pub struct MockCgroupFs<'a> {
+    pub struct MockCgroupFs {
         mounts_file: File,
-        pub proc_mount_path: &'a str,
+        pub proc_mount_path: String,
     }
 
     // Helper object that simulates the layout of the cgroup file system
     // This can be used for testing regardless of the availability of a particular
     // version of cgroups on the system
-    impl<'a> MockCgroupFs<'a> {
+    impl MockCgroupFs {
         pub fn create_file_with_contents<P: AsRef<Path> + Debug>(
             filename: P,
             contents: &str,
@@ -480,10 +473,12 @@ pub mod test_util {
             Ok(())
         }
 
-        pub fn new(
-            mock_proc_mount: &'a str,
-        ) -> std::result::Result<MockCgroupFs<'a>, std::io::Error> {
-            let mock_proc_dir = Path::new(mock_proc_mount).parent().unwrap();
+        pub fn new() -> std::result::Result<MockCgroupFs, std::io::Error> {
+            let mock_proc_mount = format!(
+                "/tmp/firecracker/test/{}/jailer/proc/mounts",
+                rand::rand_alphanumerics(4).into_string().unwrap()
+            );
+            let mock_proc_dir = Path::new(&mock_proc_mount).parent().unwrap();
```

Code after:
        Ok(CgroupV2 {
            base: CgroupBase {
                properties: Vec::new(),
                location: path,
            },
            available_controllers: Self::detect_available_controllers(unified_path),
        })
    }
}

impl Cgroup for CgroupV2 {
    fn add_property(&mut self, file: String, value: String) -> Result<(), JailerError> {
        let controller = get_controller_from_filename(&file)?;
        if self.available_controllers.contains(controller) {
            self.base.properties.push(CgroupProperty { file, value });
            Ok(())
        } else {
            Err(JailerError::CgroupControllerUnavailable(
                controller.to_string(),
            ))
        }
    }

    fn write_values(&self) -> Result<(), JailerError> {
        let mut enabled_controllers: HashSet<&str> = HashSet::new();

        // Create the cgroup directory for the controller.
        fs::create_dir_all(&self.base.location)
            .map_err(|err| JailerError::CreateDir(self.base.location.clone(), err))?;

        // Ok to unwrap since the path was just created.
        let parent = self.base.location.parent().unwrap();

        for property in self.base.properties.iter() {
            let controller = get_controller_from_filename(&property.file)?;
            // enable controllers only once
            if !enabled_controllers.contains(controller) {
                // Enable the controller in all parent directories
                CgroupV2::write_all_subtree_control(parent, controller)?;
                enabled_controllers.insert(controller);
            }
            writeln_special(&self.base.location.join(&property.file), &property.value)?;
        }

        Ok(())
    }

    fn attach_pid(&self) -> Result<(), JailerError> {
        let pid = process::id();
        let location = &self.base.location.join("cgroup.procs");