feat: lockfile

odorobo/Cargo.toml

@@ -6,7 +6,7 @@ license = "AGPL-3.0-or-later"
 
 [dependencies]
 stable-eyre = { workspace = true }
-clap = {  workspace = true, features = ["derive", "env"] }
+clap = { workspace = true, features = ["derive", "env"] }
 tokio = { workspace = true, features = ["full"] }
 reqwest = { workspace = true, features = ["json"] }
 serde_json = { workspace = true }
@@ -40,7 +40,7 @@ tower-http = { version = "0.6", features = ["trace"] }
 kameo = { version = "0.20", features = ["remote"] }
 ahash = { version = "0.8", features = ["serde"] }
 dirs = "6"
-ipnet = {version = "2.12", features = ["serde", "schemars08"]}
+ipnet = { version = "2.12", features = ["serde", "schemars08"] }
 url = { version = "2.5", features = ["serde"] }
 async-trait = "0.1"
 dashmap = "6.1"
@@ -61,8 +61,8 @@ aide = { version = "0.15", features = [
 schemars = "0.9"
 
 
-
 cloud-hypervisor-client = { git = "https://github.com/FyraStack/cloud-hypervisor-client.git" }
+signal-hook = "0.4.4"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 rtnetlink = "0.21"

odorobo/src/main.rs

@@ -8,6 +8,8 @@ pub mod types;
 mod utils;
 
 use std::fs;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use clap::Parser;
 use kameo::actor::Spawn;
@@ -20,8 +22,7 @@ use crate::config::Config;
 use crate::utils::actor_names::{HTTP_API_SERVER, SCHEDULER};
 use crate::utils::{actor_names::AGENT, connect_to_swarm, init};
 
-#[tokio::main]
-async fn main() -> Result<()> {
+fn main() -> Result<()> {
     let cli_config = config::CliConfig::parse();
     // TODO: ask infra team where they want this on the box
     let config: Config = if let Ok(file) = fs::File::open("config.json") {
@@ -31,7 +32,32 @@ async fn main() -> Result<()> {
     };
 
     init(Some("odorobo"))?;
+    let term = utils::lockfile::register_termsigs()?;
+    let _lock = utils::lockfile::init_lockfile()
+        .map_err(|e| stable_eyre::eyre::eyre!("cannot init lockfile").wrap_err(e))?;
 
+    mainloop(term, cli_config, config)
+}
+
+fn mainloop(term: Arc<AtomicBool>, cli_config: config::CliConfig, config: Config) -> Result<()> {
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("can't build tokio");
+    let handle = runtime.spawn(inner_main(cli_config, config));
+
+    loop {
+        if handle.is_finished() {
+            break runtime.block_on(handle).expect("cannot join main thread");
+        }
+        if term.load(Ordering::Relaxed) {
+            handle.abort();
+            stable_eyre::eyre::bail!("Exit due to termination signal");
+        }
+    }
+}
+
+async fn inner_main(cli_config: config::CliConfig, config: Config) -> Result<()> {
     tracing::info!(?config, "Starting odorobo");
 
     let local_peer_id = connect_to_swarm().unwrap();

odorobo/src/utils/lockfile.rs

@@ -0,0 +1,42 @@
+use std::io::Write;
+
+const LOCKFILE: &str = "/var/lock/odorobo.lock";
+
+// `Option` needed since `std::mem::take` requires `impl Default`
+pub struct Lockfile(Option<std::fs::File>);
+
+impl Drop for Lockfile {
+    fn drop(&mut self) {
+        tracing::debug!("dropping Lockfile");
+        std::mem::drop(std::mem::take(&mut self.0).expect("None in Lockfile"));
+        _ = std::fs::remove_file(LOCKFILE)
+            .inspect_err(|e| tracing::warn!(?LOCKFILE, ?e, "cannot remove lockfile"));
+    }
+}
+
+/// Create an odorobo lockfile
+///
+/// See #30, only 1 instance of odorobo should run.
+pub fn init_lockfile() -> Result<Lockfile, String> {
+    tracing::trace!("creating lockfile at {LOCKFILE}");
+    let mut f = std::fs::File::create_new(LOCKFILE)
+        .map_err(|e| format!("Cannot create lockfile at {LOCKFILE}: {e:?}"))?;
+    f.write_all(&std::process::id().to_ne_bytes())
+        .map_err(|e| format!("cannot write to {LOCKFILE}: {e:?}"))?;
+    f.flush()
+        .map_err(|e| format!("cannot flush {LOCKFILE}: {e:?}"))?;
+    Ok(Lockfile(Some(f)))
+}
+
+/// Register termination signals to watch
+///
+/// We need to tidy up lockfiles before exiting. Watching these signals can give us a chance to
+/// actually clean things up.
+pub fn register_termsigs() -> std::io::Result<std::sync::Arc<std::sync::atomic::AtomicBool>> {
+    let term = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+    for sig in signal_hook::consts::TERM_SIGNALS {
+        signal_hook::flag::register_conditional_shutdown(*sig, 1, std::sync::Arc::clone(&term))?;
+        signal_hook::flag::register(*sig, std::sync::Arc::clone(&term))?;
+    }
+    Ok(term)
+}

odorobo/src/utils/mod.rs

@@ -1,5 +1,6 @@
 pub mod actor_cache;
 pub mod actor_names;
+pub mod lockfile;
 
 use aide::OperationIo;
 use api_error::ApiError;