nucleus/resources/
limits.rs

1use crate::error::{NucleusError, Result};
2
3/// Per-device I/O throttling limit
4#[derive(Debug, Clone)]
5pub struct IoDeviceLimit {
6    /// Device identifier in "major:minor" format
7    pub device: String,
8    /// Read IOPS limit
9    pub riops: Option<u64>,
10    /// Write IOPS limit
11    pub wiops: Option<u64>,
12    /// Read bytes/sec limit
13    pub rbps: Option<u64>,
14    /// Write bytes/sec limit
15    pub wbps: Option<u64>,
16}
17
18impl IoDeviceLimit {
19    /// Parse an I/O device limit spec like "8:0 riops=1000 wbps=10485760"
20    pub fn parse(s: &str) -> Result<Self> {
21        let mut parts = s.split_whitespace();
22
23        let device = parts
24            .next()
25            .ok_or_else(|| NucleusError::InvalidResourceLimit("Empty I/O limit spec".into()))?;
26
27        // Validate device format: "major:minor"
28        let mut dev_parts = device.split(':');
29        let major = dev_parts.next().and_then(|s| s.parse::<u64>().ok());
30        let minor = dev_parts.next().and_then(|s| s.parse::<u64>().ok());
31        if major.is_none() || minor.is_none() || dev_parts.next().is_some() {
32            return Err(NucleusError::InvalidResourceLimit(format!(
33                "Invalid device format '{}', expected 'major:minor'",
34                device
35            )));
36        }
37
38        let mut limit = Self {
39            device: device.to_string(),
40            riops: None,
41            wiops: None,
42            rbps: None,
43            wbps: None,
44        };
45
46        for param in parts {
47            let (key, value) = param.split_once('=').ok_or_else(|| {
48                NucleusError::InvalidResourceLimit(format!(
49                    "Invalid I/O param '{}', expected key=value",
50                    param
51                ))
52            })?;
53            let value: u64 = value.parse().map_err(|_| {
54                NucleusError::InvalidResourceLimit(format!("Invalid I/O value: {}", value))
55            })?;
56
57            match key {
58                "riops" => limit.riops = Some(value),
59                "wiops" => limit.wiops = Some(value),
60                "rbps" => limit.rbps = Some(value),
61                "wbps" => limit.wbps = Some(value),
62                _ => {
63                    return Err(NucleusError::InvalidResourceLimit(format!(
64                        "Unknown I/O param '{}'",
65                        key
66                    )));
67                }
68            }
69        }
70
71        Ok(limit)
72    }
73
74    /// Format as cgroup v2 io.max line: "major:minor riops=X wiops=Y rbps=Z wbps=W"
75    pub fn to_io_max_line(&self) -> String {
76        let mut parts = vec![self.device.clone()];
77        if let Some(v) = self.riops {
78            parts.push(format!("riops={}", v));
79        }
80        if let Some(v) = self.wiops {
81            parts.push(format!("wiops={}", v));
82        }
83        if let Some(v) = self.rbps {
84            parts.push(format!("rbps={}", v));
85        }
86        if let Some(v) = self.wbps {
87            parts.push(format!("wbps={}", v));
88        }
89        parts.join(" ")
90    }
91}
92
93/// Resource limits configuration
94#[derive(Debug, Clone)]
95pub struct ResourceLimits {
96    /// Memory limit in bytes (None = unlimited)
97    pub memory_bytes: Option<u64>,
98    /// Memory soft limit in bytes (auto-set to 90% of memory_bytes)
99    pub memory_high: Option<u64>,
100    /// Swap limit in bytes (Some(0) = disable swap)
101    pub memory_swap_max: Option<u64>,
102    /// CPU quota in microseconds per period
103    pub cpu_quota_us: Option<u64>,
104    /// CPU period in microseconds (default: 100000 = 100ms)
105    pub cpu_period_us: u64,
106    /// CPU scheduling weight (1-10000)
107    pub cpu_weight: Option<u64>,
108    /// Maximum number of PIDs (None = unlimited)
109    pub pids_max: Option<u64>,
110    /// Per-device I/O limits
111    pub io_limits: Vec<IoDeviceLimit>,
112}
113
114impl ResourceLimits {
115    /// Create unlimited resource limits
116    pub fn unlimited() -> Self {
117        Self {
118            memory_bytes: None,
119            memory_high: None,
120            memory_swap_max: None,
121            cpu_quota_us: None,
122            cpu_period_us: 100_000, // 100ms default period
123            cpu_weight: None,
124            pids_max: None,
125            io_limits: Vec::new(),
126        }
127    }
128
129    /// Parse memory limit from string (e.g., "512M", "1G")
130    pub fn parse_memory(s: &str) -> Result<u64> {
131        let s = s.trim();
132        if s.is_empty() {
133            return Err(NucleusError::InvalidResourceLimit(
134                "Empty memory limit".to_string(),
135            ));
136        }
137
138        let (num_str, multiplier) = if s.ends_with('K') || s.ends_with('k') {
139            (&s[..s.len() - 1], 1024u64)
140        } else if s.ends_with('M') || s.ends_with('m') {
141            (&s[..s.len() - 1], 1024 * 1024)
142        } else if s.ends_with('G') || s.ends_with('g') {
143            (&s[..s.len() - 1], 1024 * 1024 * 1024)
144        } else if s.ends_with('T') || s.ends_with('t') {
145            (&s[..s.len() - 1], 1024 * 1024 * 1024 * 1024)
146        } else {
147            // No suffix, assume bytes
148            (s, 1)
149        };
150
151        let num: u64 = num_str.parse().map_err(|_| {
152            NucleusError::InvalidResourceLimit(format!("Invalid memory value: {}", s))
153        })?;
154
155        num.checked_mul(multiplier).ok_or_else(|| {
156            NucleusError::InvalidResourceLimit(format!("Memory value overflows u64: {}", s))
157        })
158    }
159
160    /// Set memory limit from string (e.g., "512M", "1G")
161    ///
162    /// Automatically sets memory_high to 90% of the hard limit and
163    /// disables swap (memory_swap_max = 0) unless swap was explicitly enabled.
164    pub fn with_memory(mut self, limit: &str) -> Result<Self> {
165        let bytes = Self::parse_memory(limit)?;
166        self.memory_bytes = Some(bytes);
167        // Auto-set soft limit to 90% of hard limit (per spec)
168        self.memory_high = Some(bytes - bytes / 10);
169        // Disable swap by default when memory limit is set
170        if self.memory_swap_max.is_none() {
171            self.memory_swap_max = Some(0);
172        }
173        Ok(self)
174    }
175
176    /// Enable swap (removes the default swap=0 restriction)
177    pub fn with_swap_enabled(mut self) -> Self {
178        self.memory_swap_max = None;
179        self
180    }
181
182    /// Set CPU limit in cores (e.g., 2.5 cores)
183    pub fn with_cpu_cores(mut self, cores: f64) -> Result<Self> {
184        const MAX_CPU_CORES: f64 = 65_536.0;
185
186        if cores <= 0.0 || cores.is_nan() || cores.is_infinite() {
187            return Err(NucleusError::InvalidResourceLimit(
188                "CPU cores must be a finite positive number".to_string(),
189            ));
190        }
191        if cores > MAX_CPU_CORES {
192            return Err(NucleusError::InvalidResourceLimit(format!(
193                "CPU cores must be <= {}",
194                MAX_CPU_CORES
195            )));
196        }
197        // Convert cores to quota: cores * period
198        let quota = (cores * self.cpu_period_us as f64) as u64;
199        self.cpu_quota_us = Some(quota);
200        Ok(self)
201    }
202
203    /// Set maximum number of PIDs
204    pub fn with_pids(mut self, max_pids: u64) -> Result<Self> {
205        if max_pids == 0 {
206            return Err(NucleusError::InvalidResourceLimit(
207                "Max PIDs must be positive".to_string(),
208            ));
209        }
210        self.pids_max = Some(max_pids);
211        Ok(self)
212    }
213
214    /// Set CPU scheduling weight (1-10000)
215    pub fn with_cpu_weight(mut self, weight: u64) -> Result<Self> {
216        if !(1..=10000).contains(&weight) {
217            return Err(NucleusError::InvalidResourceLimit(
218                "CPU weight must be between 1 and 10000".to_string(),
219            ));
220        }
221        self.cpu_weight = Some(weight);
222        Ok(self)
223    }
224
225    /// Add an I/O device limit
226    pub fn with_io_limit(mut self, limit: IoDeviceLimit) -> Self {
227        self.io_limits.push(limit);
228        self
229    }
230}
231
232impl Default for ResourceLimits {
233    fn default() -> Self {
234        Self {
235            pids_max: Some(512),
236            ..Self::unlimited()
237        }
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn test_parse_memory() {
247        assert_eq!(ResourceLimits::parse_memory("1024").unwrap(), 1024);
248        assert_eq!(ResourceLimits::parse_memory("512K").unwrap(), 512 * 1024);
249        assert_eq!(
250            ResourceLimits::parse_memory("512M").unwrap(),
251            512 * 1024 * 1024
252        );
253        assert_eq!(
254            ResourceLimits::parse_memory("2G").unwrap(),
255            2 * 1024 * 1024 * 1024
256        );
257    }
258
259    #[test]
260    fn test_parse_memory_invalid() {
261        assert!(ResourceLimits::parse_memory("").is_err());
262        assert!(ResourceLimits::parse_memory("abc").is_err());
263        assert!(ResourceLimits::parse_memory("M").is_err());
264    }
265
266    #[test]
267    fn test_parse_memory_overflow_rejected() {
268        // 18446744073709551615T would overflow u64
269        assert!(ResourceLimits::parse_memory("99999999999999T").is_err());
270        // Just under u64::MAX in bytes should work
271        assert!(ResourceLimits::parse_memory("16383P").is_err()); // not a valid suffix, treated as bytes
272    }
273
274    #[test]
275    fn test_with_cpu_cores() {
276        let limits = ResourceLimits::unlimited();
277        let limits = limits.with_cpu_cores(2.0).unwrap();
278        assert_eq!(limits.cpu_quota_us, Some(200_000)); // 2.0 * 100_000
279    }
280
281    #[test]
282    fn test_with_cpu_cores_fractional() {
283        let limits = ResourceLimits::unlimited();
284        let limits = limits.with_cpu_cores(0.5).unwrap();
285        assert_eq!(limits.cpu_quota_us, Some(50_000)); // 0.5 * 100_000
286    }
287
288    #[test]
289    fn test_with_cpu_cores_invalid() {
290        let limits = ResourceLimits::unlimited();
291        assert!(limits.with_cpu_cores(0.0).is_err());
292        assert!(ResourceLimits::unlimited().with_cpu_cores(-1.0).is_err());
293    }
294
295    #[test]
296    fn test_with_memory_auto_sets_memory_high() {
297        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
298        let expected_bytes = 1024 * 1024 * 1024u64;
299        assert_eq!(limits.memory_bytes, Some(expected_bytes));
300        // memory_high should be 90% of hard limit
301        assert_eq!(limits.memory_high, Some(expected_bytes - expected_bytes / 10));
302    }
303
304    #[test]
305    fn test_with_memory_disables_swap_by_default() {
306        let limits = ResourceLimits::unlimited().with_memory("512M").unwrap();
307        assert_eq!(limits.memory_swap_max, Some(0));
308    }
309
310    #[test]
311    fn test_swap_enabled_clears_swap_limit() {
312        let limits = ResourceLimits::unlimited()
313            .with_memory("512M")
314            .unwrap()
315            .with_swap_enabled();
316        assert!(limits.memory_swap_max.is_none());
317    }
318
319    #[test]
320    fn test_with_cpu_weight_valid() {
321        let limits = ResourceLimits::unlimited().with_cpu_weight(100).unwrap();
322        assert_eq!(limits.cpu_weight, Some(100));
323
324        let limits = ResourceLimits::unlimited().with_cpu_weight(1).unwrap();
325        assert_eq!(limits.cpu_weight, Some(1));
326
327        let limits = ResourceLimits::unlimited().with_cpu_weight(10000).unwrap();
328        assert_eq!(limits.cpu_weight, Some(10000));
329    }
330
331    #[test]
332    fn test_with_cpu_weight_invalid() {
333        assert!(ResourceLimits::unlimited().with_cpu_weight(0).is_err());
334        assert!(ResourceLimits::unlimited().with_cpu_weight(10001).is_err());
335    }
336
337    #[test]
338    fn test_io_device_limit_parse_valid() {
339        let limit = IoDeviceLimit::parse("8:0 riops=1000 wbps=10485760").unwrap();
340        assert_eq!(limit.device, "8:0");
341        assert_eq!(limit.riops, Some(1000));
342        assert_eq!(limit.wbps, Some(10485760));
343        assert!(limit.wiops.is_none());
344        assert!(limit.rbps.is_none());
345    }
346
347    #[test]
348    fn test_io_device_limit_parse_all_params() {
349        let limit = IoDeviceLimit::parse("8:0 riops=100 wiops=200 rbps=300 wbps=400").unwrap();
350        assert_eq!(limit.riops, Some(100));
351        assert_eq!(limit.wiops, Some(200));
352        assert_eq!(limit.rbps, Some(300));
353        assert_eq!(limit.wbps, Some(400));
354    }
355
356    #[test]
357    fn test_io_device_limit_parse_invalid() {
358        // Empty string
359        assert!(IoDeviceLimit::parse("").is_err());
360        // Bad device format
361        assert!(IoDeviceLimit::parse("bad").is_err());
362        assert!(IoDeviceLimit::parse("8:0:1").is_err());
363        // Bad param format
364        assert!(IoDeviceLimit::parse("8:0 riops").is_err());
365        // Unknown param
366        assert!(IoDeviceLimit::parse("8:0 foo=100").is_err());
367        // Bad value
368        assert!(IoDeviceLimit::parse("8:0 riops=abc").is_err());
369    }
370
371    #[test]
372    fn test_io_device_limit_to_io_max_line() {
373        let limit = IoDeviceLimit {
374            device: "8:0".to_string(),
375            riops: Some(1000),
376            wiops: None,
377            rbps: None,
378            wbps: Some(10485760),
379        };
380        assert_eq!(limit.to_io_max_line(), "8:0 riops=1000 wbps=10485760");
381    }
382
383    #[test]
384    fn test_unlimited_defaults() {
385        let limits = ResourceLimits::unlimited();
386        assert!(limits.memory_bytes.is_none());
387        assert!(limits.memory_high.is_none());
388        assert!(limits.memory_swap_max.is_none());
389        assert!(limits.cpu_quota_us.is_none());
390        assert!(limits.cpu_weight.is_none());
391        assert!(limits.pids_max.is_none());
392        assert!(limits.io_limits.is_empty());
393    }
394
395    #[test]
396    fn test_memory_high_uses_integer_arithmetic() {
397        // BUG-13: memory_high must use integer arithmetic, not floating point
398        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
399        let bytes = 1024u64 * 1024 * 1024;
400        let expected_high = bytes - bytes / 10; // 90% via integer
401        assert_eq!(
402            limits.memory_high,
403            Some(expected_high),
404            "memory_high must be exactly bytes - bytes/10 (integer arithmetic)"
405        );
406    }
407
408    #[test]
409    fn test_cpu_cores_rejects_extreme_values() {
410        // BUG-12: Extreme CPU core values must be rejected, not silently overflow
411        assert!(ResourceLimits::unlimited().with_cpu_cores(f64::NAN).is_err());
412        assert!(ResourceLimits::unlimited().with_cpu_cores(f64::INFINITY).is_err());
413        assert!(ResourceLimits::unlimited().with_cpu_cores(100_000.0).is_err(),
414            "CPU cores > 65536 must be rejected to prevent quota overflow");
415    }
416}
nucleus/resources/limits.rs

nucleus/resources/
limits.rs
