在 Rust 中使用 Redis 分配器

引言

在开发 redismodule-rs（用于编写 Redis 模块的 Rust API）时，我遇到了需要设置自定义内存分配器的情况。

通常，当 Rust 程序需要分配一些内存时，例如创建 String 或 Vec 实例时，它会使用程序中定义的全局分配器。由于 Redis 模块是作为共享库构建以加载到 Redis 中，Rust 将使用 System 分配器，这是操作系统提供的默认分配器（使用 libc 的 malloc(3) 函数）。

这种行为会带来几个问题。

首先，Redis 可能根本没有使用系统分配器，而是依赖 jemalloc。 jemalloc 分配器是系统 malloc 的替代品，包含许多优化以避免碎片，还有其他功能。如果模块使用系统分配器而 Redis 使用 jemalloc，则分配行为将不一致。

其次，即使 Redis 一直使用系统分配器，模块直接分配的内存对 Redis 也是不可见的：它不会显示在 info memory 等命令中，也不会受到 Redis 执行的清理操作（如键的逐出）的影响。

出于这些原因，Redis Modules API 提供了诸如 RedisModule_Alloc 和 RedisModule_Free 之类的钩子。这些钩子的用法与标准的 malloc 和 free 调用非常相似，但除了实际将调用传递给内存分配器之外，还让 Redis 了解分配的内存。

使用自定义分配器

Rust 提供了定义自定义内存分配器的选项，只需提供 GlobalAlloc trait 的自定义实现即可。

我们可以通过实现 GlobalAlloc trait，并在其自定义方法中将分配委托给 Redis 来使用它。为此，我们需要一种从 Rust 调用 Redis Module API 函数的方法。这是一个后续文章的主题，但简而言之，我们通过使用 bindgen crate 从 redismodule.h C 头文件生成 Rust 绑定来实现此目的。

头文件将函数定义如下

#define REDISMODULE_API_FUNC(x) (*x)

void *REDISMODULE_API_FUNC(RedisModule_Alloc)(size_t bytes);
void REDISMODULE_API_FUNC(RedisModule_Free)(void *ptr);

这些函数，与 Modules API 的其余部分一样，被定义为函数指针。从 Rust 调用这些函数时，我们首先需要解引用函数指针，这可以通过 unwrap() 方法完成。我们还需要进行一些类型转换以匹配指针类型。最后，由于我们解引用了裸指针，需要使用 unsafe 关键字，这在安全的 Rust 中是不允许的，出于很好的原因。

use std::alloc::{GlobalAlloc, Layout};
use std::os::raw::c_void;

struct RedisAlloc;

unsafe impl GlobalAlloc for RedisAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        RedisModule_Alloc.unwrap()(layout.size()) as *mut u8
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        RedisModule_Free.unwrap()(ptr as *mut c_void)
    }
}

崩溃

不幸的是，事情并非如此简单。当我们使用这个自定义分配器构建模块并将其加载到 Redis 中时，它会崩溃。Redis 在崩溃时会打印出一个不错的堆栈跟踪，所以让我们看看它。

$ redis-server --loadmodule ./target/debug/examples/libhello.dylib
thread panicked while processing panic. aborting.
...
Backtrace:
0   redis-server                        0x000000010adce1dc logStackTrace + 110
1   redis-server                        0x000000010adce562 sigsegvHandler + 236
2   libsystem_platform.dylib            0x00007fff5b95db3d _sigtramp + 29
3   ???                                 0x0000000000000000 0x0 + 0
4   libhello.dylib                      0x000000010af5af6d _ZN3std9panicking18continue_panic_fmt17h0e74ab2b215a1401E + 157
5   libhello.dylib                      0x000000010af5ae69 rust_begin_unwind + 9
6   libhello.dylib                      0x000000010af6ea3f _ZN4core9panicking9panic_fmt17h09741a3213dba543E + 63
7   libhello.dylib                      0x000000010af6e984 _ZN4core9panicking5panic17hb4bc64e7f35c9151E + 100
8   libhello.dylib                      0x000000010af53108 _ZN4core6option15Option$LT$T$GT$6unwrap17h66957b4d942a4d3cE + 56
9   libhello.dylib                      0x000000010af4f8f3 _ZN76_$LT$redis_module..alloc..RedisAlloc$u20$as$u20$core..alloc..GlobalAlloc$GT$5alloc17h6588ea2d7520a3ebE + 35
...

所以，看起来这里发生了空指针解引用（3 ??? 0x0000000000000000 0x0 + 0），但是这些以 _ZN… 开头的奇怪符号是什么呢？

