trybuild_internals_api/
flock.rs

1use crate::error::Result;
2use std::fs::{self, File, OpenOptions};
3use std::io;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
7use std::thread;
8use std::time::{Duration, SystemTime};
9
10static LOCK: Mutex<()> = Mutex::new(());
11
12pub struct Lock {
13    intraprocess_guard: Guard,
14    lockfile: FileLock,
15}
16
17// High-quality lock to coordinate different #[test] functions within the *same*
18// integration test crate.
19enum Guard {
20    NotLocked,
21    Locked(#[allow(dead_code)] MutexGuard<'static, ()>),
22}
23
24// Best-effort filesystem lock to coordinate different #[test] functions across
25// *different* integration tests.
26enum FileLock {
27    NotLocked,
28    Locked {
29        path: PathBuf,
30        done: Arc<AtomicBool>,
31    },
32}
33
34impl Lock {
35    pub fn acquire(path: impl AsRef<Path>) -> Result<Self> {
36        Ok(Lock {
37            intraprocess_guard: Guard::acquire(),
38            lockfile: FileLock::acquire(path)?,
39        })
40    }
41}
42
43impl Guard {
44    fn acquire() -> Self {
45        Guard::Locked(LOCK.lock().unwrap_or_else(PoisonError::into_inner))
46    }
47}
48
49impl FileLock {
50    fn acquire(path: impl AsRef<Path>) -> Result<Self> {
51        let path = path.as_ref().to_owned();
52        let Some(lockfile) = create(&path) else {
53            return Ok(FileLock::NotLocked);
54        };
55        let done = Arc::new(AtomicBool::new(false));
56        let thread = thread::Builder::new().name("trybuild-flock".to_owned());
57        thread.spawn({
58            let done = Arc::clone(&done);
59            move || poll(lockfile, done)
60        })?;
61        Ok(FileLock::Locked { path, done })
62    }
63}
64
65impl Drop for Lock {
66    fn drop(&mut self) {
67        let Lock {
68            intraprocess_guard,
69            lockfile,
70        } = self;
71        // Unlock file lock first.
72        *lockfile = FileLock::NotLocked;
73        *intraprocess_guard = Guard::NotLocked;
74    }
75}
76
77impl Drop for FileLock {
78    fn drop(&mut self) {
79        match self {
80            FileLock::NotLocked => {}
81            FileLock::Locked { path, done } => {
82                done.store(true, Ordering::Release);
83                let _ = fs::remove_file(path);
84            }
85        }
86    }
87}
88
89fn create(path: &Path) -> Option<File> {
90    loop {
91        match OpenOptions::new().write(true).create_new(true).open(path) {
92            // Acquired lock by creating lockfile.
93            Ok(lockfile) => return Some(lockfile),
94            Err(io_error) => match io_error.kind() {
95                // Lock is already held by another test.
96                io::ErrorKind::AlreadyExists => {}
97                // File based locking isn't going to work for some reason.
98                _ => return None,
99            },
100        }
101
102        // Check whether it's okay to bust the lock.
103        let metadata = match fs::metadata(path) {
104            Ok(metadata) => metadata,
105            Err(io_error) => match io_error.kind() {
106                // Other holder of the lock finished. Retry.
107                io::ErrorKind::NotFound => continue,
108                _ => return None,
109            },
110        };
111
112        let Ok(modified) = metadata.modified() else {
113            return None;
114        };
115
116        let now = SystemTime::now();
117        let considered_stale = now - Duration::from_millis(1500);
118        let considered_future = now + Duration::from_millis(1500);
119        if modified < considered_stale || considered_future < modified {
120            return File::create(path).ok();
121        }
122
123        // Try again shortly.
124        thread::sleep(Duration::from_millis(500));
125    }
126}
127
128// Bump mtime periodically while test directory is in use.
129fn poll(lockfile: File, done: Arc<AtomicBool>) {
130    loop {
131        thread::sleep(Duration::from_millis(500));
132        if done.load(Ordering::Acquire) || lockfile.set_len(0).is_err() {
133            return;
134        }
135    }
136}