Work with stdbuf environment variables

library/std/src/io/mod.rs

@@ -346,7 +346,7 @@ mod util;
 
 const DEFAULT_BUF_SIZE: usize = crate::sys::io::DEFAULT_BUF_SIZE;
 
-pub(crate) use stdio::cleanup;
+pub(crate) use stdio::{StderrRaw, StdinRaw, StdoutRaw, cleanup};
 
 struct Guard<'a> {
     buf: &'a mut Vec<u8>,

library/std/src/io/stdio.rs

@@ -4,17 +4,20 @@
 mod tests;
 
 use crate::cell::{Cell, RefCell};
-use crate::fmt;
+use crate::ffi::OsStr;
 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::str::FromStr as _;
 use crate::sync::atomic::{Atomic, AtomicBool, Ordering};
 use crate::sync::{Arc, Mutex, MutexGuard, OnceLock, ReentrantLock, ReentrantLockGuard};
 use crate::sys::stdio;
 use crate::thread::AccessError;
+use crate::{env, fmt};
 
 type LocalStream = Arc<Mutex<Vec<u8>>>;
 
@@ -43,19 +46,19 @@ static OUTPUT_CAPTURE_USED: Atomic<bool> = AtomicBool::new(false);
 ///
 /// This handle is not synchronized or buffered in any fashion. Constructed via
 /// the `std::io::stdio::stdin_raw` function.
-struct StdinRaw(stdio::Stdin);
+pub(crate) struct StdinRaw(stdio::Stdin);
 
 /// A handle to a raw instance of the standard output stream of this process.
 ///
 /// This handle is not synchronized or buffered in any fashion. Constructed via
 /// the `std::io::stdio::stdout_raw` function.
-struct StdoutRaw(stdio::Stdout);
+pub(crate) struct StdoutRaw(stdio::Stdout);
 
 /// A handle to a raw instance of the standard output stream of this process.
 ///
 /// This handle is not synchronized or buffered in any fashion. Constructed via
 /// the `std::io::stdio::stderr_raw` function.
-struct StderrRaw(stdio::Stderr);
+pub(crate) struct StderrRaw(stdio::Stderr);
 
 /// Constructs a new raw handle to the standard input of this process.
 ///
@@ -576,6 +579,88 @@ impl fmt::Debug for StdinLock<'_> {
     }
 }
 
