Review: Static Multi Pool Allocator
@fuji-184 Fu
·
2026-04-25
·
via The Rust Programming Language Forum - Latest posts
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); }
此内容由惯性聚合(RSS阅读器)自动聚合整理,仅供阅读参考。 原文来自 — 版权归原作者所有。