Rust FFI 包装推理库:unsafe 边界要像防火墙一样清楚
Rust FFI 包装推理库:unsafe 边界要像防火墙一样清楚
很多高性能推理库是 C/C++ 写的,Rust 服务要复用它们,就绕不开 FFI。FFI 本身没问题,问题在于把 unsafe 扩散到业务代码里。指针生命周期、内存释放、线程安全、错误码转换,任何一处没封好,都可能把 Rust 的安全边界打穿。
我的原则是:unsafe 只出现在最小封装层,业务层拿到的是安全 API。unsafe 边界要像防火墙一样清楚,不能到处漏风。
一、先定义 C 接口的所有权
FFI 最怕所有权说不清。谁分配,谁释放?返回指针能活多久?调用是否线程安全?这些都要写进接口约定。
flowchart TD A[Rust Safe API] --> B[FFI Wrapper] B --> C[unsafe extern call] C --> D[C/C++ Runtime] B --> E[错误码转换] B --> F[Drop 释放资源]Rust wrapper 的任务,就是把不可靠的边界收窄,并把约定固化成类型。
二、用 RAII 管理句柄
C 库常返回 handle,Rust 里应该用结构体包起来,并在 Drop 中释放。
pub struct ModelHandle { raw: NonNull<c_void>, } impl Drop for ModelHandle { fn drop(&mut self) { unsafe { ffi_model_destroy(self.raw.as_ptr()) } } }这里 unsafe 仍然存在,但范围很小。调用方不能忘记释放,也不能随便拿 raw pointer 玩。
三、输入输出要检查长度
传 tensor buffer 时,长度、对齐、dtype 都要检查。不要相信调用方。
pub fn run(&self, input: &[f32], output: &mut [f32]) -> Result<()> { if input.len() != self.input_len { return Err(Error::InvalidInputShape); } let code = unsafe { ffi_model_run(self.raw.as_ptr(), input.as_ptr(), output.as_mut_ptr()) }; Error::from_code(code) }这类检查看起来啰嗦,但它把崩溃变成了可处理错误。系统级代码最怕"相信上游"。
输入验证的边界不止于 len 检查。在推理场景中,tensor buffer 可能来自共享内存、DMA 区域或另一进程的 mmap,对齐要求往往比标准 malloc 严格——比如 256 字节对齐用于 GPU DMA 传输。如果 FFI 层不校验对齐,kernel launch 会在 CUDA 内部静默失败或产生错位结果。另一个容易被忽略的检查是 dtype 兼容性:下游 C 库期望 f32,但 Rust 侧传入了从 bf16 字节重解释的&[f32],Slice 不会报错,但计算结果完全错误。建议在 FFI 边界的前置校验中加入 alignment check(ptr as usize % required_alignment == 0)和 dtype 标签校验,用枚举而非裸整数传递数据类型,让编译器帮你挡掉类型不匹配。对于 GPU 侧的 pinned memory 输入,还要验证指针是否确实在 pinned 区域——这可以通过cudaPointerGetAttributes查询,避免 kernel 内部因非 pinned 内存的隐式拷贝导致延迟陡增。除了输入校验,输出 buffer 的治理同样重要:C 库写入的 output tensor 若有越界写行为,Rust 侧难以检测,建议在 debug 编译时用 canary page 或 AddressSanitizer 包裹输出 buffer,捕捉越界写;生产环境则在 wrapper 层加入 output 校验和,定期抽样比对,发现异常立即告警并隔离对应 handle,防止错误结果污染业务决策。
四、线程安全要显式声明
不是所有 C handle 都能跨线程。Rust 的Send、Sync不能随便 unsafe impl。只有确认底层库线程安全,才能声明。
如果底层不支持并发,就在 wrapper 里加 Mutex 或要求每线程一个 handle。不要为了通过编译器,把不确定性塞进unsafe impl Send。
FFI 还要处理 panic 边界。Rust panic 不能跨过 C ABI 边界,C++ exception 也不能随便穿进 Rust。回调函数尤其要小心,必要时用catch_unwind把 panic 转成错误码。
let result = std::panic::catch_unwind(|| { user_callback(token_id) }); if result.is_err() { return FFI_CALLBACK_PANIC; }边界代码要宁可啰嗦,也不要让未定义行为混进推理服务。
五、总结
Rust FFI 包装推理库时,unsafe 边界要小、清楚、可审查。所有权、Drop、长度检查、错误码、线程安全,都要在封装层处理。
Rust 的安全不是自动延伸到 C 库里的。边界守住,Rust 才能继续帮你挡 Bug。
