pset 0.1.0

//! Orchestrate a set of child processes as a single stream of events.
//!
//! A [`ProcessSet`] owns any number of children, each labelled with a tag of
//! the caller's choosing, and hands out what they do — output on stdout, output
//! on stderr, exiting — one [`PsetEvent`] at a time:
//!
//! ```
//! use pset::{ProcessSet, PsetEvent};
//! use std::process::{Command, Stdio};
//!
//! #[derive(Debug, Clone, Copy, PartialEq)]
//! enum Tag {
//!     Date,
//!     Uname,
//! }
//!
//! # fn main() -> std::io::Result<()> {
//! let mut pset = pset::create::<Tag>()?;
//!
//! // Piping a stream is what turns it into events: see `ProcessSet::spawn`.
//! let mut cmd1 = Command::new("sh");
//! cmd1.arg("-c").arg("date").stdout(Stdio::piped()).stderr(Stdio::piped());
//!
//! let mut cmd2 = Command::new("sh");
//! cmd2.arg("-c").arg("uname").stdout(Stdio::piped()).stderr(Stdio::piped());
//!
//! pset.spawn(Tag::Date, cmd1)?;
//! pset.spawn(Tag::Uname, cmd2)?;
//!
//! while let Some((tag, event)) = pset.wait_next()? {
//!     match event {
//!         PsetEvent::Stdout(data) => println!("{tag:?} said {} bytes", data.len()),
//!         PsetEvent::Stderr(data) => eprintln!("{tag:?} complained: {} bytes", data.len()),
//!         PsetEvent::ProcessExited(status) => println!("{tag:?} exited: {status}"),
//!     }
//! }
//! # Ok(())
//! # }
//! ```
//!
//! `wait_next` returns `None` once the set is empty, so the loop above ends
//! when the last child is gone — no bookkeeping of who is still running.
//!
//! # What the set guarantees
//!
//! - **`ProcessExited` is the last event for a process.** It is emitted only
//!   after the child has been reaped *and* every pipe the set holds for it has
//!   reached end of file, so no output can appear after it and nothing written
//!   before the exit is lost. The flip side: a child that leaves a grandchild
//!   holding its piped stdout open is only reported as exited once that
//!   grandchild lets go.
//! - **No output is lost to a full pipe.** Every pipe is drained by the same
//!   loop that waits for the exits, so a child writing megabytes cannot
//!   deadlock against a parent waiting for it to finish.
//! - **Signals cannot go astray.** [`ProcessSet::kill`] never signals through
//!   a pid that could have been recycled: on Linux a pidfd carries it, and on
//!   macOS it only ever goes to a child the set has not reaped yet, whose pid
//!   therefore still belongs to it.
//!
//! # Platform
//!
//! Linux and macOS, each on its native way of making a process exit an
//! ordinary event in a kernel queue — `pidfd_open(2)` and `epoll(7)` on
//! Linux, `kqueue(2)` and `EVFILT_PROC` on macOS. Either way one thread waits
//! on the exits and the pipes together, with no signal handler and no thread
//! per pipe. The platform-specific plumbing lives in the `sys` module behind
//! a small common surface; everything above it is shared.

use std::{
    collections::{HashMap, VecDeque},
    fmt,
    io::{self, Read},
    os::fd::AsFd,
    process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, ExitStatus},
    time::{Duration, Instant},
};

mod sys;

/// Handle to one process in a set, returned by [`ProcessSet::spawn`].
///
/// Tags identify processes in events; this identifies a single process even
/// when several share a tag.
pub type ProcId = u64;

/// Something one of the processes in the set did.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PsetEvent {
    /// Bytes read from the process's stdout. Never empty; the chunking is
    /// whatever the pipe delivered, not lines.
    ///
    /// Only ever reported for a process whose command asked for
    /// `stdout(Stdio::piped())`; see [`ProcessSet::spawn`].
    Stdout(Vec<u8>),
    /// Bytes read from the process's stderr. Never empty.
    ///
    /// Only ever reported for a process whose command asked for
    /// `stderr(Stdio::piped())`; see [`ProcessSet::spawn`].
    Stderr(Vec<u8>),
    /// The process is gone and everything it wrote has already been reported.
    /// The last event for that process; it leaves the set with this one.
    ProcessExited(ExitStatus),
}

