【Rust】在Rust里自定义全局内存分配器

tangptr@126.com

前言

Rust早期是直接使用jemalloc这个库进行堆上内存分配的。尽管jemalloc性能很高，还支持多线程，碎片回收能力也强，但是也有很明显的缺点：

• 体积过大：一个Hello-World都能超过2MB
• 和valgrind不兼容：无法验证无内存泄漏
• 兼容性差：很多架构都不支持

因此在2018年底，Rust不再使用jemalloc，而是直接用系统默认的内存分配器。实际上系统默认的内存分配器可能也是jemalloc（如FreeBSD）。
此外，在今年六月，jemalloc的创始人Jason Evans宣布停更并简介了jemalloc项目的生平。

GlobalAlloc trait

自定义内存分配器最重要的是要实现GlobalAlloc trait。它有四个方法：alloc, dealloc, alloc_zeroed和realloc。其中alloc和dealloc是必须实现的，而alloc_zeroed和realloc是不需要自己实现的。
实现了GlobalAlloc trait后，需要用#[global_allocator]宏来声明使用该全局分配器：

#[global_allocator] static WINDOWS_ALLOCATOR:WindowsAllocator=WindowsAllocator;

Layout结构体

具体如何分配内存主要看Layout结构体中的size和align方法。它们分别返回这个对象的大小和对齐粒度。
在我之前讲堆上Print的时候就给出过一个简易的包装HeapAlloc API的实现。当时我就说过：这个实现不考虑内存对齐。如果用到高对齐粒度的指令（如AVX(-512)的vmovaps指令要求32甚至64字节对齐），会在未对齐的条件下会报错。

强制对齐算法

如果你的内存分配器没有指定对齐粒度的能力（比如POSIX里的memalign函数，C11标准里的aligned_alloc函数），则需要使用强制对齐算法。
比如微软的windows-drivers-rs里的对齐内存分配器就是我实现的（暂未合并）。
这个算法的具体流程是：

1. 判断对齐粒度：如果所需要的对齐粒度小于等于分配器的固定对齐粒度，则直接进行分配。否则进入下一步。
2. 扩大分配量：你总计需要分配align+size的大小。
3. 计算对齐指针：用and运算对返回的内存地址进行对齐后，再增加一次对齐粒度的值。将这个值作为返回的指针。
4. 存放原始指针：将对齐出来的指针减去指针的大小，在这个地址存放原始指针。

在HeapAlloc上应用强制对齐

这里我们实现一下alloc, dealloc和realloc，而alloc_zeroed实在是没有单独实现的必要了。

struct WindowsAllocator;

static mut PROCESS_HEAP: LazyCell<HANDLE>=LazyCell::new(|| unsafe{GetProcessHeap()});
#[global_allocator] static WINDOWS_ALLOCATOR:WindowsAllocator=WindowsAllocator;

impl WindowsAllocator
{
    const ALIGNMENT:usize=MEMORY_ALLOCATION_ALIGNMENT as usize;

    fn require_realignment(align:usize)->bool
    {
        align>Self::ALIGNMENT
    }

    fn realign(ptr:*mut u8,align:usize)->*mut *mut u8
    {
        let align_mask=!(align-1);
        let mut q:usize=ptr as usize;
        q&=align_mask;
        q+=align;
        q as *mut *mut u8
    }
}

unsafe impl GlobalAlloc for WindowsAllocator
{
    unsafe fn alloc(&self,layout: Layout)->*mut u8
    {
        if Self::require_realignment(layout.align())
        {
            let p:*mut u8=unsafe{HeapAlloc(*PROCESS_HEAP,0,layout.size()+layout.align())}.cast();
            let q=Self::realign(p,layout.align());
            unsafe
            {
                q.sub(1).write(p);
            }
            println!("HeapAlloc returned {p:p}! Returning {q:p}...");
            q.cast()
        }
        else
        {
            unsafe
            {
                HeapAlloc(*PROCESS_HEAP,0,layout.size()).cast()
            }
        }
    }

    unsafe fn dealloc(&self,ptr:*mut u8,layout: Layout)
    {
        if Self::require_realignment(layout.align())
        {
            let q=ptr as *mut *mut u8;
            unsafe
            {
                let p=q.sub(1).read();
                println!("Freeing {q:p}! Passing {p:p} to HeapFree...");
                HeapFree(*PROCESS_HEAP,0,p.cast());
            }
        }
        else
        {
            unsafe
            {
                HeapFree(*PROCESS_HEAP,0,ptr.cast());
            }
        }
    }

    unsafe fn realloc(&self,ptr:*mut u8,layout: Layout,new_size:usize)->*mut u8
    {
        if Self::require_realignment(layout.align())
        {
            unsafe
            {
                let new=self.alloc(Layout::from_size_align_unchecked(new_size,layout.align()));
                copy_nonoverlapping(ptr,new,layout.size());
                self.dealloc(ptr,layout);
                new
            }
        }
        else
        {
            unsafe
            {
                HeapReAlloc(*PROCESS_HEAP,0,ptr.cast(),new_size).cast()
            }
        }
    }
}

其他第三方库

虽然jemalloc停更了，但也不是不能用。比如Rust的编译器就在用jemallocator这个crate。
除了jemalloc，微软出品的mimalloc也是个很好的选择。

我把经典的dlmalloc重新包装了一下，发布为portable-dlmalloc。这个库主要在可移植性上（其实我只测试了Windows, Linux和UEFI的no_std环境）做文章，理论上可移植到任意平台上（只要C编译器能编译出来即可）。但是移植需要自行实现（如封装mmap函数等）。

结语

本文详解了GlobalAlloc trait，并提出了强制对齐算法，可以使一些不支持自定义粒度的内存分配器去分配对齐的内存。
Rust其实还有个Allocator trait，这个trait允许你使用不同的内存分配器来分配内存（比如Box的new_in方法）。但是这个trait仅在nightly里可用，尚不stable，因此不在本文讨论范围。
Rust提供的大部分动态内存类型里有try_xxx方法（比如Box的try_new，Vec的try_reserve，BTreeMap的try_insert），它们会返回Result类型，可以用于检测内存分配是否失败了。但很可惜仅在nightly可用，尚不stable。不使用try_xxx时，内存分配失败会直接panic。

0xAA55

原来强制对齐的算法可以这么实现。那它分配超量内存的时候肯定必须要至少超过一个指针的长度，不然没地方存指针。

修改内存分配器只能 Nightly 么？咦说起来我好像 Stable 的 Rust 是可以直接用 cargo add 来添加 Nightly 的 crate 的。

tangptr@126.com

0xAA55 发表于 2025-8-21 15:41
原来强制对齐的算法可以这么实现。那它分配超量内存的时候肯定必须要至少超过一个指针的长度，不然没地方存 ...

使用#[global_allocator]来指定自定义的全局内存分配器时，不需要nightly。
需要nightly的东西是《同时使用多个内存分配器》（也就是那个Allocator trait，比如Box::new_in返回的类型是Box<T,A>，而不是Box::new返回的Box<T,Global>），以及《用try_xxx分配内存》（允许分配失败）。