watt: implement single-instance lock mechanism for watt daemon

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: Ideaf221f690c7bc694ca1f171385bfde6a6a6964

Author: NotAShelf

Cargo.lock

@@ -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"

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

watt/Cargo.toml

@@ -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<PathBuf>,
+
+  /// 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)
 }

watt/lib.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<File>,
+  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<Option<Self>, 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<u32> {
+    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();
+  }
+}