Mooncake 统一内存池：从默认 Evict 到 Linux Reclaim

之前那篇 Mooncake 统一内存池:AI Vibe Coding 与 Rust 里, 我主要讲的是 local master + unify memory pool 怎么把 Pod 内分裂的数据面收拢成统一共享池 G. 在我们的系统中，Unify memory pool 已经逐步替代推理/训练系统里所有需要 shared memory 交换数据的地方，在大幅降低代码复杂度/架构复杂度的同时，也带来了不少的性能提升。那篇文章写完之后, 其实还有一个问题躲不过去: 当所有本地 replica 都被收拢到 G 里之后, evict 到底该怎么做?

一开始看上去, 这像是一个“给 Mooncake 加个 LRU”的问题. 但真正把语义扒开来看就会发现, 事情远比“淘汰几个冷 key”复杂. 因为 unified pool 下被驱逐的不是普通 cache entry, 而是仍然可能被用户 handle 引用、可能带着 softpin, 甚至还可能正处在 put -> remove -> put 多代并存状态里的本地 replica.

问题不只是“淘汰几个冷 key”

在 unify memory pool 模式下, 当前节点本地内存池 G 里的对象, 大致可以分成两类:

这两类对象共享同一块 G, 但显然不能用同一把删除锤子去处理.

对 softpin=true 来说, 我们真正需要的是“把对象从当前节点内存里挪走”, 而不是“让对象从系统里消失”. 也就是说, evict 的语义首先应该是 migrate 或 offload, 最后才是释放当前节点 G 中那块内存.

对 softpin=false 来说, 语义又是另一回事. 这类对象更像 cache, 本来就允许 miss. 如果当前节点只是某个非主副本, 那么在内存压力下把它直接丢掉, 完全是合理的.

所以这里真正的问题, 从来都不是“有没有淘汰动作”, 而是:

1. 哪些对象应该继续留在本地内存里;
2. 哪些对象可以直接丢;
3. 哪些对象必须先迁移出去才能释放;
4. 这些判断能不能在不把在线 Put/Get 链路拖重的前提下完成.

默认 evict 为什么不够

Mooncake 默认的 evict, 在我看来更接近“遍历 key 空间并尝试删除本地副本”, 而不是一个真正对冷热有建模的回收器. 这意味着扫描成本和 key 空间大小更接近线性关系, 而不是和本轮真正候选的对象数量相关. 由于没有一个稳定的近似 LRU 机制. 热点和冷点并没有通过独立的数据结构持续维护起来. 所以有些对象之所以需要额外打 softpin, 本质上是因为默认驱逐器并不能自然地把“经常访问的对象应该尽量留下来”这件事做得足够稳定.

其次, softpin=true 在 lease 未过期时会被整体跳过. 这个行为放在 RL/Omni/EPD 等场景中问题尤其大. 因为 put 往往集中发生在数据生产节点, 而这些节点上的大量对象又正好是 softpin=true 的中转数据. 结果就是: 数据生产节点上的 G 很快被塞满, 但默认 evict 既不会删这些对象, 也不会把它们迁到别的节点去, 后续 PutStart 就只能分配失败. 当然本质上是我们把 softpin 语义改了，导致默认 evict 策略不能满足我们的需求。

再往后, 默认 evict 的动作本质上仍然比较接近“删除本地副本”, 而不是“给当前节点腾空间”. 这两句话看起来很像, 其实差别很大. 前者没有 remote swapout / offload fallback 的概念, 后者才真正关心“对象离开当前 G 之后, 还能不能在别处被恢复”.

最后一个问题更工程化一点: 如果驱逐逻辑一直围着 metadata 表边扫边判边删, 那么你几乎不可避免要在持锁状态下做很多判断. 要保证“看见的元数据状态”和“最后提交的删除动作”不乱掉, 很多步骤基本都只能围着锁转. 这种结构天然不利于继续往近似 LRU、远端迁移、offload fallback 这些方向演化. 说白了就是: 扫描、分类、真正提交删除这几件事没拆开, 锁范围就很难看, 在线 Put/Get 也很难真正瘦下来.

这不是说默认 evict 完全不能用, 而是它和 unify memory pool 面临的问题已经不太匹配了.

先把 evict 的语义说对

我觉得这一步很重要. 因为只要一开始把 evict 理解成“删除对象”, 后面很多设计都会走偏. 在 unify memory pool 这里, evict 的准确定义应该是: 把对象从当前节点的统一内存池 G 中移走, 而不是把对象从系统中删除. 沿着这个定义, 两类对象的行为就清楚了:

1. softpin=true: 不能直接丢. 只有 remote swapout（写入到其他节点内存中）或 offload（写入到下级存储中）成功之后, 当前节点内存才可以释放.
2. softpin=false: 保留 cache 语义. 非主副本可以直接删本地副本; 主副本仍然要先 remote swapout/offload 成功, 再释放本地内存.

