Rollup merge of #78714 - m-ou-se:simplify-local-streams, r=KodrAus

Simplify output capturing

This is a sequence of incremental improvements to the unstable/internal `set_panic` and `set_print` mechanism used by the `test` crate:

1. Remove the `LocalOutput` trait and use `Arc<Mutex<dyn Write>>` instead of `Box<dyn LocalOutput>`. In practice, all implementations of `LocalOutput` were just `Arc<Mutex<..>>`. This simplifies some logic and removes all custom `Sink` implementations such as `library/test/src/helpers/sink.rs`. Also removes a layer of indirection, as the outermost `Box` is now gone. It also means that locking now happens per `write_fmt`, not per individual `write` within. (So `"{} {}\n"` now results in one `lock()`, not four or more.)

2. Since in all cases the `dyn Write`s were just `Vec<u8>`s, replace the type with `Arc<Mutex<Vec<u8>>>`. This simplifies things more, as error handling and flushing can be removed now. This also removes the hack needed in the default panic handler to make this work with `::realstd`, as (unlike `Write`) `Vec<u8>` is from `alloc`, not `std`.

3. Replace the `RefCell`s by regular `Cell`s. The `RefCell`s were mostly used as `mem::replace(&mut *cell.borrow_mut(), something)`, which is just `Cell::replace`. This removes an unecessary bookkeeping and makes the code a bit easier to read.

4. Merge `set_panic` and `set_print` into a single `set_output_capture`. Neither the test crate nor rustc (the only users of this feature) have a use for using these separately. Merging them simplifies things even more. This uses a new function name and feature name, to make it clearer this is internal and not supposed to be used by other crates.

Might be easier to review per commit.

compiler/rustc_interface/src/lib.rs

@@ -1,6 +1,6 @@
 #![feature(bool_to_option)]
 #![feature(box_syntax)]
-#![feature(set_stdio)]
+#![feature(internal_output_capture)]
 #![feature(nll)]
 #![feature(generator_trait)]
 #![feature(generators)]

compiler/rustc_interface/src/util.rs

@@ -25,7 +25,7 @@ use rustc_span::symbol::{sym, Symbol};
 use smallvec::SmallVec;
 use std::env;
 use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};
-use std::io::{self, Write};
+use std::io;
 use std::lazy::SyncOnceCell;
 use std::mem;
 use std::ops::DerefMut;
@@ -106,21 +106,6 @@ fn get_stack_size() -> Option<usize> {
     env::var_os("RUST_MIN_STACK").is_none().then_some(STACK_SIZE)
 }
 
-struct Sink(Arc<Mutex<Vec<u8>>>);
-impl Write for Sink {
-    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        Write::write(&mut *self.0.lock().unwrap(), data)
-    }
-    fn flush(&mut self) -> io::Result<()> {
-        Ok(())
-    }
-}
-impl io::LocalOutput for Sink {
-    fn clone_box(&self) -> Box<dyn io::LocalOutput> {
-        Box::new(Self(self.0.clone()))
-    }
-}
-
 /// Like a `thread::Builder::spawn` followed by a `join()`, but avoids the need
 /// for `'static` bounds.
 #[cfg(not(parallel_compiler))]
@@ -163,9 +148,7 @@ pub fn setup_callbacks_and_run_in_thread_pool_with_globals<F: FnOnce() -> R + Se
 
     let main_handler = move || {
         rustc_span::with_session_globals(edition, || {
-            if let Some(stderr) = stderr {
-                io::set_panic(Some(box Sink(stderr.clone())));
-            }
+            io::set_output_capture(stderr.clone());
             f()
         })
     };
@@ -203,9 +186,7 @@ pub fn setup_callbacks_and_run_in_thread_pool_with_globals<F: FnOnce() -> R + Se
             // on the new threads.
             let main_handler = move |thread: rayon::ThreadBuilder| {
                 rustc_span::SESSION_GLOBALS.set(session_globals, || {
-                    if let Some(stderr) = stderr {
-                        io::set_panic(Some(box Sink(stderr.clone())));
-                    }
+                    io::set_output_capture(stderr.clone());
                     thread.run()
                 })
             };

library/std/src/io/impls.rs

@@ -209,20 +209,6 @@ impl<B: BufRead + ?Sized> BufRead for Box<B> {
     }
 }
 
