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

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

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

Review comment (suggestion):
nitpicking: the name `proc_mounts_path` is used for `CgroupBuilder` and the file path `/proc/mounts`, so `proc_mounts_path` would be preferable here as well.
```suggestion
        pub proc_mounts_path: &'a str,
```

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

        writeln_special(location, pid)?;

        Ok(())
    }
}