use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
use nix::fcntl::{Flock, FlockArg};
use nix::sys::signal::kill;
use nix::unistd::Pid;
use crate::error::{Error, Result};
const LOCK_FILE: &str = ".handoff.lock";
const PID_FILE: &str = ".handoff.pidfile";
pub struct DataDirLock {
_flock: Flock<File>,
data_dir: PathBuf,
pid_path: PathBuf,
}
impl std::fmt::Debug for DataDirLock {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DataDirLock")
.field("data_dir", &self.data_dir)
.finish()
}
}
impl DataDirLock {
pub fn data_dir(&self) -> &Path {
&self.data_dir
}
}
impl DataDirLock {
pub fn acquire(data_dir: &Path) -> Result<Self> {
std::fs::create_dir_all(data_dir)?;
let lock_path = data_dir.join(LOCK_FILE);
let pid_path = data_dir.join(PID_FILE);
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)?;
let flock = match Flock::lock(file, FlockArg::LockExclusiveNonblock) {
Ok(flock) => flock,
Err((_file, nix::errno::Errno::EWOULDBLOCK)) => {
let holder = read_pidfile(&pid_path).unwrap_or(0);
return Err(Error::LockHeld { holder_pid: holder });
}
Err((_file, errno)) => return Err(Error::Nix(errno)),
};
write_pid_atomic(&pid_path, std::process::id())?;
Ok(Self {
_flock: flock,
data_dir: data_dir.to_path_buf(),
pid_path,
})
}
pub fn acquire_or_break_stale(data_dir: &Path) -> Result<Self> {
match Self::acquire(data_dir) {
Ok(lock) => Ok(lock),
Err(Error::LockHeld { holder_pid }) => {
if holder_pid != 0 && is_pid_alive(holder_pid) {
return Err(Error::StaleLockBreakRefused { holder_pid });
}
tracing::warn!(
holder_pid,
"data-dir flock appears stale (named holder dead); retrying acquire"
);
match Self::acquire(data_dir) {
Ok(lock) => Ok(lock),
Err(Error::LockHeld { holder_pid }) => {
Err(Error::StaleLockBreakRefused { holder_pid })
}
Err(e) => Err(e),
}
}
Err(e) => Err(e),
}
}
}
impl Drop for DataDirLock {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.pid_path);
}
}
fn read_pidfile(path: &Path) -> Option<i32> {
std::fs::read_to_string(path).ok()?.trim().parse().ok()
}
fn write_pid_atomic(path: &Path, pid: u32) -> Result<()> {
let tmp = path.with_extension("pidfile.tmp");
{
let mut f = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp)?;
writeln!(f, "{pid}")?;
f.sync_all()?;
}
std::fs::rename(&tmp, path)?;
if let Some(parent) = path.parent() {
let target = if parent.as_os_str().is_empty() {
Path::new(".")
} else {
parent
};
File::open(target)?.sync_all()?;
}
Ok(())
}
fn is_pid_alive(pid: i32) -> bool {
if pid <= 0 {
return false;
}
matches!(kill(Pid::from_raw(pid), None), Ok(()))
}
#[cfg(test)]
mod tests {
use std::os::fd::AsRawFd;
use super::*;
#[test]
fn acquire_succeeds_on_empty_dir() {
let dir = tempfile::tempdir().unwrap();
let lock = DataDirLock::acquire(dir.path()).unwrap();
drop(lock);
}
#[test]
fn second_acquire_returns_lock_held() {
let dir = tempfile::tempdir().unwrap();
let _lock = DataDirLock::acquire(dir.path()).unwrap();
match DataDirLock::acquire(dir.path()) {
Err(Error::LockHeld { holder_pid }) => {
assert_eq!(holder_pid as u32, std::process::id());
}
other => panic!("expected LockHeld, got {other:?}"),
}
}
#[test]
fn release_on_drop_allows_reacquire() {
let dir = tempfile::tempdir().unwrap();
{
let _lock = DataDirLock::acquire(dir.path()).unwrap();
}
let _lock = DataDirLock::acquire(dir.path()).unwrap();
}
#[test]
fn stale_break_refuses_for_live_pid() {
let dir = tempfile::tempdir().unwrap();
let _held = DataDirLock::acquire(dir.path()).unwrap();
match DataDirLock::acquire_or_break_stale(dir.path()) {
Err(Error::StaleLockBreakRefused { holder_pid }) => {
assert_eq!(holder_pid as u32, std::process::id());
}
other => panic!("expected refusal, got {other:?}"),
}
}
#[test]
fn stale_break_succeeds_when_kernel_released_flock() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join(LOCK_FILE), b"").unwrap();
std::fs::write(dir.path().join(PID_FILE), format!("{}", i32::MAX)).unwrap();
let _new_lock = DataDirLock::acquire_or_break_stale(dir.path()).unwrap();
}
#[test]
fn stale_break_refuses_when_pidfile_lies_but_flock_held() {
let dir = tempfile::tempdir().unwrap();
let lock_path = dir.path().join(LOCK_FILE);
let pid_path = dir.path().join(PID_FILE);
let f = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)
.unwrap();
let _other_flock = Flock::lock(f, FlockArg::LockExclusiveNonblock)
.map_err(|(_, e)| e)
.unwrap();
std::fs::write(&pid_path, format!("{}", i32::MAX)).unwrap();
match DataDirLock::acquire_or_break_stale(dir.path()) {
Err(Error::StaleLockBreakRefused { .. }) => {}
other => panic!("expected StaleLockBreakRefused, got {other:?}"),
}
assert!(_other_flock.as_raw_fd() >= 0);
}
}