Implementation of Switchable Buffering for Stdout

Author: Lucretiel

library/std/src/io/buffered/mod.rs

@@ -4,6 +4,7 @@ mod bufreader;
 mod bufwriter;
 mod linewriter;
 mod linewritershim;
+mod switchwriter;
 
 #[cfg(test)]
 mod tests;
@@ -19,6 +20,9 @@ use linewritershim::LineWriterShim;
 #[stable(feature = "bufwriter_into_parts", since = "1.56.0")]
 pub use bufwriter::WriterPanicked;
 
+#[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+pub use switchwriter::{BufferMode, SwitchWriter};
+
 /// An error returned by [`BufWriter::into_inner`] which combines an error that
 /// happened while writing out the buffer, and the buffered writer object
 /// which may be used to recover from the condition.

library/std/src/io/buffered/switchwriter.rs

@@ -0,0 +1,117 @@
+use crate::fmt::Arguments;
+use crate::io::{self, buffered::LineWriterShim, BufWriter, IoSlice, Write};
+/// Different buffering modes a writer can use
+#[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum BufferMode {
+    /// Immediate: forward writes directly to the underlying writer. In some
+    /// cases, a writer may buffer temporarily to reduce the number of
+    /// underlying writes (for instance, when processing a formatted write!(),
+    /// which makes several tiny writes), but even in this case the formatted
+    /// buffer will always be forwarded immediately.
+    Immediate,
+
+    /// Block buffering: buffer writes until the buffer is full, then forward
+    /// to the underlying writer
+    Block,
+
+    /// Line buffering: same as block buffering, except that it immediately
+    /// forwards the buffered content when it encounters a newline.
+    Line,
+}
+
+/// Wraps a writer and provides a switchable buffering mode for its output
+#[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+#[derive(Debug)]
+pub struct SwitchWriter<W: Write> {
+    buffer: BufWriter<W>,
+    mode: BufferMode,
+}
+
+impl<W: Write> SwitchWriter<W> {
+    #[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+    pub fn new(writer: W, mode: BufferMode) -> Self {
+        Self { buffer: BufWriter::new(writer), mode }
+    }
+
+    // Don't forget to add with_capacity if this type ever becomes public
+
+    #[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+    pub fn mode(&self) -> BufferMode {
+        self.mode
+    }
+
+    /// Set the buffering mode. This will not attempt any io; it only changes
+    /// the mode used for subsequent writes.
+    #[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+    pub fn set_mode(&mut self, mode: BufferMode) {
+        self.mode = mode
+    }
+}
+
+/// Shared logic for io methods that need to switch over the buffering mode
+macro_rules! use_correct_writer {
+    ($this:ident, |$writer:ident| $usage:expr) => {
+        match $this.mode {
+            BufferMode::Immediate => {
+                $this.buffer.flush_buf()?;
+                let $writer = $this.buffer.get_mut();
+                $usage
+            }
+            BufferMode::Block => {
+                let $writer = &mut $this.buffer;
+                $usage
+            }
+            BufferMode::Line => {
+                let mut $writer = LineWriterShim::new(&mut $this.buffer);
+                $usage
+            }
+        }
+    };
+}
+
+impl<W: Write> Write for SwitchWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        use_correct_writer!(self, |writer| writer.write(buf))
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.buffer.flush()
+    }
+
+    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        use_correct_writer!(self, |writer| writer.write_vectored(bufs))
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        self.buffer.is_write_vectored()
+    }
+
+    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        use_correct_writer!(self, |writer| writer.write_all(buf))
+    }
+
+    fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
+        use_correct_writer!(self, |writer| writer.write_all_vectored(bufs))
+    }
+
+    fn write_fmt(&mut self, fmt: Arguments<'_>) -> io::Result<()> {
+        match self.mode {
+            BufferMode::Immediate => {
+                // write_fmt is usually going to be very numerous tiny writes
+                // from the constituent fmt calls, so even though we're in
+                // unbuffered mode we still collect it to the buffer so that
+                // we can flush it in a single write.
+                self.buffer.flush_buf()?;
+                self.buffer.write_fmt(fmt)?;
+                self.buffer.flush_buf()
+            }
+            BufferMode::Block => self.buffer.write_fmt(fmt),
+            BufferMode::Line => {
+                let mut shim = LineWriterShim::new(&mut self.buffer);
+                shim.write_fmt(fmt)
+            }
+        }
+    }
+}

