use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::Mutex;
use tokio::sync::OwnedMutexGuard;
use crate::Res;
type Registry = Mutex<HashMap<PathBuf, Arc<tokio::sync::Mutex<()>>>>;
static REGISTRY: LazyLock<Registry> = LazyLock::new(Registry::default);
#[derive(Debug)]
pub struct LockGuard {
_file: fs::File,
_held: OwnedMutexGuard<()>,
}
fn joined<T>(result: Result<std::io::Result<T>, tokio::task::JoinError>) -> Res<T> {
Ok(result.map_err(std::io::Error::other)??)
}
async fn open(path: &Path) -> Res<(fs::File, Arc<tokio::sync::Mutex<()>>)> {
let path = path.to_path_buf();
let (file, key) = joined(
tokio::task::spawn_blocking(move || {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
let file = fs::OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(&path)?;
let key = fs::canonicalize(&path)?;
Ok((file, key))
})
.await,
)?;
let mutex = Arc::clone(
REGISTRY
.lock()
.expect("lock registry")
.entry(key)
.or_default(),
);
Ok((file, mutex))
}
pub(crate) async fn lock_exclusive(path: &Path) -> Res<LockGuard> {
let (file, mutex) = open(path).await?;
let held = mutex.lock_owned().await;
let file = joined(
tokio::task::spawn_blocking(move || {
file.lock()?;
Ok(file)
})
.await,
)?;
Ok(LockGuard {
_file: file,
_held: held,
})
}
pub(crate) async fn try_lock_exclusive(path: &Path) -> Res<Option<LockGuard>> {
let (file, mutex) = open(path).await?;
let Ok(held) = mutex.try_lock_owned() else {
return Ok(None);
};
let file = joined(
tokio::task::spawn_blocking(move || match file.try_lock() {
Ok(()) => Ok(Some(file)),
Err(fs::TryLockError::WouldBlock) => Ok(None),
Err(fs::TryLockError::Error(err)) => Err(err),
})
.await,
)?;
Ok(file.map(|file| LockGuard {
_file: file,
_held: held,
}))
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use std::time::Duration;
use test_log::test;
#[test(tokio::test)]
async fn two_spellings_of_one_path_share_the_lock() -> Res {
let dir = tempfile::TempDir::new()?;
let plain = dir.path().join("data.json.lock");
let dotted = dir.path().join(".").join("data.json.lock");
let first = lock_exclusive(&plain).await?;
let second =
tokio::time::timeout(Duration::from_millis(200), lock_exclusive(&dotted)).await;
assert!(second.is_err(), "the second handle must wait");
drop(first);
tokio::time::timeout(Duration::from_secs(5), lock_exclusive(&dotted))
.await
.expect("the lock is free once the first guard drops")?;
Ok(())
}
#[test(tokio::test)]
async fn a_try_on_a_held_lock_is_busy() -> Res {
let dir = tempfile::TempDir::new()?;
let path = dir.path().join("demo.lock");
let first = lock_exclusive(&path).await?;
assert!(
try_lock_exclusive(&path).await?.is_none(),
"held in this process"
);
drop(first);
assert!(
try_lock_exclusive(&path).await?.is_some(),
"free once dropped"
);
Ok(())
}
#[test(tokio::test)]
async fn different_paths_do_not_contend() -> Res {
let dir = tempfile::TempDir::new()?;
let _a = lock_exclusive(&dir.path().join("a.lock")).await?;
tokio::time::timeout(
Duration::from_secs(5),
lock_exclusive(&dir.path().join("b.lock")),
)
.await
.expect("another path is free")?;
Ok(())
}
const CHILD_SIGNALS: &str = "QUILT_RS_TEST_CHILD_SIGNALS";
const CHILD_LOCK_FILE: &str = "QUILT_RS_TEST_CHILD_LOCK_FILE";
#[test]
#[ignore = "run only as the child of OtherProcess::holding"]
fn child_holds_the_lock() -> Res {
let (Ok(signals), Ok(lock_file)) =
(std::env::var(CHILD_SIGNALS), std::env::var(CHILD_LOCK_FILE))
else {
return Ok(());
};
let signals = PathBuf::from(signals);
let runtime = tokio::runtime::Runtime::new()?;
runtime.block_on(async {
let _held = lock_exclusive(Path::new(&lock_file)).await?;
fs::write(signals.join("locked"), b"")?;
while signals.join("hold").exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
Ok(())
})
}
pub(crate) struct OtherProcess {
child: std::process::Child,
signals: tempfile::TempDir,
}
impl OtherProcess {
pub(crate) async fn holding(lock_file: &Path) -> Res<Self> {
let signals = tempfile::TempDir::new()?;
fs::write(signals.path().join("hold"), b"")?;
let child = std::process::Command::new(std::env::current_exe()?)
.args([
"--exact",
"io::storage::lock::tests::child_holds_the_lock",
"--ignored",
"--nocapture",
])
.env(CHILD_SIGNALS, signals.path())
.env(CHILD_LOCK_FILE, lock_file)
.spawn()?;
let deadline = std::time::Instant::now() + Duration::from_secs(30);
while !signals.path().join("locked").exists() {
assert!(
std::time::Instant::now() < deadline,
"the child never took the lock"
);
tokio::time::sleep(Duration::from_millis(10)).await;
}
Ok(Self { child, signals })
}
pub(crate) fn release(mut self) -> Res {
fs::remove_file(self.signals.path().join("hold"))?;
assert!(self.child.wait()?.success());
Ok(())
}
}
#[test(tokio::test)]
async fn another_process_waits_for_the_lock() -> Res {
let dir = tempfile::TempDir::new()?;
let lock_path = dir.path().join("data.json.lock");
let other = OtherProcess::holding(&lock_path).await?;
let waiting = tokio::time::timeout(Duration::from_millis(300), lock_exclusive(&lock_path));
assert!(waiting.await.is_err(), "the other process holds the lock");
assert!(
try_lock_exclusive(&lock_path).await?.is_none(),
"a try sees the other process too"
);
other.release()?;
tokio::time::timeout(Duration::from_secs(30), lock_exclusive(&lock_path))
.await
.expect("the lock is free once the other process lets go")?;
Ok(())
}
}