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;

/// Associates a file descriptor with an optional reactor token.
pub(crate) struct FdSlot {
    /// Token assigned by the reactor for this descriptor. `None` while the fd
    /// is paused (removed from the reactor) or not yet registered.
    pub token: Option<Token>,
    /// The managed file descriptor.
    pub fd: Fd,
}

/// 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>,
}

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.
    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>,
    ) -> 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 })
            } else {
                None
            }
        } else {
            None
        };

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

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

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

    /// 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()
    }

    /// 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)
    }

    /// 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 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.buffer
                .read_from_fd(fd, is_stdout, &mut self.early_exit)?
        };

        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
                .buffer
                .read_from_fd(&slot.fd, true, &mut self.early_exit)?;
            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
                .buffer
                .read_from_fd(&slot.fd, false, &mut self.early_exit)?;
            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, leaving it in place on failure and treating a second
/// registration as a no-op so the fd and stream are never lost.
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 s.token.is_some() {
        return Ok(());
    }
    s.token = Some(reactor.add(&s.fd, readable, writable)?);
    Ok(())
}

/// Remove a slot's fd from the reactor, skipping an already-paused (tokenless)
/// slot. `ENOENT` is tolerated: the fd may already have been removed by a
/// pause or 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),
    }
}

/// Remove a slot's fd from the reactor and clear its token, keeping the slot
/// so the stream can be resumed later.
fn pause_slot(
    reactor: &mut crate::reactor::Reactor,
    slot: &mut Option<FdSlot>,
) -> Result<(), CoreError> {
    let Some(s) = slot.as_mut() else {
        return Ok(());
    };
    if s.token.is_none() {
        return Ok(());
    }
    s.token = None;
    match reactor.del(&s.fd) {
        Ok(()) => Ok(()),
        Err(e) if e.raw_os_error() == Some(libc::ENOENT) => Ok(()),
        Err(e) => Err(e),
    }
}