use crate::atomic::{atomic_write_with_options, AtomicWriteOptions, AtomicWriter, ReplaceMode};
use crate::error::{FsError, Result};
use crate::{process_liveness, ProcessLiveness};
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const HEADER: &str = "rightkit-lease v1";
const ATTEMPTS: u32 = 4;
static COUNTER: AtomicU64 = AtomicU64::new(0);
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct LeaseOptions {
pub ttl: Duration,
pub label: String,
}
impl LeaseOptions {
pub fn new(ttl: Duration) -> Self {
Self {
ttl,
label: String::new(),
}
}
pub fn label(mut self, label: impl Into<String>) -> Self {
self.label = label.into();
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Staleness {
Fresh,
HolderDead,
Expired,
}
impl Staleness {
pub fn is_stale(self) -> bool {
self != Staleness::Fresh
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct LeaseInfo {
pub pid: u32,
pub token: String,
pub label: String,
pub ttl: Duration,
pub acquired_ms: u64,
pub heartbeat_ms: u64,
}
impl LeaseInfo {
pub fn staleness(&self) -> Staleness {
self.staleness_at(now_ms())
}
fn staleness_at(&self, now_ms: u64) -> Staleness {
if self.pid != 0 && process_liveness(self.pid) == ProcessLiveness::Dead {
return Staleness::HolderDead;
}
let ttl_ms = millis(self.ttl);
if self.heartbeat_ms > now_ms.saturating_add(ttl_ms) {
return Staleness::Expired;
}
if now_ms.saturating_sub(self.heartbeat_ms) > ttl_ms {
return Staleness::Expired;
}
Staleness::Fresh
}
fn encode(&self) -> String {
format!(
"{HEADER}\npid={}\ntoken={}\nlabel={}\nttl_ms={}\nacquired_ms={}\nheartbeat_ms={}\n",
self.pid,
self.token,
self.label,
millis(self.ttl),
self.acquired_ms,
self.heartbeat_ms
)
}
fn parse(text: &str) -> Option<LeaseInfo> {
let mut lines = text.lines();
if lines.next()? != HEADER {
return None;
}
let (mut pid, mut token, mut label) = (None, None, String::new());
let (mut ttl, mut acquired, mut heartbeat) = (None, None, None);
for line in lines {
let (key, value) = line.split_once('=')?;
match key {
"pid" => pid = value.parse::<u32>().ok(),
"token" => token = Some(value.to_string()),
"label" => label = value.to_string(),
"ttl_ms" => ttl = value.parse::<u64>().ok(),
"acquired_ms" => acquired = value.parse::<u64>().ok(),
"heartbeat_ms" => heartbeat = value.parse::<u64>().ok(),
_ => {}
}
}
Some(LeaseInfo {
pid: pid?,
token: token.filter(|t| !t.is_empty())?,
label,
ttl: Duration::from_millis(ttl?),
acquired_ms: acquired?,
heartbeat_ms: heartbeat?,
})
}
}
#[derive(Debug)]
#[must_use]
pub enum Acquire {
Acquired(Lease),
Held(LeaseInfo),
}
impl Acquire {
pub fn acquired(self) -> Option<Lease> {
match self {
Acquire::Acquired(lease) => Some(lease),
Acquire::Held(_) => None,
}
}
}
#[derive(Debug)]
pub struct Lease {
path: PathBuf,
info: LeaseInfo,
held: bool,
}
impl Lease {
pub fn try_acquire(path: &Path, options: &LeaseOptions) -> Result<Acquire> {
Self::try_acquire_at(path, options, now_ms())
}
fn try_acquire_at(path: &Path, options: &LeaseOptions, now: u64) -> Result<Acquire> {
let info = LeaseInfo {
pid: std::process::id(),
token: new_token(),
label: options.label.replace(['\r', '\n'], " "),
ttl: options.ttl,
acquired_ms: now,
heartbeat_ms: now,
};
for _ in 0..ATTEMPTS {
let Some(guard) = Guard::acquire(path, options.ttl, now)? else {
continue;
};
match create_exclusive(path, &info.encode()) {
Ok(()) => {
return Ok(Acquire::Acquired(Lease {
path: path.to_path_buf(),
info,
held: true,
}))
}
Err(e) if e.kind() == Some(io::ErrorKind::AlreadyExists) => {}
Err(e) => return Err(e),
}
match judge(path, options.ttl, now)? {
Judged::Gone => {}
Judged::Live(holder) => return Ok(Acquire::Held(holder)),
Judged::Stale { text } => {
if take_down(path, &text, &info.encode(), &guard)? == Takedown::Created {
return Ok(Acquire::Acquired(Lease {
path: path.to_path_buf(),
info,
held: true,
}));
}
}
}
}
Err(FsError::io(
path,
io::Error::new(
io::ErrorKind::WouldBlock,
"the lease record kept changing while it was being acquired",
),
))
}
pub fn peek(path: &Path) -> Result<Option<LeaseInfo>> {
match read_record(path)? {
None => Ok(None),
Some(Record::Good(info)) => Ok(Some(info)),
Some(Record::Unreadable { mtime_ms, .. }) => Ok(Some(unreadable_info(mtime_ms))),
}
}
pub fn renew(&mut self) -> Result<()> {
self.renew_at(now_ms())
}
fn renew_at(&mut self, now: u64) -> Result<()> {
let Some(_guard) = Guard::acquire(&self.path, self.info.ttl, now)? else {
return Err(FsError::io(
&self.path,
io::Error::new(
io::ErrorKind::WouldBlock,
"the lease is being changed by another process; renew again",
),
));
};
match read_record(&self.path)? {
None => {
self.held = false;
Err(FsError::io(
&self.path,
io::Error::new(io::ErrorKind::NotFound, "the lease record is gone"),
))
}
Some(Record::Good(current)) if current.encode() == self.info.encode() => {
let mut next = self.info.clone();
next.heartbeat_ms = now.max(current.heartbeat_ms);
write_replace(&self.path, &next.encode())?;
self.info = next;
Ok(())
}
Some(other) => {
self.held = false;
let who = match other {
Record::Good(info) => {
format!("pid {} ({})", info.pid, info.label)
}
Record::Unreadable { .. } => "an unreadable record".to_string(),
};
Err(FsError::Locked {
path: self.path.clone(),
owner: who,
})
}
}
}
pub fn release(mut self) -> Result<()> {
self.release_inner()
}
fn release_inner(&mut self) -> Result<()> {
if !self.held {
return Ok(());
}
self.held = false;
let Some(_guard) = Guard::acquire(&self.path, self.info.ttl, now_ms())? else {
return Ok(());
};
match read_record(&self.path)? {
Some(Record::Good(current)) if current.encode() == self.info.encode() => {
match fs::remove_file(&self.path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(FsError::io(&self.path, e)),
}
}
_ => Ok(()),
}
}
pub fn info(&self) -> &LeaseInfo {
&self.info
}
pub fn token(&self) -> &str {
&self.info.token
}
pub fn path(&self) -> &Path {
&self.path
}
}
impl Drop for Lease {
fn drop(&mut self) {
let _ = self.release_inner();
}
}
enum Record {
Good(LeaseInfo),
Unreadable { text: String, mtime_ms: u64 },
}
enum Judged {
Gone,
Live(LeaseInfo),
Stale { text: String },
}
fn read_record(path: &Path) -> Result<Option<Record>> {
let text = match crate::replaced::read_to_string_replaced(path) {
Ok(text) => text,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None),
Err(e) if e.kind() == io::ErrorKind::InvalidData => fs::read(path)
.map(|b| String::from_utf8_lossy(&b).into_owned())
.map_err(|e| FsError::io(path, e))?,
Err(e) => return Err(FsError::io(path, e)),
};
Ok(Some(match LeaseInfo::parse(&text) {
Some(info) => Record::Good(info),
None => {
let mtime_ms = fs::metadata(path)
.ok()
.and_then(|m| m.modified().ok())
.and_then(|t| t.duration_since(UNIX_EPOCH).ok())
.map(|d| d.as_millis() as u64)
.unwrap_or(0);
Record::Unreadable { text, mtime_ms }
}
}))
}
fn unreadable_info(mtime_ms: u64) -> LeaseInfo {
LeaseInfo {
pid: 0,
token: String::new(),
label: String::new(),
ttl: Duration::ZERO,
acquired_ms: mtime_ms,
heartbeat_ms: mtime_ms,
}
}
fn judge(path: &Path, ttl: Duration, now: u64) -> Result<Judged> {
Ok(match read_record(path)? {
None => Judged::Gone,
Some(Record::Good(info)) => {
if info.staleness_at(now).is_stale() {
Judged::Stale {
text: info.encode(),
}
} else {
Judged::Live(info)
}
}
Some(Record::Unreadable { text, mtime_ms }) => {
let mut info = unreadable_info(mtime_ms);
info.ttl = ttl;
if info.staleness_at(now).is_stale() {
Judged::Stale { text }
} else {
Judged::Live(info)
}
}
})
}
fn record_text(record: &Record) -> String {
match record {
Record::Good(info) => info.encode(),
Record::Unreadable { text, .. } => text.clone(),
}
}
#[derive(Debug, PartialEq, Eq)]
enum Takedown {
Created,
Declined,
}
fn take_down(path: &Path, judged_text: &str, new_text: &str, _guard: &Guard) -> Result<Takedown> {
match read_record(path)? {
None => {}
Some(current) if record_text(¤t) == judged_text => match fs::remove_file(path) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::NotFound => {}
Err(e) => return Err(FsError::io(path, e)),
},
Some(_) => return Ok(Takedown::Declined),
}
match create_exclusive(path, new_text) {
Ok(()) => Ok(Takedown::Created),
Err(e) if e.kind() == Some(io::ErrorKind::AlreadyExists) => Ok(Takedown::Declined),
Err(e) => Err(e),
}
}
const GUARD_HEADER: &str = "rightkit-lease-guard v2";
struct Guard {
file: fs::File,
}
impl Guard {
fn path_for(lease: &Path) -> PathBuf {
let mut name = lease.as_os_str().to_os_string();
name.push(".takeover");
PathBuf::from(name)
}
fn acquire(lease: &Path, _ttl: Duration, now: u64) -> Result<Option<Guard>> {
Self::acquire_after_publish(lease, now, || {})
}
fn acquire_after_publish(
lease: &Path,
now: u64,
after_publish: impl FnOnce(),
) -> Result<Option<Guard>> {
let path = Self::path_for(lease);
if let Ok(meta) = fs::symlink_metadata(&path) {
if meta.file_type().is_symlink() {
return Err(FsError::Symlink(path));
}
}
if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
fs::create_dir_all(parent).map_err(|e| FsError::io(parent, e))?;
}
let mut options = fs::OpenOptions::new();
options.read(true).write(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options
.mode(0o600)
.custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC);
}
let mut file = match options.open(&path) {
Ok(file) => file,
Err(e) if e.kind() == io::ErrorKind::NotFound => {
if publish_guard(&path, now, |from, to| fs::hard_link(from, to))? {
after_publish();
}
options.open(&path).map_err(|e| FsError::io(&path, e))?
}
Err(e) => return Err(FsError::io(&path, e)),
};
match file.try_lock() {
Ok(()) => {}
Err(fs::TryLockError::WouldBlock) => return Ok(None),
Err(fs::TryLockError::Error(e)) => return Err(FsError::io(&path, e)),
}
let existing = {
use std::io::{Read, Seek, SeekFrom};
let mut text = String::new();
file.seek(SeekFrom::Start(0))
.map_err(|e| FsError::io(&path, e))?;
Read::take(&mut file, 4097)
.read_to_string(&mut text)
.map_err(|e| FsError::io(&path, e))?;
text
};
if !existing.starts_with(&format!("{GUARD_HEADER}\n"))
&& !legacy_guard_owner_is_dead(&existing)
{
let _ = file.unlock();
return Ok(None);
}
{
use std::io::{Seek, SeekFrom};
let text = format!(
"{GUARD_HEADER}\npid={}\ncreated_ms={now}\ntoken={}\n",
std::process::id(),
new_token()
);
let written = (|| -> io::Result<()> {
file.seek(SeekFrom::Start(0))?;
file.write_all(text.as_bytes())?;
file.set_len(text.len() as u64)
})();
if let Err(e) = written {
let _ = file.unlock();
return Err(FsError::io(&path, e));
}
}
Ok(Some(Guard { file }))
}
}
fn publish_guard(
path: &Path,
now: u64,
publish: impl FnOnce(&Path, &Path) -> io::Result<()>,
) -> Result<bool> {
let mut writer = AtomicWriter::create_with(
path,
AtomicWriteOptions {
private: true,
durable: false,
..AtomicWriteOptions::default()
},
)?;
let text = format!(
"{GUARD_HEADER}\npid={}\ncreated_ms={now}\ntoken={}\n",
std::process::id(),
new_token()
);
writer
.file_mut()
.write_all(text.as_bytes())
.map_err(|e| FsError::io(writer.temp_path(), e))?;
match publish(writer.temp_path(), path) {
Ok(()) => Ok(true),
Err(e) if e.kind() == io::ErrorKind::AlreadyExists => Ok(false),
Err(e) => Err(FsError::io(path, e)),
}
}
impl Drop for Guard {
fn drop(&mut self) {
let _ = self.file.unlock();
}
}
fn legacy_guard_owner_is_dead(text: &str) -> bool {
if text.len() > 4096 {
return false;
}
let mut pids = text.lines().filter_map(|line| line.strip_prefix("pid="));
let Some(pid) = pids.next().and_then(|value| value.parse::<u32>().ok()) else {
return false;
};
pids.next().is_none() && pid != 0 && process_liveness(pid) == ProcessLiveness::Dead
}
fn create_exclusive(path: &Path, text: &str) -> Result<()> {
atomic_write_with_options(
path,
AtomicWriteOptions {
replace: ReplaceMode::NoReplace,
create_parents: true,
durable: false,
..AtomicWriteOptions::default()
},
|f| f.write_all(text.as_bytes()),
)
}
fn write_replace(path: &Path, text: &str) -> Result<()> {
atomic_write_with_options(
path,
AtomicWriteOptions {
durable: false,
retry_sharing_violation: Some(crate::atomic::RetryPolicy::new(
5,
Duration::from_millis(10),
)),
..AtomicWriteOptions::default()
},
|f| f.write_all(text.as_bytes()),
)
}
fn new_token() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!(
"{}-{:x}-{}",
std::process::id(),
nanos,
COUNTER.fetch_add(1, Ordering::Relaxed)
)
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
fn millis(d: Duration) -> u64 {
u64::try_from(d.as_millis()).unwrap_or(u64::MAX)
}