Work with stdbuf environment variables

Check `_STDBUF_O` and `_STDBUF_E` environment variables for setting up
stdout and stderr. This makes Rust programs compatible with coreutils's
(and `uutils`'s) `stdbuf` program. It incidentally also solves
uutils/coreutils#7967.

This lays ground work for solving #60673 because it lets the user of a
Rust program decide the default buffering, offering the opportunity to
change the current default, letting stdout buffer blockwise if it's not
connected to a terminal, because the user can now override the behavior
using [`stdbuf`](https://linux.die.net/man/1/stdbuf) if it's needed.

Author: tbu-

library/std/src/io/stdio.rs

@@ -4,18 +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, 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>>>;
 
@@ -579,32 +581,38 @@ 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.
+/// 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 + IsTerminal> StdioBufWriter<W> {
+impl<W: Write> 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))
+    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)
     }
-}
-
-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))
+    fn default_buffering(stderr: bool, w: W) -> Self {
+        if stderr {
+            StdioBufWriter::Unbuffered(w)
+        } else {
+            StdioBufWriter::LineBuffered(LineWriter::new(w))
+        }
     }
 }
 
@@ -613,36 +621,42 @@ impl<W: Write> Write for StdioBufWriter<W> {
         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),
         }
     }
 }
@@ -785,30 +799,46 @@ static STDOUT: OnceLock<ReentrantLock<RefCell<StdioBufWriter<StdoutRaw>>>> = Onc
 #[cfg_attr(not(test), rustc_diagnostic_item = "io_stdout")]
 pub fn stdout() -> Stdout {
     Stdout {
-        inner: STDOUT
-            .get_or_init(|| ReentrantLock::new(RefCell::new(StdioBufWriter::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(StdioBufWriter::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() = StdioBufWriter::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 {
@@ -960,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.
@@ -982,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.
@@ -1033,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 {

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

@@ -1,8 +1,6 @@
 use super::*;
-use crate::assert_matches::assert_matches;
-use crate::os::fd::{FromRawFd, OwnedFd};
 use crate::panic::{RefUnwindSafe, UnwindSafe};
-use crate::sync::mpsc::{TryRecvError, sync_channel};
+use crate::sync::mpsc::sync_channel;
 use crate::thread;
 
 #[test]
@@ -166,95 +164,3 @@ 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");
-}