论坛 › 技巧探讨 › 【Rust】在Rust里自定义全局内存分配器

tangptr@126.com 发表于 2025-8-21 15:07:50

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

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

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

因此在2018年底，Rust不再使用`jemalloc`，而是直接用系统默认的内存分配器。实际上系统默认的内存分配器可能也是`jemalloc`（如FreeBSD）。
此外，在今年六月，`jemalloc`的创始人Jason Evans宣布停更[并简介了jemalloc项目的生平](https://jasone.github.io/2025/06/12/jemalloc-postmortem/)。

# `GlobalAlloc` trait
自定义内存分配器最重要的是要实现[`GlobalAlloc` trait](https://doc.rust-lang.org/alloc/alloc/trait.GlobalAlloc.html)。它有四个方法：`alloc`, `dealloc`, `alloc_zeroed`和`realloc`。其中`alloc`和`dealloc`是必须实现的，而`alloc_zeroed`和`realloc`是不需要自己实现的。
实现了`GlobalAlloc` trait后，需要用`#`宏来声明使用该全局分配器：
```Rust
# static WINDOWS_ALLOCATOR:WindowsAllocator=WindowsAllocator;
```

## `Layout`结构体
具体如何分配内存主要看`Layout`结构体中的`size`和`align`方法。它们分别返回这个对象的大小和对齐粒度。
在我之前讲[堆上Print](https://www.0xaa55.com/thread-27602-1-1.html#42379_%E5%A0%86%E4%B8%8Aprint)的时候就给出过一个简易的包装`HeapAlloc` API的实现。当时我就说过：这个实现不考虑内存对齐。如果用到高对齐粒度的指令（如AVX(-512)的`vmovaps`指令要求32甚至64字节对齐），会在未对齐的条件下会报错。

## 强制对齐算法
如果你的内存分配器没有指定对齐粒度的能力（比如(https://linux.die.net/man/3/memalign)，C11标准里的`aligned_alloc`函数），则需要使用强制对齐算法。
比如微软的(https://github.com/microsoft/windows-drivers-rs/pull/353)就是我实现的（暂未合并）。
这个算法的具体流程是：

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

## 在`HeapAlloc`上应用强制对齐
这里我们实现一下`alloc`, `dealloc`和`realloc`，而`alloc_zeroed`实在是没有单独实现的必要了。
```Rust
struct WindowsAllocator;

static mut PROCESS_HEAP: LazyCell<HANDLE>=LazyCell::new(|| unsafe{GetProcessHeap()});
# 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的编译器就在用(https://crates.io/crates/jemallocator)这个crate。
除了`jemalloc`，微软出品的(https://crates.io/crates/mimalloc)也是个很好的选择。

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

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

0xAA55 发表于 2025-8-21 15:41:58

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

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

tangptr@126.com 发表于 2025-8-21 16:19:56

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


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

我没有加特林 发表于 2025-9-7 09:37:03

复制黏贴复制黏贴复制黏贴复制黏贴复制黏贴复制黏贴

查看完整版本: 【Rust】在Rust里自定义全局内存分配器