From 6f31f20997c9202e30b362640a528386c3d2d573 Mon Sep 17 00:00:00 2001
From: Jack Wrenn <jack@wrenn.fyi>
Date: Mon, 21 Mar 2022 14:25:39 -0400
Subject: [PATCH 1/2] regression test for #95126: io::cleanup() can panic in
 unusual circumstances

---
 src/test/ui/issues/issue-95126.rs | 69 +++++++++++++++++++++++++++++++
 src/tools/tidy/src/ui_tests.rs    |  2 +-
 2 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 src/test/ui/issues/issue-95126.rs

diff --git a/src/test/ui/issues/issue-95126.rs b/src/test/ui/issues/issue-95126.rs
new file mode 100644
index 0000000000000..49f4fdc95fb6f
--- /dev/null
+++ b/src/test/ui/issues/issue-95126.rs
@@ -0,0 +1,69 @@
+// run-fail
+// known-bug
+// no-prefer-dynamic
+
+// regression test for #95126: io::cleanup() can panic in unusual circumstances
+
+#[global_allocator]
+static ALLOCATOR: allocator::TracingSystemAllocator = allocator::TracingSystemAllocator;
+
+fn main() {
+    let _ = std::io::stdout();
+    allocator::enable_tracing();
+    // panic occurs after `main()`
+}
+
+// a global allocator that prints on `alloc` and `dealloc`
+mod allocator {
+    use std::{
+        alloc::{GlobalAlloc, Layout, System},
+        cell::{RefCell, RefMut},
+        panic::catch_unwind,
+    };
+
+    pub struct TracingSystemAllocator;
+
+    unsafe impl GlobalAlloc for TracingSystemAllocator {
+        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+            let ptr = System.alloc(layout);
+
+            let _ = catch_unwind(|| {
+                maybe_with_guard(|trace_allocations| {
+                    if *trace_allocations {
+                        println!("alloc({:?}) = {}", layout, ptr as usize);
+                    }
+                })
+            });
+
+            ptr
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+            System.dealloc(ptr, layout);
+
+            let _ = catch_unwind(|| {
+                maybe_with_guard(|trace_allocations| {
+                    if *trace_allocations {
+                        println!("dealloc({}, {:?})", ptr as usize, layout);
+                    }
+                })
+            });
+        }
+    }
+
+    pub fn enable_tracing() {
+        maybe_with_guard(|mut trace| *trace = true)
+    }
+
+    // maybe run `f`, if a unique, mutable reference to `TRACE_ALLOCATOR` can be
+    // acquired.
+    fn maybe_with_guard<F>(f: F)
+    where
+        F: for<'a> FnOnce(RefMut<'a, bool>),
+    {
+        let _ = TRACE_ALLOCATOR.try_with(|guard| guard.try_borrow_mut().map(f));
+    }
+
+    // used to prevent infinitely recursive tracing
+    thread_local! { static TRACE_ALLOCATOR: RefCell<bool> = RefCell::default(); }
+}
diff --git a/src/tools/tidy/src/ui_tests.rs b/src/tools/tidy/src/ui_tests.rs
index 7b932b867f240..004172f0138e9 100644
--- a/src/tools/tidy/src/ui_tests.rs
+++ b/src/tools/tidy/src/ui_tests.rs
@@ -8,7 +8,7 @@ use std::path::Path;
 const ENTRY_LIMIT: usize = 1000;
 // FIXME: The following limits should be reduced eventually.
 const ROOT_ENTRY_LIMIT: usize = 984;
-const ISSUES_ENTRY_LIMIT: usize = 2310;
+const ISSUES_ENTRY_LIMIT: usize = 2311;
 
 fn check_entries(path: &Path, bad: &mut bool) {
     let dirs = walkdir::WalkDir::new(&path.join("test/ui"))

From 2b3a3a53668984c6aaa89dc8c06a2c69c205bd7d Mon Sep 17 00:00:00 2001
From: Jack Wrenn <jack@wrenn.fyi>
Date: Mon, 21 Mar 2022 14:46:24 -0400
Subject: [PATCH 2/2] fix #95126: io::cleanup() can panic in unusual
 circumstances

---
 library/std/src/io/stdio.rs       | 13 ++++++++++++-
 src/test/ui/issues/issue-95126.rs |  3 +--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/library/std/src/io/stdio.rs b/library/std/src/io/stdio.rs
index ac6d41e13b009..b7e431ae6e7fc 100644
--- a/library/std/src/io/stdio.rs
+++ b/library/std/src/io/stdio.rs
@@ -621,7 +621,18 @@ pub fn cleanup() {
         // 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 refmut) = lock.try_borrow_mut() {
+                // construct an unbuffered LineWriter
+                let unbuffered = LineWriter::with_capacity(0, stdout_raw());
+                // replace the current LineWriter with the unbuffered one
+                let mut buffered = crate::mem::replace(&mut *refmut, unbuffered);
+                // flush the old, buffered LineWriter
+                let _ = buffered.flush();
+                // leak the old, buffered LineWriter to avoid deallocations this
+                // ensures that writes to stdout *during* cleanup (e.g., writes
+                // emitted by a custom global allocator) won't deadlock (#95126)
+                crate::mem::forget(buffered);
+            }
         }
     }
 }
diff --git a/src/test/ui/issues/issue-95126.rs b/src/test/ui/issues/issue-95126.rs
index 49f4fdc95fb6f..28274c7ca8c1f 100644
--- a/src/test/ui/issues/issue-95126.rs
+++ b/src/test/ui/issues/issue-95126.rs
@@ -1,5 +1,4 @@
-// run-fail
-// known-bug
+// run-pass
 // no-prefer-dynamic
 
 // regression test for #95126: io::cleanup() can panic in unusual circumstances