/// A set of child processes, waited on together.
///
/// See the [module documentation](crate) for the guarantees it makes.
pub struct ProcessSet<T> {
    poller: sys::Poller,
    procs: HashMap<u64, Proc<T>>,
    /// Events produced but not handed out yet: one wait can make several
    /// descriptors ready at once.
    pending: VecDeque<(T, PsetEvent)>,
    /// Tokens the last wait reported ready, reused between waits.
    tokens: Vec<u64>,
    buf: Vec<u8>,
    next_id: u64,
}

struct Proc<T> {
    tag: T,
    child: Child,
    exit_watch: sys::ExitWatch,
    /// `None` if the command never asked for a pipe, and again once the pipe
    /// has reached end of file and been unregistered.
    stdout: Option<ChildStdout>,
    stderr: Option<ChildStderr>,
    /// `Some` once the process has been reaped.
    status: Option<ExitStatus>,
}

impl<T> Proc<T> {
    /// Reaped, and every pipe the set holds drained: nothing more can come out
    /// of this process. A stream that was never piped is drained from the
    /// start, so a child with no pipes at all is finished as soon as it is
    /// reaped.
    fn finished(&self) -> bool {
        self.status.is_some() && self.stdout.is_none() && self.stderr.is_none()
    }
}

/// How much of a pipe to take in one read, and so the largest [`PsetEvent`]
/// payload. A pipe holds 64K by default, so one read empties a full one.
const READ_CHUNK: usize = 64 * 1024;

/// Readiness of a descriptor is reported as a `u64`; the low two bits say
/// which of a process's three descriptors it was, the rest is the process id.
const STREAM_STDOUT: u64 = 0;
const STREAM_STDERR: u64 = 1;
const STREAM_EXIT: u64 = 2;

fn token(id: u64, stream: u64) -> u64 {
    (id << 2) | stream
}

fn untoken(token: u64) -> (u64, u64) {
    (token >> 2, token & 0b11)
}

/// Create an empty [`ProcessSet`] tagged with `T`.
pub fn create<T: Clone>() -> io::Result<ProcessSet<T>> {
    ProcessSet::new()
}

impl<T: Clone> ProcessSet<T> {
    /// Create an empty set. Fails only if the process cannot afford another
    /// file descriptor for the kernel queue.
    pub fn new() -> io::Result<Self> {
        Ok(Self {
            poller: sys::Poller::new()?,
            procs: HashMap::new(),
            pending: VecDeque::new(),
            tokens: Vec::new(),
            buf: vec![0; READ_CHUNK],
            next_id: 0,
        })
    }

    /// How many processes are still in the set — spawned, and not yet reported
    /// as exited.
    pub fn len(&self) -> usize {
        self.procs.len()
    }

    /// Whether the set has anything left to report: no processes *and* no
    /// events waiting to be handed out.
    pub fn is_empty(&self) -> bool {
        self.procs.is_empty() && self.pending.is_empty()
    }

