Show backtrace on allocation failures when possible

And if an allocation while printing the backtrace fails, don't try to
print another backtrace as that will never succeed.

Author: bjorn3

library/std/src/alloc.rs

@@ -57,7 +57,7 @@
 #![stable(feature = "alloc_module", since = "1.28.0")]
 
 use core::ptr::NonNull;
-use core::sync::atomic::{Atomic, AtomicPtr, Ordering};
+use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
 use core::{hint, mem, ptr};
 
 #[stable(feature = "alloc_module", since = "1.28.0")]
@@ -287,7 +287,7 @@ unsafe impl Allocator for System {
     }
 }
 
-static HOOK: Atomic<*mut ()> = AtomicPtr::new(ptr::null_mut());
+static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 
 /// Registers a custom allocation error hook, replacing any that was previously registered.
 ///
@@ -344,7 +344,12 @@ pub fn take_alloc_error_hook() -> fn(Layout) {
     if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
 }
 
+#[optimize(size)]
 fn default_alloc_error_hook(layout: Layout) {
+    if cfg!(panic = "immediate-abort") {
+        return;
+    }
+
     unsafe extern "Rust" {
         // This symbol is emitted by rustc next to __rust_alloc_error_handler.
         // Its value depends on the -Zoom={panic,abort} compiler option.
@@ -354,28 +359,80 @@ fn default_alloc_error_hook(layout: Layout) {
 
     if unsafe { __rust_alloc_error_handler_should_panic_v2() != 0 } {
         panic!("memory allocation of {} bytes failed", layout.size());
+    }
+
+    // This is the default path taken on OOM, and the only path taken on stable with std.
+    // Crucially, it does *not* call any user-defined code, and therefore users do not have to
+    // worry about allocation failure causing reentrancy issues. That makes it different from
+    // the default `__rdl_alloc_error_handler` defined in alloc (i.e., the default alloc error
+    // handler that is  called when there is no `#[alloc_error_handler]`), which triggers a
+    // regular panic and thus can invoke a user-defined panic hook, executing arbitrary
+    // user-defined code.
+
+    static PREV_ALLOC_FAILURE: AtomicBool = AtomicBool::new(false);
+    if PREV_ALLOC_FAILURE.swap(true, Ordering::Relaxed) {
+        // Don't try to print a backtrace if a previous alloc error happened. This likely means
+        // there is not enough memory to print a backtrace, although it could also mean that two
+        // threads concurrently run out of memory.
+        rtprintpanic!(
+            "memory allocation of {} bytes failed\nskipping backtrace printing to avoid potential recursion\n",
+            layout.size()
+        );
+        return;
     } else {
-        // This is the default path taken on OOM, and the only path taken on stable with std.
-        // Crucially, it does *not* call any user-defined code, and therefore users do not have to
-        // worry about allocation failure causing reentrancy issues. That makes it different from
-        // the default `__rdl_alloc_error_handler` defined in alloc (i.e., the default alloc error
-        // handler that is  called when there is no `#[alloc_error_handler]`), which triggers a
-        // regular panic and thus can invoke a user-defined panic hook, executing arbitrary
-        // user-defined code.
         rtprintpanic!("memory allocation of {} bytes failed\n", layout.size());
     }
+
+    let Some(mut out) = crate::sys::stdio::panic_output() else {
+        return;
+    };
+
+    // Use a lock to prevent mixed output in multithreading context.
+    // Some platforms also require it when printing a backtrace, like `SymFromAddr` on Windows.
+    // Make sure to not take this lock until after checking PREV_ALLOC_FAILURE to avoid deadlocks
+    // when there is too little memory to print a backtrace.
+    let mut lock = crate::sys::backtrace::lock();
+
+    match crate::panic::get_backtrace_style() {
+        // SAFETY: we took out a lock just a second ago.
+        Some(crate::panic::BacktraceStyle::Short) => {
+            drop(lock.print(&mut out, crate::backtrace_rs::PrintFmt::Short))
+        }
+        Some(crate::panic::BacktraceStyle::Full) => {
+            drop(lock.print(&mut out, crate::backtrace_rs::PrintFmt::Full))
+        }
+        Some(crate::panic::BacktraceStyle::Off) => {
+            use crate::io::Write;
+            let _ = writeln!(
+                out,
+                "note: run with `RUST_BACKTRACE=1` environment variable to display a \
+                             backtrace"
+            );
+            if cfg!(miri) {
+                let _ = writeln!(
+                    out,
+                    "note: in Miri, you may have to set `MIRIFLAGS=-Zmiri-env-forward=RUST_BACKTRACE` \
+                                for the environment variable to have an effect"
+                );
+            }
+        }
+        // If backtraces aren't supported or are forced-off, do nothing.
+        None => {}
+    }
 }
 
 #[cfg(not(test))]
 #[doc(hidden)]
 #[alloc_error_handler]
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub fn rust_oom(layout: Layout) -> ! {
-    let hook = HOOK.load(Ordering::Acquire);
-    let hook: fn(Layout) =
-        if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
-    hook(layout);
-    crate::process::abort()
+    crate::sys::backtrace::__rust_end_short_backtrace(|| {
+        let hook = HOOK.load(Ordering::Acquire);
+        let hook: fn(Layout) =
+            if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };
+        hook(layout);
+        crate::process::abort()
+    })
 }
 
 #[cfg(not(test))]

src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr

@@ -1,11 +1,15 @@
 memory allocation of 4 bytes failed
+note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
+note: in Miri, you may have to set `MIRIFLAGS=-Zmiri-env-forward=RUST_BACKTRACE` for the environment variable to have an effect
 error: abnormal termination: the program aborted execution
   --> RUSTLIB/std/src/alloc.rs:LL:CC
    |
-LL |     crate::process::abort()
-   |     ^^^^^^^^^^^^^^^^^^^^^^^ abnormal termination occurred here
+LL |         crate::process::abort()
+   |         ^^^^^^^^^^^^^^^^^^^^^^^ abnormal termination occurred here
    |
    = note: BACKTRACE:
+   = note: inside closure at RUSTLIB/std/src/alloc.rs:LL:CC
+   = note: inside `std::sys::backtrace::__rust_end_short_backtrace::<{closure@std::alloc::rust_oom::{closure#0}}, !>` at RUSTLIB/std/src/sys/backtrace.rs:LL:CC
    = note: inside `std::alloc::rust_oom` at RUSTLIB/std/src/alloc.rs:LL:CC
    = note: inside `std::alloc::_::__rust_alloc_error_handler` at RUSTLIB/std/src/alloc.rs:LL:CC
    = note: inside `std::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC

tests/ui/alloc-error/alloc-error-backtrace.rs

@@ -0,0 +1,32 @@
+//@ run-pass
+//@ ignore-android FIXME #17520
+//@ needs-subprocess
+//@ ignore-openbsd no support for libbacktrace without filename
+//@ ignore-fuchsia Backtraces not symbolized
+//@ compile-flags:-g
+//@ compile-flags:-Cstrip=none
+
+use std::alloc::{Layout, handle_alloc_error};
+use std::process::Command;
+use std::{env, str};
+
+fn main() {
+    if env::args().len() > 1 {
+        handle_alloc_error(Layout::new::<[u8; 42]>())
+    }
+
+    let me = env::current_exe().unwrap();
+    let output = Command::new(&me).env("RUST_BACKTRACE", "1").arg("next").output().unwrap();
+    assert!(!output.status.success(), "{:?} is a success", output.status);
+
+    let mut stderr = str::from_utf8(&output.stderr).unwrap();
+
+    // When running inside QEMU user-mode emulation, there will be an extra message printed by QEMU
+    // in the stderr whenever a core dump happens. Remove it before the check.
+    stderr = stderr
+        .strip_suffix("qemu: uncaught target signal 6 (Aborted) - core dumped\n")
+        .unwrap_or(stderr);
+
+    assert!(stderr.contains("memory allocation of 42 bytes failed"), "{}", stderr);
+    assert!(stderr.contains("alloc_error_backtrace::main"), "{}", stderr);
+}

tests/ui/alloc-error/default-alloc-error-hook.rs

@@ -2,17 +2,16 @@
 //@ needs-subprocess
 
 use std::alloc::{Layout, handle_alloc_error};
-use std::env;
 use std::process::Command;
-use std::str;
+use std::{env, str};
 
 fn main() {
     if env::args().len() > 1 {
         handle_alloc_error(Layout::new::<[u8; 42]>())
     }
 
     let me = env::current_exe().unwrap();
-    let output = Command::new(&me).arg("next").output().unwrap();
+    let output = Command::new(&me).env("RUST_BACKTRACE", "0").arg("next").output().unwrap();
     assert!(!output.status.success(), "{:?} is a success", output.status);
 
     let mut stderr = str::from_utf8(&output.stderr).unwrap();
@@ -23,5 +22,9 @@ fn main() {
         .strip_suffix("qemu: uncaught target signal 6 (Aborted) - core dumped\n")
         .unwrap_or(stderr);
 
-    assert_eq!(stderr, "memory allocation of 42 bytes failed\n");
+    assert_eq!(
+        stderr,
+        "memory allocation of 42 bytes failed\n\
+        note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace\n"
+    );
 }