• 时间:
  • 分类:

本文基于 Cloudflare 官方博客，介绍其 Bot 管理模块从 Lua 迁移到 Rust 过程中，如何通过消除内存分配将机器学习推理延迟降低 20%。

原文链接：https://blog.cloudflare.com/how-cloudflare-runs-ml-inference-...

每一微秒都有代价

Cloudflare 的边缘节点承载着海量请求，每一个都要经过一系列安全检查才能被转发或拦截。其中 Bot 管理模块（Bot Management）位于请求处理的关键路径上（hot path），它对每个请求打分，判断其是否来自机器人。

这个模块的工作量不轻：它需要计算多种启发式特征，并调用多个机器学习模型。问题在于，它加在每个请求上的延迟，是直接叠加到用户感知的响应时间里的。对于一个以低延迟为核心卖点的 CDN 来说，这是不能视而不见的成本。

Cloudflare 决定把 Bot 管理模块从 Lua 重写为 Rust，并在迁移过程中进行了系统性的性能优化。本文聚焦其中最有意思的一条优化主线：彻底消除热路径上的内存分配，将 P50 延迟从 388μs 降低到 309μs，减少了 79μs，降幅 20%。

为什么内存分配是问题？

在高级语言里，内存分配往往是隐形的——你写个 String::new() 或者 Vec::new()，运行时替你做了一切。但这"一切"背后，实际上包含了相当多的工作：

  • 在堆上寻找足够大的连续空闲空间
  • 处理内存碎片
  • 必要时向内核申请新的内存页

对于带有垃圾回收的语言（比如 Lua），还要额外追踪对象的生命周期，并在 GC 触发时暂停程序执行。

在每秒处理数千万请求的场景下，热路径上每一次额外的内存分配，都在以极高的频率重复消耗这些开销。Cloudflare 的目标是：让 Bot 管理模块的核心计算路径做到零内存分配（zero allocation）。

优化一：用栈代替堆

最直接的减少分配方式，是把固定大小的缓冲区放到栈上，而不是每次都在堆上重新申请。

栈分配只是在当前函数栈帧里预留空间，由编译器完成，没有任何运行时开销：

let mut buf = [0u8; BUFFER_SIZE];

如果缓冲区大小不确定，可以在初始化阶段一次性分配到位，之后反复复用：

let mut buf = Vec::with_capacity(BUFFER_SIZE);

具体有多大差距？ 用一个大小写不敏感的字符串比较来说明。

每次都新建缓冲区：

fn eq_alloc(s: &str, pat: &str) -> bool {
    let mut buf = String::with_capacity(s.len());
    buf.extend(s.chars().map(|c| c.to_ascii_lowercase()));
    buf == pat
}

复用已有缓冲区：

fn eq_reuse(s: &str, pat: &str, buf: &mut String) -> bool {
    buf.clear();
    buf.extend(s.chars().map(|c| c.to_ascii_lowercase()));
    buf == pat
}

基准测试结果：前者约 40ns/次，后者约 25ns/次——仅改变内存分配行为，速度提升 38%。

优化二：选择不需要缓冲区的算法

更进一步，有些操作根本不需要中间缓冲区。还是以大小写不敏感比较为例，用迭代器可以完全在栈上完成：

fn eq_iter(s: &str, pat: &str) -> bool {
    s.chars().map(|c| c.to_ascii_lowercase()).eq(pat.chars())
}

这个版本不仅代码最短，速度也最快：约 13ns/次，接近标准库 eq_ignore_ascii_case 的 11ns。它们快的原因是一样的——都靠迭代器在原始数据上直接操作，没有任何数据拷贝或额外分配。

这背后的思路是：在选择算法时，把"是否需要额外内存"作为一个重要维度纳入考量，而不只是看时间复杂度。

优化三：用测试守住零分配

光靠人工审查来保证"热路径零分配"是不可靠的，代码演化中随时可能引入新的分配。Cloudflare 用 Rust 的 dhat crate 写了自动化测试来做保障：

#[test]
fn zero_allocations() {
    let _profiler = dhat::Profiler::builder().testing().build();

    // 执行热路径逻辑
    run_hot_path_logic();

    let stats = dhat::HeapStats::get();
    // 断言没有发生任何内存分配
    dhat::assert_eq!(stats.total_blocks, 0);
    dhat::assert_eq!(stats.total_bytes, 0);
}

dhat 通过替换全局分配器来监控所有堆分配。只要热路径触发了任何堆分配，测试就会失败。这条测试在 CI 里常驻，成为防止性能回退的自动护栏。

