Merge pull request #54 from Furisto/cg-common

Align cgroup controller implementations
youki-dev · Jun 2, 2021 · 2d342c0 · 2d342c0
2 parents 52b9a15 + 3c707de
commit 2d342c0
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 169 deletions.
diff --git a/src/cgroups/common.rs b/src/cgroups/common.rs
@@ -38,18 +38,8 @@ impl Display for Cgroup {
     }
 }
 
-pub fn write_cgroup_file_truncate(path: &Path, data: &str) -> Result<()> {
-    fs::OpenOptions::new()
-        .create(false)
-        .write(true)
-        .truncate(true)
-        .open(path)?
-        .write_all(data.as_bytes())?;
-
-    Ok(())
-}
-
-pub fn write_cgroup_file(path: &Path, data: &str) -> Result<()> {
+#[inline]
+pub fn write_cgroup_file<P: AsRef<Path>>(path: P, data: &str) -> Result<()> {
     fs::OpenOptions::new()
         .create(false)
         .write(true)
@@ -94,7 +84,7 @@ pub fn create_cgroup_manager<P: Into<PathBuf>>(cgroup_path: P) -> Result<Box<dyn
                     )?))
                 }
                 _ => Ok(Box::new(v1::manager::Manager::new(cgroup_path.into())?)),
-            }
+            } 
         }
         _ => bail!("could not find cgroup filesystem"),
     }

diff --git a/src/cgroups/v1/blkio.rs b/src/cgroups/v1/blkio.rs
@@ -1,10 +1,9 @@
 use std::{
-    fs::{self, OpenOptions},
-    io::Write,
+    fs::{self},
     path::Path,
 };
 
-use crate::cgroups::v1::Controller;
+use crate::cgroups::{common, v1::Controller};
 use oci_spec::{LinuxBlockIo, LinuxResources};
 
 const CGROUP_BLKIO_THROTTLE_READ_BPS: &str = "blkio.throttle.read_bps_device";
