Code review for Static Pool Allocator

code review
1

Hello, I just finished my Static Pool Allocator. I checked with Miri then fixed all errors reported by Miri, now no error is reported by Miri. I need human review help

Here is the code :

#![feature(allocator_api)]

pub struct Block {
    next: Option<std::ptr::NonNull<Block>>
}

pub struct StaticPoolAllocator {
    start: std::ptr::NonNull<u8>,
    head: std::cell::UnsafeCell<Option<std::ptr::NonNull<Block>>>,
    block_size: usize,
    layout: std::alloc::Layout
}

impl StaticPoolAllocator {
    pub fn new(num_blocks: usize, block_size: usize) -> Self {
        let block_size = block_size.max(std::mem::size_of::<Block>());
        let layout = std::alloc::Layout::from_size_align(
            num_blocks * block_size,
            std::mem::align_of::<Block>()
        ).expect("Error in creating layout in fn new -> PoolAllocator");
        
        // SAFETY: the layout is valid, because if not it will trigger panic
        let start = unsafe { std::alloc::alloc(layout) };
        let start_ptr = std::ptr::NonNull::new(start).expect("Error OOM in fn new -> PoolAllocator");
        
        unsafe {
            for i in 0..num_blocks {
                let current_ptr = start.add(i * block_size).cast::<Block>();
                let next_ptr = if i < num_blocks - 1 {
                    Some(std::ptr::NonNull::new_unchecked(start.add((i + 1) * block_size).cast::<Block>()))
                } else {
                    None
                };
                
                (*current_ptr).next = next_ptr;
            }
        }
        
        Self {
            start: start_ptr,
            head: std::cell::UnsafeCell::new(Some(std::ptr::NonNull::new(start_ptr.as_ptr().cast::<Block>()).unwrap())),
            block_size,
            layout,
        }
        
    }
}

impl Drop for StaticPoolAllocator {
    fn drop(&mut self) {
        // SAFETY: deallocate using the pointer. The layout is saved, so it is the correct layout
        unsafe {
            std::alloc::dealloc(self.start.as_ptr(), self.layout);
        }
    }
}

unsafe impl std::alloc::Allocator for StaticPoolAllocator {
    #[inline(always)]
    fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        println!("allocate");
        let size = layout.size();
        if size > self.block_size {
            return Err(std::alloc::AllocError);
        }
        
        let head_ptr = unsafe { &mut *self.head.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
            ));
        }
        
        Err(std::alloc::AllocError)
    }
    
    #[inline(always)]
    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {
        println!("deallocate");
        let new_block = ptr.cast::<Block>();
        
        // dereference raw pointer is unsafe
        let head_ptr = unsafe { &mut *self.head.get() };
        unsafe { (*new_block.as_ptr()).next = *head_ptr };
        
        *head_ptr = Some(new_block);
    }
}

fn main() {
    let pool_allocator = StaticPoolAllocator::new(10, 1024);
    let mut vec = Vec::with_capacity_in(1, &pool_allocator);
    let a: i32 = 10;
    vec.push(a);
    vec.push(a);
    println!("{}", vec[0]);
}
2

allocate is not checking the alignment at all, that seems wrong to me.

3

Thank you for the review

I don't understand alignment in deep yet, my allocator above alocate memory with alignment 8 (I tried with println it)

What are the rules?

  • is it can only be used for type that has alignment <= 8?
  • aka it can't be used for type that has alignment > 8?

If yes, then I check if the type alignment > 8 or not, if yes then I return error?

So that means my allocator can't be used with all different type, is there a way to make it can be used in different allignment?

4

New lesson, I tried and yes it can not be used for alignment bigger than 8, Miri reported Undefined Behaviour :[

#[repr(align(64))]
#[derive(Debug, Clone)]
struct Tes {
    a: u64
}

fn main() {
    let pool_allocator = StaticPoolAllocator::new(10, 1024);
    let mut vec = Vec::with_capacity_in(1, &pool_allocator);
    let a = Tes { a: 10 };
    vec.push(a.clone());
    vec.push(a);
    println!("{:?}", vec[0]);
}

Smaller alignment works. But it must be power of 2, I forgot this :[

Should I use alignment of 64 in fn new? 64 is power of 8. What are the drawbacks of this? Eg will it use bigger memory that it should?

I have idea to make the alignment costumizeable, eg alignment = alignment of block * caller input (usize)

What is your solution idea?

5

I updated the allocate method to this :

fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        println!("allocate");
        let size = layout.size();
        let align = layout.align();
        let base_align = std::mem::align_of::<Block>();
        if size > self.block_size || align > base_align {
            return Err(std::alloc::AllocError);
        }
        
        let head_ptr = unsafe { &mut *self.head.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
            ));
        }
        
        Err(std::alloc::AllocError)
    }

Now it will return error type align > the base align

I learned that the layout already guarante that the alignment must be power of 2, because it already return compile error when I test with align 6. So I don't add check for power of 2 in method allocate

Edit : I saved the base align in the struct to avoid recalculate every allocating

Edit : I canceled the saving because I learned that align_of() will be replaced with constant number value at compile time, saving in struct will waste memory :[

Is there any another error?

6

You might want to read through this, as well as watch a few lectures on the subject. In my (humble) opinion, you really shouldn't tackle manual allocation until you're crystal clear on what you're doing. The topic itself is quite a fun challenge to tackle, though. Have fun with it!

7

It might be worth knowing that besides alignment being needed for "cpu reasons" ( typically up to 16), memory allocators themselves sometimes use a large alignment so they can find out what block a given allocation came from.

The way I think of it is it means the lowest n bits of an address are zero, when n is the base 2 log of the alignment.

8

Thank you, that documentation really helps. I think if we want to be able to do something, we must try and learn it, not avoid it

I got new idea, I will make other new pool allocator but this one will have pool list of different alignment, eg

[pool 1 align 8]--chain--[pool 2 align 16]--[pool 3 align 32]

Then when allocating data, I will check

If align_of data <= pool 1 -> use pool 1
Else if align <= pool 2 -> use pool 2

Is it good?