-// Used by panicking::default_hook
-#[cfg(test)]
-/// This impl is only used by printing logic, so any error returned is always
-/// of kind `Other`, and should be ignored.
-impl Write for dyn ::realstd::io::LocalOutput {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        (*self).write(buf).map_err(|_| ErrorKind::Other.into())
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        (*self).flush().map_err(|_| ErrorKind::Other.into())
-    }
-}
-
 // =============================================================================
 // In-memory buffer implementations
 

library/std/src/io/mod.rs

@@ -271,20 +271,18 @@ pub use self::copy::copy;
 pub use self::cursor::Cursor;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::error::{Error, ErrorKind, Result};
+#[unstable(feature = "internal_output_capture", issue = "none")]
+#[doc(no_inline, hidden)]
+pub use self::stdio::set_output_capture;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::stdio::{stderr, stdin, stdout, Stderr, Stdin, Stdout};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::stdio::{StderrLock, StdinLock, StdoutLock};
 #[unstable(feature = "print_internals", issue = "none")]
 pub use self::stdio::{_eprint, _print};
-#[unstable(feature = "libstd_io_internals", issue = "42788")]
-#[doc(no_inline, hidden)]
-pub use self::stdio::{set_panic, set_print, LocalOutput};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::util::{empty, repeat, sink, Empty, Repeat, Sink};
 
-pub(crate) use self::stdio::clone_io;
-
 mod buffered;
 pub(crate) mod copy;
 mod cursor;

library/std/src/io/stdio.rs

@@ -5,44 +5,38 @@ mod tests;
 
 use crate::io::prelude::*;
 
-use crate::cell::RefCell;
+use crate::cell::{Cell, RefCell};
 use crate::fmt;
 use crate::io::{self, BufReader, Initializer, IoSlice, IoSliceMut, LineWriter};
 use crate::lazy::SyncOnceCell;
 use crate::sync::atomic::{AtomicBool, Ordering};
-use crate::sync::{Mutex, MutexGuard};
+use crate::sync::{Arc, Mutex, MutexGuard};
 use crate::sys::stdio;
 use crate::sys_common;
 use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
-use crate::thread::LocalKey;
 
-thread_local! {
-    /// Used by the test crate to capture the output of the print! and println! macros.
-    static LOCAL_STDOUT: RefCell<Option<Box<dyn LocalOutput>>> = {
-        RefCell::new(None)
-    }
-}
+type LocalStream = Arc<Mutex<Vec<u8>>>;
 
 thread_local! {
-    /// Used by the test crate to capture the output of the eprint! and eprintln! macros, and panics.
-    static LOCAL_STDERR: RefCell<Option<Box<dyn LocalOutput>>> = {
-        RefCell::new(None)
+    /// Used by the test crate to capture the output of the print macros and panics.
+    static OUTPUT_CAPTURE: Cell<Option<LocalStream>> = {
+        Cell::new(None)
     }
 }
 
-/// Flag to indicate LOCAL_STDOUT and/or LOCAL_STDERR is used.
+/// Flag to indicate OUTPUT_CAPTURE is used.
 ///
-/// If both are None and were never set on any thread, this flag is set to
-/// false, and both LOCAL_STDOUT and LOCAL_STDOUT can be safely ignored on all
-/// threads, saving some time and memory registering an unused thread local.
+/// If it is None and was never set on any thread, this flag is set to false,
+/// and OUTPUT_CAPTURE can be safely ignored on all threads, saving some time
+/// and memory registering an unused thread local.
 ///
-/// Note about memory ordering: This contains information about whether two
-/// thread local variables might be in use. Although this is a global flag, the
+/// Note about memory ordering: This contains information about whether a
+/// thread local variable might be in use. Although this is a global flag, the
 /// memory ordering between threads does not matter: we only want this flag to
-/// have a consistent order between set_print/set_panic and print_to *within
+/// have a consistent order between set_output_capture and print_to *within
 /// the same thread*. Within the same thread, things always have a perfectly
 /// consistent order. So Ordering::Relaxed is fine.
-static LOCAL_STREAMS: AtomicBool = AtomicBool::new(false);
+static OUTPUT_CAPTURE_USED: AtomicBool = AtomicBool::new(false);
 
 /// A handle to a raw instance of the standard input stream of this process.
 ///
@@ -896,97 +890,24 @@ impl fmt::Debug for StderrLock<'_> {
     }
 }
 
