Hello, I just finished my fixed size multi pool allocator. Now I need code review :>
From the previous review, I got this following :
- small alignment can't be used to alloc bigger alignment
- the size must be power of the alignment size
I created new project to address that with the following solution :
- there are multiple pool
- each pool can be configured to have different size and alignment
- the pool with size and alignment bigger than the data's size and alignment is chosen
I already checked it with Miri and fixed the reported errors, is there still something wrong that I don't know?
#![feature(allocator_api)]
pub struct Block {
next: Option<std::ptr::NonNull<Block>>
}
pub struct StaticMultiPoolAlloc<const N: usize> {
start: [std::ptr::NonNull<u8>; N],
head: [std::cell::UnsafeCell<Option<std::ptr::NonNull<Block>>>; N],
block_size: [usize; N],
aligns: [usize; N],
layout: [std::alloc::Layout; N]
}
fn align_up(size: usize, align: usize) -> usize {
(size + align - 1) & !(align - 1)
}
impl<const N: usize> StaticMultiPoolAlloc<N> {
pub fn new(num_blocks: [usize; N], block_size: [usize; N], align_multiply: [usize; N]) -> Self {
_ = align_multiply.map(|val| {
if val == 0 {
panic!("Align multiply must be > 0");
}
});
let block_align = std::mem::align_of::<Block>();
let aligns: [usize; N] = align_multiply.map(|val| val * block_align);
let base_block_size = std::mem::size_of::<Block>();
let block_size: [usize; N] = std::array::from_fn(|i| {
let size = block_size[i].max(base_block_size);
align_up(size, block_align)
});
let layout: [std::alloc::Layout; N] = std::array::from_fn(|i| std::alloc::Layout::from_size_align(
num_blocks[i] * block_size[i],
aligns[i]
).expect("Error in creating layout in fn new -> PoolAllocator"));
// SAFETY: the layout is valid, because if not it will trigger panic
let starts: [std::ptr::NonNull<u8>; N] = std::array::from_fn(|i| {
let start = unsafe { std::alloc::alloc(layout[i]) };
let start_ptr = std::ptr::NonNull::new(start).expect("Error OOM in fn new -> PoolAllocator");
unsafe {
for j in 0..num_blocks[i] {
let current_ptr = start.add(j * block_size[i]).cast::<Block>();
let next_ptr = if j < num_blocks[i] - 1 {
Some(std::ptr::NonNull::new_unchecked(start.add((j + 1) * block_size[i]).cast::<Block>()))
} else {
None
};
(*current_ptr).next = next_ptr;
}
}
start_ptr
});
let head: [std::cell::UnsafeCell<Option<std::ptr::NonNull<Block>>>; N] = std::array::from_fn(|i| {
std::cell::UnsafeCell::new(Some(std::ptr::NonNull::new(starts[i].as_ptr().cast::<Block>()).unwrap()))
});
Self {
start: starts,
head,
block_size,
aligns,
layout,
}
}
}
impl<const N: usize> Drop for StaticMultiPoolAlloc<N> {
fn drop(&mut self) {
// SAFETY: deallocate using the pointer. The layout is saved, so it is the correct layout
for (i, val) in self.start.iter().enumerate() {
unsafe {
std::alloc::dealloc(val.as_ptr(), self.layout[i]);
}
}
}
}
unsafe impl<const N: usize> std::alloc::Allocator for StaticMultiPoolAlloc<N> {
#[inline(always)]
fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
let size = layout.size();
let align = layout.align();
//let base_align = std::mem::align_of::<Block>();
let index = self.aligns.iter().enumerate().position(|(i, &val)| {
size <= self.block_size[i] && align <= val
});
let i = index.ok_or(std::alloc::AllocError)?;
let head_ptr = unsafe { &mut *self.head[i].get() };
if let Some(block) = *head_ptr {
let taken_block = block;
unsafe { *head_ptr = block.as_ref().next };
let ptr = taken_block.cast::<u8>();
return Ok(std::ptr::NonNull::slice_from_raw_parts(
ptr, self.block_size[i]
));
}
Err(std::alloc::AllocError)
}
#[inline(always)]
unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
let new_block = ptr.cast::<Block>();
let align = layout.align();
let index = self.aligns.iter().enumerate().position(|(_, &val)| {
align <= val
});
let i = index.ok_or(std::alloc::AllocError).unwrap();
// dereference raw pointer is unsafe
let head_ptr = unsafe { &mut *self.head[i].get() };
unsafe { (*new_block.as_ptr()).next = *head_ptr };
*head_ptr = Some(new_block);
}
}
#[derive(Debug)]
#[repr(align(64))]
struct Tes(i32);
fn main() {
/*
[pool 1, num block: 4, each block: 80 byte, total byte: 80*4, align: 8*1]
[pool 2, num block: 2, each block: 1024 byte, total byte: 1024*2, align: 8*8]
total prealloc in entire pool: (80*4) + (1024*2)
fixed size, panic if no enough size available
single threaded usage, each thread has their own pool so there is no locking, maximizing cpu cache hit
*/
let pool: StaticMultiPoolAlloc<2> = StaticMultiPoolAlloc::new([4, 2], [80, 1024], [1, 8]);
let mut vec = Vec::with_capacity_in(2, &pool);
vec.push(Tes(10));
println!("{:?}", vec[0].0);
}