@@ -20,73 +19,55 @@ impl Controller for Blkio {
         cgroup_root: &Path,
         pid: nix::unistd::Pid,
     ) -> anyhow::Result<()> {
-        match &linux_resources.block_io {
-            None => return Ok(()),
-            Some(block_io) => {
-                fs::create_dir_all(cgroup_root)?;
-                Self::apply(cgroup_root, block_io)?;
-            }
-        }
+        log::debug!("Apply blkio cgroup config");
+        fs::create_dir_all(cgroup_root)?;
 
-        OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(false)
-            .open(cgroup_root.join("cgroup.procs"))?
-            .write_all(pid.to_string().as_bytes())?;
+        if let Some(blkio) = &linux_resources.block_io {
+            Self::apply(cgroup_root, blkio)?;
+        }
 
+        common::write_cgroup_file(&cgroup_root.join("cgroup.procs"), &pid.to_string())?;
         Ok(())
     }
 }
 
 impl Blkio {
     fn apply(root_path: &Path, blkio: &LinuxBlockIo) -> anyhow::Result<()> {
         for trbd in &blkio.blkio_throttle_read_bps_device {
-            Self::write_file(
+            common::write_cgroup_file(
                 &root_path.join(CGROUP_BLKIO_THROTTLE_READ_BPS),
                 &format!("{}:{} {}", trbd.major, trbd.minor, trbd.rate),
             )?;
         }
 
         for twbd in &blkio.blkio_throttle_write_bps_device {
-            Self::write_file(
+            common::write_cgroup_file(
                 &root_path.join(CGROUP_BLKIO_THROTTLE_WRITE_BPS),
                 &format!("{}:{} {}", twbd.major, twbd.minor, twbd.rate),
             )?;
         }
 
         for trid in &blkio.blkio_throttle_read_iops_device {
-            Self::write_file(
+            common::write_cgroup_file(
                 &root_path.join(CGROUP_BLKIO_THROTTLE_READ_IOPS),
                 &format!("{}:{} {}", trid.major, trid.minor, trid.rate),
             )?;
         }
 
         for twid in &blkio.blkio_throttle_write_iops_device {
-            Self::write_file(
+            common::write_cgroup_file(
                 &root_path.join(CGROUP_BLKIO_THROTTLE_WRITE_IOPS),
                 &format!("{}:{} {}", twid.major, twid.minor, twid.rate),
             )?;
         }
 
         Ok(())
     }
-
-    fn write_file(file_path: &Path, data: &str) -> anyhow::Result<()> {
-        fs::OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(false)
-            .open(file_path)?
-            .write_all(data.as_bytes())?;
-
-        Ok(())
-    }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::path::PathBuf;
+    use std::{io::Write, path::PathBuf};
 
     use super::*;
     use oci_spec::{LinuxBlockIo, LinuxThrottleDevice};

diff --git a/src/cgroups/v1/devices.rs b/src/cgroups/v1/devices.rs
@@ -1,12 +1,9 @@
-use std::io::Write;
-use std::{
-    fs::{create_dir_all, OpenOptions},
-    path::Path,
-};
+use std::{fs::create_dir_all, path::Path};
 
 use anyhow::Result;
 use nix::unistd::Pid;
 
+use crate::cgroups::common;
 use crate::{cgroups::v1::Controller, rootfs::default_devices};
 use oci_spec::{LinuxDeviceCgroup, LinuxDeviceType, LinuxResources};
 
@@ -30,12 +27,7 @@ impl Controller for Devices {
             Self::apply_device(&d, &cgroup_root)?;
         }
 
-        OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(false)
-            .open(cgroup_root.join("cgroup.procs"))?
-            .write_all(pid.to_string().as_bytes())?;
+        common::write_cgroup_file(cgroup_root.join("cgroup.procs"), &pid.to_string())?;
         Ok(())
     }
 }
@@ -48,12 +40,7 @@ impl Devices {
             cgroup_root.join("devices.deny")
         };
 
-        OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(false)
-            .open(path)?
-            .write_all(device.to_string().as_bytes())?;
+        common::write_cgroup_file(path, &device.to_string())?;
         Ok(())
     }
 

diff --git a/src/cgroups/v1/hugetlb.rs b/src/cgroups/v1/hugetlb.rs
@@ -1,13 +1,9 @@
-use std::{
-    fs::{self, OpenOptions},
-    io::Write,
-    path::Path,
-};
+use std::{fs, path::Path};
 
 use anyhow::anyhow;
 use regex::Regex;
 
-use crate::cgroups::v1::Controller;
+use crate::cgroups::{common, v1::Controller};
 use oci_spec::{LinuxHugepageLimit, LinuxResources};
 
 pub struct Hugetlb {}
@@ -25,12 +21,7 @@ impl Controller for Hugetlb {
             Self::apply(cgroup_root, hugetlb)?
         }
 
-        OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(false)
-            .open(cgroup_root.join("cgroup.procs"))?
-            .write_all(pid.to_string().as_bytes())?;
+        common::write_cgroup_file(cgroup_root.join("cgroup.procs"), &pid.to_string())?;
         Ok(())
     }
 }
@@ -49,32 +40,21 @@ impl Hugetlb {
             }
         }
 
-        Self::write_file(
+        common::write_cgroup_file(
             &root_path.join(format!("hugetlb.{}.limit_in_bytes", hugetlb.page_size)),
             &hugetlb.limit.to_string(),
         )?;
         Ok(())
     }
 
