Auto merge of #109507 - Amanieu:panic-oom-payload, r=davidtwco

Report allocation errors as panics

OOM is now reported as a panic but with a custom payload type (`AllocErrorPanicPayload`) which holds the layout that was passed to `handle_alloc_error`.

This should be review one commit at a time:
- The first commit adds `AllocErrorPanicPayload` and changes allocation errors to always be reported as panics.
- The second commit removes `#[alloc_error_handler]` and the `alloc_error_hook` API.

ACP: rust-lang/libs-team#192

Closes #51540
Closes #51245

Author: bors

alloc/Cargo.toml

@@ -35,3 +35,6 @@ compiler-builtins-mem = ['compiler_builtins/mem']
 compiler-builtins-c = ["compiler_builtins/c"]
 compiler-builtins-no-asm = ["compiler_builtins/no-asm"]
 compiler-builtins-mangled-names = ["compiler_builtins/mangled-names"]
+
+# Make panics and failed asserts immediately abort without formatting any message
+panic_immediate_abort = ["core/panic_immediate_abort"]

alloc/src/alloc.rs

@@ -14,6 +14,11 @@ use core::ptr::{self, NonNull};
 #[doc(inline)]
 pub use core::alloc::*;
 
+#[cfg(not(no_global_oom_handling))]
+use core::any::Any;
+#[cfg(not(no_global_oom_handling))]
+use core::panic::BoxMeUp;
+
 #[cfg(test)]
 mod tests;
 
@@ -343,28 +348,84 @@ pub(crate) unsafe fn box_free<T: ?Sized, A: Allocator>(ptr: Unique<T>, alloc: A)
     }
 }
 
-// # Allocation error handler
+/// Payload passed to the panic handler when `handle_alloc_error` is called.
+#[unstable(feature = "panic_oom_payload", issue = "none")]
+#[derive(Debug)]
+pub struct AllocErrorPanicPayload {
+    layout: Layout,
+}
+
+impl AllocErrorPanicPayload {
+    /// Internal function for the standard library to clone a payload.
+    #[unstable(feature = "std_internals", issue = "none")]
+    #[doc(hidden)]
+    pub fn internal_clone(&self) -> Self {
+        AllocErrorPanicPayload { layout: self.layout }
+    }
 
+    /// Returns the [`Layout`] of the allocation attempt that caused the error.
+    #[unstable(feature = "panic_oom_payload", issue = "none")]
+    pub fn layout(&self) -> Layout {
+        self.layout
+    }
+}
+
+#[unstable(feature = "std_internals", issue = "none")]
 #[cfg(not(no_global_oom_handling))]
