feat: seat cgroups proc number

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

container/Cargo.toml

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

container/src/cgroups.rs

@@ -0,0 +1,29 @@
+use cgroups_rs::Cgroup;
+use cgroups_rs::Resources;
+
+#[derive(Default, Clone)]
+pub struct Cgroups {
+    cgroup: Cgroup,
+    resources: Resources,
+}
+
+impl Cgroups {
+    pub fn new() -> Self {
+        let h = cgroups_rs::hierarchies::auto();
+        let cgroup: Cgroup = Cgroup::new(h, "kaps");
+        let resources = cgroups_rs::Resources::default();
+
+        Cgroups {
+            cgroup,
+            resources,
+        }
+    }
+
+    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,59 @@
+use oci_spec::runtime::LinuxCpu;
+use anyhow::{Result};
+
+const CGROUP_CPU_SHARES: &str = "cpu.shares";
+const CGROUP_CPU_QUOTA: &str = "cpu.cfs_quota_us";
+const CGROUP_CPU_PERIOD: &str = "cpu.cfs_period_us";
+const CGROUP_CPU_BURST: &str = "cpu.cfs_burst_us";
+const CGROUP_CPU_RT_RUNTIME: &str = "cpu.rt_runtime_us";
+const CGROUP_CPU_RT_PERIOD: &str = "cpu.rt_period_us";
+const CGROUP_CPU_STAT: &str = "cpu.stat";
+const CGROUP_CPU_PROCS: &str = "cgroup.procs";
+
+pub struct Cpu {}
+
+impl Cpu {
+    fn apply(cgroups: &Cgroups, cpu: &LinuxCpu) -> Result<()> {
+        if let Some(cpu_shares) = cpu.shares() {
+            if cpu_shares != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_SHARES.to_string(), cpu_shares.to_string());
+            }
+        }
+
+        if let Some(cpu_period) = cpu.period() {
+            if cpu_period != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_PERIOD.to_string(), cpu_period.to_string());
+            }
+        }
+
+        if let Some(cpu_quota) = cpu.quota() {
+            if cpu_quota != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_QUOTA.to_string(), cpu_quota.to_string());
+            }
+        }
+
+        if let Some(cpu_burst) = cpu.burst() {
+            cgroups.resources.cpu.attrs.insert(CGROUP_CPU_BURST.to_string(), cpu_burst.to_string());
+        }
+
+        if let Some(rt_runtime) = cpu.realtime_runtime() {
+            if rt_runtime != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_RT_RUNTIME.to_string(), rt_runtime.to_string());
+            }
+        }
+
+        if let Some(rt_period) = cpu.realtime_period() {
+            if rt_period != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_RT_PERIOD.to_string(), rt_period.to_string());
+            }
+        }
+
+        if let Some(cpus) = cpu.cpus() {
+            if cpus != 0 {
+                cgroups.resources.cpu.attrs.insert(CGROUP_CPU_PROCS.to_string(), cpus.to_string());
+            }
+        }
+
+        Ok(())
+    }
+}

container/src/lib.rs

@@ -2,11 +2,13 @@ use std::path::PathBuf;
 
 use oci_spec::runtime::Spec;
 
+use cgroups::Cgroups;
 use command::Command;
 use environment::Environment;
 use mounts::Mounts;
 use namespaces::Namespaces;
 
+mod cgroups;
 mod command;
 mod environment;
 mod mounts;
@@ -45,11 +47,13 @@ pub struct Container {
     environment: Environment,
     /// The command entrypoint
     command: Command,
+    /// Nomber of cpus
+    cpus: u64
 }
 
 impl Container {
     /// Build a new container with the bundle provided in parameters.
-    pub fn new(bundle_path: &str) -> Result<Self> {
+    pub fn new(bundle_path: &str, cpus: u64) -> Result<Self> {
         let bundle = PathBuf::from(bundle_path);
 
         // Load the specification from the file
@@ -78,6 +82,7 @@ impl Container {
             command: Command::from(spec.process()),
             namespaces,
             rootfs,
+            cpus,
             ..Default::default()
         })
     }
@@ -98,6 +103,10 @@ impl Container {
                     return self.mounts.cleanup(self.rootfs.clone());
                 }
             };
+
+            let mut cgroups = Cgroups::new();
+            cgroups.set_cpus(self.cpus);
+
             child.wait().map_err(Error::ContainerWaitCommand)?.code()
         };
 

src/cli/run.rs

@@ -17,12 +17,15 @@ pub struct RunCommand {
     /// The bundle used by the container.
     #[clap(short, long)]
     bundle: String,
+    /// The number of cpus used by the container.
+    #[clap(short, long)]
+    cpus: u64,
 }
 
 impl Handler for RunCommand {
     fn handler(&self) -> Result<()> {
         // Create a container by passing the bundle provided in arguments to it's constructor.
-        let container = Container::new(&self.bundle)?;
+        let container = Container::new(&self.bundle, self.cpus)?;
 
         // Run the container
         // At the moment, we don't have a detached mode for the container,