nucleus/resources/
limits.rs

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

nucleus/resources/
limits.rs
