pset 0.1.0

//! Linux: `pidfd_open(2)` makes a process exit a readable descriptor, and
//! `epoll(7)` waits on it alongside the pipes.

use std::{
    io,
    os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd},
    ptr,
    time::Duration,
};

/// Handle to one watched child: a pidfd.
///
/// A pidfd refers to the process itself, not its pid, so signalling through
/// it cannot hit an unrelated process that inherited the pid in the meantime.
#[derive(Debug)]
pub struct ExitWatch {
    pidfd: OwnedFd,
}

impl ExitWatch {
    /// `pidfd_send_signal(2)`: signal the process the descriptor refers to.
    pub fn send_signal(&self, signal: i32) -> io::Result<()> {
        // SAFETY: a null `siginfo_t` asks the kernel to synthesize one, which
        // is the documented way to send a plain signal.
        cvt!(unsafe {
            libc::syscall(
                libc::SYS_pidfd_send_signal,
                self.pidfd.as_raw_fd(),
                signal,
                ptr::null_mut::<libc::siginfo_t>(),
                0,
            )
        });
        Ok(())
    }
}

/// An `epoll` instance watching descriptors for readability.
#[derive(Debug)]
pub struct Poller {
    epoll: OwnedFd,
    /// Scratch for `epoll_wait`, so `wait` does not allocate.
    events: Vec<libc::epoll_event>,
}

impl Poller {
    pub fn new() -> io::Result<Self> {
        // SAFETY: no arguments beyond a flag word.
        let fd = cvt!(unsafe { libc::epoll_create1(libc::EPOLL_CLOEXEC) });
        Ok(Self {
            epoll: unsafe { OwnedFd::from_raw_fd(fd) },
            events: vec![libc::epoll_event { events: 0, u64: 0 }; 64],
        })
    }

    /// Watch `fd` for readability, tagging its readiness with `token`.
    ///
    /// Registrations are level triggered: a descriptor stays ready until it
    /// is drained or removed.
    pub fn watch_read(&self, fd: BorrowedFd<'_>, token: u64) -> io::Result<()> {
        self.add(fd, 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.delete(fd)
    }

    /// `pidfd_open(2)` the process and watch the descriptor: it becomes
    /// readable once the process exits.
    ///
    /// Opening the pidfd is race free as long as the process has not been
    /// reaped yet — an unreaped zombie still owns its pid, so the pid cannot
    /// have been recycled behind our back.
    pub fn watch_exit(&self, pid: libc::pid_t, token: u64) -> io::Result<ExitWatch> {
        // SAFETY: the syscall takes a pid and a flag word, no memory is
        // shared with the kernel. A non-negative return is a fresh, owned
        // descriptor.
        let fd = cvt!(unsafe { libc::syscall(libc::SYS_pidfd_open, pid, 0) });
        let pidfd = unsafe { OwnedFd::from_raw_fd(fd as RawFd) };
        self.add(pidfd.as_fd(), token)?;
        Ok(ExitWatch { pidfd })
    }

    /// Remove an exit watch. An exited pidfd stays readable forever, so this
    /// must happen once the exit has been seen.
    pub fn unwatch_exit(&self, watch: &ExitWatch) -> io::Result<()> {
        self.delete(watch.pidfd.as_fd())
    }

    fn add(&self, fd: BorrowedFd<'_>, token: u64) -> io::Result<()> {
        let mut event = libc::epoll_event {
            events: libc::EPOLLIN as u32,
            u64: token,
        };
        self.ctl(libc::EPOLL_CTL_ADD, fd, &mut event)
    }

    fn delete(&self, fd: BorrowedFd<'_>) -> io::Result<()> {
        let mut event = libc::epoll_event { events: 0, u64: 0 };
        self.ctl(libc::EPOLL_CTL_DEL, fd, &mut event)
    }

    fn ctl(&self, op: i32, fd: BorrowedFd<'_>, event: &mut libc::epoll_event) -> io::Result<()> {
        // SAFETY: both descriptors are live for the call and `event` points
        // at a valid, initialized `epoll_event`.
        cvt!(unsafe { libc::epoll_ctl(self.epoll.as_raw_fd(), op, fd.as_raw_fd(), event) });
        Ok(())
    }

    /// Block until at least one watched descriptor 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 millis = match timeout {
            None => -1,
            // A zero timeout is a poll: come straight back.
            Some(d) if d.is_zero() => 0,
            // Round up, so a wait never comes back before the deadline. A
            // truncating conversion would return up to a millisecond early,
            // and would turn a sub-millisecond wait into a busy poll.
            Some(d) => {
                let millis = d.as_nanos().div_ceil(1_000_000);
                millis.clamp(1, i32::MAX as u128) as i32
            }
        };
        // SAFETY: the event list is a valid slice of `epoll_event`s and its
        // length is what the kernel is allowed to write.
        let n = cvt!(unsafe {
            libc::epoll_wait(
                self.epoll.as_raw_fd(),
                self.events.as_mut_ptr(),
                self.events.len() as i32,
                millis,
            )
        });
        tokens.clear();
        tokens.extend(self.events[..n as usize].iter().map(|ev| ev.u64));
        Ok(n as usize)
    }
}