smarm/src/io.rs

//! Off-scheduler IO: blocking-work offload and epoll-based fd readiness.
//!
//! `block_on_io(closure)` runs `closure` on a dedicated worker OS thread,
//! parks the calling actor in the meantime, and returns the closure's
//! value when it completes. Lets actors call into blocking C libraries,
//! synchronous file IO, or anything else that doesn't fit the readiness
//! model.
//!
//! `wait_readable(fd)` / `wait_writable(fd)` register interest in an fd
//! with epoll and park the calling actor. When the fd becomes ready, the
//! epoll thread unparks the actor. The actual `read(2)`/`write(2)` syscall
//! runs back on the scheduler thread, *inside* the actor — buffer never
//! leaves the actor, no copying through an intermediary thread. Built on
//! these are the conveniences `read(fd, &mut buf)` and `write(fd, &buf)`.
//!
//! Architecture
//! ============
//! Per `run()`, two OS threads:
//!   - **epoll thread**: owns the epollfd. Loops in `epoll_wait`. On a
//!     ready fd, pushes `Completion::FdReady { pid, fd, events }` to the
//!     shared completion queue and writes the scheduler-wake pipe. On the
//!     shutdown pipe (also registered in epollfd), exits.
//!   - **pool thread**: blocks on the request mpsc. Runs the closure
//!     inside `catch_unwind`, pushes `Completion::Blocking { pid, result }`,
//!     writes the scheduler-wake pipe.
//!
//! Both threads share a single `completions: Arc<Mutex<VecDeque<Completion>>>`
//! and the same scheduler-wake pipe.
//!
//! `epoll_ctl` (register/unregister fd interest) is called by the
//! scheduler thread *directly* on the epollfd. That's well-defined per
//! `epoll_ctl(2)`: a thread may be calling `epoll_wait` on the epollfd
//! while another thread calls `epoll_ctl`. Avoids needing a second mpsc
//! and a second wake mechanism.
//!
//! Epoll mode
//! ==========
//! Level-triggered with EPOLLONESHOT. After a wakeup the kernel
//! auto-disarms the fd, so we never get two wakeups for one
//! `wait_readable` call. The scheduler explicitly `EPOLL_CTL_DEL`s the fd
//! on completion to free the slot for re-registration. Net effect: each
//! `wait_readable(fd)` is one ADD, one wakeup, one DEL — symmetric and
//! stateless between calls.
//!
//! Fd hygiene
//! ==========
//! An actor stopped while waiting on an fd unwinds out of `wait_fd`'s park;
//! a drop guard there (armed after a successful register, forgotten on a
//! normal wake) removes the `waiters` entry iff it is still that wait's
//! `(pid, epoch)` and only then `EPOLL_CTL_DEL`s the fd — an entry already
//! consumed by a racing `FdReady` means the fd may carry someone else's
//! fresh registration, which must be left alone. `epoll_register` keeps a
//! defensive bare DEL before ADD as belt-and-braces.
//!
//! Buffers used with `read`/`write` should be on fds opened with
//! `O_NONBLOCK`. If they aren't, the syscall may block the scheduler
//! thread despite the readiness notification (the fd reporting readable
//! doesn't guarantee the syscall completes without blocking — e.g. a
//! signal could be delivered). Documented; not enforced.
//!
//! Panic handling
//! ==============
//! The pool worker runs the closure inside `catch_unwind` and ships either
//! the return value or the panic payload back to the scheduler.
//! `block_on_io` resumes the panic on the calling actor's stack, so the
//! actor's supervisor sees a real `Signal::Panic` as if the work had run
//! inline. Fd-wait primitives don't run user code on the IO thread, so
//! they have no equivalent panic-propagation path.

use crate::pid::Pid;
use std::any::Any;
use std::collections::{HashMap, VecDeque};
use std::io;
use std::os::fd::RawFd;
use std::panic;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle as OsJoinHandle;