-/// A writer than can be cloned to new threads.
-#[unstable(
-    feature = "set_stdio",
-    reason = "this trait may disappear completely or be replaced \
-                     with a more general mechanism",
-    issue = "none"
-)]
-#[doc(hidden)]
-pub trait LocalOutput: Write + Send {
-    fn clone_box(&self) -> Box<dyn LocalOutput>;
-}
-
-/// Resets the thread-local stderr handle to the specified writer
-///
-/// This will replace the current thread's stderr handle, returning the old
-/// handle. All future calls to `panic!` and friends will emit their output to
-/// this specified handle.
-///
-/// Note that this does not need to be called for all new threads; the default
-/// output handle is to the process's stderr stream.
-#[unstable(
-    feature = "set_stdio",
-    reason = "this function may disappear completely or be replaced \
-                     with a more general mechanism",
-    issue = "none"
-)]
-#[doc(hidden)]
-pub fn set_panic(sink: Option<Box<dyn LocalOutput>>) -> Option<Box<dyn LocalOutput>> {
-    use crate::mem;
-    if sink.is_none() && !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        // LOCAL_STDERR is definitely None since LOCAL_STREAMS is false.
-        return None;
-    }
-    let s = LOCAL_STDERR.with(move |slot| mem::replace(&mut *slot.borrow_mut(), sink)).and_then(
-        |mut s| {
-            let _ = s.flush();
-            Some(s)
-        },
-    );
-    LOCAL_STREAMS.store(true, Ordering::Relaxed);
-    s
-}
-
-/// Resets the thread-local stdout handle to the specified writer
-///
-/// This will replace the current thread's stdout handle, returning the old
-/// handle. All future calls to `print!` and friends will emit their output to
-/// this specified handle.
-///
-/// Note that this does not need to be called for all new threads; the default
-/// output handle is to the process's stdout stream.
+/// Sets the thread-local output capture buffer and returns the old one.
 #[unstable(
-    feature = "set_stdio",
-    reason = "this function may disappear completely or be replaced \
-                     with a more general mechanism",
+    feature = "internal_output_capture",
+    reason = "this function is meant for use in the test crate \
+        and may disappear in the future",
     issue = "none"
 )]
 #[doc(hidden)]
-pub fn set_print(sink: Option<Box<dyn LocalOutput>>) -> Option<Box<dyn LocalOutput>> {
-    use crate::mem;
-    if sink.is_none() && !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        // LOCAL_STDOUT is definitely None since LOCAL_STREAMS is false.
+pub fn set_output_capture(sink: Option<LocalStream>) -> Option<LocalStream> {
+    if sink.is_none() && !OUTPUT_CAPTURE_USED.load(Ordering::Relaxed) {
+        // OUTPUT_CAPTURE is definitely None since OUTPUT_CAPTURE_USED is false.
         return None;
     }
-    let s = LOCAL_STDOUT.with(move |slot| mem::replace(&mut *slot.borrow_mut(), sink)).and_then(
-        |mut s| {
-            let _ = s.flush();
-            Some(s)
-        },
-    );
-    LOCAL_STREAMS.store(true, Ordering::Relaxed);
-    s
-}
-
-pub(crate) fn clone_io() -> (Option<Box<dyn LocalOutput>>, Option<Box<dyn LocalOutput>>) {
-    // Don't waste time when LOCAL_{STDOUT,STDERR} are definitely None.
-    if !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        return (None, None);
-    }
-
-    LOCAL_STDOUT.with(|stdout| {
-        LOCAL_STDERR.with(|stderr| {
-            (
-                stdout.borrow().as_ref().map(|o| o.clone_box()),
-                stderr.borrow().as_ref().map(|o| o.clone_box()),
-            )
-        })
-    })
+    OUTPUT_CAPTURE_USED.store(true, Ordering::Relaxed);
+    OUTPUT_CAPTURE.with(move |slot| slot.replace(sink))
 }
 
-/// Write `args` to output stream `local_s` if possible, `global_s`
+/// Write `args` to the capture buffer if enabled and possible, or `global_s`
 /// otherwise. `label` identifies the stream in a panic message.
 ///
 /// This function is used to print error messages, so it takes extra
