2 files changed, 41 insertions, 23 deletions

diff --git a/library/std/src/io/stdio.rs b/library/std/src/io/stdio.rs
index 61ccc6f13c8..814e0dfda54 100644
--- a/library/std/src/io/stdio.rs
+++ b/library/std/src/io/stdio.rs
@@ -9,21 +9,24 @@ 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;
 use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
 use crate::thread::LocalKey;
 
+static LOCAL_STREAMS: AtomicBool = AtomicBool::new(false);
+
 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)
     }
@@ -890,10 +893,14 @@ 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)
-    })
+    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::Release);
+    s
 }
 
 /// Resets the thread-local stdout handle to the specified writer
@@ -913,10 +920,14 @@ 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)
-    })
+    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::Release);
+    s
 }
 
 /// Write `args` to output stream `local_s` if possible, `global_s`
@@ -937,20 +948,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::Acquire)
+        .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 d03428dd082..410860b4a5d 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)]