Skip to main content

lean_ctx/core/
startup_guard.rs

1use std::path::PathBuf;
2use std::time::Duration;
3
4pub struct StartupLockGuard {
5    path: PathBuf,
6}
7
8impl Drop for StartupLockGuard {
9    fn drop(&mut self) {
10        let _ = std::fs::remove_file(&self.path);
11    }
12}
13
14fn sanitize_lock_name(name: &str) -> String {
15    name.chars()
16        .map(|c| {
17            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
18                c
19            } else {
20                '_'
21            }
22        })
23        .collect()
24}
25
26/// Best-effort cross-process lock (create_new + stale eviction).
27///
28/// Returns `None` if the data dir can't be resolved or if the lock can't be acquired
29/// within `timeout`.
30pub fn try_acquire_lock(
31    name: &str,
32    timeout: Duration,
33    stale_after: Duration,
34) -> Option<StartupLockGuard> {
35    let dir = crate::core::data_dir::lean_ctx_data_dir().ok()?;
36    let _ = std::fs::create_dir_all(&dir);
37
38    let name = sanitize_lock_name(name);
39    let path = dir.join(format!(".{name}.lock"));
40
41    let deadline = std::time::Instant::now().checked_add(timeout)?;
42    let mut sleep_ms: u64 = 10;
43
44    loop {
45        if std::fs::OpenOptions::new()
46            .write(true)
47            .create_new(true)
48            .open(&path)
49            .is_ok()
50        {
51            return Some(StartupLockGuard { path });
52        }
53
54        if let Ok(meta) = std::fs::metadata(&path) {
55            if let Ok(modified) = meta.modified() {
56                if modified
57                    .elapsed()
58                    .unwrap_or_default()
59                    .saturating_sub(stale_after)
60                    > Duration::from_secs(0)
61                {
62                    let _ = std::fs::remove_file(&path);
63                }
64            }
65        }
66
67        if std::time::Instant::now() >= deadline {
68            return None;
69        }
70
71        std::thread::sleep(Duration::from_millis(sleep_ms));
72        sleep_ms = (sleep_ms.saturating_mul(2)).min(120);
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    struct EnvVarGuard {
81        key: &'static str,
82        prev: Option<String>,
83    }
84
85    impl EnvVarGuard {
86        fn set(key: &'static str, value: &std::path::Path) -> Self {
87            let prev = std::env::var(key).ok();
88            std::env::set_var(key, value);
89            Self { key, prev }
90        }
91    }
92
93    impl Drop for EnvVarGuard {
94        fn drop(&mut self) {
95            match self.prev.as_deref() {
96                Some(v) => std::env::set_var(self.key, v),
97                None => std::env::remove_var(self.key),
98            }
99        }
100    }
101
102    #[test]
103    fn lock_acquire_and_release() {
104        let _env = crate::core::data_dir::test_env_lock();
105        let dir = tempfile::tempdir().unwrap();
106        let _guard = EnvVarGuard::set("LEAN_CTX_DATA_DIR", dir.path());
107
108        let g = try_acquire_lock(
109            "unit-test",
110            Duration::from_millis(200),
111            Duration::from_secs(30),
112        );
113        assert!(g.is_some());
114
115        let lock_path = dir.path().join(".unit-test.lock");
116        assert!(lock_path.exists());
117
118        drop(g);
119        assert!(!lock_path.exists());
120    }
121
122    #[test]
123    fn lock_times_out_while_held() {
124        let _env = crate::core::data_dir::test_env_lock();
125        let dir = tempfile::tempdir().unwrap();
126        let _guard = EnvVarGuard::set("LEAN_CTX_DATA_DIR", dir.path());
127
128        let g1 = try_acquire_lock(
129            "unit-test-2",
130            Duration::from_millis(200),
131            Duration::from_secs(30),
132        )
133        .expect("first lock should acquire");
134        let g2 = try_acquire_lock(
135            "unit-test-2",
136            Duration::from_millis(60),
137            Duration::from_secs(30),
138        );
139        assert!(g2.is_none());
140
141        drop(g1);
142        let g3 = try_acquire_lock(
143            "unit-test-2",
144            Duration::from_millis(200),
145            Duration::from_secs(30),
146        );
147        assert!(g3.is_some());
148    }
149}