stdio: make stdout block-buffered when not associated to a terminal

Author: andreacorbellini

library/std/src/io/stdio.rs

@@ -8,7 +8,8 @@ use crate::fmt;
 use crate::fs::File;
 use crate::io::prelude::*;
 use crate::io::{
- self, BorrowedCursor, BufReader, IoSlice, IoSliceMut, LineWriter, Lines, SpecReadByte,
+ self, BorrowedCursor, BufReader, BufWriter, IoSlice, IoSliceMut, LineWriter, Lines,
+ SpecReadByte,
 };
 use crate::panic::{RefUnwindSafe, UnwindSafe};
 use crate::sync::atomic::{Atomic, AtomicBool, Ordering};
@@ -168,6 +169,20 @@ impl Write for StdoutRaw {
  }
 }
 
+#[cfg(any(
+ unix,
+ target_os = "hermit",
+ target_os = "trusty",
+ target_os = "wasi",
+ target_os = "motor",
+))]
+impl crate::os::fd::AsFd for StdoutRaw {
+ #[inline]
+ fn as_fd(&self) -> crate::os::fd::BorrowedFd<'_> {
+ unsafe { crate::os::fd::BorrowedFd::borrow_raw(1) }
+ }
+}
+
 impl Write for StderrRaw {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  handle_ebadf(self.0.write(buf), || Ok(buf.len()))
@@ -576,6 +591,76 @@ impl fmt::Debug for StdinLock<'_> {
  }
 }
 
+/// A buffered writer for stdout and stderr.
+///
+/// This writer may be either [line-buffered](LineWriter) or [block-buffered](BufWriter), depending
+/// on whether the underlying file is a terminal or not.
+#[derive(Debug)]
+enum StdioBufWriter<W: Write> {
+ LineBuffered(LineWriter<W>),
+ BlockBuffered(BufWriter<W>),
+}
+
+impl<W: Write + IsTerminal> StdioBufWriter<W> {
+ /// Wraps a writer using the most appropriate buffering method.
+ ///
+ /// If `w` is a terminal, then the resulting `StdioBufWriter` will be line-buffered, otherwise
+ /// it will be block-buffered.
+ fn new(w: W) -> Self {
+ if w.is_terminal() {
+ Self::LineBuffered(LineWriter::new(w))
+ } else {
+ Self::BlockBuffered(BufWriter::new(w))
+ }
+ }
+}
+
+impl<W: Write> StdioBufWriter<W> {
+ /// Wraps a writer using a block-buffer with the given capacity.
+ fn with_capacity(cap: usize, w: W) -> Self {
+ Self::BlockBuffered(BufWriter::with_capacity(cap, w))
+ }
+}
+
+impl<W: Write> Write for StdioBufWriter<W> {
+ fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+ match self {
+ Self::LineBuffered(w) => w.write(buf),
+ Self::BlockBuffered(w) => w.write(buf),
+ }
+ }
+ fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+ match self {
+ Self::LineBuffered(w) => w.write_vectored(bufs),
+ Self::BlockBuffered(w) => w.write_vectored(bufs),
+ }
+ }
+ fn is_write_vectored(&self) -> bool {
+ match self {
+ Self::LineBuffered(w) => w.is_write_vectored(),
+ Self::BlockBuffered(w) => w.is_write_vectored(),
+ }
+ }
+ fn flush(&mut self) -> io::Result<()> {
+ match self {
+ Self::LineBuffered(w) => w.flush(),
+ Self::BlockBuffered(w) => w.flush(),
+ }
+ }
+ fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+ match self {
+ Self::LineBuffered(w) => w.write_all(buf),
+ Self::BlockBuffered(w) => w.write_all(buf),
+ }
+ }
+ fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+ match self {
+ Self::LineBuffered(w) => w.write_all_vectored(bufs),
+ Self::BlockBuffered(w) => w.write_all_vectored(bufs),
+ }
+ }
+}
+
 /// A handle to the global standard output stream of the current process.
 ///
 /// Each handle shares a global buffer of data to be written to the standard
@@ -606,10 +691,9 @@ impl fmt::Debug for StdinLock<'_> {
 /// [`io::stdout`]: stdout
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Stdout {
- // FIXME: this should be LineWriter or BufWriter depending on the state of
- // stdout (tty or not). Note that if this is not line buffered it
- // should also flush-on-panic or some form of flush-on-abort.
- inner: &'static ReentrantLock<RefCell<LineWriter<StdoutRaw>>>,
+ // FIXME: if this is not line buffered it should flush-on-panic or some
+ // form of flush-on-abort.
+ inner: &'static ReentrantLock<RefCell<StdioBufWriter<StdoutRaw>>>,
 }
 
 /// A locked reference to the [`Stdout`] handle.
@@ -638,10 +722,10 @@ pub struct Stdout {
 #[must_use = "if unused stdout will immediately unlock"]
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct StdoutLock<'a> {
- inner: ReentrantLockGuard<'a, RefCell<LineWriter<StdoutRaw>>>,
+ inner: ReentrantLockGuard<'a, RefCell<StdioBufWriter<StdoutRaw>>>,
 }
 
-static STDOUT: OnceLock<ReentrantLock<RefCell<LineWriter<StdoutRaw>>>> = OnceLock::new();
+static STDOUT: OnceLock<ReentrantLock<RefCell<StdioBufWriter<StdoutRaw>>>> = OnceLock::new();
 
 /// Constructs a new handle to the standard output of the current process.
 ///
@@ -716,7 +800,7 @@ static STDOUT: OnceLock<ReentrantLock<RefCell<LineWriter<StdoutRaw>>>> = OnceLoc
 pub fn stdout() -> Stdout {
  Stdout {
  inner: STDOUT
- .get_or_init(|| ReentrantLock::new(RefCell::new(LineWriter::new(stdout_raw())))),
+ .get_or_init(|| ReentrantLock::new(RefCell::new(StdioBufWriter::new(stdout_raw())))),
  }
 }
 