-    fn write_file(file_path: &Path, data: &str) -> anyhow::Result<()> {
-        fs::OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(true)
-            .open(file_path)?
-            .write_all(data.as_bytes())?;
-
-        Ok(())
-    }
-
     fn is_power_of_two(number: u64) -> bool {
         (number != 0) && (number & (number - 1)) == 0
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::path::PathBuf;
+    use std::{io::Write, path::PathBuf};
 
     use super::*;
     use oci_spec::LinuxHugepageLimit;

diff --git a/src/cgroups/v1/manager.rs b/src/cgroups/v1/manager.rs
@@ -12,7 +12,7 @@ use super::{
     Controller,
 };
 
-use crate::{cgroups::v1::ControllerType, cgroups::common::CgroupManager, utils::PathBufExt};
+use crate::{cgroups::common::CgroupManager, cgroups::v1::ControllerType, utils::PathBufExt};
 use oci_spec::LinuxResources;
 
 const CONTROLLERS: &[ControllerType] = &[

diff --git a/src/cgroups/v1/memory.rs b/src/cgroups/v1/memory.rs
@@ -7,6 +7,7 @@ use std::{
 use anyhow::{Result, *};
 use nix::{errno::Errno, unistd::Pid};
 
+use crate::cgroups::common;
 use crate::cgroups::v1::Controller;
 use oci_spec::{LinuxMemory, LinuxResources};
 
@@ -64,14 +65,9 @@ impl Controller for Memory {
             if let Some(tcp_mem) = memory.kernel_tcp {
                 Self::set(tcp_mem, &cgroup_root.join(CGROUP_KERNEL_TCP_MEMORY_LIMIT))?;
             }
-
-            OpenOptions::new()
-                .create(false)
-                .write(true)
-                .truncate(false)
-                .open(cgroup_root.join("cgroup.procs"))?
-                .write_all(pid.to_string().as_bytes())?;
         }
+
+        common::write_cgroup_file(cgroup_root.join("cgroup.procs"), &pid.to_string())?;
         Ok(())
     }
 }
@@ -180,7 +176,6 @@ impl Memory {
         }
 
         let path = cgroup_root.join(CGROUP_MEMORY_SWAP_LIMIT);
-
         Self::set(val, &path)?;
 
         Ok(())

diff --git a/src/cgroups/v1/network_classifier.rs b/src/cgroups/v1/network_classifier.rs
@@ -1,12 +1,9 @@
-use std::io::Write;
-use std::{
-    fs::{create_dir_all, OpenOptions},
-    path::Path,
-};
+use std::{fs::create_dir_all, path::Path};
 
 use anyhow::Result;
 use nix::unistd::Pid;
 
+use crate::cgroups::common;
 use crate::cgroups::v1::Controller;
 use oci_spec::{LinuxNetwork, LinuxResources};
 
@@ -19,43 +16,26 @@ impl Controller for NetworkClassifier {
 
         if let Some(network) = linux_resources.network.as_ref() {
             Self::apply(cgroup_root, network)?;
-
-            OpenOptions::new()
-                .create(false)
-                .write(true)
-                .truncate(true)
-                .open(cgroup_root.join("cgroup.procs"))?
-                .write_all(pid.to_string().as_bytes())?;
         }
 
+        common::write_cgroup_file(cgroup_root.join("cgroup.procs"), &pid.to_string())?;
         Ok(())
     }
 }
 
 impl NetworkClassifier {
     fn apply(root_path: &Path, network: &LinuxNetwork) -> Result<()> {
         if let Some(class_id) = network.class_id {
-            Self::write_file(&root_path.join("net_cls.classid"), &class_id.to_string())?;
+            common::write_cgroup_file(&root_path.join("net_cls.classid"), &class_id.to_string())?;
         }
 
         Ok(())
     }
-
-    fn write_file(file_path: &Path, data: &str) -> Result<()> {
-        OpenOptions::new()
-            .create(false)
-            .write(true)
-            .truncate(true)
-            .open(file_path)?
-            .write_all(data.as_bytes())?;
-
-        Ok(())
-    }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::path::PathBuf;
+    use std::{io::Write, path::PathBuf};
 
     use super::*;