也正是因为这个语义变化, softpin=true 就不再等价于“永远不能参与回收”, 而变成了“不能无损语义地直接丢弃”. 我觉得这是这次改造里最关键的一步. 之前的问题恰恰就在于, Mooncake 把 softpin 理解成“别碰”, 我们这里则把它重新解释成“可以迁走, 但不能消失”.

借 Linux reclaim, 但不是照抄 Linux

Mooncake 在这里要管理的, 其实就是一块有容量上限的本地共享内存池. 它有后台线程, 有前台分配失败, 有冷热差别, 也有暂时不能回收的对象. 从职责上看, 这和 Linux reclaim 已经很像了. 所以我们最后借的, 不是某个具体内核函数的实现细节, 而是 Linux 那套已经被反复验证过的职责划分:

1. 用 active list / inactive list 维护冷热度, 而不是每次从全量 key 空间里重新猜.
2. 把后台 evict 和前台 direct reclaim 分开, 但让它们共享同一套冷热和候选选择逻辑, 对应 Linux 里的 kswapd 和 direct reclaim.
3. 把“访问一次”和“真正变成热点”区分开, 思路上对应 folio_mark_accessed.
4. 在真正检查对象是否可驱逐之前, 先把一批候选从共享链表里 isolate 出来, 然后锁外扫描, 这件事基本就是在借 isolate_lru_pages 的味道.

真正像 Linux 的地方, 不是名字里有了 active / inactive 就算完事, 而是我们也把“先隔离候选, 再在锁外慢慢看”这件事学了过来. 这背后的动机其实很朴素: 我们希望 Put/Get 这些在线链路尽可能只做最轻的事情, 而把“找谁该被回收”“该迁到哪里”“offload 要不要兜底”这些重活脏活, 尽量塞到 reclaim 路径里.

让 Put/Get 在线链路尽量少做事

这一点是我这次最在意的地方. 在线链路最怕的, 不是多一两个字段更新, 而是顺手做太多事情: 顺手扫一遍 key, 顺手做远端迁移, 顺手写 offload, 顺手长时间拿锁. 这些“顺手”最后基本都会变成 tail latency. 所以我们把在线路径压得很薄:

PutEnd
  -> 为当前 COMPLETE memory replica 创建 KeyEvictInfo
  -> 插入 inactive 头部

AcquireReplica / GetReplicaList hit local
  -> mark_accessed
  -> 若对象仍在 inactive 且再次被访问, promote 到 active

LocalRemove
  -> 走完现有 remove 语义
  -> 幂等地从 evict list 摘掉

也就是说:

• PutEnd 不负责决定谁该被驱逐, 它只负责把新对象纳入冷热系统.
• Get 命中本地时不做扫描, 不做 I/O, 只做一次非常轻量的热度更新.
• Remove 也不顺手做一堆 reclaim 工作, 它只把当前代从 evict 结构里摘掉.

这里最像 Linux folio_mark_accessed 的一点是: 一次访问不会立刻把对象抬成长期热点. 我们这里也是类似思路. 新对象先进入 inactive, 第一次访问只是打个 referenced 标记, 只有第二次访问且对象还在 inactive 时, 才把它 promote 到 active. 这样一次性的 get 不会轻易污染热点判断.

唯一的例外是 PutStart 真的分配不出内存的时候. 这时允许触发 direct reclaim, 但即便如此我们也不想让它什么都干. 所以前台 direct reclaim 只碰 inactive, 不去顺手 shrink active. 原因也很简单: 前台分配失败时, 第一目标是先腾出空间, 而不是顺便把热点集合打散.

更进一步, direct reclaim 也不是所有路径都能开. 单次读 miss 这种没有 fallback 的路径, 会更希望 PutStart 尽量成功, 因此可以允许 allow_direct_reclaim=true; batch read 和 remote swapout 写远端时则不行, 否则很容易把 reclaim 做成环路.

我们的 evict

最后落地出来, 这套 evict 大概可以分成四层:

回收器维护的不是全局大一统的一条 LRU, 而是每个 MetadataShard 一组 active/inactive list. 一开始我也想过是不是做成全局更直接, 但最后还是觉得全局链表的锁冲突太危险了. 既然 metadata 本来就是 shard 化的, 那热度结构也跟着 shard 化更自然. 这样做之后, 每个 shard 只关心自己的本地 COMPLETE memory replica. 整体需要回收多少字节, 再按 shard.used_bytes / sum(all shard.used_bytes) 去分 quota. 这就把全局压力和局部锁冲突拆开了. 如果只讲 active/inactive list, 很容易让人误以为这次实现不过是“给 Mooncake 加一个 LRU”. 实际上真正麻烦的部分, 根本不在链表, 而在并发状态机.

