coreshift_core/io/
drain.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/
4
5//! High-level process I/O management.
6//!
7//! This module provides the [`DrainState`] structure, which coordinates the
8//! simultaneous reading from process output pipes and writing to process
9//! input pipes.
10//!
11//! This is an advanced helper for callers that already own child-process file
12//! descriptors and want non-blocking drain semantics without reimplementing
13//! the bookkeeping.
14
15use crate::CoreError;
16use crate::fd::{Fd, Token};
17use crate::io::buffer::{BufferState, ChunkSink, ReadState};
18use crate::io::writer::WriterState;
19
20/// Associates a file descriptor with an optional reactor token.
21pub(crate) struct FdSlot {
22    /// Token assigned by the reactor for this descriptor. `None` while the fd
23    /// is paused (removed from the reactor) or not yet registered.
24    pub token: Option<Token>,
25    /// The managed file descriptor.
26    pub fd: Fd,
27}
28
29/// Orchestrates non-blocking process I/O.
30///
31/// `DrainState` tracks the state of stdin, stdout, and stderr pipes for a
32/// single process. It handles the multiplexing of data between these pipes
33/// and internal buffers.
34///
35/// # Example
36/// ```no_run
37/// # use coreshift_core::io::DrainState;
38/// # use coreshift_core::reactor::Reactor;
39/// # fn example(mut drain: DrainState<fn(&[u8]) -> bool>, mut reactor: Reactor) -> Result<(), Box<dyn std::error::Error>> {
40/// while !drain.is_done() {
41///     let mut events = Vec::new();
42///     reactor.wait(&mut events, 64, -1)?;
43///     for ev in events {
44///         // Map event tokens to drain calls...
45///     }
46/// }
47/// # Ok(())
48/// # }
49/// ```
50#[repr(align(64))]
51pub struct DrainState<F>
52where
53    F: FnMut(&[u8]) -> bool,
54{
55    pub(crate) stdout_slot: Option<FdSlot>,
56    pub(crate) stderr_slot: Option<FdSlot>,
57    pub(crate) stdin_slot: Option<FdSlot>,
58
59    pub(crate) buffer: BufferState,
60    pub(crate) writer: WriterState,
61
62    pub(crate) early_exit: Option<F>,
63
64    /// `true` when the stdout slot is a pty master rather than a pipe. A pty
65    /// master reports EOF as `EIO` (returned once the session leader and all
66    /// slave holders have closed), so the read path maps a stdout `EIO` to a
67    /// clean [`ReadState::Eof`] instead of surfacing it as an I/O error.
68    pub(crate) pty_master: bool,
69}
70
71impl<F> DrainState<F>
72where
73    F: FnMut(&[u8]) -> bool,
74{
75    /// Initialize a new drain state for the provided descriptors.
76    ///
77    /// This consumes the descriptors and sets them to non-blocking mode.
78    /// `chunk_sink` enables streaming mode: every retained output chunk is
79    /// forwarded to the sink instead of being accumulated into the internal
80    /// buffers.
81    ///
82    /// ### Errors
83    /// - `EBADF`: One of the provided file descriptors is invalid.
84    pub fn new(
85        stdin_fd: Option<Fd>,
86        stdin_buf: Option<Box<[u8]>>,
87        stdout_fd: Option<Fd>,
88        stderr_fd: Option<Fd>,
89        limit: usize,
90        early_exit: Option<F>,
91        chunk_sink: Option<ChunkSink>,
92        pty_master: bool,
93    ) -> Result<Self, CoreError> {
94        let stdin_slot = if stdin_buf.is_some() {
95            if let Some(fd) = stdin_fd {
96                fd.set_nonblock()?;
97                Some(FdSlot { token: None, fd })
98            } else {
99                None
100            }
101        } else {
102            None
103        };
104
105        let stdout_slot = if let Some(fd) = stdout_fd {
106            fd.set_nonblock()?;
107            Some(FdSlot { token: None, fd })
108        } else {
109            None
110        };
111
112        let stderr_slot = if let Some(fd) = stderr_fd {
113            fd.set_nonblock()?;
114            Some(FdSlot { token: None, fd })
115        } else {
116            None
117        };
118
119        Ok(Self {
120            stdin_slot,
121            stdout_slot,
122            stderr_slot,
123            buffer: BufferState::new(limit, chunk_sink),
124            writer: WriterState::new(stdin_buf),
125            early_exit,
126            pty_master,
127        })
128    }
129
130    /// Returns `true` if all pipes have been closed or fully drained.
131    #[inline(always)]
132    pub fn is_done(&self) -> bool {
133        self.stdin_slot.is_none() && self.stdout_slot.is_none() && self.stderr_slot.is_none()
134    }
135
136    /// Apply a new window size (`TIOCSWINSZ`) to the pty master's terminal.
137    ///
138    /// Sending a signal to the foreground process group after a resize is the
139    /// caller's job (SIGWINCH); this only updates the kernel's `winsize` so a
140    /// subsequent `TIOCGWINSZ`/`SIGWINCH`-driven refresh reads the new size.
141    ///
142    /// ### Errors
143    /// - `EINVAL`: No pty master is present (non-pty spawn or stream already
144    ///   closed), or `rows`/`cols` is zero.
145    /// - `ENOTTY`: The stdout descriptor is not a terminal.
146    pub(crate) fn resize_pty(&self, rows: u16, cols: u16) -> Result<(), CoreError> {
147        if rows == 0 || cols == 0 {
148            return Err(CoreError::sys(
149                libc::EINVAL,
150                "resize_pty: rows and cols must be non-zero",
151            ));
152        }
153        let Some(slot) = &self.stdout_slot else {
154            return Err(CoreError::sys(libc::EINVAL, "resize_pty: no pty master"));
155        };
156        let ws = libc::winsize {
157            ws_row: rows,
158            ws_col: cols,
159            ws_xpixel: 0,
160            ws_ypixel: 0,
161        };
162        let r = unsafe { libc::ioctl(slot.fd.raw(), libc::TIOCSWINSZ as libc::Ioctl, &ws) };
163        crate::error::syscall_ret(r, "TIOCSWINSZ")
164    }
165
166    /// Perform a non-blocking write to stdin if pending.
167    ///
168    /// Returns `Ok(true)` if the write buffer is empty or the descriptor is
169    /// closed.
170    ///
171    /// ### Errors
172    /// - `EPIPE`: The child process closed its reading end of the pipe.
173    /// - `EIO`: Low-level I/O error.
174    #[inline(always)]
175    pub fn write_stdin(&mut self) -> Result<bool, CoreError> {
176        let fd = if let Some(s) = &self.stdin_slot {
177            &s.fd
178        } else {
179            return Ok(true);
180        };
181
182        let done = self.writer.write_to_fd(fd)?;
183        if done {
184            self.stdin_slot.take();
185            return Ok(true);
186        }
187        Ok(false)
188    }
189
190    /// Read from a slot's descriptor, mapping a pty-master `EIO` EOF to a
191    /// clean [`ReadState::Eof`]. All other errors propagate.
192    ///
193    /// Associated fn (no `self` receiver) so callers can borrow `buffer` and
194    /// `early_exit` mutably while a `stdout_slot`/`stderr_slot` borrow is
195    /// still live — disjoint field borrows the compiler can see.
196    #[inline(always)]
197    fn read_from_slot(
198        buffer: &mut BufferState,
199        fd: &Fd,
200        is_stdout: bool,
201        early_exit: &mut Option<F>,
202        pty_master: bool,
203    ) -> Result<ReadState, CoreError> {
204        match buffer.read_from_fd(fd, is_stdout, early_exit) {
205            Err(e) if pty_master && is_stdout && e.raw_os_error() == Some(libc::EIO) => {
206                Ok(ReadState::Eof)
207            }
208            other => other,
209        }
210    }
211
212    /// Perform a non-blocking read from stdout or stderr.
213    ///
214    /// Returns `Ok(true)` if the stream reached EOF, the early-exit condition
215    /// was met, or the stream is paused on a full sink queue (the caller
216    /// resumes it later). In the paused case the slot is retained.
217    ///
218    /// ### Errors
219    /// - `EOVERFLOW`: The captured output exceeded the specified limit.
220    /// - `EIO`: Low-level I/O error.
221    #[inline(always)]
222    pub fn read_fd(&mut self, is_stdout: bool) -> Result<bool, CoreError> {
223        let pty_master = self.pty_master;
224        let read_state = {
225            let slot = if is_stdout {
226                &self.stdout_slot
227            } else {
228                &self.stderr_slot
229            };
230            let fd = if let Some(s) = slot {
231                &s.fd
232            } else {
233                return Ok(true);
234            };
235            Self::read_from_slot(&mut self.buffer, fd, is_stdout, &mut self.early_exit, pty_master)?
236        };
237
238        match read_state {
239            ReadState::Open => Ok(false),
240            ReadState::Paused => Ok(false),
241            ReadState::Eof | ReadState::EarlyExit => {
242                if is_stdout {
243                    self.stdout_slot.take();
244                } else {
245                    self.stderr_slot.take();
246                }
247                Ok(true)
248            }
249        }
250    }
251
252    /// Extract all active slots for cleanup or reactor removal.
253    pub(crate) fn take_all_slots(&mut self) -> Vec<FdSlot> {
254        let mut slots = Vec::new();
255        if let Some(slot) = self.stdin_slot.take() {
256            slots.push(slot);
257        }
258        if let Some(slot) = self.stdout_slot.take() {
259            slots.push(slot);
260        }
261        if let Some(slot) = self.stderr_slot.take() {
262            slots.push(slot);
263        }
264        slots
265    }
266
267    pub(crate) fn register_with_reactor(
268        &mut self,
269        reactor: &mut crate::reactor::Reactor,
270    ) -> Result<(), CoreError> {
271        register_slot(reactor, &mut self.stdin_slot, false, true)?;
272        register_slot(reactor, &mut self.stdout_slot, true, false)?;
273        register_slot(reactor, &mut self.stderr_slot, true, false)?;
274        Ok(())
275    }
276
277    pub(crate) fn stdout_matches(&self, token: Token) -> bool {
278        self.stdout_slot
279            .as_ref()
280            .is_some_and(|slot| slot.token == Some(token))
281    }
282
283    pub(crate) fn stderr_matches(&self, token: Token) -> bool {
284        self.stderr_slot
285            .as_ref()
286            .is_some_and(|slot| slot.token == Some(token))
287    }
288
289    pub(crate) fn stdin_matches(&self, token: Token) -> bool {
290        self.stdin_slot
291            .as_ref()
292            .is_some_and(|slot| slot.token == Some(token))
293    }
294
295    pub(crate) fn drop_stdout(
296        &mut self,
297        reactor: &mut crate::reactor::Reactor,
298    ) -> Result<(), CoreError> {
299        if let Some(slot) = self.stdout_slot.take() {
300            del_slot(reactor, &slot)?;
301        }
302        Ok(())
303    }
304
305    pub(crate) fn drop_stderr(
306        &mut self,
307        reactor: &mut crate::reactor::Reactor,
308    ) -> Result<(), CoreError> {
309        if let Some(slot) = self.stderr_slot.take() {
310            del_slot(reactor, &slot)?;
311        }
312        Ok(())
313    }
314
315    pub(crate) fn drop_stdin(
316        &mut self,
317        reactor: &mut crate::reactor::Reactor,
318    ) -> Result<(), CoreError> {
319        if let Some(slot) = self.stdin_slot.take() {
320            del_slot(reactor, &slot)?;
321        }
322        self.writer.buf = None;
323        Ok(())
324    }
325
326    pub(crate) fn handle_stdout_ready(
327        &mut self,
328        reactor: &mut crate::reactor::Reactor,
329    ) -> Result<(), CoreError> {
330        if let Some(slot) = &self.stdout_slot {
331            let read_state = Self::read_from_slot(
332                &mut self.buffer,
333                &slot.fd,
334                true,
335                &mut self.early_exit,
336                self.pty_master,
337            )?;
338            match read_state {
339                ReadState::Open => {}
340                ReadState::Paused => {
341                    // Sink queue full: remove the fd from the reactor so the
342                    // edge-triggered readiness does not spin the loop; the
343                    // caller re-registers via `resume_stdout` when it has
344                    // drained the queue.
345                    self.pause_stdout(reactor)?;
346                }
347                ReadState::Eof | ReadState::EarlyExit => {
348                    self.drop_stdout(reactor)?;
349                }
350            }
351        }
352        Ok(())
353    }
354
355    pub(crate) fn handle_stderr_ready(
356        &mut self,
357        reactor: &mut crate::reactor::Reactor,
358    ) -> Result<(), CoreError> {
359        if let Some(slot) = &self.stderr_slot {
360            let read_state = Self::read_from_slot(
361                &mut self.buffer,
362                &slot.fd,
363                false,
364                &mut self.early_exit,
365                self.pty_master,
366            )?;
367            match read_state {
368                ReadState::Open => {}
369                ReadState::Paused => {
370                    self.pause_stderr(reactor)?;
371                }
372                ReadState::Eof | ReadState::EarlyExit => {
373                    self.drop_stderr(reactor)?;
374                }
375            }
376        }
377        Ok(())
378    }
379
380    /// Remove the stdout fd from the reactor while its sink queue is full.
381    /// The slot is retained (token cleared) so the stream can be resumed.
382    pub(crate) fn pause_stdout(
383        &mut self,
384        reactor: &mut crate::reactor::Reactor,
385    ) -> Result<(), CoreError> {
386        pause_slot(reactor, &mut self.stdout_slot)
387    }
388
389    /// Remove the stderr fd from the reactor while its sink queue is full.
390    pub(crate) fn pause_stderr(
391        &mut self,
392        reactor: &mut crate::reactor::Reactor,
393    ) -> Result<(), CoreError> {
394        pause_slot(reactor, &mut self.stderr_slot)
395    }
396
397    /// Return whether the stdout stream is paused on a full sink queue.
398    pub fn stdout_paused(&self) -> bool {
399        self.buffer.stdout_paused()
400    }
401
402    /// Return whether the stderr stream is paused on a full sink queue.
403    pub fn stderr_paused(&self) -> bool {
404        self.buffer.stderr_paused()
405    }
406
407    /// Re-deliver the held stdout chunk (if any) and re-register the fd when
408    /// the sink has room again. Returns `true` when the stream is resumed,
409    /// `false` when the sink is still full and the stream stays paused.
410    pub fn resume_stdout(
411        &mut self,
412        reactor: &mut crate::reactor::Reactor,
413    ) -> Result<bool, CoreError> {
414        if !self.buffer.deliver_pending_stdout()? {
415            return Ok(false);
416        }
417        register_slot(reactor, &mut self.stdout_slot, true, false)?;
418        Ok(true)
419    }
420
421    /// Re-deliver the held stderr chunk (if any) and re-register the fd when
422    /// the sink has room again. Returns `true` when the stream is resumed,
423    /// `false` when the sink is still full and the stream stays paused.
424    pub fn resume_stderr(
425        &mut self,
426        reactor: &mut crate::reactor::Reactor,
427    ) -> Result<bool, CoreError> {
428        if !self.buffer.deliver_pending_stderr()? {
429            return Ok(false);
430        }
431        register_slot(reactor, &mut self.stderr_slot, true, false)?;
432        Ok(true)
433    }
434
435    /// Take the un-delivered stdout chunk (streaming mode), if any.
436    pub(crate) fn take_stdout_pending(&mut self) -> Option<Vec<u8>> {
437        self.buffer.take_stdout_pending()
438    }
439
440    /// Take the un-delivered stderr chunk (streaming mode), if any.
441    pub(crate) fn take_stderr_pending(&mut self) -> Option<Vec<u8>> {
442        self.buffer.take_stderr_pending()
443    }
444
445    pub(crate) fn handle_stdin_writable(
446        &mut self,
447        reactor: &mut crate::reactor::Reactor,
448    ) -> Result<(), CoreError> {
449        if let Some(slot) = &self.stdin_slot {
450            let done = self.writer.write_to_fd(&slot.fd)?;
451            if done {
452                self.drop_stdin(reactor)?;
453            }
454        }
455        Ok(())
456    }
457
458    /// Consume the state and return (stdout, stderr) buffers.
459    pub fn into_parts(mut self) -> (Vec<u8>, Vec<u8>) {
460        let (stdout, stderr, _, _) = std::mem::take(&mut self.buffer).into_parts();
461        (stdout, stderr)
462    }
463
464    /// Return whether the combined stdout+stderr output limit was exceeded.
465    #[inline(always)]
466    pub fn output_limit_exceeded(&self) -> bool {
467        self.buffer.output_limit_exceeded()
468    }
469
470    /// Return whether stdout was explicitly stopped by the early-exit predicate.
471    #[inline(always)]
472    pub fn stdout_early_exited(&self) -> bool {
473        self.buffer.stdout_early_exited()
474    }
475
476    /// Consume the state and return buffers plus drain flags.
477    pub(crate) fn into_parts_with_state(mut self) -> (Vec<u8>, Vec<u8>, bool, bool) {
478        std::mem::take(&mut self.buffer).into_parts()
479    }
480}
481
482/// Register one slot, leaving it in place on failure and treating a second
483/// registration as a no-op so the fd and stream are never lost.
484fn register_slot(
485    reactor: &mut crate::reactor::Reactor,
486    slot: &mut Option<FdSlot>,
487    readable: bool,
488    writable: bool,
489) -> Result<(), CoreError> {
490    let Some(s) = slot.as_mut() else {
491        return Ok(());
492    };
493    if s.token.is_some() {
494        return Ok(());
495    }
496    s.token = Some(reactor.add(&s.fd, readable, writable)?);
497    Ok(())
498}
499
500/// Remove a slot's fd from the reactor, skipping an already-paused (tokenless)
501/// slot. `ENOENT` is tolerated: the fd may already have been removed by a
502/// pause or by reactor teardown.
503fn del_slot(
504    reactor: &crate::reactor::Reactor,
505    slot: &FdSlot,
506) -> Result<(), CoreError> {
507    if slot.token.is_none() {
508        return Ok(());
509    }
510    match reactor.del(&slot.fd) {
511        Ok(()) => Ok(()),
512        Err(e) if e.raw_os_error() == Some(libc::ENOENT) => Ok(()),
513        Err(e) => Err(e),
514    }
515}
516
517/// Remove a slot's fd from the reactor and clear its token, keeping the slot
518/// so the stream can be resumed later.
519fn pause_slot(
520    reactor: &mut crate::reactor::Reactor,
521    slot: &mut Option<FdSlot>,
522) -> Result<(), CoreError> {
523    let Some(s) = slot.as_mut() else {
524        return Ok(());
525    };
526    if s.token.is_none() {
527        return Ok(());
528    }
529    s.token = None;
530    match reactor.del(&s.fd) {
531        Ok(()) => Ok(()),
532        Err(e) if e.raw_os_error() == Some(libc::ENOENT) => Ok(()),
533        Err(e) => Err(e),
534    }
535}
coreshift_core/io/drain.rs

coreshift_core/io/
drain.rs
