From 0f4ef003ac1691d04f0ce519d1d78696689534aa Mon Sep 17 00:00:00 2001 From: Mike Hommey Date: Tue, 15 May 2018 09:56:46 +0900 Subject: Pass a `Layout` to `oom` As discussed in https://github.com/rust-lang/rust/issues/49668#issuecomment-384893456 and subsequent, there are use-cases where the OOM handler needs to know the size of the allocation that failed. The alignment might also be a cause for allocation failure, so providing it as well can be useful. --- src/liballoc/alloc.rs | 13 ++-- src/liballoc/arc.rs | 2 +- src/liballoc/raw_vec.rs | 156 +++++++++++++++++++++++++----------------------- src/liballoc/rc.rs | 2 +- 4 files changed, 91 insertions(+), 82 deletions(-) (limited to 'src/liballoc') diff --git a/src/liballoc/alloc.rs b/src/liballoc/alloc.rs index 4ae8fc649dd..8753c495737 100644 --- a/src/liballoc/alloc.rs +++ b/src/liballoc/alloc.rs @@ -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) } } } @@ -134,12 +134,13 @@ pub(crate) unsafe fn box_free(ptr: Unique) { } #[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)] @@ -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::().as_ptr(); let end = i.offset(layout.size() as isize); diff --git a/src/liballoc/arc.rs b/src/liballoc/arc.rs index d0950bff9ce..f7513248784 100644 --- a/src/liballoc/arc.rs +++ b/src/liballoc/arc.rs @@ -553,7 +553,7 @@ impl Arc { 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; diff --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs index 5c6f6b22aae..07bb7f1a3eb 100644 --- a/src/liballoc/raw_vec.rs +++ b/src/liballoc/raw_vec.rs @@ -96,14 +96,15 @@ impl RawVec { NonNull::::dangling().as_opaque() } else { let align = mem::align_of::(); + 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), } }; @@ -318,7 +319,7 @@ impl RawVec { 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 => { @@ -327,7 +328,7 @@ impl RawVec { let new_cap = if elem_size > (!0) / 8 { 1 } else { 4 }; match self.a.alloc_array::(new_cap) { Ok(ptr) => (new_cap, ptr.into()), - Err(_) => oom(), + Err(_) => oom(Layout::array::(new_cap).unwrap()), } } }; @@ -389,37 +390,7 @@ impl RawVec { 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::(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 @@ -443,9 +414,9 @@ impl RawVec { /// /// 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 */ } } } @@ -467,37 +438,7 @@ impl RawVec { /// 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::(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 @@ -553,12 +494,12 @@ impl RawVec { /// # } /// ``` 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 @@ -670,7 +611,7 @@ impl RawVec { 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; @@ -678,6 +619,73 @@ impl RawVec { } } +enum Fallibility { + Fallible, + Infallible, +} + +use self::Fallibility::*; + +enum ReserveStrategy { + Exact, + Amortized, +} + +use self::ReserveStrategy::*; + +impl RawVec { + 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::(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 RawVec { /// Converts the entire buffer into `Box<[T]>`. /// diff --git a/src/liballoc/rc.rs b/src/liballoc/rc.rs index d0188c6e828..1648fc6b7ef 100644 --- a/src/liballoc/rc.rs +++ b/src/liballoc/rc.rs @@ -668,7 +668,7 @@ impl Rc { 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; -- cgit 1.4.1-3-g733a5