    /// Send `signal` to one process ([`libc::SIGTERM`] and friends).
    ///
    /// The process stays in the set: dying is reported as an ordinary
    /// [`PsetEvent::ProcessExited`], with the signal in its [`ExitStatus`].
    /// Fails with [`io::ErrorKind::NotFound`] for a process that has already
    /// left the set.
    pub fn kill(&self, id: ProcId, signal: i32) -> io::Result<()> {
        let proc = self
            .procs
            .get(&id)
            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "no such process in the set"))?;
        if proc.status.is_some() {
            // Already reaped, and only waiting for its pipes to drain. There
            // is no process left to signal — and no pid to trust.
            return Ok(());
        }
        proc.exit_watch.send_signal(signal)
    }

    /// Send `signal` to every process in the set that is still running.
    pub fn kill_all(&self, signal: i32) -> io::Result<()> {
        let mut first_err = None;
        for proc in self.procs.values() {
            if proc.status.is_some() {
                continue;
            }
            if let Err(e) = proc.exit_watch.send_signal(signal) {
                first_err.get_or_insert(e);
            }
        }
        match first_err {
            Some(e) => Err(e),
            None => Ok(()),
        }
    }

    /// Spawn `command` into the set under `tag`.
    ///
    /// The command is run exactly as the caller built it — program, arguments,
    /// environment, working directory, and all three standard streams. Nothing
    /// is reported synchronously; the process shows up as events from
    /// [`wait_next`](Self::wait_next).
    ///
    /// Returns the process's [`ProcId`] and, for a command that asked for
    /// `Stdio::piped()` stdin, the write end of that pipe, to feed the child
    /// input with. Any other stdin — the inherited default, `Stdio::null()`,
    /// a file — gives `None`; the set neither adds a stdin pipe nor takes one
    /// away.
    ///
    /// The pipe is handed over rather than kept by the set, because only the
    /// caller knows when the input is complete: the child sees end of file
    /// when the returned [`ChildStdin`] is dropped, and a child reading to end
    /// of file waits until then.
    ///
    /// ```
    /// use pset::PsetEvent;
    /// use std::{io::Write, process::{Command, Stdio}};
    ///
    /// # fn main() -> std::io::Result<()> {
    /// let mut pset = pset::create::<&str>()?;
    ///
    /// let mut cmd = Command::new("wc");
    /// cmd.arg("-c").stdin(Stdio::piped()).stdout(Stdio::piped());
    ///
    /// let (_id, stdin) = pset.spawn("counter", cmd)?;
    /// let mut stdin = stdin.expect("stdin was set to a pipe");
    /// stdin.write_all(b"five!")?;
    /// drop(stdin); // end of file: `wc` counts what it has and exits
    ///
    /// let mut counted = Vec::new();
    /// while let Some((_tag, event)) = pset.wait_next()? {
    ///     if let PsetEvent::Stdout(data) = event {
    ///         counted.extend_from_slice(&data);
    ///     }
    /// }
    /// assert_eq!(String::from_utf8_lossy(&counted).trim(), "5");
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Writing more than a pipe holds (64K, typically) blocks until the child
    /// reads, and the child may itself be blocked writing to a stdout that
    /// only [`wait_next`](Self::wait_next) drains — so feed a child more than
    /// that from another thread, or interleave the writes with `wait_next`
    /// calls, rather than filling stdin from the loop that drains the output.
    pub fn spawn(
        &mut self,
        tag: T,
        mut command: Command,
    ) -> io::Result<(ProcId, Option<ChildStdin>)> {
        let mut child = command.spawn()?;
        // Out of the `Child` and off to the caller: a stdin pipe the set held
        // on to would close only once the process left the set, and that is an
        // end of file the child may well be waiting for in order to get there.
        let stdin = child.stdin.take();
        let id = self.next_id;

        match self.attach(id, &mut child) {
            Ok((exit_watch, stdout, stderr)) => {
                self.next_id += 1;
                self.procs.insert(
                    id,
                    Proc {
                        tag,
                        child,
                        exit_watch,
                        stdout,
                        stderr,
                        status: None,
                    },
                );
                Ok((id, stdin))
            }
            Err(e) => {
                // The child is spawned but unwatched, and no caller will ever
                // hear about it: take it down rather than leak it.
                let _ = child.kill();
                let _ = child.wait();
                Err(e)
            }
        }
    }

    /// Take whichever of the child's output streams are pipes and put them,
    /// plus the child's exit, under the set's poller. A stream the command did
    /// not ask to have piped is none of the set's business, and stays `None`.
    /// On failure the descriptors are dropped, and closing them unregisters
    /// whatever was registered already.
    fn attach(
        &mut self,
        id: u64,
        child: &mut Child,
    ) -> io::Result<(sys::ExitWatch, Option<ChildStdout>, Option<ChildStderr>)> {
        let stdout = child.stdout.take();
        let stderr = child.stderr.take();

        if let Some(stdout) = stdout.as_ref() {
            // A pipe that is merely *reported* ready must never block the loop.
            sys::set_nonblocking(stdout.as_fd())?;
            self.poller
                .watch_read(stdout.as_fd(), token(id, STREAM_STDOUT))?;
        }
        if let Some(stderr) = stderr.as_ref() {
            sys::set_nonblocking(stderr.as_fd())?;
            self.poller
                .watch_read(stderr.as_fd(), token(id, STREAM_STDERR))?;
        }
        // Watching the exit before anything is reaped is what makes it safe:
        // the pid still belongs to this child, running or zombie.
        let exit_watch = self
            .poller
            .watch_exit(child.id() as libc::pid_t, token(id, STREAM_EXIT))?;
        Ok((exit_watch, stdout, stderr))
    }

    /// Wait for the next thing any process in the set does.
    ///
    /// Blocks until there is something to report. `None` means the set is
    /// empty: every process has exited and every exit has been handed out.
    pub fn wait_next(&mut self) -> io::Result<Option<(T, PsetEvent)>> {
        self.next_event(None)
    }

    /// [`wait_next`](Self::wait_next), giving up after `timeout`.
    ///
    /// `None` means there was nothing to report in time — or that the set is
    /// empty, which [`is_empty`](Self::is_empty) tells apart.
    pub fn wait_next_timeout(&mut self, timeout: Duration) -> io::Result<Option<(T, PsetEvent)>> {
        self.next_event(Some(Instant::now() + timeout))
    }

    fn next_event(&mut self, deadline: Option<Instant>) -> io::Result<Option<(T, PsetEvent)>> {
        loop {
            if let Some(event) = self.pending.pop_front() {
                return Ok(Some(event));
            }
            if self.procs.is_empty() {
                return Ok(None);
            }

            let timeout = deadline.map(|d| d.saturating_duration_since(Instant::now()));
            let n = match self.poller.wait(&mut self.tokens, timeout) {
                Ok(n) => n,
                // A signal is not an answer: wait again for what is left of
                // the timeout.
                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => return Err(e),
            };
            if n == 0 {
                return Ok(None); // timed out
            }
            self.dispatch(n)?;
        }
    }

    /// Turn the first `n` readiness notifications into pending events.
    fn dispatch(&mut self, n: usize) -> io::Result<()> {
        for i in 0..n {
            let (id, stream) = untoken(self.tokens[i]);
            match stream {
                STREAM_STDOUT | STREAM_STDERR => self.read_stream(id, stream)?,
                STREAM_EXIT => self.reap(id)?,
                _ => unreachable!("token carries one of three streams"),
            }
        }
        self.collect_finished();
        Ok(())
    }

    /// Read one chunk out of a ready pipe, or retire the stream on end of file.
    fn read_stream(&mut self, id: u64, stream: u64) -> io::Result<()> {
        let Some(proc) = self.procs.get_mut(&id) else {
            return Ok(());
        };
        let reader: Option<&mut dyn Read> = match stream {
            STREAM_STDOUT => match proc.stdout.as_mut() {
                Some(stdout) => Some(stdout),
                None => None,
            },
            _ => match proc.stderr.as_mut() {
                Some(stderr) => Some(stderr),
                None => None,
            },
        };
        let Some(reader) = reader else {
            return Ok(());
        };

        let read = match reader.read(&mut self.buf) {
            Ok(read) => read,
            Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()),
            Err(e) if e.kind() == io::ErrorKind::Interrupted => return Ok(()),
            // A pipe whose other end is gone in a way plain EOF cannot express
            // (EIO from a pty, say) has nothing more to give either: treat any
            // hard read error as the end of the stream rather than failing the
            // whole set over one process.
            Err(_) => 0,
        };

        if read == 0 {
            // End of file. Unregister before the descriptor is closed.
            match stream {
                STREAM_STDOUT => {
                    let stdout = proc.stdout.take().expect("checked above");
                    self.poller.unwatch_read(stdout.as_fd())?;
                }
                _ => {
                    let stderr = proc.stderr.take().expect("checked above");
                    self.poller.unwatch_read(stderr.as_fd())?;
                }
            }
            return Ok(());
        }

        let data = self.buf[..read].to_vec();
        let event = match stream {
            STREAM_STDOUT => PsetEvent::Stdout(data),
            _ => PsetEvent::Stderr(data),
        };
        self.pending.push_back((proc.tag.clone(), event));
        Ok(())
    }

    /// The exit watch fired, so the process has exited: collect its status.
    fn reap(&mut self, id: u64) -> io::Result<()> {
        let Some(proc) = self.procs.get_mut(&id) else {
            return Ok(());
        };
        if proc.status.is_some() {
            return Ok(());
        }
        // A fired exit watch must leave the queue before anything else can be
        // waited for: on Linux an exited pidfd stays readable forever.
        self.poller.unwatch_exit(&proc.exit_watch)?;
        // The process is known to be dead, so this does not block.
        proc.status = Some(proc.child.wait()?);
        Ok(())
    }

    /// Report — and remove — every process that is both reaped and drained.
    fn collect_finished(&mut self) {
        let mut done: Vec<u64> = self
            .procs
            .iter()
            .filter(|(_, proc)| proc.finished())
            .map(|(id, _)| *id)
            .collect();
        // Spawn order, so a batch of exits is reported deterministically.
        done.sort_unstable();
        for id in done {
            let proc = self.procs.remove(&id).expect("just listed");
            let status = proc.status.expect("finished implies reaped");
            self.pending
                .push_back((proc.tag, PsetEvent::ProcessExited(status)));
        }
    }
}

