From a69dca8e6ef7d368a74c0fbee2780cc31b6f37c0 Mon Sep 17 00:00:00 2001 From: Nathan West Date: Fri, 6 Aug 2021 01:00:03 -0400 Subject: [PATCH] Implementation of Switchable Buffering for Stdout --- library/std/src/io/buffered/mod.rs | 4 + library/std/src/io/buffered/switchwriter.rs | 117 ++++++++++++++++++++ library/std/src/io/mod.rs | 2 + library/std/src/io/stdio.rs | 41 +++++-- library/std/src/sys/hermit/stdio.rs | 7 +- library/std/src/sys/sgx/stdio.rs | 6 +- library/std/src/sys/unix/stdio.rs | 8 +- library/std/src/sys/unsupported/stdio.rs | 6 +- library/std/src/sys/wasi/stdio.rs | 6 +- library/std/src/sys/windows/stdio.rs | 6 +- 10 files changed, 185 insertions(+), 18 deletions(-) create mode 100644 library/std/src/io/buffered/switchwriter.rs diff --git a/library/std/src/io/buffered/mod.rs b/library/std/src/io/buffered/mod.rs index 100dab1e2493c..56fa74a9c088a 100644 --- a/library/std/src/io/buffered/mod.rs +++ b/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. diff --git a/library/std/src/io/buffered/switchwriter.rs b/library/std/src/io/buffered/switchwriter.rs new file mode 100644 index 0000000000000..c9bad0f7d2b57 --- /dev/null +++ b/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 { + buffer: BufWriter, + mode: BufferMode, +} + +impl SwitchWriter { + #[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 Write for SwitchWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + 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 { + 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) + } + } + } +} diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs index 94812e3fe3b2c..ecbffef01c5ce 100644 --- a/library/std/src/io/mod.rs +++ b/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")] diff --git a/library/std/src/io/stdio.rs b/library/std/src/io/stdio.rs index ae16015e35a1a..37b0027294fd7 100644 --- a/library/std/src/io/stdio.rs +++ b/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>>>, + inner: Pin<&'static ReentrantMutex>>>, } /// 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>>, + inner: ReentrantMutexGuard<'a, RefCell>>, } -static STDOUT: SyncOnceCell>>> = SyncOnceCell::new(); +static STDOUT: SyncOnceCell>>> = SyncOnceCell::new(); /// Constructs a new handle to the standard output of the current process. /// @@ -605,7 +602,12 @@ static STDOUT: SyncOnceCell>>> = 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 { @@ -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")] diff --git a/library/std/src/sys/hermit/stdio.rs b/library/std/src/sys/hermit/stdio.rs index 514de1df6f9c3..44820faf6ef60 100644 --- a/library/std/src/sys/hermit/stdio.rs +++ b/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 { Some(Stderr::new()) } + +pub fn default_stdout_buffer_mode() -> BufferMode { + BufferMode::Line +} diff --git a/library/std/src/sys/sgx/stdio.rs b/library/std/src/sys/sgx/stdio.rs index 2e680e740fde3..4c40e85ade962 100644 --- a/library/std/src/sys/sgx/stdio.rs +++ b/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 { 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))] diff --git a/library/std/src/sys/unix/stdio.rs b/library/std/src/sys/unix/stdio.rs index e4d83ba0ffd13..895e26f4bd026 100644 --- a/library/std/src/sys/unix/stdio.rs +++ b/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 +} diff --git a/library/std/src/sys/unsupported/stdio.rs b/library/std/src/sys/unsupported/stdio.rs index b5e3f5be9885b..38fccb27e9c59 100644 --- a/library/std/src/sys/unsupported/stdio.rs +++ b/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> { None } + +pub fn default_stdout_buffer_mode() -> BufferMode { + BufferMode::Line +} diff --git a/library/std/src/sys/wasi/stdio.rs b/library/std/src/sys/wasi/stdio.rs index 4cc0e4ed5a45a..d76b06eb93c25 100644 --- a/library/std/src/sys/wasi/stdio.rs +++ b/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 { Some(Stderr::new()) } + +pub fn default_stdout_buffer_mode() -> BufferMode { + BufferMode::Line +} diff --git a/library/std/src/sys/windows/stdio.rs b/library/std/src/sys/windows/stdio.rs index a001d6b985823..1ff1b9048ca62 100644 --- a/library/std/src/sys/windows/stdio.rs +++ b/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 { Some(Stderr::new()) } + +pub fn default_stdout_buffer_mode() -> BufferMode { + BufferMode::Line +}