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    /// RLIMIT_MEMLOCK in bytes (None = use default 64KB).
113    /// io_uring requires a larger limit (e.g. 8M) for ring buffers.
114    pub memlock_bytes: Option<u64>,
115}
116
117impl ResourceLimits {
118    /// Create unlimited resource limits
119    pub fn unlimited() -> Self {
120        Self {
121            memory_bytes: None,
122            memory_high: None,
123            memory_swap_max: None,
124            cpu_quota_us: None,
125            cpu_period_us: 100_000, // 100ms default period
126            cpu_weight: None,
127            pids_max: None,
128            io_limits: Vec::new(),
129            memlock_bytes: None,
130        }
131    }
132
133    /// Parse memory limit from string (e.g., "512M", "1G")
134    pub fn parse_memory(s: &str) -> Result<u64> {
135        let s = s.trim();
136        if s.is_empty() {
137            return Err(NucleusError::InvalidResourceLimit(
138                "Empty memory limit".to_string(),
139            ));
140        }
141
142        let (num_str, multiplier) = if s.ends_with('K') || s.ends_with('k') {
143            (&s[..s.len() - 1], 1024u64)
144        } else if s.ends_with('M') || s.ends_with('m') {
145            (&s[..s.len() - 1], 1024 * 1024)
146        } else if s.ends_with('G') || s.ends_with('g') {
147            (&s[..s.len() - 1], 1024 * 1024 * 1024)
148        } else if s.ends_with('T') || s.ends_with('t') {
149            (&s[..s.len() - 1], 1024 * 1024 * 1024 * 1024)
150        } else {
151            // No suffix, assume bytes
152            (s, 1)
153        };
154
155        let num: u64 = num_str.parse().map_err(|_| {
156            NucleusError::InvalidResourceLimit(format!("Invalid memory value: {}", s))
157        })?;
158
159        num.checked_mul(multiplier).ok_or_else(|| {
160            NucleusError::InvalidResourceLimit(format!("Memory value overflows u64: {}", s))
161        })
162    }
163
164    /// Set memory limit from string (e.g., "512M", "1G")
165    ///
166    /// Automatically sets memory_high to 90% of the hard limit and
167    /// disables swap (memory_swap_max = 0) unless swap was explicitly enabled.
168    pub fn with_memory(mut self, limit: &str) -> Result<Self> {
169        let bytes = Self::parse_memory(limit)?;
170        self.memory_bytes = Some(bytes);
171        // Auto-set soft limit to 90% of hard limit (per spec)
172        self.memory_high = Some(bytes - bytes / 10);
173        // Disable swap by default when memory limit is set
174        if self.memory_swap_max.is_none() {
175            self.memory_swap_max = Some(0);
176        }
177        Ok(self)
178    }
179
180    /// Enable swap (removes the default swap=0 restriction)
181    pub fn with_swap_enabled(mut self) -> Self {
182        self.memory_swap_max = None;
183        self
184    }
185
186    /// Set CPU limit in cores (e.g., 2.5 cores)
187    pub fn with_cpu_cores(mut self, cores: f64) -> Result<Self> {
188        const MAX_CPU_CORES: f64 = 65_536.0;
189
190        if cores <= 0.0 || cores.is_nan() || cores.is_infinite() {
191            return Err(NucleusError::InvalidResourceLimit(
192                "CPU cores must be a finite positive number".to_string(),
193            ));
194        }
195        if cores > MAX_CPU_CORES {
196            return Err(NucleusError::InvalidResourceLimit(format!(
197                "CPU cores must be <= {}",
198                MAX_CPU_CORES
199            )));
200        }
201        // Convert cores to quota: cores * period
202        let quota = (cores * self.cpu_period_us as f64) as u64;
203        self.cpu_quota_us = Some(quota);
204        Ok(self)
205    }
206
207    /// Set maximum number of PIDs
208    pub fn with_pids(mut self, max_pids: u64) -> Result<Self> {
209        if max_pids == 0 {
210            return Err(NucleusError::InvalidResourceLimit(
211                "Max PIDs must be positive".to_string(),
212            ));
213        }
214        self.pids_max = Some(max_pids);
215        Ok(self)
216    }
217
218    /// Set CPU scheduling weight (1-10000)
219    pub fn with_cpu_weight(mut self, weight: u64) -> Result<Self> {
220        if !(1..=10000).contains(&weight) {
221            return Err(NucleusError::InvalidResourceLimit(
222                "CPU weight must be between 1 and 10000".to_string(),
223            ));
224        }
225        self.cpu_weight = Some(weight);
226        Ok(self)
227    }
228
229    /// Add an I/O device limit
230    pub fn with_io_limit(mut self, limit: IoDeviceLimit) -> Self {
231        self.io_limits.push(limit);
232        self
233    }
234
235    /// Set RLIMIT_MEMLOCK (e.g. "8M" for io_uring ring buffers)
236    pub fn with_memlock(mut self, limit: &str) -> Result<Self> {
237        self.memlock_bytes = Some(Self::parse_memory(limit)?);
238        Ok(self)
239    }
240}
241
242impl Default for ResourceLimits {
243    fn default() -> Self {
244        Self {
245            pids_max: Some(512),
246            ..Self::unlimited()
247        }
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[test]
256    fn test_parse_memory() {
257        assert_eq!(ResourceLimits::parse_memory("1024").unwrap(), 1024);
258        assert_eq!(ResourceLimits::parse_memory("512K").unwrap(), 512 * 1024);
259        assert_eq!(
260            ResourceLimits::parse_memory("512M").unwrap(),
261            512 * 1024 * 1024
262        );
263        assert_eq!(
264            ResourceLimits::parse_memory("2G").unwrap(),
265            2 * 1024 * 1024 * 1024
266        );
267    }
268
269    #[test]
270    fn test_parse_memory_invalid() {
271        assert!(ResourceLimits::parse_memory("").is_err());
272        assert!(ResourceLimits::parse_memory("abc").is_err());
273        assert!(ResourceLimits::parse_memory("M").is_err());
274    }
275
276    #[test]
277    fn test_parse_memory_overflow_rejected() {
278        // 18446744073709551615T would overflow u64
279        assert!(ResourceLimits::parse_memory("99999999999999T").is_err());
280        // Just under u64::MAX in bytes should work
281        assert!(ResourceLimits::parse_memory("16383P").is_err()); // not a valid suffix, treated as bytes
282    }
283
284    #[test]
285    fn test_with_cpu_cores() {
286        let limits = ResourceLimits::unlimited();
287        let limits = limits.with_cpu_cores(2.0).unwrap();
288        assert_eq!(limits.cpu_quota_us, Some(200_000)); // 2.0 * 100_000
289    }
290
291    #[test]
292    fn test_with_cpu_cores_fractional() {
293        let limits = ResourceLimits::unlimited();
294        let limits = limits.with_cpu_cores(0.5).unwrap();
295        assert_eq!(limits.cpu_quota_us, Some(50_000)); // 0.5 * 100_000
296    }
297
298    #[test]
299    fn test_with_cpu_cores_invalid() {
300        let limits = ResourceLimits::unlimited();
301        assert!(limits.with_cpu_cores(0.0).is_err());
302        assert!(ResourceLimits::unlimited().with_cpu_cores(-1.0).is_err());
303    }
304
305    #[test]
306    fn test_with_memory_auto_sets_memory_high() {
307        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
308        let expected_bytes = 1024 * 1024 * 1024u64;
309        assert_eq!(limits.memory_bytes, Some(expected_bytes));
310        // memory_high should be 90% of hard limit
311        assert_eq!(
312            limits.memory_high,
313            Some(expected_bytes - expected_bytes / 10)
314        );
315    }
316
317    #[test]
318    fn test_with_memory_disables_swap_by_default() {
319        let limits = ResourceLimits::unlimited().with_memory("512M").unwrap();
320        assert_eq!(limits.memory_swap_max, Some(0));
321    }
322
323    #[test]
324    fn test_swap_enabled_clears_swap_limit() {
325        let limits = ResourceLimits::unlimited()
326            .with_memory("512M")
327            .unwrap()
328            .with_swap_enabled();
329        assert!(limits.memory_swap_max.is_none());
330    }
331
332    #[test]
333    fn test_with_cpu_weight_valid() {
334        let limits = ResourceLimits::unlimited().with_cpu_weight(100).unwrap();
335        assert_eq!(limits.cpu_weight, Some(100));
336
337        let limits = ResourceLimits::unlimited().with_cpu_weight(1).unwrap();
338        assert_eq!(limits.cpu_weight, Some(1));
339
340        let limits = ResourceLimits::unlimited().with_cpu_weight(10000).unwrap();
341        assert_eq!(limits.cpu_weight, Some(10000));
342    }
343
344    #[test]
345    fn test_with_cpu_weight_invalid() {
346        assert!(ResourceLimits::unlimited().with_cpu_weight(0).is_err());
347        assert!(ResourceLimits::unlimited().with_cpu_weight(10001).is_err());
348    }
349
350    #[test]
351    fn test_io_device_limit_parse_valid() {
352        let limit = IoDeviceLimit::parse("8:0 riops=1000 wbps=10485760").unwrap();
353        assert_eq!(limit.device, "8:0");
354        assert_eq!(limit.riops, Some(1000));
355        assert_eq!(limit.wbps, Some(10485760));
356        assert!(limit.wiops.is_none());
357        assert!(limit.rbps.is_none());
358    }
359
360    #[test]
361    fn test_io_device_limit_parse_all_params() {
362        let limit = IoDeviceLimit::parse("8:0 riops=100 wiops=200 rbps=300 wbps=400").unwrap();
363        assert_eq!(limit.riops, Some(100));
364        assert_eq!(limit.wiops, Some(200));
365        assert_eq!(limit.rbps, Some(300));
366        assert_eq!(limit.wbps, Some(400));
367    }
368
369    #[test]
370    fn test_io_device_limit_parse_invalid() {
371        // Empty string
372        assert!(IoDeviceLimit::parse("").is_err());
373        // Bad device format
374        assert!(IoDeviceLimit::parse("bad").is_err());
375        assert!(IoDeviceLimit::parse("8:0:1").is_err());
376        // Bad param format
377        assert!(IoDeviceLimit::parse("8:0 riops").is_err());
378        // Unknown param
379        assert!(IoDeviceLimit::parse("8:0 foo=100").is_err());
380        // Bad value
381        assert!(IoDeviceLimit::parse("8:0 riops=abc").is_err());
382    }
383
384    #[test]
385    fn test_io_device_limit_to_io_max_line() {
386        let limit = IoDeviceLimit {
387            device: "8:0".to_string(),
388            riops: Some(1000),
389            wiops: None,
390            rbps: None,
391            wbps: Some(10485760),
392        };
393        assert_eq!(limit.to_io_max_line(), "8:0 riops=1000 wbps=10485760");
394    }
395
396    #[test]
397    fn test_unlimited_defaults() {
398        let limits = ResourceLimits::unlimited();
399        assert!(limits.memory_bytes.is_none());
400        assert!(limits.memory_high.is_none());
401        assert!(limits.memory_swap_max.is_none());
402        assert!(limits.cpu_quota_us.is_none());
403        assert!(limits.cpu_weight.is_none());
404        assert!(limits.pids_max.is_none());
405        assert!(limits.io_limits.is_empty());
406    }
407
408    #[test]
409    fn test_memory_high_uses_integer_arithmetic() {
410        // BUG-13: memory_high must use integer arithmetic, not floating point
411        let limits = ResourceLimits::unlimited().with_memory("1G").unwrap();
412        let bytes = 1024u64 * 1024 * 1024;
413        let expected_high = bytes - bytes / 10; // 90% via integer
414        assert_eq!(
415            limits.memory_high,
416            Some(expected_high),
417            "memory_high must be exactly bytes - bytes/10 (integer arithmetic)"
418        );
419    }
420
421    #[test]
422    fn test_cpu_cores_rejects_extreme_values() {
423        // BUG-12: Extreme CPU core values must be rejected, not silently overflow
424        assert!(ResourceLimits::unlimited()
425            .with_cpu_cores(f64::NAN)
426            .is_err());
427        assert!(ResourceLimits::unlimited()
428            .with_cpu_cores(f64::INFINITY)
429            .is_err());
430        assert!(
431            ResourceLimits::unlimited()
432                .with_cpu_cores(100_000.0)
433                .is_err(),
434            "CPU cores > 65536 must be rejected to prevent quota overflow"
435        );
436    }
437}
nucleus/resources/limits.rs

nucleus/resources/
limits.rs
