coreshift_core/io/

drain.rs

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/

//! High-level process I/O management.
//!
//! This module provides the [`DrainState`] structure, which coordinates the
//! simultaneous reading from process output pipes and writing to process
//! input pipes.
//!
//! This is an advanced helper for callers that already own child-process file
//! descriptors and want non-blocking drain semantics without reimplementing
//! the bookkeeping.

use crate::CoreError;
use crate::fd::{Fd, Token};
use crate::io::buffer::{BufferState, ChunkSink, ReadState};
use crate::io::writer::WriterState;

/// Time bound (ms) for a [`DrainState::write_input`] `poll(POLLOUT)` wait. The
/// pty master is `O_NONBLOCK`, so a write returns `EAGAIN` once the child's tty
/// input buffer is full; the write then waits for writability. Bounded so a
/// wedged child (never draining its stdin) cannot stall the caller forever.
const WRITE_INPUT_POLL_TIMEOUT_MS: i32 = 2_000;

#[inline(always)]
fn errno() -> i32 {
    std::io::Error::last_os_error().raw_os_error().unwrap_or(0)
}

/// Associates a file descriptor with an optional reactor token and the current
/// per-direction interest registered with the reactor.
pub(crate) struct FdSlot {
    /// Token assigned by the reactor for this descriptor. `None` while the fd
    /// is not registered (never added, or removed wholesale on teardown).
    pub token: Option<Token>,
    /// The managed file descriptor.
    pub fd: Fd,
    /// Current readable interest (`EPOLLIN`) registered with the reactor.
    pub readable: bool,
    /// Current writable interest (`EPOLLOUT`) registered with the reactor.
    pub writable: bool,
}

/// Orchestrates non-blocking process I/O.
///
/// `DrainState` tracks the state of stdin, stdout, and stderr pipes for a
/// single process. It handles the multiplexing of data between these pipes
/// and internal buffers.
///
/// # Example
/// ```no_run
/// # use coreshift_core::io::DrainState;
/// # use coreshift_core::reactor::Reactor;
/// # fn example(mut drain: DrainState<fn(&[u8]) -> bool>, mut reactor: Reactor) -> Result<(), Box<dyn std::error::Error>> {
/// while !drain.is_done() {
///     let mut events = Vec::new();
///     reactor.wait(&mut events, 64, -1)?;
///     for ev in events {
///         // Map event tokens to drain calls...
///     }
/// }
/// # Ok(())
/// # }
/// ```
#[repr(align(64))]
pub struct DrainState<F>
where
    F: FnMut(&[u8]) -> bool,
{
    pub(crate) stdout_slot: Option<FdSlot>,
    pub(crate) stderr_slot: Option<FdSlot>,
    pub(crate) stdin_slot: Option<FdSlot>,

    pub(crate) buffer: BufferState,
    pub(crate) writer: WriterState,

    pub(crate) early_exit: Option<F>,

    /// `true` when the stdout slot is a pty master rather than a pipe. A pty
    /// master reports EOF as `EIO` (returned once the session leader and all
    /// slave holders have closed), so the read path maps a stdout `EIO` to a
    /// clean [`ReadState::Eof`] instead of surfacing it as an I/O error.
    pub(crate) pty_master: bool,
}