经过一番搜索，我们发现这是 Rust 进行名称修饰的方式：与 C 不同，类似于 C++，在 Rust 中，多个同名函数可以共存，因为有各种命名空间机制，如模块和 trait 来区分它们。为了生成与 C 兼容的唯一符号，编译器将它们修饰成长而丑陋的唯一名称。要将这些名称还原回原始形式，我们可以通过 rustfilt 过滤输出。

这给了我们以下堆栈跟踪（已移除不重要的部分）

3   ???                                 0x0000000000000000 0x0 + 0
...
31  libhello.dylib                      0x000000010af6e984 core::panicking::panic + 100
32  libhello.dylib                      0x000000010af53108 core::option::Option<T>::unwrap + 56
33  libhello.dylib                      0x000000010af4f8f3 <redis_module::alloc::RedisAlloc as core::alloc::GlobalAlloc>::alloc + 35
34  libhello.dylib                      0x000000010af4cc8c __rg_alloc + 60
35  libhello.dylib                      0x000000010af6e2f6 <alloc::vec::Vec<u8> as core::convert::From<&str>>::from + 38
36  libhello.dylib                      0x000000010af4de54 <T as core::convert::Into<U>>::into + 36
37  libhello.dylib                      0x000000010af4f57f std::ffi::c_str::CString::new + 47
38  libhello.dylib                      0x000000010af40daa RedisModule_OnLoad + 58
39  redis-server                        0x000000010adf97d9 moduleLoad + 118
40  redis-server                        0x000000010adf9735 moduleLoadFromQueue + 69
41  redis-server                        0x000000010ad94428 main + 1190
...

我仍然挠头并进行了大量实验才弄明白，但事情是这样的：

Redis 模块 API 的函数通过 C 函数指针访问。它们不是依赖动态链接器来初始化，而是由 Redis 在模块初始化过程中显式初始化。

正如堆栈跟踪所示，在加载模块期间，我们调用了 CString::new 函数。这个标准库函数为字符串分配内存。这反过来会调用我们的分配器，然后由我们的分配器调用 RedisModule_Alloc.unwrap()… 来实际执行分配。这就造成了一个鸡生蛋蛋生鸡的问题。Redis 模块尚未准备好，意味着我们的函数指针尚未初始化，因此我们无法调用相关的 API 来执行分配。

解决方案

我尝试了各种方法来解决这个问题，但似乎没有一种干净的方式可以避免模块初始化期间的内存分配。次优的方法是先使用标准分配器直到模块准备好，然后切换到自定义分配器。然而，Rust 不允许在运行时更改分配器，所以我们无法做到这一点。

我最终在自定义分配器中添加了一个标志，使得在启动时将分配请求传递给系统分配器。模块初始化完成后，该标志被切换，以便后续的分配通过 Redis 分配器执行。这个解决方案仍然存在一些边缘情况——最重要的是要求在切换之前释放所有先前分配的内存，否则会导致内存泄漏。然而，对于我们的目的来说，这已经足够了。

最终的代码如下所示

use ...;
use std::sync::atomic::{AtomicBool, Ordering::SeqCst};

pub struct RedisAlloc;

static USE_REDIS_ALLOC: AtomicBool = AtomicBool::new(false);

unsafe impl GlobalAlloc for RedisAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let use_redis = USE_REDIS_ALLOC.load(SeqCst);
        if use_redis {
            return raw::RedisModule_Alloc.unwrap()(layout.size()) as *mut u8;
        }
        System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let use_redis = USE_REDIS_ALLOC.load(SeqCst);
        if use_redis {
            return raw::RedisModule_Free.unwrap()(ptr as *mut c_void);
        }
        System.dealloc(ptr, layout);
    }
}

pub fn use_redis_alloc() {
    USE_REDIS_ALLOC.store(true, SeqCst);
    eprintln!("Now using Redis allocator");
}

我们添加了一个名为 USE_REDIS_ALLOC 的 static 标志，它决定了我们应该使用 Redis 分配器还是系统分配器。在修改静态数据时保证安全性非常重要，所以我们在这里使用了一个 AtomicBool，它默认为 false。

在模块初始化代码中，当模块准备好使用时，我们调用 use_redis_alloc。此时，我们可以安全地开始使用 Redis 分配器，并且所有未来的分配都将由 Redis 统计。

这解决了崩溃问题，并最终被纳入了 redis-module crate。欢迎查看并告诉我你的使用体验！