Choose memory allocator based on enviorment variable

help
1

Self-explanatory. Something like how cpython can choose it's allocator based on an environment variable. Essentially I'd prefer mimalloc by default, but on the chance it breaks due to some esoteric config the user should be able to fallback to malloc.

2

not quite self-explanatory, do you want the decision to happen at compile time or runtime?

if you want it to happen at runtime, you'll need to use libc::getenv to avoid allocations, and do that inside a LazyLock.

you are going to be adding a branch to every allocation, which may be noticeable depending on your application.

3

Decision at Runtime, I think it would be easy to get it working at compile time with env!.

Ideally I would prefer not to branch however. I think I'll look a little more into how cpython handles it.

4

What you can do is declare your “allocator” as just a wrapper around some inner, unknown allocator:

struct ConfigurableAllocator {
    inner: &'static (dyn GlobalAlloc + Sync),
}

To allow changing the allocator you use after static initialization, you’ll need to make a Sync wrapper around UnsafeCell:

struct MakeMutable<T> {
    inner: UnsafeCell<T>,
}

unsafe impl<T: Sync> Sync for MakeMutable<T> {}

Then declare your global allocator.

#[global_allocator]
static ALLOC: MakeMutable<ConfigurableAllocator> = MakeMutable {
    inner: UnsafeCell::new(ConfigurableAllocator {
        inner: &Mimalloc,
    }),
};

unsafe impl GlobalAlloc for MakeMutable<ConfigurableAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        unsafe { (*self.inner.get()).inner.alloc(layout) }
    }
    
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { (*self.inner.get()).inner.dealloc(ptr, layout) }
    }
    
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        unsafe { (*self.inner.get()).inner.alloc_zeroed(layout) }
    }
    
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        unsafe { (*self.inner.get()).inner.realloc(ptr, layout, new_size) }
    }
}

Now, in main(), change the allocator:

fn main() {
    let ptr = unsafe {
        libc::getenv(c"PROGRAM_ALLOCATOR".as_ptr())
    };
    if !ptr.is_null() {
        let string = unsafe { CStr::from_ptr(ptr) };
        if string == c"malloc" {
            unsafe {
                (*ALLOC.inner.get()).inner = &Malloc;
            }
        }
    }
    you_can_now_allocate();
    do_whatever_your_program_is_supposed_to();
}

The dynamic dispatch looks slow because it is. However, using the global allocator in Rust always incurs dynamic dispatch cost through the magic __rust_alloc/__rust_dealloc family[1] of functions. The only way to not have dynamic dispatch happening is to patch the code of <ConfigurableAllocator as GlobalAlloc> at runtime.[2]

It would probably be faster to have ConfigurableAllocator hold function pointers to each function and then use those instead of using dyn Traits. Using dyn GlobalAlloc makes the example for the forum quicker to write though.

Here’s all the code pasted into the Rust Playground.

  1. The other members are __rust_realloc and __rust_alloc_zeroed IIRC ↩︎

  2. The idea is to replace the code in the function bodies (or even worse, the magic __rust functions, which are definitely not stable API) with just a jmp to <Mimalloc as GlobalAlloc>::alloc or <Malloc as GlobalAlloc>::alloc. This is a horrible code crime and probably wouldn’t even work because the code pages aren’t writable. ↩︎