From ac2b66a3395af6a7b19d79120fbce450e3d99f0d Mon Sep 17 00:00:00 2001 From: madonuko Date: Sun, 5 Jul 2026 17:05:30 +0800 Subject: [PATCH 1/2] feat: create lockfile on startup Close #30 --- odorobo/Cargo.toml | 5 ++--- odorobo/src/main.rs | 12 +++++------- odorobo/src/utils/lockfile.rs | 27 +++++++++++++++++++++++++++ odorobo/src/utils/mod.rs | 34 +++++++++++++++------------------- 4 files changed, 49 insertions(+), 29 deletions(-) create mode 100644 odorobo/src/utils/lockfile.rs diff --git a/odorobo/Cargo.toml b/odorobo/Cargo.toml index 91cee6c..cecda2f 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 } @@ -42,7 +42,7 @@ ahash = { version = "0.8", features = ["serde"] } ulid = { version = "1.2", features = ["serde", "uuid"] } dirs = "6" bytesize = { version = "2.3", features = ["serde"] } -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" @@ -63,5 +63,4 @@ aide = { version = "0.15", features = [ schemars = "0.9" - cloud-hypervisor-client = { git = "https://github.com/FyraStack/cloud-hypervisor-client.git" } diff --git a/odorobo/src/main.rs b/odorobo/src/main.rs index 350843f..4a4c359 100644 --- a/odorobo/src/main.rs +++ b/odorobo/src/main.rs @@ -1,12 +1,12 @@ pub mod actors; mod ch_api; +pub mod config; pub mod http_api; +pub mod messages; pub mod networking; mod state; -mod utils; -pub mod messages; -pub mod config; pub mod types; +mod utils; use std::fs; @@ -14,15 +14,16 @@ use clap::Parser; use kameo::actor::Spawn; use stable_eyre::Result; +use crate::actors::agent_actor::AgentActor; use crate::actors::http_actor::HTTPActor; use crate::actors::scheduler_actor::SchedulerActor; use crate::config::Config; use crate::utils::actor_names::{HTTP_API_SERVER, SCHEDULER}; use crate::utils::{actor_names::AGENT, connect_to_swarm, init}; -use crate::actors::agent_actor::AgentActor; #[tokio::main] async fn main() -> Result<()> { + let _ = utils::lockfile::init_lockfile(); 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") { @@ -35,11 +36,9 @@ async fn main() -> Result<()> { tracing::info!(?config, "Starting odorobo"); - let local_peer_id = connect_to_swarm().unwrap(); tracing::info!(?local_peer_id, "Peer ID"); - // start agents let agent_actor = AgentActor::spawn(config.clone()); agent_actor.register(AGENT).await?; @@ -55,7 +54,6 @@ async fn main() -> Result<()> { http_actor.wait_for_shutdown().await; } - agent_actor.wait_for_shutdown().await; Ok(()) diff --git a/odorobo/src/utils/lockfile.rs b/odorobo/src/utils/lockfile.rs new file mode 100644 index 0000000..1890f75 --- /dev/null +++ b/odorobo/src/utils/lockfile.rs @@ -0,0 +1,27 @@ +use tokio::io::AsyncWriteExt; + +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) { + 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 async fn init_lockfile() -> Result { + let mut f = tokio::fs::File::create_new(LOCKFILE) + .await + .map_err(|e| format!("Cannot create lockfile at {LOCKFILE}: {e:?}"))?; + f.write_all(&std::process::id().to_ne_bytes()) + .await + .map_err(|e| format!("cannot write to {LOCKFILE}: {e:?}"))?; + Ok(Lockfile(Some(f))) +} diff --git a/odorobo/src/utils/mod.rs b/odorobo/src/utils/mod.rs index 01c7125..f85de5f 100644 --- a/odorobo/src/utils/mod.rs +++ b/odorobo/src/utils/mod.rs @@ -1,17 +1,18 @@ -pub mod actor_names; pub mod actor_cache; +pub mod actor_names; +pub mod lockfile; use aide::OperationIo; -use stable_eyre::{Result, Report}; -use tracing::level_filters::LevelFilter; -use tracing_subscriber::EnvFilter; -use thiserror::Error; +use api_error::ApiError; use kameo::prelude::*; use libp2p::futures::StreamExt; use libp2p::swarm::{NetworkBehaviour, SwarmEvent}; use libp2p::{PeerId, mdns, noise, tcp, yamux}; +use stable_eyre::{Report, Result}; +use thiserror::Error; +use tracing::level_filters::LevelFilter; use tracing::{debug, error, info, trace, warn}; -use api_error::ApiError; +use tracing_subscriber::EnvFilter; // todo: wrap with axum-responses, return this type on request failure #[derive(Error, Debug, ApiError, OperationIo)] @@ -26,21 +27,18 @@ impl From> for OdoroboError { fn from(value: kameo::error::SendError) -> Self { let kameo_error = value.to_string(); error!(?value); - OdoroboError::Report(value.err().unwrap_or_else(|| { - Report::msg(format!("could not unwrap kameo error: {kameo_error}")) - })) + OdoroboError::Report( + value.err().unwrap_or_else(|| { + Report::msg(format!("could not unwrap kameo error: {kameo_error}")) + }), + ) } } #[cfg(test)] mod tests { use super::*; - use axum::{ - http::StatusCode, - body::Body, http::Request, - routing::get, - Router, - }; + use axum::{Router, body::Body, http::Request, http::StatusCode, routing::get}; use http_body_util::BodyExt; use tower::ServiceExt; @@ -50,7 +48,8 @@ mod tests { #[tokio::test] async fn test_error() { - let response = Router::new().route("/", get(handler)) + let response = Router::new() + .route("/", get(handler)) .oneshot(Request::get("/").body(Body::empty()).unwrap()) .await .unwrap(); @@ -64,7 +63,6 @@ mod tests { } } - pub fn env_filter(debug_target: Option<&str>) -> EnvFilter { let env = std::env::var("ODOROBO_LOG").unwrap_or_else(|_| "".into()); @@ -110,8 +108,6 @@ pub fn init_default() -> Result<()> { init(None) } - - #[derive(NetworkBehaviour)] pub struct ProductionBehaviour { kameo: remote::Behaviour, From e7f08f590db00b3d4646ba7fba686a37097a4c32 Mon Sep 17 00:00:00 2001 From: madonuko Date: Sun, 5 Jul 2026 18:11:17 +0800 Subject: [PATCH 2/2] feat: watch for termsigs --- Cargo.lock | 11 +++++++++++ odorobo/Cargo.toml | 1 + odorobo/src/main.rs | 31 ++++++++++++++++++++++++++++--- odorobo/src/utils/lockfile.rs | 27 +++++++++++++++++++++------ 4 files changed, 61 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2a573c2..8265cec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2922,6 +2922,7 @@ dependencies = [ "schemars 0.9.0", "serde", "serde_json", + "signal-hook", "stable-eyre", "sysinfo", "thiserror 2.0.18", @@ -3958,6 +3959,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 cecda2f..9890606 100644 --- a/odorobo/Cargo.toml +++ b/odorobo/Cargo.toml @@ -64,3 +64,4 @@ schemars = "0.9" cloud-hypervisor-client = { git = "https://github.com/FyraStack/cloud-hypervisor-client.git" } +signal-hook = "0.4.4" diff --git a/odorobo/src/main.rs b/odorobo/src/main.rs index 4a4c359..85f209f 100644 --- a/odorobo/src/main.rs +++ b/odorobo/src/main.rs @@ -9,6 +9,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; @@ -21,9 +23,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<()> { - let _ = utils::lockfile::init_lockfile(); +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") { @@ -33,7 +33,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 index 1890f75..cd70447 100644 --- a/odorobo/src/utils/lockfile.rs +++ b/odorobo/src/utils/lockfile.rs @@ -1,12 +1,13 @@ -use tokio::io::AsyncWriteExt; +use std::io::Write; const LOCKFILE: &str = "/var/lock/odorobo.lock"; // `Option` needed since `std::mem::take` requires `impl Default` -pub struct Lockfile(Option); +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")); @@ -16,12 +17,26 @@ impl Drop for Lockfile { /// Create an odorobo lockfile /// /// See #30, only 1 instance of odorobo should run. -pub async fn init_lockfile() -> Result { - let mut f = tokio::fs::File::create_new(LOCKFILE) - .await +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()) - .await .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) +}