1use std::str::FromStr;
2
3use ccp_shared::types::CPUIdType;
4use ccp_shared::types::LogicalCoreId;
5
6pub const LINUX_CPUSET_CPUS_PATH: &str = "/sys/fs/cgroup/cpuset.cpus";
7
8#[derive(Debug, thiserror::Error, PartialEq)]
9pub enum CpuListParserError {
10 #[error("invalid CPU list {list:?}: {error}")]
11 InvalidNumber {
12 list: String,
13 error: <CPUIdType as FromStr>::Err,
14 },
15}
16
17pub fn read_cpu_list_string() -> std::io::Result<String> {
18 std::fs::read_to_string(LINUX_CPUSET_CPUS_PATH)
19}
20
21pub fn parse_cpu_list(content: &str) -> Result<Vec<LogicalCoreId>, CpuListParserError> {
22 if content.is_empty() {
23 return Ok(vec![]);
24 }
25 let ranges = content
26 .split(',')
27 .map(|group| match group.split_once('-') {
28 Some((a, b)) => a
29 .trim()
30 .parse::<CPUIdType>()
31 .and_then(|a| b.trim().parse().map(move |b| (a..=b))),
32 None => group.trim().parse().map(|a| (a..=a)),
33 })
34 .collect::<Result<Vec<_>, _>>()
35 .map_err(|error| CpuListParserError::InvalidNumber {
36 list: content.to_owned(),
37 error,
38 })?;
39
40 let mut cpu_ids: Vec<_> = ranges
41 .into_iter()
42 .flatten()
43 .map(LogicalCoreId::new)
44 .collect();
45
46 cpu_ids.sort_unstable();
48 cpu_ids.dedup();
49
50 Ok(cpu_ids)
51}
52
53#[cfg(test)]
54mod tests {
55 use super::*;
56
57 #[test]
58 fn test_parser() {
59 let list = parse_cpu_list("4,0-2,6-8,1").unwrap();
60 assert_eq!(
61 list,
62 vec![
63 LogicalCoreId::new(0),
64 LogicalCoreId::new(1),
65 LogicalCoreId::new(2),
66 LogicalCoreId::new(4),
67 LogicalCoreId::new(6),
68 LogicalCoreId::new(7),
69 LogicalCoreId::new(8),
70 ]
71 );
72 }
73
74 #[test]
75 fn test_parser_spaces() {
76 let list = parse_cpu_list(" 4 , 0 - 2 , 6-8, 1 ").unwrap();
77 assert_eq!(
78 list,
79 vec![
80 LogicalCoreId::new(0),
81 LogicalCoreId::new(1),
82 LogicalCoreId::new(2),
83 LogicalCoreId::new(4),
84 LogicalCoreId::new(6),
85 LogicalCoreId::new(7),
86 LogicalCoreId::new(8),
87 ]
88 );
89 }
90
91 #[test]
92 fn test_empty() {
93 let list = parse_cpu_list("").unwrap();
94 assert!(list.is_empty());
95 }
96
97 #[test]
98 fn test_invalid_1() {
99 let content = "1-a";
100 let err = parse_cpu_list(content).unwrap_err();
101 let expected_err = "a".parse::<CPUIdType>().unwrap_err();
102 assert_eq!(
103 err,
104 CpuListParserError::InvalidNumber {
105 list: content.to_owned(),
106 error: expected_err
107 }
108 );
109 }
110
111 #[test]
112 fn test_invalid_empty_range_1() {
113 let content = "3,";
114 let err = parse_cpu_list(content).unwrap_err();
115 let expected_err = "".parse::<CPUIdType>().unwrap_err();
116 assert_eq!(
117 err,
118 CpuListParserError::InvalidNumber {
119 list: content.to_owned(),
120 error: expected_err
121 }
122 );
123 }
124
125 #[test]
126 fn test_invalid_empty_range_2() {
127 let content = ",3";
128 let err = parse_cpu_list(content).unwrap_err();
129 let expected_err = "".parse::<CPUIdType>().unwrap_err();
130 assert_eq!(
131 err,
132 CpuListParserError::InvalidNumber {
133 list: content.to_owned(),
134 error: expected_err
135 }
136 );
137 }
138}