需要注意的是，dhat 只能检测 Rust 侧的分配，通过 FFI 调用 C 库时产生的分配不在其监控范围内——这正是下面要讲的优化重点。

优化四：改造 CatBoost 的 C++ 核心

Bot 管理模块使用 CatBoost 这个开源机器学习库来执行决策树模型。CatBoost 的核心是 C++ 实现，通过 FFI 供 Rust 调用。

原始的 CatBoost API 被设计成批量处理：一次评估多个样本，因此在内部会分配 vector 来存储中间结果：

TVector<TConstArrayRef<float>> floatFeaturesVec(docCount);
TVector<TConstArrayRef<int>> catFeaturesVec(docCount);
for (size_t i = 0; i < docCount; ++i) {
    floatFeaturesVec[i] = TConstArrayRef<float>(floatFeatures[i], floatFeaturesSize);
    catFeaturesVec[i] = TConstArrayRef<int>(catFeatures[i], catFeaturesSize);
}
FULL_MODEL_PTR(modelHandle)->Calc(floatFeaturesVec, catFeaturesVec, result);

但 Cloudflare 的场景是每次只评估一个请求，根本不需要批处理逻辑。评估单个样本时，只需要一个指向连续内存的引用，完全不需要那些 vector：

TConstArrayRef<float> floatFeaturesArray(floatFeatures, floatFeaturesSize);
TConstArrayRef<int> catFeaturesArray(catFeatures, catFeaturesSize);
FULL_MODEL_PTR(modelHandle)->Calc(floatFeaturesArray, catFeaturesArray, result);

这个改动直接消除了 CatBoost 内部的动态分配，生产模型的推理速度因此提升了约 15%。

优化五：重构 Rust 绑定层，彻底复用缓冲区

问题不只在 C++ 层，Rust 绑定层也存在类似的设计缺陷。

原始的 Rust 绑定在处理多文档时会分配一个指针 vector：

let mut float_features_ptr = float_features
    .iter()
    .map(|x| x.as_ptr())
    .collect::<Vec<_>>();

对于单文档场景，中间那层 vector 完全多余，直接取内层指针即可：

let float_features_ptr = float_features.as_ptr();

更重要的是，原始 API 接受的参数类型是所有权形式（Vec<Vec<f32>>、Vec<Vec<String>>），意味着每次调用前都必须分配并填充新的数据结构：

pub fn calc_model_prediction(
    &self,
    float_features: Vec<Vec<f32>>,
    cat_features: Vec<Vec<String>>,
) -> CatBoostResult<Vec<f64>> { ... }

重构后改为借用（&[f32]、&[&str]），调用方可以直接传入已有数据的引用，无需任何拷贝或额外分配：

pub fn calc_model_prediction_single(
    &self,
    float_features: &[f32],
    cat_features: &[&str],
) -> CatBoostResult<f64> { ... }

还有一个细节：每个请求需要调用多个模型，这些模型往往共享相同的类别特征（categorical features）。特征值转哈希的操作只需要做一次，然后把哈希值传给每个模型复用：

pub fn calc_model_prediction_single_with_hashed_cat_features(
    &self,
    float_features: &[f32],
    hashed_cat_features: &[i32],
) -> CatBoostResult<f64> { ... }

结果：数字说明一切

上述所有优化叠加之后，效果如下：

指标优化前优化后降幅
P50 延迟388μs309μs20%
P99 延迟940μs813μs14%

79μs 的 P50 改善听起来不多，但对于每个请求都要经过这条路径的全球 CDN 来说，这个数字背后是极其可观的用户体验收益。

几点值得借鉴的工程思路

这次优化背后有几个值得提炼的工程判断：

不要假设库的默认行为适合你的场景。 CatBoost 为批量处理设计的 API 对 Cloudflare 的单请求场景完全是多余的开销。当你接入一个外部库时，值得问一句：它的默认路径是否和你的使用场景匹配？

API 设计影响性能上限。 接受所有权参数（Vec<T>）的 API，天然迫使调用方每次都做分配和拷贝；接受引用（&[T]）的 API 才能让调用方自由控制内存生命周期。这个差别在热路径上会被指数级放大。

用测试把性能约束固化下来。 "零分配"不是一个靠 code review 就能长期维持的属性，它需要自动化测试来守护，否则随着代码演进，迟早会悄悄被破坏。

语言选择很重要，但不是全部。 从 Lua 换到 Rust 提供了精细控制内存的可能性，但这种可能性需要通过具体的设计决策才能兑现成实际收益。换了语言本身不会自动让代码变快。

标签: none

添加新评论