1use 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 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 reserved: u64,
37}
38
39static STORES: LazyLock<Mutex<Registry>> =
43 LazyLock::new(|| Mutex::new(Registry { stores: BTreeMap::new(), reserved: 0 }));
44
45const REGISTRY: &str = "the process's store lock registry";
50
51#[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#[derive(Debug)]
68pub struct StoreClosingGuard(Option<Arc<LockEntry>>);
69
70impl StoreClosingGuard {
71 pub fn closing_guard(&self) -> Self {
73 Self(self.0.clone())
74 }
75
76 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 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 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 _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 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
183fn 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
191pub(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 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
258pub(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 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 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 #[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
319pub(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
341pub(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;