// ---------------------------------------------------------------------------
// Wire types
// ---------------------------------------------------------------------------

/// What the pool stores while computing a result. `Ok` is the closure's
/// return value (boxed as `Any`); `Err` is the panic payload.
pub type IoResult = Result<Box<dyn Any + Send>, Box<dyn Any + Send>>;

struct Request {
    /// The submitter's park-epoch — carried through to the `Blocking`
    /// completion so the wake is epoch-matched.
    epoch: u32,
    pid: Pid,
    /// The work to perform. Returns the wire-form result directly.
    work: Box<dyn FnOnce() -> IoResult + Send>,
}

/// Completion message from either IO thread back to the scheduler.
pub enum Completion {
    /// A `block_on_io` closure has finished (Ok = return value, Err = panic
    /// payload).
    Blocking { pid: Pid, epoch: u32, result: IoResult },
    /// An fd registered via `wait_readable`/`wait_writable` is ready. The
    /// scheduler looks up the parked pid in `waiters`, unparks it, and
    /// removes the entry. `pid` isn't in this variant because the epoll
    /// thread doesn't have access to the `waiters` map; the scheduler
    /// thread owns that.
    FdReady { fd: RawFd, events: u32 },
}

// ---------------------------------------------------------------------------
// IoThread — created per `run()`, owned by `SchedulerState`.
// ---------------------------------------------------------------------------

pub struct IoThread {
    // ----- Channels & queues -----

    /// Submission queue into the blocking-work pool.
    tx: mpsc::Sender<Request>,
    /// Shared completion queue, fed by both the pool and the epoll thread.
    completions: Arc<Mutex<VecDeque<Completion>>>,
    /// Pipe the scheduler polls in its idle path. Both IO threads write to
    /// `wake_write` after pushing a completion.
    wake_read: RawFd,
    wake_write: RawFd,

    // ----- Epoll machinery -----

    /// The epollfd, owned by `IoThread`. Callable cross-thread via
    /// `epoll_ctl` per the man page.
    epollfd: RawFd,
    /// Pipe used to signal the epoll thread to exit. Registered inside the
    /// epollfd so a single `epoll_wait` covers both fd readiness and
    /// shutdown.
    shutdown_read: RawFd,
    shutdown_write: RawFd,
    /// One parked actor per registered fd. Populated by `wait_readable` /
    /// `wait_writable` and drained by the scheduler when a `FdReady`
    /// completion is processed.
    pub waiters: HashMap<RawFd, (Pid, u32)>,

    // ----- Threads -----

    pool_thread: Option<OsJoinHandle<()>>,
    epoll_thread: Option<OsJoinHandle<()>>,

    /// Number of `block_on_io` requests in-flight. Used by the scheduler's
    /// idle path to decide whether to wait on the pipe or exit. Fd waits
    /// are not counted here; they're counted by `waiters.len()`.
    pub outstanding: u32,
}

