coreshift_core/io/

buffer.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/

//! Asynchronous I/O buffering.
//!
//! This module provides the [`BufferState`] structure which accumulates
//! stdout and stderr data from monitored processes.

use crate::CoreError;
use crate::fd::Fd;
use std::sync::Arc;

const READ_CHUNK: usize = 65536;

/// Outcome of delivering one chunk to a [`ChunkSink`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SinkResult {
    /// The chunk was consumed and draining may continue.
    Accept,
    /// The consumer's bounded queue is full. The stream pauses; the chunk is
    /// retained and re-delivered when the consumer resumes the stream. Bytes
    /// are never dropped on this path.
    Pause,
}

/// Streaming drain consumer.
///
/// The sink is invoked for every retained output chunk while draining (per
/// stream: `is_stdout` distinguishes stdout from stderr). Returning
/// [`SinkResult::Pause`] stops the drain without losing the chunk — the drain
/// holds it and re-delivers on resume, so the stream stays lossless and the
/// read loop never blocks the reactor thread. The caller (e.g. a daemon
/// exec backend) implements the sink against its own bounded per-session
/// queue and drives resume when the queue drains.
pub type ChunkSink = Arc<dyn Fn(bool, &[u8]) -> SinkResult + Send + Sync>;

/// Accumulates output from process streams.
///
/// `BufferState` manages the collection of bytes from stdout and stderr pipes.
/// It enforces a combined memory limit to prevent runaway memory usage by
/// misbehaving processes. When a [`ChunkSink`] is attached the buffer stops
/// accumulating: chunks are forwarded to the sink as they are read and the
/// retained output stays empty (streaming mode).
#[derive(Default)]
#[repr(align(64))]
pub(crate) struct BufferState {
    stdout: Vec<u8>,
    stderr: Vec<u8>,
    limit: usize,
    output_limit_exceeded: bool,
    stdout_early_exited: bool,
    sink: Option<ChunkSink>,
    scratch: Vec<u8>,
    stdout_pending: Option<Vec<u8>>,
    stderr_pending: Option<Vec<u8>>,
    stdout_paused: bool,
    stderr_paused: bool,
}

/// Result of one non-blocking drain attempt.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ReadState {
    /// The descriptor would block before EOF.
    Open,
    /// Kernel EOF was reached.
    Eof,
    /// The caller-provided early-exit predicate requested stdout shutdown.
    EarlyExit,
    /// The chunk sink returned [`SinkResult::Pause`]; the stream is paused and
    /// the un-delivered chunk is held for re-delivery on resume.
    Paused,
}

impl BufferState {
    /// Create a new buffer state with the specified memory limit and an
    /// optional streaming sink.
    pub(crate) fn new(limit: usize, sink: Option<ChunkSink>) -> Self {
        Self {
            stdout: Vec::with_capacity(1024),
            stderr: Vec::with_capacity(1024),
            limit,
            output_limit_exceeded: false,
            stdout_early_exited: false,
            sink,
            scratch: Vec::with_capacity(READ_CHUNK),
            stdout_pending: None,
            stderr_pending: None,
            stdout_paused: false,
            stderr_paused: false,
        }
    }

    /// Drain available data from a file descriptor into internal storage (or
    /// forward it to the chunk sink when one is attached).
    ///
    /// # Returns
    /// * `Ok(ReadState::Eof)` if EOF was reached.
    /// * `Ok(ReadState::EarlyExit)` if stdout matched the early-exit callback.
    /// * `Ok(ReadState::Paused)` if the chunk sink's queue is full (streaming
    ///   mode; the un-delivered chunk is retained for resume).
    /// * `Ok(ReadState::Open)` if the operation would block (`EAGAIN`).
    #[inline(always)]
    pub(crate) fn read_from_fd(
        &mut self,
        fd: &Fd,
        is_stdout: bool,
        early_exit: &mut Option<impl FnMut(&[u8]) -> bool>,
    ) -> Result<ReadState, CoreError> {
        if self.sink.is_some() {
            return self.read_from_fd_streaming(fd, is_stdout, early_exit);
        }
        loop {
            let current_total = self.stdout.len().saturating_add(self.stderr.len());
            let remaining_limit = self.limit.saturating_sub(current_total);

            if remaining_limit == 0 {
                let mut drop_buf = [0u8; 8192];
                match fd.read_slice(&mut drop_buf) {
                    Ok(Some(n)) if n > 0 => {
                        self.output_limit_exceeded = true;
                        continue;
                    }
                    Ok(Some(_)) => return Ok(ReadState::Eof),
                    Ok(None) => return Ok(ReadState::Open),
                    Err(e) => return Err(e),
                }
            }

            let dest = if is_stdout {
                &mut self.stdout
            } else {
                &mut self.stderr
            };
            let len = dest.len();

            // Ensure space and read directly into the Vec.
            // We resize with 0s to remain safe (no UB with uninitialized memory).
            let to_read = remaining_limit.min(READ_CHUNK);
            dest.resize(len + to_read, 0);

            match fd.read_slice(&mut dest[len..len + to_read]) {
                Ok(Some(n)) if n > 0 => {
                    dest.truncate(len + n);

                    if is_stdout
                        && let Some(f) = early_exit
                        && f(&dest[len..len + n])
                    {
                        self.stdout_early_exited = true;
                        return Ok(ReadState::EarlyExit);
                    }
                }
                Ok(Some(_)) => {
                    dest.truncate(len);
                    return Ok(ReadState::Eof);
                }
                Ok(None) => {
                    dest.truncate(len);
                    return Ok(ReadState::Open);
                }
                Err(e) => {
                    dest.truncate(len);
                    return Err(e);
                }
            }
        }
    }

