From 07fd17f7014de6dc209c7dc4de159a2a5acea173 Mon Sep 17 00:00:00 2001
From: Mara Bos <m-ou.se@m-ou.se>
Date: Sun, 27 Sep 2020 15:19:42 +0200
Subject: [PATCH] Only use LOCAL_{STDOUT,STDERR} when set_{print/panic} is
 used.

The thread local LOCAL_STDOUT and LOCAL_STDERR are only used by the test
crate to capture output from tests when running them in the same process
in differen threads. However, every program will check these variables
on every print, even outside of testing.

This involves allocating a thread local key, and registering a thread
local destructor. This can be somewhat expensive.

This change keeps a global flag (LOCAL_STREAMS) which will be set to
true when either of these local streams is used. (So, effectively only
in test and benchmark runs.) When this flag is off, these thread locals
are not even looked at and therefore will not be initialized on the
first output on every thread, which also means no thread local
destructors will be registered.
---
 library/std/src/io/stdio.rs | 63 +++++++++++++++++++++++--------------
 library/std/src/lib.rs      |  1 +
 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)]