impl<T> fmt::Debug for ProcessSet<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ProcessSet")
            .field("processes", &self.procs.len())
            .field("pending_events", &self.pending.len())
            .finish()
    }
}

impl<T> Drop for ProcessSet<T> {
    /// Take every surviving child down with the set, rather than leave it
    /// orphaned or lingering as a zombie.
    fn drop(&mut self) {
        for proc in self.procs.values_mut() {
            if proc.status.is_some() {
                continue;
            }
            let _ = proc.exit_watch.send_signal(libc::SIGKILL);
            let _ = proc.child.wait();
        }
    }
}

/// Converting libc error into [`io::Error`]
macro_rules! cvt {
    ($e:expr) => {{
        let n = $e;
        if n < 0 {
            return Err(io::Error::last_os_error());
        }
        n
    }};
}
pub(crate) use cvt;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn token_round_trips() {
        for id in [0u64, 1, 7, 4096, u64::MAX >> 2] {
            for stream in [STREAM_STDOUT, STREAM_STDERR, STREAM_EXIT] {
                assert_eq!(untoken(token(id, stream)), (id, stream));
            }
        }
    }

    #[test]
    fn tokens_of_one_process_differ() {
        let id = 3;
        let tokens = [
            token(id, STREAM_STDOUT),
            token(id, STREAM_STDERR),
            token(id, STREAM_EXIT),
        ];
        assert_eq!(
            tokens
                .iter()
                .collect::<std::collections::HashSet<_>>()
                .len(),
            3
        );
    }

    #[test]
    fn a_fresh_set_is_empty() {
        let set = create::<()>().unwrap();
        assert!(set.is_empty());
        assert_eq!(set.len(), 0);
    }

    #[test]
    fn waiting_on_an_empty_set_returns_nothing() {
        let mut set = create::<()>().unwrap();
        assert!(set.wait_next().unwrap().is_none());
    }

    #[test]
    fn killing_an_unknown_process_is_not_found() {
        let set = create::<()>().unwrap();
        let err = set.kill(42, libc::SIGTERM).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::NotFound);
    }
}