use std::fs::OpenOptions;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use crate::error::{Error, Result};
pub const WRITER_LOCK: &str = "writer.lock";
pub const WATCH_LEASE: &str = "watch.lock";
#[must_use]
#[allow(unsafe_code)]
pub fn pid_alive(pid: i64) -> bool {
let Ok(pid) = libc::pid_t::try_from(pid) else {
return false;
};
if pid <= 0 {
return false;
}
let rc = unsafe { libc::kill(pid, 0) };
if rc == 0 {
return true;
}
std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
#[allow(unsafe_code)]
pub(crate) fn send_sigterm(pid: i64) {
let Ok(pid) = libc::pid_t::try_from(pid) else {
return;
};
if pid <= 0 {
return;
}
unsafe {
libc::kill(pid, libc::SIGTERM);
}
}
fn now_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
.unwrap_or(0)
}
#[must_use]
pub fn read_holder_pid(path: &Path) -> Option<i64> {
let text = std::fs::read_to_string(path).ok()?;
let v: serde_json::Value = serde_json::from_str(&text).ok()?;
v.get("pid")?.as_i64()
}
fn try_create(path: &Path) -> Result<bool> {
match OpenOptions::new().write(true).create_new(true).open(path) {
Ok(mut f) => {
let record = serde_json::json!({ "pid": std::process::id(), "ts": now_ms() });
f.write_all(record.to_string().as_bytes())
.map_err(|e| Error::io("cannot write lock", path, e))?;
Ok(true)
}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
let holder = read_holder_pid(path);
let stale = holder.is_none_or(|pid| !pid_alive(pid));
if !stale {
return Ok(false);
}
if std::fs::remove_file(path).is_err() {
return Ok(false);
}
match OpenOptions::new().write(true).create_new(true).open(path) {
Ok(mut f) => {
let record = serde_json::json!({ "pid": std::process::id(), "ts": now_ms() });
f.write_all(record.to_string().as_bytes())
.map_err(|e| Error::io("cannot write lock", path, e))?;
Ok(true)
}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(false),
Err(e) => Err(Error::io("cannot create lock", path, e)),
}
}
Err(e) => Err(Error::io("cannot create lock", path, e)),
}
}
fn ensure_dir(dir: &Path) -> Result<()> {
std::fs::create_dir_all(dir).map_err(|e| Error::io("cannot create", dir, e))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WriterLockOptions {
pub timeout: Duration,
pub poll: Duration,
}
impl Default for WriterLockOptions {
fn default() -> Self {
Self {
timeout: Duration::from_millis(5000),
poll: Duration::from_millis(25),
}
}
}
#[derive(Debug)]
pub struct WriterLock {
path: PathBuf,
held: bool,
}
impl WriterLock {
#[must_use]
pub fn path_in(omgbase_dir: &Path) -> PathBuf {
omgbase_dir.join(WRITER_LOCK)
}
pub fn try_acquire(omgbase_dir: &Path) -> Result<Option<Self>> {
ensure_dir(omgbase_dir)?;
let path = Self::path_in(omgbase_dir);
Ok(try_create(&path)?.then(|| Self { path, held: true }))
}
pub fn acquire(omgbase_dir: &Path, opts: WriterLockOptions) -> Result<Self> {
ensure_dir(omgbase_dir)?;
let path = Self::path_in(omgbase_dir);
let deadline = Instant::now() + opts.timeout;
loop {
if try_create(&path)? {
return Ok(Self { path, held: true });
}
if Instant::now() >= deadline {
return Err(Error::WriterLockTimeout {
holder_pid: read_holder_pid(&path),
lock_path: path,
});
}
std::thread::sleep(opts.poll);
}
}
#[must_use]
pub fn is_free(omgbase_dir: &Path) -> bool {
let path = Self::path_in(omgbase_dir);
if !path.exists() {
return true;
}
read_holder_pid(&path).is_none_or(|pid| !pid_alive(pid))
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
pub fn release(&mut self) {
if !self.held {
return;
}
self.held = false;
let _ = std::fs::remove_file(&self.path);
}
}
impl Drop for WriterLock {
fn drop(&mut self) {
self.release();
}
}
pub fn with_writer_lock<T>(
omgbase_dir: &Path,
opts: WriterLockOptions,
f: impl FnOnce() -> Result<T>,
) -> Result<T> {
let mut lock = WriterLock::acquire(omgbase_dir, opts)?;
let out = f();
lock.release();
out
}
#[derive(Debug)]
pub struct WatchLease {
path: PathBuf,
held: bool,
}
impl WatchLease {
#[must_use]
pub fn path_in(omgbase_dir: &Path) -> PathBuf {
omgbase_dir.join(WATCH_LEASE)
}
pub fn try_acquire(omgbase_dir: &Path) -> Result<Option<Self>> {
ensure_dir(omgbase_dir)?;
let path = Self::path_in(omgbase_dir);
Ok(try_create(&path)?.then(|| Self { path, held: true }))
}
#[must_use]
pub fn live(omgbase_dir: &Path) -> bool {
let path = Self::path_in(omgbase_dir);
if !path.exists() {
return false;
}
read_holder_pid(&path).is_some_and(pid_alive)
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
pub fn release(&mut self) {
if !self.held {
return;
}
self.held = false;
let _ = std::fs::remove_file(&self.path);
}
}
impl Drop for WatchLease {
fn drop(&mut self) {
self.release();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fs::TempDir;
fn dead_pid() -> i64 {
let mut child = std::process::Command::new("true")
.spawn()
.expect("spawn true");
let pid = i64::from(child.id());
child.wait().unwrap();
pid
}
#[test]
fn liveness_probe() {
assert!(pid_alive(i64::from(std::process::id())));
assert!(pid_alive(1), "pid 1 is alive (EPERM counts as alive)");
assert!(!pid_alive(dead_pid()));
assert!(!pid_alive(0));
assert!(!pid_alive(-5));
assert!(!pid_alive(i64::MAX));
}
#[test]
fn writer_lock_round_trip_and_timeout() {
let tmp = TempDir::new("writer");
let dir = tmp.path().join(".omgbase");
assert!(WriterLock::is_free(&dir));
let fast = WriterLockOptions {
timeout: Duration::from_millis(80),
poll: Duration::from_millis(10),
};
{
let lock = WriterLock::acquire(&dir, fast).unwrap();
assert!(lock.path().is_file());
let body: serde_json::Value =
serde_json::from_str(&std::fs::read_to_string(lock.path()).unwrap()).unwrap();
assert_eq!(body["pid"], std::process::id());
assert!(body["ts"].is_number());
assert!(!WriterLock::is_free(&dir));
assert!(WriterLock::try_acquire(&dir).unwrap().is_none());
let err = WriterLock::acquire(&dir, fast).unwrap_err();
match err {
Error::WriterLockTimeout {
holder_pid,
lock_path,
} => {
assert_eq!(holder_pid, Some(i64::from(std::process::id())));
assert_eq!(lock_path, lock.path());
}
other => panic!("{other}"),
}
assert!(err_string_names_pid(
&WriterLock::acquire(&dir, fast).unwrap_err()
));
}
assert!(WriterLock::is_free(&dir), "released on drop");
assert!(!WriterLock::path_in(&dir).exists());
let r: Result<()> = with_writer_lock(&dir, fast, || Err(Error::Other("boom".into())));
assert!(r.is_err());
assert!(WriterLock::is_free(&dir));
assert_eq!(with_writer_lock(&dir, fast, || Ok(7)).unwrap(), 7);
std::fs::write(
WriterLock::path_in(&dir),
format!("{{\"pid\":{}}}", dead_pid()),
)
.unwrap();
assert!(WriterLock::is_free(&dir));
let l = WriterLock::try_acquire(&dir).unwrap().expect("stolen");
drop(l);
std::fs::write(WriterLock::path_in(&dir), "garbage").unwrap();
assert!(WriterLock::is_free(&dir));
assert!(WriterLock::try_acquire(&dir).unwrap().is_some());
assert!(WriterLock::is_free(&dir));
let mut l = WriterLock::try_acquire(&dir).unwrap().unwrap();
l.release();
l.release();
}
fn err_string_names_pid(e: &Error) -> bool {
e.to_string()
.contains(&format!("held by pid {}", std::process::id()))
}
#[test]
fn watch_lease_is_try_only_and_probeable() {
let tmp = TempDir::new("lease");
let dir = tmp.path().join(".omgbase");
assert!(!WatchLease::live(&dir));
let lease = WatchLease::try_acquire(&dir).unwrap().expect("free");
assert!(WatchLease::live(&dir));
assert!(WatchLease::try_acquire(&dir).unwrap().is_none());
assert_eq!(
read_holder_pid(lease.path()),
Some(i64::from(std::process::id()))
);
drop(lease);
assert!(!WatchLease::live(&dir));
std::fs::write(
WatchLease::path_in(&dir),
format!("{{\"pid\":{},\"ts\":0}}", dead_pid()),
)
.unwrap();
assert!(!WatchLease::live(&dir), "a dead holder is not live");
let mut l = WatchLease::try_acquire(&dir)
.unwrap()
.expect("stolen from the dead");
assert!(WatchLease::live(&dir));
l.release();
assert!(!WatchLease::path_in(&dir).exists());
std::fs::write(WatchLease::path_in(&dir), "{}").unwrap();
assert!(!WatchLease::live(&dir));
assert!(WatchLease::try_acquire(&dir).unwrap().is_some());
let _w = WriterLock::try_acquire(&dir).unwrap().unwrap();
assert!(!WatchLease::live(&dir));
}
}