    /// Streaming-mode drain: read one chunk at a time into the scratch buffer
    /// and forward it to the sink. Never accumulates into the retained output
    /// and never drops bytes: when the sink reports the queue full, the chunk
    /// is held (`stdout_pending`/`stderr_pending`) and the stream pauses.
    fn read_from_fd_streaming(
        &mut self,
        fd: &Fd,
        is_stdout: bool,
        early_exit: &mut Option<impl FnMut(&[u8]) -> bool>,
    ) -> Result<ReadState, CoreError> {
        let Some(sink) = &self.sink else {
            return Ok(ReadState::Open);
        };
        loop {
            // Re-deliver a held chunk before reading new bytes, preserving
            // stream order across a pause/resume cycle.
            let pending = if is_stdout {
                self.stdout_pending.take()
            } else {
                self.stderr_pending.take()
            };
            if let Some(pending) = pending {
                match sink(is_stdout, &pending) {
                    SinkResult::Accept => {
                        if is_stdout {
                            self.stdout_paused = false;
                        } else {
                            self.stderr_paused = false;
                        }
                    }
                    SinkResult::Pause => {
                        if is_stdout {
                            self.stdout_pending = Some(pending);
                            self.stdout_paused = true;
                        } else {
                            self.stderr_pending = Some(pending);
                            self.stderr_paused = true;
                        }
                        return Ok(ReadState::Paused);
                    }
                }
            }

            self.scratch.resize(READ_CHUNK, 0);
            match fd.read_slice(&mut self.scratch) {
                Ok(Some(n)) if n > 0 => {
                    self.scratch.truncate(n);

                    if is_stdout
                        && let Some(f) = early_exit
                        && f(&self.scratch)
                    {
                        self.stdout_early_exited = true;
                        return Ok(ReadState::EarlyExit);
                    }

                    match sink(is_stdout, &self.scratch) {
                        SinkResult::Accept => continue,
                        SinkResult::Pause => {
                            if is_stdout {
                                self.stdout_pending = Some(self.scratch.clone());
                                self.stdout_paused = true;
                            } else {
                                self.stderr_pending = Some(self.scratch.clone());
                                self.stderr_paused = true;
                            }
                            return Ok(ReadState::Paused);
                        }
                    }
                }
                Ok(Some(_)) => {
                    self.scratch.clear();
                    return Ok(ReadState::Eof);
                }
                Ok(None) => {
                    self.scratch.clear();
                    return Ok(ReadState::Open);
                }
                Err(e) => {
                    self.scratch.clear();
                    return Err(e);
                }
            }
        }
    }

    /// Attempt to re-deliver the held stdout chunk. Returns `true` when the
    /// stream is resumed (chunk delivered, or nothing held), `false` when the
    /// sink is still full and the stream stays paused.
    pub(crate) fn deliver_pending_stdout(&mut self) -> Result<bool, CoreError> {
        if !self.stdout_paused {
            return Ok(true);
        }
        let Some(sink) = &self.sink else {
            self.stdout_paused = false;
            self.stdout_pending = None;
            return Ok(true);
        };
        match self.stdout_pending.take() {
            None => {
                self.stdout_paused = false;
                Ok(true)
            }
            Some(pending) => match sink(true, &pending) {
                SinkResult::Accept => {
                    self.stdout_paused = false;
                    Ok(true)
                }
                SinkResult::Pause => {
                    self.stdout_pending = Some(pending);
                    Ok(false)
                }
            },
        }
    }

    /// Attempt to re-deliver the held stderr chunk. Returns `true` when the
    /// stream is resumed, `false` when the sink is still full.
    pub(crate) fn deliver_pending_stderr(&mut self) -> Result<bool, CoreError> {
        if !self.stderr_paused {
            return Ok(true);
        }
        let Some(sink) = &self.sink else {
            self.stderr_paused = false;
            self.stderr_pending = None;
            return Ok(true);
        };
        match self.stderr_pending.take() {
            None => {
                self.stderr_paused = false;
                Ok(true)
            }
            Some(pending) => match sink(false, &pending) {
                SinkResult::Accept => {
                    self.stderr_paused = false;
                    Ok(true)
                }
                SinkResult::Pause => {
                    self.stderr_pending = Some(pending);
                    Ok(false)
                }
            },
        }
    }

    /// Return whether the stdout stream is paused on a full sink queue.
    #[inline(always)]
    pub(crate) fn stdout_paused(&self) -> bool {
        self.stdout_paused
    }

    /// Return whether the stderr stream is paused on a full sink queue.
    #[inline(always)]
    pub(crate) fn stderr_paused(&self) -> bool {
        self.stderr_paused
    }

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

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

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

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

    /// Consume the state and return the accumulated buffers.
    pub(crate) fn into_parts(mut self) -> (Vec<u8>, Vec<u8>, bool, bool) {
        (
            std::mem::take(&mut self.stdout),
            std::mem::take(&mut self.stderr),
            self.output_limit_exceeded,
            self.stdout_early_exited,
        )
    }
}