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, Serialize)]
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    /// Return true when at least one cgroup v2 control will be configured.
278    pub fn has_cgroup_control(&self) -> bool {
279        self.memory_bytes.is_some()
280            || self.memory_high.is_some()
281            || self.memory_swap_max.is_some()
282            || self.cpu_quota_us.is_some()
283            || self.cpu_weight.is_some()
284            || self.pids_max.is_some()
285            || !self.io_limits.is_empty()
286    }
287}
288
289impl Default for ResourceLimits {
290    fn default() -> Self {
291        Self {
292            pids_max: Some(512),
293            ..Self::unlimited()
294        }
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn test_parse_memory() {
304        assert_eq!(ResourceLimits::parse_memory("1024").unwrap(), 1024);
305        assert_eq!(ResourceLimits::parse_memory("512K").unwrap(), 512 * 1024);
306        assert_eq!(
307            ResourceLimits::parse_memory("512M").unwrap(),
308            512 * 1024 * 1024
309        );
310        assert_eq!(
311            ResourceLimits::parse_memory("2G").unwrap(),
312            2 * 1024 * 1024 * 1024
313        );
314    }
315
316    #[test]
317    fn test_parse_memory_invalid() {
318        assert!(ResourceLimits::parse_memory("").is_err());
319        assert!(ResourceLimits::parse_memory("abc").is_err());
320        assert!(ResourceLimits::parse_memory("M").is_err());
321    }
322
323    #[test]
324    fn test_parse_memory_overflow_rejected() {
325        // 18446744073709551615T would overflow u64
326        assert!(ResourceLimits::parse_memory("99999999999999T").is_err());
327        // Just under u64::MAX in bytes should work
328        assert!(ResourceLimits::parse_memory("16383P").is_err()); // not a valid suffix, treated as bytes
329    }
330
331    #[test]
332    fn test_with_cpu_cores() {
333        let limits = ResourceLimits::unlimited();
334        let limits = limits.with_cpu_cores(2.0).unwrap();
335        assert_eq!(limits.cpu_quota_us, Some(200_000)); // 2.0 * 100_000
336    }
337
338    #[test]
339    fn test_with_cpu_cores_fractional() {
340        let limits = ResourceLimits::unlimited();
341        let limits = limits.with_cpu_cores(0.5).unwrap();
342        assert_eq!(limits.cpu_quota_us, Some(50_000)); // 0.5 * 100_000
343    }
344
345    #[test]
346    fn test_with_cpu_cores_invalid() {
347        let limits = ResourceLimits::unlimited();
348        assert!(limits.with_cpu_cores(0.0).is_err());
349        assert!(ResourceLimits::unlimited().with_cpu_cores(-1.0).is_err());
350    }
351
352    #[test]
353    fn test_with_cpu_cores_rejects_zero_period() {
354        let mut limits = ResourceLimits::unlimited();
355        limits.cpu_period_us = 0;
356        assert!(limits.with_cpu_cores(1.0).is_err());
357    }
358
359    #[test]
360    fn test_with_memory_auto_sets_memory_high() {
361        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
362        let expected_bytes = 1024 * 1024 * 1024u64;
363        assert_eq!(limits.memory_bytes, Some(expected_bytes));
364        // memory_high should be 90% of hard limit
365        assert_eq!(
366            limits.memory_high,
367            Some(expected_bytes - expected_bytes / 10)
368        );
369    }
370
371    #[test]
372    fn test_with_memory_disables_swap_by_default() {
373        let limits = ResourceLimits::unlimited().with_memory("512M").unwrap();
374        assert_eq!(limits.memory_swap_max, Some(0));
375    }
376
377    #[test]
378    fn test_swap_enabled_clears_swap_limit() {
379        let limits = ResourceLimits::unlimited()
380            .with_memory("512M")
381            .unwrap()
382            .with_swap_enabled();
383        assert!(limits.memory_swap_max.is_none());
384    }
385
386    #[test]
387    fn test_with_cpu_weight_valid() {
388        let limits = ResourceLimits::unlimited().with_cpu_weight(100).unwrap();
389        assert_eq!(limits.cpu_weight, Some(100));
390
391        let limits = ResourceLimits::unlimited().with_cpu_weight(1).unwrap();
392        assert_eq!(limits.cpu_weight, Some(1));
393
394        let limits = ResourceLimits::unlimited().with_cpu_weight(10000).unwrap();
395        assert_eq!(limits.cpu_weight, Some(10000));
396    }
397
398    #[test]
399    fn test_with_cpu_weight_invalid() {
400        assert!(ResourceLimits::unlimited().with_cpu_weight(0).is_err());
401        assert!(ResourceLimits::unlimited().with_cpu_weight(10001).is_err());
402    }
403
404    #[test]
405    fn test_with_pids_rejects_rlim_infinity() {
406        assert!(ResourceLimits::unlimited()
407            .with_pids(libc::RLIM_INFINITY)
408            .is_err());
409    }
410
411    #[test]
412    fn test_io_device_limit_parse_valid() {
413        let limit = IoDeviceLimit::parse("8:0 riops=1000 wbps=10485760").unwrap();
414        assert_eq!(limit.device, "8:0");
415        assert_eq!(limit.riops, Some(1000));
416        assert_eq!(limit.wbps, Some(10485760));
417        assert!(limit.wiops.is_none());
418        assert!(limit.rbps.is_none());
419    }
420
421    #[test]
422    fn test_io_device_limit_parse_all_params() {
423        let limit = IoDeviceLimit::parse("8:0 riops=100 wiops=200 rbps=300 wbps=400").unwrap();
424        assert_eq!(limit.riops, Some(100));
425        assert_eq!(limit.wiops, Some(200));
426        assert_eq!(limit.rbps, Some(300));
427        assert_eq!(limit.wbps, Some(400));
428    }
429
430    #[test]
431    fn test_io_device_limit_parse_invalid() {
432        // Empty string
433        assert!(IoDeviceLimit::parse("").is_err());
434        // Bad device format
435        assert!(IoDeviceLimit::parse("bad").is_err());
436        assert!(IoDeviceLimit::parse("8:0:1").is_err());
437        // Bad param format
438        assert!(IoDeviceLimit::parse("8:0 riops").is_err());
439        // Unknown param
440        assert!(IoDeviceLimit::parse("8:0 foo=100").is_err());
441        // Bad value
442        assert!(IoDeviceLimit::parse("8:0 riops=abc").is_err());
443    }
444
445    #[test]
446    fn test_io_device_limit_to_io_max_line() {
447        let limit = IoDeviceLimit {
448            device: "8:0".to_string(),
449            riops: Some(1000),
450            wiops: None,
451            rbps: None,
452            wbps: Some(10485760),
453        };
454        assert_eq!(limit.to_io_max_line(), "8:0 riops=1000 wbps=10485760");
455    }
456
457    #[test]
458    fn test_unlimited_defaults() {
459        let limits = ResourceLimits::unlimited();
460        assert!(limits.memory_bytes.is_none());
461        assert!(limits.memory_high.is_none());
462        assert!(limits.memory_swap_max.is_none());
463        assert!(limits.cpu_quota_us.is_none());
464        assert!(limits.cpu_weight.is_none());
465        assert!(limits.pids_max.is_none());
466        assert!(limits.io_limits.is_empty());
467        assert!(!limits.has_cgroup_control());
468        assert!(ResourceLimits::default().has_cgroup_control());
469    }
470
471    #[test]
472    fn test_memory_high_uses_integer_arithmetic() {
473        // BUG-13: memory_high must use integer arithmetic, not floating point
474        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
475        let bytes = 1024u64 * 1024 * 1024;
476        let expected_high = bytes - bytes / 10; // 90% via integer
477        assert_eq!(
478            limits.memory_high,
479            Some(expected_high),
480            "memory_high must be exactly bytes - bytes/10 (integer arithmetic)"
481        );
482    }
483
484    #[test]
485    fn test_cpu_cores_rejects_extreme_values() {
486        // BUG-12: Extreme CPU core values must be rejected, not silently overflow
487        assert!(ResourceLimits::unlimited()
488            .with_cpu_cores(f64::NAN)
489            .is_err());
490        assert!(ResourceLimits::unlimited()
491            .with_cpu_cores(f64::INFINITY)
492            .is_err());
493        assert!(
494            ResourceLimits::unlimited()
495                .with_cpu_cores(100_000.0)
496                .is_err(),
497            "CPU cores > 65536 must be rejected to prevent quota overflow"
498        );
499    }
500
501    #[test]
502    fn test_validate_runtime_sanity_rejects_invalid_deserialized_values() {
503        let mut limits = ResourceLimits::unlimited();
504        limits.cpu_period_us = 0;
505        limits.pids_max = Some(libc::RLIM_INFINITY);
506
507        assert!(limits.validate_runtime_sanity().is_err());
508    }
509}
nucleus/resources/limits.rs

nucleus/resources/
limits.rs