+/// A buffered writer for stdout and stderr.
+///
+/// This writer may be either [line-buffered](LineWriter),
+/// [block-buffered](BufWriter), or unbuffered.
+#[derive(Debug)]
+enum StdioBufWriter<W: Write> {
+    LineBuffered(LineWriter<W>),
+    BlockBuffered(BufWriter<W>),
+    Unbuffered(W),
+}
+
+impl<W: Write> StdioBufWriter<W> {
+    /// Wraps a writer using the most appropriate buffering method.
+    fn from_env_value(stderr: bool, w: W, value: Option<&OsStr>) -> Self {
+        if let Some(value) = value {
+            if value == "L" {
+                return StdioBufWriter::LineBuffered(LineWriter::new(w));
+            }
+            if let Some(size) = value.to_str().and_then(|v| usize::from_str(v).ok()) {
+                if size == 0 {
+                    return StdioBufWriter::Unbuffered(w);
+                } else {
+                    return StdioBufWriter::BlockBuffered(BufWriter::with_capacity(size, w));
+                }
+            }
+        }
+        Self::default_buffering(stderr, w)
+    }
+    fn default_buffering(stderr: bool, w: W) -> Self {
+        if stderr {
+            StdioBufWriter::Unbuffered(w)
+        } else {
+            StdioBufWriter::LineBuffered(LineWriter::new(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),
+            Self::Unbuffered(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),
+            Self::Unbuffered(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(),
+            Self::Unbuffered(w) => w.is_write_vectored(),
+        }
+    }
+    fn flush(&mut self) -> io::Result<()> {
+        match self {
+            Self::LineBuffered(w) => w.flush(),
+            Self::BlockBuffered(w) => w.flush(),
+            Self::Unbuffered(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),
+            Self::Unbuffered(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),
+            Self::Unbuffered(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.
 ///
@@ -715,30 +799,46 @@ static STDOUT: OnceLock<ReentrantLock<RefCell<LineWriter<StdoutRaw>>>> = OnceLoc
 #[cfg_attr(not(test), rustc_diagnostic_item = "io_stdout")]
 pub fn stdout() -> Stdout {
     Stdout {
-        inner: STDOUT
-            .get_or_init(|| ReentrantLock::new(RefCell::new(LineWriter::new(stdout_raw())))),
+        inner: STDOUT.get_or_init(|| {
+            ReentrantLock::new(RefCell::new(StdioBufWriter::from_env_value(
+                false,
+                stdout_raw(),
+                env::var_os("_STDBUF_O").as_deref(),
+            )))
+        }),
     }
 }
 
 // Flush the data and disable buffering during shutdown
 // by replacing the line writer by one with zero
 // buffering capacity.
 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())))
-    });
-
-    if !initialized {
-        // The buffer was previously initialized, overwrite it here.
-        // We use try_lock() instead of lock(), because someone
-        // 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());
+    fn cleanup<W: Write, F: Fn() -> W>(
+        global: &'static OnceLock<ReentrantLock<RefCell<StdioBufWriter<W>>>>,
+        init_writer: F,
+    ) {
+        let mut initialized = false;
+        let global = global.get_or_init(|| {
+            initialized = true;
+            ReentrantLock::new(RefCell::new(StdioBufWriter::Unbuffered(init_writer())))
+        });
+
+        if !initialized {
+            // The buffer was previously initialized, overwrite it here.
+            // We use try_lock() instead of lock(), because someone
+            // might have leaked a StdoutLock, which would
+            // otherwise cause a deadlock here.
+            if let Some(lock) = global.try_lock() {
+                let mut lock = lock.borrow_mut();
+                if !matches!(*lock, StdioBufWriter::Unbuffered(_)) {
+                    *lock = StdioBufWriter::Unbuffered(init_writer());
+                }
+            }
         }
     }
+
+    cleanup(&STDERR, stderr_raw);
+    cleanup(&STDOUT, stdout_raw);
 }
 
 impl Stdout {
@@ -890,7 +990,9 @@ impl fmt::Debug for StdoutLock<'_> {
 /// standard library or via raw Windows API calls, will fail.
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Stderr {
-    inner: &'static ReentrantLock<RefCell<StderrRaw>>,
+    // FIXME: if this is not line buffered it should flush-on-panic or some
+    //        form of flush-on-abort.
+    inner: &'static ReentrantLock<RefCell<StdioBufWriter<StderrRaw>>>,
 }
 
 /// A locked reference to the [`Stderr`] handle.
@@ -912,9 +1014,11 @@ pub struct Stderr {
 #[must_use = "if unused stderr will immediately unlock"]
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct StderrLock<'a> {
-    inner: ReentrantLockGuard<'a, RefCell<StderrRaw>>,
+    inner: ReentrantLockGuard<'a, RefCell<StdioBufWriter<StderrRaw>>>,
 }
 
+static STDERR: OnceLock<ReentrantLock<RefCell<StdioBufWriter<StderrRaw>>>> = OnceLock::new();
+
 /// Constructs a new handle to the standard error of the current process.
 ///
 /// This handle is not buffered.
@@ -963,13 +1067,15 @@ pub struct StderrLock<'a> {
 #[stable(feature = "rust1", since = "1.0.0")]
 #[cfg_attr(not(test), rustc_diagnostic_item = "io_stderr")]
 pub fn stderr() -> Stderr {
-    // Note that unlike `stdout()` we don't use `at_exit` here to register a
-    // destructor. Stderr is not buffered, so there's no need to run a
-    // destructor for flushing the buffer
-    static INSTANCE: ReentrantLock<RefCell<StderrRaw>> =
-        ReentrantLock::new(RefCell::new(stderr_raw()));
-
-    Stderr { inner: &INSTANCE }
+    Stderr {
+        inner: STDERR.get_or_init(|| {
+            ReentrantLock::new(RefCell::new(StdioBufWriter::from_env_value(
+                true,
+                stderr_raw(),
+                env::var_os("_STDBUF_E").as_deref(),
+            )))
+        }),
+    }
 }
 
 impl Stderr {
@@ -1262,7 +1368,18 @@ macro_rules! impl_is_terminal {
     )*}
 }
 
-impl_is_terminal!(File, Stdin, StdinLock<'_>, Stdout, StdoutLock<'_>, Stderr, StderrLock<'_>);
+impl_is_terminal!(
+    File,
+    Stdin,
+    StdinLock<'_>,
+    StdinRaw,
+    Stdout,
+    StdoutLock<'_>,
+    StdoutRaw,
+    Stderr,
+    StderrLock<'_>,
+    StderrRaw,
+);
 
 #[unstable(
     feature = "print_internals",

library/std/src/os/fd/owned.rs

@@ -489,6 +489,13 @@ impl<'a> AsFd for io::StdinLock<'a> {
     }
 }
 
+impl AsFd for io::StdinRaw {
+    #[inline]
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        unsafe { BorrowedFd::borrow_raw(0) }
+    }
+}
+
 #[stable(feature = "io_safety", since = "1.63.0")]
 impl AsFd for io::Stdout {
     #[inline]
@@ -506,6 +513,13 @@ impl<'a> AsFd for io::StdoutLock<'a> {
     }
 }
 
