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}