impl<F> DrainState<F>
where
    F: FnMut(&[u8]) -> bool,
{
    /// Initialize a new drain state for the provided descriptors.
    ///
    /// This consumes the descriptors and sets them to non-blocking mode.
    /// `chunk_sink` enables streaming mode: every retained output chunk is
    /// forwarded to the sink instead of being accumulated into the internal
    /// buffers.
    ///
    /// ### Errors
    /// - `EBADF`: One of the provided file descriptors is invalid.
    #[allow(clippy::too_many_arguments)] // published constructor; grouped params would break the API
    pub fn new(
        stdin_fd: Option<Fd>,
        stdin_buf: Option<Box<[u8]>>,
        stdout_fd: Option<Fd>,
        stderr_fd: Option<Fd>,
        limit: usize,
        early_exit: Option<F>,
        chunk_sink: Option<ChunkSink>,
        pty_master: bool,
    ) -> Result<Self, CoreError> {
        let stdin_slot = if stdin_buf.is_some() {
            if let Some(fd) = stdin_fd {
                fd.set_nonblock()?;
                Some(FdSlot {
                    token: None,
                    fd,
                    readable: false,
                    writable: false,
                })
            } else {
                None
            }
        } else {
            None
        };

        let stdout_slot = if let Some(fd) = stdout_fd {
            fd.set_nonblock()?;
            Some(FdSlot {
                token: None,
                fd,
                readable: false,
                writable: false,
            })
        } else {
            None
        };

        let stderr_slot = if let Some(fd) = stderr_fd {
            fd.set_nonblock()?;
            Some(FdSlot {
                token: None,
                fd,
                readable: false,
                writable: false,
            })
        } else {
            None
        };

        Ok(Self {
            stdin_slot,
            stdout_slot,
            stderr_slot,
            buffer: BufferState::new(limit, chunk_sink),
            writer: WriterState::new(stdin_buf),
            early_exit,
            pty_master,
        })
    }

    /// Returns `true` if all pipes have been closed or fully drained.
    #[inline(always)]
    pub fn is_done(&self) -> bool {
        self.stdin_slot.is_none() && self.stdout_slot.is_none() && self.stderr_slot.is_none()
    }

    /// Apply a new window size (`TIOCSWINSZ`) to the pty master's terminal.
    ///
    /// Sending a signal to the foreground process group after a resize is the
    /// caller's job (SIGWINCH); this only updates the kernel's `winsize` so a
    /// subsequent `TIOCGWINSZ`/`SIGWINCH`-driven refresh reads the new size.
    ///
    /// ### Errors
    /// - `EINVAL`: No pty master is present (non-pty spawn or stream already
    ///   closed), or `rows`/`cols` is zero.
    /// - `ENOTTY`: The stdout descriptor is not a terminal.
    pub(crate) fn resize_pty(&self, rows: u16, cols: u16) -> Result<(), CoreError> {
        if rows == 0 || cols == 0 {
            return Err(CoreError::sys(
                libc::EINVAL,
                "resize_pty: rows and cols must be non-zero",
            ));
        }
        let Some(slot) = &self.stdout_slot else {
            return Err(CoreError::sys(libc::EINVAL, "resize_pty: no pty master"));
        };
        let ws = libc::winsize {
            ws_row: rows,
            ws_col: cols,
            ws_xpixel: 0,
            ws_ypixel: 0,
        };
        let r = unsafe { libc::ioctl(slot.fd.raw(), libc::TIOCSWINSZ as libc::Ioctl, &ws) };
        crate::error::syscall_ret(r, "TIOCSWINSZ")
    }

    /// Write bytes to the pty master — the child's stdin on a pty spawn.
    ///
    /// The master is `O_NONBLOCK`, so a full write waits for `POLLOUT` (bounded
    /// by [`WRITE_INPUT_POLL_TIMEOUT_MS`]) when the tty input buffer is full,
    /// and returns once every byte has been accepted by the line discipline.
    ///
    /// ### Errors
    /// - `EINVAL`: Not a pty spawn, or the pty master is already closed.
    /// - `EIO`: All slave holders have closed (master-side write failure).
    /// - `ETIMEDOUT`: The child did not drain its input within the poll bound.
    pub(crate) fn write_input(&self, bytes: &[u8]) -> Result<usize, CoreError> {
        if !self.pty_master {
            return Err(CoreError::sys(libc::EINVAL, "write_input: not a pty spawn"));
        }
        let Some(slot) = &self.stdout_slot else {
            return Err(CoreError::sys(
                libc::EINVAL,
                "write_input: pty master closed",
            ));
        };
        let fd = slot.fd.raw();
        let mut written = 0usize;
        while written < bytes.len() {
            let n = unsafe {
                libc::write(
                    fd,
                    bytes[written..].as_ptr() as *const libc::c_void,
                    bytes.len() - written,
                )
            };
            if n < 0 {
                let e = errno();
                if e == libc::EINTR {
                    continue;
                }
                if e == libc::EAGAIN {
                    let mut pfd = libc::pollfd {
                        fd,
                        events: libc::POLLOUT,
                        revents: 0,
                    };
                    let rc = unsafe { libc::poll(&mut pfd, 1, WRITE_INPUT_POLL_TIMEOUT_MS) };
                    if rc < 0 {
                        let pe = errno();
                        if pe == libc::EINTR {
                            continue;
                        }
                        return Err(CoreError::sys(pe, "write_input:poll"));
                    }
                    if rc == 0 {
                        return Err(CoreError::sys(
                            libc::ETIMEDOUT,
                            "write_input: tty input buffer stayed full",
                        ));
                    }
                    continue;
                }
                return Err(CoreError::sys(e, "write_input"));
            }
            written += n as usize;
        }
        Ok(written)
    }

    /// Write bytes to the pty master without blocking.
    ///
    /// Returns `Ok(Some(n))` with the number of bytes written (which may be
    /// less than `bytes.len()` when the tty input buffer fills mid-write), or
    /// `Ok(None)` when `EAGAIN` on the first byte (buffer full). Never blocks
    /// — the caller owns the input queue and re-arms `POLLOUT` interest on
    /// `EAGAIN` (register writable), writing again on readiness.
    ///
    /// ### Errors
    /// - `EINVAL`: Not a pty spawn, or the pty master is already closed.
    /// - `EIO`: All slave holders have closed (master-side write failure).
    pub(crate) fn write_input_nonblock(&self, bytes: &[u8]) -> Result<Option<usize>, CoreError> {
        if !self.pty_master {
            return Err(CoreError::sys(
                libc::EINVAL,
                "write_input_nonblock: not a pty spawn",
            ));
        }
        let Some(slot) = &self.stdout_slot else {
            return Err(CoreError::sys(
                libc::EINVAL,
                "write_input_nonblock: pty master closed",
            ));
        };
        let fd = slot.fd.raw();
        let mut written = 0usize;
        while written < bytes.len() {
            let n = unsafe {
                libc::write(
                    fd,
                    bytes[written..].as_ptr() as *const libc::c_void,
                    bytes.len() - written,
                )
            };
            if n < 0 {
                let e = errno();
                if e == libc::EINTR {
                    continue;
                }
                if e == libc::EAGAIN {
                    return Ok(if written == 0 { None } else { Some(written) });
                }
                return Err(CoreError::sys(e, "write_input_nonblock"));
            }
            written += n as usize;
        }
        Ok(Some(written))
    }

    /// Perform a non-blocking write to stdin if pending.
    ///
    /// Returns `Ok(true)` if the write buffer is empty or the descriptor is
    /// closed.
    ///
    /// ### Errors
    /// - `EPIPE`: The child process closed its reading end of the pipe.
    /// - `EIO`: Low-level I/O error.
    #[inline(always)]
    pub fn write_stdin(&mut self) -> Result<bool, CoreError> {
        let fd = if let Some(s) = &self.stdin_slot {
            &s.fd
        } else {
            return Ok(true);
        };

        let done = self.writer.write_to_fd(fd)?;
        if done {
            self.stdin_slot.take();
            return Ok(true);
        }
        Ok(false)
    }

    /// Read from a slot's descriptor, mapping a pty-master `EIO` EOF to a
    /// clean [`ReadState::Eof`]. All other errors propagate.
    ///
    /// Associated fn (no `self` receiver) so callers can borrow `buffer` and
    /// `early_exit` mutably while a `stdout_slot`/`stderr_slot` borrow is
    /// still live — disjoint field borrows the compiler can see.
    #[inline(always)]
    fn read_from_slot(
        buffer: &mut BufferState,
        fd: &Fd,
        is_stdout: bool,
        early_exit: &mut Option<F>,
        pty_master: bool,
    ) -> Result<ReadState, CoreError> {
        match buffer.read_from_fd(fd, is_stdout, early_exit) {
            Err(e) if pty_master && is_stdout && e.raw_os_error() == Some(libc::EIO) => {
                Ok(ReadState::Eof)
            }
            other => other,
        }
    }

    /// Perform a non-blocking read from stdout or stderr.
    ///
    /// Returns `Ok(true)` if the stream reached EOF, the early-exit condition
    /// was met, or the stream is paused on a full sink queue (the caller
    /// resumes it later). In the paused case the slot is retained.
    ///
    /// ### Errors
    /// - `EOVERFLOW`: The captured output exceeded the specified limit.
    /// - `EIO`: Low-level I/O error.
    #[inline(always)]
    pub fn read_fd(&mut self, is_stdout: bool) -> Result<bool, CoreError> {
        let pty_master = self.pty_master;
        let read_state = {
            let slot = if is_stdout {
                &self.stdout_slot
            } else {
                &self.stderr_slot
            };
            let fd = if let Some(s) = slot {
                &s.fd
            } else {
                return Ok(true);
            };
            Self::read_from_slot(
                &mut self.buffer,
                fd,
                is_stdout,
                &mut self.early_exit,
                pty_master,
            )?
        };

        match read_state {
            ReadState::Open => Ok(false),
            ReadState::Paused => Ok(false),
            ReadState::Eof | ReadState::EarlyExit => {
                if is_stdout {
                    self.stdout_slot.take();
                } else {
                    self.stderr_slot.take();
                }
                Ok(true)
            }
        }
    }

    /// Extract all active slots for cleanup or reactor removal.
    pub(crate) fn take_all_slots(&mut self) -> Vec<FdSlot> {
        let mut slots = Vec::new();
        if let Some(slot) = self.stdin_slot.take() {
            slots.push(slot);
        }
        if let Some(slot) = self.stdout_slot.take() {
            slots.push(slot);
        }
        if let Some(slot) = self.stderr_slot.take() {
            slots.push(slot);
        }
        slots
    }

    pub(crate) fn register_with_reactor(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        register_slot(reactor, &mut self.stdin_slot, false, true)?;
        register_slot(reactor, &mut self.stdout_slot, true, false)?;
        register_slot(reactor, &mut self.stderr_slot, true, false)?;
        Ok(())
    }

    pub(crate) fn stdout_matches(&self, token: Token) -> bool {
        self.stdout_slot
            .as_ref()
            .is_some_and(|slot| slot.token == Some(token))
    }

    pub(crate) fn stderr_matches(&self, token: Token) -> bool {
        self.stderr_slot
            .as_ref()
            .is_some_and(|slot| slot.token == Some(token))
    }

    pub(crate) fn stdin_matches(&self, token: Token) -> bool {
        self.stdin_slot
            .as_ref()
            .is_some_and(|slot| slot.token == Some(token))
    }

    pub(crate) fn drop_stdout(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = self.stdout_slot.take() {
            del_slot(reactor, &slot)?;
        }
        Ok(())
    }

    pub(crate) fn drop_stderr(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = self.stderr_slot.take() {
            del_slot(reactor, &slot)?;
        }
        Ok(())
    }

    pub(crate) fn drop_stdin(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = self.stdin_slot.take() {
            del_slot(reactor, &slot)?;
        }
        self.writer.buf = None;
        Ok(())
    }

    pub(crate) fn handle_stdout_ready(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = &self.stdout_slot {
            let read_state = Self::read_from_slot(
                &mut self.buffer,
                &slot.fd,
                true,
                &mut self.early_exit,
                self.pty_master,
            )?;
            match read_state {
                ReadState::Open => {}
                ReadState::Paused => {
                    // Sink queue full: remove the fd from the reactor so the
                    // edge-triggered readiness does not spin the loop; the
                    // caller re-registers via `resume_stdout` when it has
                    // drained the queue.
                    self.pause_stdout(reactor)?;
                }
                ReadState::Eof | ReadState::EarlyExit => {
                    self.drop_stdout(reactor)?;
                }
            }
        }
        Ok(())
    }

    pub(crate) fn handle_stderr_ready(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = &self.stderr_slot {
            let read_state = Self::read_from_slot(
                &mut self.buffer,
                &slot.fd,
                false,
                &mut self.early_exit,
                self.pty_master,
            )?;
            match read_state {
                ReadState::Open => {}
                ReadState::Paused => {
                    self.pause_stderr(reactor)?;
                }
                ReadState::Eof | ReadState::EarlyExit => {
                    self.drop_stderr(reactor)?;
                }
            }
        }
        Ok(())
    }

    /// Remove the stdout fd from the reactor while its sink queue is full.
    /// The slot is retained (token cleared) so the stream can be resumed.
    pub(crate) fn pause_stdout(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        pause_slot(reactor, &mut self.stdout_slot)
    }

    /// Remove the stderr fd from the reactor while its sink queue is full.
    pub(crate) fn pause_stderr(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        pause_slot(reactor, &mut self.stderr_slot)
    }

    /// Return whether the stdout stream is paused on a full sink queue.
    pub fn stdout_paused(&self) -> bool {
        self.buffer.stdout_paused()
    }

    /// Return whether the stderr stream is paused on a full sink queue.
    pub fn stderr_paused(&self) -> bool {
        self.buffer.stderr_paused()
    }

    /// Re-deliver the held stdout chunk (if any) and re-register the fd when
    /// the sink has room again. Returns `true` when the stream is resumed,
    /// `false` when the sink is still full and the stream stays paused.
    pub fn resume_stdout(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<bool, CoreError> {
        if !self.buffer.deliver_pending_stdout()? {
            return Ok(false);
        }
        register_slot(reactor, &mut self.stdout_slot, true, false)?;
        Ok(true)
    }

    /// Re-deliver the held stderr chunk (if any) and re-register the fd when
    /// the sink has room again. Returns `true` when the stream is resumed,
    /// `false` when the sink is still full and the stream stays paused.
    pub fn resume_stderr(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<bool, CoreError> {
        if !self.buffer.deliver_pending_stderr()? {
            return Ok(false);
        }
        register_slot(reactor, &mut self.stderr_slot, true, false)?;
        Ok(true)
    }

    /// Take the un-delivered stdout chunk (streaming mode), if any.
    pub(crate) fn take_stdout_pending(&mut self) -> Option<Vec<u8>> {
        self.buffer.take_stdout_pending()
    }

    /// Take the un-delivered stderr chunk (streaming mode), if any.
    pub(crate) fn take_stderr_pending(&mut self) -> Option<Vec<u8>> {
        self.buffer.take_stderr_pending()
    }

    pub(crate) fn handle_stdin_writable(
        &mut self,
        reactor: &mut crate::reactor::Reactor,
    ) -> Result<(), CoreError> {
        if let Some(slot) = &self.stdin_slot {
            let done = self.writer.write_to_fd(&slot.fd)?;
            if done {
                self.drop_stdin(reactor)?;
            }
        }
        Ok(())
    }

    /// Consume the state and return (stdout, stderr) buffers.
    pub fn into_parts(mut self) -> (Vec<u8>, Vec<u8>) {
        let (stdout, stderr, _, _) = std::mem::take(&mut self.buffer).into_parts();
        (stdout, stderr)
    }

    /// Return whether the combined stdout+stderr output limit was exceeded.
    #[inline(always)]
    pub fn output_limit_exceeded(&self) -> bool {
        self.buffer.output_limit_exceeded()
    }

    /// Return whether stdout was explicitly stopped by the early-exit predicate.
    #[inline(always)]
    pub fn stdout_early_exited(&self) -> bool {
        self.buffer.stdout_early_exited()
    }

    /// Consume the state and return buffers plus drain flags.
    pub(crate) fn into_parts_with_state(mut self) -> (Vec<u8>, Vec<u8>, bool, bool) {
        std::mem::take(&mut self.buffer).into_parts()
    }
}

/// Register one slot with the given per-direction interest, leaving it in
/// place on failure. When the fd is already registered the interest is updated
/// with `EPOLL_CTL_MOD` (direction-preserving: pause/resume toggle ONE
/// registration instead of del + add), which re-arms edge-triggered
/// readiness for the same token.
fn register_slot(
    reactor: &mut crate::reactor::Reactor,
    slot: &mut Option<FdSlot>,
    readable: bool,
    writable: bool,
) -> Result<(), CoreError> {
    let Some(s) = slot.as_mut() else {
        return Ok(());
    };
    if let Some(token) = s.token {
        if s.readable == readable && s.writable == writable {
            return Ok(());
        }
        reactor.mod_(&s.fd, token, readable, writable)?;
        s.readable = readable;
        s.writable = writable;
        return Ok(());
    }
    s.token = Some(reactor.add(&s.fd, readable, writable)?);
    s.readable = readable;
    s.writable = writable;
    Ok(())
}

/// Remove a slot's fd from the reactor, skipping an already-deregistered
/// (tokenless) slot. `ENOENT` is tolerated: the fd may already have been
/// removed by reactor teardown.
fn del_slot(reactor: &crate::reactor::Reactor, slot: &FdSlot) -> Result<(), CoreError> {
    if slot.token.is_none() {
        return Ok(());
    }
    match reactor.del(&slot.fd) {
        Ok(()) => Ok(()),
        Err(e) if e.raw_os_error() == Some(libc::ENOENT) => Ok(()),
        Err(e) => Err(e),
    }
}

/// Pause a slot's READABLE interest while keeping any writable interest and
/// the registration itself (direction-preserving `EPOLL_CTL_MOD`). The token
/// is retained so a later resume re-arms the same registration. `ENOENT` is
/// tolerated for the teardown race.
fn pause_slot(
    reactor: &mut crate::reactor::Reactor,
    slot: &mut Option<FdSlot>,
) -> Result<(), CoreError> {
    let Some(s) = slot.as_mut() else {
        return Ok(());
    };
    let Some(token) = s.token else {
        return Ok(());
    };
    if !s.readable {
        return Ok(());
    }
    match reactor.mod_(&s.fd, token, false, s.writable) {
        Ok(()) => {
            s.readable = false;
            Ok(())
        }
        Err(e) if e.raw_os_error() == Some(libc::ENOENT) => {
            // Registration already gone (teardown race); drop the stale token.
            s.token = None;
            s.readable = false;
            Ok(())
        }
        Err(e) => Err(e),
    }
}