Skip to main content

memory_budget/
rss.rs

1//! Process-RSS reader abstraction.
2//!
3//! Production code uses [`SystemRss`] which delegates to `sysinfo`.
4//! Tests substitute a hand-rolled source so the budget logic can be
5//! driven through specific RSS values without spawning a sub-process.
6
7use std::sync::Mutex;
8
9use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System};
10
11/// Abstraction over "what is the current process RSS, in bytes".
12///
13/// Implementations may cache their reading; the budget calls
14/// [`RssSource::read_rss_bytes`] once per tick and is happy with a
15/// best-effort value.
16pub trait RssSource: Send + Sync {
17    /// Current process RSS in bytes.
18    fn read_rss_bytes(&self) -> u64;
19
20    /// Installed system RAM in bytes. Used at startup when auto-deriving
21    /// the budget target. Implementations must return a positive value.
22    fn total_ram_bytes(&self) -> u64 {
23        panic!("RssSource::total_ram_bytes must be implemented");
24    }
25}
26
27/// Live RSS reader backed by [`sysinfo`].
28pub struct SystemRss {
29    inner: Mutex<SystemRssInner>,
30}
31
32struct SystemRssInner {
33    system: System,
34    pid: Pid,
35}
36
37impl SystemRss {
38    /// Construct a reader bound to the current process.
39    ///
40    /// # Panics
41    ///
42    /// Panics if the current process id is not representable as a
43    /// [`Pid`] on this platform — a "cannot happen" condition.
44    #[must_use]
45    pub fn new() -> Self {
46        let pid = sysinfo::get_current_pid().expect("sysinfo cannot resolve current pid");
47        let mut system = System::new();
48        system.refresh_memory();
49        Self {
50            inner: Mutex::new(SystemRssInner { system, pid }),
51        }
52    }
53}
54
55impl Default for SystemRss {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl RssSource for SystemRss {
62    fn read_rss_bytes(&self) -> u64 {
63        let mut guard = self.inner.lock().expect("system-rss lock");
64        let pid = guard.pid;
65        let refresh = ProcessRefreshKind::nothing().with_memory();
66        let refreshed = guard.system.refresh_processes_specifics(
67            ProcessesToUpdate::Some(&[pid]),
68            true,
69            refresh,
70        );
71        assert_eq!(refreshed, 1, "sysinfo did not refresh the current process");
72        guard
73            .system
74            .process(pid)
75            .map(sysinfo::Process::memory)
76            .expect("sysinfo lost the current process")
77    }
78
79    fn total_ram_bytes(&self) -> u64 {
80        let mut guard = self.inner.lock().expect("system-rss lock");
81        guard.system.refresh_memory();
82        guard.system.total_memory()
83    }
84}