Rust: An easy way to count memory allocations

Real time waits for nothing. That’s why avoiding memory allocations while generating audio samples is important. Today I implemented a quick hack to make sure my SNES SPC player does not allocate memory after the initial setup is done.

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};

struct CheckAlloc;

static ALLOCATED: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for CheckAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOCATED.fetch_add(1, SeqCst);
        System.alloc(layout)
    }

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

#[global_allocator]
static A: CheckAlloc = CheckAlloc;

#[test]
fn test_no_allocation_while_running() {
    let mut synth = Synth::new();
    synth.init();
    let before_count = ALLOCATED.load(SeqCst);
    synth.generate_samples(NUM_SAMPLES_TO_GENERATE);
    let after_count = ALLOCATED.load(SeqCst);
    assert_eq!(before_count, after_count);
}

The idea is simple. Create a custom allocator which wraps the System allocator and counts the number of allocations. After the initial setup, the number of allocations should be the same. I borrowed the idea from a blog post from redislabs.