libcontainer/process/
cpu_affinity.rs

1use nix::sched::{CpuSet, sched_getaffinity, sched_setaffinity};
2use nix::unistd::Pid;
3use tracing::{Level, enabled};
4
5#[derive(Debug, thiserror::Error)]
6pub enum CPUAffinityError {
7    #[error("invalid CPU string: {0}")]
8    ParseError(String),
9    #[error("values larger than {max} are not supported")]
10    CpuOutOfRange { cpu: usize, max: usize },
11    #[error("failed to set CPU for CPU {cpu}: {source}")]
12    CpuSet {
13        cpu: usize,
14        #[source]
15        source: nix::Error,
16    },
17    #[error("failed to setaffinity")]
18    SetAffinity(#[source] nix::Error),
19    #[error("failed to getaffinity")]
20    GetAffinity(#[source] nix::Error),
21}
22
23type Result<T> = std::result::Result<T, CPUAffinityError>;
24
25pub fn to_cpuset(cpuset_str: &str) -> Result<CpuSet> {
26    let mut cpuset = CpuSet::new();
27    let max_cpu = CpuSet::count();
28
29    for part in cpuset_str
30        .trim()
31        .split(',')
32        .map(str::trim)
33        .filter(|s| !s.is_empty())
34    {
35        match part.split_once('-') {
36            Some((start_str, end_str)) => {
37                let start = parse_cpu_index(start_str, max_cpu)?;
38                let end = parse_cpu_index(end_str, max_cpu)?;
39                if start > end {
40                    return Err(CPUAffinityError::ParseError(format!(
41                        "invalid range: {}-{}",
42                        start, end
43                    )));
44                }
45                for cpu in start..=end {
46                    cpuset
47                        .set(cpu)
48                        .map_err(|e| CPUAffinityError::CpuSet { cpu, source: e })?;
49                }
50            }
51            None => {
52                let cpu = parse_cpu_index(part, max_cpu)?;
53                cpuset
54                    .set(cpu)
55                    .map_err(|e| CPUAffinityError::CpuSet { cpu, source: e })?;
56            }
57        }
58    }
59    Ok(cpuset)
60}
61
62fn parse_cpu_index(s: &str, max_cpu: usize) -> Result<usize> {
63    let cpu: usize = s
64        .parse()
65        .map_err(|_| CPUAffinityError::ParseError(s.to_string()))?;
66    if cpu >= max_cpu {
67        return Err(CPUAffinityError::CpuOutOfRange {
68            cpu,
69            max: max_cpu - 1,
70        });
71    }
72    Ok(cpu)
73}
74
75pub fn set_cpuset_affinity_from_string(pid: Pid, cpuset_str: &str) -> Result<()> {
76    tracing::debug!(?cpuset_str, "setting CPU affinity for tenant container");
77    sched_setaffinity(pid, &to_cpuset(cpuset_str)?).map_err(CPUAffinityError::SetAffinity)
78}
79
80// Logs a compact CPU affinity bitmask similar to runc's nsexec.c (see: https://github.com/opencontainers/runc/blob/main/libcontainer/nsenter/nsexec.c#L676).
81// This helps in debugging which CPUs the current process is allowed to run on.
82// Only logs when DEBUG level is enabled.
83pub fn log_cpu_affinity() -> Result<()> {
84    if !enabled!(Level::DEBUG) {
85        return Ok(());
86    }
87    let cpuset = sched_getaffinity(Pid::this()).map_err(CPUAffinityError::GetAffinity)?;
88    let mask = (0..usize::BITS as usize)
89        .filter(|&i| cpuset.is_set(i).unwrap_or(false))
90        .fold(0usize, |mask, i| mask | (1usize << i));
91    tracing::debug!("affinity: 0x{:x}", mask);
92    Ok(())
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_to_cpuset_single_values() {
101        let cpuset = to_cpuset("0,1,2").unwrap();
102        for cpu in [0, 1, 2] {
103            assert!(cpuset.is_set(cpu).unwrap());
104        }
105    }
106
107    #[test]
108    fn test_to_cpuset_range() {
109        let cpuset = to_cpuset("3-5").unwrap();
110        for cpu in [3, 4, 5] {
111            assert!(cpuset.is_set(cpu).unwrap());
112        }
113    }
114
115    #[test]
116    fn test_to_cpuset_mixed() {
117        let cpuset = to_cpuset("0, 2-4, 6").unwrap();
118        for cpu in [0, 2, 3, 4, 6] {
119            assert!(cpuset.is_set(cpu).unwrap());
120        }
121        for cpu in [1, 5, 7] {
122            assert!(!cpuset.is_set(cpu).unwrap_or(false));
123        }
124    }
125
126    #[test]
127    fn test_to_cpuset_spaces_and_empty() {
128        let cpuset = to_cpuset("  , 1 , 3 , 5-7 , ").unwrap();
129        for cpu in [1, 3, 5, 6, 7] {
130            assert!(cpuset.is_set(cpu).unwrap());
131        }
132    }
133
134    #[test]
135    fn test_to_cpuset_invalid_range() {
136        let err = to_cpuset("5-3").unwrap_err();
137        matches!(err, CPUAffinityError::ParseError(_));
138    }
139
140    #[test]
141    fn test_to_cpuset_invalid_value() {
142        let err = to_cpuset("a,b,c").unwrap_err();
143        matches!(err, CPUAffinityError::ParseError(_));
144    }
145
146    #[test]
147    fn test_to_cpuset_max_allowed_cpu() {
148        let max = CpuSet::count();
149        let highest = max - 1;
150        let cpuset = to_cpuset(&highest.to_string()).unwrap();
151        assert!(cpuset.is_set(highest).unwrap());
152    }
153
154    #[test]
155    fn test_to_cpuset_exceeds_max_cpu() {
156        let max = CpuSet::count();
157        let result = to_cpuset(&max.to_string());
158        assert!(matches!(
159            result,
160            Err(CPUAffinityError::CpuOutOfRange { .. })
161        ));
162    }
163}
libcontainer/process/cpu_affinity.rs

libcontainer/process/
cpu_affinity.rs
