Writing arenas in Rust from scratch
Custom memory allocations and arenas are notoriously hard to deal with in Rust due to the ownership model.
In a lot of projects, they can significantly improve performance. For example, let's take databases. Sending a query to a database produces a lot of temporary objects, and if the default allocator is used, you will need to allocate and deallocate thousands of small objects for each query, which is wasteful.
Almost every database engine uses custom memory pools to improve performance. Some go so far as to statically allocate all the memory before the program starts.
Another good example is scripting languages. I have an article on how Python uses arenas to reduce the number of allocations for Python types. A simple web app in Python can allocate more than a million of objects after 100 HTTP requests. Using a system allocator for such a case is wasteful.
And yet, to this day, there is no proper way to deal with custom allocations and arenas in Rust.
The allocator_api, which allows using a custom allocator, has been nightly-only (experimental) since 2016.
Recently, I wanted to learn how to implement an arena in Rust. This is my explanation of how to implement a basic version.
What is an arena?
An arena is a preallocated memory block that you allocate only once. Instead of allocating 1000 small objects, you can allocate one big block of memory and put the objects there when needed. This is usually done by simply moving the pointer forward in the allocated memory block. Usually, blocks allocated inside the arena are called slots or cells.
When you are done working with the allocated block, you can reuse the allocated space for the next set of objects instead of deallocating it. You don't even need to clear the memory - all you need to do is move the pointer back to the beginning of the block.
By using arenas, you get faster allocations, less memory fragmentation, and better cache locality.
Basic implementation
The simplest possible arena can be implemented as a vector that has a fixed capacity.
struct Arena<T> {
items: Vec<T>,
capacity: usize,
}
impl<T> Arena<T> {
fn with_capacity(capacity: usize) -> Self {
Self {
items: Vec::with_capacity(capacity),
capacity,
}
}
fn alloc(&mut self, value: T) -> Option<usize> {
if self.items.len() >= self.capacity {
return None;
}
let id = self.items.len();
self.items.push(value);
Some(id)
}
fn get(&self, id: usize) -> Option<&T> {
self.items.get(id)
}
fn get_mut(&mut self, id: usize) -> Option<&mut T> {
self.items.get_mut(id)
}
fn len(&self) -> usize {
self.items.len()
}
}
This can even be simplified to:
type Arena<T> = Vec<T>;
fn alloc<T>(arena: &mut Arena<T>, value: T) -> Option<usize> {
if arena.len() == arena.capacity() {
return None;
}
let id = arena.len();
arena.push(value);
Some(id)
}
fn main() {
let mut arena: Arena<i32> = Vec::with_capacity(3);
}
It uses a fixed capacity because we don't want the arena to be reallocated. The whole point of arenas is to minimize memory allocations and fragmentation.
Right now, this code is safe, but it has a lot of limitations and is only used for simple use cases.
- You can only store one data type
- You can't remove objects
- It can't hold collections or any dynamically sized objects (the actual values will be allocated outside of the arena)
- You can't have multiple mutable references to different memory blocks
- It's not thread-safe, even when different threads mutate different slots
The borrow checker doesn't let you hold multiple mutable (&mut) references into the same vector:
fn main() {
#[derive(Debug)]
struct Point {
x: i32,
y: i32,
}
let mut arena = Arena::with_capacity(10);
let p1 = arena.alloc(Point { x: 10, y: 20 }).unwrap();
let p2 = arena.alloc(Point { x: 30, y: 40 }).unwrap();
let a = arena.get_mut(p1).unwrap();
let b = arena.get_mut(p2).unwrap();
println!("{:?}", a);
println!("{:?}", b);
}
error[E0499]: cannot borrow `arena` as mutable more than once at a time
--> src/main.rs:49:13
|
48 | let a = arena.get_mut(p1).unwrap();
| ----- first mutable borrow occurs here
49 | let b = arena.get_mut(p2).unwrap();
| ^^^^^ second mutable borrow occurs here
50 |
51 | println!("{:?}", a);
| - first borrow later used here
When we store simple objects, they are stored directly inside the arena:
let mut arena = Arena::<i32>::with_capacity(4);
arena.alloc(10);
arena.alloc(20);
arena.alloc(30);
The memory layout looks something like this:
arena
|
+-- Vec
|
+-- ptr ────────────────┐
+-- len = 3 |
+-- capacity = 4 |
v
Vector based buffer:
+--------+--------+--------+--------+
| i32 10| i32 20| i32 30| unused |
+--------+--------+--------+--------+
Now let's look at the strings:
let mut arena = Arena::<String>::with_capacity(3);
arena.alloc("hello".to_string());
arena.alloc("world".to_string());
arena.alloc("arena".to_string());
arena
|
+-- Vec<String>
|
+-- ptr ────────────────────────────────┐
+-- len = 3 |
+-- capacity = 3 |
v
Vector based buffer:
+----------------+----------------+----------------+
| String #0 | String #1 | String #2 |
| | | |
| ptr ───────┐ | ptr ───────┐ | ptr ───────┐ |
| len = 5 | | len = 5 | | len = 5 | |
| cap = 5 | | cap = 5 | | cap = 5 | |
+-------------|--+-------------|--+-------------|--+
| | |
v v v
Separate heap allocations:
+-------+
| hello |
+-------+
+-------+
| world |
+-------+
+-------+
| arena |
+-------+
As you can see, only string headers are stored inside the arena. This defeats the whole purpose of arenas. We want to store such objects in the same memory block.
Because of this, most arena implementations use unsafe.
Advanced implementation
To support different data types and multiple mutable references, we will need to use memory buffers.
Here is what the Arena implementation now looks like:
use std::{
alloc::{self, Layout},
cell::UnsafeCell,
fmt,
marker::PhantomData,
ops::{Deref, DerefMut},
ptr::{self, NonNull},
};
struct Slot {
offset: usize,
drop: unsafe fn(*mut u8),
}
pub struct ArenaBox<'a, T> {
ptr: NonNull<T>,
// PhantomData tells the compiler to treat this pointer as owning a T.
_marker: PhantomData<&'a mut T>,
}
pub struct Arena {
buf: NonNull<u8>,
capacity: usize,
cursor: UnsafeCell<usize>,
slots: UnsafeCell<Vec<Slot>>,
}
impl Arena {
pub fn new(capacity: usize) -> Self {
let layout = Layout::from_size_align(capacity, 1).unwrap();
let buf = unsafe { NonNull::new(alloc::alloc(layout)).expect("alloc failed") };
Self {
buf,
capacity,
cursor: UnsafeCell::new(0),
slots: UnsafeCell::new(Vec::new()),
}
}
pub fn alloc<'a, T>(&'a self, value: T) -> ArenaBox<'a, T> {
// Compute the offset and proper alignment
let layout = Layout::new::<T>();
let cursor = unsafe { *self.cursor.get() };
let base = self.buf.as_ptr() as usize;
let aligned = base
.checked_add(cursor)
.unwrap()
.next_multiple_of(layout.align());
let offset = aligned - base;
assert!(offset + layout.size() <= self.capacity);
// Register a drop fn for this type
// This generates a proper deallocation function for collection types (Vec, String)
unsafe fn drop_ptr<T>(ptr: *mut u8) {
unsafe { ptr::drop_in_place(ptr.cast::<T>()) };
}
unsafe {
(*self.slots.get()).push(Slot {
offset,
drop: drop_ptr::<T>,
});
// Write its value into the raw buffer
ptr::write(self.buf.as_ptr().add(offset).cast::<T>(), value);
*self.cursor.get() = offset + layout.size();
}
// Return a typed handle that preserves the type, even though it's a raw pointer into the buffer.
ArenaBox {
ptr: unsafe { NonNull::new_unchecked(self.buf.as_ptr().add(offset).cast()) },
_marker: PhantomData,
}
}
}
impl Drop for Arena {
fn drop(&mut self) {
unsafe {
for slot in (*self.slots.get()).iter().rev() {
(slot.drop)(self.buf.as_ptr().add(slot.offset));
}
alloc::dealloc(
self.buf.as_ptr(),
Layout::from_size_align(self.capacity, 1).unwrap(),
);
}
}
}
impl<'a, T> Deref for ArenaBox<'a, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { self.ptr.as_ref() }
}
}
impl<'a, T> DerefMut for ArenaBox<'a, T> {
fn deref_mut(&mut self) -> &mut T {
unsafe { self.ptr.as_mut() }
}
}
impl<'a, T: fmt::Debug> fmt::Debug for ArenaBox<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
(**self).fmt(f)
}
}
fn main() {
#[derive(Debug)]
struct Point {
x: i32,
y: i32,
}
#[derive(Debug)]
struct Color {
r: u8,
g: u8,
b: u8,
}
let arena = Arena::new(4096);
let mut p = arena.alloc(Point { x: 10, y: 20 });
let mut c = arena.alloc(Color { r: 255, g: 0, b: 0 });
p.x = 99;
c.r = 128;
println!("p = {:?}", p);
println!("c = {:?}", c);
}
Our Arena struct now holds a pointer to the memory block as well as a vector that keeps track of occupied slots.
Each Slot holds an offset and a drop function. We need to keep track of the drop functions in case we want to store
dynamic objects that also allocate outside of the arena (e.g., String).
Without this, we would leak memory when the arena is dropped.
The Slot struct can be entirely avoided if we only allow simple types with the Copy trait, which do not allocate memory outside of the arena.
The ArenaBox is a wrapper around a pointer to the object stored inside the arena.
It's used to implement the Deref and DerefMut traits, so we can use it like a normal reference and provide type safety.
In new, instead of a vector, we now allocate a raw memory buffer with a fixed capacity.
When we call alloc, it computes the proper offset and alignment for the specified type.
The Drop implementation calls the drop functions for all slots when the arena is destroyed,
so there are no memory leaks for dynamic objects.
When dereferencing an ArenaBox, we now get an appropriate reference to the object stored inside the arena.
Also, we can now have two mutable references to different objects stored inside the arena.
cargo run
p = Point { x: 99, y: 20 }
c = Color { r: 128, g: 0, b: 0 }
UnsafeCell
The alloc method takes an immutable reference, but we still can mutate the cursor and slots, because
we wrapped it in UnsafeCell.
If alloc were mutable (fn alloc<'a, T>(&'a mut self, value: T)), we could only have one mutable reference even when
alloc points to different slots because of the common 'a lifetime.
That's a limitation of the Rust borrow checker, because the Rust compiler does not understand that different
calls to alloc are independent.
We could also remove lifetimes with mutable alloc, but that would result in UB.
ArenaBox would outlive the arena, and we would have dangling pointers.
The UnsafeCell type allows interior mutability, but it's our responsibility now to prevent undefined behavior.
Handling dynamic types
Our current implementation solved two problems (different types and mutable references),
but dynamic types such as String or a Vec will still use
regular heap allocations to store the underlying values. Only the header (pointer, length, capacity) is stored inside the arena.
There are two ways to allocate dynamic types inside an arena. One option would be to reimplement collection types from the std.
This option is used by some of the arena crates, such as bumpalo.
It requires a full reimplementation. In bumpalo, the String reimplementation is 2000 lines long and Vec is 3000 lines long.
In the end, you may also lose some of the optimizations that the std collections have, and you will have to maintain your own collection types
once new features are added to the std.
Another option would be to use the nightly allocator_api feature.
By using a nightly feature, you are locked into nightly Rust.
Here is how we can implement this:
#![feature(allocator_api)]
use std::{
alloc::{AllocError, Allocator, Layout},
cell::UnsafeCell,
ptr::NonNull,
};
pub struct Arena {
buf: NonNull<u8>,
capacity: usize,
cursor: UnsafeCell<usize>,
}
unsafe impl Allocator for &Arena {
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
let cursor = unsafe { *self.cursor.get() };
let base = self.buf.as_ptr() as usize;
let current = base.checked_add(cursor).ok_or(AllocError)?;
let aligned = current.next_multiple_of(layout.align());
let offset = aligned.checked_sub(base).ok_or(AllocError)?;
let end = offset.checked_add(layout.size()).ok_or(AllocError)?;
if end > self.capacity {
return Err(AllocError);
}
unsafe {
*self.cursor.get() = end;
let ptr = self.buf.as_ptr().add(offset);
let slice = std::ptr::slice_from_raw_parts_mut(ptr, layout.size());
Ok(NonNull::new_unchecked(slice))
}
}
unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
// no-op, we free everything at once
}
}
impl Arena {
pub fn new(capacity: usize) -> Self {
let layout = Layout::from_size_align(capacity, 1).unwrap();
let buf = unsafe {
let ptr = std::alloc::alloc(layout);
NonNull::new(ptr).expect("arena allocation failed")
};
Self {
buf,
capacity,
cursor: UnsafeCell::new(0),
}
}
}
impl Drop for Arena {
fn drop(&mut self) {
unsafe {
std::alloc::dealloc(
self.buf.as_ptr(),
Layout::from_size_align(self.capacity, 1).unwrap(),
);
}
}
}
#[derive(Debug)]
struct Point { x: i32, y: i32 }
fn main() {
let arena = Arena::new(4096);
let mut v: Vec<i32, &Arena> = Vec::new_in(&arena);
v.push(10);
v.push(20);
v.push(30);
let mut w: Vec<f64, &Arena> = Vec::new_in(&arena);
w.push(1.5);
w.push(2.5);
let point = Box::new_in(Point { x: 10, y: 20 }, &arena);
println!("vec = {:?}", v);
println!("w = {:?}", w);
println!("point = {:?}", point);
}
In this code, we also need UnsafeCell because the Allocator trait requires &self, but we need to mutate the cursor.
This is the design of the Allocator trait, and we can't change it.
Since this code is not thread-safe, there is no UB.
You cannot use it from multiple threads.
On nightly, each std collection accepts an allocator that must implement the Allocator trait.
When a vector resizes, old memory will be left allocated because we did not implement deallocation.
So, such an approach works best when you know the upper bound on how much memory your dynamic collections will need.
Implementing deallocation is tricky. It requires maintaining a list of free slots. When dealing with free slots, we would also need to check if an object would fit in one of the free slots. After a lot of deallocations, our memory can become fragmented and the entire buffer would need to be compacted.
One nice thing about this approach is that all memory is freed at once when the arena is dropped, so we don't need to deal
with the Slot struct and drop functions.
The current limitation of our approach is that our arena is not thread-safe. We did not implement any locking mechanism. Implementing thread-safe arenas is another hard problem and requires different approaches depending on the use case.
Thread safety
The easiest way to add thread-safety is to use Mutex and Arc:
struct Arena {
buf: NonNull<u8>,
capacity: usize,
cursor: Mutex<usize>,
}
unsafe impl Send for Arena {}
unsafe impl Sync for Arena {}
#[derive(Clone)]
pub struct ArenaAllocator(Arc<Arena>);
impl ArenaAllocator {
pub fn new(capacity: usize) -> Self {
Self(Arc::new(Arena::new(capacity)))
}
}
Here we use Arc (reference counting) to keep the arena alive when other threads are using it, and Mutex to protect the cursor from concurrent access.
This adds overhead and contention.
When we access a vector (from the earlier allocator_api example) from the arena inside a thread, its header (pointer, length, capacity) gets moved, but
the underlying data still stays in the arena. So, we now can have multiple threads accessing different dynamic objects inside the arena.
Each new allocation requires locking the mutex, so that multiple threads do not allocate at the same time and corrupt the cursor.
Our Send and Sync implementations are unsafe, so Rust can't guarantee thread-safety here.
Two threads mutating the same logical object will result in UB.
Everything becomes much harder if you want to be able to resize your arena or deallocate objects, which a lot of popular crates support. When working with arenas, a lot of trade-offs need to be made.
Can the mutex be avoided? Yes.
While I was looking at how other projects do this, I found a clever approach in the Turso source code that is lock-free, i.e., it does not use mutexes.
Turso is a Rust rewrite of SQLite.
In databases, rows are often stored in fixed-size pages. Each page (usually 4 KB) stores multiple rows.
Turso devs have implemented an arena for fixed-size byte buffers.
Each arena gets divided into 4 KB slots.
Slot availability is tracked using an atomic slot bitmap.
It's an array of AtomicU64 where each bit represents a slot. If the bit is set, the slot is occupied, otherwise it's free.
This approach is lock-free because atomic integers are guaranteed to be thread-safe and can be modified without locking.
They also use some extra tricks:
mmapwithMADV_HUGEPAGEto read buffers into memory lazilyio_uringto read the fixed buffers from disk to memory faster- When the fixed arena is full, they fall back to regular heap allocations
- When deallocating the data, they just mark the slot empty and do not remove the actual data.
The list of improvemnts that we can implement can go on and on. If you want to learn more, I suggest studying the source code of bumpalo. It had UB problems in the past, but they are fixed now. Studying fixes is also a good way to learn more about UBs.
If you have any questions, feel free to ask them via e-mail displayed in the footer.
All articles on this website are written by a human.