Fix #95126: io::cleanup() can panic in unusual circumstances

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);
+            }
         }
     }
 }

src/test/ui/issues/issue-95126.rs

@@ -0,0 +1,68 @@
+// run-pass
+// 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(); }
+}

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"))