asc_cache/
procs.rs

1//! Registry of the resident hosts (MCP servers, daemons) one user runs, so each
2//! can size its memory budget as a share of one machine-wide budget instead
3//! of every process claiming a quarter of the free memory for itself, and
4//! split the CPUs with the others that are busy.
5//!
6//! Each process keeps a small file `<dir>/<pid>` holding its start time, the
7//! bytes its memory cache holds and whether it is serving a request, rewritten
8//! when a request starts or ends and while it loads entries. Entries of processes that are gone
9//! (killed without unregistering, or whose pid was reused) are recognised by
10//! the start time and removed by whoever reads them.
11
12use std::fs;
13use std::io::Write;
14use std::path::{Path, PathBuf};
15use std::time::Duration;
16
17/// Environment variable overriding [`ProcessRegistry::default_dir`].
18pub const RUNTIME_DIR_ENV: &str = "ASC_RUNTIME_DIR";
19
20/// Another registered process.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct ProcessUsage {
23    pub pid: u32,
24    /// Bytes its memory cache held at its last update.
25    pub bytes: u64,
26    /// Time since that update.
27    pub idle: Duration,
28    /// Whether it is serving a request; `None` from versions that did not
29    /// say (only the update's age tells).
30    pub busy: Option<bool>,
31}
32
33/// This process's entry in a registry directory; removed on drop.
34pub struct ProcessRegistry {
35    dir: PathBuf,
36    pid: u32,
37    start: u64,
38}
39
40impl ProcessRegistry {
41    /// `$ASC_RUNTIME_DIR`, else `$XDG_RUNTIME_DIR/asc`, else
42    /// `<temp dir>/asc-<uid>/procs`.
43    pub fn default_dir() -> Option<PathBuf> {
44        if let Some(d) = std::env::var_os(RUNTIME_DIR_ENV) {
45            return Some(PathBuf::from(d));
46        }
47        if let Some(d) = std::env::var_os("XDG_RUNTIME_DIR") {
48            return Some(PathBuf::from(d).join("asc"));
49        }
50        Some(std::env::temp_dir().join(format!("asc-{}", uid()?)).join("procs"))
51    }
52
53    /// Register this process in `dir` (created private to the user). `None`
54    /// where processes cannot be told apart reliably (not Linux or macOS) or the
55    /// directory is unusable, e.g. owned by someone else.
56    pub fn open(dir: impl Into<PathBuf>) -> Option<ProcessRegistry> {
57        let dir = dir.into();
58        let pid = std::process::id();
59        let start = start_time(pid)?;
60        create_private_dir(&dir)?;
61        let reg = ProcessRegistry { dir, pid, start };
62        reg.update(0, false).ok()?;
63        Some(reg)
64    }
65
66    pub fn dir(&self) -> &Path {
67        &self.dir
68    }
69
70    /// Publish the bytes this process's memory cache holds now, and whether
71    /// it is serving a request.
72    pub fn update(&self, bytes: u64, busy: bool) -> std::io::Result<()> {
73        let path = self.dir.join(self.pid.to_string());
74        let tmp = self.dir.join(format!(".{}.tmp", self.pid));
75        fs::File::create(&tmp)?.write_all(format!("{} {bytes} {}\n", self.start, u8::from(busy)).as_bytes())?;
76        fs::rename(&tmp, path)
77    }
78
79    /// Every other live registered process; entries of dead ones are removed.
80    pub fn others(&self) -> Vec<ProcessUsage> {
81        Self::others_in(&self.dir)
82    }
83
84    /// [`ProcessRegistry::others`] of `dir` for a process that does not
85    /// register itself (a one-shot command).
86    pub fn others_in(dir: &Path) -> Vec<ProcessUsage> {
87        let mut out = Vec::new();
88        let Ok(rd) = fs::read_dir(dir) else { return out };
89        let me = std::process::id();
90        for f in rd.flatten() {
91            let Some(pid) = f.file_name().to_str().and_then(|n| n.parse::<u32>().ok()) else { continue };
92            if pid == me {
93                continue;
94            }
95            let Ok(text) = fs::read_to_string(f.path()) else { continue };
96            let mut fields = text.split_whitespace().map(|v| v.parse::<u64>().ok());
97            let (Some(Some(start)), Some(Some(bytes))) = (fields.next(), fields.next()) else { continue };
98            let busy = fields.next().flatten().map(|b| b != 0);
99            if start_time(pid) == Some(start) {
100                let idle = f.metadata().and_then(|m| m.modified()).ok().and_then(|t| t.elapsed().ok()).unwrap_or_default();
101                out.push(ProcessUsage { pid, bytes, idle, busy });
102            } else {
103                let _ = fs::remove_file(f.path());
104            }
105        }
106        out
107    }
108}
109
110impl Drop for ProcessRegistry {
111    fn drop(&mut self) {
112        let _ = fs::remove_file(self.dir.join(self.pid.to_string()));
113    }
114}
115
116/// Start time of `pid` in clock ticks since boot (field 22 of
117/// `/proc/<pid>/stat`): with the pid it identifies a process for good.
118#[cfg(target_os = "linux")]
119fn start_time(pid: u32) -> Option<u64> {
120    let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
121    // The command name (field 2) may hold spaces and parentheses; the fields
122    // after its closing parenthesis start at field 3.
123    stat.rsplit_once(')')?.1.split_whitespace().nth(19)?.parse().ok()
124}
125
126/// Start time of `pid` in microseconds since the epoch.
127#[cfg(target_os = "macos")]
128fn start_time(pid: u32) -> Option<u64> {
129    let pid = libc::c_int::try_from(pid).ok()?;
130    let size = libc::c_int::try_from(std::mem::size_of::<libc::proc_bsdinfo>()).ok()?;
131    let mut info = std::mem::MaybeUninit::<libc::proc_bsdinfo>::zeroed();
132    // SAFETY: `info` is a writable buffer of `size` bytes; proc_pidinfo
133    // writes at most that many and returns how many it wrote.
134    let written = unsafe { libc::proc_pidinfo(pid, libc::PROC_PIDTBSDINFO, 0, info.as_mut_ptr().cast(), size) };
135    if written != size {
136        return None;
137    }
138    // SAFETY: proc_pidinfo filled the whole struct (checked above), and it
139    // is plain integers, valid for any bytes.
140    let info = unsafe { info.assume_init() };
141    Some(info.pbi_start_tvsec * 1_000_000 + info.pbi_start_tvusec)
142}
143
144#[cfg(not(any(target_os = "linux", target_os = "macos")))]
145fn start_time(_pid: u32) -> Option<u64> {
146    None
147}
148
149#[cfg(unix)]
150fn uid() -> Option<u32> {
151    // SAFETY: getuid has no preconditions.
152    Some(unsafe { libc::getuid() })
153}
154
155#[cfg(not(unix))]
156fn uid() -> Option<u32> {
157    None
158}
159
160/// Create `dir` readable by the user only, and refuse one owned by someone
161/// else (a world-writable temp dir could hold a planted one).
162#[cfg(unix)]
163fn create_private_dir(dir: &Path) -> Option<()> {
164    use std::os::unix::fs::{DirBuilderExt, MetadataExt};
165    let _ = fs::DirBuilder::new().recursive(true).mode(0o700).create(dir);
166    let meta = fs::symlink_metadata(dir).ok()?;
167    (meta.is_dir() && Some(meta.uid()) == uid()).then_some(())
168}
169
170#[cfg(not(unix))]
171fn create_private_dir(dir: &Path) -> Option<()> {
172    fs::create_dir_all(dir).ok()
173}
174
175#[cfg(all(test, any(target_os = "linux", target_os = "macos")))]
176mod tests {
177    use super::*;
178
179    // pid 1 (init) is not readable by every user elsewhere.
180    #[cfg(target_os = "linux")]
181    #[test]
182    fn processes_see_each_other_and_stale_entries_go() {
183        let dir = tempfile::tempdir().unwrap();
184        let me = ProcessRegistry::open(dir.path()).unwrap();
185        me.update(123, true).unwrap();
186        assert!(me.others().is_empty(), "never lists itself");
187        // A live process (pid 1 is init) and a dead one (start time mismatch).
188        let init_start = start_time(1).unwrap();
189        fs::write(dir.path().join("1"), format!("{init_start} 456 1\n")).unwrap();
190        fs::write(dir.path().join("2"), "1 789\n").unwrap();
191        fs::write(dir.path().join("notes"), "x").unwrap();
192        // Another registry for the same pid stands in for a second process.
193        let others: Vec<_> = me.others().iter().map(|o| (o.pid, o.bytes, o.busy)).collect();
194        assert_eq!(others, vec![(1, 456, Some(true))]);
195        // Entries of older versions say nothing about being busy.
196        fs::write(dir.path().join("1"), format!("{init_start} 456\n")).unwrap();
197        assert_eq!(me.others()[0].busy, None);
198        assert_eq!(ProcessRegistry::others_in(dir.path()).len(), 1, "readable without registering");
199        assert!(me.others()[0].idle < Duration::from_secs(60), "just written");
200        assert!(!dir.path().join("2").exists(), "dead entry removed");
201        assert!(dir.path().join(std::process::id().to_string()).exists());
202        drop(me);
203        assert!(!dir.path().join(std::process::id().to_string()).exists(), "unregistered on drop");
204    }
205
206    #[test]
207    fn start_time_survives_odd_command_names() {
208        assert!(start_time(std::process::id()).is_some());
209        assert!(start_time(u32::MAX).is_none());
210    }
211}
asc_cache/procs.rs

asc_cache/
procs.rs
