rlmctl_common/
limit.rs

1use crate::{Error, Result};
2use serde::{Deserialize, Serialize};
3
4/// A memory limit below this is rejected: a process cannot start, let alone
5/// make progress, in less than 8 MiB.
6pub const MIN_MEMORY_BYTES: u64 = 8 * 1024 * 1024;
7
8/// An I/O bandwidth limit below this is rejected: below 64 KiB/s a process
9/// cannot make meaningful progress.
10pub const MIN_IO_BPS: u64 = 64 * 1024;
11
12/// Parse a byte size: optional decimal fraction, optional unit K/M/G/T with
13/// an optional "B" or "iB" suffix, case-insensitive, all binary multiples.
14/// A bare number is bytes.
15pub fn parse_size(input: &str) -> Result<u64> {
16    let s = input.trim();
17    let bad = || Error::InvalidMemory(s.to_string());
18    let split = s
19        .find(|c: char| !(c.is_ascii_digit() || c == '.'))
20        .unwrap_or(s.len());
21    let (num, unit) = s.split_at(split);
22    let mult: u64 = match unit.trim().to_ascii_lowercase().as_str() {
23        "" | "b" => 1,
24        "k" | "kb" | "kib" => 1 << 10,
25        "m" | "mb" | "mib" => 1 << 20,
26        "g" | "gb" | "gib" => 1 << 30,
27        "t" | "tb" | "tib" => 1 << 40,
28        _ => return Err(bad()),
29    };
30    if num.is_empty() || num.starts_with('.') || num.ends_with('.') || num.matches('.').count() > 1
31    {
32        return Err(bad());
33    }
34    let overflow = || Error::InvalidMemory("value too large (overflow)".into());
35    let bytes = match num.split_once('.') {
36        None => num
37            .parse::<u64>()
38            .map_err(|_| bad())?
39            .checked_mul(mult)
40            .ok_or_else(overflow)?,
41        Some((whole, frac)) => {
42            let whole: u64 = whole.parse().map_err(|_| bad())?;
43            let digits = frac.len().min(9);
44            let frac_val: u128 = frac[..digits].parse().map_err(|_| bad())?;
45            let frac_bytes = (frac_val * u128::from(mult) / 10u128.pow(digits as u32)) as u64;
46            whole
47                .checked_mul(mult)
48                .and_then(|w| w.checked_add(frac_bytes))
49                .ok_or_else(overflow)?
50        }
51    };
52    if bytes == 0 {
53        return Err(Error::InvalidMemory("value cannot be zero".into()));
54    }
55    Ok(bytes)
56}
57
58/// Resource limits to apply to a process
59#[derive(Debug, Clone, Default, Serialize, Deserialize)]
60pub struct Limit {
61    pub memory: Option<MemoryLimit>,
62    pub cpu: Option<CpuLimit>,
63    pub io: Option<IoLimit>,
64}
65
66impl Limit {
67    /// True if none of memory, cpu or io carry a value.
68    pub fn is_empty(&self) -> bool {
69        self.memory.is_none() && self.cpu.is_none() && self.io.is_none_or(|io| io.is_empty())
70    }
71
72    /// Overlay explicit `over` values on top of `self` (e.g. a profile),
73    /// field by field; io read/write bandwidth overlay independently.
74    pub fn overlay(self, over: &Limit) -> Limit {
75        let io = match (self.io, over.io) {
76            (Some(a), Some(b)) => Some(IoLimit {
77                read_bps: b.read_bps.or(a.read_bps),
78                write_bps: b.write_bps.or(a.write_bps),
79            }),
80            (a, b) => b.or(a),
81        };
82        Limit {
83            memory: over.memory.or(self.memory),
84            cpu: over.cpu.or(self.cpu),
85            io,
86        }
87    }
88}
89
90/// I/O bandwidth limit in bytes per second
91#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
92pub struct IoLimit {
93    /// Read bandwidth limit (bytes/sec)
94    pub read_bps: Option<u64>,
95    /// Write bandwidth limit (bytes/sec)
96    pub write_bps: Option<u64>,
97}
98
99impl IoLimit {
100    pub fn parse_bps(s: &str) -> Result<u64> {
101        let bytes = parse_size(s)?;
102        if bytes < MIN_IO_BPS {
103            return Err(Error::InvalidMemory(format!(
104                "{} per second is below the 64K minimum",
105                s.trim()
106            )));
107        }
108        Ok(bytes)
109    }
110
111    pub fn is_empty(&self) -> bool {
112        self.read_bps.is_none() && self.write_bps.is_none()
113    }
114}
115
116/// Memory limit in bytes
117#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
118pub struct MemoryLimit(u64);
119
120impl MemoryLimit {
121    pub fn bytes(self) -> u64 {
122        self.0
123    }
124
125    /// Parse human-readable memory string (e.g., "2G", "512M", "1.5GiB").
126    /// Rejects values below the 8 MiB floor: a process cannot run in less.
127    pub fn parse(s: &str) -> Result<Self> {
128        let bytes = parse_size(s)?;
129        if bytes < MIN_MEMORY_BYTES {
130            return Err(Error::InvalidMemory(format!(
131                "{} is below the 8M minimum; a process cannot run in less",
132                s.trim()
133            )));
134        }
135        Ok(Self(bytes))
136    }
137}
138
139/// CPU limit as percentage (0-100 per core, can exceed 100 for multiple cores)
140#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
141pub struct CpuLimit(u32);
142
143impl CpuLimit {
144    pub fn percent(self) -> u32 {
145        self.0
146    }
147
148    /// Parse CPU percentage string (e.g., "50%", "150%")
149    /// Maximum is 10000% (100 cores)
150    pub fn parse(s: &str) -> Result<Self> {
151        let s = s.trim().trim_end_matches('%');
152        let percent: u32 = s.parse().map_err(|_| Error::InvalidCpu(s.into()))?;
153        if percent == 0 {
154            return Err(Error::InvalidCpu("value cannot be zero".into()));
155        }
156        if percent > 10000 {
157            return Err(Error::InvalidCpu(
158                "value too large (max 10000% = 100 cores)".into(),
159            ));
160        }
161        Ok(Self(percent))
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn parse_size_table() {
171        let mib = 1024 * 1024u64;
172        let ok: &[(&str, u64)] = &[
173            ("4096", 4096),
174            ("1K", 1024),
175            ("512M", 512 * mib),
176            ("512MB", 512 * mib),
177            ("512 MB", 512 * mib),
178            ("512MiB", 512 * mib),
179            ("2g", 2048 * mib),
180            ("2GiB", 2048 * mib),
181            ("1.5G", 1536 * mib),
182            ("0.5g", 512 * mib),
183            ("  8m  ", 8 * mib),
184            ("1T", 1024 * 1024 * mib),
185        ];
186        for (s, want) in ok {
187            assert_eq!(parse_size(s).unwrap(), *want, "{s}");
188        }
189        for s in [
190            "",
191            "abc",
192            "-1G",
193            "1e3M",
194            "1..5G",
195            ".5G",
196            "5.G",
197            "0",
198            "0M",
199            "12X",
200            "999999999999999999T",
201        ] {
202            assert!(parse_size(s).is_err(), "{s:?} should be rejected");
203        }
204    }
205
206    #[test]
207    fn memory_limits_have_a_floor() {
208        assert!(MemoryLimit::parse("1K").is_err());
209        assert!(MemoryLimit::parse("4M").is_err());
210        assert_eq!(MemoryLimit::parse("8M").unwrap().bytes(), MIN_MEMORY_BYTES);
211        let e = MemoryLimit::parse("1024").unwrap_err().to_string();
212        assert!(e.contains("8M"), "{e}");
213    }
214
215    #[test]
216    fn io_limits_have_a_floor() {
217        assert!(IoLimit::parse_bps("1K").is_err());
218        assert_eq!(IoLimit::parse_bps("64K").unwrap(), MIN_IO_BPS);
219    }
220
221    #[test]
222    fn explicit_values_override_profile_values() {
223        let profile = crate::build_limit(Some("512M"), Some("25%"), Some("10M"), None).unwrap();
224        let flags = crate::build_limit(Some("1G"), None, None, Some("5M")).unwrap();
225        let l = profile.overlay(&flags);
226        assert_eq!(l.memory.unwrap().bytes(), 1024 * 1024 * 1024);
227        assert_eq!(l.cpu.unwrap().percent(), 25);
228        let io = l.io.unwrap();
229        assert_eq!(
230            (io.read_bps, io.write_bps),
231            (Some(10 * 1024 * 1024), Some(5 * 1024 * 1024))
232        );
233    }
234
235    #[test]
236    fn parse_memory_overflow() {
237        // Value too large for u64
238        assert!(MemoryLimit::parse("999999999999999999T").is_err());
239    }
240
241    #[test]
242    fn parse_cpu_percent() {
243        assert_eq!(CpuLimit::parse("50%").unwrap().percent(), 50);
244        assert_eq!(CpuLimit::parse("150").unwrap().percent(), 150);
245        assert_eq!(CpuLimit::parse("  75%  ").unwrap().percent(), 75);
246    }
247
248    #[test]
249    fn parse_cpu_errors() {
250        assert!(CpuLimit::parse("abc").is_err());
251        assert!(CpuLimit::parse("-50%").is_err());
252    }
253
254    #[test]
255    fn io_limit_is_empty() {
256        let empty = IoLimit::default();
257        assert!(empty.is_empty());
258
259        let with_read = IoLimit {
260            read_bps: Some(1000),
261            write_bps: None,
262        };
263        assert!(!with_read.is_empty());
264
265        let with_write = IoLimit {
266            read_bps: None,
267            write_bps: Some(1000),
268        };
269        assert!(!with_write.is_empty());
270    }
271
272    #[test]
273    fn parse_io_bps() {
274        assert_eq!(IoLimit::parse_bps("100M").unwrap(), 100 * 1024 * 1024);
275        assert_eq!(IoLimit::parse_bps("1G").unwrap(), 1024 * 1024 * 1024);
276    }
277}
rlmctl_common/limit.rs

rlmctl_common/
limit.rs
