kaps: add memory and cpu cgroups support

Signed-off-by: Hugo Amalric <[email protected]>

container/Cargo.toml

@@ -10,6 +10,8 @@ oci-spec = "0.5.3"
 unshare = { git = "https://github.com/virt-do/unshare", branch = "main" }
 serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
+controlgroup = "0.3.0"
+anyhow = "1.0.57"
 
 [dev-dependencies]
 proc-mounts = "0.3.0"

container/src/cgroups.rs

@@ -0,0 +1,19 @@
+
+/*pub struct Cgroups {
+    cgroup: Cgroup,
+    resources: Resources,
+}
+
+impl Cgroups {
+    pub fn new() -> Self {
+
+    }
+
+    pub fn set_cpus(&mut self, cpus: u64) {
+        self.resources
+            .cpu
+            .attrs
+            .insert("cgroup.procs".to_string(), cpus.to_string());
+        self.cgroup.apply(&self.resources).unwrap();
+    }
+}*/

container/src/cpu.rs

@@ -0,0 +1,68 @@
+use anyhow::Result;
+use controlgroup::{
+    v1::{cpu, Cgroup, CgroupPath, SubsystemKind},
+    Pid,
+};
+use oci_spec::runtime::LinuxCpu;
+use std::path::PathBuf;
+
+pub struct Cpu {
+    cpu_cgroup: cpu::Subsystem,
+}
+
+impl Cpu {
+    pub fn new() -> Self {
+        // Define and create a new cgroup controlled by the CPU subsystem.
+        let mut cpu_cgroup =
+            cpu::Subsystem::new(CgroupPath::new(SubsystemKind::Cpu, PathBuf::from("kaps")));
+        cpu_cgroup.create().unwrap();
+        Cpu { cpu_cgroup }
+    }
+
+    pub fn apply(&mut self, cpu: &LinuxCpu) -> Result<()> {
+        if let Some(cpu_shares) = cpu.shares() {
+            if cpu_shares != 0 {
+                let _ = &self.cpu_cgroup.set_shares(cpu_shares);
+            }
+        }
+
+        if let Some(cpu_period) = cpu.period() {
+            if cpu_period != 0 {
+                let _ = &self.cpu_cgroup.set_cfs_period_us(cpu_period);
+            }
+        }
+
+        if let Some(cpu_quota) = cpu.quota() {
+            if cpu_quota != 0 {
+                let _ = &self.cpu_cgroup.set_cfs_quota_us(cpu_quota);
+            }
+        }
+
+        if let Some(rt_runtime) = cpu.realtime_runtime() {
+            if rt_runtime != 0 {
+                let _ = &self.cpu_cgroup.set_rt_runtime_us(rt_runtime);
+            }
+        }
+
+        if let Some(rt_period) = cpu.realtime_period() {
+            if rt_period != 0 {
+                let _ = &self.cpu_cgroup.set_rt_period_us(rt_period);
+            }
+        }
+
+        // Attach the self process to the cgroup.
+        let pid = Pid::from(std::process::id());
+        self.cpu_cgroup.add_task(pid).unwrap();
+
+        Ok(())
+    }
+
+    pub fn delete(&mut self) -> Result<()> {
+        // Removing self process from the cgroup
+        let pid = Pid::from(std::process::id());
+        self.cpu_cgroup.remove_task(pid)?;
+        // and deleting the cgroup.
+        self.cpu_cgroup.delete()?;
+        Ok(())
+    }
+}

container/src/lib.rs

@@ -5,15 +5,21 @@ use std::path::PathBuf;
 use oci_spec::runtime::Spec;
 
 use command::Command;
+use cpu::Cpu;
 use environment::Environment;
+use memory::Memory;
 use mounts::Mounts;
 use namespaces::Namespaces;
+use oci_spec::runtime::LinuxResources;
 use state::{ContainerState, Status};
 
 mod command;
+mod cpu;
 mod environment;
+mod memory;
 mod mounts;
 mod namespaces;
+mod resources;
 pub mod spec;
 mod state;
 
@@ -67,6 +73,8 @@ pub struct Container {
     command: Command,
     /// The container state
     state: ContainerState,
+    /// The container resources,
+    resources: LinuxResources,
 }
 
 impl Container {
@@ -95,6 +103,19 @@ impl Container {
                 Namespaces::from(linux.namespaces())
             });
 
+        // Get the container resources if the linux block is defined into the specification.
+        let resources = spec
+            .linux()
+            .as_ref()
+            .map_or(LinuxResources::default(), |linux| {
+                linux
+                    .resources()
+                    .as_ref()
+                    .map_or(LinuxResources::default(), |resources| {
+                        LinuxResources::from(resources.clone())
+                    })
+            });
+
         // Set the state of the container
         let state = ContainerState::new(id, bundle_path)?;
 
