From 4af0a6a3fad806df5af850bbf655042ae47d6369 Mon Sep 17 00:00:00 2001 From: unrelentingtech Date: Thu, 29 Sep 2022 21:37:48 +0300 Subject: [PATCH] Foreground process group management, again (#6584) * Revert "Revert "Try again: in unix like system, set foreground process while running external command (#6273)" (#6542)" This reverts commit 2bb367f570502cd26d6b7a81425d9f3b3b9eb432. * Make foreground job control hopefully work correctly These changes are mostly inspired by the glibc manual. * Fix typo in external command description * Only restore tty control to shell when no fg procs are left; reuse pgrp * Rework terminal acquirement code to be like fish Co-authored-by: Darren Schroeder <343840+fdncred@users.noreply.github.com> --- Cargo.lock | 5 + Cargo.toml | 5 + crates/nu-command/src/system/run_external.rs | 33 ++-- crates/nu-protocol/src/engine/engine_state.rs | 7 +- crates/nu-system/Cargo.toml | 4 + crates/nu-system/src/foreground.rs | 180 ++++++++++++++++++ crates/nu-system/src/lib.rs | 2 + src/main.rs | 100 ++++++++++ 8 files changed, 323 insertions(+), 13 deletions(-) create mode 100644 crates/nu-system/src/foreground.rs diff --git a/Cargo.lock b/Cargo.lock index a31753c71..3424faa6a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2511,6 +2511,7 @@ dependencies = [ "bitflags", "cfg-if 1.0.0", "libc", + "memoffset", ] [[package]] @@ -2580,6 +2581,7 @@ name = "nu" version = "0.69.2" dependencies = [ "assert_cmd", + "atty", "chrono", "crossterm 0.24.0", "ctrlc", @@ -2588,6 +2590,7 @@ dependencies = [ "itertools", "log", "miette 5.1.1", + "nix", "nu-ansi-term", "nu-cli", "nu-color-config", @@ -2873,11 +2876,13 @@ dependencies = [ name = "nu-system" version = "0.69.2" dependencies = [ + "atty", "chrono", "errno", "libc", "libproc", "mach2", + "nix", "ntapi", "once_cell", "procfs", diff --git a/Cargo.toml b/Cargo.toml index 5f0b70db7..2043aeae8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -122,3 +122,8 @@ debug = false [[bin]] name = "nu" path = "src/main.rs" + +[target.'cfg(target_family = "unix")'.dependencies] +nix = "0.24" +atty = "0.2" + diff --git a/crates/nu-command/src/system/run_external.rs b/crates/nu-command/src/system/run_external.rs index 788181f23..6e97823f8 100644 --- a/crates/nu-command/src/system/run_external.rs +++ b/crates/nu-command/src/system/run_external.rs @@ -7,6 +7,7 @@ use nu_protocol::did_you_mean; use nu_protocol::engine::{EngineState, Stack}; use nu_protocol::{ast::Call, engine::Command, ShellError, Signature, SyntaxShape, Value}; use nu_protocol::{Category, Example, ListStream, PipelineData, RawStream, Span, Spanned}; +use nu_system::ForegroundProcess; use pathdiff::diff_paths; use std::collections::HashMap; use std::io::{BufRead, BufReader, Write}; @@ -34,7 +35,7 @@ impl Command for External { Signature::build(self.name()) .switch("redirect-stdout", "redirect-stdout", None) .switch("redirect-stderr", "redirect-stderr", None) - .required("command", SyntaxShape::Any, "external comamdn to run") + .required("command", SyntaxShape::Any, "external command to run") .rest("args", SyntaxShape::Any, "arguments for external command") .category(Category::System) } @@ -141,7 +142,10 @@ impl ExternalCommand { let ctrlc = engine_state.ctrlc.clone(); - let mut process = self.create_process(&input, false, head)?; + let mut fg_process = ForegroundProcess::new( + self.create_process(&input, false, head)?, + engine_state.pipeline_externals_state.clone(), + ); // mut is used in the windows branch only, suppress warning on other platforms #[allow(unused_mut)] let mut child; @@ -156,8 +160,7 @@ impl ExternalCommand { // fails to be run as a normal executable: // 1. "shell out" to cmd.exe if the command is a known cmd.exe internal command // 2. Otherwise, use `which-rs` to look for batch files etc. then run those in cmd.exe - - match process.spawn() { + match fg_process.spawn() { Err(err) => { // set the default value, maybe we'll override it later child = Err(err); @@ -174,7 +177,10 @@ impl ExternalCommand { .any(|&cmd| command_name_upper == cmd); if looks_like_cmd_internal { - let mut cmd_process = self.create_process(&input, true, head)?; + let mut cmd_process = ForegroundProcess::new( + self.create_process(&input, true, head)?, + engine_state.pipeline_externals_state.clone(), + ); child = cmd_process.spawn(); } else { #[cfg(feature = "which-support")] @@ -202,8 +208,11 @@ impl ExternalCommand { item: file_name.to_string_lossy().to_string(), span: self.name.span, }; - let mut cmd_process = new_command - .create_process(&input, true, head)?; + let mut cmd_process = ForegroundProcess::new( + new_command + .create_process(&input, true, head)?, + engine_state.pipeline_externals_state.clone(), + ); child = cmd_process.spawn(); } } @@ -221,7 +230,7 @@ impl ExternalCommand { #[cfg(not(windows))] { - child = process.spawn() + child = fg_process.spawn() } match child { @@ -273,7 +282,7 @@ impl ExternalCommand { engine_state.config.use_ansi_coloring = false; // if there is a string or a stream, that is sent to the pipe std - if let Some(mut stdin_write) = child.stdin.take() { + if let Some(mut stdin_write) = child.as_mut().stdin.take() { std::thread::spawn(move || { let input = crate::Table::run( &crate::Table, @@ -314,7 +323,7 @@ impl ExternalCommand { // and we create a ListStream that can be consumed if redirect_stderr { - let stderr = child.stderr.take().ok_or_else(|| { + let stderr = child.as_mut().stderr.take().ok_or_else(|| { ShellError::ExternalCommand( "Error taking stderr from external".to_string(), "Redirects need access to stderr of an external command" @@ -353,7 +362,7 @@ impl ExternalCommand { } if redirect_stdout { - let stdout = child.stdout.take().ok_or_else(|| { + let stdout = child.as_mut().stdout.take().ok_or_else(|| { ShellError::ExternalCommand( "Error taking stdout from external".to_string(), "Redirects need access to stdout of an external command" @@ -391,7 +400,7 @@ impl ExternalCommand { } } - match child.wait() { + match child.as_mut().wait() { Err(err) => Err(ShellError::ExternalCommand( "External command exited with error".into(), err.to_string(), diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs index 0cf719ad0..d8020900a 100644 --- a/crates/nu-protocol/src/engine/engine_state.rs +++ b/crates/nu-protocol/src/engine/engine_state.rs @@ -10,7 +10,10 @@ use std::path::Path; use std::path::PathBuf; use std::{ collections::{HashMap, HashSet, VecDeque}, - sync::{atomic::AtomicBool, Arc, Mutex}, + sync::{ + atomic::{AtomicBool, AtomicU32}, + Arc, Mutex, + }, }; static PWD_ENV: &str = "PWD"; @@ -80,6 +83,7 @@ pub struct EngineState { pub env_vars: EnvVars, pub previous_env_vars: HashMap, pub config: Config, + pub pipeline_externals_state: Arc<(AtomicU32, AtomicU32)>, pub repl_buffer_state: Arc>>, pub repl_operation_queue: Arc>>, #[cfg(feature = "plugin")] @@ -121,6 +125,7 @@ impl EngineState { env_vars: EnvVars::from([(DEFAULT_OVERLAY_NAME.to_string(), HashMap::new())]), previous_env_vars: HashMap::new(), config: Config::default(), + pipeline_externals_state: Arc::new((AtomicU32::new(0), AtomicU32::new(0))), repl_buffer_state: Arc::new(Mutex::new(None)), repl_operation_queue: Arc::new(Mutex::new(VecDeque::new())), #[cfg(feature = "plugin")] diff --git a/crates/nu-system/Cargo.toml b/crates/nu-system/Cargo.toml index 1e0d95acc..a1951fb9b 100644 --- a/crates/nu-system/Cargo.toml +++ b/crates/nu-system/Cargo.toml @@ -16,6 +16,10 @@ path = "src/main.rs" [dependencies] libc = "0.2" +[target.'cfg(target_family = "unix")'.dependencies] +nix = "0.24" +atty = "0.2" + [target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies] procfs = "0.14.0" diff --git a/crates/nu-system/src/foreground.rs b/crates/nu-system/src/foreground.rs new file mode 100644 index 000000000..9c9a704e0 --- /dev/null +++ b/crates/nu-system/src/foreground.rs @@ -0,0 +1,180 @@ +use std::{ + process::{Child, Command}, + sync::{ + atomic::{AtomicU32, Ordering}, + Arc, + }, +}; + +/// A simple wrapper for `std::process::Command` +/// +/// ## Spawn behavior +/// ### Unix +/// +/// The spawned child process will get its own process group id, and it's going to foreground (by making stdin belong's to child's process group). +/// +/// On drop, the calling process's group will become the foreground process group once again. +/// +/// ### Windows +/// It does nothing special on windows system, `spawn` is the same as [std::process::Command::spawn](std::process::Command::spawn) +pub struct ForegroundProcess { + inner: Command, + pipeline_state: Arc<(AtomicU32, AtomicU32)>, +} + +/// A simple wrapper for `std::process::Child` +/// +/// It can only be created by `ForegroundProcess::spawn`. +pub struct ForegroundChild { + inner: Child, + pipeline_state: Arc<(AtomicU32, AtomicU32)>, +} + +impl ForegroundProcess { + pub fn new(cmd: Command, pipeline_state: Arc<(AtomicU32, AtomicU32)>) -> Self { + Self { + inner: cmd, + pipeline_state, + } + } + + pub fn spawn(&mut self) -> std::io::Result { + let (ref pgrp, ref pcnt) = *self.pipeline_state; + let existing_pgrp = pgrp.load(Ordering::SeqCst); + fg_process_setup::prepare_to_foreground(&mut self.inner, existing_pgrp); + self.inner + .spawn() + .map(|child| { + fg_process_setup::set_foreground(&child, existing_pgrp); + let _ = pcnt.fetch_add(1, Ordering::SeqCst); + if existing_pgrp == 0 { + pgrp.store(child.id(), Ordering::SeqCst); + } + ForegroundChild { + inner: child, + pipeline_state: self.pipeline_state.clone(), + } + }) + .map_err(|e| { + fg_process_setup::reset_foreground_id(); + e + }) + } +} + +impl AsMut for ForegroundChild { + fn as_mut(&mut self) -> &mut Child { + &mut self.inner + } +} + +impl Drop for ForegroundChild { + fn drop(&mut self) { + let (ref pgrp, ref pcnt) = *self.pipeline_state; + if pcnt.fetch_sub(1, Ordering::SeqCst) == 1 { + pgrp.store(0, Ordering::SeqCst); + fg_process_setup::reset_foreground_id() + } + } +} + +// It's a simpler version of fish shell's external process handling. +#[cfg(target_family = "unix")] +mod fg_process_setup { + use nix::{ + sys::signal, + unistd::{self, Pid}, + }; + use std::os::unix::prelude::{CommandExt, RawFd}; + + // TODO: when raising MSRV past 1.63.0, switch to OwnedFd + struct TtyHandle(RawFd); + + impl Drop for TtyHandle { + fn drop(&mut self) { + let _ = unistd::close(self.0); + } + } + + pub(super) fn prepare_to_foreground( + external_command: &mut std::process::Command, + existing_pgrp: u32, + ) { + let tty = TtyHandle(unistd::dup(nix::libc::STDIN_FILENO).expect("dup")); + unsafe { + // Safety: + // POSIX only allows async-signal-safe functions to be called. + // `sigprocmask`, `setpgid` and `tcsetpgrp` are async-signal-safe according to: + // https://manpages.ubuntu.com/manpages/bionic/man7/signal-safety.7.html + external_command.pre_exec(move || { + // When this callback is run, std::process has already done: + // - pthread_sigmask(SIG_SETMASK) with an empty sigset + // - signal(SIGPIPE, SIG_DFL) + // However, we do need TTOU/TTIN blocked again during this setup. + let mut sigset = signal::SigSet::empty(); + sigset.add(signal::Signal::SIGTSTP); + sigset.add(signal::Signal::SIGTTOU); + sigset.add(signal::Signal::SIGTTIN); + sigset.add(signal::Signal::SIGCHLD); + signal::sigprocmask(signal::SigmaskHow::SIG_BLOCK, Some(&sigset), None) + .expect("signal mask"); + + // According to glibc's job control manual: + // https://www.gnu.org/software/libc/manual/html_node/Launching-Jobs.html + // This has to be done *both* in the parent and here in the child due to race conditions. + set_foreground_pid(unistd::getpid(), existing_pgrp, tty.0); + + // Now let the child process have all the signals by resetting with SIG_SETMASK. + let mut sigset = signal::SigSet::empty(); + sigset.add(signal::Signal::SIGTSTP); // for now not really all: we don't support background jobs, so keep this one blocked + signal::sigprocmask(signal::SigmaskHow::SIG_SETMASK, Some(&sigset), None) + .expect("signal mask"); + + Ok(()) + }); + } + } + + pub(super) fn set_foreground(process: &std::process::Child, existing_pgrp: u32) { + // called from the parent shell process - do the stdin tty check here + if atty::is(atty::Stream::Stdin) { + set_foreground_pid( + Pid::from_raw(process.id() as i32), + existing_pgrp, + nix::libc::STDIN_FILENO, + ); + } + } + + // existing_pgrp is 0 when we don't have an existing foreground process in the pipeline. + // Conveniently, 0 means "current pid" to setpgid. But not to tcsetpgrp. + fn set_foreground_pid(pid: Pid, existing_pgrp: u32, tty: RawFd) { + let _ = unistd::setpgid(pid, Pid::from_raw(existing_pgrp as i32)); + let _ = unistd::tcsetpgrp( + tty, + if existing_pgrp == 0 { + pid + } else { + Pid::from_raw(existing_pgrp as i32) + }, + ); + } + + /// Reset the foreground process group to the shell + pub(super) fn reset_foreground_id() { + if atty::is(atty::Stream::Stdin) { + if let Err(e) = nix::unistd::tcsetpgrp(nix::libc::STDIN_FILENO, unistd::getpgrp()) { + println!("ERROR: reset foreground id failed, tcsetpgrp result: {e:?}"); + } + } + } +} + +#[cfg(not(target_family = "unix"))] +mod fg_process_setup { + pub(super) fn prepare_to_foreground(_: &mut std::process::Command, _: u32) {} + + pub(super) fn set_foreground(_: &std::process::Child, _: u32) {} + + pub(super) fn reset_foreground_id() {} +} diff --git a/crates/nu-system/src/lib.rs b/crates/nu-system/src/lib.rs index 1ad89d73d..2283a94b7 100644 --- a/crates/nu-system/src/lib.rs +++ b/crates/nu-system/src/lib.rs @@ -1,3 +1,4 @@ +mod foreground; #[cfg(any(target_os = "android", target_os = "linux"))] mod linux; #[cfg(target_os = "macos")] @@ -5,6 +6,7 @@ mod macos; #[cfg(target_os = "windows")] mod windows; +pub use self::foreground::{ForegroundChild, ForegroundProcess}; #[cfg(any(target_os = "android", target_os = "linux"))] pub use self::linux::*; #[cfg(target_os = "macos")] diff --git a/src/main.rs b/src/main.rs index a16c4fd1c..0c66e7ed0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,6 +37,100 @@ use std::{ thread_local! { static IS_PERF: RefCell = RefCell::new(false) } +// Inspired by fish's acquire_tty_or_exit +#[cfg(unix)] +fn take_control(interactive: bool) { + use nix::{ + errno::Errno, + sys::signal::{self, SaFlags, SigAction, SigHandler, SigSet, Signal}, + unistd::{self, Pid}, + }; + + let shell_pgid = unistd::getpgrp(); + let owner_pgid = unistd::tcgetpgrp(nix::libc::STDIN_FILENO).expect("tcgetpgrp"); + + // Common case, nothing to do + if owner_pgid == shell_pgid { + return; + } + + // This can apparently happen with sudo: https://github.com/fish-shell/fish-shell/issues/7388 + if owner_pgid == unistd::getpid() { + let _ = unistd::setpgid(owner_pgid, owner_pgid); + return; + } + + // Reset all signal handlers to default + for sig in Signal::iterator() { + unsafe { + if let Ok(old_act) = signal::sigaction( + sig, + &SigAction::new(SigHandler::SigDfl, SaFlags::empty(), SigSet::empty()), + ) { + // fish preserves ignored SIGHUP, presumably for nohup support, so let's do the same + if sig == Signal::SIGHUP && old_act.handler() == SigHandler::SigIgn { + let _ = signal::sigaction(sig, &old_act); + } + } + } + } + + let mut success = false; + for _ in 0..4096 { + match unistd::tcgetpgrp(nix::libc::STDIN_FILENO) { + Ok(owner_pgid) if owner_pgid == shell_pgid => { + success = true; + break; + } + Ok(owner_pgid) if owner_pgid == Pid::from_raw(0) => { + // Zero basically means something like "not owned" and we can just take it + let _ = unistd::tcsetpgrp(nix::libc::STDIN_FILENO, shell_pgid); + } + Err(Errno::ENOTTY) => { + if !interactive { + // that's fine + return; + } + eprintln!("ERROR: no TTY for interactive shell"); + std::process::exit(1); + } + _ => { + // fish also has other heuristics than "too many attempts" for the orphan check, but they're optional + if signal::killpg(Pid::from_raw(-shell_pgid.as_raw()), Signal::SIGTTIN).is_err() { + eprintln!("ERROR: failed to SIGTTIN ourselves"); + std::process::exit(1); + } + } + } + } + if !success { + eprintln!("ERROR: failed take control of the terminal, we might be orphaned"); + std::process::exit(1); + } +} + +#[cfg(unix)] +fn acquire_terminal(interactive: bool) { + use nix::sys::signal::{signal, SigHandler, Signal}; + + if !atty::is(atty::Stream::Stdin) { + return; + } + + take_control(interactive); + + unsafe { + // SIGINT and SIGQUIT have special handling above + signal(Signal::SIGTSTP, SigHandler::SigIgn).expect("signal ignore"); + signal(Signal::SIGTTIN, SigHandler::SigIgn).expect("signal ignore"); + signal(Signal::SIGTTOU, SigHandler::SigIgn).expect("signal ignore"); + // signal::signal(Signal::SIGCHLD, SigHandler::SigIgn).expect("signal ignore"); // needed for std::command's waitpid usage + } +} + +#[cfg(not(unix))] +fn acquire_terminal(_: bool) {} + fn main() -> Result<()> { // miette::set_panic_hook(); let miette_hook = std::panic::take_hook(); @@ -155,6 +249,12 @@ fn main() -> Result<()> { match parsed_nu_cli_args { Ok(binary_args) => { + // keep this condition in sync with the branches below + acquire_terminal( + binary_args.commands.is_none() + && (script_name.is_empty() || binary_args.interactive_shell.is_some()), + ); + if let Some(t) = binary_args.threads { // 0 means to let rayon decide how many threads to use let threads = t.as_i64().unwrap_or(0);