@@ -996,36 +917,26 @@ pub(crate) fn clone_io() -> (Option<Box<dyn LocalOutput>>, Option<Box<dyn LocalO
 /// thread, it will just fall back to the global stream.
 ///
 /// However, if the actual I/O causes an error, this function does panic.
-fn print_to<T>(
-    args: fmt::Arguments<'_>,
-    local_s: &'static LocalKey<RefCell<Option<Box<dyn LocalOutput>>>>,
-    global_s: fn() -> T,
-    label: &str,
-) where
+fn print_to<T>(args: fmt::Arguments<'_>, global_s: fn() -> T, label: &str)
+where
     T: Write,
 {
-    let result = LOCAL_STREAMS
-        .load(Ordering::Relaxed)
-        .then(|| {
-            local_s
-                .try_with(|s| {
-                    // Note that we completely remove a local sink to write to in case
-                    // our printing recursively panics/prints, so the recursive
-                    // panic/print goes to the global sink instead of our local sink.
-                    let prev = s.borrow_mut().take();
-                    if let Some(mut w) = prev {
-                        let result = w.write_fmt(args);
-                        *s.borrow_mut() = Some(w);
-                        return result;
-                    }
-                    global_s().write_fmt(args)
-                })
-                .ok()
-        })
-        .flatten()
-        .unwrap_or_else(|| global_s().write_fmt(args));
-
-    if let Err(e) = result {
+    if OUTPUT_CAPTURE_USED.load(Ordering::Relaxed)
+        && OUTPUT_CAPTURE.try_with(|s| {
+            // Note that we completely remove a local sink to write to in case
+            // our printing recursively panics/prints, so the recursive
+            // panic/print goes to the global sink instead of our local sink.
+            s.take().map(|w| {
+                let _ = w.lock().unwrap_or_else(|e| e.into_inner()).write_fmt(args);
+                s.set(Some(w));
+            })
+        }) == Ok(Some(()))
+    {
+        // Succesfully wrote to capture buffer.
+        return;
+    }
+
+    if let Err(e) = global_s().write_fmt(args) {
         panic!("failed printing to {}: {}", label, e);
     }
 }
@@ -1038,7 +949,7 @@ fn print_to<T>(
 #[doc(hidden)]
 #[cfg(not(test))]
 pub fn _print(args: fmt::Arguments<'_>) {
-    print_to(args, &LOCAL_STDOUT, stdout, "stdout");
+    print_to(args, stdout, "stdout");
 }
 
 #[unstable(
@@ -1049,7 +960,7 @@ pub fn _print(args: fmt::Arguments<'_>) {
 #[doc(hidden)]
 #[cfg(not(test))]
 pub fn _eprint(args: fmt::Arguments<'_>) {
-    print_to(args, &LOCAL_STDERR, stderr, "stderr");
+    print_to(args, stderr, "stderr");
 }
 
 #[cfg(test)]

library/std/src/lib.rs

@@ -207,7 +207,7 @@
 // std may use features in a platform-specific way
 #![allow(unused_features)]
 #![cfg_attr(not(bootstrap), feature(rustc_allow_const_fn_unstable))]
-#![cfg_attr(test, feature(print_internals, set_stdio, update_panic_count))]
+#![cfg_attr(test, feature(internal_output_capture, print_internals, update_panic_count))]
 #![cfg_attr(
     all(target_vendor = "fortanix", target_env = "sgx"),
     feature(slice_index_methods, coerce_unsized, sgx_platform)
@@ -298,6 +298,7 @@
 #![feature(raw)]
 #![feature(raw_ref_macros)]
 #![feature(ready_macro)]
+#![feature(refcell_take)]
 #![feature(rustc_attrs)]
 #![feature(rustc_private)]
 #![feature(shrink_to)]

library/std/src/panicking.rs

@@ -24,11 +24,11 @@ use crate::sys_common::{thread_info, util};
 use crate::thread;
 
 #[cfg(not(test))]
-use crate::io::set_panic;
+use crate::io::set_output_capture;
 // make sure to use the stderr output configured
 // by libtest in the real copy of std
 #[cfg(test)]
-use realstd::io::set_panic;
+use realstd::io::set_output_capture;
 
 // Binary interface to the panic runtime that the standard library depends on.
 //
@@ -218,11 +218,9 @@ fn default_hook(info: &PanicInfo<'_>) {
         }
     };
 
-    if let Some(mut local) = set_panic(None) {
-        // NB. In `cfg(test)` this uses the forwarding impl
-        // for `dyn ::realstd::io::LocalOutput`.
-        write(&mut local);
-        set_panic(Some(local));
+    if let Some(local) = set_output_capture(None) {
+        write(&mut *local.lock().unwrap_or_else(|e| e.into_inner()));
+        set_output_capture(Some(local));
     } else if let Some(mut out) = panic_output() {
         write(&mut out);
     }