1use std::fs::{self, OpenOptions};
10use std::io::{ErrorKind, Write};
11use std::path::{Path, PathBuf};
12use std::thread;
13use std::time::{Duration, Instant, SystemTime};
14
15#[derive(Debug)]
17pub enum LockError {
18 Locked,
20 Io(std::io::Error),
22}
23
24impl std::fmt::Display for LockError {
25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26 match self {
27 LockError::Locked => write!(f, "directory is locked by another live process"),
28 LockError::Io(e) => write!(f, "{e}"),
29 }
30 }
31}
32
33impl std::error::Error for LockError {}
34
35impl From<std::io::Error> for LockError {
36 fn from(e: std::io::Error) -> Self {
37 LockError::Io(e)
38 }
39}
40
41pub struct WorkdirLock {
43 lock_path: PathBuf,
44}
45
46impl WorkdirLock {
47 pub fn acquire(dir: &Path) -> Result<Self, LockError> {
51 fs::create_dir_all(dir)?;
52 let lock_path = dir.join("wdlock");
53 let pid = std::process::id();
54
55 loop {
56 match OpenOptions::new()
57 .write(true)
58 .create_new(true)
59 .open(&lock_path)
60 {
61 Ok(mut file) => {
62 if let Err(e) = write!(file, "{pid}").and_then(|()| file.sync_all()) {
63 let _ = fs::remove_file(&lock_path);
64 return Err(LockError::Io(e));
65 }
66 return Ok(WorkdirLock { lock_path });
67 }
68 Err(e) if e.kind() == ErrorKind::AlreadyExists => {
69 if let Some(snapshot) = stale_lock_snapshot(&lock_path)? {
70 remove_if_unchanged(&lock_path, &snapshot)?;
74 continue;
75 }
76 return Err(LockError::Locked);
77 }
78 Err(e) => return Err(LockError::Io(e)),
79 }
80 }
81 }
82
83 pub fn acquire_wait(dir: &Path, timeout: Duration) -> Result<Self, LockError> {
87 let start = Instant::now();
88 loop {
89 match Self::acquire(dir) {
90 Ok(lock) => return Ok(lock),
91 Err(LockError::Locked) if start.elapsed() < timeout => {
92 thread::sleep(Duration::from_millis(5));
93 }
94 Err(err) => return Err(err),
95 }
96 }
97 }
98}
99
100impl Drop for WorkdirLock {
101 fn drop(&mut self) {
102 if let Ok(contents) = fs::read_to_string(&self.lock_path) {
105 if contents.trim() == std::process::id().to_string() {
106 let _ = fs::remove_file(&self.lock_path);
107 }
108 }
109 }
110}
111
112#[derive(Debug)]
113struct LockSnapshot {
114 contents: String,
115 len: u64,
116 modified: Option<SystemTime>,
117}
118
119fn stale_lock_snapshot(lock_path: &Path) -> Result<Option<LockSnapshot>, LockError> {
124 let gone = || LockSnapshot {
125 contents: String::new(),
126 len: 0,
127 modified: None,
128 };
129 let meta = match fs::metadata(lock_path) {
130 Ok(meta) => meta,
131 Err(e) if e.kind() == ErrorKind::NotFound => return Ok(Some(gone())),
132 Err(e) => return Err(LockError::Io(e)),
133 };
134 let contents = match fs::read_to_string(lock_path) {
135 Ok(contents) => contents,
136 Err(e) if e.kind() == ErrorKind::NotFound => return Ok(Some(gone())),
137 Err(e) => return Err(LockError::Io(e)),
138 };
139 let snapshot = LockSnapshot {
140 len: meta.len(),
141 modified: meta.modified().ok(),
142 contents,
143 };
144 let Ok(pid) = snapshot.contents.trim().parse::<u32>() else {
145 return Ok(lock_age(lock_path)
146 .is_some_and(|age| age > Duration::from_secs(30))
147 .then_some(snapshot));
148 };
149 Ok((!is_process_alive(pid)).then_some(snapshot))
150}
151
152fn remove_if_unchanged(lock_path: &Path, snapshot: &LockSnapshot) -> Result<(), LockError> {
156 if snapshot.modified.is_none() && snapshot.contents.is_empty() {
157 return Ok(());
158 }
159 let meta = match fs::metadata(lock_path) {
160 Ok(meta) => meta,
161 Err(e) if e.kind() == ErrorKind::NotFound => return Ok(()),
162 Err(e) => return Err(LockError::Io(e)),
163 };
164 let contents = match fs::read_to_string(lock_path) {
165 Ok(contents) => contents,
166 Err(e) if e.kind() == ErrorKind::NotFound => return Ok(()),
167 Err(e) => return Err(LockError::Io(e)),
168 };
169 if contents != snapshot.contents
170 || meta.len() != snapshot.len
171 || meta.modified().ok() != snapshot.modified
172 {
173 return Ok(());
174 }
175 match fs::remove_file(lock_path) {
176 Ok(()) | Err(_) => Ok(()),
177 }
178}
179
180fn lock_age(lock_path: &Path) -> Option<Duration> {
181 fs::metadata(lock_path)
182 .and_then(|m| m.modified())
183 .ok()
184 .and_then(|modified| modified.elapsed().ok())
185}
186
187fn is_process_alive(pid: u32) -> bool {
188 #[cfg(target_os = "linux")]
189 {
190 Path::new(&format!("/proc/{pid}")).exists()
191 }
192 #[cfg(target_os = "macos")]
193 {
194 std::process::Command::new("kill")
196 .args(["-0", &pid.to_string()])
197 .stdout(std::process::Stdio::null())
198 .stderr(std::process::Stdio::null())
199 .status()
200 .map(|s| s.success())
201 .unwrap_or(false)
202 }
203 #[cfg(not(any(target_os = "linux", target_os = "macos")))]
204 {
205 let _ = pid;
206 true
209 }
210}
211
212#[cfg(test)]
213mod tests {
214 use super::*;
215
216 fn scratch(name: &str) -> PathBuf {
217 let dir = std::env::temp_dir().join(format!("choir-fs-{name}-{}", std::process::id()));
218 std::fs::remove_dir_all(&dir).ok();
219 std::fs::create_dir_all(&dir).unwrap();
220 dir
221 }
222
223 #[test]
224 fn second_acquire_is_refused_and_release_frees() {
225 let dir = scratch("lock-basic");
226 let lock = WorkdirLock::acquire(&dir).unwrap();
227 assert!(dir.join("wdlock").is_file());
228 match WorkdirLock::acquire(&dir) {
229 Err(LockError::Locked) => {}
230 other => panic!("expected Locked, got {:?}", other.map(|_| ())),
231 }
232 drop(lock);
233 assert!(!dir.join("wdlock").exists());
234 let _relock = WorkdirLock::acquire(&dir).unwrap();
235 std::fs::remove_dir_all(&dir).ok();
236 }
237
238 #[test]
239 fn a_dead_owners_lock_is_reaped() {
240 let dir = scratch("lock-reap");
241 let dead = std::process::Command::new("true")
243 .status()
244 .map(|_| ())
245 .and_then(|()| {
246 let child = std::process::Command::new("true").spawn()?;
247 let pid = child.id();
248 child.wait_with_output()?;
249 Ok(pid)
250 })
251 .unwrap();
252 fs::create_dir_all(&dir).unwrap();
253 fs::write(dir.join("wdlock"), dead.to_string()).unwrap();
254
255 let _lock = WorkdirLock::acquire(&dir).expect("dead owner's lock reaped");
256 assert_eq!(
257 fs::read_to_string(dir.join("wdlock")).unwrap().trim(),
258 std::process::id().to_string()
259 );
260 std::fs::remove_dir_all(&dir).ok();
261 }
262
263 #[test]
264 fn fresh_malformed_lock_is_respected() {
265 let dir = scratch("lock-malformed");
268 fs::write(dir.join("wdlock"), "not-a-pid").unwrap();
269 match WorkdirLock::acquire(&dir) {
270 Err(LockError::Locked) => {}
271 other => panic!("expected Locked, got {:?}", other.map(|_| ())),
272 }
273 std::fs::remove_dir_all(&dir).ok();
274 }
275
276 #[test]
277 fn acquire_wait_gets_the_lock_once_the_holder_releases() {
278 let dir = scratch("lock-wait");
279 let held = WorkdirLock::acquire(&dir).unwrap();
280 let dir2 = dir.clone();
281 let waiter =
282 thread::spawn(move || WorkdirLock::acquire_wait(&dir2, Duration::from_secs(5)).is_ok());
283 thread::sleep(Duration::from_millis(30));
284 drop(held);
285 assert!(waiter.join().unwrap());
286 std::fs::remove_dir_all(&dir).ok();
287 }
288}