pset 0.1.0

Orchestrate a set of child processes: one event stream for their output and their exits
//! macOS: `kqueue(2)` waits on the pipes with `EVFILT_READ` and on the exits
//! with `EVFILT_PROC`/`NOTE_EXIT` — no pidfd, but the same one-queue shape.
use crate::cvt;
use std::{
    io,
    os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd},
    ptr,
    time::Duration,
};

/// Handle to one watched child: its pid, plus which filter the exit watch
/// went in under.
///
/// macOS has no pidfd, but signalling by pid is just as safe here: the set
/// only signals children it has not reaped yet, and an unreaped child —
/// running or zombie — still owns its pid, so the pid cannot have been
/// recycled behind our back.
#[derive(Debug)]
pub struct ExitWatch {
    pid: libc::pid_t,
    /// `EVFILT_PROC` normally; `EVFILT_USER` when the process was already
    /// gone as the watch was made and the exit event had to be synthesized.
    filter: i16,
}

impl ExitWatch {
    /// `kill(2)`: signal the process. See the type docs for why the pid can
    /// still be trusted.
    pub fn send_signal(&self, signal: i32) -> io::Result<()> {
        // SAFETY: takes a pid and a signal number, no memory is shared.
        cvt!(unsafe { libc::kill(self.pid, signal) });
        Ok(())
    }
}

/// A `kqueue` instance watching descriptors for readability and processes
/// for exiting.
#[derive(Debug)]
pub struct Poller {
    kq: OwnedFd,
    /// Scratch for `kevent`, so `wait` does not allocate.
    events: Vec<libc::kevent>,
}

fn kev(ident: usize, filter: i16, flags: u16, fflags: u32, token: u64) -> libc::kevent {
    libc::kevent {
        ident,
        filter,
        flags,
        fflags,
        data: 0,
        udata: token as *mut libc::c_void,
    }
}

impl Poller {
    pub fn new() -> io::Result<Self> {
        // SAFETY: no arguments.
        let fd = cvt!(unsafe { libc::kqueue() });
        let kq = unsafe { OwnedFd::from_raw_fd(fd) };
        // `kqueue` has no CLOEXEC flag of its own.
        // SAFETY: plain fcntl on a live descriptor.
        cvt!(unsafe { libc::fcntl(kq.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) });
        Ok(Self {
            kq,
            events: vec![kev(0, 0, 0, 0, 0); 64],
        })
    }

    /// Watch `fd` for readability, tagging its readiness with `token`.
    ///
    /// Registrations are level triggered (no `EV_CLEAR`): a descriptor stays
    /// ready until it is drained or removed.
    pub fn watch_read(&self, fd: BorrowedFd<'_>, token: u64) -> io::Result<()> {
        self.change(kev(
            fd.as_raw_fd() as usize,
            libc::EVFILT_READ,
            libc::EV_ADD,
            0,
            token,
        ))
    }

    /// Stop watching `fd`. Must happen before the descriptor is closed, or
    /// the registration goes away silently and this call reports `ENOENT`.
    pub fn unwatch_read(&self, fd: BorrowedFd<'_>) -> io::Result<()> {
        self.change(kev(
            fd.as_raw_fd() as usize,
            libc::EVFILT_READ,
            libc::EV_DELETE,
            0,
            0,
        ))
    }

    /// Watch the process for exiting: its `token` comes out of `wait` at
    /// least once after the exit.
    ///
    /// Attaching is race free as long as the process has not been reaped —
    /// its pid cannot have been recycled. It may well have *exited* already,
    /// though, and `EVFILT_PROC` refuses to attach to a zombie with `ESRCH`;
    /// the exit still has to be reported, so in that case a pre-triggered
    /// user event goes on the queue under the same token instead.
    pub fn watch_exit(&self, pid: libc::pid_t, token: u64) -> io::Result<ExitWatch> {
        let proc_ev = kev(
            pid as usize,
            libc::EVFILT_PROC,
            libc::EV_ADD,
            libc::NOTE_EXIT,
            token,
        );
        match self.change(proc_ev) {
            Ok(()) => Ok(ExitWatch {
                pid,
                filter: libc::EVFILT_PROC,
            }),
            Err(e) if e.raw_os_error() == Some(libc::ESRCH) => {
                self.change(kev(
                    pid as usize,
                    libc::EVFILT_USER,
                    libc::EV_ADD | libc::EV_ONESHOT,
                    0,
                    token,
                ))?;
                self.change(kev(
                    pid as usize,
                    libc::EVFILT_USER,
                    0,
                    libc::NOTE_TRIGGER,
                    token,
                ))?;
                Ok(ExitWatch {
                    pid,
                    filter: libc::EVFILT_USER,
                })
            }
            Err(e) => Err(e),
        }
    }

    /// Remove an exit watch. The kernel drops some of these on its own — a
    /// oneshot user event once delivered, a proc watch when the process is
    /// reaped — so "already gone" is success here, not an error.
    pub fn unwatch_exit(&self, watch: &ExitWatch) -> io::Result<()> {
        let ev = kev(watch.pid as usize, watch.filter, libc::EV_DELETE, 0, 0);
        match self.change(ev) {
            Err(e) if matches!(e.raw_os_error(), Some(libc::ENOENT) | Some(libc::ESRCH)) => Ok(()),
            result => result,
        }
    }

    /// Submit one change to the queue.
    fn change(&self, ev: libc::kevent) -> io::Result<()> {
        // SAFETY: the change list is one valid entry long and no events are
        // asked back, so nothing is written.
        cvt!(unsafe { libc::kevent(self.kq.as_raw_fd(), &ev, 1, ptr::null_mut(), 0, ptr::null()) });
        Ok(())
    }

    /// Block until at least one watched thing is ready, filling `tokens`
    /// with their tokens and returning how many there are. `None` waits
    /// forever; a timeout that expires first yields `0`.
    pub fn wait(&mut self, tokens: &mut Vec<u64>, timeout: Option<Duration>) -> io::Result<usize> {
        let ts;
        let ts_ptr = match timeout {
            None => ptr::null(),
            Some(d) => {
                ts = libc::timespec {
                    tv_sec: d.as_secs().min(libc::time_t::MAX as u64) as libc::time_t,
                    tv_nsec: d.subsec_nanos() as libc::c_long,
                };
                &ts
            }
        };
        // SAFETY: the event list is a valid slice of `kevent`s and its length
        // is what the kernel is allowed to write; the timeout is null or
        // points at a valid `timespec`.
        let n = cvt!(unsafe {
            libc::kevent(
                self.kq.as_raw_fd(),
                ptr::null(),
                0,
                self.events.as_mut_ptr(),
                self.events.len() as libc::c_int,
                ts_ptr,
            )
        });
        tokens.clear();
        tokens.extend(self.events[..n as usize].iter().map(|ev| ev.udata as u64));
        Ok(n as usize)
    }
}