impl IoThread {
    pub fn start() -> io::Result<Self> {
        // Scheduler-facing wake pipe.
        let (wake_read, wake_write) = make_pipe()?;
        // Pool submission channel + shared completion queue.
        let (tx, rx) = mpsc::channel::<Request>();
        let completions: Arc<Mutex<VecDeque<Completion>>> =
            Arc::new(Mutex::new(VecDeque::new()));

        // Epoll machinery.
        let epollfd = unsafe { libc::epoll_create1(libc::EPOLL_CLOEXEC) };
        if epollfd < 0 {
            // Best-effort fd cleanup before bailing.
            unsafe {
                libc::close(wake_read);
                libc::close(wake_write);
            }
            return Err(io::Error::last_os_error());
        }

        let (shutdown_read, shutdown_write) = match make_pipe() {
            Ok(p) => p,
            Err(e) => {
                unsafe {
                    libc::close(epollfd);
                    libc::close(wake_read);
                    libc::close(wake_write);
                }
                return Err(e);
            }
        };

        // Register the shutdown pipe in epollfd. We use a sentinel `data`
        // value to recognise shutdown events. RawFd values are non-negative,
        // so u64::MAX is unambiguously not a real fd-data encoding.
        let mut shutdown_ev = libc::epoll_event {
            events: libc::EPOLLIN as u32,
            u64: SHUTDOWN_EPOLL_TOKEN,
        };
        if unsafe {
            libc::epoll_ctl(
                epollfd,
                libc::EPOLL_CTL_ADD,
                shutdown_read,
                &mut shutdown_ev as *mut _,
            )
        } < 0
        {
            let e = io::Error::last_os_error();
            unsafe {
                libc::close(epollfd);
                libc::close(shutdown_read);
                libc::close(shutdown_write);
                libc::close(wake_read);
                libc::close(wake_write);
            }
            return Err(e);
        }

        // Spawn pool thread.
        let pool_comps = completions.clone();
        let pool_thread = std::thread::Builder::new()
            .name("smarm-io-pool".into())
            .spawn(move || pool_loop(rx, pool_comps, wake_write))?;

        // Spawn epoll thread.
        let epoll_comps = completions.clone();
        let epoll_thread = std::thread::Builder::new()
            .name("smarm-io-epoll".into())
            .spawn(move || epoll_loop(epollfd, epoll_comps, wake_write))?;

        Ok(Self {
            tx,
            completions,
            wake_read,
            wake_write,
            epollfd,
            shutdown_read,
            shutdown_write,
            waiters: HashMap::new(),
            pool_thread: Some(pool_thread),
            epoll_thread: Some(epoll_thread),
            outstanding: 0,
        })
    }

    /// Hand a request to the pool. Increments `outstanding`.
    pub fn submit(&mut self, pid: Pid, epoch: u32, work: Box<dyn FnOnce() -> IoResult + Send>) {
        self.outstanding += 1;
        // Send can only fail if the pool has hung up, which only happens
        // on shutdown. submit during shutdown is a bug.
        self.tx
            .send(Request { pid, epoch, work })
            .expect("io pool hung up unexpectedly");
    }

    /// Drain every available completion. Caller (the scheduler) routes the
    /// results and updates `outstanding` / `waiters` accordingly.
    pub fn drain_completions(&mut self) -> Vec<Completion> {
        let mut q = self.completions.lock().unwrap();
        let mut out = Vec::with_capacity(q.len());
        while let Some(c) = q.pop_front() {
            out.push(c);
        }
        out
    }

    pub fn wake_fd(&self) -> RawFd {
        self.wake_read
    }

    /// Write the wake pipe directly: rouse every scheduler thread blocked in
    /// its idle `poll_wake`. Used by the terminal (AllDone) path — an idle
    /// sibling may be blocked on a snapshot that nothing will ever refresh
    /// (an orphaned timer deadline, or `io_outstanding` from a waiter that
    /// was stop-cancelled and so never produces a completion).
    pub fn wake(&self) {
        wake_scheduler(self.wake_write);
    }

