Skip to main content

lean_ctx/core/
startup_guard.rs

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