library/std/src/io/mod.rs

@@ -260,6 +260,8 @@ use crate::str;
 use crate::sys;
 use crate::sys_common::memchr;
 
+#[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+pub use self::buffered::BufferMode;
 #[stable(feature = "bufwriter_into_parts", since = "1.56.0")]
 pub use self::buffered::WriterPanicked;
 #[unstable(feature = "internal_output_capture", issue = "none")]

library/std/src/io/stdio.rs

@@ -7,7 +7,7 @@ use crate::io::prelude::*;
 
 use crate::cell::{Cell, RefCell};
 use crate::fmt;
-use crate::io::{self, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
+use crate::io::{self, buffered::SwitchWriter, BufReader, BufferMode, IoSlice, IoSliceMut, Lines};
 use crate::lazy::SyncOnceCell;
 use crate::pin::Pin;
 use crate::sync::atomic::{AtomicBool, Ordering};
@@ -524,10 +524,7 @@ 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: Pin<&'static ReentrantMutex<RefCell<LineWriter<StdoutRaw>>>>,
+    inner: Pin<&'static ReentrantMutex<RefCell<SwitchWriter<StdoutRaw>>>>,
 }
 
 /// A locked reference to the [`Stdout`] handle.
@@ -549,10 +546,10 @@ pub struct Stdout {
 #[must_use = "if unused stdout will immediately unlock"]
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct StdoutLock<'a> {
-    inner: ReentrantMutexGuard<'a, RefCell<LineWriter<StdoutRaw>>>,
+    inner: ReentrantMutexGuard<'a, RefCell<SwitchWriter<StdoutRaw>>>,
 }
 
-static STDOUT: SyncOnceCell<ReentrantMutex<RefCell<LineWriter<StdoutRaw>>>> = SyncOnceCell::new();
+static STDOUT: SyncOnceCell<ReentrantMutex<RefCell<SwitchWriter<StdoutRaw>>>> = SyncOnceCell::new();
 
 /// Constructs a new handle to the standard output of the current process.
 ///
@@ -605,7 +602,12 @@ static STDOUT: SyncOnceCell<ReentrantMutex<RefCell<LineWriter<StdoutRaw>>>> = Sy
 pub fn stdout() -> Stdout {
     Stdout {
         inner: Pin::static_ref(&STDOUT).get_or_init_pin(
-            || unsafe { ReentrantMutex::new(RefCell::new(LineWriter::new(stdout_raw()))) },
+            || unsafe {
+                ReentrantMutex::new(RefCell::new(SwitchWriter::new(
+                    stdout_raw(),
+                    stdio::default_stdout_buffer_mode(),
+                )))
+            },
             |mutex| unsafe { mutex.init() },
         ),
     }
@@ -614,13 +616,14 @@ pub fn stdout() -> Stdout {
 pub fn cleanup() {
     if let Some(instance) = STDOUT.get() {
         // Flush the data and disable buffering during shutdown
-        // by replacing the line writer by one with zero
-        // buffering capacity.
         // 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) = Pin::static_ref(instance).try_lock() {
-            *lock.borrow_mut() = LineWriter::with_capacity(0, stdout_raw());
+            if let Ok(mut instance) = lock.try_borrow_mut() {
+                let _ = instance.flush();
+                instance.set_mode(BufferMode::Immediate);
+            }
         }
     }
 }
@@ -713,6 +716,20 @@ impl Write for &Stdout {
     }
 }
 
