Seperate adding tasks to cgroups and applying resource restrictions

src/cgroups/common.rs

@@ -19,7 +19,11 @@ pub const CGROUP_PROCS: &str = "cgroup.procs";
 pub const DEFAULT_CGROUP_ROOT: &str = "/sys/fs/cgroup";
 
 pub trait CgroupManager {
-    fn apply(&self, linux_resources: &LinuxResources, pid: Pid) -> Result<()>;
+    /// 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, linux_resources: &LinuxResources) -> Result<()>;
+    /// Removes the cgroup
     fn remove(&self) -> Result<()>;
 }
 

src/cgroups/v1/blkio.rs

@@ -1,12 +1,7 @@
-use std::{
-    fs::{self},
-    path::Path,
-};
-
-use crate::cgroups::{
-    common::{self, CGROUP_PROCS},
-    v1::Controller,
-};
+use std::path::Path;
+
+use crate::cgroups::{common, v1::Controller};
+use anyhow::Result;
 use oci_spec::{LinuxBlockIo, LinuxResources};
 
 const CGROUP_BLKIO_THROTTLE_READ_BPS: &str = "blkio.throttle.read_bps_device";
@@ -17,25 +12,19 @@ const CGROUP_BLKIO_THROTTLE_WRITE_IOPS: &str = "blkio.throttle.write_iops_device
 pub struct Blkio {}
 
 impl Controller for Blkio {
-    fn apply(
-        linux_resources: &LinuxResources,
-        cgroup_root: &Path,
-        pid: nix::unistd::Pid,
-    ) -> anyhow::Result<()> {
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path) -> Result<()> {
         log::debug!("Apply blkio cgroup config");
-        fs::create_dir_all(cgroup_root)?;
 
         if let Some(blkio) = &linux_resources.block_io {
             Self::apply(cgroup_root, blkio)?;
         }
 
-        common::write_cgroup_file(cgroup_root.join(CGROUP_PROCS), pid)?;
         Ok(())
     }
 }
 
 impl Blkio {
-    fn apply(root_path: &Path, blkio: &LinuxBlockIo) -> anyhow::Result<()> {
+    fn apply(root_path: &Path, blkio: &LinuxBlockIo) -> Result<()> {
         for trbd in &blkio.blkio_throttle_read_bps_device {
             common::write_cgroup_file_str(
                 &root_path.join(CGROUP_BLKIO_THROTTLE_READ_BPS),
@@ -70,6 +59,8 @@ impl Blkio {
 
 #[cfg(test)]
 mod tests {
+    use std::fs;
+
     use super::*;
     use crate::cgroups::test::setup;
     use oci_spec::{LinuxBlockIo, LinuxThrottleDevice};

src/cgroups/v1/controller.rs

@@ -1,10 +1,18 @@
-use std::path::Path;
+use std::{fs, path::Path};
 
 use anyhow::Result;
 use nix::unistd::Pid;
 
 use oci_spec::LinuxResources;
 
+use crate::cgroups::common::{self, CGROUP_PROCS};
+
 pub trait Controller {
-    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path, pid: Pid) -> Result<()>;
+    fn add_task(pid: Pid, cgroup_path: &Path) -> Result<()> {
+        fs::create_dir_all(cgroup_path)?;
+        common::write_cgroup_file(cgroup_path.join(CGROUP_PROCS), pid)?;
+        Ok(())
+    }
+
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path) -> Result<()>;
 }

src/cgroups/v1/controller_type.rs

@@ -1,5 +1,6 @@
 use std::fmt::Display;
 
+#[derive(Hash, PartialEq, Eq, Debug, Clone)]
 pub enum ControllerType {
     Cpu,
     CpuAcct,

src/cgroups/v1/cpu.rs

@@ -1,10 +1,9 @@
-use std::{fs, path::Path};
+use std::path::Path;
 
 use anyhow::Result;
-use nix::unistd::Pid;
 use oci_spec::{LinuxCpu, LinuxResources};
 
-use crate::cgroups::common::{self, CGROUP_PROCS};
+use crate::cgroups::common;
 
 use super::Controller;
 
@@ -17,14 +16,13 @@ const CGROUP_CPU_RT_PERIOD: &str = "cpu.rt_period_us";
 pub struct Cpu {}
 
 impl Controller for Cpu {
-    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path, pid: Pid) -> Result<()> {
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path) -> Result<()> {
         log::debug!("Apply Cpu cgroup config");
-        fs::create_dir_all(cgroup_root)?;
+
         if let Some(cpu) = &linux_resources.cpu {
             Self::apply(cgroup_root, cpu)?;
         }
 
-        common::write_cgroup_file(cgroup_root.join(CGROUP_PROCS), pid)?;
         Ok(())
     }
 }

src/cgroups/v1/cpuacct.rs

@@ -1,37 +1,33 @@
-use std::{fs, path::Path};
+use std::path::Path;
 
 use anyhow::Result;
-use nix::unistd::Pid;
 use oci_spec::LinuxResources;
 
-use crate::cgroups::common::{self, CGROUP_PROCS};
-
 use super::Controller;
 
 pub struct CpuAcct {}
 
 impl Controller for CpuAcct {
-    fn apply(_linux_resources: &LinuxResources, cgroup_path: &Path, pid: Pid) -> Result<()> {
-        log::debug!("Apply cpuacct cgroup config");
-        fs::create_dir_all(cgroup_path)?;
-
-        common::write_cgroup_file(cgroup_path.join(CGROUP_PROCS), pid)?;
+    fn apply(_linux_resources: &LinuxResources, _cgroup_path: &Path) -> Result<()> {
         Ok(())
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::fs;
+
+    use nix::unistd::Pid;
+
     use super::*;
-    use crate::cgroups::test::setup;
+    use crate::cgroups::{common::CGROUP_PROCS, test::setup};
 
     #[test]
-    fn test_apply() {
+    fn test_add_task() {
         let (tmp, procs) = setup("test_cpuacct_apply", CGROUP_PROCS);
-        let resource = LinuxResources::default();
         let pid = Pid::from_raw(1000);
 
-        CpuAcct::apply(&resource, &tmp, pid).expect("apply cpuacct");
+        CpuAcct::add_task(pid, &tmp).expect("apply cpuacct");
 
         let content = fs::read_to_string(&procs)
             .unwrap_or_else(|_| panic!("read {} file content", CGROUP_PROCS));

src/cgroups/v1/cpuset.rs

@@ -1,8 +1,9 @@
 use std::{fs, path::Path};
 
 use anyhow::{bail, Result};
-use nix::unistd::Pid;
+use nix::unistd;
 use oci_spec::{LinuxCpu, LinuxResources};
+use unistd::Pid;
 
 use crate::cgroups::common::{self, CGROUP_PROCS};
 
@@ -14,18 +15,23 @@ const CGROUP_CPUSET_MEMS: &str = "cpuset.mems";
 pub struct CpuSet {}
 
 impl Controller for CpuSet {
-    fn apply(linux_resources: &LinuxResources, cgroup_path: &Path, pid: Pid) -> Result<()> {
-        log::debug!("Apply CpuSet cgroup config");
+    fn add_task(pid: Pid, cgroup_path: &Path) -> Result<()> {
         fs::create_dir_all(cgroup_path)?;
 
         Self::ensure_not_empty(cgroup_path, CGROUP_CPUSET_CPUS)?;
         Self::ensure_not_empty(cgroup_path, CGROUP_CPUSET_MEMS)?;
 
+        common::write_cgroup_file(cgroup_path.join(CGROUP_PROCS), pid)?;
+        Ok(())
+    }
+
+    fn apply(linux_resources: &LinuxResources, cgroup_path: &Path) -> Result<()> {
+        log::debug!("Apply CpuSet cgroup config");
+
         if let Some(cpuset) = &linux_resources.cpu {
             Self::apply(cgroup_path, cpuset)?;
         }
 
-        common::write_cgroup_file(cgroup_path.join(CGROUP_PROCS), pid)?;
         Ok(())
     }
 }

src/cgroups/v1/devices.rs

@@ -1,18 +1,16 @@
-use std::{fs::create_dir_all, path::Path};
+use std::path::Path;
 
 use anyhow::Result;
-use nix::unistd::Pid;
 
-use crate::cgroups::common::{self, CGROUP_PROCS};
+use crate::cgroups::common;
 use crate::{cgroups::v1::Controller, rootfs::default_devices};
 use oci_spec::{LinuxDeviceCgroup, LinuxDeviceType, LinuxResources};
 
 pub struct Devices {}
 
 impl Controller for Devices {
-    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path, pid: Pid) -> Result<()> {
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path) -> Result<()> {
         log::debug!("Apply Devices cgroup config");
-        create_dir_all(&cgroup_root)?;
 
         for d in &linux_resources.devices {
             Self::apply_device(d, cgroup_root)?;
@@ -27,7 +25,6 @@ impl Controller for Devices {
             Self::apply_device(&d, &cgroup_root)?;
         }
 
-        common::write_cgroup_file(cgroup_root.join(CGROUP_PROCS), pid)?;
         Ok(())
     }
 }

src/cgroups/v1/freezer.rs

@@ -6,9 +6,8 @@ use std::{
 };
 
 use anyhow::{Result, *};
-use nix::unistd::Pid;
 
-use crate::cgroups::common::{self, CGROUP_PROCS};
+use crate::cgroups::common;
 use crate::cgroups::v1::Controller;
 use oci_spec::{FreezerState, LinuxResources};
 
@@ -20,15 +19,14 @@ const FREEZER_STATE_FREEZING: &str = "FREEZING";
 pub struct Freezer {}
 
 impl Controller for Freezer {
-    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path, pid: Pid) -> Result<()> {
+    fn apply(linux_resources: &LinuxResources, cgroup_root: &Path) -> Result<()> {
         log::debug!("Apply Freezer cgroup config");
         create_dir_all(&cgroup_root)?;
 
         if let Some(freezer_state) = linux_resources.freezer {
             Self::apply(freezer_state, cgroup_root)?;
         }
 
-        common::write_cgroup_file(cgroup_root.join(CGROUP_PROCS), pid)?;
         Ok(())
     }
 }
@@ -116,8 +114,10 @@ impl Freezer {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::cgroups::common::CGROUP_PROCS;
     use crate::cgroups::test::set_fixture;
     use crate::utils::create_temp_dir;
+    use nix::unistd::Pid;
     use oci_spec::FreezerState;
 
     #[test]
@@ -160,10 +160,9 @@ mod tests {
     }
 
     #[test]
-    fn test_apply() {
-        let tmp =
-            create_temp_dir("test_set_freezer_state").expect("create temp directory for test");
-        set_fixture(&tmp, CGROUP_FREEZER_STATE, "").expect("Set fixure for freezer state");
+    fn test_add_and_apply() {
+        let tmp = create_temp_dir("test_add_task").expect("create temp directory for test");
+        set_fixture(&tmp, CGROUP_FREEZER_STATE, "").expect("set fixure for freezer state");
         set_fixture(&tmp, CGROUP_PROCS, "").expect("set fixture for proc file");
 
         // set Thawed state.
@@ -182,13 +181,13 @@ mod tests {
             };
 
             let pid = Pid::from_raw(1000);
-            let _ =
-                <Freezer as Controller>::apply(&linux_resources, &tmp, pid).expect("freezer apply");
+            Freezer::add_task(pid, &tmp).expect("freezer add task");
+            <Freezer as Controller>::apply(&linux_resources, &tmp).expect("freezer apply");
             let state_content =
-                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("read to string");
             assert_eq!(FREEZER_STATE_THAWED, state_content);
             let pid_content =
-                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("read to string");
             assert_eq!(pid_content, "1000");
         }
 
@@ -208,13 +207,13 @@ mod tests {
             };
 
             let pid = Pid::from_raw(1001);
-            let _ =
-                <Freezer as Controller>::apply(&linux_resources, &tmp, pid).expect("freezer apply");
+            Freezer::add_task(pid, &tmp).expect("freezer add task");
+            <Freezer as Controller>::apply(&linux_resources, &tmp).expect("freezer apply");
             let state_content =
-                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("read to string");
             assert_eq!(FREEZER_STATE_FROZEN, state_content);
             let pid_content =
-                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("read to string");
             assert_eq!(pid_content, "1001");
         }
 
@@ -233,16 +232,16 @@ mod tests {
                 freezer: Some(FreezerState::Undefined),
             };
 
-            let old_state_content =
-                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("Read to string");
             let pid = Pid::from_raw(1002);
-            let _ =
-                <Freezer as Controller>::apply(&linux_resources, &tmp, pid).expect("freezer apply");
+            let old_state_content =
+                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("read to string");
+            Freezer::add_task(pid, &tmp).expect("freezer add task");
+            <Freezer as Controller>::apply(&linux_resources, &tmp).expect("freezer apply");
             let state_content =
-                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_FREEZER_STATE)).expect("read to string");
             assert_eq!(old_state_content, state_content);
             let pid_content =
-                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("Read to string");
+                std::fs::read_to_string(tmp.join(CGROUP_PROCS)).expect("read to string");
             assert_eq!(pid_content, "1002");
         }
     }