Auto merge of #50880 - glandium:oom, r=SimonSapin

OOM handling changes

As discussed in https://github.com/rust-lang/rust/issues/49668#issuecomment-384893456 and subsequent.

This does have codegen implications. Even without the hooks, and with a handler that ignores the arguments, the compiler doesn't eliminate calling `rust_oom` with the `Layout`. Even if it managed to eliminate that, with the hooks, I don't know if the compiler would be able to figure out it can skip it if the hook is never set.

A couple implementation notes:
- I went with explicit enums rather than bools because it makes it clearer in callers what is being requested.
- I didn't know what `feature` to put the hook setting functions behind. (and surprisingly, the compile went through without any annotation on the functions)
- There's probably some bikeshedding to do on the naming.

Cc: @Simonsapin, @sfackler

1 files changed, 47 insertions, 3 deletions

diff --git a/src/libstd/alloc.rs b/src/libstd/alloc.rs
index 78d3d6d5e60..4f9dffc7c95 100644
--- a/src/libstd/alloc.rs
+++ b/src/libstd/alloc.rs
@@ -13,15 +13,59 @@
 #![unstable(issue = "32838", feature = "allocator_api")]
 
 #[doc(inline)] #[allow(deprecated)] pub use alloc_crate::alloc::Heap;
-#[doc(inline)] pub use alloc_crate::alloc::{Global, oom};
+#[doc(inline)] pub use alloc_crate::alloc::{Global, Layout, oom};
 #[doc(inline)] pub use alloc_system::System;
 #[doc(inline)] pub use core::alloc::*;
 
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::{mem, ptr};
+
+static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
+
+/// Registers a custom OOM hook, replacing any that was previously registered.
+///
+/// The OOM hook is invoked when an infallible memory allocation fails.
+/// The default hook prints a message to standard error and aborts the
+/// execution, but this behavior can be customized with the [`set_oom_hook`]
+/// and [`take_oom_hook`] functions.
+///
+/// The hook is provided with a `Layout` struct which contains information
+/// about the allocation that failed.
+///
+/// The OOM hook is a global resource.
+pub fn set_oom_hook(hook: fn(Layout) -> !) {
+    HOOK.store(hook as *mut (), Ordering::SeqCst);
+}
+
+/// Unregisters the current OOM hook, returning it.
+///
+/// *See also the function [`set_oom_hook`].*
+///
+/// If no custom hook is registered, the default hook will be returned.
+pub fn take_oom_hook() -> fn(Layout) -> ! {
+    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
+    if hook.is_null() {
+        default_oom_hook
+    } else {
+        unsafe { mem::transmute(hook) }
+    }
+}
+
+fn default_oom_hook(layout: Layout) -> ! {
+    rtabort!("memory allocation of {} bytes failed", layout.size())
+}
+
 #[cfg(not(test))]
 #[doc(hidden)]
 #[lang = "oom"]
-pub extern fn rust_oom() -> ! {
-    rtabort!("memory allocation failed");
+pub extern fn rust_oom(layout: Layout) -> ! {
+    let hook = HOOK.load(Ordering::SeqCst);
+    let hook: fn(Layout) -> ! = if hook.is_null() {
+        default_oom_hook
+    } else {
+        unsafe { mem::transmute(hook) }
+    };
+    hook(layout)
 }
 
 #[cfg(not(test))]