From a4d899b4a1248f885563e241fa56fe9f69616dc2 Mon Sep 17 00:00:00 2001
From: Mike Hommey <mh@glandium.org>
Date: Wed, 16 May 2018 09:16:37 +0900
Subject: Add hooks allowing to override the `oom` behavior

---
 src/libstd/alloc.rs | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

(limited to 'src/libstd/alloc.rs')

diff --git a/src/libstd/alloc.rs b/src/libstd/alloc.rs
index 0c95ceff2e3..4f9dffc7c95 100644
--- a/src/libstd/alloc.rs
+++ b/src/libstd/alloc.rs
@@ -17,11 +17,55 @@
 #[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(_: Layout) -> ! {
-    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))]
-- 
cgit 1.4.1-3-g733a5