Add set_oom_handler and use it print a message when out of memory

2 files changed, 45 insertions, 11 deletions

diff --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs
index 93b84cdedd4..dd4bf5174ee 100644
--- a/src/liballoc/lib.rs
+++ b/src/liballoc/lib.rs
@@ -92,6 +92,7 @@
 #![feature(unsize)]
 #![feature(drop_in_place)]
 #![feature(fn_traits)]
+#![feature(const_fn)]
 
 #![feature(needs_allocator)]
 
@@ -127,15 +128,6 @@ mod boxed_test;
 pub mod arc;
 pub mod rc;
 pub mod raw_vec;
+pub mod oom;
 
-/// Common out-of-memory routine
-#[cold]
-#[inline(never)]
-#[unstable(feature = "oom", reason = "not a scrutinized interface",
-           issue = "27700")]
-pub fn oom() -> ! {
-    // FIXME(#14674): This really needs to do something other than just abort
-    //                here, but any printing done must be *guaranteed* to not
-    //                allocate.
-    unsafe { core::intrinsics::abort() }
-}
+pub use oom::oom;
diff --git a/src/liballoc/oom.rs b/src/liballoc/oom.rs
new file mode 100644
index 00000000000..d355d59185e
--- /dev/null
+++ b/src/liballoc/oom.rs
@@ -0,0 +1,42 @@
+// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::mem;
+use core::intrinsics;
+
+static OOM_HANDLER: AtomicPtr<()> = AtomicPtr::new(default_oom_handler as *mut ());
+
+fn default_oom_handler() -> ! {
+    // The default handler can't do much more since we can't assume the presence
+    // of libc or any way of printing an error message.
+    unsafe { intrinsics::abort() }
+}
+
+/// Common out-of-memory routine
+#[cold]
+#[inline(never)]
+#[unstable(feature = "oom", reason = "not a scrutinized interface",
+           issue = "27700")]
+pub fn oom() -> ! {
+    let value = OOM_HANDLER.load(Ordering::SeqCst);
+    let handler: fn() -> ! = unsafe { mem::transmute(value) };
+    handler();
+}
+
+/// Set a custom handler for out-of-memory conditions
+///
+/// To avoid recursive OOM failures, it is critical that the OOM handler does
+/// not allocate any memory itself.
+#[unstable(feature = "oom", reason = "not a scrutinized interface",
+           issue = "27700")]
+pub fn set_oom_handler(handler: fn() -> !) {
+    OOM_HANDLER.store(handler as *mut (), Ordering::SeqCst);
+}