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    let rc = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
79    rc == 0
80}
81
82#[cfg(not(unix))]
83fn try_flock(_file: &File) -> bool {
84    true
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn acquire_and_release() {
93        let guard = acquire();
94        assert!(guard.is_some(), "should acquire first slot");
95        drop(guard);
96    }
97
98    #[cfg(unix)]
99    #[test]
100    fn active_count_reflects_held_slots() {
101        let g1 = acquire();
102        assert!(g1.is_some());
103        let count = active_count();
104        assert!(count >= 1, "at least one slot held, got {count}");
105        drop(g1);
106    }
107}