lean_ctx/core/
process_guard.rs

1//! Global concurrency limiter for lean-ctx processes.
2//!
3//! Prevents runaway CPU usage by limiting the number of concurrent lean-ctx
4//! processes to `MAX_CONCURRENT`. Each process acquires a numbered lock slot
5//! under `~/.lean-ctx/locks/`. If all slots are taken, the caller gets `None`.
6
7use std::fs::File;
8use std::path::PathBuf;
9
10const MAX_CONCURRENT: usize = 4;
11
12pub struct ProcessGuard {
13    _file: File,
14    path: PathBuf,
15}
16
17impl Drop for ProcessGuard {
18    fn drop(&mut self) {
19        let _ = std::fs::remove_file(&self.path);
20    }
21}
22
23fn lock_dir() -> Option<PathBuf> {
24    let dir = crate::core::data_dir::lean_ctx_data_dir()
25        .ok()?
26        .join("locks");
27    let _ = std::fs::create_dir_all(&dir);
28    Some(dir)
29}
30
31/// Try to acquire one of N concurrent process slots.
32/// Returns `None` if all slots are occupied (= too many lean-ctx already running).
33pub fn acquire() -> Option<ProcessGuard> {
34    let dir = lock_dir()?;
35
36    for slot in 0..MAX_CONCURRENT {
37        let path = dir.join(format!("slot-{slot}.lock"));
38
39        let Ok(file) = std::fs::OpenOptions::new()
40            .write(true)
41            .create(true)
42            .truncate(false)
43            .open(&path)
44        else {
45            continue;
46        };
47
48        if try_flock(&file) {
49            use std::io::Write;
50            let mut f = file;
51            let _ = f.write_all(format!("{}", std::process::id()).as_bytes());
52            return Some(ProcessGuard { _file: f, path });
53        }
54    }
55
56    None
57}
58
59/// Checks how many slots are currently held (best-effort).
60pub fn active_count() -> usize {
61    let Some(dir) = lock_dir() else { return 0 };
62    let mut count = 0;
63    for slot in 0..MAX_CONCURRENT {
64        let path = dir.join(format!("slot-{slot}.lock"));
65        if let Ok(f) = std::fs::OpenOptions::new().read(true).open(&path)
66            && !try_flock(&f)
67        {
68            count += 1;
69        }
70    }
71    count
72}
73
74#[cfg(unix)]
75fn try_flock(file: &File) -> bool {
76    use std::os::unix::io::AsRawFd;
77    let fd = file.as_raw_fd();
78    // SAFETY: `fd` is a valid, open descriptor owned by `file`, which outlives
79    // this call; `flock` performs no pointer dereference and reports errors via
80    // its return value.
81    let rc = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
82    rc == 0
83}
84
85#[cfg(not(unix))]
86fn try_flock(_file: &File) -> bool {
87    true
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    /// Restores `LEAN_CTX_DATA_DIR` to its previous value on drop (panic-safe).
95    struct EnvRestore(Option<String>);
96    impl Drop for EnvRestore {
97        fn drop(&mut self) {
98            match &self.0 {
99                Some(v) => crate::test_env::set_var("LEAN_CTX_DATA_DIR", v),
100                None => crate::test_env::remove_var("LEAN_CTX_DATA_DIR"),
101            }
102        }
103    }
104
105    /// Runs `body` against a private, empty lock directory.
106    ///
107    /// `acquire()` and `active_count()` both resolve the lock dir from
108    /// `LEAN_CTX_DATA_DIR`. Serializing on `test_env_lock` stops a concurrent
109    /// test from repointing that variable between the two calls (which made
110    /// `active_count` inspect a different, empty dir and miss the held slot), and
111    /// the private temp dir keeps slots independent of any real lean-ctx process
112    /// (daemon/proxy) that might otherwise occupy them.
113    fn with_isolated_lock_dir(body: impl FnOnce()) {
114        let _env = crate::core::data_dir::test_env_lock();
115        let tmp = tempfile::tempdir().expect("tempdir");
116        // Restore runs before `tmp` is removed and while the lock is still held.
117        let _restore = EnvRestore(std::env::var("LEAN_CTX_DATA_DIR").ok());
118        crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
119        body();
120    }
121
122    #[test]
123    fn acquire_and_release() {
124        with_isolated_lock_dir(|| {
125            let guard = acquire();
126            assert!(guard.is_some(), "should acquire first slot");
127            drop(guard);
128        });
129    }
130
131    #[cfg(unix)]
132    #[test]
133    fn active_count_reflects_held_slots() {
134        with_isolated_lock_dir(|| {
135            let g1 = acquire();
136            assert!(g1.is_some());
137            let count = active_count();
138            assert!(count >= 1, "at least one slot held, got {count}");
139            drop(g1);
140        });
141    }
142}
lean_ctx/core/process_guard.rs

lean_ctx/core/
process_guard.rs