-extern "Rust" {
-    // This is the magic symbol to call the global alloc error handler. rustc generates
-    // it to call `__rg_oom` if there is a `#[alloc_error_handler]`, or to call the
-    // default implementations below (`__rdl_oom`) otherwise.
-    fn __rust_alloc_error_handler(size: usize, align: usize) -> !;
+unsafe impl BoxMeUp for AllocErrorPanicPayload {
+    fn take_box(&mut self) -> *mut (dyn Any + Send) {
+        use crate::boxed::Box;
+        Box::into_raw(Box::new(self.internal_clone()))
+    }
+
+    fn get(&mut self) -> &(dyn Any + Send) {
+        self
+    }
+}
+
+// # Allocation error handler
+
+#[cfg(all(not(no_global_oom_handling), not(test)))]
+fn rust_oom(layout: Layout) -> ! {
+    if cfg!(feature = "panic_immediate_abort") {
+        core::intrinsics::abort()
+    }
+
+    extern "Rust" {
+        // NOTE This function never crosses the FFI boundary; it's a Rust-to-Rust call
+        // that gets resolved to the `#[panic_handler]` function.
+        #[lang = "panic_impl"]
+        fn panic_impl(pi: &core::panic::PanicInfo<'_>) -> !;
+
+        // This symbol is emitted by rustc .
+        // Its value depends on the -Zoom={unwind,abort} compiler option.
+        static __rust_alloc_error_handler_should_panic: u8;
+    }
+
+    // Hack to work around issues with the lifetime of Arguments.
+    match format_args!("memory allocation of {} bytes failed", layout.size()) {
+        fmt => {
+            // Create a PanicInfo with a custom payload for the panic handler.
+            let can_unwind = unsafe { __rust_alloc_error_handler_should_panic != 0 };
+            let mut pi = core::panic::PanicInfo::internal_constructor(
+                Some(&fmt),
+                core::panic::Location::caller(),
+                can_unwind,
+            );
+            let payload = AllocErrorPanicPayload { layout };
+            pi.set_payload(&payload);
+
+            // SAFETY: `panic_impl` is defined in safe Rust code and thus is safe to call.
+            unsafe { panic_impl(&pi) }
+        }
+    }
 }
 
 /// Abort on memory allocation error or failure.
 ///
 /// Callers of memory allocation APIs wishing to abort computation
 /// in response to an allocation error are encouraged to call this function,
 /// rather than directly invoking `panic!` or similar.
-///
-/// The default behavior of this function is to print a message to standard error
-/// and abort the process.
-/// It can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
-///
-/// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
-/// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
 #[stable(feature = "global_alloc", since = "1.28.0")]
 #[rustc_const_unstable(feature = "const_alloc_error", issue = "92523")]
 #[cfg(all(not(no_global_oom_handling), not(test)))]
@@ -375,9 +436,7 @@ pub const fn handle_alloc_error(layout: Layout) -> ! {
     }
 
     fn rt_error(layout: Layout) -> ! {
-        unsafe {
-            __rust_alloc_error_handler(layout.size(), layout.align());
-        }
+        rust_oom(layout);
     }
 
     unsafe { core::intrinsics::const_eval_select((layout,), ct_error, rt_error) }
@@ -387,6 +446,7 @@ pub const fn handle_alloc_error(layout: Layout) -> ! {
 #[cfg(all(not(no_global_oom_handling), test))]
 pub use std::alloc::handle_alloc_error;
 
+#[cfg(bootstrap)]
 #[cfg(all(not(no_global_oom_handling), not(test)))]
 #[doc(hidden)]
 #[allow(unused_attributes)]
@@ -398,7 +458,7 @@ pub mod __alloc_error_handler {
     pub unsafe fn __rdl_oom(size: usize, _align: usize) -> ! {
         extern "Rust" {
             // This symbol is emitted by rustc next to __rust_alloc_error_handler.
-            // Its value depends on the -Zoom={panic,abort} compiler option.
+            // Its value depends on the -Zoom={unwind,abort} compiler option.
             static __rust_alloc_error_handler_should_panic: u8;
         }
 

alloc/src/lib.rs

@@ -135,6 +135,7 @@
 #![feature(maybe_uninit_slice)]
 #![feature(maybe_uninit_uninit_array)]
 #![feature(maybe_uninit_uninit_array_transpose)]
+#![feature(panic_internals)]
 #![feature(pattern)]
 #![feature(pointer_byte_offsets)]
 #![feature(provide_any)]

core/src/macros/mod.rs

@@ -1531,16 +1531,6 @@ pub(crate) mod builtin {
         /* compiler built-in */
     }
 
-    /// Attribute macro applied to a function to register it as a handler for allocation failure.
-    ///
-    /// See also [`std::alloc::handle_alloc_error`](../../../std/alloc/fn.handle_alloc_error.html).
-    #[unstable(feature = "alloc_error_handler", issue = "51540")]
-    #[allow_internal_unstable(rustc_attrs)]
-    #[rustc_builtin_macro]
-    pub macro alloc_error_handler($item:item) {
-        /* compiler built-in */
-    }
-
     /// Keeps the item it's applied to if the passed path is accessible, and removes it otherwise.
     #[unstable(
         feature = "cfg_accessible",

core/src/prelude/v1.rs

@@ -76,9 +76,7 @@ pub use crate::macros::builtin::{RustcDecodable, RustcEncodable};
 // Do not `doc(no_inline)` so that they become doc items on their own
 // (no public module for them to be re-exported from).
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
-pub use crate::macros::builtin::{
-    alloc_error_handler, bench, derive, global_allocator, test, test_case,
-};
+pub use crate::macros::builtin::{bench, derive, global_allocator, test, test_case};
 
 #[unstable(feature = "derive_const", issue = "none")]
 pub use crate::macros::builtin::derive_const;

std/Cargo.toml

@@ -67,7 +67,7 @@ llvm-libunwind = ["unwind/llvm-libunwind"]
 system-llvm-libunwind = ["unwind/system-llvm-libunwind"]
 
 # Make panics and failed asserts immediately abort without formatting any message
-panic_immediate_abort = ["core/panic_immediate_abort"]
+panic_immediate_abort = ["alloc/panic_immediate_abort"]
 
 # Enable std_detect default features for stdarch/crates/std_detect:
 # https://github.com/rust-lang/stdarch/blob/master/crates/std_detect/Cargo.toml

std/src/alloc.rs

@@ -57,9 +57,8 @@
 #![stable(feature = "alloc_module", since = "1.28.0")]
 
 use core::intrinsics;
+use core::ptr;
 use core::ptr::NonNull;
-use core::sync::atomic::{AtomicPtr, Ordering};
-use core::{mem, ptr};
 
 #[stable(feature = "alloc_module", since = "1.28.0")]
 #[doc(inline)]
@@ -286,76 +285,6 @@ unsafe impl Allocator for System {
     }
 }
 
-static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
-
-/// Registers a custom allocation error hook, replacing any that was previously registered.
-///
-/// The allocation error hook is invoked when an infallible memory allocation fails, before
-/// the runtime aborts. The default hook prints a message to standard error,
-/// but this behavior can be customized with the [`set_alloc_error_hook`] and
-/// [`take_alloc_error_hook`] functions.
-///
-/// The hook is provided with a `Layout` struct which contains information
-/// about the allocation that failed.
-///
-/// The allocation error hook is a global resource.
-///
-/// # Examples
-///
-/// ```
-/// #![feature(alloc_error_hook)]
-///
-/// use std::alloc::{Layout, set_alloc_error_hook};
-///
-/// fn custom_alloc_error_hook(layout: Layout) {
-///    panic!("memory allocation of {} bytes failed", layout.size());
-/// }
-///
-/// set_alloc_error_hook(custom_alloc_error_hook);
-/// ```
-#[unstable(feature = "alloc_error_hook", issue = "51245")]
-pub fn set_alloc_error_hook(hook: fn(Layout)) {
-    HOOK.store(hook as *mut (), Ordering::SeqCst);
-}
-
-/// Unregisters the current allocation error hook, returning it.
-///
-/// *See also the function [`set_alloc_error_hook`].*
-///
-/// If no custom hook is registered, the default hook will be returned.
-#[unstable(feature = "alloc_error_hook", issue = "51245")]
-pub fn take_alloc_error_hook() -> fn(Layout) {
-    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
-    if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } }
-}
-
-fn default_alloc_error_hook(layout: Layout) {
-    extern "Rust" {
-        // This symbol is emitted by rustc next to __rust_alloc_error_handler.
-        // Its value depends on the -Zoom={panic,abort} compiler option.
-        static __rust_alloc_error_handler_should_panic: u8;
-    }
-
-    #[allow(unused_unsafe)]
-    if unsafe { __rust_alloc_error_handler_should_panic != 0 } {
-        panic!("memory allocation of {} bytes failed", layout.size());
-    } else {
-        rtprintpanic!("memory allocation of {} bytes failed\n", layout.size());
-    }
-}
-
-#[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::SeqCst);
-    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))]
 #[doc(hidden)]
 #[allow(unused_attributes)]