    /// Register interest in `fd` becoming readable/writable; record `pid`
    /// as the parked waiter. The epoll thread will push a `FdReady`
    /// completion when the kernel signals.
    ///
    /// EPOLLONESHOT: one wakeup per registration. The scheduler must
    /// `epoll_del` on completion to free the slot for re-registration.
    pub fn epoll_register(
        &mut self,
        fd: RawFd,
        pid: Pid,
        epoch: u32,
        readable: bool,
        writable: bool,
    ) -> io::Result<()> {
        // Two actors waiting on the same fd would be a misuse: the kernel
        // delivers exactly one EPOLLONESHOT wakeup, so the second waiter
        // would hang. Reject up front.
        if self.waiters.contains_key(&fd) {
            return Err(io::Error::new(
                io::ErrorKind::AlreadyExists,
                "fd already has a parked waiter",
            ));
        }

        // Belt-and-braces: the unwind guard in `wait_fd` is responsible for
        // cleaning up a stopped waiter's registration, but a bare DEL is
        // harmless if the fd isn't registered (ENOENT) and removes any leak
        // a path we haven't thought of might leave behind.
        unsafe {
            libc::epoll_ctl(self.epollfd, libc::EPOLL_CTL_DEL, fd, std::ptr::null_mut());
        }

        let mut events: u32 = libc::EPOLLONESHOT as u32;
        if readable {
            events |= libc::EPOLLIN as u32;
        }
        if writable {
            events |= libc::EPOLLOUT as u32;
        }
        let mut ev = libc::epoll_event {
            events,
            u64: fd as u64,
        };
        let r = unsafe {
            libc::epoll_ctl(self.epollfd, libc::EPOLL_CTL_ADD, fd, &mut ev as *mut _)
        };
        if r < 0 {
            return Err(io::Error::last_os_error());
        }
        self.waiters.insert(fd, (pid, epoch));
        Ok(())
    }

    /// Remove `fd` from the epollfd. Called by the scheduler after a
    /// `FdReady` completion, so the next `wait_readable(fd)` can ADD again.
    ///
    /// Does NOT touch `waiters` — that's the scheduler's bookkeeping; this
    /// is purely the kernel-side cleanup.
    pub fn epoll_deregister(&mut self, fd: RawFd) {
        // EPOLL_CTL_DEL of an already-removed fd returns ENOENT; ignore.
        unsafe {
            libc::epoll_ctl(self.epollfd, libc::EPOLL_CTL_DEL, fd, std::ptr::null_mut());
        }
    }
}

impl Drop for IoThread {
    fn drop(&mut self) {
        // 1. Signal the epoll thread to exit by writing the shutdown pipe.
        unsafe {
            let buf: [u8; 1] = [0];
            // Single byte; we don't care about EINTR retry here — worst
            // case the epoll thread blocks until process exit, which is
            // fine because we then close fds out from under it.
            libc::write(self.shutdown_write, buf.as_ptr() as *const _, 1);
        }

        // 2. Hang up the pool's request channel so the pool thread exits.
        let (dead_tx, _) = mpsc::channel::<Request>();
        let real_tx = std::mem::replace(&mut self.tx, dead_tx);
        drop(real_tx);

        // 3. Join both threads.
        if let Some(h) = self.epoll_thread.take() {
            let _ = h.join();
        }
        if let Some(h) = self.pool_thread.take() {
            let _ = h.join();
        }

        // 4. Close fds.
        unsafe {
            libc::close(self.epollfd);
            libc::close(self.shutdown_read);
            libc::close(self.shutdown_write);
            libc::close(self.wake_read);
            libc::close(self.wake_write);
        }
    }
}

/// Sentinel `epoll_event.u64` distinguishing the shutdown pipe from
/// registered actor fds. RawFd values fit in i32, so the high bits are
/// available for a marker; we use u64::MAX which can't be a valid fd.
const SHUTDOWN_EPOLL_TOKEN: u64 = u64::MAX;

// ---------------------------------------------------------------------------
// Pool loop
// ---------------------------------------------------------------------------

fn pool_loop(
    rx: mpsc::Receiver<Request>,
    completions: Arc<Mutex<VecDeque<Completion>>>,
    wake_write: RawFd,
) {
    while let Ok(Request { pid, epoch, work }) = rx.recv() {
        let result: IoResult = match panic::catch_unwind(panic::AssertUnwindSafe(work)) {
            Ok(r) => r,
            Err(payload) => Err(payload),
        };
        completions
            .lock()
            .unwrap()
            .push_back(Completion::Blocking { pid, epoch, result });
        wake_scheduler(wake_write);
    }
}

// ---------------------------------------------------------------------------
// Epoll loop
// ---------------------------------------------------------------------------