+impl StdoutLock<'_> {
+    /// Get the current global stdout buffering mode
+    #[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+    pub fn buffer_mode(&self) -> BufferMode {
+        self.inner.borrow().mode()
+    }
+
+    /// Set the current global stdout buffering mode
+    #[unstable(feature = "stdout_switchable_buffering", issue = "78515")]
+    pub fn set_buffer_mode(&mut self, mode: BufferMode) {
+        self.inner.borrow_mut().set_mode(mode)
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl Write for StdoutLock<'_> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
@@ -734,6 +751,8 @@ impl Write for StdoutLock<'_> {
     fn write_all_vectored(&mut self, bufs: &mut [IoSlice<'_>]) -> io::Result<()> {
         self.inner.borrow_mut().write_all_vectored(bufs)
     }
+    // Don't specialize write_fmt because we need to be sure we don't
+    // reentrantly call borrow_mut
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]

library/std/src/sys/hermit/stdio.rs

@@ -1,5 +1,4 @@
-use crate::io;
-use crate::io::{IoSlice, IoSliceMut};
+use crate::io::{self, BufferMode, IoSlice, IoSliceMut};
 use crate::sys::hermit::abi;
 
 pub struct Stdin;
@@ -118,3 +117,7 @@ pub fn is_ebadf(_err: &io::Error) -> bool {
 pub fn panic_output() -> Option<impl io::Write> {
     Some(Stderr::new())
 }
+
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}

library/std/src/sys/sgx/stdio.rs

@@ -1,6 +1,6 @@
 use fortanix_sgx_abi as abi;
 
-use crate::io;
+use crate::io::{self, BufferMode};
 #[cfg(not(test))]
 use crate::slice;
 #[cfg(not(test))]
@@ -73,6 +73,10 @@ pub fn panic_output() -> Option<impl io::Write> {
     super::abi::panic::SgxPanicOutput::new()
 }
 
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}
+
 // This function is needed by libunwind. The symbol is named in pre-link args
 // for the target specification, so keep that in sync.
 #[cfg(not(test))]

library/std/src/sys/unix/stdio.rs

@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BufferMode, IoSlice, IoSliceMut};
 use crate::mem::ManuallyDrop;
 use crate::os::unix::io::{AsFd, BorrowedFd, FromRawFd};
 use crate::sys::fd::FileDesc;
@@ -139,3 +139,9 @@ impl<'a> AsFd for io::StderrLock<'a> {
         unsafe { BorrowedFd::borrow_raw(libc::STDERR_FILENO) }
     }
 }
+
+/// Get the default stdout buffermode. In the future, this should be determined
+/// via `isatty` or some similar check on stdout
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}

library/std/src/sys/unsupported/stdio.rs

@@ -1,4 +1,4 @@
-use crate::io;
+use crate::io::{self, BufferMode};
 
 pub struct Stdin;
 pub struct Stdout;
@@ -57,3 +57,7 @@ pub fn is_ebadf(_err: &io::Error) -> bool {
 pub fn panic_output() -> Option<Vec<u8>> {
     None
 }
+
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}

library/std/src/sys/wasi/stdio.rs

@@ -1,7 +1,7 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 
 use super::fd::WasiFd;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BufferMode, IoSlice, IoSliceMut};
 use crate::mem::ManuallyDrop;
 use crate::os::raw;
 use crate::os::wasi::io::{AsRawFd, FromRawFd};
@@ -110,3 +110,7 @@ pub fn is_ebadf(err: &io::Error) -> bool {
 pub fn panic_output() -> Option<impl io::Write> {
     Some(Stderr::new())
 }
+
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}

library/std/src/sys/windows/stdio.rs

@@ -2,7 +2,7 @@
 
 use crate::char::decode_utf16;
 use crate::cmp;
-use crate::io;
+use crate::io::{self, BufferMode};
 use crate::os::windows::io::{FromRawHandle, IntoRawHandle};
 use crate::ptr;
 use crate::str;
@@ -420,3 +420,7 @@ pub fn is_ebadf(err: &io::Error) -> bool {
 pub fn panic_output() -> Option<impl io::Write> {
     Some(Stderr::new())
 }
+
+pub fn default_stdout_buffer_mode() -> BufferMode {
+    BufferMode::Line
+}