第一层麻烦是, evict 维护的链表状态和 metadata 状态不能共用一把大锁. 如果把这两件事全塞进同一把锁里, reclaim 路径很快就会变成一个大临界区. 所以最后是 ShardEvictInfo 自己管 active/inactive list, metadata 继续走自己的 shard 锁. 这当然会带来短暂不一致, 但这是我们主动接受的复杂度.

第二层麻烦是, 一旦你想把扫描阶段挪到锁外, 你就必须先把一批候选 isolate 出来. 我们这里 shrink_active_list() / shrink_inactive_list() 的核心做法, 就是在 seinfo.mutex 下把整条链表 swap 到线程私有临时 list, 然后把 epoch 自增. 之后其他线程即使拿到了 seinfo.mutex, 只要发现某个项的 epoch 不对, 就知道它现在正处于 isolate 状态, 不该乱改. 这一步我觉得和 Linux isolate_lru_pages 很像, 也是这次实现里最有意思的一段. 它的好处很直接: 真正慢的那部分检查, 比如去 metadata 里确认这个对象是不是还存在, 还能不能驱逐, 有没有 lease, 可不可以删, 都可以放到锁外去做. 锁内只保留 list 的组织动作.

第三层麻烦是, 哪些字段可以锁外碰, 哪些字段绝对不能乱动, 必须卡得很严. 比如 refered_ 只是冷热近似判定, 锁外 relaxed 读写问题不大; 但 lru_、epoch_、iter_ 这些字段如果锁外改了, 很容易直接把链表状态搞坏. 再比如 list 节点在不同临时 list 之间迁移时, 必须用 std::list::splice, 不能图省事 erase + insert, 否则 iterator 立刻失效.

第四层麻烦是, local master 本来就允许同 key 多代共存. 所以 evict 不能天真地假设“同一个 key 只有一条真实状态”. 有些已经 remove 掉的旧代 KeyEvictInfo, 甚至可能因为并发 AcquireReplica 拿着 stale 引用, 暂时又被挂回到 evict 结构里. 这件事在 v1 里我们没有额外上 tombstone 或 generation state 去彻底堵死, 而是接受它, 并在下一轮 evict 里依据 metadata 不存在或者 replica id 不匹配再把它清掉. 这听上去有点土, 但它是一个非常现实的工程取舍.

换句话说, 这次实现最难的部分不是“选个 LRU 还是 FIFO”, 而是如何在不把锁扩大成灾难的前提下, 让 reclaim 状态机还能保持自洽. 有了冷热链表之后, 回收入口就可以统一成 LocalBatchEvict(), 但回收场景并不只有一种.

后台线程对应的是 Linux kswapd 那类角色. 它的目标是尽量提前把冷对象迁走或者落到 offload, 让前台 PutStart 在稳态下只做本地分配. 这条路径会做完整的 shrink_active_list + shrink_inactive_list, 因为后台有条件更温和地重新整理冷热状态.

另一条路径是 direct reclaim. 也就是 PutStart 分配失败之后, 前台同步发起回收, 回收到足够字节之后再重试分配. 这条路径和后台 reclaim 共用同一套数据结构, 但不能粗暴复用同一条执行流程. 一个很关键的区别是: direct reclaim 只收 inactive, 不主动 shrink active. 否则如果前台压力一大, active 很容易被频繁打散, 原本应该稳稳留在热点集里的对象也会被过早地下放. 另外多个线程同时因为 PutStart 失败而发起 direct reclaim 时, 我们这里会共用一次 prepare 阶段, 这个味道其实有点像 group commit.

又一次从 Linux 借设计

后来想想, 这其实也不是我第一次从 Linux 借设计了.

之前写 深入浅出 KuiBaDB: 使用 SharedBuffer 时, 我就一边看 PostgreSQL, 一边把问题往 Linux address_space 上靠. 再后来折腾 在 tokio 上几个失败尝试 时, 参考的也是 Linux CFS scheduler、PELT、load balance 那一套思路.

这次轮到 Mooncake unify memory pool 的 evict, 借的则是 reclaim. 挺有意思的. 从 shared buffer, 到 scheduler, 再到内存回收, 兜兜转转又回到 Linux. 这倒也不奇怪. 只要问题开始涉及缓存、冷热、回收、调度等这些主题, Linux 内核里往往都已经有一套被反复打磨过很多年的抽象在那里等着你. 你不一定照抄, 但大概率可以从里面借到一个足够结实的起点. 从这个角度看, 这次 evict 做到最后又借到了 Linux reclaim, 反而说明之前这些内核笔记和折腾没有白看.

• Previous
  Rapid transfer_kv：只为更快的 PD 分离
• Next
  Mooncake Evict: 一次 std::make_pair 让 iter_ 悄悄失效