1use std::{
20 collections::HashMap,
21 fs::File,
22 path::{Path, PathBuf},
23 sync::{Arc, Condvar, Mutex, MutexGuard, OnceLock},
24 thread::{self, ThreadId},
25};
26
27use fs2::FileExt;
28pub use heddle_object_model::error::LockError;
29
30pub type Result<T> = std::result::Result<T, LockError>;
31
32struct GateState {
35 owner: Option<ThreadId>,
36 depth: usize,
37 flock: Option<File>,
38}
39
40struct Entry {
41 gate: Mutex<GateState>,
42 cv: Condvar,
43}
44
45impl Entry {
46 fn new() -> Self {
47 Self {
48 gate: Mutex::new(GateState {
49 owner: None,
50 depth: 0,
51 flock: None,
52 }),
53 cv: Condvar::new(),
54 }
55 }
56}
57
58static REGISTRY: OnceLock<Mutex<HashMap<PathBuf, Arc<Entry>>>> = OnceLock::new();
62
63fn registry() -> &'static Mutex<HashMap<PathBuf, Arc<Entry>>> {
64 REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
65}
66
67fn entry_for(key: PathBuf) -> Arc<Entry> {
68 let mut map = registry().lock().unwrap_or_else(|e| e.into_inner());
69 Arc::clone(map.entry(key).or_insert_with(|| Arc::new(Entry::new())))
70}
71
72fn lock_gate(entry: &Entry) -> MutexGuard<'_, GateState> {
73 entry.gate.lock().unwrap_or_else(|e| e.into_inner())
74}
75
76pub struct ReadLockGuard {
77 _file: Option<File>,
80}
81
82impl Drop for ReadLockGuard {
83 fn drop(&mut self) {
84 if let Some(file) = &self._file {
85 let _ = file.unlock();
86 }
87 }
88}
89
90pub struct WriteLockGuard {
91 entry: Arc<Entry>,
92}
93
94impl Drop for WriteLockGuard {
95 fn drop(&mut self) {
96 let mut state = lock_gate(&self.entry);
97 if state.depth > 0 {
98 state.depth -= 1;
99 }
100 if state.depth == 0 {
101 state.owner = None;
102 state.flock = None;
104 self.entry.cv.notify_one();
105 }
106 }
107}
108
109pub struct RepoLock {
110 lock_path: PathBuf,
111}
112
113impl RepoLock {
114 pub fn new(repo_root: &Path) -> Self {
115 let lock_path = repo_root.join(".heddle/locks/repo.lock");
116 Self { lock_path }
117 }
118
119 pub fn at(lock_path: PathBuf) -> Self {
120 Self { lock_path }
121 }
122
123 pub fn read(&self) -> Result<ReadLockGuard> {
124 self.ensure_lock_dir()?;
125 let entry = entry_for(self.registry_key());
126
127 {
131 let state = lock_gate(&entry);
132 if state.owner == Some(thread::current().id()) {
133 return Ok(ReadLockGuard { _file: None });
134 }
135 }
136
137 let file = self.open_lock_file()?;
138 file.lock_shared().map_err(LockError::Acquire)?;
139 Ok(ReadLockGuard { _file: Some(file) })
140 }
141
142 pub fn write(&self) -> Result<WriteLockGuard> {
143 self.ensure_lock_dir()?;
144 let entry = entry_for(self.registry_key());
145 let tid = thread::current().id();
146 let mut state = lock_gate(&entry);
147 loop {
148 match state.owner {
149 Some(owner) if owner == tid => {
150 state.depth += 1;
151 return Ok(WriteLockGuard {
152 entry: Arc::clone(&entry),
153 });
154 }
155 None => {
156 let file = self.open_lock_file()?;
161 file.lock_exclusive().map_err(LockError::Acquire)?;
162 state.owner = Some(tid);
163 state.depth = 1;
164 state.flock = Some(file);
165 return Ok(WriteLockGuard {
166 entry: Arc::clone(&entry),
167 });
168 }
169 Some(_) => {
170 state = entry.cv.wait(state).unwrap_or_else(|e| e.into_inner());
171 }
172 }
173 }
174 }
175
176 pub fn try_read(&self) -> Result<Option<ReadLockGuard>> {
177 self.ensure_lock_dir()?;
178 let file = self.open_lock_file()?;
179
180 match file.try_lock_shared() {
181 Ok(()) => Ok(Some(ReadLockGuard { _file: Some(file) })),
182 Err(_) => Ok(None),
183 }
184 }
185
186 pub fn try_write(&self) -> Result<Option<WriteLockGuard>> {
187 self.ensure_lock_dir()?;
188 let entry = entry_for(self.registry_key());
189 let mut state = lock_gate(&entry);
190 match state.owner {
197 Some(_) => Ok(None),
198 None => {
199 let file = self.open_lock_file()?;
200 match file.try_lock_exclusive() {
201 Ok(()) => {
202 state.owner = Some(thread::current().id());
203 state.depth = 1;
204 state.flock = Some(file);
205 Ok(Some(WriteLockGuard {
206 entry: Arc::clone(&entry),
207 }))
208 }
209 Err(_) => Ok(None),
210 }
211 }
212 }
213 }
214
215 fn ensure_lock_dir(&self) -> Result<()> {
216 if let Some(parent) = self.lock_path.parent() {
217 crate::fs_atomic::create_dir_all_durable(parent).map_err(LockError::Io)?;
218 }
219 Ok(())
220 }
221
222 fn registry_key(&self) -> PathBuf {
226 match self.lock_path.parent() {
227 Some(parent) => {
228 let canon_parent = parent
229 .canonicalize()
230 .unwrap_or_else(|_| parent.to_path_buf());
231 match self.lock_path.file_name() {
232 Some(name) => canon_parent.join(name),
233 None => canon_parent,
234 }
235 }
236 None => self.lock_path.clone(),
237 }
238 }
239
240 fn open_lock_file(&self) -> Result<File> {
241 File::create(&self.lock_path).map_err(LockError::Io)
242 }
243}
244
245pub trait RepositoryLockExt {
246 fn locker(&self) -> RepoLock;
247}
248
249#[cfg(test)]
250mod tests {
251 use std::{
252 sync::{
253 Arc,
254 mpsc::{self},
255 },
256 thread,
257 };
258
259 use tempfile::TempDir;
260
261 use super::*;
262
263 #[test]
264 fn test_read_lock_acquired() {
265 let temp = TempDir::new().unwrap();
266 let lock = RepoLock::new(temp.path());
267
268 let guard = lock.read().unwrap();
269 assert!(std::mem::size_of_val(&guard) > 0);
270 }
271
272 #[test]
273 fn test_write_lock_acquired() {
274 let temp = TempDir::new().unwrap();
275 let lock = RepoLock::new(temp.path());
276
277 let guard = lock.write().unwrap();
278 assert!(std::mem::size_of_val(&guard) > 0);
279 }
280
281 #[test]
282 fn test_multiple_readers() {
283 let temp = TempDir::new().unwrap();
284 let lock = Arc::new(RepoLock::new(temp.path()));
285
286 let mut handles = vec![];
287 for _ in 0..10 {
288 let lock = Arc::clone(&lock);
289 let handle = thread::spawn(move || {
290 let _guard = lock.read().unwrap();
291 thread::sleep(std::time::Duration::from_millis(10));
292 });
293 handles.push(handle);
294 }
295
296 for handle in handles {
297 handle.join().unwrap();
298 }
299 }
300
301 #[test]
302 fn test_writer_excludes_reader() {
303 let temp = TempDir::new().unwrap();
304 let lock = Arc::new(RepoLock::new(temp.path()));
305
306 let _write_guard = lock.write().unwrap();
307 let read_result = lock.try_read().unwrap();
308 assert!(read_result.is_none(), "Reader should be blocked by writer");
309 }
310
311 #[test]
312 fn test_reader_excludes_writer() {
313 let temp = TempDir::new().unwrap();
314 let lock = Arc::new(RepoLock::new(temp.path()));
315
316 let _read_guard = lock.read().unwrap();
317 let write_result = lock.try_write().unwrap();
318 assert!(write_result.is_none(), "Writer should be blocked by reader");
319 }
320
321 #[test]
322 fn test_lock_released_on_drop() {
323 let temp = TempDir::new().unwrap();
324 let lock = RepoLock::new(temp.path());
325
326 {
327 let _guard = lock.write().unwrap();
328 }
329
330 let _guard2 = lock.read().unwrap();
331 }
332
333 #[test]
336 fn same_thread_write_is_reentrant() {
337 let temp = TempDir::new().unwrap();
338 let lock = RepoLock::new(temp.path());
339
340 let _a = lock.write().unwrap();
341 let _b = lock.write().unwrap();
342 }
345
346 #[test]
349 fn same_thread_read_under_write_does_not_deadlock() {
350 let temp = TempDir::new().unwrap();
351 let lock = RepoLock::new(temp.path());
352
353 let _w = lock.write().unwrap();
354 let _r = lock.read().unwrap();
355 }
356
357 #[test]
360 fn distinct_threads_still_exclude() {
361 let temp = TempDir::new().unwrap();
362 let lock = Arc::new(RepoLock::new(temp.path()));
363
364 let (acquired_tx, acquired_rx) = mpsc::channel();
365 let (release_tx, release_rx) = mpsc::channel();
366 let lock_a = Arc::clone(&lock);
367 let handle = thread::spawn(move || {
368 let _g = lock_a.write().unwrap();
369 acquired_tx.send(()).unwrap();
370 release_rx.recv().unwrap();
371 });
372
373 acquired_rx.recv().unwrap();
374 assert!(
375 lock.try_write().unwrap().is_none(),
376 "a second thread must not acquire the write lock"
377 );
378
379 release_tx.send(()).unwrap();
380 handle.join().unwrap();
381
382 assert!(
383 lock.try_write().unwrap().is_some(),
384 "write lock is available once the owning thread releases"
385 );
386 }
387
388 #[test]
391 fn reentrant_release_keeps_lock_until_outermost_drop() {
392 let temp = TempDir::new().unwrap();
393 let lock = Arc::new(RepoLock::new(temp.path()));
394
395 let a1 = lock.write().unwrap();
396 let a2 = lock.write().unwrap();
397
398 let other = |lock: &Arc<RepoLock>| {
399 let lock = Arc::clone(lock);
400 thread::spawn(move || lock.try_write().unwrap().is_none())
401 .join()
402 .unwrap()
403 };
404
405 assert!(other(&lock), "excluded while held at depth 2");
406 drop(a2);
407 assert!(other(&lock), "still excluded while held at depth 1");
408 drop(a1);
409
410 let lock_b = Arc::clone(&lock);
411 let now_available = thread::spawn(move || lock_b.try_write().unwrap().is_some())
412 .join()
413 .unwrap();
414 assert!(now_available, "available after the outermost guard drops");
415 }
416
417 #[test]
425 fn try_write_is_non_reentrant_even_for_owner() {
426 let temp = TempDir::new().unwrap();
427 let lock = RepoLock::new(temp.path());
428
429 let _held = lock.write().unwrap();
430 assert!(
431 lock.try_write().unwrap().is_none(),
432 "try_write must report contention even for the lock's own owner thread"
433 );
434 }
435}