use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
#[derive(Debug)]
pub enum LockError {
HeldBy {
pid: u32,
path: PathBuf,
},
Io(io::Error),
}
impl From<io::Error> for LockError {
fn from(e: io::Error) -> Self {
Self::Io(e)
}
}
impl std::fmt::Display for LockError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::HeldBy { pid, path } => {
write!(f, "pidfile {} held by live process {pid}", path.display())
}
Self::Io(e) => write!(f, "pidfile io: {e}"),
}
}
}
impl std::error::Error for LockError {}
#[derive(Debug)]
pub struct PidFile {
path: PathBuf,
pid: u32,
released: bool,
}
impl PidFile {
pub fn acquire(path: impl Into<PathBuf>) -> Result<Self, LockError> {
let path = path.into();
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(existing) = read_pidfile(&path) {
if is_process_alive(existing) && existing != std::process::id() {
return Err(LockError::HeldBy {
pid: existing,
path,
});
}
}
let pid = std::process::id();
let temp = path.with_extension("pid.tmp");
{
let mut f = fs::File::create(&temp)?;
writeln!(f, "{pid}")?;
f.sync_all()?;
}
fs::rename(&temp, &path)?;
Ok(Self {
path,
pid,
released: false,
})
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn pid(&self) -> u32 {
self.pid
}
pub fn release(mut self) -> io::Result<()> {
self.released = true;
remove_if_ours(&self.path, self.pid)
}
}
impl Drop for PidFile {
fn drop(&mut self) {
if self.released {
return;
}
let _ = remove_if_ours(&self.path, self.pid);
}
}
pub fn read_pidfile(path: &Path) -> Option<u32> {
let raw = fs::read_to_string(path).ok()?;
raw.trim().parse::<u32>().ok()
}
pub fn is_process_alive(pid: u32) -> bool {
if pid == 0 {
return false;
}
let rc = unsafe { libc::kill(pid as libc::pid_t, 0) };
if rc == 0 {
return true;
}
match io::Error::last_os_error().raw_os_error() {
Some(libc::EPERM) => true, _ => false,
}
}
fn remove_if_ours(path: &Path, our_pid: u32) -> io::Result<()> {
match read_pidfile(path) {
Some(pid) if pid == our_pid => fs::remove_file(path),
_ => Ok(()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
fn unique_tmp(label: &str) -> PathBuf {
static COUNTER: AtomicUsize = AtomicUsize::new(0);
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
std::env::temp_dir().join(format!("ao-rs-lock-{label}-{nanos}-{n}.pid"))
}
#[test]
fn acquire_when_no_file_writes_our_pid() {
let path = unique_tmp("fresh");
let lock = PidFile::acquire(&path).unwrap();
assert!(path.exists());
assert_eq!(read_pidfile(&path), Some(std::process::id()));
assert_eq!(lock.pid(), std::process::id());
drop(lock);
assert!(!path.exists(), "drop should remove the pidfile");
}
#[test]
fn acquire_replaces_stale_pidfile() {
let stale_pid: u32 = 999_999;
assert!(!is_process_alive(stale_pid), "sanity: {stale_pid} is dead");
let path = unique_tmp("stale");
fs::create_dir_all(path.parent().unwrap()).unwrap();
fs::write(&path, format!("{stale_pid}\n")).unwrap();
let lock = PidFile::acquire(&path).unwrap();
assert_eq!(read_pidfile(&path), Some(std::process::id()));
drop(lock);
}
#[test]
fn acquire_rejects_live_other_pid() {
let path = unique_tmp("held");
fs::create_dir_all(path.parent().unwrap()).unwrap();
fs::write(&path, "1\n").unwrap();
match PidFile::acquire(&path) {
Err(LockError::HeldBy { pid, .. }) => assert_eq!(pid, 1),
other => panic!("expected HeldBy(1), got {other:?}"),
}
assert_eq!(read_pidfile(&path), Some(1));
fs::remove_file(&path).ok();
}
#[test]
fn drop_does_not_remove_file_if_stolen() {
let path = unique_tmp("stolen");
let lock = PidFile::acquire(&path).unwrap();
fs::write(&path, "1\n").unwrap();
drop(lock);
assert_eq!(read_pidfile(&path), Some(1));
fs::remove_file(&path).ok();
}
#[test]
fn is_process_alive_returns_true_for_self() {
assert!(is_process_alive(std::process::id()));
}
#[test]
fn is_process_alive_returns_false_for_zero() {
assert!(!is_process_alive(0));
}
}