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

asc_cache/
procs.rs
