diff --git a/Cargo.lock b/Cargo.lock index e065c4f..1cf763f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -206,7 +206,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73736a89c4aff73035ba2ed2e565061954da00d4970fc9ac25dcc85a2a20d790" dependencies = [ "dispatch2", - "nix", + "nix 0.30.1", "windows-sys", ] @@ -413,6 +413,18 @@ dependencies = [ "libc", ] +[[package]] +name = "nix" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225e7cfe711e0ba79a68baeddb2982723e4235247aefce1482f2f16c27865b66" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -826,6 +838,7 @@ dependencies = [ "derive_more", "env_logger", "log", + "nix 0.31.1", "num_cpus", "proptest", "serde", diff --git a/Cargo.toml b/Cargo.toml index e0f96cc..11ac21a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,10 +18,11 @@ clap = { features = [ "derive", "env" ], version = "4.5.50" } clap-verbosity-flag = "3.0.4" clap_complete = "4.5.59" clap_complete_nushell = "4.5.9" -ctrlc = "3.5.0" +ctrlc = "3.5.1" derive_more = { features = [ "full" ], version = "2.0.1" } env_logger = "0.11.8" log = "0.4.28" +nix = { features = [ "fs" ], version = "0.31.1" } num_cpus = "1.17.0" serde = { features = [ "derive" ], version = "1.0.228" } thiserror = "2.0.17" diff --git a/watt/Cargo.toml b/watt/Cargo.toml index 7f3ed74..0532288 100644 --- a/watt/Cargo.toml +++ b/watt/Cargo.toml @@ -21,6 +21,7 @@ ctrlc.workspace = true derive_more.workspace = true env_logger.workspace = true log.workspace = true +nix.workspace = true num_cpus.workspace = true serde.workspace = true thiserror.workspace = true diff --git a/watt/lib.rs b/watt/lib.rs index 491da5c..2ccf227 100644 --- a/watt/lib.rs +++ b/watt/lib.rs @@ -1,4 +1,7 @@ -use std::path::PathBuf; +use std::{ + env, + path::PathBuf, +}; use anyhow::Context as _; use clap::Parser as _; @@ -11,6 +14,8 @@ pub mod fs; pub mod config; +pub mod lock; + #[derive(clap::Parser, Debug)] #[command(version, about)] pub struct Cli { @@ -20,6 +25,11 @@ pub struct Cli { /// The daemon config path. #[arg(long, env = "WATT_CONFIG")] config: Option, + + /// Force running even if another instance is already running. Potentially + /// destructive. + #[arg(long)] + force: bool, } pub fn main() -> anyhow::Result<()> { @@ -38,5 +48,13 @@ pub fn main() -> anyhow::Result<()> { log::info!("starting watt daemon"); + let lock_path = env::var("XDG_RUNTIME_DIR") + .map(|dir| PathBuf::from(dir).join("watt.pid")) + .unwrap_or_else(|_| PathBuf::from("/run/watt.pid")); + + let _lock = lock::LockFile::acquire(&lock_path, cli.force).context( + format!("failed to acquire pid lock at {}", lock_path.display()), + )?; + system::run_daemon(config) } diff --git a/watt/lock.rs b/watt/lock.rs new file mode 100644 index 0000000..d060d1a --- /dev/null +++ b/watt/lock.rs @@ -0,0 +1,179 @@ +use std::{ + error::Error, + fmt, + fs::{ + self, + File, + OpenOptions, + }, + io::Write, + ops, + path::{ + Path, + PathBuf, + }, + process, +}; + +#[cfg(unix)] use nix::fcntl::{ + Flock, + FlockArg, +}; + +#[cfg(not(unix))] +compile_error!("watt is only supported on Unix-like systems"); + +pub struct LockFile { + lock: Flock, + path: PathBuf, +} + +#[derive(Debug)] +pub struct LockFileError { + pub path: PathBuf, + pid: u32, +} + +impl fmt::Display for LockFileError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.pid == 0 { + write!(f, "failed to acquire lock at {}", self.path.display()) + } else { + write!(f, "another watt daemon is running (PID: {})", self.pid) + } + } +} + +impl Error for LockFileError {} + +impl ops::Deref for LockFile { + type Target = File; + + fn deref(&self) -> &Self::Target { + &self.lock + } +} + +impl ops::DerefMut for LockFile { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.lock + } +} + +impl LockFile { + pub fn path(&self) -> &Path { + &self.path + } + + pub fn acquire( + lock_path: &Path, + force: bool, + ) -> Result, LockFileError> { + let pid = process::id(); + + #[allow(clippy::suspicious_open_options)] + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(lock_path) + .map_err(|error| { + log::error!( + "failed to open lock file at {}: {}", + lock_path.display(), + error + ); + LockFileError { + path: lock_path.to_owned(), + pid: 0, + } + })?; + + let mut lock = match Flock::lock(file, FlockArg::LockExclusiveNonblock) { + Ok(lock) => lock, + Err((_, nix::errno::Errno::EWOULDBLOCK)) => { + let Some(existing_pid) = Self::read_pid(lock_path) else { + if force { + log::warn!( + "could not determine PID of existing watt instance, starting \ + anyway", + ); + return Ok(None); + } + + return Err(LockFileError { + path: lock_path.to_owned(), + pid: 0, + }); + }; + + if force { + log::warn!( + "another watt instance is running (PID: {existing_pid}), starting \ + anyway", + ); + return Ok(None); + } + + return Err(LockFileError { + path: lock_path.to_owned(), + pid: existing_pid, + }); + }, + + Err((_, error)) => { + log::error!("failed to acquire lock: {}", error); + return Err(LockFileError { + path: lock_path.to_owned(), + pid: 0, + }); + }, + }; + + if let Err(e) = lock.set_len(0) { + log::error!("failed to truncate lock file: {}", e); + return Err(LockFileError { + path: lock_path.to_owned(), + pid: 0, + }); + } + + if let Err(e) = lock.write_all(format!("{pid}\n").as_bytes()) { + log::error!("failed to write PID to lock file: {}", e); + return Err(LockFileError { + path: lock_path.to_owned(), + pid: 0, + }); + } + + if let Err(e) = lock.sync_all() { + log::error!("failed to sync lock file: {}", e); + return Err(LockFileError { + path: lock_path.to_owned(), + pid: 0, + }); + } + + Ok(Some(LockFile { + lock, + path: lock_path.to_owned(), + })) + } + + fn read_pid(lock_path: &Path) -> Option { + match fs::read_to_string(lock_path) { + Ok(content) => content.trim().parse().ok(), + Err(_) => None, + } + } + + pub fn release(&mut self) { + let _ = fs::remove_file(&self.path); + } +} + +impl Drop for LockFile { + fn drop(&mut self) { + self.release(); + } +}