Jobs (#14883)

# Description

This is an attempt to improve the nushell situation with regard to issue
#247.

This PR implements:
- [X] spawning jobs: `job spawn { do_background_thing }`
Jobs will be implemented as threads and not forks, to maintain a
consistent behavior between unix and windows.

- [X] listing running jobs: `job list`
This should allow users to list what background tasks they currently
have running.

- [X] killing jobs: `job kill <id>`
- [X] interupting nushell code in the job's background thread
- [X] interrupting the job's currently-running process, if any.

Things that should be taken into consideration for implementation:
- [X] (unix-only) Handling `TSTP` signals while executing code and
turning the current program into a background job, and unfreezing them
in foreground `job unfreeze`.

- [X] Ensuring processes spawned by background jobs get distinct process
groups from the nushell shell itself

This PR originally aimed to implement some of the following, but it is
probably ideal to be left for another PR (scope creep)
- Disowning external process jobs (`job dispatch`)
- Inter job communication (`job send/recv`)

Roadblocks encountered so far:
- Nushell does some weird terminal sequence magics which make so that
when a background process or thread prints something to stderr and the
prompt is idle, the stderr output ends up showing up weirdly

crates/nu-protocol/src/engine/engine_state.rs

@@ -31,6 +31,8 @@ type PoisonDebuggerError<'a> = PoisonError<MutexGuard<'a, Box<dyn Debugger>>>;
 #[cfg(feature = "plugin")]
 use crate::{PluginRegistryFile, PluginRegistryItem, RegisteredPlugin};
 
+use super::{Jobs, ThreadJob};
+
 #[derive(Clone, Debug)]
 pub enum VirtualPath {
     File(FileId),
@@ -111,6 +113,19 @@ pub struct EngineState {
     startup_time: i64,
     is_debugging: IsDebugging,
     pub debugger: Arc<Mutex<Box<dyn Debugger>>>,
+
+    pub jobs: Arc<Mutex<Jobs>>,
+
+    // The job being executed with this engine state, or None if main thread
+    pub current_thread_job: Option<ThreadJob>,
+
+    // When there are background jobs running, the interactive behavior of `exit` changes depending on
+    // the value of this flag:
+    // - if this is false, then a warning about running jobs is shown and `exit` enables this flag
+    // - if this is true, then `exit` will `std::process::exit`
+    //
+    // This ensures that running exit twice will terminate the program correctly
+    pub exit_warning_given: Arc<AtomicBool>,
 }
 
 // The max number of compiled regexes to keep around in a LRU cache, arbitrarily chosen
@@ -180,6 +195,9 @@ impl EngineState {
             startup_time: -1,
             is_debugging: IsDebugging::new(false),
             debugger: Arc::new(Mutex::new(Box::new(NoopDebugger))),
+            jobs: Arc::new(Mutex::new(Jobs::default())),
+            current_thread_job: None,
+            exit_warning_given: Arc::new(AtomicBool::new(false)),
         }
     }
 
@@ -1036,6 +1054,9 @@ impl EngineState {
                 cursor_pos: 0,
             }));
         }