@@ -727,7 +811,7 @@ pub fn cleanup() {
  let mut initialized = false;
  let stdout = STDOUT.get_or_init(|| {
  initialized = true;
- ReentrantLock::new(RefCell::new(LineWriter::with_capacity(0, stdout_raw())))
+ ReentrantLock::new(RefCell::new(StdioBufWriter::with_capacity(0, stdout_raw())))
  });
 
  if !initialized {
@@ -736,7 +820,7 @@ pub fn cleanup() {
  // might have leaked a StdoutLock, which would
  // otherwise cause a deadlock here.
  if let Some(lock) = stdout.try_lock() {
- *lock.borrow_mut() = LineWriter::with_capacity(0, stdout_raw());
+ *lock.borrow_mut() = StdioBufWriter::with_capacity(0, stdout_raw());
  }
  }
 }
@@ -1262,7 +1346,16 @@ macro_rules! impl_is_terminal {
  )*}
 }
 
-impl_is_terminal!(File, Stdin, StdinLock<'_>, Stdout, StdoutLock<'_>, Stderr, StderrLock<'_>);
+impl_is_terminal!(
+ File,
+ Stdin,
+ StdinLock<'_>,
+ Stdout,
+ StdoutLock<'_>,
+ StdoutRaw,
+ Stderr,
+ StderrLock<'_>,
+);
 
 #[unstable(
  feature = "print_internals",

library/std/src/io/stdio/tests.rs

@@ -1,6 +1,8 @@
 use super::*;
+use crate::assert_matches::assert_matches;
+use crate::os::fd::{FromRawFd, OwnedFd};
 use crate::panic::{RefUnwindSafe, UnwindSafe};
-use crate::sync::mpsc::sync_channel;
+use crate::sync::mpsc::{TryRecvError, sync_channel};
 use crate::thread;
 
 #[test]
@@ -164,3 +166,95 @@ where
  [Start1, Acquire1, Start2, Release1, Acquire2, Release2, Acquire1, Release1]
  );
 }
+
+#[test]
+#[cfg(unix)]
+fn stdiobufwriter_line_buffered() {
+ let mut fd1 = -1;
+ let mut fd2 = -1;
+
+ if unsafe {
+ libc::openpty(
+ &raw mut fd1,
+ &raw mut fd2,
+ crate::ptr::null_mut(),
+ crate::ptr::null(),
+ crate::ptr::null(),
+ )
+ } < 0
+ {
+ panic!("openpty() failed with {:?}", io::Error::last_os_error());
+ }
+
+ let writer = unsafe { File::from_raw_fd(fd1) };
+ let mut reader = unsafe { File::from_raw_fd(fd2) };
+
+ assert!(writer.is_terminal(), "file descriptor returned by openpty() must be a terminal");
+ assert!(reader.is_terminal(), "file descriptor returned by openpty() must be a terminal");
+
+ let (sender, receiver) = sync_channel(64);
+
+ thread::spawn(move || {
+ loop {
+ let mut buf = vec![0u8; 1024];
+ let size = reader.read(&mut buf[..]).expect("read failed");
+ buf.truncate(size);
+ sender.send(buf);
+ }
+ });
+
+ let mut writer = StdioBufWriter::new(writer);
+ assert_matches!(
+ writer,
+ StdioBufWriter::LineBuffered(_),
+ "StdioBufWriter should be line-buffered when created from a terminal"
+ );
+
+ writer.write_all(b"line 1\n").expect("write failed");
+ assert_eq!(receiver.recv().expect("recv failed"), b"line 1\n");
+
+ writer.write_all(b"line 2\nextra ").expect("write failed");
+ assert_eq!(receiver.recv().expect("recv failed"), b"line 2\n");
+
+ writer.write_all(b"line 3\n").expect("write failed");
+ assert_eq!(receiver.recv().expect("recv failed"), b"extra line 3\n");
+}
+
+#[test]
+fn stdiobufwriter_block_buffered() {
+ let (mut reader, mut writer) = io::pipe().expect("pipe() failed");
+
+ // Need to convert to an OwnedFd and then into a File because PipeReader/PipeWriter don't
+ // implement IsTerminal, but that is required by StdioBufWriter
+ let mut reader = File::from(OwnedFd::from(reader));
+ let mut writer = File::from(OwnedFd::from(writer));
+
+ assert!(!reader.is_terminal(), "file returned by pipe() cannot be a terminal");
+ assert!(!writer.is_terminal(), "file returned by pipe() cannot be a terminal");
+
+ let (sender, receiver) = sync_channel(64);
+
+ thread::spawn(move || {
+ loop {
+ let mut buf = vec![0u8; 1024];
+ let size = reader.read(&mut buf[..]).expect("read failed");
+ buf.truncate(size);
+ sender.send(buf);
+ }
+ });
+
+ let mut writer = StdioBufWriter::new(writer);
+ assert_matches!(
+ writer,
+ StdioBufWriter::BlockBuffered(_),
+ "StdioBufWriter should be block-buffered when created from a file that is not a terminal"
+ );
+
+ writer.write_all(b"line 1\n").expect("write failed");
+ writer.write_all(b"line 2\n").expect("write failed");
+ writer.write_all(b"line 3\n").expect("write failed");
+ assert_matches!(receiver.try_recv(), Err(TryRecvError::Empty));
+
+ writer.flush().expect("flush failed");
+ assert_eq!(receiver.recv().expect("recv failed"), b"line 1\nline 2\nline 3\n");
+}