trybuild/
flock.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
use crate::error::Result;
use std::fs::{self, File, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::thread;
use std::time::{Duration, SystemTime};

static LOCK: Mutex<()> = Mutex::new(());

pub(crate) struct Lock {
    intraprocess_guard: Guard,
    lockfile: FileLock,
}

// High-quality lock to coordinate different #[test] functions within the *same*
// integration test crate.
enum Guard {
    NotLocked,
    Locked(#[allow(dead_code)] MutexGuard<'static, ()>),
}

// Best-effort filesystem lock to coordinate different #[test] functions across
// *different* integration tests.
enum FileLock {
    NotLocked,
    Locked {
        path: PathBuf,
        done: Arc<AtomicBool>,
    },
}

impl Lock {
    pub fn acquire(path: impl AsRef<Path>) -> Result<Self> {
        Ok(Lock {
            intraprocess_guard: Guard::acquire(),
            lockfile: FileLock::acquire(path)?,
        })
    }
}

impl Guard {
    fn acquire() -> Self {
        Guard::Locked(LOCK.lock().unwrap_or_else(PoisonError::into_inner))
    }
}

impl FileLock {
    fn acquire(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_owned();
        let lockfile = match create(&path) {
            None => return Ok(FileLock::NotLocked),
            Some(lockfile) => lockfile,
        };
        let done = Arc::new(AtomicBool::new(false));
        let thread = thread::Builder::new().name("trybuild-flock".to_owned());
        thread.spawn({
            let done = Arc::clone(&done);
            move || poll(lockfile, done)
        })?;
        Ok(FileLock::Locked { path, done })
    }
}

impl Drop for Lock {
    fn drop(&mut self) {
        let Lock {
            intraprocess_guard,
            lockfile,
        } = self;
        // Unlock file lock first.
        *lockfile = FileLock::NotLocked;
        *intraprocess_guard = Guard::NotLocked;
    }
}

impl Drop for FileLock {
    fn drop(&mut self) {
        match self {
            FileLock::NotLocked => {}
            FileLock::Locked { path, done } => {
                done.store(true, Ordering::Release);
                let _ = fs::remove_file(path);
            }
        }
    }
}

fn create(path: &Path) -> Option<File> {
    loop {
        match OpenOptions::new().write(true).create_new(true).open(path) {
            // Acquired lock by creating lockfile.
            Ok(lockfile) => return Some(lockfile),
            Err(io_error) => match io_error.kind() {
                // Lock is already held by another test.
                io::ErrorKind::AlreadyExists => {}
                // File based locking isn't going to work for some reason.
                _ => return None,
            },
        }

        // Check whether it's okay to bust the lock.
        let metadata = match fs::metadata(path) {
            Ok(metadata) => metadata,
            Err(io_error) => match io_error.kind() {
                // Other holder of the lock finished. Retry.
                io::ErrorKind::NotFound => continue,
                _ => return None,
            },
        };

        let Ok(modified) = metadata.modified() else {
            return None;
        };

        let now = SystemTime::now();
        let considered_stale = now - Duration::from_millis(1500);
        let considered_future = now + Duration::from_millis(1500);
        if modified < considered_stale || considered_future < modified {
            return File::create(path).ok();
        }

        // Try again shortly.
        thread::sleep(Duration::from_millis(500));
    }
}

// Bump mtime periodically while test directory is in use.
fn poll(lockfile: File, done: Arc<AtomicBool>) {
    loop {
        thread::sleep(Duration::from_millis(500));
        if done.load(Ordering::Acquire) || lockfile.set_len(0).is_err() {
            return;
        }
    }
}