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

Author: andreacorbellini

library/std/src/io/stdio.rs

@@ -8,8 +8,10 @@ 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::os::fd::{AsFd, BorrowedFd};
 use crate::panic::{RefUnwindSafe, UnwindSafe};
 use crate::sync::atomic::{Atomic, AtomicBool, Ordering};
 use crate::sync::{Arc, Mutex, MutexGuard, OnceLock, ReentrantLock, ReentrantLockGuard};
@@ -168,6 +170,13 @@ impl Write for StdoutRaw {
     }
 }
 
+impl AsFd for StdoutRaw {
+    #[inline]
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        unsafe { 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()))
@@ -200,6 +209,13 @@ impl Write for StderrRaw {
     }
 }
 
+impl AsFd for StderrRaw {
+    #[inline]
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        unsafe { BorrowedFd::borrow_raw(2) }
+    }
+}
+
 fn handle_ebadf<T>(r: io::Result<T>, default: impl FnOnce() -> io::Result<T>) -> io::Result<T> {
     match r {
         Err(ref e) if stdio::is_ebadf(e) => default(),
@@ -576,6 +592,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 +692,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 +723,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 +801,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 +812,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 +821,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 +1347,17 @@ 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<'_>,
+    StderrRaw
+);
 
 #[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,106 @@ where
         [Start1, Acquire1, Start2, Release1, Acquire2, Release2, Acquire1, Release1]
     );
 }
+
+#[cfg(unix)]
+unsafe extern "C" {
+    fn openpty(
+        amaster: *mut crate::ffi::c_int,
+        aslave: *mut crate::ffi::c_int,
+        name: *mut crate::ffi::c_char,
+        termp: *const crate::ffi::c_void,
+        winsize: *const crate::ffi::c_void,
+    ) -> crate::ffi::c_int;
+}
+
+#[test]
+#[cfg(unix)]
+fn stdiobufwriter_line_buffered() {
+    let mut fd1 = -1;
+    let mut fd2 = -1;
+
+    if unsafe {
+        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");
+}