+        if Mutex::is_poisoned(&self.jobs) {
+            self.jobs = Arc::new(Mutex::new(Jobs::default()));
+        }
         if Mutex::is_poisoned(&self.regex_cache) {
             self.regex_cache = Arc::new(Mutex::new(LruCache::new(
                 NonZeroUsize::new(REGEX_CACHE_SIZE).expect("tried to create cache of size zero"),
@@ -1056,6 +1077,11 @@ impl EngineState {
             .position(|sp| sp == &span)
             .map(SpanId::new)
     }
+
+    // Determines whether the current state is being held by a background job
+    pub fn is_background_job(&self) -> bool {
+        self.current_thread_job.is_some()
+    }
 }
 
 impl GetSpan for &EngineState {

crates/nu-protocol/src/engine/jobs.rs

@@ -0,0 +1,219 @@
+use std::{
+    collections::{HashMap, HashSet},
+    sync::{Arc, Mutex},
+};
+
+use nu_system::{kill_by_pid, UnfreezeHandle};
+
+use crate::Signals;
+
+use crate::JobId;
+
+pub struct Jobs {
+    next_job_id: usize,
+
+    // this is the ID of the most recently added frozen job in the jobs table.
+    // the methods of this struct must ensure the invariant of this always
+    // being None or pointing to a valid job in the table
+    last_frozen_job_id: Option<JobId>,
+    jobs: HashMap<JobId, Job>,
+}
+
+impl Default for Jobs {
+    fn default() -> Self {
+        Self {
+            next_job_id: 1,
+            last_frozen_job_id: None,
+            jobs: HashMap::default(),
+        }
+    }
+}
+
+impl Jobs {
+    pub fn iter(&self) -> impl Iterator<Item = (JobId, &Job)> {
+        self.jobs.iter().map(|(k, v)| (*k, v))
+    }
+
+    pub fn lookup(&self, id: JobId) -> Option<&Job> {
+        self.jobs.get(&id)
+    }
+
+    pub fn remove_job(&mut self, id: JobId) -> Option<Job> {
+        if self.last_frozen_job_id.is_some_and(|last| id == last) {
+            self.last_frozen_job_id = None;
+        }
+
+        self.jobs.remove(&id)
+    }
+
+    fn assign_last_frozen_id_if_frozen(&mut self, id: JobId, job: &Job) {
+        if let Job::Frozen(_) = job {
+            self.last_frozen_job_id = Some(id);
+        }
+    }
+
+    pub fn add_job(&mut self, job: Job) -> JobId {
+        let this_id = JobId::new(self.next_job_id);
+
+        self.assign_last_frozen_id_if_frozen(this_id, &job);
+
+        self.jobs.insert(this_id, job);
+        self.next_job_id += 1;
+
+        this_id
+    }
+
+    pub fn most_recent_frozen_job_id(&mut self) -> Option<JobId> {
+        self.last_frozen_job_id
+    }
+
+    // this is useful when you want to remove a job from the list and add it back later
+    pub fn add_job_with_id(&mut self, id: JobId, job: Job) -> Result<(), &'static str> {
+        self.assign_last_frozen_id_if_frozen(id, &job);
+
+        if let std::collections::hash_map::Entry::Vacant(e) = self.jobs.entry(id) {
+            e.insert(job);
+            Ok(())
+        } else {
+            Err("job already exists")
+        }
+    }
+
+    /// This function tries to forcefully kill a job from this job table,
+    /// removes it from the job table. It always succeeds in removing the job
+    /// from the table, but may fail in killing the job's active processes.
+    pub fn kill_and_remove(&mut self, id: JobId) -> std::io::Result<()> {
+        if let Some(job) = self.jobs.get(&id) {
+            let err = job.kill();
+
+            self.remove_job(id);
+
+            err?
+        }
+
+        Ok(())
+    }
+
+    /// This function tries to forcefully kill all the background jobs and
+    /// removes all of them from the job table.
+    ///
+    /// It returns an error if any of the job killing attempts fails, but always
+    /// succeeds in removing the jobs from the table.
+    pub fn kill_all(&mut self) -> std::io::Result<()> {
+        self.last_frozen_job_id = None;
+
+        self.jobs.clear();
+
+        let first_err = self
+            .iter()
+            .map(|(_, job)| job.kill().err())
+            .fold(None, |acc, x| acc.or(x));
+
+        if let Some(err) = first_err {
+            Err(err)
+        } else {
+            Ok(())
+        }
+    }
+}
+
+pub enum Job {
+    Thread(ThreadJob),
+    Frozen(FrozenJob),
+}
+
+// A thread job represents a job that is currently executing as a background thread in nushell.
+// This is an Arc-y type, cloning it does not uniquely clone the information of this particular
+// job.
+
+// Although rust's documentation does not document the acquire-release semantics of Mutex, this
+// is a direct undocumentented requirement of its soundness, and is thus assumed by this
+// implementaation.
+// see issue https://github.com/rust-lang/rust/issues/126239.
+#[derive(Clone)]
+pub struct ThreadJob {
+    signals: Signals,
+    pids: Arc<Mutex<HashSet<u32>>>,
+}
+
+impl ThreadJob {
+    pub fn new(signals: Signals) -> Self {
+        ThreadJob {
+            signals,
+            pids: Arc::new(Mutex::new(HashSet::default())),
+        }
+    }
+
+    /// Tries to add the provided pid to the active pid set of the current job.
+    ///
+    /// Returns true if the pid was added successfully, or false if the
+    /// current job is interrupted.
+    pub fn try_add_pid(&self, pid: u32) -> bool {
+        let mut pids = self.pids.lock().expect("PIDs lock was poisoned");
+
+        // note: this signals check must occur after the pids lock has been locked.
+        if self.signals.interrupted() {
+            false
+        } else {
+            pids.insert(pid);
+            true
+        }
+    }
+
+    pub fn collect_pids(&self) -> Vec<u32> {
+        let lock = self.pids.lock().expect("PID lock was poisoned");
+
+        lock.iter().copied().collect()
+    }
+
+    pub fn kill(&self) -> std::io::Result<()> {
+        // it's okay to make this interrupt outside of the mutex, since it has acquire-release
+        // semantics.
+
+        self.signals.trigger();
+
+        let mut pids = self.pids.lock().expect("PIDs lock was poisoned");
+
+        for pid in pids.iter() {
+            kill_by_pid((*pid).into())?;
+        }
+
+        pids.clear();
+
+        Ok(())
+    }
+
+    pub fn remove_pid(&self, pid: u32) {
+        let mut pids = self.pids.lock().expect("PID lock was poisoned");
+
+        pids.remove(&pid);
+    }
+}
+
+impl Job {
+    pub fn kill(&self) -> std::io::Result<()> {
+        match self {
+            Job::Thread(thread_job) => thread_job.kill(),
+            Job::Frozen(frozen_job) => frozen_job.kill(),
+        }
+    }
+}
+
+pub struct FrozenJob {
+    pub unfreeze: UnfreezeHandle,
+}
+
+impl FrozenJob {
+    pub fn kill(&self) -> std::io::Result<()> {
+        #[cfg(unix)]
+        {
+            kill_by_pid(self.unfreeze.pid() as i64)
+        }
+
+        // it doesn't happen outside unix.
+        #[cfg(not(unix))]
+        {
+            Ok(())
+        }
+    }
+}

crates/nu-protocol/src/engine/mod.rs

@@ -8,6 +8,7 @@ mod command;
 mod description;
 mod engine_state;
 mod error_handler;
+mod jobs;
 mod overlay;
 mod pattern_match;
 mod sequence;
@@ -26,6 +27,7 @@ pub use capture_block::*;
 pub use command::*;
 pub use engine_state::*;
 pub use error_handler::*;
+pub use jobs::*;
 pub use overlay::*;
 pub use pattern_match::*;
 pub use sequence::*;

crates/nu-protocol/src/errors/shell_error/mod.rs

@@ -1326,6 +1326,40 @@ On Windows, this would be %USERPROFILE%\AppData\Roaming"#
         span: Option<Span>,
     },
 
+    #[error("Job {id} not found")]
+    #[diagnostic(
+        code(nu::shell::job_not_found),
+        help(
+            "The operation could not be completed, there is no job currently running with this id"
+        )
+    )]
+    JobNotFound {
+        id: usize,
+        #[label = "job not found"]
+        span: Span,
+    },
+
+    #[error("No frozen job to unfreeze")]
+    #[diagnostic(
+        code(nu::shell::no_frozen_job),
+        help("There is currently no frozen job to unfreeze")
+    )]
+    NoFrozenJob {
+        #[label = "no frozen job"]
+        span: Span,
+    },
+
+    #[error("Job {id} is not frozen")]
+    #[diagnostic(
+        code(nu::shell::os_disabled),
+        help("You tried to unfreeze a job which is not frozen")
+    )]
+    JobNotFrozen {
+        id: usize,
+        #[label = "job not frozen"]
+        span: Span,
+    },
+
     #[error(transparent)]
     #[diagnostic(transparent)]
     ChainedError(ChainedError),

