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 lockfile = match create(&path) {
53            None => return Ok(FileLock::NotLocked),
54            Some(lockfile) => lockfile,
55        };
56        let done = Arc::new(AtomicBool::new(false));
57        let thread = thread::Builder::new().name("trybuild-flock".to_owned());
58        thread.spawn({
59            let done = Arc::clone(&done);
60            move || poll(lockfile, done)
61        })?;
62        Ok(FileLock::Locked { path, done })
63    }
64}
65
66impl Drop for Lock {
67    fn drop(&mut self) {
68        let Lock {
69            intraprocess_guard,
70            lockfile,
71        } = self;
72        // Unlock file lock first.
73        *lockfile = FileLock::NotLocked;
74        *intraprocess_guard = Guard::NotLocked;
75    }
76}
77
78impl Drop for FileLock {
79    fn drop(&mut self) {
80        match self {
81            FileLock::NotLocked => {}
82            FileLock::Locked { path, done } => {
83                done.store(true, Ordering::Release);
84                let _ = fs::remove_file(path);
85            }
86        }
87    }
88}
89
90fn create(path: &Path) -> Option<File> {
91    loop {
92        match OpenOptions::new().write(true).create_new(true).open(path) {
93            // Acquired lock by creating lockfile.
94            Ok(lockfile) => return Some(lockfile),
95            Err(io_error) => match io_error.kind() {
96                // Lock is already held by another test.
97                io::ErrorKind::AlreadyExists => {}
98                // File based locking isn't going to work for some reason.
99                _ => return None,
100            },
101        }
102
103        // Check whether it's okay to bust the lock.
104        let metadata = match fs::metadata(path) {
105            Ok(metadata) => metadata,
106            Err(io_error) => match io_error.kind() {
107                // Other holder of the lock finished. Retry.
108                io::ErrorKind::NotFound => continue,
109                _ => return None,
110            },
111        };
112
113        let Ok(modified) = metadata.modified() else {
114            return None;
115        };
116
117        let now = SystemTime::now();
118        let considered_stale = now - Duration::from_millis(1500);
119        let considered_future = now + Duration::from_millis(1500);
120        if modified < considered_stale || considered_future < modified {
121            return File::create(path).ok();
122        }
123
124        // Try again shortly.
125        thread::sleep(Duration::from_millis(500));
126    }
127}
128
129// Bump mtime periodically while test directory is in use.
130fn poll(lockfile: File, done: Arc<AtomicBool>) {
131    loop {
132        thread::sleep(Duration::from_millis(500));
133        if done.load(Ordering::Acquire) || lockfile.set_len(0).is_err() {
134            return;
135        }
136    }
137}