add exec-cpu-affinity

Signed-off-by: Yusuke Sakurai <yusuke.sakurai@3-shake.com>

Author: saku3

Cargo.lock

@@ -447,6 +447,12 @@ impl TenantContainerBuilder {
                 process_builder = process_builder.cwd(cwd);
             }
 
+            if let Some(process) = spec.process() {
+                if let Some(cpu_affinity) = process.exec_cpu_affinity() {
+                    process_builder = process_builder.exec_cpu_affinity(cpu_affinity.clone());
+                }
+            }
+
             if let Some(no_new_priv) = self.get_no_new_privileges() {
                 process_builder = process_builder.no_new_privileges(no_new_priv);
             }

crates/libcontainer/src/container/tenant_builder.rs

@@ -11,7 +11,7 @@ use super::fork::CloneCb;
 use super::init::process as init_process;
 use crate::error::MissingSpecError;
 use crate::namespaces::Namespaces;
-use crate::process::{channel, fork};
+use crate::process::{channel, cpu_affinity, fork};
 
 #[derive(Debug, thiserror::Error)]
 pub enum IntermediateProcessError {
@@ -31,6 +31,8 @@ pub enum IntermediateProcessError {
     ExecNotify(#[source] nix::Error),
     #[error(transparent)]
     MissingSpec(#[from] crate::error::MissingSpecError),
+    #[error("CPU affinity error {0}")]
+    CpuAffinity(#[from] cpu_affinity::CPUAffinityError),
     #[error("other error")]
     Other(String),
 }
@@ -52,6 +54,21 @@ pub fn container_intermediate_process(
     let cgroup_manager = libcgroups::common::create_cgroup_manager(args.cgroup_config.to_owned())
         .map_err(|e| IntermediateProcessError::Cgroup(e.to_string()))?;
 
+    let current_pid = Pid::this();
+    // setting CPU affinity for tenant container before cgroup move
+    if matches!(args.container_type, ContainerType::TenantContainer { .. }) {
+        if let Some(exec_cpu_affinity) = spec
+            .process()
+            .as_ref()
+            .and_then(|p| p.exec_cpu_affinity().as_ref())
+        {
+            if let Some(initial) = exec_cpu_affinity.initial() {
+                cpu_affinity::set_cpuset_affinity_from_string(current_pid, initial)?;
+            }
+        }
+    }
+    let _ = cpu_affinity::log_cpu_affinity();
+
     // this needs to be done before we create the init process, so that the init
     // process will already be captured by the cgroup. It also needs to be done
     // before we enter the user namespace because if a privileged user starts a
@@ -68,6 +85,19 @@ pub fn container_intermediate_process(
         matches!(args.container_type, ContainerType::InitContainer),
     )?;
 
+    // setting CPU affinity for tenant container after cgroup move
+    if matches!(args.container_type, ContainerType::TenantContainer { .. }) {
+        if let Some(exec_cpu_affinity) = spec
+            .process()
+            .as_ref()
+            .and_then(|p| p.exec_cpu_affinity().as_ref())
+        {
+            if let Some(cpu_affinity_final) = exec_cpu_affinity.cpu_affinity_final() {
+                cpu_affinity::set_cpuset_affinity_from_string(current_pid, cpu_affinity_final)?;
+            }
+        }
+    }
+
     // if new user is specified in specification, this will be true and new
     // namespace will be created, check
     // https://man7.org/linux/man-pages/man7/user_namespaces.7.html for more

crates/libcontainer/src/process/container_intermediate_process.rs

@@ -0,0 +1,156 @@
+use nix::sched::{sched_getaffinity, sched_setaffinity, CpuSet};
+use nix::unistd::Pid;
+
+#[derive(Debug, thiserror::Error)]
+pub enum CPUAffinityError {
+    #[error("invalid CPU string: {0}")]
+    ParseError(String),
+    #[error("values larger than {max} are not supported")]
+    CpuOutOfRange { cpu: usize, max: usize },
+    #[error("failed to set CPU for CPU {cpu}: {source}")]
+    CpuSet {
+        cpu: usize,
+        #[source]
+        source: nix::Error,
+    },
+    #[error("failed to setaffinity")]
+    SetAffinity(#[source] nix::Error),
+    #[error("failed to getaffinity")]
+    GetAffinity(#[source] nix::Error),
+}
+
+type Result<T> = std::result::Result<T, CPUAffinityError>;
+
+pub fn to_cpuset(cpuset_str: &str) -> Result<CpuSet> {
+    let mut cpuset = CpuSet::new();
+    let max_cpu = CpuSet::count();
+
+    for part in cpuset_str
+        .trim()
+        .split(',')
+        .map(str::trim)
+        .filter(|s| !s.is_empty())
+    {
+        match part.split_once('-') {
+            Some((start_str, end_str)) => {
+                let start = parse_cpu_index(start_str, max_cpu)?;
+                let end = parse_cpu_index(end_str, max_cpu)?;
+                if start > end {
+                    return Err(CPUAffinityError::ParseError(format!(
+                        "invalid range: {}-{}",
+                        start, end
+                    )));
+                }
+                for cpu in start..=end {
+                    cpuset
+                        .set(cpu)
+                        .map_err(|e| CPUAffinityError::CpuSet { cpu, source: e })?;
+                }
+            }
+            None => {
+                let cpu = parse_cpu_index(part, max_cpu)?;
+                cpuset
+                    .set(cpu)
+                    .map_err(|e| CPUAffinityError::CpuSet { cpu, source: e })?;
+            }
+        }
+    }
+    Ok(cpuset)
+}
+
+fn parse_cpu_index(s: &str, max_cpu: usize) -> Result<usize> {
+    let cpu: usize = s
+        .parse()
+        .map_err(|_| CPUAffinityError::ParseError(s.to_string()))?;
+    if cpu >= max_cpu {
+        return Err(CPUAffinityError::CpuOutOfRange {
+            cpu,
+            max: max_cpu - 1,
+        });
+    }
+    Ok(cpu)
+}
+
+pub fn set_cpuset_affinity_from_string(pid: Pid, cpuset_str: &str) -> Result<()> {
+    tracing::debug!(?cpuset_str, "setting CPU affinity for tenant container");
+    sched_setaffinity(pid, &to_cpuset(cpuset_str)?).map_err(CPUAffinityError::SetAffinity)
+}
+
+pub fn log_cpu_affinity() -> Result<()> {
+    let cpuset = sched_getaffinity(Pid::this()).map_err(CPUAffinityError::GetAffinity)?;
+    let mask = (0..usize::BITS as usize)
+        .filter(|&i| cpuset.is_set(i).unwrap_or(false))
+        .fold(0, |mask, i| mask | (1 << i));
+    tracing::debug!("affinity: 0x{:x}", mask);
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_to_cpuset_single_values() {
+        let cpuset = to_cpuset("0,1,2").unwrap();
+        for cpu in [0, 1, 2] {
+            assert!(cpuset.is_set(cpu).unwrap());
+        }
+    }
+
+    #[test]
+    fn test_to_cpuset_range() {
+        let cpuset = to_cpuset("3-5").unwrap();
+        for cpu in [3, 4, 5] {
+            assert!(cpuset.is_set(cpu).unwrap());
+        }
+    }
+
+    #[test]
+    fn test_to_cpuset_mixed() {
+        let cpuset = to_cpuset("0, 2-4, 6").unwrap();
+        for cpu in [0, 2, 3, 4, 6] {
+            assert!(cpuset.is_set(cpu).unwrap());
+        }
+        for cpu in [1, 5, 7] {
+            assert!(!cpuset.is_set(cpu).unwrap_or(false));
+        }
+    }
+
+    #[test]
+    fn test_to_cpuset_spaces_and_empty() {
+        let cpuset = to_cpuset("  , 1 , 3 , 5-7 , ").unwrap();
+        for cpu in [1, 3, 5, 6, 7] {
+            assert!(cpuset.is_set(cpu).unwrap());
+        }
+    }
+
+    #[test]
+    fn test_to_cpuset_invalid_range() {
+        let err = to_cpuset("5-3").unwrap_err();
+        matches!(err, CPUAffinityError::ParseError(_));
+    }
+
+    #[test]
+    fn test_to_cpuset_invalid_value() {
+        let err = to_cpuset("a,b,c").unwrap_err();
+        matches!(err, CPUAffinityError::ParseError(_));
+    }
+
+    #[test]
+    fn test_to_cpuset_max_allowed_cpu() {
+        let max = CpuSet::count();
+        let highest = max - 1;
+        let cpuset = to_cpuset(&highest.to_string()).unwrap();
+        assert!(cpuset.is_set(highest).unwrap());
+    }
+
+    #[test]
+    fn test_to_cpuset_exceeds_max_cpu() {
+        let max = CpuSet::count();
+        let result = to_cpuset(&max.to_string());
+        assert!(matches!(
+            result,
+            Err(CPUAffinityError::CpuOutOfRange { .. })
+        ));
+    }
+}

crates/libcontainer/src/process/cpu_affinity.rs

@@ -5,6 +5,7 @@ pub mod args;
 pub mod channel;
 pub mod container_intermediate_process;
 pub mod container_main_process;
+pub mod cpu_affinity;
 mod fork;
 pub mod init;
 pub mod intel_rdt;

crates/libcontainer/src/process/mod.rs

@@ -25,6 +25,7 @@ tempfile = "3"
 scopeguard = "1.2.0"
 tracing = { version = "0.1.41", features = ["attributes"]}
 tracing-subscriber = { version = "0.3.19", features = ["json", "env-filter"] }
+regex = "1"
 
 [dependencies.clap]
 version = "4.1.6"

tests/contest/contest/Cargo.toml

@@ -13,6 +13,7 @@ use crate::tests::delete::get_delete_test;
 use crate::tests::devices::get_devices_test;
 use crate::tests::domainname::get_domainname_tests;
 use crate::tests::example::get_example_test;
+use crate::tests::exec_cpu_affinity::get_exec_cpu_affinity_test;
 use crate::tests::fd_control::get_fd_control_test;
 use crate::tests::hooks::get_hooks_tests;
 use crate::tests::hostname::get_hostname_test;
@@ -142,6 +143,7 @@ fn main() -> Result<()> {
     let rootfs_propagation = get_rootfs_propagation_test();
     let process_capabilities_fail = get_process_capabilities_fail_test();
     let uid_mappings = get_uid_mappings_test();
+    let exec_cpu_affinity = get_exec_cpu_affinity_test();
 
     tm.add_test_group(Box::new(cl));
     tm.add_test_group(Box::new(cc));
@@ -180,6 +182,7 @@ fn main() -> Result<()> {
     tm.add_test_group(Box::new(rootfs_propagation));
     tm.add_test_group(Box::new(process_capabilities_fail));
     tm.add_test_group(Box::new(uid_mappings));
+    tm.add_test_group(Box::new(exec_cpu_affinity));
 
     tm.add_test_group(Box::new(io_priority_test));
     tm.add_cleanup(Box::new(cgroups::cleanup_v1));