use std::io::{self, Read};
const READ_CHUNK: usize = 4096;
#[cfg(windows)]
pub(crate) trait PipeSource: Read + std::os::windows::io::AsRawHandle {}
#[cfg(windows)]
impl<T: Read + std::os::windows::io::AsRawHandle> PipeSource for T {}
#[cfg(unix)]
pub(crate) trait PipeSource: Read + std::os::unix::io::AsRawFd {}
#[cfg(unix)]
impl<T: Read + std::os::unix::io::AsRawFd> PipeSource for T {}
enum Outcome {
Data(usize),
Empty,
Eof,
}
pub(crate) struct PipeDrain<R: PipeSource> {
reader: Option<R>,
buf: [u8; READ_CHUNK],
}
impl<R: PipeSource> PipeDrain<R> {
pub(crate) fn new(reader: R) -> Self {
#[cfg(unix)]
{
if !set_nonblocking(&reader) {
return Self {
reader: None,
buf: [0u8; READ_CHUNK],
};
}
}
Self {
reader: Some(reader),
buf: [0u8; READ_CHUNK],
}
}
pub(crate) fn finished(&self) -> bool {
self.reader.is_none()
}
pub(crate) fn release(&mut self) {
self.reader = None;
}
pub(crate) fn drain(&mut self, budget: usize, sink: &mut impl FnMut(&[u8])) -> usize {
let mut moved = 0;
while moved < budget {
let Some(reader) = self.reader.as_mut() else {
break;
};
let want = READ_CHUNK.min(budget - moved);
match read_available(reader, &mut self.buf[..want]) {
Outcome::Data(0) | Outcome::Eof => {
self.reader = None;
break;
}
Outcome::Data(n) => {
sink(&self.buf[..n]);
moved += n;
}
Outcome::Empty => break,
}
}
moved
}
}
#[cfg(windows)]
fn read_available<R: PipeSource>(reader: &mut R, buf: &mut [u8]) -> Outcome {
use std::ffi::c_void;
#[link(name = "kernel32")]
extern "system" {
fn PeekNamedPipe(
handle: *mut c_void,
buffer: *mut c_void,
buffer_size: u32,
bytes_read: *mut u32,
total_available: *mut u32,
bytes_left_this_message: *mut u32,
) -> i32;
}
const ERROR_BROKEN_PIPE: i32 = 109;
let mut available: u32 = 0;
let ok = unsafe {
PeekNamedPipe(
reader.as_raw_handle(),
std::ptr::null_mut(),
0,
std::ptr::null_mut(),
&mut available,
std::ptr::null_mut(),
)
};
if ok == 0 {
let err = io::Error::last_os_error();
return match err.raw_os_error() {
Some(ERROR_BROKEN_PIPE) => Outcome::Eof,
_ => Outcome::Eof,
};
}
if available == 0 {
return Outcome::Empty;
}
let want = buf.len().min(available as usize);
match reader.read(&mut buf[..want]) {
Ok(0) => Outcome::Eof,
Ok(n) => Outcome::Data(n),
Err(e) if e.kind() == io::ErrorKind::Interrupted => Outcome::Empty,
Err(_) => Outcome::Eof,
}
}
#[cfg(unix)]
fn read_available<R: PipeSource>(reader: &mut R, buf: &mut [u8]) -> Outcome {
match reader.read(buf) {
Ok(0) => Outcome::Eof,
Ok(n) => Outcome::Data(n),
Err(e) if e.kind() == io::ErrorKind::WouldBlock => Outcome::Empty,
Err(e) if e.kind() == io::ErrorKind::Interrupted => Outcome::Empty,
Err(_) => Outcome::Eof,
}
}
#[cfg(unix)]
fn set_nonblocking<R: PipeSource>(reader: &R) -> bool {
let fd = reader.as_raw_fd();
unsafe {
let flags = libc::fcntl(fd, libc::F_GETFL);
if flags < 0 {
return false;
}
libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) >= 0
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::process::Stdio;
use crate::process::shell_command;
fn collect(drain: &mut PipeDrain<impl PipeSource>, budget: usize) -> Vec<u8> {
let mut out = Vec::new();
drain.drain(budget, &mut |chunk: &[u8]| out.extend_from_slice(chunk));
out
}
#[test]
fn a_finished_command_ends_its_pipe() {
let mut child = shell_command("echo drained_ok")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()
.expect("spawn");
let mut pipe = PipeDrain::new(child.stdout.take().expect("stdout is piped"));
let _ = child.wait();
let mut got = Vec::new();
for _ in 0..200 {
got.extend_from_slice(&collect(&mut pipe, 64 * 1024));
if pipe.finished() {
break;
}
std::thread::sleep(std::time::Duration::from_millis(5));
}
assert!(
pipe.finished(),
"the pipe of an exited command must reach its end"
);
assert!(
String::from_utf8_lossy(&got).contains("drained_ok"),
"everything written before the end must be drained: {got:?}"
);
}
#[test]
fn an_idle_pipe_reports_empty_without_blocking() {
#[cfg(windows)]
let line = "ping -n 4 127.0.0.1 >nul";
#[cfg(unix)]
let line = "sleep 3";
let mut child = shell_command(line)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()
.expect("spawn");
let mut pipe = PipeDrain::new(child.stdout.take().expect("stdout is piped"));
let start = std::time::Instant::now();
let moved = collect(&mut pipe, 64 * 1024);
let elapsed = start.elapsed();
assert!(
moved.is_empty(),
"a silent command wrote nothing: {moved:?}"
);
assert!(
!pipe.finished(),
"a pipe held open by a live process has not ended"
);
assert!(
elapsed < std::time::Duration::from_secs(1),
"draining an idle pipe must return at once, not wait on it: {elapsed:?}"
);
let _ = child.kill();
let _ = child.wait();
}
#[test]
fn a_pass_stops_at_its_budget() {
const PAYLOAD: usize = 2048;
const BUDGET: usize = 1024;
let mut child = shell_command(&format!("echo {}", "x".repeat(PAYLOAD)))
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()
.expect("spawn");
let mut pipe = PipeDrain::new(child.stdout.take().expect("stdout is piped"));
let _ = child.wait();
let first = collect(&mut pipe, BUDGET);
assert_eq!(
first.len(),
BUDGET,
"one pass must stop at its budget, with the payload already buffered"
);
let rest = collect(&mut pipe, 64 * 1024);
assert_eq!(
first.len() + rest.len(),
PAYLOAD + line_ending_len(),
"the bytes past the budget are kept for the next pass"
);
}
const fn line_ending_len() -> usize {
if cfg!(windows) {
2
} else {
1
}
}
}