From 3657ca09676ebeb555d4c4cf36af0f2e52e8be6f Mon Sep 17 00:00:00 2001
From: utam0k <k0ma@utam0k.jp>
Date: Thu, 23 Sep 2021 21:13:39 +0900
Subject: [PATCH] add a unit test for applying cgroup in builder_impl().

---
 cgroups/src/common.rs         |  5 ++++
 cgroups/src/lib.rs            |  3 +-
 cgroups/src/test_manager.rs   | 56 +++++++++++++++++++++++++++++++++++
 src/container/builder_impl.rs | 53 ++++++++++++++++++++++-----------
 4 files changed, 98 insertions(+), 19 deletions(-)
 create mode 100644 cgroups/src/test_manager.rs

diff --git a/cgroups/src/common.rs b/cgroups/src/common.rs
index 5e67bb445..334576ad0 100644
--- a/cgroups/src/common.rs
+++ b/cgroups/src/common.rs
@@ -29,14 +29,19 @@ pub const DEFAULT_CGROUP_ROOT: &str = "/sys/fs/cgroup";
 pub trait CgroupManager {
     /// Adds a task specified by its pid to the cgroup
     fn add_task(&self, pid: Pid) -> Result<()>;
+
     /// Applies resource restrictions to the cgroup
     fn apply(&self, controller_opt: &ControllerOpt) -> Result<()>;
+
     /// Removes the cgroup
     fn remove(&self) -> Result<()>;
+
     // Sets the freezer cgroup to the specified state
     fn freeze(&self, state: FreezerState) -> Result<()>;
+
     /// Retrieve statistics for the cgroup
     fn stats(&self) -> Result<Stats>;
+
     // Gets the PIDs inside the cgroup
     fn get_all_pids(&self) -> Result<Vec<Pid>>;
 }
diff --git a/cgroups/src/lib.rs b/cgroups/src/lib.rs
index 35eca40ea..fe0778ca4 100644
--- a/cgroups/src/lib.rs
+++ b/cgroups/src/lib.rs
@@ -4,9 +4,10 @@
 #[cfg(test)]
 #[macro_use]
 extern crate quickcheck;
+mod test;
 
 pub mod common;
 pub mod stats;
-mod test;
+pub mod test_manager;
 pub mod v1;
 pub mod v2;
diff --git a/cgroups/src/test_manager.rs b/cgroups/src/test_manager.rs
new file mode 100644
index 000000000..2631045bb
--- /dev/null
+++ b/cgroups/src/test_manager.rs
@@ -0,0 +1,56 @@
+use std::cell::RefCell;
+
+use anyhow::Result;
+use nix::unistd::Pid;
+
+use crate::{
+    common::{CgroupManager, ControllerOpt, FreezerState},
+    stats::Stats,
+};
+
+#[derive(Debug)]
+pub struct TestManager {
+    add_task_args: RefCell<Vec<Pid>>,
+}
+
+impl Default for TestManager {
+    fn default() -> Self {
+        Self {
+            add_task_args: RefCell::new(vec![]),
+        }
+    }
+}
+
+impl CgroupManager for TestManager {
+    fn add_task(&self, pid: Pid) -> Result<()> {
+        self.add_task_args.borrow_mut().push(pid);
+        Ok(())
+    }
+
+    // NOTE: The argument cannot be stored due to lifetime.
+    fn apply(&self, _controller_opt: &ControllerOpt) -> Result<()> {
+        Ok(())
+    }
+
+    fn remove(&self) -> Result<()> {
+        unimplemented!()
+    }
+
+    fn freeze(&self, _state: FreezerState) -> Result<()> {
+        unimplemented!()
+    }
+
+    fn stats(&self) -> anyhow::Result<Stats> {
+        unimplemented!()
+    }
+
+    fn get_all_pids(&self) -> Result<Vec<Pid>> {
+        unimplemented!()
+    }
+}
+
+impl TestManager {
+    pub fn get_add_task_args(&self) -> Vec<Pid> {
+        self.add_task_args.borrow_mut().clone()
+    }
+}
diff --git a/src/container/builder_impl.rs b/src/container/builder_impl.rs
index 84f09fa19..a685b1e49 100644
--- a/src/container/builder_impl.rs
+++ b/src/container/builder_impl.rs
@@ -7,9 +7,9 @@ use crate::{
     utils,
 };
 use anyhow::{Context, Result};
-use cgroups;
+use cgroups::{self, common::CgroupManager};
 use nix::unistd::Pid;
-use oci_spec::runtime::Spec;
+use oci_spec::runtime::{LinuxResources, Spec};
 use std::{fs, io::Write, os::unix::prelude::RawFd, path::PathBuf};
 
 use super::{Container, ContainerStatus};
@@ -44,9 +44,7 @@ pub(super) struct ContainerBuilderImpl<'a> {
 
 impl<'a> ContainerBuilderImpl<'a> {
     pub(super) fn create(&mut self) -> Result<()> {
-        self.run_container()?;
-
-        Ok(())
+        self.run_container()
     }
 
     fn run_container(&mut self) -> Result<()> {
@@ -151,19 +149,7 @@ impl<'a> ContainerBuilderImpl<'a> {
 
         if self.rootless.is_none() && linux.resources.is_some() && self.init {
             if let Some(resources) = linux.resources.as_ref() {
-                let controller_opt = cgroups::common::ControllerOpt {
-                    resources,
-                    freezer_state: None,
-                    oom_score_adj: None,
-                    disable_oom_killer: false,
-                };
-                cmanager
-                    .add_task(init_pid)
-                    .context("Failed to add tasks to cgroup manager")?;
-
-                cmanager
-                    .apply(&controller_opt)
-                    .context("Failed to apply resource limits through cgroup")?;
+                apply_cgroups(resources, init_pid, cmanager.as_ref())?;
             }
         }
 
@@ -202,8 +188,29 @@ fn setup_mapping(rootless: &Rootless, pid: Pid) -> Result<()> {
     Ok(())
 }
 
+fn apply_cgroups<C: CgroupManager + ?Sized>(
+    resources: &LinuxResources,
+    pid: Pid,
+    cmanager: &C,
+) -> Result<()> {
+    let controller_opt = cgroups::common::ControllerOpt {
+        resources,
+        freezer_state: None,
+        oom_score_adj: None,
+        disable_oom_killer: false,
+    };
+    cmanager
+        .add_task(pid)
+        .context("Failed to add tasks to cgroup manager")?;
+    cmanager
+        .apply(&controller_opt)
+        .context("Failed to apply resource limits through cgroup")?;
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
+    use cgroups::test_manager::TestManager;
     use nix::{
         sched::{unshare, CloneFlags},
         unistd::{self, getgid, getuid},
@@ -309,4 +316,14 @@ mod tests {
         }
         Ok(())
     }
+
+    #[test]
+    fn apply_cgroup_successed() -> Result<()> {
+        let cmanager = TestManager::default();
+        let sample_pid = Pid::from_raw(1000);
+        let resources = LinuxResources::default();
+        apply_cgroups(&resources, sample_pid, &cmanager)?;
+        assert_eq!(cmanager.get_add_task_args(), vec![sample_pid]);
+        Ok(())
+    }
 }