pass signals to tasks (#1527)

* pass signals to tasks * [MegaLinter] Apply linters fixes --------- Co-authored-by: jdx <[email protected]>
jdx · Jan 25, 2024 · 59c7216 · 59c7216
1 parent bdb9014
commit 59c7216
Show file tree

Hide file tree

Showing 11 changed files with 123 additions and 60 deletions.
diff --git a/.mise.toml b/.mise.toml
@@ -91,3 +91,6 @@ run = "cargo release"
 
 [tasks.lint-fix]
 run = "just lint-fix"
+
+[tasks.signal-test]
+run = "node ./test/fixtures/signal-test.js"
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -48,7 +48,6 @@ color-print = "0.3.5"
 confique = { version = "0.2.5", default-features = false }
 console = "0.15.8"
 contracts = "0.6.3"
-ctrlc = "3.4.2"
 demand = "1.0.1"
 dotenvy = "0.15.7"
 duct = "0.13.7"
@@ -87,6 +86,7 @@ serde_json = { version = "1.0.111", features = [] }
 sha2 = "0.10.8"
 shell-escape = "0.1.5"
 shell-words = "1.1.0"
+signal-hook = "0.3.17"
 simplelog = { version = "0.12.1" }
 strum = { version = "0.25.0", features = ["derive"] }
 sys-info = "0.9.1"

diff --git a/src/cli/run.rs b/src/cli/run.rs
@@ -24,6 +24,7 @@ use crate::errors::Error::ScriptFailed;
 use crate::file::display_path;
 use crate::task::{Deps, Task};
 use crate::toolset::{InstallOptions, ToolsetBuilder};
+use crate::ui::ctrlc;
 use crate::ui::style;
 use crate::{env, file, ui};
 
@@ -63,7 +64,7 @@ pub struct Run {
     pub task: String,
 
     /// Arguments to pass to the task. Use ":::" to separate tasks.
-    #[clap()]
+    #[clap(allow_hyphen_values = true)]
     pub args: Vec<String>,
 
     /// Change to this directory before executing the command
@@ -269,6 +270,7 @@ impl Run {
     ) -> Result<()> {
         let program = program.to_executable();
         let mut cmd = CmdLineRunner::new(program.clone()).args(args).envs(env);
+        cmd.with_pass_signals();
         match &self.output(task)? {
             TaskOutput::Prefix => cmd = cmd.prefix(format!("{prefix} ")),
             TaskOutput::Interleave => {
@@ -345,7 +347,7 @@ impl Run {
         for name in task_names {
             s = s.option(DemandOption::new(name));
         }
-        ui::handle_ctrlc();
+        let _ = ctrlc::handle_ctrlc()?;
         let name = s.run()?;
         match tasks.get(name) {
             Some(task) => Ok(task.clone()),

diff --git a/src/cli/upgrade.rs b/src/cli/upgrade.rs
@@ -126,7 +126,7 @@ impl Upgrade {
     }
 
     fn get_interactive_tool_set(&self, outdated: &OutputVec) -> Result<HashSet<ToolVersion>> {
-        ui::handle_ctrlc();
+        let _ = ui::ctrlc::handle_ctrlc()?;
         let mut ms = demand::MultiSelect::new("mise upgrade")
             .description("Select tools to upgrade")
             .filterable(true)

diff --git a/src/cmd.rs b/src/cmd.rs
@@ -10,6 +10,8 @@ use std::thread;
 use color_eyre::Result;
 use duct::{Expression, IntoExecutablePath};
 use eyre::Context;
+use signal_hook::consts::{SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGUSR1, SIGUSR2};
+use signal_hook::iterator::Signals;
 
 use crate::config::Settings;
 use crate::env;
@@ -95,6 +97,7 @@ pub struct CmdLineRunner<'a> {
     stdin: Option<String>,
     prefix: String,
     raw: bool,
+    pass_signals: bool,
 }
 
 static OUTPUT_LOCK: Mutex<()> = Mutex::new(());
@@ -112,6 +115,7 @@ impl<'a> CmdLineRunner<'a> {
             stdin: None,
             prefix: String::new(),
             raw: false,
+            pass_signals: false,
         }
     }
 
@@ -208,6 +212,11 @@ impl<'a> CmdLineRunner<'a> {
         self
     }
 
+    pub fn with_pass_signals(&mut self) -> &mut Self {
+        self.pass_signals = true;
+        self
+    }
+
     pub fn stdin_string(mut self, input: impl Into<String>) -> Self {
         self.cmd.stdin(Stdio::piped());
         self.stdin = Some(input.into());
@@ -237,7 +246,6 @@ impl<'a> CmdLineRunner<'a> {
                         let line = line.unwrap();
                         tx.send(ChildProcessOutput::Stdout(line)).unwrap();
                     }
-                    tx.send(ChildProcessOutput::Done).unwrap();
                 }
             });
         }
@@ -249,7 +257,6 @@ impl<'a> CmdLineRunner<'a> {
                         let line = line.unwrap();
                         tx.send(ChildProcessOutput::Stderr(line)).unwrap();
                     }
-                    tx.send(ChildProcessOutput::Done).unwrap();
                 }
             });
         }
