shadow_rs/core/
cpu.rs

1pub struct RangeListIter<'a> {
2    current_range: Option<std::ops::RangeInclusive<u32>>,
3    remaining: &'a str,
4}
5
6impl Iterator for RangeListIter<'_> {
7    type Item = u32;
8
9    fn next(&mut self) -> Option<Self::Item> {
10        loop {
11            // if we have an active range to process
12            if let Some(current_range) = &mut self.current_range {
13                // get the next num in the range
14                let Some(rv) = current_range.next() else {
15                    // we've returned all numbers in the range, so clear the range and start over
16                    self.current_range = None;
17                    continue;
18                };
19
20                // return the next num in the range
21                break Some(rv);
22            // we don't have an active range, but there are more ranges to process
23            } else if !self.remaining.is_empty() {
24                let (next_range, remaining) = match self.remaining.split_once(',') {
25                    // there was at least one comma
26                    Some(x) => x,
27                    // there are no more commas
28                    None => (self.remaining, ""),
29                };
30
31                self.remaining = remaining;
32
33                if next_range.is_empty() {
34                    continue;
35                }
36
37                let mut split = next_range.split('-');
38                let start = split.next().unwrap();
39                let end = split.next();
40                assert!(split.next().is_none());
41
42                let start = start.parse().unwrap();
43                let end = end.map(|x| x.parse().unwrap()).unwrap_or(start);
44
45                self.current_range = Some(std::ops::RangeInclusive::new(start, end));
46
47                continue;
48            // no active range and no more ranges remaining
49            } else {
50                // the iterator is complete
51                break None;
52            }
53        }
54    }
55}
56
57/// Take an input of a list of ranges like '1-3,5,7-10' and return an iterator of integers like
58/// \[1,2,3,5,7,8,9,10\].
59///
60/// The returned iterator will panic if the input is not nicely formatted (no whitespace, etc) or
61/// contains invalid characters.
62///
63/// The iterator will return items in the order of the list, meaning that they are not guaranteed to
64/// be returned in increasing order and there may be duplicates. For example "1,2,3,3,2" would
65/// return items \[1, 2, 3, 3, 2\].
66pub fn parse_range_list(range_list: &str) -> RangeListIter<'_> {
67    RangeListIter {
68        current_range: None,
69        remaining: range_list,
70    }
71}
72
73/// Get the nodes from `/sys/devices/system/node/possible`.
74pub fn nodes() -> Vec<u32> {
75    let name = "/sys/devices/system/node/possible";
76    parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
77}
78
79/// Get the CPUs in a node from `/sys/devices/system/node/node{node}/cpulist`.
80pub fn cpus(node: u32) -> Vec<u32> {
81    let name = format!("/sys/devices/system/node/node{node}/cpulist");
82    parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
83}
84
85/// Get the core ID from `/sys/devices/system/cpu/cpu{cpu}/topology/core_id`.
86pub fn core(cpu: u32) -> u32 {
87    let name = format!("/sys/devices/system/cpu/cpu{cpu}/topology/core_id");
88    std::fs::read_to_string(name)
89        .unwrap()
90        .trim()
91        .parse()
92        .unwrap()
93}
94
95/// Get the online CPUs from `/sys/devices/system/cpu/online`.
96pub fn online() -> Vec<u32> {
97    let name = "/sys/devices/system/cpu/online";
98    parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
99}
100
101/// Count the number of physical cores available. Uses `sched_getaffinity` so should take into
102/// account CPU affinity and cgroups.
103pub fn count_physical_cores() -> u32 {
104    let affinity = nix::sched::sched_getaffinity(nix::unistd::Pid::from_raw(0)).unwrap();
105
106    let mut physical_cores = std::collections::HashSet::new();
107
108    for cpu in online() {
109        if affinity.is_set(cpu.try_into().unwrap()).unwrap() {
110            physical_cores.insert(core(cpu));
111        }
112    }
113
114    assert!(!physical_cores.is_empty());
115    physical_cores.len().try_into().unwrap()
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    fn check(list: &str, array: &[u32]) {
123        let list: Vec<_> = parse_range_list(list).collect();
124        assert_eq!(list, array);
125    }
126
127    #[test]
128    fn test_range_list() {
129        check("", &[]);
130        check("1", &[1]);
131        check("1,2", &[1, 2]);
132        check("1-2", &[1, 2]);
133        check("1-1", &[1]);
134        check("1,2,3", &[1, 2, 3]);
135        check("1-3", &[1, 2, 3]);
136        check("1,2-3,4", &[1, 2, 3, 4]);
137        check("1,2-4,5", &[1, 2, 3, 4, 5]);
138        check(
139            "0-5,7-9,13,15-19",
140            &[0, 1, 2, 3, 4, 5, 7, 8, 9, 13, 15, 16, 17, 18, 19],
141        );
142        check("1,,5", &[1, 5]);
143        check("1,1,5", &[1, 1, 5]);
144        check("1-1,5", &[1, 5]);
145        check("1-1,0,5", &[1, 0, 5]);
146        check("1-0", &[]);
147    }
148}
shadow_rs/core/cpu.rs

shadow_rs/core/
cpu.rs
