Skip to main content

coven_foundation/files/
lock_registry.rs

1//! Process-local lock lifetimes. Only closing instances cause an async wait.
2
3use super::file_lock::FileLock;
4use super::lock::{lock_bootstrap, lock_existing, StoreLockError};
5use crate::id_source::StoreId;
6use crate::task::{blocking_completing, BugSinkRef};
7use std::{
8    collections::BTreeMap,
9    path::{Path, PathBuf},
10    sync::{Arc, LazyLock, Mutex},
11};
12use tokio::sync::watch;
13
14#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
15pub(super) enum LockKind {
16    Writer,
17    Bootstrap,
18    /// A shared readers' lock: any number of instances hold it at once, and
19    /// no writer or bootstrap waits on one closing.
20    Reader,
21}
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24enum LockPhase {
25    Live,
26    Closing,
27    Released,
28}
29
30type StoreKey = (PathBuf, StoreId);
31
32#[derive(Debug)]
33struct Registry {
34    stores: BTreeMap<StoreKey, BTreeMap<(LockKind, LockInstance), Arc<LockEntry>>>,
35    /// How many locks were reserved: each reservation is the next instance.
36    reserved: u64,
37}
38
39// Independently constructed layouts must see the same closing instance. The
40// registry mutex serializes classification and reservation. A reservation stays
41// registered until its acquisition worker finishes and its OS handle closes.
42static STORES: LazyLock<Mutex<Registry>> =
43    LazyLock::new(|| Mutex::new(Registry { stores: BTreeMap::new(), reserved: 0 }));
44
45/// What a lock of `STORES` names. The registry is shared by every store
46/// this process opens, and a change to it checks the registry's rules as it
47/// goes, so a panic can leave it half-changed, and no one store's handle
48/// covers it: the process ends (api.md E15).
49const REGISTRY: &str = "the process's store lock registry";
50
51/// One reservation of a lock: no other reservation in the process, of any
52/// store, has the same instance.
53#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
54pub(super) struct LockInstance(u64);
55
56#[derive(Debug)]
57struct LockEntry {
58    store: StoreKey,
59    kind: LockKind,
60    instance: LockInstance,
61    phase: watch::Sender<LockPhase>,
62}
63
64/// Marks an owning scope as closing on drop, without retaining its lock.
65/// Share an `Arc` of this guard across clones of that owner. A distinct owner
66/// scope can use `closing_guard`; ending either scope marks the instance.
67#[derive(Debug)]
68pub struct StoreClosingGuard(Option<Arc<LockEntry>>);
69
70impl StoreClosingGuard {
71    /// A separate owning scope for the same lock instance.
72    pub fn closing_guard(&self) -> Self {
73        Self(self.0.clone())
74    }
75
76    /// Begin closing before awaiting shutdown of workers and connections.
77    /// A guard from a released instance never changes a replacement instance.
78    pub fn begin_close(&self) {
79        let Some(entry) = &self.0 else { return };
80        let registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
81        if *entry.phase.borrow() == LockPhase::Released {
82            return;
83        }
84        let Some(locks) = registry.stores.get(&entry.store) else {
85            crate::locks::abort_broken(REGISTRY, "a held lock is registered")
86        };
87        // A dropped bootstrap owns the installation, including a writer that
88        // may still be opening or running blocking work on its behalf. Any
89        // other lock closes only its own instance: one reader closing leaves
90        // the others live.
91        for held in locks.values() {
92            if entry.kind == LockKind::Bootstrap || Arc::ptr_eq(held, entry) {
93                held.phase.send_replace(LockPhase::Closing);
94            }
95        }
96    }
97
98    /// Transfer lifetime responsibility to the successfully constructed owner.
99    pub fn disarm(mut self) {
100        self.0.take();
101    }
102}
103
104impl Drop for StoreClosingGuard {
105    fn drop(&mut self) {
106        self.begin_close();
107    }
108}
109
110#[derive(Debug)]
111pub(super) struct RegisteredLock {
112    // Field order makes notification follow both unlock and file close.
113    _file: FileLock,
114    registration: LockRegistration,
115}
116
117#[derive(Debug)]
118struct LockRegistration {
119    entry: Arc<LockEntry>,
120}
121
122impl RegisteredLock {
123    pub(super) fn closing_guard(&self) -> StoreClosingGuard {
124        self.registration.closing_guard()
125    }
126
127    pub(super) fn instance(&self) -> LockInstance {
128        self.registration.entry.instance
129    }
130
131    #[cfg(test)]
132    pub(super) fn inherited_file(&self) -> std::fs::File {
133        self._file.inherited_file()
134    }
135}
136
137impl LockRegistration {
138    fn closing_guard(&self) -> StoreClosingGuard {
139        StoreClosingGuard(Some(self.entry.clone()))
140    }
141
142    fn acquire_file(self) -> Result<RegisteredLock, StoreLockError> {
143        let (directory, id) = &self.entry.store;
144        let file = match self.entry.kind {
145            LockKind::Writer => lock_existing(directory, *id, false)?.1,
146            LockKind::Reader => lock_existing(directory, *id, true)?.1,
147            LockKind::Bootstrap => lock_bootstrap(directory, *id)?,
148        };
149        Ok(RegisteredLock { _file: file, registration: self })
150    }
151
152    /// Take the lock file on the blocking pool. A caller that stops waiting
153    /// leaves the work to finish, and its panic goes to `bugs` (E15).
154    async fn open_on_worker(self, bugs: &BugSinkRef) -> Result<RegisteredLock, StoreLockError> {
155        let closing = self.closing_guard();
156        let result =
157            blocking_completing("lock the store", bugs.clone(), move || self.acquire_file())
158                .await
159                .map_err(|failed| StoreLockError::from(crate::files::FileError::from(failed)))
160                .and_then(|result| result);
161        closing.disarm();
162        result
163    }
164}
165
166impl Drop for LockRegistration {
167    fn drop(&mut self) {
168        let mut registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
169        let Some(locks) = registry.stores.get_mut(&self.entry.store) else {
170            crate::locks::abort_broken(REGISTRY, "a held lock is registered")
171        };
172        match locks.remove(&(self.entry.kind, self.entry.instance)) {
173            Some(removed) if Arc::ptr_eq(&removed, &self.entry) => {}
174            _ => crate::locks::abort_broken(REGISTRY, "a held lock is registered as itself"),
175        }
176        if locks.is_empty() {
177            registry.stores.remove(&self.entry.store);
178        }
179        self.entry.phase.send_replace(LockPhase::Released);
180    }
181}
182
183/// The registered instances of `kind` among one store's locks.
184fn held(
185    locks: &BTreeMap<(LockKind, LockInstance), Arc<LockEntry>>,
186    kind: LockKind,
187) -> impl Iterator<Item = &Arc<LockEntry>> {
188    locks.iter().filter(move |((held, _), _)| *held == kind).map(|(_, entry)| entry)
189}
190
191/// Take a shared readers' lock now, registered so that a deletion in this
192/// process waits for it once its holder begins closing.
193pub(super) fn acquire_reader(
194    directory: &Path,
195    id: StoreId,
196) -> Result<RegisteredLock, StoreLockError> {
197    let registration = {
198        let mut registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
199        reserve(&mut registry, (directory.to_owned(), id), LockKind::Reader)
200    };
201    registration.acquire_file()
202}
203
204fn reserve(registry: &mut Registry, store: StoreKey, kind: LockKind) -> LockRegistration {
205    reserve_in(registry, store, kind, LockPhase::Live)
206}
207
208fn reserve_in(
209    registry: &mut Registry,
210    store: StoreKey,
211    kind: LockKind,
212    phase: LockPhase,
213) -> LockRegistration {
214    registry.reserved += 1;
215    let entry = Arc::new(LockEntry {
216        store: store.clone(),
217        kind,
218        instance: LockInstance(registry.reserved),
219        phase: watch::channel(phase).0,
220    });
221    let key = (kind, entry.instance);
222    assert!(registry.stores.entry(store).or_default().insert(key, entry.clone()).is_none());
223    LockRegistration { entry }
224}
225
226#[cfg(any(test, feature = "test-utils"))]
227pub(super) fn try_acquire(
228    directory: &Path,
229    id: StoreId,
230    kind: LockKind,
231) -> Result<RegisteredLock, StoreLockError> {
232    let registration = {
233        let mut registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
234        let store = (directory.to_owned(), id);
235        if registry.stores.get(&store).is_some_and(|locks| held(locks, kind).next().is_some()) {
236            return Err(StoreLockError::AlreadyOpen(id));
237        }
238        reserve(&mut registry, store, kind)
239    };
240    registration.acquire_file()
241}
242
243enum Acquisition {
244    /// The kind of the closing instance it waits for, and its phase.
245    Waiting(LockKind, watch::Receiver<LockPhase>),
246    Reserved(LockRegistration),
247}
248
249pub(super) async fn acquire_registered(
250    directory: &Path,
251    id: StoreId,
252    kind: LockKind,
253    bugs: &BugSinkRef,
254) -> Result<RegisteredLock, StoreLockError> {
255    acquire_registered_in(directory, id, kind, LockPhase::Live, bugs).await
256}
257
258/// Take the writer lock as `acquire_registered` does, registered as closing
259/// from the start: its holder only reads the closed store and lets go, so an
260/// open in this process waits for it rather than refusing.
261pub(super) async fn acquire_closing_writer(
262    directory: &Path,
263    id: StoreId,
264    bugs: &BugSinkRef,
265) -> Result<RegisteredLock, StoreLockError> {
266    acquire_registered_in(directory, id, LockKind::Writer, LockPhase::Closing, bugs).await
267}
268
269async fn acquire_registered_in(
270    directory: &Path,
271    id: StoreId,
272    kind: LockKind,
273    phase: LockPhase,
274    bugs: &BugSinkRef,
275) -> Result<RegisteredLock, StoreLockError> {
276    let store = (directory.to_owned(), id);
277    loop {
278        let acquisition = {
279            let mut registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
280            let locks = registry.stores.get(&store);
281            // A live holder of the requested lock always refuses immediately,
282            // even if a different lock on this directory is closing.
283            if locks.is_some_and(|locks| {
284                held(locks, kind).any(|entry| *entry.phase.borrow() == LockPhase::Live)
285            }) {
286                return Err(StoreLockError::AlreadyOpen(id));
287            }
288            // Readers share the store with a writer and a bootstrap, so
289            // neither waits on one closing.
290            match locks.and_then(|locks| {
291                locks.values().find(|entry| {
292                    entry.kind != LockKind::Reader && *entry.phase.borrow() == LockPhase::Closing
293                })
294            }) {
295                Some(entry) => Acquisition::Waiting(entry.kind, entry.phase.subscribe()),
296                None => {
297                    Acquisition::Reserved(reserve_in(&mut registry, store.clone(), kind, phase))
298                }
299            }
300        };
301        let mut released = match acquisition {
302            Acquisition::Reserved(registration) => return registration.open_on_worker(bugs).await,
303            Acquisition::Waiting(closing, released) => {
304                // Told at the lock file of the instance it waits for.
305                #[cfg(any(test, feature = "test-utils"))]
306                super::reached::reached(&super::lock::kind_lock(directory, id, closing));
307                #[cfg(not(any(test, feature = "test-utils")))]
308                let _: LockKind = closing;
309                released
310            }
311        };
312        released
313            .wait_for(|phase| *phase == LockPhase::Released)
314            .await
315            .expect("registered lock publishes release before dropping its sender");
316    }
317}
318
319/// Wait until no instance of the store at `directory` in this process is
320/// closing, a read-only one included. A live instance is not waited for:
321/// whatever lock the caller takes next refuses it.
322pub(super) async fn closed(directory: &Path, id: StoreId) {
323    let store = (directory.to_owned(), id);
324    loop {
325        let mut released = {
326            let registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
327            match registry.stores.get(&store).and_then(|locks| {
328                locks.values().find(|entry| *entry.phase.borrow() == LockPhase::Closing)
329            }) {
330                Some(entry) => entry.phase.subscribe(),
331                None => return,
332            }
333        };
334        released
335            .wait_for(|phase| *phase == LockPhase::Released)
336            .await
337            .expect("registered lock publishes release before dropping its sender");
338    }
339}
340
341/// Whether `instance` of the store's bootstrap lock is held, and its holder
342/// not closing: the bootstrap that took it is running.
343pub(super) fn bootstrap_running(directory: &Path, id: StoreId, instance: LockInstance) -> bool {
344    let registry = crate::locks::lock_or_abort(&STORES, REGISTRY);
345    registry
346        .stores
347        .get(&(directory.to_owned(), id))
348        .and_then(|locks| locks.get(&(LockKind::Bootstrap, instance)))
349        .is_some_and(|entry| *entry.phase.borrow() == LockPhase::Live)
350}
351
352#[cfg(test)]
353#[path = "lock_registry_tests.rs"]
354mod tests;