Rust Custom Allocators

Posted on:November 20, 2024

On Tuesday, I had knee surgery, which meant I was in bed for about 36 hours with very little to do. I decided to take a dive into Rust allocators, self-referential structures, memory pinning, and custom concurrency executors. I’m not sure that’s what they meant when they said take it easy.

I learned a number of technical details during this deep dive. In this post, I will dive into Rust’s Custom Allocators.

Custom Allocators

My deep dive into Rust memory started with building a custom allocator. While you might not need them every day, custom allocators can be a real game-changer when you need to improve performance or handle specific memory limitations.

Let’s explore what custom allocators are, why you might need one, and how to build a simple bump allocator in Rust. We’ll also discuss how to use this allocator for specific parts of your application and provide an option for using it globally.

What Are Custom Allocators?

In most Rust programs, memory allocation happens automatically using the global allocator. By default, Rust uses std::alloc::System, which calls the system’s allocator (e.g., malloc on Linux or HeapAlloc on Windows). This works fine for most cases, but sometimes it’s not the best choice.

Custom allocators let you replace the global allocator or create specialized allocators for specific situations. They let you control how memory is used, which can help reduce waste, improve performance, or make memory usage more predictable. This can be especially helpful in:

Building a Simple Bump Allocator for Local Use

One of the simplest custom allocators is a bump allocator. It works by pre-allocating a block of memory and giving out parts of it in a straight line. It doesn’t support freeing individual allocations; instead, all memory is released at once when you reset the allocator. This makes bump allocators extremely fast, but they only work well in certain situations.

Here’s how we can create one in Rust for local use—specifically for allocating memory within a particular scope or for a single data structure.

1. Define the Allocator

Note: The Allocator API is currently unstable, so you need to use the nightly version of Rust and enable the feature allocator_api.

First, add the following to your Cargo.toml to use the nightly version:

[dependencies]
#![feature(allocator_api)]

Then, define the allocator:

#![feature(allocator_api)]
use std::alloc::{AllocError, Allocator, Layout};
use std::ptr::NonNull;
use std::sync::Mutex;

pub struct BumpAllocator {
    memory: Mutex<BumpMemory>,
}

struct BumpMemory {
    buffer: [u8; 1024], // Pre-allocated memory buffer
    offset: usize,      // Current allocation offset
}

impl BumpAllocator {
    pub fn new() -> Self {
        Self {
            memory: Mutex::new(BumpMemory {
                buffer: [0; 1024],
                offset: 0,
            }),
        }
    }
}

unsafe impl Allocator for BumpAllocator {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        let mut memory = self.memory.lock().unwrap();
        let start = memory.offset;
        let end = start + layout.size();

        if end > memory.buffer.len() {
            Err(AllocError)
        } else {
            memory.offset = end;
            println!("Allocated {} from {start} to {}", end-start, end-1);
            let slice = &mut memory.buffer[start..end];
            Ok(NonNull::from(slice))
        }
    }

    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
        // No-op: deallocation is unsupported in a bump allocator.
    }
}

Using the Allocator Locally

Now that we have a bump allocator, let’s see how we can use it locally for specific data structures instead of globally.

Example: Scoped Allocation with a Vec

Here’s how you can use the bump allocator for a single Vec:

#![feature(allocator_api)]

use allocator::BumpAllocator;

fn main() {
    let bump_allocator = BumpAllocator::new();
    let mut my_vec: Vec<u8, &BumpAllocator> = Vec::with_capacity_in(1, &bump_allocator);
    for i in 0u32..128 {
        my_vec.push((i % 255).try_into().unwrap());
    }
    println!("{:?}", my_vec); // Outputs: [1, 2, 3, 4, 5]
}

Advantages of Localized Allocator Usage

Using the Allocator Globally

If you want to use the bump allocator for the entire program, you can declare it as the global allocator using the #[global_allocator] attribute:

#![feature(allocator_api)]

/// Set the global allocator.
#[global_allocator]
static GLOBAL_ALLOCATOR: SimpleBumpAllocator = SimpleBumpAllocator;

fn main() {
    let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; // Allocates from the bump allocator
    println!("{:?}", v);
    let total_memory_allocated = OFFSET.load(Ordering::Relaxed);
    println!("Total memory allocated: {} bytes", total_memory_allocated);
}

By declaring the allocator globally, all allocations in your program—such as Vec, Box, or String—will use the bump allocator. There are lots of challenges to consider when using a bump allocator globally, such as fragmentation and memory leaks. Also, you cannot allocate memory during alloc, which makes debugging harder.

The full code for this allocator is available on GitHub.

Using Bumpalo for Efficient Bump Allocations

Another option for using bump allocation in Rust is the Bumpalo crate. Bumpalo is a popular, well-tested library that provides an easy-to-use bump allocator for efficient memory management in certain situations.

Example: Using Bumpalo

To use Bumpalo, add it to your Cargo.toml:

[dependencies]
bumpalo = "3"

Here’s a simple example of using Bumpalo for scoped memory allocations:

use bumpalo::Bump;

fn main() {
    let bump = Bump::new();
    // Allocate a vector using the bump allocator
    let numbers = bump.alloc_slice_copy(&[1, 2, 3, 4, 5]);
    println!("{:?}", numbers); // Outputs: [1, 2, 3, 4, 5]
}

In this example:

Advantages of Using Bumpalo

Key Takeaways

So the next time you’re working on a performance-critical or memory-constrained part of your app, try using a localized bump allocator or Bumpalo. It’s a powerful tool that gives you control—without sacrificing Rust’s safety guarantees!

All the code examples in this post are available on GitHub.