agentos_kernel/
resource_accounting.rs

1use crate::fd_table::FdTableManager;
2use crate::pipe_manager::PipeManager;
3use crate::process_table::{ProcessStatus, ProcessTable};
4use crate::pty::PtyManager;
5use crate::socket_table::{SocketState, SocketTable};
6use agentos_bridge::queue_tracker::{register_limit, QueueGauge, TrackedLimit};
7use std::collections::BTreeMap;
8use std::error::Error;
9use std::fmt;
10use std::sync::Arc;
11use vfs::posix::usage::RootFilesystemResourceLimits;
12
13pub use vfs::posix::usage::{
14    measure_filesystem_usage, FileSystemUsage, DEFAULT_MAX_FILESYSTEM_BYTES,
15    DEFAULT_MAX_INODE_COUNT,
16};
17
18pub const DEFAULT_MAX_PROCESSES: usize = 256;
19pub const DEFAULT_MAX_OPEN_FDS: usize = 256;
20pub const DEFAULT_MAX_PIPES: usize = 128;
21pub const DEFAULT_MAX_PTYS: usize = 128;
22pub const DEFAULT_MAX_SOCKETS: usize = 256;
23pub const DEFAULT_MAX_CONNECTIONS: usize = 256;
24pub const DEFAULT_MAX_SOCKET_BUFFERED_BYTES: usize = 4 * 1024 * 1024;
25pub const DEFAULT_MAX_SOCKET_DATAGRAM_QUEUE_LEN: usize = 1_024;
26pub const DEFAULT_BLOCKING_READ_TIMEOUT_MS: u64 = 5_000;
27pub const DEFAULT_MAX_PREAD_BYTES: usize = 64 * 1024 * 1024;
28pub const DEFAULT_MAX_FD_WRITE_BYTES: usize = 64 * 1024 * 1024;
29pub const DEFAULT_MAX_PROCESS_ARGV_BYTES: usize = 1024 * 1024;
30pub const DEFAULT_MAX_PROCESS_ENV_BYTES: usize = 1024 * 1024;
31pub const DEFAULT_MAX_READDIR_ENTRIES: usize = 4_096;
32pub const DEFAULT_MAX_RECURSIVE_FS_DEPTH: usize = 128;
33pub const DEFAULT_MAX_RECURSIVE_FS_ENTRIES: usize = 65_536;
34pub const DEFAULT_VIRTUAL_CPU_COUNT: usize = 1;
35pub const DEFAULT_MAX_WASM_MEMORY_BYTES: u64 = 128 * 1024 * 1024;
36
37#[derive(Debug, Clone, PartialEq, Eq, Default)]
38pub struct ResourceSnapshot {
39    pub running_processes: usize,
40    pub exited_processes: usize,
41    pub fd_tables: usize,
42    pub open_fds: usize,
43    pub pipes: usize,
44    pub pipe_buffered_bytes: usize,
45    pub ptys: usize,
46    pub pty_buffered_input_bytes: usize,
47    pub pty_buffered_output_bytes: usize,
48    pub sockets: usize,
49    pub socket_listeners: usize,
50    pub socket_connections: usize,
51    pub socket_buffered_bytes: usize,
52    pub socket_datagram_queue_len: usize,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct ResourceLimits {
57    pub virtual_cpu_count: Option<usize>,
58    pub max_processes: Option<usize>,
59    pub max_open_fds: Option<usize>,
60    pub max_pipes: Option<usize>,
61    pub max_ptys: Option<usize>,
62    pub max_sockets: Option<usize>,
63    pub max_connections: Option<usize>,
64    pub max_socket_buffered_bytes: Option<usize>,
65    pub max_socket_datagram_queue_len: Option<usize>,
66    pub max_filesystem_bytes: Option<u64>,
67    pub max_inode_count: Option<usize>,
68    pub max_blocking_read_ms: Option<u64>,
69    pub max_pread_bytes: Option<usize>,
70    pub max_fd_write_bytes: Option<usize>,
71    pub max_process_argv_bytes: Option<usize>,
72    pub max_process_env_bytes: Option<usize>,
73    pub max_readdir_entries: Option<usize>,
74    pub max_recursive_fs_depth: Option<usize>,
75    pub max_recursive_fs_entries: Option<usize>,
76    pub max_wasm_fuel: Option<u64>,
77    pub max_wasm_memory_bytes: Option<u64>,
78    pub max_wasm_stack_bytes: Option<usize>,
79}
80
81impl Default for ResourceLimits {
82    fn default() -> Self {
83        Self {
84            virtual_cpu_count: Some(DEFAULT_VIRTUAL_CPU_COUNT),
85            max_processes: Some(DEFAULT_MAX_PROCESSES),
86            max_open_fds: Some(DEFAULT_MAX_OPEN_FDS),
87            max_pipes: Some(DEFAULT_MAX_PIPES),
88            max_ptys: Some(DEFAULT_MAX_PTYS),
89            max_sockets: Some(DEFAULT_MAX_SOCKETS),
90            max_connections: Some(DEFAULT_MAX_CONNECTIONS),
91            max_socket_buffered_bytes: Some(DEFAULT_MAX_SOCKET_BUFFERED_BYTES),
92            max_socket_datagram_queue_len: Some(DEFAULT_MAX_SOCKET_DATAGRAM_QUEUE_LEN),
93            max_filesystem_bytes: Some(DEFAULT_MAX_FILESYSTEM_BYTES),
94            max_inode_count: Some(DEFAULT_MAX_INODE_COUNT),
95            max_blocking_read_ms: Some(DEFAULT_BLOCKING_READ_TIMEOUT_MS),
96            max_pread_bytes: Some(DEFAULT_MAX_PREAD_BYTES),
97            max_fd_write_bytes: Some(DEFAULT_MAX_FD_WRITE_BYTES),
98            max_process_argv_bytes: Some(DEFAULT_MAX_PROCESS_ARGV_BYTES),
99            max_process_env_bytes: Some(DEFAULT_MAX_PROCESS_ENV_BYTES),
100            max_readdir_entries: Some(DEFAULT_MAX_READDIR_ENTRIES),
101            max_recursive_fs_depth: Some(DEFAULT_MAX_RECURSIVE_FS_DEPTH),
102            max_recursive_fs_entries: Some(DEFAULT_MAX_RECURSIVE_FS_ENTRIES),
103            max_wasm_fuel: None,
104            // Match the Workers-style default memory envelope where sensible:
105            // guests are bounded unless the trusted VM config raises the cap.
106            max_wasm_memory_bytes: Some(DEFAULT_MAX_WASM_MEMORY_BYTES),
107            max_wasm_stack_bytes: None,
108        }
109    }
110}
111
112#[derive(Debug, Clone, PartialEq, Eq)]
113pub struct ResourceError {
114    code: &'static str,
115    message: String,
116}
117
118impl RootFilesystemResourceLimits for ResourceLimits {
119    fn max_filesystem_bytes(&self) -> Option<u64> {
120        self.max_filesystem_bytes
121    }
122
123    fn max_inode_count(&self) -> Option<usize> {
124        self.max_inode_count
125    }
126}
127
128impl ResourceError {
129    pub fn code(&self) -> &'static str {
130        self.code
131    }
132
133    fn exhausted(message: impl Into<String>) -> Self {
134        Self {
135            code: "EAGAIN",
136            message: message.into(),
137        }
138    }
139
140    fn file_table_full(message: impl Into<String>) -> Self {
141        Self {
142            code: "ENFILE",
143            message: message.into(),
144        }
145    }
146
147    fn filesystem_full(message: impl Into<String>) -> Self {
148        Self {
149            code: "ENOSPC",
150            message: message.into(),
151        }
152    }
153
154    fn invalid_input(message: impl Into<String>) -> Self {
155        Self {
156            code: "EINVAL",
157            message: message.into(),
158        }
159    }
160
161    fn out_of_memory(message: impl Into<String>) -> Self {
162        Self {
163            code: "ENOMEM",
164            message: message.into(),
165        }
166    }
167}
168
169impl fmt::Display for ResourceError {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        write!(f, "{}: {}", self.code, self.message)
172    }
173}
174
175impl Error for ResourceError {}
176
177#[derive(Debug, Clone, Default)]
178/// Per-VM gauges for the saturating resource limits, registered with the central
179/// limit registry so their usage is inspectable and they emit an edge-triggered
180/// approach warning (~80%) before the guest hits the hard cap. Only limits that
181/// are actually set get a gauge; unbounded (`None`) limits are skipped.
182struct ResourceGauges {
183    processes: Option<Arc<QueueGauge>>,
184    open_fds: Option<Arc<QueueGauge>>,
185    pipes: Option<Arc<QueueGauge>>,
186    ptys: Option<Arc<QueueGauge>>,
187    sockets: Option<Arc<QueueGauge>>,
188    connections: Option<Arc<QueueGauge>>,
189    socket_buffered_bytes: Option<Arc<QueueGauge>>,
190    socket_datagram_queue_len: Option<Arc<QueueGauge>>,
191    filesystem_bytes: Option<Arc<QueueGauge>>,
192    inodes: Option<Arc<QueueGauge>>,
193    recursive_fs_depth: Option<Arc<QueueGauge>>,
194    recursive_fs_entries: Option<Arc<QueueGauge>>,
195}
196
197fn register_resource_gauge(name: TrackedLimit, limit: Option<usize>) -> Option<Arc<QueueGauge>> {
198    limit.map(|capacity| register_limit(name, capacity))
199}
200
201fn register_resource_gauge_u64(name: TrackedLimit, limit: Option<u64>) -> Option<Arc<QueueGauge>> {
202    limit.map(|capacity| register_limit(name, usize_saturating_from_u64(capacity)))
203}
204
205fn usize_saturating_from_u64(value: u64) -> usize {
206    usize::try_from(value).unwrap_or(usize::MAX)
207}
208
209impl ResourceGauges {
210    fn new(limits: &ResourceLimits) -> Self {
211        Self {
212            processes: register_resource_gauge(TrackedLimit::VmProcesses, limits.max_processes),
213            open_fds: register_resource_gauge(TrackedLimit::VmOpenFds, limits.max_open_fds),
214            pipes: register_resource_gauge(TrackedLimit::VmPipes, limits.max_pipes),
215            ptys: register_resource_gauge(TrackedLimit::VmPtys, limits.max_ptys),
216            sockets: register_resource_gauge(TrackedLimit::VmSockets, limits.max_sockets),
217            connections: register_resource_gauge(
218                TrackedLimit::VmConnections,
219                limits.max_connections,
220            ),
221            socket_buffered_bytes: register_resource_gauge(
222                TrackedLimit::VmSocketBufferedBytes,
223                limits.max_socket_buffered_bytes,
224            ),
225            socket_datagram_queue_len: register_resource_gauge(
226                TrackedLimit::VmSocketDatagramQueueLen,
227                limits.max_socket_datagram_queue_len,
228            ),
229            filesystem_bytes: register_resource_gauge_u64(
230                TrackedLimit::VmFilesystemBytes,
231                limits.max_filesystem_bytes,
232            ),
233            inodes: register_resource_gauge(TrackedLimit::VmInodes, limits.max_inode_count),
234            recursive_fs_depth: register_resource_gauge(
235                TrackedLimit::VmRecursiveFsDepth,
236                limits.max_recursive_fs_depth,
237            ),
238            recursive_fs_entries: register_resource_gauge(
239                TrackedLimit::VmRecursiveFsEntries,
240                limits.max_recursive_fs_entries,
241            ),
242        }
243    }
244}
245
246pub struct ResourceAccountant {
247    limits: ResourceLimits,
248    gauges: ResourceGauges,
249}
250
251impl ResourceAccountant {
252    pub fn new(limits: ResourceLimits) -> Self {
253        let gauges = ResourceGauges::new(&limits);
254        Self { limits, gauges }
255    }
256
257    /// Sample the saturating-resource gauges from a fresh snapshot so the central
258    /// registry tracks usage and warns before any cap is reached.
259    fn observe_resource_gauges(&self, snapshot: &ResourceSnapshot) {
260        if let Some(gauge) = &self.gauges.processes {
261            gauge.observe_depth(snapshot.running_processes + snapshot.exited_processes);
262        }
263        if let Some(gauge) = &self.gauges.open_fds {
264            gauge.observe_depth(snapshot.open_fds);
265        }
266        if let Some(gauge) = &self.gauges.pipes {
267            gauge.observe_depth(snapshot.pipes);
268        }
269        if let Some(gauge) = &self.gauges.ptys {
270            gauge.observe_depth(snapshot.ptys);
271        }
272        if let Some(gauge) = &self.gauges.sockets {
273            gauge.observe_depth(snapshot.sockets);
274        }
275        if let Some(gauge) = &self.gauges.connections {
276            gauge.observe_depth(snapshot.socket_connections);
277        }
278        if let Some(gauge) = &self.gauges.socket_buffered_bytes {
279            gauge.observe_depth(snapshot.socket_buffered_bytes);
280        }
281        if let Some(gauge) = &self.gauges.socket_datagram_queue_len {
282            gauge.observe_depth(snapshot.socket_datagram_queue_len);
283        }
284    }
285
286    pub fn limits(&self) -> &ResourceLimits {
287        &self.limits
288    }
289
290    pub fn snapshot(
291        &self,
292        processes: &ProcessTable,
293        fd_tables: &FdTableManager,
294        pipes: &PipeManager,
295        ptys: &PtyManager,
296        sockets: &SocketTable,
297    ) -> ResourceSnapshot {
298        let process_list = processes.list_processes();
299        let running_processes = process_list
300            .values()
301            .filter(|process| process.status == ProcessStatus::Running)
302            .count();
303        let exited_processes = process_list
304            .values()
305            .filter(|process| process.status == ProcessStatus::Exited)
306            .count();
307        let socket_snapshot = sockets.snapshot();
308
309        let snapshot = ResourceSnapshot {
310            running_processes,
311            exited_processes,
312            fd_tables: fd_tables.len(),
313            open_fds: fd_tables.total_open_fds(),
314            pipes: pipes.pipe_count(),
315            pipe_buffered_bytes: pipes.buffered_bytes(),
316            ptys: ptys.pty_count(),
317            pty_buffered_input_bytes: ptys.buffered_input_bytes(),
318            pty_buffered_output_bytes: ptys.buffered_output_bytes(),
319            sockets: socket_snapshot.sockets,
320            socket_listeners: socket_snapshot.listeners,
321            socket_connections: socket_snapshot.connections,
322            socket_buffered_bytes: socket_snapshot.buffered_bytes,
323            socket_datagram_queue_len: socket_snapshot.datagram_queue_len,
324        };
325        self.observe_resource_gauges(&snapshot);
326        snapshot
327    }
328
329    pub fn check_process_spawn(
330        &self,
331        snapshot: &ResourceSnapshot,
332        additional_fds: usize,
333    ) -> Result<(), ResourceError> {
334        if let Some(limit) = self.limits.max_processes {
335            if snapshot.running_processes + snapshot.exited_processes >= limit {
336                return Err(ResourceError::exhausted("maximum process limit reached"));
337            }
338        }
339
340        self.check_open_fds(snapshot, additional_fds)
341    }
342
343    pub fn check_process_argv_bytes(
344        &self,
345        command: &str,
346        args: &[String],
347    ) -> Result<(), ResourceError> {
348        if let Some(limit) = self.limits.max_process_argv_bytes {
349            let total = argv_payload_bytes(command, args);
350            if total > limit {
351                return Err(ResourceError::invalid_input(format!(
352                    "process argv payload {total} bytes exceeds configured limit {limit}"
353                )));
354            }
355        }
356
357        Ok(())
358    }
359
360    pub fn check_process_env_bytes(
361        &self,
362        inherited_env: &BTreeMap<String, String>,
363        overrides: &BTreeMap<String, String>,
364    ) -> Result<(), ResourceError> {
365        if let Some(limit) = self.limits.max_process_env_bytes {
366            let total = merged_env_payload_bytes(inherited_env, overrides);
367            if total > limit {
368                return Err(ResourceError::invalid_input(format!(
369                    "process environment payload {total} bytes exceeds configured limit {limit}"
370                )));
371            }
372        }
373
374        Ok(())
375    }
376
377    pub fn check_pipe_allocation(&self, snapshot: &ResourceSnapshot) -> Result<(), ResourceError> {
378        if let Some(limit) = self.limits.max_pipes {
379            if snapshot.pipes >= limit {
380                return Err(ResourceError::exhausted("maximum pipe count reached"));
381            }
382        }
383
384        self.check_open_fds(snapshot, 2)
385    }
386
387    pub fn check_pty_allocation(&self, snapshot: &ResourceSnapshot) -> Result<(), ResourceError> {
388        if let Some(limit) = self.limits.max_ptys {
389            if snapshot.ptys >= limit {
390                return Err(ResourceError::exhausted("maximum PTY count reached"));
391            }
392        }
393
394        self.check_open_fds(snapshot, 2)
395    }
396
397    pub fn check_socket_allocation(
398        &self,
399        snapshot: &ResourceSnapshot,
400    ) -> Result<(), ResourceError> {
401        if let Some(limit) = self.limits.max_sockets {
402            if snapshot.sockets >= limit {
403                return Err(ResourceError::exhausted("maximum socket count reached"));
404            }
405        }
406
407        Ok(())
408    }
409
410    pub fn check_socket_state_transition(
411        &self,
412        snapshot: &ResourceSnapshot,
413        current: SocketState,
414        next: SocketState,
415    ) -> Result<(), ResourceError> {
416        if !current.counts_as_connection() && next.counts_as_connection() {
417            if let Some(limit) = self.limits.max_connections {
418                if snapshot.socket_connections >= limit {
419                    return Err(ResourceError::exhausted("maximum connection count reached"));
420                }
421            }
422        }
423
424        Ok(())
425    }
426
427    pub fn check_socket_buffer_growth(
428        &self,
429        snapshot: &ResourceSnapshot,
430        additional_bytes: usize,
431    ) -> Result<(), ResourceError> {
432        if let Some(limit) = self.limits.max_socket_buffered_bytes {
433            if snapshot
434                .socket_buffered_bytes
435                .saturating_add(additional_bytes)
436                > limit
437            {
438                return Err(ResourceError::exhausted(
439                    "maximum socket buffered byte limit reached",
440                ));
441            }
442        }
443
444        Ok(())
445    }
446
447    pub fn check_socket_datagram_enqueue(
448        &self,
449        snapshot: &ResourceSnapshot,
450        additional_bytes: usize,
451    ) -> Result<(), ResourceError> {
452        self.check_socket_buffer_growth(snapshot, additional_bytes)?;
453        if let Some(limit) = self.limits.max_socket_datagram_queue_len {
454            if snapshot.socket_datagram_queue_len.saturating_add(1) > limit {
455                return Err(ResourceError::exhausted(
456                    "maximum socket datagram queue length reached",
457                ));
458            }
459        }
460
461        Ok(())
462    }
463
464    pub fn check_pread_length(&self, length: usize) -> Result<(), ResourceError> {
465        if let Some(limit) = self.limits.max_pread_bytes {
466            if length > limit {
467                return Err(ResourceError::invalid_input(format!(
468                    "pread length {length} exceeds configured limit {limit}"
469                )));
470            }
471        }
472
473        Ok(())
474    }
475
476    pub fn check_fd_write_size(&self, size: usize) -> Result<(), ResourceError> {
477        if let Some(limit) = self.limits.max_fd_write_bytes {
478            if size > limit {
479                return Err(ResourceError::invalid_input(format!(
480                    "write size {size} exceeds configured limit {limit}"
481                )));
482            }
483        }
484
485        Ok(())
486    }
487
488    pub fn check_fd_allocation(
489        &self,
490        snapshot: &ResourceSnapshot,
491        additional_fds: usize,
492    ) -> Result<(), ResourceError> {
493        self.check_open_fds(snapshot, additional_fds)
494    }
495
496    pub fn max_readdir_entries(&self) -> Option<usize> {
497        self.limits.max_readdir_entries
498    }
499
500    pub fn max_recursive_fs_depth(&self) -> Option<usize> {
501        self.limits.max_recursive_fs_depth
502    }
503
504    pub fn max_recursive_fs_entries(&self) -> Option<usize> {
505        self.limits.max_recursive_fs_entries
506    }
507
508    pub fn check_readdir_entries(&self, entries: usize) -> Result<(), ResourceError> {
509        if let Some(limit) = self.limits.max_readdir_entries {
510            if entries > limit {
511                return Err(ResourceError::out_of_memory(format!(
512                    "directory listing with {entries} entries exceeds configured limit {limit}"
513                )));
514            }
515        }
516
517        Ok(())
518    }
519
520    pub fn check_recursive_fs_depth(&self, depth: usize) -> Result<(), ResourceError> {
521        if let Some(gauge) = self.gauges.recursive_fs_depth.as_ref() {
522            gauge.observe_depth(depth);
523        }
524        if let Some(limit) = self.limits.max_recursive_fs_depth {
525            if depth > limit {
526                return Err(ResourceError::out_of_memory(format!(
527                    "recursive filesystem operation depth {depth} exceeds configured limit {limit}"
528                )));
529            }
530        }
531
532        Ok(())
533    }
534
535    pub fn check_recursive_fs_entries(&self, entries: usize) -> Result<(), ResourceError> {
536        if let Some(gauge) = self.gauges.recursive_fs_entries.as_ref() {
537            gauge.observe_depth(entries);
538        }
539        if let Some(limit) = self.limits.max_recursive_fs_entries {
540            if entries > limit {
541                return Err(ResourceError::out_of_memory(format!(
542                    "recursive filesystem operation with {entries} entries exceeds configured limit {limit}"
543                )));
544            }
545        }
546
547        Ok(())
548    }
549
550    fn check_open_fds(
551        &self,
552        snapshot: &ResourceSnapshot,
553        additional_fds: usize,
554    ) -> Result<(), ResourceError> {
555        if let Some(limit) = self.limits.max_open_fds {
556            if snapshot.open_fds.saturating_add(additional_fds) > limit {
557                return Err(ResourceError::file_table_full(
558                    "maximum open file descriptor limit reached",
559                ));
560            }
561        }
562
563        Ok(())
564    }
565
566    pub fn check_filesystem_usage(
567        &self,
568        _usage: &FileSystemUsage,
569        resulting_bytes: u64,
570        resulting_inodes: usize,
571    ) -> Result<(), ResourceError> {
572        if let Some(limit) = self.limits.max_filesystem_bytes {
573            if resulting_bytes > limit {
574                return Err(ResourceError::filesystem_full(
575                    "maximum filesystem size limit reached",
576                ));
577            }
578        }
579
580        if let Some(limit) = self.limits.max_inode_count {
581            if resulting_inodes > limit {
582                return Err(ResourceError::filesystem_full(
583                    "maximum inode count limit reached",
584                ));
585            }
586        }
587
588        // Sample the gauges only on the success path: observing the *projected*
589        // value before the bounds check would latch a spurious near/at-capacity
590        // warning for a write that is then rejected and never actually applied.
591        if let Some(gauge) = &self.gauges.filesystem_bytes {
592            gauge.observe_depth(usize_saturating_from_u64(resulting_bytes));
593        }
594        if let Some(gauge) = &self.gauges.inodes {
595            gauge.observe_depth(resulting_inodes);
596        }
597        Ok(())
598    }
599}
600
601fn argv_payload_bytes(command: &str, args: &[String]) -> usize {
602    let command_bytes = command.len().saturating_add(1);
603    command_bytes.saturating_add(
604        args.iter()
605            .map(|arg| arg.len().saturating_add(1))
606            .sum::<usize>(),
607    )
608}
609
610fn env_entry_payload_bytes(key: &str, value: &str) -> usize {
611    key.len()
612        .saturating_add(1)
613        .saturating_add(value.len())
614        .saturating_add(1)
615}
616
617fn merged_env_payload_bytes(
618    inherited_env: &BTreeMap<String, String>,
619    overrides: &BTreeMap<String, String>,
620) -> usize {
621    let mut total = inherited_env
622        .iter()
623        .map(|(key, value)| env_entry_payload_bytes(key, value))
624        .sum::<usize>();
625
626    for (key, value) in overrides {
627        if let Some(previous) = inherited_env.get(key) {
628            total = total.saturating_sub(env_entry_payload_bytes(key, previous));
629        }
630        total = total.saturating_add(env_entry_payload_bytes(key, value));
631    }
632
633    total
634}
635
636#[cfg(test)]
637mod gauge_tests {
638    use super::*;
639    use agentos_bridge::queue_tracker::{set_limit_warning_handler, LimitWarning, TrackedLimit};
640    use std::sync::{Arc, Mutex};
641
642    #[test]
643    fn resource_gauges_track_usage_and_warn_on_approach() {
644        let captured: Arc<Mutex<Vec<LimitWarning>>> = Arc::new(Mutex::new(Vec::new()));
645        let sink = Arc::clone(&captured);
646        // Filter by name so a gauge from a concurrently-running test can't pollute.
647        set_limit_warning_handler(Box::new(move |warning| {
648            if warning.name == TrackedLimit::VmOpenFds {
649                sink.lock().expect("sink mutex").push(warning.clone());
650            }
651        }));
652
653        let limits = ResourceLimits {
654            max_open_fds: Some(10),
655            ..ResourceLimits::default()
656        };
657        let accountant = ResourceAccountant::new(limits);
658        let snapshot = ResourceSnapshot {
659            open_fds: 9, // 90% of the cap
660            ..ResourceSnapshot::default()
661        };
662        accountant.observe_resource_gauges(&snapshot);
663
664        // The gauge reflects the sampled usage...
665        let gauge = accountant
666            .gauges
667            .open_fds
668            .as_ref()
669            .expect("open_fds gauge registered when the limit is set");
670        assert_eq!(gauge.depth(), 9);
671        assert_eq!(gauge.capacity(), 10);
672        assert_eq!(gauge.high_water(), 9);
673
674        // ...and crossing ~80% emits the approach warning to the host sink.
675        assert!(
676            captured
677                .lock()
678                .unwrap()
679                .iter()
680                .any(|warning| warning.name == TrackedLimit::VmOpenFds),
681            "open_fds at 90% of cap must emit an approach warning"
682        );
683    }
684
685    #[test]
686    fn unset_limit_registers_no_gauge() {
687        let limits = ResourceLimits {
688            max_ptys: None,
689            ..ResourceLimits::default()
690        };
691        let accountant = ResourceAccountant::new(limits);
692        assert!(
693            accountant.gauges.ptys.is_none(),
694            "an unbounded (None) limit must not register a gauge"
695        );
696    }
697
698    #[test]
699    fn filesystem_gauge_not_latched_by_rejected_write() {
700        let limits = ResourceLimits {
701            max_filesystem_bytes: Some(1000),
702            max_inode_count: Some(100),
703            ..ResourceLimits::default()
704        };
705        let accountant = ResourceAccountant::new(limits);
706        let usage = FileSystemUsage::default();
707
708        // A write that would exceed the byte cap is rejected and must NOT latch
709        // the gauge to the projected (never-applied) value.
710        let rejected = accountant.check_filesystem_usage(&usage, 2000, 0);
711        assert!(rejected.is_err());
712        let bytes_gauge = accountant
713            .gauges
714            .filesystem_bytes
715            .as_ref()
716            .expect("filesystem_bytes gauge registered");
717        assert_eq!(
718            bytes_gauge.depth(),
719            0,
720            "a rejected over-limit write must not bump the gauge"
721        );
722
723        // A successful write does update it.
724        accountant
725            .check_filesystem_usage(&usage, 500, 7)
726            .expect("under-limit write is accepted");
727        assert_eq!(bytes_gauge.depth(), 500);
728        assert_eq!(
729            accountant.gauges.inodes.as_ref().unwrap().depth(),
730            7,
731            "inode gauge tracks the accepted value"
732        );
733    }
734}
agentos_kernel/resource_accounting.rs

agentos_kernel/
resource_accounting.rs