fn epoll_loop(
    epollfd: RawFd,
    completions: Arc<Mutex<VecDeque<Completion>>>,
    wake_write: RawFd,
) {
    // Buffer for epoll_wait. 64 is plenty for our scale; if a real load
    // appears that needs more, this is a one-line change.
    const MAX_EVENTS: usize = 64;
    let mut events: [libc::epoll_event; MAX_EVENTS] = unsafe { std::mem::zeroed() };

    loop {
        let n = unsafe {
            libc::epoll_wait(
                epollfd,
                events.as_mut_ptr(),
                MAX_EVENTS as libc::c_int,
                -1,
            )
        };

        if n < 0 {
            let e = unsafe { *libc::__errno_location() };
            if e == libc::EINTR {
                continue;
            }
            // Anything else here is a programming error (EBADF on epollfd
            // after we've closed it from Drop — the close races with us).
            // Treat as shutdown.
            return;
        }

        let mut shutdown_requested = false;
        let mut pushed_any = false;
        {
            let mut q = completions.lock().unwrap();
            for ev in events.iter().take(n as usize) {
                if ev.u64 == SHUTDOWN_EPOLL_TOKEN {
                    shutdown_requested = true;
                    continue;
                }
                let fd = ev.u64 as RawFd;
                let evs = ev.events;
                q.push_back(Completion::FdReady {
                    fd,
                    events: evs,
                });
                pushed_any = true;
            }
        }

        if pushed_any {
            wake_scheduler(wake_write);
        }
        if shutdown_requested {
            return;
        }
    }
}

/// Write one byte to the scheduler's wake pipe. Retries on EINTR; ignores
/// EAGAIN (pipe full means there's already an outstanding wake we haven't
/// consumed yet, which is sufficient).
fn wake_scheduler(wake_write: RawFd) {
    let buf: [u8; 1] = [0];
    unsafe {
        loop {
            let n = libc::write(wake_write, buf.as_ptr() as *const _, 1);
            if n < 0 {
                let e = *libc::__errno_location();
                if e == libc::EINTR {
                    continue;
                }
            }
            break;
        }
    }
}

// ---------------------------------------------------------------------------
// Pipe helpers (unchanged from v0.2)
// ---------------------------------------------------------------------------

fn make_pipe() -> io::Result<(RawFd, RawFd)> {
    let mut fds: [libc::c_int; 2] = [0; 2];
    let r = unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC | libc::O_NONBLOCK) };
    if r != 0 {
        return Err(io::Error::last_os_error());
    }
    Ok((fds[0], fds[1]))
}

/// Drain pending bytes from the wake pipe. The scheduler calls this after
/// a `poll` wakeup so the next idle call sees an empty pipe.
pub fn drain_wake_pipe(fd: RawFd) {
    let mut buf = [0u8; 64];
    loop {
        let n = unsafe { libc::read(fd, buf.as_mut_ptr() as *mut _, buf.len()) };
        if n <= 0 {
            break;
        }
    }
}

/// Block on `fd` for up to `timeout`, returning when either there's data
/// to read or the timeout elapses. `None` for `timeout` means wait forever.
pub fn poll_wake(fd: RawFd, timeout: Option<std::time::Duration>) {
    let timeout_ms: libc::c_int = match timeout {
        None => -1,
        Some(d) => {
            let ms = d.as_millis();
            if ms > i32::MAX as u128 {
                i32::MAX
            } else {
                ms as i32
            }
        }
    };
    let mut pfd = libc::pollfd {
        fd,
        events: libc::POLLIN,
        revents: 0,
    };
    loop {
        let r = unsafe { libc::poll(&mut pfd as *mut _, 1, timeout_ms) };
        if r < 0 {
            let e = unsafe { *libc::__errno_location() };
            if e == libc::EINTR {
                continue;
            }
        }
        break;
    }
}