2 files changed, 61 insertions, 23 deletions

diff --git a/library/std/src/io/stdio.rs b/library/std/src/io/stdio.rs
index 61ccc6f13c8..36b49401591 100644
--- a/library/std/src/io/stdio.rs
+++ b/library/std/src/io/stdio.rs
@@ -9,6 +9,7 @@ use crate::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::sys::stdio;
 use crate::sys_common;
@@ -16,19 +17,33 @@ use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
 use crate::thread::LocalKey;
 
 thread_local! {
-    /// Stdout used by print! and println! macros
+    /// Used by the test crate to capture the output of the print! and println! macros.
     static LOCAL_STDOUT: RefCell<Option<Box<dyn Write + Send>>> = {
         RefCell::new(None)
     }
 }
 
 thread_local! {
-    /// Stderr used by eprint! and eprintln! macros, and panics
+    /// Used by the test crate to capture the output of the eprint! and eprintln! macros, and panics.
     static LOCAL_STDERR: RefCell<Option<Box<dyn Write + Send>>> = {
         RefCell::new(None)
     }
 }
 
+/// Flag to indicate LOCAL_STDOUT and/or LOCAL_STDERR 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.
+///
+/// Note about memory ordering: This contains information about whether two
+/// thread local variables 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
+/// 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);
+
 /// A handle to a raw instance of the standard input stream of this process.
 ///
 /// This handle is not synchronized or buffered in any fashion. Constructed via
@@ -890,10 +905,18 @@ impl fmt::Debug for StderrLock<'_> {
 #[doc(hidden)]
 pub fn set_panic(sink: Option<Box<dyn Write + Send>>) -> Option<Box<dyn Write + Send>> {
     use crate::mem;
-    LOCAL_STDERR.with(move |slot| mem::replace(&mut *slot.borrow_mut(), sink)).and_then(|mut s| {
-        let _ = s.flush();
-        Some(s)
-    })
+    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
@@ -913,10 +936,18 @@ pub fn set_panic(sink: Option<Box<dyn Write + Send>>) -> Option<Box<dyn Write +
 #[doc(hidden)]
 pub fn set_print(sink: Option<Box<dyn Write + Send>>) -> Option<Box<dyn Write + Send>> {
     use crate::mem;
-    LOCAL_STDOUT.with(move |slot| mem::replace(&mut *slot.borrow_mut(), sink)).and_then(|mut s| {
-        let _ = s.flush();
-        Some(s)
-    })
+    if sink.is_none() && !LOCAL_STREAMS.load(Ordering::Relaxed) {
+        // LOCAL_STDOUT is definitely None since LOCAL_STREAMS 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
 }
 
 /// Write `args` to output stream `local_s` if possible, `global_s`
@@ -937,20 +968,26 @@ fn print_to<T>(
 ) where
     T: Write,
 {
-    let result = 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)
+    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()
         })
-        .unwrap_or_else(|_| global_s().write_fmt(args));
+        .flatten()
+        .unwrap_or_else(|| global_s().write_fmt(args));
 
     if let Err(e) = result {
         panic!("failed printing to {}: {}", label, e);
diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs
index e343eef9112..b2bd5f4da50 100644
--- a/library/std/src/lib.rs
+++ b/library/std/src/lib.rs
@@ -226,6 +226,7 @@
 #![feature(asm)]
 #![feature(associated_type_bounds)]
 #![feature(atomic_mut_ptr)]
+#![feature(bool_to_option)]
 #![feature(box_syntax)]
 #![feature(c_variadic)]
 #![feature(cfg_accessible)]