在 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 特性的自定义实现来定义自定义内存分配器的选项

我们可以通过使用 GlobalAlloc 特性来实现，方法是我们使用自己的将分配委托给 Redis 的方法。为此，我们需要一种方法来从 Rust 中调用 Redis 模块 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 中，具有相同名称的多个函数可以共存，因为有各种命名空间机制（例如模块和特征）来区分它们。为了生成与 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
...

我仍花了很长时间苦思冥想和尝试才弄明白，但以下便是发生的事情

通过 C 函数指针访问 Redis 模块 API 的函数。这些指针不会依赖于动态链接程序进行初始化，而是由 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，该 AtomicBool 默认状态为 false。

在模块初始化代码中，在模块准备好使用时，我们调用 use_redis_alloc 。这时，我们可以安全地开始使用 Redis 分配器了，Redis 将会管理以后所有的分配。

这将处理崩溃，并最终在 redis-module 中结束。请随意查看并告诉我你的感受！