std/src/lib.rs

@@ -236,7 +236,6 @@
 //
 // Language features:
 // tidy-alphabetical-start
-#![feature(alloc_error_handler)]
 #![feature(allocator_internals)]
 #![feature(allow_internal_unsafe)]
 #![feature(allow_internal_unstable)]
@@ -319,6 +318,7 @@
 #![feature(get_mut_unchecked)]
 #![feature(map_try_insert)]
 #![feature(new_uninit)]
+#![feature(panic_oom_payload)]
 #![feature(slice_concat_trait)]
 #![feature(thin_box)]
 #![feature(try_reserve_kind)]

std/src/panicking.rs

@@ -245,19 +245,24 @@ fn default_hook(info: &PanicInfo<'_>) {
 
     // The current implementation always returns `Some`.
     let location = info.location().unwrap();
-
-    let msg = match info.payload().downcast_ref::<&'static str>() {
-        Some(s) => *s,
-        None => match info.payload().downcast_ref::<String>() {
-            Some(s) => &s[..],
-            None => "Box<dyn Any>",
-        },
-    };
     let thread = thread_info::current_thread();
     let name = thread.as_ref().and_then(|t| t.name()).unwrap_or("<unnamed>");
 
     let write = |err: &mut dyn crate::io::Write| {
-        let _ = writeln!(err, "thread '{name}' panicked at '{msg}', {location}");
+        // Use the panic message directly if available, otherwise take it from
+        // the payload.
+        if let Some(msg) = info.message() {
+            let _ = writeln!(err, "thread '{name}' panicked at '{msg}', {location}");
+        } else {
+            let msg = if let Some(s) = info.payload().downcast_ref::<&'static str>() {
+                *s
+            } else if let Some(s) = info.payload().downcast_ref::<String>() {
+                &s[..]
+            } else {
+                "Box<dyn Any>"
+            };
+            let _ = writeln!(err, "thread '{name}' panicked at '{msg}', {location}");
+        }
 
         static FIRST_PANIC: AtomicBool = AtomicBool::new(true);
 