@@ -104,13 +125,25 @@ impl Container {
             namespaces,
             rootfs,
             state,
+            resources,
             ..Default::default()
         })
     }
 
     /// Run the container.
     pub fn run(&mut self) -> Result<()> {
         let mounts = self.mounts.clone();
+        let mut cpu_cgroup = Cpu::new();
+        let mut memory_cgroup = Memory::new();
+
+        if let Some(resources_cpu) = &self.resources.cpu() {
+            cpu_cgroup.apply(&resources_cpu.clone()).unwrap();
+        }
+
+        if let Some(resources_memory) = &self.resources.memory() {
+            memory_cgroup.apply(&resources_memory.clone()).unwrap();
+        }
+
         let code = unsafe {
             let mut child = match unshare::Command::from(&self.command)
                 .chroot_dir(&self.rootfs)
@@ -126,11 +159,14 @@ impl Container {
             };
 
             self.state.pid = child.pid();
-            self.state.set_status(Status::Running)?;
+            self.state.set_status(Status::Running).unwrap();
 
             child.wait().map_err(Error::ContainerWaitCommand)?.code()
         };
 
+        cpu_cgroup.delete().unwrap();
+        memory_cgroup.delete().unwrap();
+
         self.mounts.cleanup(self.rootfs.clone())?;
         self.state.remove()?;
 

container/src/memory.rs

@@ -0,0 +1,74 @@
+use anyhow::Result;
+use controlgroup::{
+    v1::{memory, Cgroup, CgroupPath, SubsystemKind},
+    Pid,
+};
+use oci_spec::runtime::LinuxMemory;
+use std::path::PathBuf;
+
+pub struct Memory {
+    memory_cgroup: memory::Subsystem,
+}
+
+impl Memory {
+    pub fn new() -> Self {
+        // Define and create a new cgroup controlled by the Memory subsystem.
+        let mut memory_cgroup = memory::Subsystem::new(CgroupPath::new(
+            SubsystemKind::Memory,
+            PathBuf::from("kaps"),
+        ));
+        memory_cgroup.create().unwrap();
+        Memory { memory_cgroup }
+    }
+
+    pub fn apply(&mut self, memory: &LinuxMemory) -> Result<()> {
+        if let Some(limit) = memory.limit() {
+            if limit != 0 {
+                let _ = &self.memory_cgroup.set_limit_in_bytes(limit);
+            }
+        }
+
+        if let Some(swappiness) = memory.swappiness() {
+            if swappiness != 0 {
+                let _ = &self.memory_cgroup.set_swappiness(swappiness);
+            }
+        }
+
+        if let Some(kernel) = memory.kernel() {
+            if kernel != 0 {
+                let _ = &self.memory_cgroup.set_kmem_limit_in_bytes(kernel);
+            }
+        }
+
+        if let Some(kernel_tcp) = memory.kernel_tcp() {
+            if kernel_tcp != 0 {
+                let _ = &self.memory_cgroup.set_kmem_tcp_limit_in_bytes(kernel_tcp);
+            }
+        }
+
+        if let Some(reservation) = memory.reservation() {
+            if reservation != 0 {
+                let _ = &self.memory_cgroup.set_soft_limit_in_bytes(reservation);
+            }
+        }
+
+        if let Some(disable_oom_killer) = memory.disable_oom_killer() {
+            let _ = &self.memory_cgroup.disable_oom_killer(disable_oom_killer);
+        }
+
+        // Attach the self process to the cgroup.
+        let pid = Pid::from(std::process::id());
+        self.memory_cgroup.add_task(pid).unwrap();
+
+        Ok(())
+    }
+
+    pub fn delete(&mut self) -> Result<()> {
+        // Removing self process from the cgroup
+        let pid = Pid::from(std::process::id());
+        self.memory_cgroup.remove_task(pid)?;
+        // and deleting the cgroup.
+        self.memory_cgroup.delete()?;
+        Ok(())
+    }
+}

src/cli/run.rs

@@ -26,7 +26,7 @@ pub struct RunCommand {
 #[async_trait]
 impl Handler for RunCommand {
     async fn handler(&self, _: &mut env_logger::Builder) -> Result<()> {
-        // Create a container by passing the bundle and the id provided in arguments to it's constructor.
+        // Create a container by passing the bundle provided in arguments to it's constructor.
         let mut container = Container::new(&self.bundle, &self.name)?;
 
         // Run the container