compose_spec/service/
cpuset.rs

1//! Provides [`CpuSet`] for the `cpuset` field of [`Service`](super::Service).
2
3use std::{
4    collections::BTreeSet,
5    fmt::{self, Display, Formatter, Write},
6    num::ParseIntError,
7    str::FromStr,
8};
9
10use compose_spec_macros::{DeserializeFromStr, SerializeDisplay};
11use thiserror::Error;
12
13/// CPUs in which to allow execution.
14///
15/// [compose-spec](https://github.com/compose-spec/compose-spec/blob/master/05-services.md#cpuset)
16#[derive(SerializeDisplay, DeserializeFromStr, Debug, Default, Clone, PartialEq, Eq)]
17#[serde(expecting = "a comma-separated list (0,1), a range (0-3), or a combination (0-3,5,7-9)")]
18pub struct CpuSet(pub BTreeSet<u64>);
19
20impl CpuSet {
21    /// Returns `true` if the set is empty.
22    #[must_use]
23    pub fn is_empty(&self) -> bool {
24        self.0.is_empty()
25    }
26}
27
28impl Display for CpuSet {
29    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
30        let mut iter = self.0.iter();
31
32        let Some(cpu) = iter.next() else {
33            return f.write_str("");
34        };
35
36        let mut range = (*cpu, *cpu);
37
38        let mut first = true;
39        for cpu in iter {
40            let (start, end) = &mut range;
41            if *cpu == *end + 1 {
42                *end = *cpu;
43            } else {
44                write_range(f, first, *start, *end)?;
45                first = false;
46                range = (*cpu, *cpu);
47            }
48        }
49
50        let (start, end) = range;
51        write_range(f, first, start, end)
52    }
53}
54
55/// Write range to a [`Formatter`].
56fn write_range(f: &mut Formatter, first: bool, start: u64, end: u64) -> fmt::Result {
57    if !first {
58        f.write_char(',')?;
59    }
60
61    let mut buffer = itoa::Buffer::new();
62
63    f.write_str(buffer.format(start))?;
64
65    if start != end {
66        f.write_char('-')?;
67        f.write_str(buffer.format(end))?;
68    }
69
70    Ok(())
71}
72
73impl FromStr for CpuSet {
74    type Err = ParseCpuSetError;
75
76    fn from_str(s: &str) -> Result<Self, Self::Err> {
77        let mut inner = BTreeSet::new();
78
79        for range in s.split_terminator(',') {
80            if let Some((start, end)) = range.split_once('-') {
81                let start: u64 = start.parse().map_err(parse_int_err(start))?;
82                let end = end.parse().map_err(parse_int_err(end))?;
83                inner.extend(start..=end);
84            } else {
85                let cpu = range.parse().map_err(parse_int_err(range))?;
86                inner.insert(cpu);
87            }
88        }
89
90        Ok(Self(inner))
91    }
92}
93
94/// Closure which constructs a [`ParseCpuSetError`] from a [`ParseIntError`] and a `value`.
95fn parse_int_err(value: &str) -> impl FnOnce(ParseIntError) -> ParseCpuSetError {
96    let value = value.to_owned();
97    |source| ParseCpuSetError { value, source }
98}
99
100/// Error returned when parsing a [`CpuSet`] from a string.
101#[derive(Error, Debug, Clone, PartialEq, Eq)]
102#[error("could not parse `{value}` as an integer")]
103pub struct ParseCpuSetError {
104    /// Value attempted to parse.
105    value: String,
106    /// Parse error.
107    source: ParseIntError,
108}
109
110impl TryFrom<&str> for CpuSet {
111    type Error = ParseCpuSetError;
112
113    fn try_from(value: &str) -> Result<Self, Self::Error> {
114        value.parse()
115    }
116}
117
118impl From<BTreeSet<u64>> for CpuSet {
119    fn from(value: BTreeSet<u64>) -> Self {
120        Self(value)
121    }
122}
123
124impl From<CpuSet> for BTreeSet<u64> {
125    fn from(value: CpuSet) -> Self {
126        value.0
127    }
128}
129
130#[cfg(test)]
131#[allow(clippy::unwrap_used)]
132mod tests {
133    use proptest::{prop_assert_eq, proptest};
134
135    use super::*;
136
137    mod display {
138        use super::*;
139
140        #[test]
141        fn individual() {
142            let test = CpuSet(BTreeSet::from([1, 3, 5]));
143            assert_eq!(test.to_string(), "1,3,5");
144        }
145
146        #[test]
147        fn range() {
148            let test = CpuSet(BTreeSet::from([1, 2, 3]));
149            assert_eq!(test.to_string(), "1-3");
150        }
151
152        #[test]
153        fn combination() {
154            let test = CpuSet(BTreeSet::from([1, 2, 3, 5, 7, 8, 9]));
155            assert_eq!(test.to_string(), "1-3,5,7-9");
156        }
157    }
158
159    mod from_str {
160        use super::*;
161
162        #[test]
163        fn individual() {
164            let test = CpuSet(BTreeSet::from([1, 3, 5]));
165            assert_eq!(test, "1,3,5".parse().unwrap());
166        }
167
168        #[test]
169        fn range() {
170            let test = CpuSet(BTreeSet::from([1, 2, 3]));
171            assert_eq!(test, "1-3".parse().unwrap());
172        }
173
174        #[test]
175        fn combination() {
176            let test = CpuSet(BTreeSet::from([1, 2, 3, 5, 7, 8, 9]));
177            assert_eq!(test, "1-3,5,7-9".parse().unwrap());
178        }
179    }
180
181    proptest! {
182        #[test]
183        fn to_string_no_panic(set: BTreeSet<u64>) {
184            CpuSet(set).to_string();
185        }
186
187        #[test]
188        fn parse_no_panic(string: String) {
189            let _ = string.parse::<CpuSet>();
190        }
191
192        #[test]
193        fn round_trip(set: BTreeSet<u64>) {
194            let test = CpuSet(set);
195            let test2 = test.to_string().parse()?;
196            prop_assert_eq!(test, test2);
197        }
198    }
199}
compose_spec/service/cpuset.rs

compose_spec/service/
cpuset.rs