@@ -524,6 +529,8 @@ pub fn panicking() -> bool {
 #[cfg(not(test))]
 #[panic_handler]
 pub fn begin_panic_handler(info: &PanicInfo<'_>) -> ! {
+    use alloc::alloc::AllocErrorPanicPayload;
+
     struct PanicPayload<'a> {
         inner: &'a fmt::Arguments<'a>,
         string: Option<String>,
@@ -550,8 +557,7 @@ pub fn begin_panic_handler(info: &PanicInfo<'_>) -> ! {
     unsafe impl<'a> BoxMeUp for PanicPayload<'a> {
         fn take_box(&mut self) -> *mut (dyn Any + Send) {
             // We do two allocations here, unfortunately. But (a) they're required with the current
-            // scheme, and (b) we don't handle panic + OOM properly anyway (see comment in
-            // begin_panic below).
+            // scheme, and (b) OOM uses its own separate payload type which doesn't allocate.
             let contents = mem::take(self.fill());
             Box::into_raw(Box::new(contents))
         }
@@ -576,7 +582,14 @@ pub fn begin_panic_handler(info: &PanicInfo<'_>) -> ! {
     let loc = info.location().unwrap(); // The current implementation always returns Some
     let msg = info.message().unwrap(); // The current implementation always returns Some
     crate::sys_common::backtrace::__rust_end_short_backtrace(move || {
-        if let Some(msg) = msg.as_str() {
+        if let Some(payload) = info.payload().downcast_ref::<AllocErrorPanicPayload>() {
+            rust_panic_with_hook(
+                &mut payload.internal_clone(),
+                info.message(),
+                loc,
+                info.can_unwind(),
+            );
+        } else if let Some(msg) = msg.as_str() {
             rust_panic_with_hook(&mut StrPanicPayload(msg), info.message(), loc, info.can_unwind());
         } else {
             rust_panic_with_hook(
@@ -623,11 +636,7 @@ pub const fn begin_panic<M: Any + Send>(msg: M) -> ! {
 
     unsafe impl<A: Send + 'static> BoxMeUp for PanicPayload<A> {
         fn take_box(&mut self) -> *mut (dyn Any + Send) {
-            // Note that this should be the only allocation performed in this code path. Currently
-            // this means that panic!() on OOM will invoke this code path, but then again we're not
-            // really ready for panic on OOM anyway. If we do start doing this, then we should
-            // propagate this allocation to be performed in the parent of this thread instead of the
-            // thread that's panicking.
+            // Note that this should be the only allocation performed in this code path.
             let data = match self.inner.take() {
                 Some(a) => Box::new(a) as Box<dyn Any + Send>,
                 None => process::abort(),

std/src/prelude/v1.rs

@@ -60,9 +60,7 @@ pub use core::prelude::v1::{RustcDecodable, RustcEncodable};
 // Do not `doc(no_inline)` so that they become doc items on their own
 // (no public module for them to be re-exported from).
 #[stable(feature = "builtin_macro_prelude", since = "1.38.0")]
-pub use core::prelude::v1::{
-    alloc_error_handler, bench, derive, global_allocator, test, test_case,
-};
+pub use core::prelude::v1::{bench, derive, global_allocator, test, test_case};
 
 #[unstable(feature = "derive_const", issue = "none")]
 pub use core::prelude::v1::derive_const;