Initial support for cgroups v2

oci_spec/src/lib.rs

@@ -363,9 +363,9 @@ pub struct LinuxCpu {
     pub realtime_runtime: Option<i64>,
     pub realtime_period: Option<u64>,
     #[serde(default)]
-    pub cpus: String,
+    pub cpus: Option<String>,
     #[serde(default)]
-    pub mems: String,
+    pub mems: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -453,7 +453,6 @@ pub struct LinuxResources {
     pub disable_oom_killer: bool,
     pub oom_score_adj: Option<i32>,
     pub memory: Option<LinuxMemory>,
-    #[serde(rename = "LinuxCPU")]
     pub cpu: Option<LinuxCpu>,
     pub pids: Option<LinuxPids>,
     #[serde(rename = "blockIO")]

src/capabilities.rs

@@ -1,6 +1,4 @@
-use crate::{
-    command::Command,
-};
+use crate::command::Command;
 use caps::*;
 
 use anyhow::Result;

src/cgroups/common.rs

@@ -0,0 +1,103 @@
+use std::{
+    env,
+    fmt::{Debug, Display},
+    fs,
+    io::Write,
+    path::{Path, PathBuf},
+};
+
+use anyhow::{anyhow, Result};
+use nix::unistd::Pid;
+use oci_spec::LinuxResources;
+use procfs::process::Process;
+
+use crate::cgroups::v1;
+use crate::cgroups::v2;
+
+pub const DEFAULT_CGROUP_ROOT: &str = "/sys/fs/cgroup";
+
+pub trait CgroupManager {
+    fn apply(&self, linux_resources: &LinuxResources, pid: Pid) -> Result<()>;
+    fn remove(&self) -> Result<()>;
+}
+
+#[derive(Debug)]
+pub enum Cgroup {
+    V1,
+    V2,
+}
+
+impl Display for Cgroup {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let print = match *self {
+            Cgroup::V1 => "v1",
+            Cgroup::V2 => "v2",
+        };
+
+        write!(f, "{}", print)
+    }
+}
+
+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<()> {
+    fs::OpenOptions::new()
+        .create(false)
+        .write(true)
+        .truncate(false)
+        .open(path)?
+        .write_all(data.as_bytes())?;
+
+    Ok(())
+}
+
+pub fn create_cgroup_manager<P: Into<PathBuf>>(cgroup_path: P) -> Result<Box<dyn CgroupManager>> {
+    let cgroup_mount = Process::myself()?
+        .mountinfo()?
+        .into_iter()
+        .find(|m| m.fs_type == "cgroup");
+
+    let cgroup2_mount = Process::myself()?
+        .mountinfo()?
+        .into_iter()
+        .find(|m| m.fs_type == "cgroup2");
+
+    match (cgroup_mount, cgroup2_mount) {
+        (Some(_), None) => {
+            log::info!("cgroup manager V1 will be used");
+            return Ok(Box::new(v1::manager::Manager::new(cgroup_path.into())?));
+        }
+        (None, Some(cgroup2)) => {
+            log::info!("cgroup manager V2 will be used");
+            return Ok(Box::new(v2::manager::Manager::new(
+                cgroup2.mount_point,
+                cgroup_path.into(),
+            )?));
+        }
+        (Some(_), Some(cgroup2)) => {
+            let cgroup_override = env::var("YOUKI_PREFER_CGROUPV2");
+            match cgroup_override {
+                Ok(v) if v == "true" => {
+                    log::info!("cgroup manager V2 will be used");
+                    return Ok(Box::new(v2::manager::Manager::new(
+                        cgroup2.mount_point,
+                        cgroup_path.into(),
+                    )?));
+                }
+                _ => {
+                    return Ok(Box::new(v1::manager::Manager::new(cgroup_path.into())?))
+                }
+            }
+        }
+        _ => return Err(anyhow!("could not find cgroup filesystem")),
+    }
+}

src/cgroups/mod.rs

@@ -1,13 +1,4 @@
-mod controller;
-mod controller_type;
-mod devices;
-mod hugetlb;
-mod blkio;
-mod manager;
-mod memory;
-mod network_classifier;
-mod network_priority;
-mod pids;
-pub use controller::Controller;
-pub use controller_type::ControllerType;
-pub use manager::Manager;
+pub mod common;
+mod test;
+pub mod v1;
+pub mod v2;

src/cgroups/test.rs

@@ -0,0 +1,86 @@
+#![cfg(test)]
+
+use anyhow::Result;
+use std::{
+    io::Write,
+    path::{Path, PathBuf},
+};
+
+use oci_spec::LinuxCpu;
+
+pub fn set_fixture(temp_dir: &Path, filename: &str, val: &str) -> Result<PathBuf> {
+    let full_path = temp_dir.join(filename);
+
+    std::fs::OpenOptions::new()
+        .create(true)
+        .write(true)
+        .truncate(true)
+        .open(&full_path)?
+        .write_all(val.as_bytes())?;
+
+    Ok(full_path)
+}
+
+pub fn create_temp_dir(test_name: &str) -> Result<PathBuf> {
+    std::fs::create_dir_all(std::env::temp_dir().join(test_name))?;
+    Ok(std::env::temp_dir().join(test_name))
+}
+
+pub struct LinuxCpuBuilder {
+    resource: LinuxCpu,
+}
+
+impl LinuxCpuBuilder {
+    pub fn new() -> Self {
+        Self {
+            resource: LinuxCpu {
+                shares: None,
+                quota: None,
+                period: None,
+                realtime_runtime: None,
+                realtime_period: None,
+                cpus: None,
+                mems: None,
+            },
+        }
+    }
+
+    pub fn with_shares(mut self, shares: u64) -> Self {
+        self.resource.shares = Some(shares);
+        self
+    }
+
+    pub fn with_quota(mut self, quota: i64) -> Self {
+        self.resource.quota = Some(quota);
+        self
+    }
+
+    pub fn with_period(mut self, period: u64) -> Self {
+        self.resource.period = Some(period);
+        self
+    }
+
+    pub fn with_realtime_runtime(mut self, runtime: i64) -> Self {
+        self.resource.realtime_runtime = Some(runtime);
+        self
+    }
+
+    pub fn with_realtime_period(mut self, period: u64) -> Self {
+        self.resource.realtime_period = Some(period);
+        self
+    }
+
+    pub fn with_cpus(mut self, cpus: String) -> Self {
+        self.resource.cpus = Some(cpus);
+        self
+    }
+
+    pub fn with_mems(mut self, mems: String) -> Self {
+        self.resource.mems = Some(mems);
+        self
+    }
+
+    pub fn build(self) -> LinuxCpu {
+        self.resource
+    }
+}

src/cgroups/blkio.rs → src/cgroups/v1/blkio.rs

@@ -4,9 +4,7 @@ use std::{
     path::Path,
 };
 
-use crate::{
-    cgroups::Controller,
-};
+use crate::cgroups::v1::Controller;
 use oci_spec::{LinuxBlockIo, LinuxResources};
 
 const CGROUP_BLKIO_THROTTLE_READ_BPS: &str = "blkio.throttle.read_bps_device";

src/cgroups/devices.rs → src/cgroups/v1/devices.rs

@@ -7,10 +7,7 @@ use std::{
 use anyhow::Result;
 use nix::unistd::Pid;
 
-use crate::{
-    cgroups::Controller,
-    rootfs::default_devices,
-};
+use crate::{cgroups::v1::Controller, rootfs::default_devices};
 use oci_spec::{LinuxDeviceCgroup, LinuxDeviceType, LinuxResources};
 
 pub struct Devices {}

src/cgroups/hugetlb.rs → src/cgroups/v1/hugetlb.rs

@@ -7,9 +7,7 @@ use std::{
 use anyhow::anyhow;
 use regex::Regex;
 
-use crate::{
-    cgroups::Controller,
-};
+use crate::cgroups::v1::Controller;
 use oci_spec::{LinuxHugepageLimit, LinuxResources};
 
 pub struct Hugetlb {}

src/cgroups/manager.rs → src/cgroups/v1/manager.rs

@@ -3,14 +3,16 @@ use std::{fs::remove_dir, path::Path};
 
 use anyhow::Result;
 use nix::unistd::Pid;
+
 use procfs::process::Process;
 
 use super::{
     blkio::Blkio, devices::Devices, hugetlb::Hugetlb, memory::Memory,
     network_classifier::NetworkClassifier, network_priority::NetworkPriority, pids::Pids,
     Controller,
 };
-use crate::{cgroups::ControllerType, utils::PathBufExt};
+
+use crate::{cgroups::v1::ControllerType, cgroups::common::CgroupManager, utils::PathBufExt};
 use oci_spec::LinuxResources;
 
 const CONTROLLERS: &[ControllerType] = &[
@@ -28,46 +30,18 @@ pub struct Manager {
 }
 
 impl Manager {
-    pub fn new(cgroup_path: &Path) -> Result<Self> {
+    pub fn new(cgroup_path: PathBuf) -> Result<Self> {
         let mut subsystems = HashMap::<String, PathBuf>::new();
         for subsystem in CONTROLLERS.iter().map(|c| c.to_string()) {
             subsystems.insert(
                 subsystem.to_owned(),
-                Self::get_subsystem_path(cgroup_path, &subsystem)?,
+                Self::get_subsystem_path(&cgroup_path, &subsystem)?,
             );
         }
 
         Ok(Manager { subsystems })
     }
 
-    pub fn apply(&self, linux_resources: &LinuxResources, pid: Pid) -> Result<()> {
-        for subsys in &self.subsystems {
-            match subsys.0.as_str() {
-                "devices" => Devices::apply(linux_resources, &subsys.1, pid)?,
-                "hugetlb" => Hugetlb::apply(linux_resources, &subsys.1, pid)?,
-                "memory" => Memory::apply(linux_resources, &subsys.1, pid)?,
-                "pids" => Pids::apply(linux_resources, &subsys.1, pid)?,
-                "blkio" => Blkio::apply(linux_resources, &subsys.1, pid)?,
-                "net_prio" => NetworkPriority::apply(linux_resources, &subsys.1, pid)?,
-                "net_cls" => NetworkClassifier::apply(linux_resources, &subsys.1, pid)?,
-                _ => continue,
-            }
-        }
-
-        Ok(())
-    }
-
-    pub fn remove(&self) -> Result<()> {
-        for cgroup_path in &self.subsystems {
-            if cgroup_path.1.exists() {
-                log::debug!("remove cgroup {:?}", cgroup_path.1);
-                remove_dir(&cgroup_path.1)?;
-            }
-        }
-
-        Ok(())
-    }
-
     fn get_subsystem_path(cgroup_path: &Path, subsystem: &str) -> anyhow::Result<PathBuf> {
         log::debug!("Get path for subsystem: {}", subsystem);
         let mount = Process::myself()?
@@ -104,3 +78,33 @@ impl Manager {
         Ok(p)
     }
 }
+
+impl CgroupManager for Manager {
+    fn apply(&self, linux_resources: &LinuxResources, pid: Pid) -> Result<()> {
+        for subsys in &self.subsystems {
+            match subsys.0.as_str() {
+                "devices" => Devices::apply(linux_resources, &subsys.1, pid)?,
+                "hugetlb" => Hugetlb::apply(linux_resources, &subsys.1, pid)?,
+                "memory" => Memory::apply(linux_resources, &subsys.1, pid)?,
+                "pids" => Pids::apply(linux_resources, &subsys.1, pid)?,
+                "blkio" => Blkio::apply(linux_resources, &subsys.1, pid)?,
+                "net_prio" => NetworkPriority::apply(linux_resources, &subsys.1, pid)?,
+                "net_cls" => NetworkClassifier::apply(linux_resources, &subsys.1, pid)?,
+                _ => continue,
+            }
+        }
+
+        Ok(())
+    }
+
+    fn remove(&self) -> Result<()> {
+        for cgroup_path in &self.subsystems {
+            if cgroup_path.1.exists() {
+                log::debug!("remove cgroup {:?}", cgroup_path.1);
+                remove_dir(&cgroup_path.1)?;
+            }
+        }
+
+        Ok(())
+    }
+}

src/cgroups/memory.rs → src/cgroups/v1/memory.rs

@@ -7,9 +7,7 @@ use std::{
 use anyhow::{Result, *};
 use nix::{errno::Errno, unistd::Pid};
 
-use crate::{
-    cgroups::Controller,
-};
+use crate::cgroups::v1::Controller;
 use oci_spec::{LinuxMemory, LinuxResources};
 
 const CGROUP_MEMORY_SWAP_LIMIT: &str = "memory.memsw.limit_in_bytes";

src/cgroups/v1/mod.rs

@@ -0,0 +1,13 @@
+mod blkio;
+mod controller;
+mod controller_type;
+mod devices;
+mod hugetlb;
+pub mod manager;
+mod memory;
+mod network_classifier;
+mod network_priority;
+mod pids;
+pub use controller::Controller;
+pub use controller_type::ControllerType;
+pub use manager::Manager;