diff --git a/Cargo.lock b/Cargo.lock index a7b25ea..f245ab3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2923,6 +2923,7 @@ dependencies = [ "schemars 0.9.0", "serde", "serde_json", + "signal-hook", "stable-eyre", "sysinfo", "thiserror 2.0.18", @@ -3975,6 +3976,16 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a0c28ca5908dbdbcd52e6fdaa00358ab88637f8ab33e1f188dd510eb44b53d" +dependencies = [ + "libc", + "signal-hook-registry", +] + [[package]] name = "signal-hook-registry" version = "1.4.8" diff --git a/odorobo/Cargo.toml b/odorobo/Cargo.toml index 3566ebc..2fe856d 100644 --- a/odorobo/Cargo.toml +++ b/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" diff --git a/odorobo/src/main.rs b/odorobo/src/main.rs index ffb57c5..894c0dd 100644 --- a/odorobo/src/main.rs +++ b/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, 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(); diff --git a/odorobo/src/utils/lockfile.rs b/odorobo/src/utils/lockfile.rs new file mode 100644 index 0000000..cd70447 --- /dev/null +++ b/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); + +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 { + 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> { + 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) +} diff --git a/odorobo/src/utils/mod.rs b/odorobo/src/utils/mod.rs index a874761..b3ab3d9 100644 --- a/odorobo/src/utils/mod.rs +++ b/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;