cgroups_rs/systemd/
cpuset.rs

1// Copyright (c) 2025 Ant Group
2//
3// SPDX-License-Identifier: Apache-2.0 or MIT
4//
5
6use bit_vec::BitVec;
7
8use crate::systemd::error::{Error, Result};
9use crate::systemd::{ALLOWED_CPUS, ALLOWED_MEMORY_NODES, CPUSET_SYSTEMD_VERSION};
10
11const BYTE_IN_BITS: usize = 8;
12
13/// Returns the property for cpuset CPUs.
14pub fn cpus(cpus: &str, systemd_version: usize) -> Result<(&'static str, Vec<u8>)> {
15    if systemd_version < CPUSET_SYSTEMD_VERSION {
16        return Err(Error::ObsoleteSystemd);
17    }
18
19    let mask = convert_list_to_mask(cpus)?;
20
21    Ok((ALLOWED_CPUS, mask))
22}
23
24/// Returns the property for cpuset memory nodes.
25pub fn mems(mems: &str, systemd_version: usize) -> Result<(&'static str, Vec<u8>)> {
26    if systemd_version < CPUSET_SYSTEMD_VERSION {
27        return Err(Error::ObsoleteSystemd);
28    }
29
30    let mask = convert_list_to_mask(mems)?;
31
32    Ok((ALLOWED_MEMORY_NODES, mask))
33}
34
35/// Convert cpuset cpus/mems from the string in comma-separated list format
36/// to bitmask restored in `Vec<u8>`, see [1].
37///
38/// 1: https://man7.org/linux/man-pages/man7/cpuset.7.html
39///
40/// # Arguments
41///
42/// * `list` - A string slice that holds the list of CPUs in the format
43///   "0-3,5,7".
44fn convert_list_to_mask(list: &str) -> Result<Vec<u8>> {
45    let mut bit_vec = BitVec::from_elem(8, false);
46
47    let local_idx =
48        |index: usize| -> usize { index / BYTE_IN_BITS * BYTE_IN_BITS + 7 - index % BYTE_IN_BITS };
49
50    for part1 in list.split(',') {
51        let range: Vec<&str> = part1.split('-').collect();
52        match range.len() {
53            // x-
54            1 => {
55                let left: usize = range[0].parse().map_err(|_| Error::InvalidArgument)?;
56
57                while left >= bit_vec.len() {
58                    bit_vec.grow(BYTE_IN_BITS, false);
59                }
60                bit_vec.set(local_idx(left), true);
61            }
62            // x-y
63            2 => {
64                let left: usize = range[0].parse().map_err(|_| Error::InvalidArgument)?;
65                let right: usize = range[1].parse().map_err(|_| Error::InvalidArgument)?;
66
67                while right >= bit_vec.len() {
68                    bit_vec.grow(BYTE_IN_BITS, false);
69                }
70
71                for index in left..=right {
72                    bit_vec.set(local_idx(index), true);
73                }
74            }
75            _ => {
76                return Err(Error::InvalidArgument);
77            }
78        }
79    }
80
81    let mut mask = bit_vec.to_bytes();
82    mask.reverse();
83
84    Ok(mask)
85}
86
87#[cfg(test)]
88mod tests {
89    use crate::systemd::cpuset::convert_list_to_mask;
90
91    #[test]
92    fn test_convert_list_to_mask() {
93        let mask = convert_list_to_mask("2-4").unwrap();
94        assert_eq!(vec![0b00011100_u8], mask);
95
96        let mask = convert_list_to_mask("1,7").unwrap();
97        assert_eq!(vec![0b10000010_u8], mask);
98
99        let mask = convert_list_to_mask("0-4,9").unwrap();
100        assert_eq!(vec![0b00000010_u8, 0b00011111_u8], mask);
101
102        assert!(convert_list_to_mask("1-3-4").is_err());
103
104        assert!(convert_list_to_mask("1-3,,").is_err());
105    }
106}
cgroups_rs/systemd/cpuset.rs

cgroups_rs/systemd/
cpuset.rs
