Split set_alloc_error_hook

library/alloc/src/alloc.rs

@@ -517,21 +517,21 @@ unsafe extern "Rust" {
 /// in response to an allocation error are encouraged to call this function,
 /// rather than directly invoking [`panic!`] or similar.
 ///
-/// This function is guaranteed to diverge (not return normally with a value), but depending on
-/// global configuration, it may either panic (resulting in unwinding or aborting as per
-/// configuration for all panics), or abort the process (with no unwinding).
+/// Callers may assume that this function always aborts execution. However, implementors that do not
+/// rely on this invariant are encouraged to document this in order to be compatible with users of
+/// [`set_alloc_error_hook_unwinding`].
 ///
 /// The default behavior is:
 ///
 ///  * If the binary links against `std` (typically the case), then
 ///   print a message to standard error and abort the process.
 ///   This behavior can be replaced with [`set_alloc_error_hook`] and [`take_alloc_error_hook`].
-///   Future versions of Rust may panic by default instead.
 ///
 /// * If the binary does not link against `std` (all of its crates are marked
 ///   [`#![no_std]`][no_std]), then call [`panic!`] with a message.
 ///   [The panic handler] applies as to any panic.
 ///
+/// [`set_alloc_error_hook_unwinding`]: ../../std/alloc/fn.set_alloc_error_hook_unwinding.html
 /// [`set_alloc_error_hook`]: ../../std/alloc/fn.set_alloc_error_hook.html
 /// [`take_alloc_error_hook`]: ../../std/alloc/fn.take_alloc_error_hook.html
 /// [The panic handler]: https://doc.rust-lang.org/reference/runtime.html#the-panic_handler-attribute

library/alloctests/tests/c_str_alloc_error.rs

@@ -14,7 +14,7 @@
 #![cfg(not(all(miri, windows)))]
 #![feature(alloc_error_hook)]
 
-use std::alloc::{GlobalAlloc, Layout, System, set_alloc_error_hook};
+use std::alloc::{GlobalAlloc, Layout, System, set_alloc_error_hook_unwinding};
 use std::ffi::CString;
 use std::panic::{AssertUnwindSafe, catch_unwind};
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -52,7 +52,9 @@ static ALLOC: OneShotFailingAlloc = OneShotFailingAlloc;
 #[test]
 #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
 fn clone_into_alloc_failure_leaves_target_valid() {
-    set_alloc_error_hook(|_| panic!("alloc error"));
+    unsafe {
+        set_alloc_error_hook_unwinding(|_| panic!("alloc error"));
+    }
 
     let src = CString::new("a fairly long value").unwrap();
     let mut target = CString::new("x").unwrap();

library/alloctests/tests/vec_deque_alloc_error.rs

@@ -1,6 +1,6 @@
 #![feature(alloc_error_hook, allocator_api)]
 
-use std::alloc::{AllocError, Allocator, Layout, System, set_alloc_error_hook};
+use std::alloc::{AllocError, Allocator, Layout, System, set_alloc_error_hook_unwinding};
 use std::collections::VecDeque;
 use std::panic::{AssertUnwindSafe, catch_unwind};
 use std::ptr::NonNull;
@@ -36,7 +36,9 @@ fn test_shrink_to_unwind() {
         }
     }
 
-    set_alloc_error_hook(|_| panic!("alloc error"));
+    unsafe {
+        set_alloc_error_hook_unwinding(|_| panic!("alloc error"));
+    }
 
     let mut v = VecDeque::with_capacity_in(15, BadAlloc);
     v.push_back(1);

library/std/src/alloc.rs

@@ -296,6 +296,8 @@ unsafe impl Allocator for System {
 unsafe impl GlobalAllocator for System {}
 
 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
+/// Do we abort on unwind from the hook?
+static ABORT: AtomicBool = AtomicBool::new(true);
 
 /// Registers a custom allocation error hook, replacing any that was previously registered.
 ///
@@ -310,8 +312,9 @@ static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 /// The hook function is provided with a [`Layout`] struct which contains information
 /// about the allocation that failed.
 ///
-/// The hook function may choose to panic or abort; in the event that it returns normally, this
-/// will cause an immediate abort.
+/// Regardless of whether the hook aborts, unwinds, or returns, the process will always abort
+/// immediately after the hook is executed; an unwind will never propagate out of
+/// `handle_alloc_error`.
 ///
 /// Since [`take_alloc_error_hook`] is a safe function that allows retrieving the hook, the hook
 /// function must be _sound_ to call even if no memory allocations were attempted.
@@ -338,6 +341,29 @@ static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
 /// ```
 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 pub fn set_alloc_error_hook(hook: fn(Layout)) {
+    ABORT.store(true, Ordering::Release);
+    HOOK.store(hook as *mut (), Ordering::Release);
+}
+
+/// Registers a custom allocation error hook, replacing any that was previously registered.
+///
+/// Unlike [`set_alloc_error_hook`], unwinds from the hook may propagate.
+///
+/// # Safety
+///
+/// Unwinding from the allocation error hook is not unsafe per se; however, callers of
+/// [`handle_alloc_error`] may assume that that method will never unwind for safety and/or
+/// correctness. Note that this [also includes][example-1] the [standard library][example-2]
+/// at times. Thus, callers must ensure that only code which explicitly promises *not* to
+/// rely on this guarantee of `handle_alloc_error` may be called between setting an unwinding
+/// hook and unsetting it via [`take_alloc_error_hook`]. This requirement applies across threads,
+/// as the hooks are global.
+///
+/// [example-1]: https://github.com/rust-lang/rust/issues/157203
+/// [example-2]: https://github.com/rust-lang/rust/issues/156490
+#[unstable(feature = "alloc_error_hook", issue = "51245")]
+pub unsafe fn set_alloc_error_hook_unwinding(hook: fn(Layout)) {
+    ABORT.store(false, Ordering::Release);
     HOOK.store(hook as *mut (), Ordering::Release);
 }
 
@@ -423,6 +449,14 @@ fn default_alloc_error_hook(layout: Layout) {
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub fn rust_oom(layout: Layout) -> ! {
     crate::sys::backtrace::__rust_end_short_backtrace(|| {
+        use core::mem::DropGuard;
+
+        let guard = DropGuard::new((), |_| crate::process::abort());
+        let abort = ABORT.load(Ordering::Acquire);
+        // Unwinds permitted, don't arm the guard.
+        if !abort {
+            DropGuard::dismiss(guard);
+        }
         let hook = HOOK.load(Ordering::Acquire);
         let hook: fn(Layout) =
             if hook.is_null() { default_alloc_error_hook } else { unsafe { mem::transmute(hook) } };

src/tools/miri/tests/panic/alloc_error_handler_hook.rs

@@ -12,7 +12,9 @@ impl Drop for Bomb {
 #[allow(unreachable_code, unused_variables)]
 fn main() {
     // This is a particularly tricky hook, since it unwinds, which the default one does not.
-    set_alloc_error_hook(|_layout| panic!("alloc error hook called"));
+    unsafe {
+        set_alloc_error_hook_unwinding(|_layout| panic!("alloc error hook called"));
+    }
 
     let bomb = Bomb;
     handle_alloc_error(Layout::for_value(&0));

tests/ui/panics/alloc_error_hook-unwind.rs

@@ -1,4 +1,5 @@
-//! Test that out-of-memory conditions trigger catchable panics with `set_alloc_error_hook`.
+//! Test that out-of-memory conditions trigger catchable panics
+//! with `set_alloc_error_hook_unwinding`.
 
 //@ run-pass
 //@ needs-unwind
@@ -12,7 +13,9 @@ use std::mem::forget;
 use std::panic::catch_unwind;
 
 fn main() {
-    std::alloc::set_alloc_error_hook(|_| panic!());
+    unsafe {
+        std::alloc::set_alloc_error_hook_unwinding(|_| panic!());
+    }
 
     let panic = catch_unwind(|| {
         // This is guaranteed to exceed even the size of the address space