Additional Signal handling

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"

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"] }

src/main.rs

@@ -52,23 +52,26 @@ async fn main() -> anyhow::Result<ExitCode> {
     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::<Vec<String>>())
         .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<ExitCode> {
                     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<ExitCode> {
                     },
                 }
             }
-            _ = 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<ExitCode> {
                 _ = 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));
             }
         }
     }