Skip to main content

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            if !try_flock(&f) {
67                count += 1;
68            }
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    #[test]
95    fn acquire_and_release() {
96        let guard = acquire();
97        assert!(guard.is_some(), "should acquire first slot");
98        drop(guard);
99    }
100
101    #[cfg(unix)]
102    #[test]
103    fn active_count_reflects_held_slots() {
104        let g1 = acquire();
105        assert!(g1.is_some());
106        let count = active_count();
107        assert!(count >= 1, "at least one slot held, got {count}");
108        drop(g1);
109    }
110}