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

Code before:
    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(())
        }

        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();

            // create a mock /proc/mounts file in a temporary directory
            fs::create_dir_all(mock_proc_dir)?;
            let file = OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .truncate(true)
                .open(mock_proc_mount)?;
            Ok(MockCgroupFs {
                mounts_file: file,
                proc_mount_path: mock_proc_mount,
            })
        }

        // Populate the mocked proc/mounts file with cgroupv2 entries
        // Also create a directory structure that simulates cgroupsv2 layout
        pub fn add_v2_mounts(&mut self) -> std::result::Result<(), std::io::Error> {
            let mock_sys_cgroups_dir = format!(
                "{}/{}",
                Path::new(self.proc_mount_path)
                    .parent()
                    .unwrap()

Review comment (suggestion):
nitpicking: See https://github.com/firecracker-microvm/firecracker/pull/4438#discussion_r1523604943.
```suggestion
            mock_proc_mounts: &'a str,
```

Code after:
            .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");

        writeln_special(location, pid)?;

        Ok(())
    }
}

pub fn setup_cgroup_conf(conf: &HashMap<String, impl Cgroup>) -> Result<(), JailerError> {
    // cgroups are iterated two times as some cgroups may require others (e.g cpuset requires
    // cpuset.mems and cpuset.cpus) to be set before attaching any pid.
    for cgroup in conf.values() {
        cgroup.write_values()?;
    }
    for cgroup in conf.values() {
        cgroup.attach_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;

    #[derive(Debug)]