crates/nu-protocol/src/id.rs

@@ -94,6 +94,8 @@ pub mod marker {
     pub struct Span;
     #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
     pub struct Reg;
+    #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub struct Job;
 }
 
 pub type VarId = Id<marker::Var>;
@@ -104,6 +106,7 @@ pub type OverlayId = Id<marker::Overlay>;
 pub type FileId = Id<marker::File>;
 pub type VirtualPathId = Id<marker::VirtualPath>;
 pub type SpanId = Id<marker::Span>;
+pub type JobId = Id<marker::Job>;
 
 /// An ID for an [IR](crate::ir) register.
 ///

crates/nu-protocol/src/process/child.rs

@@ -1,5 +1,6 @@
 use crate::{byte_stream::convert_file, shell_error::io::IoError, ShellError, Span};
-use nu_system::{ExitStatus, ForegroundChild};
+use nu_system::{ExitStatus, ForegroundChild, ForegroundWaitStatus};
+
 use os_pipe::PipeReader;
 use std::{
     fmt::Debug,
@@ -8,7 +9,7 @@ use std::{
     thread,
 };
 
-fn check_ok(status: ExitStatus, ignore_error: bool, span: Span) -> Result<(), ShellError> {
+pub fn check_ok(status: ExitStatus, ignore_error: bool, span: Span) -> Result<(), ShellError> {
     match status {
         ExitStatus::Exited(exit_code) => {
             if ignore_error {
@@ -165,12 +166,21 @@ pub struct ChildProcess {
     span: Span,
 }
 
+pub struct PostWaitCallback(pub Box<dyn FnOnce(ForegroundWaitStatus) + Send>);
+
+impl Debug for PostWaitCallback {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "<wait_callback>")
+    }
+}
+
 impl ChildProcess {
     pub fn new(
         mut child: ForegroundChild,
         reader: Option<PipeReader>,
         swap: bool,
         span: Span,
+        callback: Option<PostWaitCallback>,
     ) -> Result<Self, ShellError> {
         let (stdout, stderr) = if let Some(combined) = reader {
             (Some(combined), None)
@@ -190,7 +200,32 @@ impl ChildProcess {
 
         thread::Builder::new()
             .name("exit status waiter".into())
-            .spawn(move || exit_status_sender.send(child.wait()))
+            .spawn(move || {
+                let matched = match child.wait() {
+                    // there are two possible outcomes when we `wait` for a process to finish:
+                    // 1. the process finishes as usual
+                    // 2. (unix only) the process gets signaled with SIGTSTP
+                    //
+                    // in the second case, although the process may still be alive in a
+                    // cryonic state, we explicitly treat as it has finished with exit code 0
+                    // for the sake of the current pipeline
+                    Ok(wait_status) => {
+                        let next = match &wait_status {
+                            ForegroundWaitStatus::Frozen(_) => ExitStatus::Exited(0),
+                            ForegroundWaitStatus::Finished(exit_status) => *exit_status,
+                        };
+
+                        if let Some(callback) = callback {
+                            (callback.0)(wait_status);
+                        }
+
+                        Ok(next)
+                    }
+                    Err(err) => Err(err),
+                };
+
+                exit_status_sender.send(matched)
+            })
             .map_err(|err| {
                 IoError::new_with_additional_context(
                     err.kind(),