Skip to main content

coop/
lock.rs

1//! Fair, per-host serialization for access to coop's private ssh channel.
2
3use std::fs::{self, OpenOptions};
4use std::io::ErrorKind;
5use std::path::{Path, PathBuf};
6use std::thread;
7use std::time::{Duration, Instant};
8
9use anyhow::{Context, Result};
10
11/// How often a waiter re-checks whether its ticket is being served.
12///
13/// This is added to *every* handoff, so it is a direct tax on the whole queue,
14/// not just on one waiter. Measured with 4 callers x 5 rounds x 50ms of work
15/// (1s of serialised work, so a perfectly fair queue peaks near 200ms):
16///
17/// | poll | p50 | worst observed |
18/// | --- | --- | --- |
19/// | 20ms | ~275ms | **443ms**, and it flaked a 600ms bound 1 run in 15 |
20/// | 5ms | ~237ms | 261ms across 6 runs |
21///
22/// 5ms costs a few more syscalls on a critical section that is sub-second by
23/// construction, and buys back most of the granularity overhead.
24const POLL_INTERVAL: Duration = Duration::from_millis(5);
25
26/// How long to wait before telling the user the queue is not dead.
27///
28/// Never a timeout: aborting would reintroduce the failure the ticket lock
29/// exists to remove.
30const WARN_AFTER: Duration = Duration::from_secs(5);
31
32/// The state directory for one host's ticket lock.
33pub fn lock_path(host: &str) -> PathBuf {
34    let state = std::env::var_os("XDG_STATE_HOME")
35        .map(PathBuf::from)
36        .or_else(|| directories::BaseDirs::new().map(|dirs| dirs.home_dir().join(".local/state")))
37        .unwrap_or_else(|| PathBuf::from(".local/state"));
38    state.join("coop").join(format!("{host}.lock"))
39}
40
41/// Run `f` while holding the fair lock for `host`.
42pub fn with_lock<T>(host: &str, f: impl FnOnce() -> T) -> Result<T> {
43    let path = lock_path(host);
44    fs::create_dir_all(&path)
45        .with_context(|| format!("cannot create lock directory {}", path.display()))?;
46
47    let ticket = with_counter_lock(&path, || {
48        let next = read_counter(&path.join("next"))?;
49        write_counter(&path.join("next"), next + 1)?;
50        if !path.join("serving").exists() {
51            write_counter(&path.join("serving"), 0)?;
52        }
53        // Record that a live process is waiting on this ticket. Without it, a
54        // caller that dies *before* its turn arrives leaves a gap nothing can
55        // detect: `serving` reaches its number, no `holder` was ever written,
56        // and every later caller queues behind a ticket that will never be
57        // claimed. Reproduced -- it wedged the host indefinitely.
58        fs::write(
59            waiter_file(&path, next),
60            format!("{}\n", std::process::id()),
61        )
62        .context("cannot record lock waiter")?;
63        Ok(next)
64    })?;
65
66    wait_for_turn(&path, ticket)?;
67    let guard = HeldLock { path };
68    let result = f();
69    drop(guard);
70    Ok(result)
71}
72
73fn waiter_file(path: &Path, ticket: u64) -> PathBuf {
74    path.join(format!("waiter.{ticket}"))
75}
76
77fn wait_for_turn(path: &Path, ticket: u64) -> Result<()> {
78    let started = Instant::now();
79    let mut warned = false;
80
81    loop {
82        let acquired = with_counter_lock(path, || {
83            let serving = read_counter(&path.join("serving"))?;
84            if serving == ticket {
85                fs::write(path.join("holder"), format!("{}\n", std::process::id()))
86                    .context("cannot record lock holder")?;
87                // Our turn: we are the holder now, not a waiter.
88                fs::remove_file(waiter_file(path, ticket)).ok();
89                return Ok(true);
90            }
91
92            if serving < ticket {
93                match read_pid(&path.join("holder"))? {
94                    // A holder that died mid-work. Step over it.
95                    Some(pid) if !pid_is_alive(pid) => {
96                        fs::remove_file(path.join("holder")).ok();
97                        write_counter(&path.join("serving"), serving + 1)?;
98                    }
99                    Some(_) => {}
100                    // Nobody holds the lock, so whoever owns `serving` is
101                    // either waiting for it or gone. Deciding by pid rather
102                    // than by a timeout keeps this deterministic: a live
103                    // waiter writes `holder` inside the same counter lock we
104                    // are inside now, so it cannot be mid-claim here.
105                    None => {
106                        let waiter = waiter_file(path, serving);
107                        let abandoned = match read_pid(&waiter)? {
108                            Some(pid) => !pid_is_alive(pid),
109                            // No waiter file at all: handed out, then lost.
110                            None => true,
111                        };
112                        if abandoned {
113                            fs::remove_file(&waiter).ok();
114                            write_counter(&path.join("serving"), serving + 1)?;
115                        }
116                    }
117                }
118            }
119            Ok(false)
120        })?;
121
122        if acquired {
123            return Ok(());
124        }
125
126        if !warned && started.elapsed() >= WARN_AFTER {
127            let serving = read_counter(&path.join("serving")).unwrap_or(ticket);
128            let holder = read_pid(&path.join("holder")).ok().flatten();
129            eprintln!(
130                "coop: waiting for lock ticket {ticket} ({} ahead, holder pid {})",
131                ticket.saturating_sub(serving),
132                holder.map_or_else(|| "unknown".into(), |pid| pid.to_string())
133            );
134            warned = true;
135        }
136        thread::sleep(POLL_INTERVAL);
137    }
138}
139
140struct HeldLock {
141    path: PathBuf,
142}
143
144impl Drop for HeldLock {
145    fn drop(&mut self) {
146        let result = with_counter_lock(&self.path, || {
147            let serving = read_counter(&self.path.join("serving"))?;
148            fs::remove_file(self.path.join("holder")).ok();
149            write_counter(&self.path.join("serving"), serving + 1)
150        });
151        if let Err(error) = result {
152            eprintln!("coop: could not release lock: {error:#}");
153        }
154    }
155}
156
157#[cfg(unix)]
158fn with_counter_lock<T>(path: &Path, f: impl FnOnce() -> Result<T>) -> Result<T> {
159    use std::os::fd::AsRawFd;
160
161    unsafe extern "C" {
162        fn flock(fd: i32, operation: i32) -> i32;
163    }
164
165    let file = OpenOptions::new()
166        .read(true)
167        .write(true)
168        .create(true)
169        .truncate(false)
170        .open(path.join("counter.lock"))
171        .context("cannot open counter lock")?;
172    // The OS releases flock when a process dies. This avoids a stale sentinel,
173    // while keeping the ticket read-modify-write atomic without a dependency.
174    if unsafe { flock(file.as_raw_fd(), 2) } != 0 {
175        return Err(std::io::Error::last_os_error()).context("cannot acquire counter lock");
176    }
177    let result = f();
178    drop(file);
179    result
180}
181
182#[cfg(not(unix))]
183fn with_counter_lock<T>(path: &Path, f: impl FnOnce() -> Result<T>) -> Result<T> {
184    let sentinel = path.join("counter.lock");
185    loop {
186        match OpenOptions::new()
187            .write(true)
188            .create_new(true)
189            .open(&sentinel)
190        {
191            Ok(file) => {
192                let result = f();
193                drop(file);
194                fs::remove_file(&sentinel).ok();
195                return result;
196            }
197            Err(error) if error.kind() == ErrorKind::AlreadyExists => {
198                thread::sleep(Duration::from_millis(1));
199            }
200            Err(error) => return Err(error).context("cannot acquire counter lock"),
201        }
202    }
203}
204
205fn read_counter(path: &Path) -> Result<u64> {
206    match fs::read_to_string(path) {
207        Ok(value) => value
208            .trim()
209            .parse()
210            .with_context(|| format!("invalid counter in {}", path.display())),
211        Err(error) if error.kind() == ErrorKind::NotFound => Ok(0),
212        Err(error) => Err(error).with_context(|| format!("cannot read {}", path.display())),
213    }
214}
215
216fn write_counter(path: &Path, value: u64) -> Result<()> {
217    fs::write(path, format!("{value}\n"))
218        .with_context(|| format!("cannot write {}", path.display()))
219}
220
221fn read_pid(path: &Path) -> Result<Option<u32>> {
222    match fs::read_to_string(path) {
223        Ok(value) => {
224            Ok(Some(value.trim().parse().with_context(|| {
225                format!("invalid pid in {}", path.display())
226            })?))
227        }
228        Err(error) if error.kind() == ErrorKind::NotFound => Ok(None),
229        Err(error) => Err(error).with_context(|| format!("cannot read {}", path.display())),
230    }
231}
232
233#[cfg(unix)]
234fn pid_is_alive(pid: u32) -> bool {
235    unsafe extern "C" {
236        fn kill(pid: i32, signal: i32) -> i32;
237    }
238
239    let Ok(pid) = i32::try_from(pid) else {
240        return false;
241    };
242    // Signal 0 performs the existence/permission check without sending a signal.
243    if unsafe { kill(pid, 0) } == 0 {
244        return true;
245    }
246    std::io::Error::last_os_error().raw_os_error() != Some(3)
247}
248
249#[cfg(not(unix))]
250fn pid_is_alive(_pid: u32) -> bool {
251    true
252}