Towards a Custom Memory Allocator in Rust

I’ve never quite understood why you might want to make your own memory allocator. I’m into Rust, and Rust is a systems programming language, meaning you can use it to build low level things as you would in C or C++. The benefit of Rust, of course, are it’s guarantees around memory safety, enabling you to have more confidence that what you create will be safer in the long run.

This is a great basis for a memory allocator.

So why would you want your own? After all, the operating system will do it for you, helping you find memory for any heap allocated variables that you need in your program. This results in a lot of overhead though. Your program needs to make many calls in order to get the needed resources. What if there was a way to shortcut that and build something right into your program.

This is where the power of building your own shines. You can make your own memory allocator, as a crate within the larger program, and leveraging Rust’s type system, expose it as an API to call when you need a chunk of memory of specified size.

As an exercise, I started very simple. What if our memory allocator was nothing more than a buffer and a position marker? You can carve off chunks from the buffer and pass them out, keeping track of the position associated with the chunk. You can do it with a struct type:

struct MemoryArena {
     buffer: Vec<u8>
     position: usize,
}

There’s our starting point. We’ve got a struct type with a field called buffer (the actual memory you want to give out) and a field called position which marks the point where a new chunk will start.

With this done, we can now make some methods for interacting with our “memory arena”, like so:

impl MemoryArena {
    fn new(capacity: usize) -> Self {
        Self {
            buffer: vec![0u8; capacity],
            position: 0,
        }
    }

    fn alloc(&mut self, size: usize) -> Option<&mut [u8]> {
        let end = self.position + size;
        if end > self.buffer.len() {
            None
        } else {
            let start = self.position;
            self.position = end;
            Some(&mut self.buffer[start..end])
        }
    }

    fn reset(&mut self) {
        self.position = 0;
    }

    fn remaining(&self) -> usize {
        self.buffer.len() - self.position
    }
}

We’ve got:

• a new() method which initializes an empty instance of our memory arena
• an alloc() method which actually does the work, it takes the size of the chunk we want to give out and uses that to return an Option<&mut [u8] which is a changeable slice of memory to use
• a reset() method, which returns the position marker to 0
• a remaining() method which returns the amount of memory left in that’s available to give out

Conceptually, it’s all pretty straightforward. You could take this code and work it into any project that would benefit from a fixed pool of memory. The performance benefit in not having to go through the OS would be pretty tangible.

Despite it’s simplicity, there’s one big pitfall to the code above. Rust’s ownership rules prevent our MemoryArena type from having more than one mutable reference at a time. This limits the usefulness of the code above a little too much. In order to make this truly something you might make into production ready code, you’re going to want to dive into the arcane world of unsafe Rust. I’ll explain more about that in a future post.