@@ -259,13 +266,27 @@ impl<'a> CmdLineRunner<'a> {
                 stdin.write_all(text.as_bytes()).unwrap();
             });
         }
+        let mut sighandle = None;
+        if self.pass_signals {
+            let mut signals =
+                Signals::new([SIGINT, SIGTERM, SIGTERM, SIGHUP, SIGQUIT, SIGUSR1, SIGUSR2])?;
+            sighandle = Some(signals.handle());
+            let tx = tx.clone();
+            thread::spawn(move || {
+                for sig in &mut signals {
+                    tx.send(ChildProcessOutput::Signal(sig)).unwrap();
+                }
+            });
+        }
+        let id = cp.id();
         thread::spawn(move || {
             let status = cp.wait().unwrap();
+            if let Some(sighandle) = sighandle {
+                sighandle.close();
+            }
             tx.send(ChildProcessOutput::ExitStatus(status)).unwrap();
-            tx.send(ChildProcessOutput::Done).unwrap();
         });
         let mut combined_output = vec![];
-        let mut wait_for_count = 3;
         let mut status = None;
         for line in rx {
             match line {
@@ -280,10 +301,9 @@ impl<'a> CmdLineRunner<'a> {
                 ChildProcessOutput::ExitStatus(s) => {
                     status = Some(s);
                 }
-                ChildProcessOutput::Done => {
-                    wait_for_count -= 1;
-                    if wait_for_count == 0 {
-                        break;
+                ChildProcessOutput::Signal(sig) => {
+                    if sig != SIGINT {
+                        cmd!("kill", format!("-{sig}"), id.to_string()).run()?;
                     }
                 }
             }
@@ -375,7 +395,7 @@ enum ChildProcessOutput {
     Stdout(String),
     Stderr(String),
     ExitStatus(ExitStatus),
-    Done,
+    Signal(i32),
 }
 
 #[cfg(test)]

diff --git a/src/ui/ctrlc.rs b/src/ui/ctrlc.rs
@@ -0,0 +1,36 @@
+use console::Term;
+use signal_hook::consts::SIGINT;
+use signal_hook::iterator::{Handle, Signals};
+use std::process::exit;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::thread;
+
+#[derive(Debug)]
+pub struct HandleGuard(Handle);
+
+/// ensures cursor is displayed on ctrl-c
+pub fn handle_ctrlc() -> eyre::Result<Option<HandleGuard>> {
+    static HANDLED: AtomicBool = AtomicBool::new(false);
+    let handled = HANDLED.swap(true, Ordering::Relaxed);
+    if handled {
+        return Ok(None);
+    }
+
+    let mut signals = Signals::new([SIGINT])?;
+    let handle = HandleGuard(signals.handle());
+    thread::spawn(move || {
+        if signals.into_iter().next().is_some() {
+            let _ = Term::stderr().show_cursor();
+            debug!("Ctrl-C pressed, exiting...");
+            exit(1);
+        }
+        HANDLED.store(false, Ordering::Relaxed);
+    });
+    Ok(Some(handle))
+}
+
+impl Drop for HandleGuard {
+    fn drop(&mut self) {
+        self.0.close();
+    }
+}
diff --git a/src/ui/mod.rs b/src/ui/mod.rs
@@ -1,26 +1,9 @@
-use std::process::exit;
-use std::sync::Once;
-
-use console::{user_attended_stderr, Term};
-
 pub use prompt::confirm;
 
+pub mod ctrlc;
 pub mod multi_progress_report;
 pub mod progress_report;
 pub mod prompt;
 pub mod style;
 pub mod table;
 pub mod tree;
-
-pub fn handle_ctrlc() {
-    static ONCE: Once = Once::new();
-    ONCE.call_once(|| {
-        if user_attended_stderr() {
-            let _ = ctrlc::set_handler(move || {
-                let _ = Term::stderr().show_cursor();
-                debug!("Ctrl-C pressed, exiting...");
-                exit(1);
-            });
-        }
-    });
-}
diff --git a/src/ui/progress_report.rs b/src/ui/progress_report.rs
@@ -28,6 +28,8 @@ pub struct ProgressReport {
     pub pb: ProgressBar,
     prefix: String,
     pad: usize,
+    #[allow(dead_code)]
+    handle: Option<ui::ctrlc::HandleGuard>,
 }
 
 static LONGEST_PLUGIN_NAME: Lazy<usize> = Lazy::new(|| {
@@ -56,13 +58,18 @@ fn success_prefix(pad: usize, prefix: &str) -> String {
 
 impl ProgressReport {
     pub fn new(prefix: String) -> ProgressReport {
-        ui::handle_ctrlc();
+        let handle = ui::ctrlc::handle_ctrlc().unwrap_or_default();
         let pad = *LONGEST_PLUGIN_NAME;
         let pb = ProgressBar::new(100)
             .with_style(PROG_TEMPLATE.clone())
             .with_prefix(normal_prefix(pad, &prefix));
         pb.enable_steady_tick(Duration::from_millis(250));
-        ProgressReport { prefix, pb, pad }
+        ProgressReport {
+            prefix,
+            pb,
+            pad,
+            handle,
+        }
     }
 }
 

diff --git a/src/ui/prompt.rs b/src/ui/prompt.rs
@@ -1,18 +1,18 @@
-use std::io;
 use std::sync::Mutex;
 
 use demand::Confirm;
 
-use crate::ui;
+use crate::ui::ctrlc;
 
 static MUTEX: Mutex<()> = Mutex::new(());
 
-pub fn confirm<S: Into<String>>(message: S) -> io::Result<bool> {
+pub fn confirm<S: Into<String>>(message: S) -> eyre::Result<bool> {
     let _lock = MUTEX.lock().unwrap(); // Prevent multiple prompts at once
-    ui::handle_ctrlc();
+    let _ = ctrlc::handle_ctrlc()?;
 
     if !console::user_attended_stderr() {
         return Ok(false);
     }
-    Confirm::new(message).run()
+    let result = Confirm::new(message).run()?;
+    Ok(result)
 }
diff --git a/test/fixtures/signal-test.js b/test/fixtures/signal-test.js
@@ -0,0 +1,14 @@
+let i = 3
+
+process.on('SIGINT', function () {
+  if (i > 0) {
+    console.log(`Got SIGINT.  Press Control-D to exit. ${i} times left`)
+    i--
+  } else {
+    process.exit()
+  }
+})
+
+// wait for 60 seconds
+setTimeout(function () {}, 60000)
+console.log('Running.  Press Control-C to test.')