Skip to content

Commit

Permalink
feat: global allocators post
Browse files Browse the repository at this point in the history
  • Loading branch information
BD103 committed Jun 30, 2023
1 parent 3a32cd8 commit ca75328
Showing 1 changed file with 180 additions and 0 deletions.
180 changes: 180 additions & 0 deletions content/blog/1.2023-06-27-global-allocators.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,180 @@
# Intercepting Allocations with the Global Allocator

::callout{kind="abstract" title="PREFACE"}
This post is directed towards programmers experienced with the [Rust](https://www.rust-lang.org) programming language that are interested in manual memory management. I will not explain some concepts like pointers, statics, etc. If you want to learn how to program in Rust, you can find a helpful page [here](https://www.rust-lang.org/learn).
::

It is possible to replace the global heap allocator used by [`Box<T>`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html), [`Vec<T>`](https://doc.rust-lang.org/stable/std/vec/struct.Vec.html), and more. This is useful if you want to use a custom allocator like [jemalloc](https://jemalloc.net) for the features it provides, or if you are working in a `#![no_std]` context without an OS allocator.

This article is going to cover:

- The [`GlobalAlloc`](https://doc.rust-lang.org/stable/std/alloc/trait.GlobalAlloc.html) trait
- The [`System`](https://doc.rust-lang.org/stable/std/alloc/struct.System.html) allocator
- Wrapping the `System` allocator with custom code
- The `wg-allocators`' plans for the future

## The `GlobalAlloc` Trait

[`GlobalAlloc`](https://doc.rust-lang.org/stable/std/alloc/trait.GlobalAlloc.html) is a trait that was stabilized in Rust 1.28. It is meant to be implemented to any struct that wishes to replace the default allocator. Here's a simplified version of its definition:

```rust
pub unsafe trait GlobalAlloc {
// Required methods
unsafe fn alloc(&self, layout: Layout) -> *mut u8;
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

// Provided methods (out of scope for this article)
...
}
```

Any struct that wants to become a global allocator has to be able to both allocate and deallocate regions of memory. It can approach this in many different ways, but all implementations use a pointer and a size to represent memory.

Take `alloc()` for example. It takes a [`Layout`](https://doc.rust-lang.org/stable/std/alloc/struct.Layout.html), which represents a size and alignment[^alignment] for a block of memory. It then returns a pointer to the first byte (`u8`) of the freshly allocated memory.

`dealloc()` is the inverse of this. It takes the pointer and the layout, then internally deallocates the memory block based on the given information.

[^alignment]: Alignment is used for structure packing. While it can be useful for implementing an allocator from scratch, that is out of scope for this article. For more information, you may find this [Wikipedia article](https://en.wikipedia.org/wiki/Data_structure_alignment) interesting.

## The `System` Allocator

If you are running on a computer that the standard library supports, then you have access to the [`System`](https://doc.rust-lang.org/stable/std/alloc/struct.System.html) default global allocator. If you don't specify a global allocator, then this will be used instead.

::callout{kind="tip"}
`System` uses libc's `malloc` and `free` functions on Unix platforms, while it uses `HeapAlloc` and associated functions on Windows.
::

Lucky for us, `System` already implements `GlobalAlloc`. Without changing anything, here's an overly verbose program that manually registers the `System` allocator:

```rust
// Import the System allocator.
use std::alloc::System;

// This attribute registers the global allocater.
#[global_allocator]
static A: System = System;

fn main() {
// Let's allocate on the heap to prove that it works.
let _ = Box::new(103);
}
```

To register a global allocator, you initialize it in a [static item](https://doc.rust-lang.org/reference/items/static-items.html) and annotate it with the `#[global_allocator]` attribute. There may only be one global allocator in a binary, but there is no restriction on where the static is and what visibility it has.

::callout{kind="warning"}
It is actually possible for a dependency crate to register a global allocator without the consumer knowing of it. This is **very** bad practice, as the end-user has no control over it. Please don't do this, or at least [feature-gate](https://doc.rust-lang.org/cargo/reference/features.html) it.
::

## Wrapping `System`

Since `GlobalAlloc` is a public trait, it is very possible to define our own allocator. Instead of going into the nitty-gritty of writing one from scratch, let's instead pass these allocation calls to `System`. This can let us intercept information, which we will see in the next section. Here's our new code:

```rust
use std::alloc::{GlobalAlloc, Layout, System};

// This is our custom allocator!
pub struct MyAlloc;

unsafe impl GlobalAlloc for MyAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
// Pass everything to System.
System.alloc(layout)
}

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
System.dealloc(ptr, layout)
}
}

// Register our custom allocator.
#[global_allocator]
static A: MyAlloc = MyAlloc;

fn main() {
// String, like Box, also allocates on the heap.
let _ = String::from("Boo!");
}
```

When running this program, you'll notice that it functions no different than before. We still use `System`, but we can now run our own code for each allocation. (Which is exactly what we're going to do next!)

## Counting Allocations

Say you are profiling some code and want to track how many heap allocations were made. This can easily be implemented with a custom allocator.

```rust
use std::{
alloc::{GlobalAlloc, Layout, System},
sync::atomic::{AtomicU64, Ordering},
};

// This is our counting allocator. It wraps a u64 that stores our actual count.
pub struct Counter(AtomicU64);

impl Counter {
// A const initializer that starts the count at 0.
pub const fn new() -> Self {
Counter(AtomicU64::new(0))
}

// Returns the current count.
pub fn count(&self) -> u64 {
// We're using Relaxed since there is only 1 synchronization primitive.
self.0.load(Ordering::Relaxed)
}
}

unsafe impl GlobalAlloc for Counter {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
// Increment our counter by 1. See other comment on Ordering.
self.0.fetch_add(1, Ordering::Relaxed);
System.alloc(layout)
}

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
// No modifications here! :)
System.dealloc(ptr, layout)
}
}

#[global_allocator]
static A: Counter = Counter::new();

fn main() {
// Track initial count and count after allocating once.
let count = A.count();
let _ = Box::new(1);
let new_count = A.count();

// count = 3, new_count = 4.
dbg!(count);
dbg!(new_count);
}
```

Wait, wait! Don't go! I know it uses atomics, but we can work through this together.

This allocator keeps a bit of internal state, a `u64` number to be exact. This number starts at 0 and increments by 1 every time `alloc()` is called. In order `main()` function, we track how many allocations were made before and after a `Box` was constructed. We know our allocator works because `new_count` is greater by 1 than `count`.

Realistically that is all the knowledge you need to understand this allocator. You may have spotted the `AtomicU64` though with its `load()` and `fetch_add()` calls. This is an advanced synchronization primitive that prevents data races in concurrency. Since the global allocator is used by the entire program, it is perfectly possible for it to be used by multiple threads. Our `AtomicU64` ensures that every allocation is counted, no matter where it was called from.

As much as I would love to talk about them, threads and data races are out of the scope of this article. (It's getting long enough as it is!) If you want to learn more about this fascinating subject, please see the [`std::thread`](https://doc.rust-lang.org/stable/std/thread/index.html) module and [Mara Bos's fantastic book](https://marabos.nl/atomics/).

::callout{kind="note"}
If we wanted to follow best practices, the `AtomicU64` would be stored in a separate static and `Counter` would remain a zero-sized-type (ZST). I felt that the example provided would be easier for a beginner to understand, even if it is not perfect.
::

## Towards the Future

And that's it! Thank you for reading this article, I hope you found it interesting. Global allocators can be a handy tool, though they do cover a relatively niche surface of Rust programming. There is a lot of work being put in by the `wg-allocators` working group to streamline this process and other related parts. Here are a few links that you may find interesting:

- The [`Allocator`](https://doc.rust-lang.org/stable/std/alloc/trait.Allocator.html) trait, for non-global allocators
- [Structs that support the `Allocator` API](https://github.com/rust-lang/wg-allocators/issues/7)
- The [`Global`](https://doc.rust-lang.org/stable/std/alloc/struct.Global.html) struct, which redirects all allocation calls to the registered global allocator
- [Pre-RFC: Storage API](https://internals.rust-lang.org/t/pre-rfc-storage-api/18822?u=bd103)
- [The `wg-allocators` Roadmap](https://github.com/rust-lang/wg-allocators/issues/48)

If you have any questions, feel free to comment on [my post]() in the Rust Users' Forum or [create a new issue](https://github.com/BD103/BD103/issues) in my Github repository. The source code for all examples is available [here](https://github.com/BD103/Allogators).

Have a great day!

0 comments on commit ca75328

Please sign in to comment.