diff --git a/Cargo.lock b/Cargo.lock index 4f40f9b..b73a586 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -403,8 +403,10 @@ dependencies = [ "semver", "serde", "serde_json", + "signal-hook", "structopt", "tempfile", + "tokio", "toml", "toml_edit", "urlencoding", @@ -1197,6 +1199,25 @@ dependencies = [ "serde", ] +[[package]] +name = "signal-hook" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + [[package]] name = "similar" version = "2.2.1" @@ -1385,10 +1406,23 @@ dependencies = [ "mio", "num_cpus", "pin-project-lite", + "signal-hook-registry", "socket2 0.5.3", + "tokio-macros", "windows-sys 0.48.0", ] +[[package]] +name = "tokio-macros" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.29", +] + [[package]] name = "tokio-native-tls" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index b015216..9dfcc05 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,13 @@ toml_edit = "0.14.4" urlencoding = "2.1.0" zip = "0.5" +[target.'cfg(windows)'.dependencies] +command-group = "1.0.8" + +[target.'cfg(unix)'.dependencies] +tokio = { version = "1.18.2", features = ["macros", "sync", "process"] } +signal-hook = "0.3.14" + [dev_dependencies] assert_cmd = "2.0.2" insta = "1.14" diff --git a/src/main.rs b/src/main.rs index bf7c65e..23d9a58 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,10 +6,11 @@ mod config; mod error; mod fs; mod paths; +mod process; mod tool_cache; mod tool_provider; -use std::{env, ffi::OsStr, process}; +use std::{env, ffi::OsStr}; use paths::ForemanPaths; use structopt::StructOpt; @@ -67,7 +68,7 @@ impl ToolInvocation { let exit_code = tool_cache.run(tool_spec, &version, self.args)?; if exit_code != 0 { - process::exit(exit_code); + std::process::exit(exit_code); } Ok(()) @@ -116,7 +117,7 @@ fn main() { fn exit_with_error(error: ForemanError) -> ! { eprintln!("{}", error); - process::exit(1); + std::process::exit(1); } #[derive(Debug, StructOpt)] diff --git a/src/process/mod.rs b/src/process/mod.rs new file mode 100644 index 0000000..00713a4 --- /dev/null +++ b/src/process/mod.rs @@ -0,0 +1,11 @@ +#[cfg(windows)] +mod windows; + +#[cfg(windows)] +pub use windows::run; + +#[cfg(unix)] +mod unix; + +#[cfg(unix)] +pub use unix::run; diff --git a/src/process/unix.rs b/src/process/unix.rs new file mode 100644 index 0000000..ccb7380 --- /dev/null +++ b/src/process/unix.rs @@ -0,0 +1,69 @@ +//! On Unix, we use tokio to spawn processes so that we can listen for signals +//! and wait for process completion at the same time. + +use std::io::{Error, ErrorKind}; +use std::path::Path; +use std::thread; + +use signal_hook::consts::signal::{SIGABRT, SIGINT, SIGQUIT, SIGTERM}; +use signal_hook::iterator::Signals; +use tokio::process::Command; +use tokio::sync::oneshot; + +pub fn run(exe_path: &Path, args: Vec) -> Result { + let (kill_tx, kill_rx) = oneshot::channel(); + + // Spawn a thread dedicated to listening for signals and relaying them to + // our async runtime. + let (signal_thread, signal_handle) = { + let mut signals = Signals::new(&[SIGABRT, SIGINT, SIGQUIT, SIGTERM]).unwrap(); + let signal_handle = signals.handle(); + + let thread = thread::spawn(move || { + if let Some(signal) = signals.into_iter().next() { + kill_tx.send(signal).ok(); + } + }); + + (thread, signal_handle) + }; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .map_err(|_| Error::new(ErrorKind::Other, "could not create tokio runtime"))?; + + let _guard = runtime.enter(); + + let mut child = Command::new(exe_path).args(args).spawn().map_err(|_| { + Error::new( + ErrorKind::Other, + format!("could not spawn {}", exe_path.display()), + ) + })?; + + let code = runtime.block_on(async move { + tokio::select! { + // If the child exits cleanly, we can return its exit code directly. + // I wish everything were this tidy. + status = child.wait() => { + let code = status.ok().and_then(|s| s.code()).unwrap_or(1); + signal_handle.close(); + signal_thread.join().unwrap(); + + code + } + + // If we received a signal while the process was running, murder it + // and exit immediately with the correct error code. + code = kill_rx => { + child.kill().await.ok(); + signal_handle.close(); + signal_thread.join().unwrap(); + std::process::exit(128 + code.unwrap_or(0)); + } + } + }); + + Ok(code) +} diff --git a/src/process/windows.rs b/src/process/windows.rs new file mode 100644 index 0000000..8af9561 --- /dev/null +++ b/src/process/windows.rs @@ -0,0 +1,24 @@ +//! On Windows, we use command_group to spawn processes in a job group that will +//! be automatically cleaned up when this process exits. + +use std::io::{Error, ErrorKind}; +use std::path::Path; +use std::process::Command; + +use command_group::CommandGroup; + +pub fn run(exe_path: &Path, args: Vec) -> Result { + // On Windows, using a job group here will cause the subprocess to terminate + // automatically when Aftman is terminated. + let mut child = Command::new(exe_path) + .args(args) + .group_spawn() + .map_err(|_| { + Error::new( + ErrorKind::Other, + format!("Could not spawn {}", exe_path.display()), + ) + })?; + let status = child.wait()?; + Ok(status.code().unwrap_or(1)) +} diff --git a/src/tool_cache.rs b/src/tool_cache.rs index 3379fce..ceb1726 100644 --- a/src/tool_cache.rs +++ b/src/tool_cache.rs @@ -3,10 +3,8 @@ use std::{ env::consts::EXE_SUFFIX, io::Cursor, path::PathBuf, - process, }; -use command_group::CommandGroup; use semver::Version; use serde::{Deserialize, Serialize}; use zip::ZipArchive; @@ -18,6 +16,7 @@ use crate::{ error::{ForemanError, ForemanResult}, fs, paths::ForemanPaths, + process, tool_provider::{Release, ToolProvider}, }; @@ -61,20 +60,17 @@ impl ToolCache { log::debug!("Running tool {} ({})", tool, tool_path.display()); - let status = process::Command::new(&tool_path) - .args(args) - .group_status() - .map_err(|err| { - ForemanError::io_error_with_context(err, - format!( - "an error happened trying to run `{}` at `{}` (this is an error in Foreman)", - tool, - tool_path.display() - ) - ) - })?; - - Ok(status.code().unwrap_or(1)) + let code = process::run(&tool_path, args).map_err(|err| { + ForemanError::io_error_with_context( + err, + format!( + "an error happened trying to run `{}` at `{}` (this is an error in Foreman)", + tool, + tool_path.display() + ), + ) + })?; + Ok(code) } pub fn download_if_necessary(