+impl AsFd for io::StdoutRaw {
+    #[inline]
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        unsafe { BorrowedFd::borrow_raw(1) }
+    }
+}
+
 #[stable(feature = "io_safety", since = "1.63.0")]
 impl AsFd for io::Stderr {
     #[inline]
@@ -523,6 +537,13 @@ impl<'a> AsFd for io::StderrLock<'a> {
     }
 }
 
+impl AsFd for io::StderrRaw {
+    #[inline]
+    fn as_fd(&self) -> BorrowedFd<'_> {
+        unsafe { BorrowedFd::borrow_raw(2) }
+    }
+}
+
 #[stable(feature = "anonymous_pipe", since = "1.87.0")]
 #[cfg(not(target_os = "trusty"))]
 impl AsFd for io::PipeReader {

library/std/src/os/windows/io/handle.rs

@@ -562,6 +562,13 @@ impl<'a> AsHandle for crate::io::StdinLock<'a> {
     }
 }
 
+impl AsHandle for crate::io::StdinRaw {
+    #[inline]
+    fn as_handle(&self) -> BorrowedHandle<'_> {
+        unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) }
+    }
+}
+
 #[stable(feature = "io_safety", since = "1.63.0")]
 impl AsHandle for crate::io::Stdout {
     #[inline]
@@ -578,6 +585,13 @@ impl<'a> AsHandle for crate::io::StdoutLock<'a> {
     }
 }
 
+impl AsHandle for crate::io::StdoutRaw {
+    #[inline]
+    fn as_handle(&self) -> BorrowedHandle<'_> {
+        unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) }
+    }
+}
+
 #[stable(feature = "io_safety", since = "1.63.0")]
 impl AsHandle for crate::io::Stderr {
     #[inline]
@@ -594,6 +608,13 @@ impl<'a> AsHandle for crate::io::StderrLock<'a> {
     }
 }
 
+impl AsHandle for crate::io::StderrRaw {
+    #[inline]
+    fn as_handle(&self) -> BorrowedHandle<'_> {
+        unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) }
+    }
+}
+
 #[stable(feature = "io_safety", since = "1.63.0")]
 impl AsHandle for crate::process::ChildStdin {
     #[inline]