Skip to content

Commit

Permalink
Fix allocations on bumps with small limits exceeding limits
Browse files Browse the repository at this point in the history 
To ensure that allocations do not exceed the allocation limit imposed on
a `Bump`, the logic for calculating the final size of an allocated chunk
was factored out into a separate function so that the limit enforcing
code higher up in the control flow could more strictly enforce the
allocation limit, as previously this code did not have visibility into
the final size of an allocation request.

To ensure that allocations in arenas with small limits could succeed,
`alloc_layout_slow` was changed to allow bypassing the minimum chunk
size for new chunks when trying to allocate in new arenas when the limit
was lower than the default minimum chunk size.
  • Loading branch information
@Eliasin
Eliasin committed Jun 30, 2022
1 parent 8efa890 commit 7f47641
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 59 deletions.
155 changes: 116 additions & 39 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -452,6 +452,13 @@ const FIRST_ALLOCATION_GOAL: usize = 1 << 9;
// take the alignment into account.
const DEFAULT_CHUNK_SIZE_WITHOUT_FOOTER: usize = FIRST_ALLOCATION_GOAL - OVERHEAD;

#[derive(Debug)]
struct NewChunkMemoryDetails {
new_size_without_footer: usize,
align: usize,
size: usize,
}

/// Wrapper around `Layout::from_size_align` that adds debug assertions.
#[inline]
unsafe fn layout_from_size_align(size: usize, align: usize) -> Layout {
Expand Down Expand Up @@ -589,16 +596,77 @@ impl Bump {
.unwrap_or(None)
}

/// Whether a new allocation fits within the headroom before the allocation limit.
fn new_allocation_fits_under_limit(
/// Whether a request to allocate a new chunk with a given size for a given
/// requested layout will fit under the allocation limit set on a `Bump`.
fn chunk_alloc_request_fits_under_limit(
allocation_limit_remaining: Option<usize>,
allocation_size: usize,
chunk_size: usize,
requested_layout: Layout,
) -> bool {
let NewChunkMemoryDetails {
new_size_without_footer,
align: _,
size: _,
} = match Bump::new_chunk_memory_details(Some(chunk_size), requested_layout) {
Some(memory_details) => memory_details,
None => {
return false;
}
};

allocation_limit_remaining
.map(|allocation_limit_left| allocation_limit_left >= allocation_size)
.map(|allocation_limit_left| allocation_limit_left >= new_size_without_footer)
.unwrap_or(true)
}

/// Determine the memory details including final size, alignment and
/// final size without footer for a new chunk that would be allocated
/// to fulfill an allocation request.
fn new_chunk_memory_details(
new_size_without_footer: Option<usize>,
requested_layout: Layout,
) -> Option<NewChunkMemoryDetails> {
let mut new_size_without_footer =
new_size_without_footer.unwrap_or(DEFAULT_CHUNK_SIZE_WITHOUT_FOOTER);

// We want to have CHUNK_ALIGN or better alignment
let mut align = CHUNK_ALIGN;

// If we already know we need to fulfill some request,
// make sure we allocate at least enough to satisfy it
align = align.max(requested_layout.align());
let requested_size =
round_up_to(requested_layout.size(), align).unwrap_or_else(allocation_size_overflow);
new_size_without_footer = new_size_without_footer.max(requested_size);

// We want our allocations to play nice with the memory allocator,
// and waste as little memory as possible.
// For small allocations, this means that the entire allocation
// including the chunk footer and mallocs internal overhead is
// as close to a power of two as we can go without going over.
// For larger allocations, we only need to get close to a page
// boundary without going over.
if new_size_without_footer < PAGE_STRATEGY_CUTOFF {
new_size_without_footer =
(new_size_without_footer + OVERHEAD).next_power_of_two() - OVERHEAD;
} else {
new_size_without_footer =
round_up_to(new_size_without_footer + OVERHEAD, 0x1000)? - OVERHEAD;
}

debug_assert_eq!(align % CHUNK_ALIGN, 0);
debug_assert_eq!(new_size_without_footer % CHUNK_ALIGN, 0);
let size = new_size_without_footer
.checked_add(FOOTER_SIZE)
.unwrap_or_else(allocation_size_overflow);

Some(NewChunkMemoryDetails {
new_size_without_footer,
size,
align,
})
}

/// Allocate a new chunk and return its initialized footer.
///
/// If given, `layouts` is a tuple of the current chunk size and the
Expand All @@ -610,39 +678,11 @@ impl Bump {
prev: NonNull<ChunkFooter>,
) -> Option<NonNull<ChunkFooter>> {
unsafe {
let mut new_size_without_footer =
new_size_without_footer.unwrap_or(DEFAULT_CHUNK_SIZE_WITHOUT_FOOTER);

// We want to have CHUNK_ALIGN or better alignment
let mut align = CHUNK_ALIGN;

// If we already know we need to fulfill some request,
// make sure we allocate at least enough to satisfy it
align = align.max(requested_layout.align());
let requested_size = round_up_to(requested_layout.size(), align)
.unwrap_or_else(allocation_size_overflow);
new_size_without_footer = new_size_without_footer.max(requested_size);

// We want our allocations to play nice with the memory allocator,
// and waste as little memory as possible.
// For small allocations, this means that the entire allocation
// including the chunk footer and mallocs internal overhead is
// as close to a power of two as we can go without going over.
// For larger allocations, we only need to get close to a page
// boundary without going over.
if new_size_without_footer < PAGE_STRATEGY_CUTOFF {
new_size_without_footer =
(new_size_without_footer + OVERHEAD).next_power_of_two() - OVERHEAD;
} else {
new_size_without_footer =
round_up_to(new_size_without_footer + OVERHEAD, 0x1000)? - OVERHEAD;
}

debug_assert_eq!(align % CHUNK_ALIGN, 0);
debug_assert_eq!(new_size_without_footer % CHUNK_ALIGN, 0);
let size = new_size_without_footer
.checked_add(FOOTER_SIZE)
.unwrap_or_else(allocation_size_overflow);
let NewChunkMemoryDetails {
new_size_without_footer,
align,
size,
} = Bump::new_chunk_memory_details(new_size_without_footer, requested_layout)?;

let layout = layout_from_size_align(size, align);

Expand Down Expand Up @@ -1453,7 +1493,19 @@ impl Bump {
.checked_mul(2)?
.max(min_new_chunk_size);
let sizes = iter::from_fn(|| {
if base_size >= min_new_chunk_size {
let bypass_min_chunk_size_for_small_limits = match self.allocation_limit() {
Some(limit)
if layout.size() < limit
&& base_size >= layout.size()
&& limit < DEFAULT_CHUNK_SIZE_WITHOUT_FOOTER
&& self.allocated_bytes() == 0 =>
{
true
}
_ => false,
};

if base_size >= min_new_chunk_size || bypass_min_chunk_size_for_small_limits {
let size = base_size;
base_size = base_size / 2;
Some(size)
Expand All @@ -1464,7 +1516,11 @@ impl Bump {

let new_footer = sizes
.filter_map(|size| {
if Bump::new_allocation_fits_under_limit(allocation_limit_remaining, size) {
if Bump::chunk_alloc_request_fits_under_limit(
allocation_limit_remaining,
size,
layout,
) {
Bump::new_chunk(Some(size), layout, current_footer)
} else {
None
Expand Down Expand Up @@ -1887,6 +1943,7 @@ unsafe impl<'a> Allocator for &'a Bump {
#[cfg(test)]
mod tests {
use super::*;
use quickcheck::quickcheck;

#[test]
fn chunk_footer_is_five_words() {
Expand Down Expand Up @@ -1960,4 +2017,24 @@ mod tests {
b.realloc(p1, l3, 48000).unwrap();
}
}

quickcheck! {
fn limit_is_never_exceeded(limit: usize) -> bool {
let bump = Bump::new();

bump.set_allocation_limit(Some(limit));

// The exact numbers here on how much to allocate are a bit murky but we
// have two main goals.
//
// - Attempt to allocate over the allocation limit imposed
// - Allocate in increments small enough that at least a few allocations succeed
let layout = Layout::array::<u8>(limit / 16).unwrap();
for _ in 0..32 {
let _ = bump.try_alloc_layout(layout);
}

limit >= bump.allocated_bytes()
}
}
}
29 changes: 9 additions & 20 deletions tests/allocation_limit.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,4 @@
use std::alloc::Layout;

use bumpalo::Bump;
use quickcheck::quickcheck;

#[test]
fn allocation_limit_trivial() {
Expand Down Expand Up @@ -85,22 +82,14 @@ fn new_bump_allocated_bytes_is_zero() {
assert_eq!(bump.allocated_bytes(), 0);
}

quickcheck! {
fn limit_is_never_exceeded(limit: usize) -> bool {
let b = Bump::new();

b.set_allocation_limit(Some(limit));

// The exact numbers here on how much to allocate are a bit murky but we
// have two main goals.
//
// - Attempt to allocate over the allocation limit imposed
// - Allocate in increments small enough that at least a few allocations succeed
let layout = Layout::array::<u8>(limit / 16).unwrap();
for _ in 0..32 {
let _ = b.try_alloc_layout(layout);
}
#[test]
fn small_allocation_limit() {
let bump = Bump::new();

limit >= b.allocated_bytes()
}
// 64 is chosen because any limit below 64 is
// unreasonably tight when trying to optimize the
// allocation and will always fail due to our attempts
// to allocate chunks in 'nice' sizes
bump.set_allocation_limit(Some(64));
assert!(bump.try_alloc([0; 1]).is_ok());
}

0 comments on commit 7f47641

Please sign in to comment.