Skip to main content

choir_fs/
workdir_lock.rs

1//! A process-level lock on a directory: a PID-bearing lock file created
2//! with `create_new`, reaped when its owner is provably dead. Guards
3//! state a single process may write, e.g. a node's `.choir` dir, where
4//! two writers appending one `ops.jsonl` would fork the chain.
5//!
6//! Ported from Oak `cli/src/workdir_lock.rs` (v0.102.1, commit
7//! `8de9515`, Apache-2.0); error type localized to [`LockError`].
8
9use 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/// Why a [`WorkdirLock`] could not be acquired.
16#[derive(Debug)]
17pub enum LockError {
18    /// Another live process holds the lock.
19    Locked,
20    /// The lock file could not be created, read, or removed.
21    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
41/// Exclusive lock on a directory, released on drop.
42pub struct WorkdirLock {
43    lock_path: PathBuf,
44}
45
46impl WorkdirLock {
47    /// Acquire the lock, reaping a lock file whose recorded PID is no
48    /// longer alive. Fails with [`LockError::Locked`] if a live process
49    /// holds it.
50    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                        // Reap only if the file is still byte-identical to
71                        // what we judged stale; either way, loop so
72                        // create_new gets the next word.
73                        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    /// Acquire, waiting up to `timeout` while a live process holds the
84    /// lock, so short concurrent writers don't force callers into
85    /// process-level retries.
86    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        // Remove only if this process still owns it: a reaper may have
103        // replaced the file after e.g. a laptop sleep.
104        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
119/// `Some(snapshot)` if the lock file looks reapable: missing, owned by a
120/// dead PID, or malformed for over 30 seconds (a crashed partial
121/// acquisition; a *fresh* malformed file may be a live contender that
122/// created the file but has not written its PID yet).
123fn 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
152/// Remove the lock file if it still matches `snapshot`; a mismatch means
153/// another process re-created it between our read and now, so it is
154/// theirs, not stale.
155fn 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        // kill -0: exit 0 iff the process exists (or EPERM, also alive).
195        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        // No liveness probe: fail safe by treating the holder as alive
207        // rather than admitting a second writer.
208        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        // A real PID that is provably dead: a child we already waited on.
242        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        // A just-created file with no parseable PID may be a live
266        // contender mid-acquisition; it must read as locked, not stale.
267        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}