diff --git a/src/tools.rs b/src/tools.rs index 8eaebaa0..eb3c410b 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -30,6 +30,7 @@ pub mod borrow; pub mod fs; pub mod tty; pub mod signalfd; +pub mod daemon; #[macro_use] mod file_logger; diff --git a/src/tools/daemon.rs b/src/tools/daemon.rs new file mode 100644 index 00000000..bdd8450b --- /dev/null +++ b/src/tools/daemon.rs @@ -0,0 +1,115 @@ +//! Helpers for daemons/services. + +use std::ffi::CString; +use std::os::unix::ffi::OsStrExt; +use std::panic::UnwindSafe; + +use failure::*; + +// Unfortunately FnBox is nightly-only and Box is unusable, so just use Box... +pub type BoxedStoreFunc = Box Result + UnwindSafe + Send>; + +/// Helper trait to "store" something in the environment to be re-used after re-executing the +/// service on a reload. +pub trait ReexecContinue: Sized { + fn restore(var: &str) -> Result; + fn get_store_func(&self) -> BoxedStoreFunc; +} + +/// Manages things to be stored and reloaded upon reexec. +/// Anything which should be restorable should be instantiated via this struct's `restore` method, +pub struct ReexecStore { + pre_exec: Vec, +} + +// Currently we only need environment variables for storage, but in theory we could also add +// variants which need temporary files or pipes... +struct PreExecEntry { + name: &'static str, // Feel free to change to String if necessary... + store_fn: BoxedStoreFunc, +} + +impl ReexecStore { + pub fn new() -> Self { + Self { + pre_exec: Vec::new(), + } + } + + /// Restore an object from an environment variable of the given name, or, if none exists, uses + /// the function provided in the `or_create` parameter to instantiate the new "first" instance. + /// + /// Values created via this method will be remembered for later re-execution. + pub fn restore(&mut self, name: &'static str, or_create: F) -> Result + where + T: ReexecContinue, + F: FnOnce() -> Result, + { + let res = match std::env::var(name) { + Ok(varstr) => T::restore(&varstr)?, + Err(std::env::VarError::NotPresent) => or_create()?, + Err(_) => bail!("variable {} has invalid value", name), + }; + + self.pre_exec.push(PreExecEntry { + name, + store_fn: res.get_store_func(), + }); + Ok(res) + } + + fn pre_exec(self) -> Result<(), Error> { + for item in self.pre_exec { + std::env::set_var(item.name, (item.store_fn)()?); + } + Ok(()) + } + + pub fn fork_restart(self) -> Result<(), Error> { + // Get the path to our executable as CString + let exe = CString::new( + std::fs::read_link("/proc/self/exe")? + .into_os_string() + .as_bytes() + )?; + + // Get our parameters as Vec + let args = std::env::args_os(); + let mut new_args = Vec::with_capacity(args.len()); + for arg in args { + new_args.push(CString::new(arg.as_bytes())?); + } + + // Start ourselves in the background: + use nix::unistd::{fork, ForkResult}; + match fork() { + Ok(ForkResult::Child) => { + // At this point we call pre-exec helpers. We must be certain that if they fail for + // whatever reason we can still call `_exit()`, so use catch_unwind. + match std::panic::catch_unwind(move || self.do_exec(exe, new_args)) { + Ok(_) => eprintln!("do_exec returned unexpectedly!"), + Err(_) => eprintln!("panic in re-exec"), + } + // No matter how we managed to get here, this is the time where we bail out quickly: + unsafe { + libc::_exit(-1) + } + } + Ok(ForkResult::Parent { child }) => { + eprintln!("forked off a new server (pid: {})", child); + Ok(()) + } + Err(e) => { + eprintln!("fork() failed, restart delayed: {}", e); + Ok(()) + } + } + } + + fn do_exec(self, exe: CString, args: Vec) -> Result<(), Error> { + self.pre_exec()?; + nix::unistd::setsid()?; + nix::unistd::execvp(&exe, &args)?; + Ok(()) + } +}