Skip to content

Commit

Permalink
Auto merge of #50880 - glandium:oom, r=SimonSapin
Browse files Browse the repository at this point in the history
OOM handling changes

As discussed in #49668 (comment) 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
  • Loading branch information
bors committed May 30, 2018
2 parents 20af72b + a4d899b commit 4f99f37
Show file tree
Hide file tree
Showing 10 changed files with 206 additions and 113 deletions.
13 changes: 7 additions & 6 deletions src/liballoc/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
if !ptr.is_null() {
ptr as *mut u8
} else {
oom()
oom(layout)
}
}
}
Expand All @@ -134,12 +134,13 @@ pub(crate) unsafe fn box_free<T: ?Sized>(ptr: Unique<T>) {
}

#[rustc_allocator_nounwind]
pub fn oom() -> ! {
extern {
pub fn oom(layout: Layout) -> ! {
#[allow(improper_ctypes)]
extern "Rust" {
#[lang = "oom"]
fn oom_impl() -> !;
fn oom_impl(layout: Layout) -> !;
}
unsafe { oom_impl() }
unsafe { oom_impl(layout) }
}

#[cfg(test)]
Expand All @@ -154,7 +155,7 @@ mod tests {
unsafe {
let layout = Layout::from_size_align(1024, 1).unwrap();
let ptr = Global.alloc_zeroed(layout.clone())
.unwrap_or_else(|_| oom());
.unwrap_or_else(|_| oom(layout));

let mut i = ptr.cast::<u8>().as_ptr();
let end = i.offset(layout.size() as isize);
Expand Down
2 changes: 1 addition & 1 deletion src/liballoc/arc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -553,7 +553,7 @@ impl<T: ?Sized> Arc<T> {
let layout = Layout::for_value(&*fake_ptr);

let mem = Global.alloc(layout)
.unwrap_or_else(|_| oom());
.unwrap_or_else(|_| oom(layout));

// Initialize the real ArcInner
let inner = set_data_ptr(ptr as *mut T, mem.as_ptr() as *mut u8) as *mut ArcInner<T>;
Expand Down
156 changes: 82 additions & 74 deletions src/liballoc/raw_vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,14 +96,15 @@ impl<T, A: Alloc> RawVec<T, A> {
NonNull::<T>::dangling().as_opaque()
} else {
let align = mem::align_of::<T>();
let layout = Layout::from_size_align(alloc_size, align).unwrap();
let result = if zeroed {
a.alloc_zeroed(Layout::from_size_align(alloc_size, align).unwrap())
a.alloc_zeroed(layout)
} else {
a.alloc(Layout::from_size_align(alloc_size, align).unwrap())
a.alloc(layout)
};
match result {
Ok(ptr) => ptr,
Err(_) => oom(),
Err(_) => oom(layout),
}
};

Expand Down Expand Up @@ -318,7 +319,7 @@ impl<T, A: Alloc> RawVec<T, A> {
new_size);
match ptr_res {
Ok(ptr) => (new_cap, ptr.cast().into()),
Err(_) => oom(),
Err(_) => oom(Layout::from_size_align_unchecked(new_size, cur.align())),
}
}
None => {
Expand All @@ -327,7 +328,7 @@ impl<T, A: Alloc> RawVec<T, A> {
let new_cap = if elem_size > (!0) / 8 { 1 } else { 4 };
match self.a.alloc_array::<T>(new_cap) {
Ok(ptr) => (new_cap, ptr.into()),
Err(_) => oom(),
Err(_) => oom(Layout::array::<T>(new_cap).unwrap()),
}
}
};
Expand Down Expand Up @@ -389,37 +390,7 @@ impl<T, A: Alloc> RawVec<T, A> {
pub fn try_reserve_exact(&mut self, used_cap: usize, needed_extra_cap: usize)
-> Result<(), CollectionAllocErr> {

unsafe {
// NOTE: we don't early branch on ZSTs here because we want this
// to actually catch "asking for more than usize::MAX" in that case.
// If we make it past the first branch then we are guaranteed to
// panic.

// Don't actually need any more capacity.
// Wrapping in case they gave a bad `used_cap`.
if self.cap().wrapping_sub(used_cap) >= needed_extra_cap {
return Ok(());
}

// Nothing we can really do about these checks :(
let new_cap = used_cap.checked_add(needed_extra_cap).ok_or(CapacityOverflow)?;
let new_layout = Layout::array::<T>(new_cap).map_err(|_| CapacityOverflow)?;

alloc_guard(new_layout.size())?;

let res = match self.current_layout() {
Some(layout) => {
debug_assert!(new_layout.align() == layout.align());
self.a.realloc(NonNull::from(self.ptr).as_opaque(), layout, new_layout.size())
}
None => self.a.alloc(new_layout),
};

self.ptr = res?.cast().into();
self.cap = new_cap;

Ok(())
}
self.reserve_internal(used_cap, needed_extra_cap, Fallible, Exact)
}

/// Ensures that the buffer contains at least enough space to hold
Expand All @@ -443,9 +414,9 @@ impl<T, A: Alloc> RawVec<T, A> {
///
/// Aborts on OOM
pub fn reserve_exact(&mut self, used_cap: usize, needed_extra_cap: usize) {
match self.try_reserve_exact(used_cap, needed_extra_cap) {
match self.reserve_internal(used_cap, needed_extra_cap, Infallible, Exact) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocErr) => oom(),
Err(AllocErr) => unreachable!(),
Ok(()) => { /* yay */ }
}
}
Expand All @@ -467,37 +438,7 @@ impl<T, A: Alloc> RawVec<T, A> {
/// The same as `reserve`, but returns on errors instead of panicking or aborting.
pub fn try_reserve(&mut self, used_cap: usize, needed_extra_cap: usize)
-> Result<(), CollectionAllocErr> {
unsafe {
// NOTE: we don't early branch on ZSTs here because we want this
// to actually catch "asking for more than usize::MAX" in that case.
// If we make it past the first branch then we are guaranteed to
// panic.

// Don't actually need any more capacity.
// Wrapping in case they give a bad `used_cap`
if self.cap().wrapping_sub(used_cap) >= needed_extra_cap {
return Ok(());
}

let new_cap = self.amortized_new_size(used_cap, needed_extra_cap)?;
let new_layout = Layout::array::<T>(new_cap).map_err(|_| CapacityOverflow)?;

// FIXME: may crash and burn on over-reserve
alloc_guard(new_layout.size())?;

let res = match self.current_layout() {
Some(layout) => {
debug_assert!(new_layout.align() == layout.align());
self.a.realloc(NonNull::from(self.ptr).as_opaque(), layout, new_layout.size())
}
None => self.a.alloc(new_layout),
};

self.ptr = res?.cast().into();
self.cap = new_cap;

Ok(())
}
self.reserve_internal(used_cap, needed_extra_cap, Fallible, Amortized)
}

/// Ensures that the buffer contains at least enough space to hold
Expand Down Expand Up @@ -553,12 +494,12 @@ impl<T, A: Alloc> RawVec<T, A> {
/// # }
/// ```
pub fn reserve(&mut self, used_cap: usize, needed_extra_cap: usize) {
match self.try_reserve(used_cap, needed_extra_cap) {
match self.reserve_internal(used_cap, needed_extra_cap, Infallible, Amortized) {
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocErr) => oom(),
Err(AllocErr) => unreachable!(),
Ok(()) => { /* yay */ }
}
}
}
}
/// Attempts to ensure that the buffer contains at least enough space to hold
/// `used_cap + needed_extra_cap` elements. If it doesn't already have
/// enough capacity, will reallocate in place enough space plus comfortable slack
Expand Down Expand Up @@ -670,14 +611,81 @@ impl<T, A: Alloc> RawVec<T, A> {
old_layout,
new_size) {
Ok(p) => self.ptr = p.cast().into(),
Err(_) => oom(),
Err(_) => oom(Layout::from_size_align_unchecked(new_size, align)),
}
}
self.cap = amount;
}
}
}

enum Fallibility {
Fallible,
Infallible,
}

use self::Fallibility::*;

enum ReserveStrategy {
Exact,
Amortized,
}

use self::ReserveStrategy::*;

impl<T, A: Alloc> RawVec<T, A> {
fn reserve_internal(
&mut self,
used_cap: usize,
needed_extra_cap: usize,
fallibility: Fallibility,
strategy: ReserveStrategy,
) -> Result<(), CollectionAllocErr> {
unsafe {
use alloc::AllocErr;

// NOTE: we don't early branch on ZSTs here because we want this
// to actually catch "asking for more than usize::MAX" in that case.
// If we make it past the first branch then we are guaranteed to
// panic.

// Don't actually need any more capacity.
// Wrapping in case they gave a bad `used_cap`.
if self.cap().wrapping_sub(used_cap) >= needed_extra_cap {
return Ok(());
}

// Nothing we can really do about these checks :(
let new_cap = match strategy {
Exact => used_cap.checked_add(needed_extra_cap).ok_or(CapacityOverflow)?,
Amortized => self.amortized_new_size(used_cap, needed_extra_cap)?,
};
let new_layout = Layout::array::<T>(new_cap).map_err(|_| CapacityOverflow)?;

alloc_guard(new_layout.size())?;

let res = match self.current_layout() {
Some(layout) => {
debug_assert!(new_layout.align() == layout.align());
self.a.realloc(NonNull::from(self.ptr).as_opaque(), layout, new_layout.size())
}
None => self.a.alloc(new_layout),
};

match (&res, fallibility) {
(Err(AllocErr), Infallible) => oom(new_layout),
_ => {}
}

self.ptr = res?.cast().into();
self.cap = new_cap;

Ok(())
}
}

}

impl<T> RawVec<T, Global> {
/// Converts the entire buffer into `Box<[T]>`.
///
Expand Down
2 changes: 1 addition & 1 deletion src/liballoc/rc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -668,7 +668,7 @@ impl<T: ?Sized> Rc<T> {
let layout = Layout::for_value(&*fake_ptr);

let mem = Global.alloc(layout)
.unwrap_or_else(|_| oom());
.unwrap_or_else(|_| oom(layout));

// Initialize the real RcBox
let inner = set_data_ptr(ptr as *mut T, mem.as_ptr() as *mut u8) as *mut RcBox<T>;
Expand Down
50 changes: 47 additions & 3 deletions src/libstd/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))]
Expand Down
Loading

0 comments on commit 4f99f37

Please sign in to comment.