diff --git a/Cargo.lock b/Cargo.lock index 4b479e9..0db4940 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,12 @@ dependencies = [ "rustc-demangle", ] +[[package]] +name = "bitflags" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf4b9d6a944f767f8e5e0db018570623c85f3d925ac718db4e06d0187adb21c1" + [[package]] name = "bytes" version = "1.6.0" @@ -105,6 +111,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "clap" version = "4.5.4" @@ -187,6 +199,7 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "nix", "tokio", ] @@ -216,6 +229,18 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "num_cpus" version = "1.16.0" diff --git a/Cargo.toml b/Cargo.toml index 178a001..de6b87a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,4 +8,5 @@ edition = "2021" [dependencies] anyhow = "1.0.86" clap = { version = "4.5.4", features = ["derive"] } +nix = { version = "0.29.0", features = ["signal", "process"] } tokio = { version = "1.38.0", features = ["process", "signal", "rt", "rt-multi-thread", "macros", "io-util", "fs"] } diff --git a/src/main.rs b/src/main.rs index e8f3eb7..25e4c78 100644 --- a/src/main.rs +++ b/src/main.rs @@ -52,23 +52,26 @@ async fn main() -> anyhow::Result { let stdout_path = &args.stdout_path; // Setup our signal hook. let handled_sig: SignalKind = (&args.rotated_signal).into(); - let mut signal_stream = signal(handled_sig)?; + let mut rotation_signal_stream = signal(handled_sig)?; + let mut sigterm_stream = signal(SignalKind::terminate())?; + let mut sigkill_stream = signal(SignalKind::from_raw(9))?; + let mut sigquit_stream = signal(SignalKind::quit())?; // Setup our output wiring. let app_name = match args.cmd.first() { Some(n) => n, None => return Err(anyhow::anyhow!("No command specified")), }; - let child = Command::new(app_name) + let mut child = Command::new(app_name) .args(args.cmd.into_iter().skip(1).collect::>()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn()?; let mut stdout_reader = child - .stdout + .stdout.take() .expect("no valid stdout from command available"); let mut stdout_buffer = [0; 8 * 1024]; let mut stderr_reader = child - .stderr + .stderr.take() .expect("no valid stderr from command available"); let mut stderr_buffer = [0; 8 * 1024]; @@ -92,7 +95,6 @@ async fn main() -> anyhow::Result { Err(e) => { // TODO(zaphar): This likely means the command has broken badly. We should // do the right thing here.. - todo!() }, } } @@ -113,7 +115,7 @@ async fn main() -> anyhow::Result { }, } } - _ = signal_stream.recv() => { + _ = rotation_signal_stream.recv() => { // on sighub sync and reopen our files // NOTE(zaphar): This will cause the previously opened handles to get // dropped which will cause them to close assuming all the io has finished. This is why we sync @@ -123,7 +125,43 @@ async fn main() -> anyhow::Result { _ = stdout_writer.sync_all().await; stderr_writer = File::options().append(true).open(stderr_path).await?; stdout_writer = File::options().append(true).open(stdout_path).await?; - // wait for a signal + } + _ = sigterm_stream.recv() => { + // NOTE(zaphar): This is a giant hack. + // If https://github.com/tokio-rs/tokio/issues/3379 ever get's implemented it will become + // unnecessary. + use nix::{ + sys::signal::{kill, Signal::SIGTERM}, + unistd::Pid, + }; + if let Some(pid) = child.id() { + // If the child hasn't already completed, send a SIGTERM. + if let Err(e) = kill(Pid::from_raw(pid.try_into().expect("Invalid PID")), SIGTERM) { + eprintln!("Failed to forward SIGTERM to child process: {}", e); + } + } + } + _ = sigquit_stream.recv() => { + // NOTE(zaphar): This is a giant hack. + // If https://github.com/tokio-rs/tokio/issues/3379 ever get's implemented it will become + // unnecessary. + use nix::{ + sys::signal::{kill, Signal::SIGQUIT}, + unistd::Pid, + }; + if let Some(pid) = child.id() { + // If the child hasn't already completed, send a SIGTERM. + if let Err(e) = kill(Pid::from_raw(pid.try_into().expect("Invalid PID")), SIGQUIT) { + eprintln!("Failed to forward SIGTERM to child process: {}", e); + } + } + } + _ = sigkill_stream.recv() => { + child.start_kill()?; + } + result = child.wait() => { + // The child has finished + return Ok(ExitCode::from(result?.code().expect("No exit code for process") as u8)); } } }