diff --git a/.gitignore b/.gitignore index d198634..fa0d6cb 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,6 @@ bench_suite.lock deployment/local/ setup_and_pack/pack_fluxonkv_pylib_env.yaml fluxon_rs/moka/ +fluxon_rs/*.btr +fluxon_rs/moka.bak_non_git_*/ +build_*.rc diff --git "a/fluxon_doc_cn/design/sglang_fluxon_kv_PegaFlow\345\257\271\347\205\247\344\270\216Get\350\267\257\345\276\204\347\256\200\345\214\226\345\210\206\346\236\220.md" "b/fluxon_doc_cn/design/sglang_fluxon_kv_PegaFlow\345\257\271\347\205\247\344\270\216Get\350\267\257\345\276\204\347\256\200\345\214\226\345\210\206\346\236\220.md" new file mode 100644 index 0000000..95e9aab --- /dev/null +++ "b/fluxon_doc_cn/design/sglang_fluxon_kv_PegaFlow\345\257\271\347\205\247\344\270\216Get\350\267\257\345\276\204\347\256\200\345\214\226\345\210\206\346\236\220.md" @@ -0,0 +1,469 @@ +# Fluxon Batch Get / singleflight：PegaFlow 对照与路径简化分析 + +本文评审当前工作区新增的 owner 侧 Batch Get / per-key singleflight 路径，并与 PegaFlow 的 prefix query、prefetch、lease 和批量传输实现对照。评审范围从 `ExternalBatchGetStartReq` 进入 owner 开始，到 master `BatchGetStart`、payload transfer、`BatchGetDone` / `BatchGetRevoke` 以及 external handle 消费结束。 + +原始评审基于 2026-07-16 的 Fluxon 工作区和 PegaFlow `master` 的提交 +`939363f198a0ffded9e3f30d8af9cdd74439f16c`。2026-07-17 已按当前工作树回填实现状态； +没有把尚未执行的性能对比写成实测结论。 + +## 2026-07-17 实现回填（当前权威状态） + +| 原问题 | 当前实现 | 状态 | +| --- | --- | --- | +| exact-batch 与 per-key 两层 shared op | exact-batch phase/waiter/result cache 已删除；`ExternalGetStartEntry` 直接持请求自己的 `keys + Local/Interest items` | 已解决 | +| `decision_registered` 手工归还 | `ExternalGetKeyInterest` 在 Drop 中归还 `undecided`，prefix decision 由 guard ownership 表达 | 已解决 | +| process-wide `owner_key_control` 扫完整 batch/global states | 改为 256-shard `OwnerKeyControlTable`；local/join/leader 每个 key 独立线性化，任何同步锁均不跨 batch 或 `.await`；metrics 使用独立 Weak flight index | 已解决 | +| 规划等待 Revoke 时取消会遗留未启动 leader | `ExternalGetPlanningLeadersGuard` 在 BatchGetStart handoff 前拥有新 leader；Drop 将仍为 `Starting` 的 op 发布为 safe miss、按 Arc identity 清 marker 并唤醒 joiner | 已解决 | +| prepared slot 在 start/revoke 异常路径失去 owner | claim guard、同 operation identity Revoke retry 与 local slot release guard 已覆盖；响应不确定不猜测释放 | 已解决（仍需端到端故障注入） | +| external handle 无界遗留 | handle sweeper 已接入，超时与显式 cancel 复用 plan Drop/interest release 路径 | 已解决（仍需进程死亡测试） | +| payload 每 key 一次 transfer submission | 控制面仍批量，payload backend submission 仍需运行期计数与 benchmark 后决定是否扩展 descriptor batch API | 未解决/待测 | + +下文 §2～§4 保留为“2026-07-16 修复前快照”，用于解释为什么做上述改动；其中出现的 +“当前实现”均指该历史快照，不得覆盖本节和集成设计主文档的当前规约。 + +## 原始结论（2026-07-16 修复前） + +**per-key singleflight 的目标应当保留，当前“两层共享操作”应当收缩。** 重叠但不完全相同的 batch 确实需要按 key 选出 local、joiner 和 leader，同时继续批量执行 master RPC。当前实现又在它外面保留了 exact-batch `ExternalGetStartSharedOp`，因此一次请求同时支付 exact-batch 去重和 per-key 去重的状态、锁、通知、结果缓存与生命周期成本。完全相同的 batch 本来就会自然 join 同一组 per-key flight；第二层去重没有形成独立的正确性边界。 + +**当前首要风险是取消安全。** 常数级性能排在生命周期闭环之后。`decision_registered: bool`、`undecided`、`retained` 和 `waiter_count` 都依赖后续代码手工归还。future 在计数增加后的任一 `.await` 被取消时，没有 `Drop` guard 关闭这次 interest。静态控制流已经足以证明计数可能永久不归零，进而让 leader task、相同 batch、reclaim fence 或 exact-batch registry 一直等待。 + +**结构性 overhead 已经存在，但实际延迟幅度仍需 benchmark。** 当前路径包含多份 key 向量和结果向量复制、每个新 leader key 的独立共享对象、两层 `Mutex + Notify`、通常两个后台 task、全局 owner key fence 内的线性扫描，以及每个远端 key 一次 `transfer_data_no_copy` 调用。它们在短 value、高并发、大 batch 场景更可能可见；没有 allocator、锁等待和端到端延迟数据前，不应写成“已经造成 X% 回退”。 + +**PegaFlow 最值得学习的是所有权与生命周期的组织方式。** 它没有解决不同请求之间重叠 block 的 per-key singleflight，算法不能直接搬到 Fluxon。可借鉴的是：一次 prefix scan 直接取得 `Arc` pin、一个 ready prefix 对应一个 opaque lease、资源计数由 RAII / TTL 收口、request-level prefetch 只保留一个 task，以及数据面先汇总 descriptor 再提交。PegaFlow 当前 `QueryPrefetch` 在 probe 阶段就 pin、active prefetch 只按 `req_id` 标识、lease TTL 为 600 秒，这些都不应直接复制到 Fluxon。 + +建议按以下顺序推进： + +1. **P0，先封闭生命周期**：把 decision、waiter、prepared slot 和 handle 改成 cancellation-safe guard；为 owner handle / operation 增加有界 TTL 或 generation cleanup；补齐取消与 revoke 不确定性测试。 +2. **P1，删除 exact-batch shared-op 层**：每个 external handle 直接持有一个 `BatchPlan`，其中的 item 引用 local holder 或 per-key flight。保留当前公开的 `get_start -> get_transfer | cancel` 契约。 +3. **P2，补全 master 收敛语义**：现有 `(key, requester)` RAII fence 保留，但把通用 `KeyBeingWritten`、`InvalidArgument`、`Unknown` 收敛成有限且可处理的 leader / join / already-committed / stale 结果。 +4. **P3，依据数据继续收缩**：若 per-key 对象和 wakeup 仍是热点，再把同一 leader cohort 折叠成一个 `BatchFlight`；若 suffix start / revoke 成本显著，再评估 side-effect-free probe 与 reserve 分离。 + +## 1. 评审边界与证据等级 + +本文沿三条轴评估当前路径，后文也按同样的轴给出验收条件： + +| 评估轴 | 覆盖内容 | 本文不声称的内容 | +| --- | --- | --- | +| 正确性与生命周期 | leader / joiner 归属、取消、revoke、prepared slot、master fence、reclaim 交互 | 不证明未检查模块的全系统回收行为 | +| 控制面与数据面成本 | 对象、复制、锁、task、future、RPC 和 transfer 提交边界 | 不把静态计数换算成延迟或 QPS 百分比 | +| 可验证性 | 并发测试、故障注入、metrics 和 benchmark 维度 | 不以单轮最终 QPS 代替路径级证据 | + +文中的判断分为三类： + +- **实现事实**：可由当前符号和控制流直接确认。 +- **风险推断**：由 Rust future 取消、所有权或锁作用域推导，仍需测试复现和量化影响。 +- **待测结论**：只能通过 benchmark、fault injection 或运行期 metrics 判断。 + +PegaFlow 对照只用于提炼设计模式。它的 cache 层次、vLLM 接口和跨节点协议与 Fluxon 不同，因此本文不会从 PegaFlow 的局部实现推出 Fluxon 的全路径性能结论。 + +## 2. 修复前 Fluxon 路径快照 + +### 2.1 两层共享状态（已删除 exact-batch 层） + +修复前 owner 同时维护下列共享状态： + +| 层次 | 主要类型或容器 | 去重粒度 | 持有的状态 | +| --- | --- | --- | --- | +| external handle | `external_get_start_registry` / `ExternalGetStartEntry` | 一个返回给 external 的 handle | `req_node_id` 和 exact-batch shared op | +| exact-batch | `external_get_start_by_key` / `ExternalGetStartSharedOp` | `keys + atomic_group_lens + prefix_best_effort` 完全相同 | `Starting/Running/Ready/Failed`、`waiter_count`、完整 dedup key、prefix、keys、items、`Mutex`、`Notify` | +| per-key | `OwnerKeyControlState.external_get` / `ExternalGetKeySharedOp` | owner 内单个 key | `Starting/Started/Finishing/Revoking/Ready/Failed`、`undecided`、`retained`、key、`Mutex`、`Notify` | +| master fence | `PreparedGetRequesterTable` / `PreparedGetRequesterLease` | `(key, requester)` | active `get_id` 和 RAII release | +| prepared memory | `OwnerLocalReserveSlotState::Prepared` | 一个 local-reserve slot | 精确 grant、slot、地址和大小 | + +其中 master 的 `(key, requester)` fence 是这轮更新里已经落地的有效改进。它允许不同 GPU owner 同时 materialize 同一 key，同时阻止同一 requester 的两个 prepared Get 并行发布。当前不足是冲突结果仍不可 join，也不能直接取得已经提交的 canonical holder。 + +### 2.2 修复前调用与数据流 + +```mermaid +flowchart TD + A[ExternalBatchGetStartReq] --> B{exact-batch DashMap} + B -->|已有| C[waiter_count + 1
等待同一 ExternalGetStartSharedOp] + B -->|新建| D[创建 exact-batch shared op] + D --> E[owner_key_control 下逐 key 扫描] + E --> F[Local: clone MemoryInfo] + E --> G[Join: per-key undecided + 1] + E --> H[Leader: 创建 per-key shared op] + H --> I[一次 BatchGetStart 处理 leader keys] + I --> J[task 1: external_get_key_singleflight] + J --> K{每个 key 的 prefix interest} + K -->|retained > 0| L[per-key transfer + BatchGetDone] + K -->|retained = 0| M[BatchGetRevoke + 释放 slot] + L --> N[per-key Ready] + M --> N + F --> O[task 2: external_get_start_transfer] + G --> O + N --> O + O --> P[复制为 exact-batch Ready keys/items] + C --> Q[返回独立 handle] + P --> Q + Q --> R[get_transfer 再复制 Ready keys/items
并安装 external holdings] +``` + +这里的 batch 化只覆盖部分边界： + +- owner 到 master 使用一次 `BatchGetStart`、一次 `BatchGetDone` 或 `BatchGetRevoke` RPC。 +- master 的 `handle_batch_get_start`、`handle_batch_get_done` 和 `handle_batch_get_revoke` 仍在进程内逐 item 调用单 key handler。 +- payload 阶段的 `batch_get_finish_started` 为每个需要传输的 key 创建一个 future；每个 future 调用一次 `transfer_data_no_copy`。该调用继续形成一次 closed-runtime transfer 请求。当前代码没有在这个边界先按 peer / transport 合并 descriptor。 + +因此，“RPC 是 batch”不能直接推广成“payload 提交和底层 DMA 已经是一个 batch”。底层 backend 可能还有自己的优化，但不在本次已追踪到的调用链内。 + +### 2.3 当前资源终止路径 + +| 资源或计数 | 获得位置 | 正常终止 | 当前取消或不确定性缺口 | +| --- | --- | --- | --- | +| per-key `undecided` | `plan_external_get_key_items` 看到 `Starting/Started` | `decide_external_get_key_item` | 计数与 `decision_registered: bool` 分离；future Drop 不会自动 decide / abandon | +| per-key `retained` | prefix 计算后手工增加 | leader finish / revoke 后发布 terminal | 依赖所有 `undecided` 先归零；一个丢失 decision 会阻塞整个 cohort | +| exact-batch `waiter_count` | exact-batch 命中或创建 | transfer / cancel 末尾 `release_external_get_start_waiter` | start 或 transfer RPC future 中途取消时没有 guard | +| external handle entry | `get_start` 返回前插入 `external_get_start_registry` | `get_transfer` 或 `cancel` remove | 容器没有 TTL；调用方进程死亡或消息丢失时没有本层兜底 | +| prepared local slot | leader `BatchGetStart` 前 claim | Done commit，或 Revoke 成功后 release | `OwnerLocalReserveSlotLease` 没有 Drop cleanup；claim 后的 future 取消和 Revoke RPC 不确定性都可能失去释放路径 | +| master requester fence | master 接受 prepared Get | `InflightGetInfo` Drop | 已有 RAII，且 inflight cache 有 60 秒 TTL；它不能替 owner 归还本地 `Prepared` slot | + +## 3. 正确性与生命周期发现 + +### 3.1 手工 decision 不是 cancellation-safe 所有权 + +`plan_external_get_key_items` 在 owner key fence 内先增加 `undecided`，随后调用链至少会经过 master RPC、逐 key start-code 等待和 prefix 计算。`ExternalGetStartOwnerItem::Shared` 只保存一个 `decision_registered: bool`，没有在 `Drop` 中减少计数。 + +如果 future 在增加计数后、执行 `decide_external_get_key_item` 或 `abandon_external_get_key_decisions` 前被取消： + +1. `undecided` 永久保留本次 interest。 +2. `classify_external_get_key_leader` 一直等待 `undecided == 0`。 +3. `finish_external_get_key_leaders` 按 leader 顺序 await；一个 key 卡住会阻止同 cohort 后续 key 进入 finish / revoke。 +4. per-key marker 继续占用 owner key fence，reclaim 对该 key 返回 `Busy`。 +5. exact-batch creator 仍可能停在 `Starting`，相同 batch 后续全部等待同一未完成操作。 + +这是 P0 正确性问题。修复应让“interest 存在”对应一个实际拥有 Drop 语义的 guard；bool 与计数器的约定不再承担所有权。 + +同一问题也覆盖 prepared slot。`batch_get_start_with_local_reserve_targets` 先取得 `OwnerLocalReserveSlotLease`，再 `.await` master RPC；返回 `Err` 的显式分支会释放 lease，但 future 在 await 中被 Drop 时不会执行该分支。`OwnerLocalReserveSlotLease` 本身没有 `Drop` cleanup，所以本次 claim 的 slots 可能继续停在 `Prepared`。per-key leader task 又是在该 RPC 成功返回后才 spawn；若请求在 await 中被取消，master 可能已经接受操作，owner 却没有 executor 接管它。同步计数可以直接在 Drop 中归还；需要 RPC 的 cleanup 应由 Drop 把工作提交给一个有界、可观测的 owner cleanup actor，并由 TTL 处理 actor 或进程失效，不能依赖一个无所有者的临时 task。 + +### 3.2 exact-batch waiter 和 handle 也存在同类缺口 + +exact-batch 层的 `waiter_count` 同样依赖显式 release： + +- `external_batch_get_start` 注册 waiter 后可能在 prepare 或等待 prefix 时被取消。 +- `external_batch_get_transfer` 先从 handle registry remove，再等待 shared result；等待期间被取消后，调用方已经没有 handle 可以补发 cancel，函数尾部的 waiter release 也不会执行。 +- terminal exact-batch op 只有在 `waiter_count == 0` 时才从 `external_get_start_by_key` 删除。 + +因此这一层不仅增加常数成本，也增加了一个独立的泄漏和永久等待面。删除该层会直接减少一种必须证明正确的生命周期。 + +### 3.3 exact-batch 的收益有限，参数归属也不清晰 + +exact-batch 层并非完全没有收益。它能让 100% 相同的并发请求只计算一次 prefix、只注册一次 per-key decision，并共享聚合后的 result vector。删除它以后，每个 handle 都需要构造自己的轻量 `BatchPlan`，因此必须把 100% identical workload 纳入 benchmark。 + +这项收益没有形成独立的正确性边界：真正避免重复 slot、master Start 和 payload transfer 的仍是 per-key flight。exact-batch 层为节省重复 plan 工作，引入了第二套 phase、waiter、result cache 和 cleanup 协议。对当前重点覆盖的“前缀大量重叠但 batch 长度或 atomic groups 不完全相同”场景，它又无法命中。 + +当前参数归属也暴露了这层抽象的歧义：`ExternalGetStartDedupKey` 不包含 `transfer_concurrency`，`ExternalGetStartSharedOp` 却保存第一个 creator 的值。相同 batch 使用不同 `transfer_concurrency` 时，后来的显式参数不会决定共享工作的并发度。per-key sharing 本身也意味着 shared keys 服从各自 leader cohort 的并发策略。这里需要保留一个规范契约：优先把 transfer concurrency 变成 owner / transport 的统一策略；如果仍保留请求参数，文档必须明确它只影响本请求新建的 leader cohort，不能维持未说明的 first-creator-wins 行为。 + +### 3.4 Revoke 不确定性没有 owner 侧闭环 + +当 suffix key 无人 retain 时，leader 进入 `Revoking`。`BatchGetRevoke` 成功后，代码才释放精确 prepared target。RPC 返回错误时，当前实现发布 per-key `Failed` 并清除 marker，但没有保存可重试的 `get_id` 和 target，也没有为 owner `Prepared` slot 建立 TTL。 + +master 的 `InflightGetInfo` 过期后会 Drop `PreparedGetRequesterLease`，所以 master fence 最迟可以重新开放。owner 本地 slot 是否最终返回 `Free`，无法从当前路径证明：释放它所需的精确 target 已经随 terminal op 丢失。更安全的状态应保留 `Revoking { get_id, target, next_retry }`，直到满足下列有限终态之一： + +- master 明确确认 Revoke，owner 释放 slot； +- master 返回 Done 已胜出，owner 按 canonical committed backing 收敛； +- master operation TTL 明确过期，owner 核对 generation 后释放 slot； +- owner generation 结束，整个 pool 随进程生命周期销毁。 + +在终态前清除 per-key marker 会把“master 是否仍在执行”和“owner slot 是否仍被占用”重新变成两个失去关联的状态。 + +### 3.5 master 防重已经部分完成，收敛语义仍缺失 + +当前 `PreparedGetRequesterTable` 已正确使用 `(key, requester)` 作为 identity，并通过 `PreparedGetRequesterLease::Drop` 释放。这一部分应保留。 + +仍需收敛的分支包括： + +| 场景 | 当前行为 | 需要的有限结果 | +| --- | --- | --- | +| 同 requester 已有 prepared Get | `KeyBeingWritten` | `Join { operation_id, immutable_start_item }`，或明确可重试的 transient 结果 | +| requester 已有同版本 live replica | `InvalidArgument: cannot replace a live replica` | `AlreadyCommitted`，释放新 slot 后取得 canonical local holder | +| GetDone 时同版本 route 已被另一个操作发布 | `Unknown: could not publish current route` | `AlreadyCommitted`，输家不覆盖 backing | +| GetDone 的 `put_id` 已落后 | 依赖通用错误路径 | `Stale`，revoke / release 后重新 probe | + +owner singleflight 可以是主要的快速路径，但不能成为唯一能解释冲突的正确性边界。master 已经维护唯一性 fence，下一步应返回可枚举、可测试的结果，避免 owner 从错误字符串猜测状态。 + +## 4. 修复前控制面与数据面成本 + +下表只描述当前源码可以确认的成本。最后一列中的性能影响仍需测量。 + +| 成本维度 | 当前实现事实 | 随什么增长 | 判断 | +| --- | --- | --- | --- | +| key 所有权复制 | exact-batch DashMap key 和 `ExternalGetStartSharedOp.dedup_key` 各自拥有完整 key 向量；transfer prefix 又被构造并在 `Running` / task 间复制；每个 per-key op 还拥有自己的 `String` | batch key 数和 key 长度 | 明确存在，可通过删除 exact-batch 层减少；延迟占比待测 | +| 共享对象 | 每个新 leader key 有一个 `Arc`、`Mutex`、独立 `Arc` 和 terminal result；每个 exact batch 再有一套 shared op | leader key 数和并发 batch 数 | 对短 value / 高并发更敏感，实际 allocator 成本待测 | +| 全局 fence | `owner_key_control` 在同一临界区内先扫描 revoking keys，再逐 key 做 local / join / leader 规划，并穿插 per-key state lock 和索引查询 | required key 数及并发请求数 | 可能形成 p99 lock contention；需要 lock-wait histogram | +| per-key lock / wakeup | start publish、decision、classification、terminal publish 和每个 waiter 都访问 per-key state；exact-batch 又有第二层 lock / notify | joiner 数和状态变化次数 | 两层 wakeup 没有独立语义，优先删除外层 | +| task | 非全 local 的 creator 通常生成 `external_get_key_singleflight` 和 `external_get_start_transfer` 两个 task | creator batch 数 | 第二个 task 只聚合 per-key terminal 到 exact-batch terminal，可删除 | +| waiter future | `finish_external_get_start_transfer` 对 transferable items 建立 `join_all` future；revoke 等待也按 key 建 future | transferable / revoking key 数 | 调度与 wakeup 成本明确，幅度待测 | +| result materialize | per-key `Ready` 保存结果；聚合 task 再构造 exact-batch `Ready`；每个 handle transfer 再 clone `keys/items` | transferable key 数和相同 batch waiter 数 | exact-batch result cache 造成额外向量与 `Arc` refcount 流量 | +| payload 提交 | `batch_get_finish_started` 为每个远端 key 调用一次 `transfer_data_no_copy`；每次继续形成独立 closed-runtime transfer request | 远端 hit key 数 | 控制面 RPC 已 batch，transfer 提交尚未在该边界合并 | +| master batch handler | `BatchGetStart/Done/Revoke` handler 在 master 内逐 item await 单 key handler | leader / terminal key 数 | 网络 RPC 数受控，master 本地调度仍为 O(keys) | +| suffix speculation | 所有 leader 先 Start，prefix 计算后再按 `retained` 选择 Finish 或 Revoke | raw hit 之后的 leader suffix | 避免 probe RPC，但增加 slot claim、master work 和 revoke；需和 probe/reserve 实测比较 | + +当前优先级更高的风险依次是取消后不终止、全局 fence 临界区、双层状态 wakeup 和逐 key transfer submission。key clone 也明确存在，但不应先为它引入新的缓存或兼容层。 + +## 5. PegaFlow 可以借鉴的设计模式 + +PegaFlow 对照源码固定在提交 `939363f198a0ffded9e3f30d8af9cdd74439f16c`： + +| PegaFlow 模式 | 实现证据 | 对 Fluxon 的启发 | +| --- | --- | --- | +| 一次 prefix scan 取得稳定引用 | [`ReadCache::get_prefix_blocks`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/storage/read_cache.rs#L35-L50) 在一个 cache mutex 下按序扫描并 clone `Arc` | Fluxon 的 local branch 已在 owner fence 下 clone `MemoryInfo`；让 `BatchPlan` 直接拥有这些 pin 即可，无需再包装 exact-batch state | +| 一个 ready prefix 对应一个 opaque lease | [`QueryLeaseManager`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/lease.rs#L42-L184) 让 lease 持有整个 `Vec>`，支持 consume、release、consumer count 和 sweep | Fluxon handle 应拥有整批 plan / lease，并在 transfer 或 cancel 时一次消费；进程死亡由 TTL 收口 | +| request-level prefetch state | [`PrefetchState`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/storage/prefetch.rs#L79-L155) 用一个 `HashMap` 和一个 mutex 管理后台任务 | per-key singleflight 只保留跨请求去重真正需要的状态；请求级聚合不要再复制一套 phase machine | +| probe 与 commit 分离 | [`vllm-request-state-machine.md`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/docs/vllm-request-state-machine.md#L318-L461) 明确提出 `QueryPrefetch -> ReserveLoadBlocks -> Load -> ReleaseReservation -> TTL` | Fluxon 可用 side-effect-free probe 先确定 prefix，再只 reserve transferable prefix；是否值得多一次 RPC 必须 benchmark | +| 取消安全的资源 guard | [`TransferLockGuard`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/backing/transfer_lock_guard.rs#L13-L70) 和 [`SsdPrefetchReservation`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/storage/prefetch.rs#L141-L155) 都在 `Drop` 中归还资源 | `undecided`、waiter、slot 和 master operation retry ownership 都应由 guard 表达 | +| descriptor-first 的批量数据面 | [`transfer`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/transfer/mod.rs#L1-L16) 与 [`gpu_worker`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/src/gpu_worker.rs#L350-L496) 先收集所有 layer / segment descriptor，再交给一次 backend batch 并同步一次 | Fluxon 应统计并减少实际 transfer submission，而不只统计 BatchGet RPC 数 | +| 路径契约测试和 metrics | [`prefix_semantics`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/tests/prefix_semantics.rs)、[`prefetch_lease`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/pegaflow-core/tests/prefetch_lease.rs) 和 [`test_connector_fault_tolerance`](https://github.com/novitalabs/pegaflow/blob/939363f198a0ffded9e3f30d8af9cdd74439f16c/python/tests/test_connector_fault_tolerance.py) 分别覆盖 prefix、lease consumer、失败与 cleanup | Fluxon 需要把 overlap、取消和资源归零作为 merge gate，纯 prefix helper 不足以覆盖生命周期 | + +### 5.1 不应照搬的部分 + +- **PegaFlow 没有跨请求 per-block singleflight**：active prefetch 按 `req_id` 组织；两个不同请求的重叠 block 仍不共享同一个 per-block future。 +- **当前 probe 仍有副作用**：`QueryPrefetch` 已经 pin ready blocks。PegaFlow 自己的设计文档也把 probe / reserve 分离列为未来协议。 +- **`req_id` 不是完整 identity**：同一请求在 preemption 或调度重试后可能对应不同 hash slice；PegaFlow 文档建议加入 `digest(block_hashes)`。 +- **600 秒 lease TTL 不适合直接复制**：Fluxon prepared local-reserve slot 是稀缺资源，TTL 应由 slot 容量、最大合法 transfer 时长和故障恢复目标决定。 +- **语义边界不同**：PegaFlow 的简单 prefix cache scan 不承担 Fluxon 的 atomic group、master route、owner reclaim fence 和跨 requester 唯一性。 + +## 6. 建议的收缩设计 + +### 6.1 第一阶段：handle 直接持有 `BatchPlan` + +第一阶段保留现有 per-key shared op，先删除 exact-batch shared op。内部结构可以收敛为： + +```rust +struct ExternalGetBatchPlan { + req_node_id: String, + keys: Vec, + prefix: ExternalGetStartPrefixResult, + items: Vec, +} + +enum ExternalGetPlanItem { + Local(Arc), + Shared(ExternalGetInterest), +} +``` + +这只是内部设计草图，名称应在实现时与仓库现有命名统一。关键契约是： + +- `ExternalGetInterest` 由 RAII guard 表达。创建时注册一次 prefix decision；prefix 计算会把它消费成 abandon 或 plan 持有的 result interest。未决定就 Drop 时自动 abandon；进入 `Finishing` 后的 Drop 只释放 waiter，不回滚已经提交的数据面终态。 +- `external_get_start_registry` 直接持有 `BatchPlan`，不再间接持有 `ExternalGetStartSharedOp`。 +- `get_transfer` 消费 plan 并等待其中的 per-key terminal；`cancel` Drop plan。两个入口共享同一条释放路径。 +- 完全相同的两个 batch 仍会在 per-key registry 中 join 相同 flight，因此不会重复 Start / transfer。 +- leader cohort 的现有后台 task 已经负责 transfer 和 Done。删除 `external_get_start_transfer` 聚合 task 不会失去 start 与 GPU allocation 的重叠，只会删除第二次等待和结果 materialize。 +- handle registry 需要 TTL / generation cleanup。TTL 回调与显式 cancel 必须调用同一个幂等终止函数。 + +可以删除的重复 surface 包括： + +- `ExternalGetStartDedupKey` +- `ExternalGetStartSharedOp`、`ExternalGetStartSharedState`、`ExternalGetStartSharedPhase` +- `external_get_start_by_key` +- `register_external_get_start_waiter`、`release_external_get_start_waiter` +- `wait_external_get_start_prefix`、`wait_external_get_start_transfer` +- `publish_external_get_start_ready` 及第二份 batch terminal result +- `external_get_start_transfer` 聚合 task + +`ExternalGetStartSharedItemResult` 此后只属于 per-key terminal，应改成唯一且明确的 `ExternalGetKeyResult`，避免名字继续暗示 exact-batch sharing。 + +```mermaid +flowchart LR + A[ExternalBatchGetStartReq] --> B[owner fence 下构造 BatchPlan] + B --> C[Local
Arc MemoryInfo] + B --> D[Join
KeyInterest guard] + B --> E[Leader cohort] + E --> F[一次 BatchGetStart] + F --> G[一个 cohort task
transfer + Done/Revoke] + G --> H[per-key terminal] + C --> I[handle registry
BatchPlan + TTL] + D --> I + H --> I + I -->|get_transfer| J[消费 plan，安装 holdings] + I -->|cancel / TTL| K[幂等 Drop plan] +``` + +### 6.2 第二阶段：有数据再折叠成 `BatchFlight` + +删除 exact-batch 层后，如果 metrics 证明 per-key `Arc + Mutex + Notify` 仍是热点，可以让同一次扫描产生的 leader keys 共用一个 `BatchFlight`： + +- key registry 保存 `(Weak, slot_index, generation)`。 +- `BatchFlight` 持有一个 slots 数组、一个 mutex、一个 notify 和一个 cohort task。 +- joiner 只保存 `Arc + slot_index + InterestGuard`。 +- 每个 slot 有独立 outcome 和 interest count，但同步原语属于整个 cohort。 +- registry 删除必须同时校验 key、flight identity、slot 和 generation，旧 waiter 的 Drop 不能删除后来创建的 flight。 + +这一步能把同步对象从 O(leader keys) 降到 O(leader cohorts)，代价是一个 notify 可能唤醒同 cohort 的无关 slot waiter。应先测 per-key 对象和 wakeup 是否真的占主导，再决定是否接受该 trade-off。不要在删除重复层之前再叠加第三套 batch state machine。 + +### 6.3 master 返回可收敛的有限结果 + +master 的 prepared Get start 应返回一个可穷举结果集合。下面只描述分支语义，不是当前已存在的 Rust 类型： + +```text +Leader { start_item } +Join { operation_id, immutable_start_item } +AlreadyCommitted { put_id, backing_identity } +Stale +``` + +具体 wire 类型需要单独设计，但分支必须保持有限： + +- `Leader` 获得本次 prepared target 的唯一提交权。 +- `Join` 返回足够重建原 operation 的 immutable start 信息。owner 仍只能选出一个 transfer executor，其他 waiter 观察同一终态。 +- `AlreadyCommitted` 返回 version 和 backing identity，让 owner 释放候选 slot，并在 reclaim fence 下取得 canonical local holder。 +- `Stale` 终止旧 operation，重新读取当前 route / version。 + +master fence 的 RAII lease 和 60 秒兜底可以继续保留，但 owner slot 的终止不能只依赖 master cache Drop。 + +### 6.4 probe / reserve 作为可选协议，不先假设它更快 + +当前 speculative Start 的优点是少一次网络往返，并能尽早开始远端 materialization；成本是 raw prefix 之后的 key 也 claim slot、进入 master、再走 Revoke，还需要 `undecided/retained/Revoking` 协调多个 batch 的 prefix interest。 + +可选的两阶段协议为： + +1. `BatchGetProbe(keys)`：只返回每个 key 的 local / route 可用性和 version，不占 prepared slot。 +2. owner 按原始顺序和 atomic groups 算出 `transferable_len`。 +3. `BatchGetReserve(keys[..transferable_len], versions)`：全量成功或返回 stale / retry，不留下部分 reservation。 +4. transfer 消费 reservation；cancel / failure / TTL 释放 reservation。 + +该协议能删除大部分 suffix revoke 和 decision accounting，但多一次 RPC，并引入 probe 与 reserve 之间的 eviction / version race。正确处理方式是 reserve 时原子复核，把 race 降级为 cache miss / retry。是否采用，应比较： + +- 当前 speculative Start + suffix Revoke； +- probe + reserve； +- 不同 batch 长度、miss 位置、overlap ratio 和 RTT 下的端到端结果。 + +### 6.5 保留 batch 语义，并把 batch 延伸到 transfer submission + +无论控制面采用 per-key op 还是 `BatchFlight`，都应维持： + +1. owner 以逐 key/shard 短 fence 完成 required batch 的 local / join / leader 分类；不得用 + process-wide lock 取得整批快照，也不得让任何同步锁跨 `.await`； +2. leader keys 压缩成一次 `BatchGetStart`； +3. terminal 使用一次 `BatchGetDone` 或 `BatchGetRevoke` 收敛； +4. 结果按原 index scatter，再计算 raw prefix 和 atomic-group prefix； +5. payload descriptor 先按 peer、方向和 transport 能力分组，再调用 batch transfer surface。 + +PegaFlow 的 descriptor-first 做法适合借鉴到第 5 步。Fluxon 至少应新增 `transfer_submissions`、`descriptors_per_submission` 和 `bytes_per_submission`，先确认当前每 key 调用是否真的成为瓶颈，再设计 batch transfer API。 + +## 7. 迁移顺序 + +| 阶段 | 改动 | 完成判据 | +| --- | --- | --- | +| P0 生命周期 | 引入 decision / waiter / prepared operation guard；统一显式 cancel、future Drop、TTL cleanup；保留 revoke retry identity | fault injection 取消任一 await 后，active handle、interest、marker 和 prepared slot 都回到基线 | +| P1 删除重复层 | handle 改持 `BatchPlan`；删除 exact-batch map、phase、notify、waiter 和聚合 task | identical / overlapping batch 的 leader Start 数不增加；公开 API 不变 | +| P2 master 收敛 | wire 层增加有限 outcome；owner 实现 join、already-committed 和 stale 分支 | 同 requester 冲突不再暴露通用 `KeyBeingWritten` / `InvalidArgument` / `Unknown` | +| P3 数据驱动优化 | 测量并选择 `BatchFlight`、probe/reserve、batch transfer submission | 每个新增机制都有对应热点证据和独立回退对照 | + +截至 2026-07-17，P0 的 Interest/planning leader/prepared-slot guard、P1 exact-batch 删除、 +P2 的 `(key, requester)` fence 与 overlap/ABA 防线已进入当前工作树并通过定向单元测试； +P2 的完整有限 outcome 仍需故障注入确认。P3 的 +payload descriptor batch 仍未实现，必须由部署后的 submission metrics/QPS 决定，不能仅凭 +静态结构继续加层。 + +P0 与 P1 可以先在现有 per-key 实现上完成。不要为了等待理想的 `BatchFlight` 一次性重写而继续保留已知的取消缺口。 + +## 8. 测试、观测与验收门槛 + +### 8.1 必须补齐的测试 + +最新更新已经增加三类有价值的 unit test：两个非完全相同 batch 复用同一 per-key marker、旧 marker cleanup 不删除新 generation，以及 master `(key, requester)` fence 的同 owner / 跨 owner 与 ABA 行为。它们验证了同步注册和 identity 规则，但没有启动真实 RPC / transfer / cancel 生命周期。下列测试仍需直接运行 async 流程，显式检查 exit code / timeout 和资源终态： + +| 场景 | 必须断言 | +| --- | --- | +| overlap：`[a,b,c]` 与 `[b,c,d]` | `b/c` 对同 owner 各只有一个 leader Start 和一份 committed backing；两个 batch 各自 prefix 正确 | +| identical batch | 删除 exact-batch 层后仍只产生一组 per-key leader；两个 handle 可独立 transfer / cancel | +| duplicate key：`[a,a,b]` | `a` 只有一个 flight；两个原 index 都得到一致结果；interest 不下溢或泄漏 | +| atomic-group suffix | group 不被切开；suffix local pin、join interest 和 leader prepared slot 全部释放 | +| task cancellation matrix | 在注册 decision、Start 前后、等待 prefix、等待 result、handle remove 后逐点 abort；后续同 key 请求可在 deadline 内完成 | +| Revoke 响应丢失 | owner 保留 retry identity；master TTL / terminal query 后 slot 最终回到 `Free`，不会直接清 marker 丢状态 | +| Done / Revoke race | 只有一个终态释放资源；Done 胜出时 Revoke 不释放 committed slot | +| reclaim race | Local、Starting、Finishing、Revoking 和 terminal 各阶段与 reclaim 并发，不出现地址复用或永久 `Busy` | +| master requester fence | 同 requester 第二个 Get 返回 join / already；不同 requester 可以同时成为 leader | +| stale version | 旧 `put_id` 不覆盖新 route；候选 slot 被回收并重新 probe | +| owner / external generation change | 旧 handle、旧 flight 和旧 cleanup 不能删除或消费新 generation 状态 | + +### 8.2 必须暴露的路径 metrics + +| 类别 | 指标 | +| --- | --- | +| 分类 | `required_keys`、`local_pinned`、`inflight_joined`、`leader_started`、`duplicate_indices` | +| 去重收益 | `get_start_avoided`、`transfer_avoided`、`identical_batch_joined` | +| 生命周期 | active handles / flights / interests / prepared slots、各自 age、cancel / TTL / Drop cleanup 次数 | +| suffix | local pin drop、join interest drop、leader revoke、revoke retry、revoke unknown | +| master 收敛 | leader、join、already-committed、stale、TTL expired，按结果分别计数 | +| 等待与锁 | owner fence wait / hold、join wait、flight lifetime 的 histogram | +| 数据面 | BatchGet RPC items、transfer submissions、descriptors / bytes per submission、Done / Revoke batch size | + +所有 gauge 都必须能在 workload 结束后的有界时间内回到基线。只增加累计 counter 无法发现永久挂住的 handle、interest 或 prepared slot。 + +### 8.3 benchmark 矩阵 + +至少比较三个实现：变更前的 batch path、当前双层 shared-op path、删除 exact-batch 层后的 `BatchPlan` path。输入矩阵应覆盖： + +- key 数：`1 / 16 / 64 / 256`； +- overlap ratio：`0% / 50% / 90% / 100%`； +- first miss：首部、中部、尾部和 all-hit； +- 并发：单请求、同 owner 多请求、多个 requester owner； +- value / page 大小：控制面占主导的小 value，以及数据面占主导的实际 KV page； +- local / same-host / remote transport； +- fault-free、cancel storm 和 Revoke 响应延迟 / 丢失。 + +采集 `get_start` p50 / p99、完整 transfer p50 / p99、CPU、allocation bytes / count、task spawn、lock wait、RPC items、transfer submission 和最终 slot / handle gauge。验收时先检查语义计数和资源归零，再比较延迟；最终 QPS 只能作为补充结果。 + +## 9. Fluxon 证据索引 + +| 事实 | 当前代码位置 | +| --- | --- | +| per-key flight 与 request plan 类型 | `fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs`：`ExternalGetKeySharedOp`、`ExternalGetKeyInterest`、`ExternalGetStartEntry` | +| per-key 规划和 RAII decision | `fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs`：`plan_external_get_key_items`、`ExternalGetKeyInterest::decide/Drop` | +| leader finish / revoke | 同文件：`classify_external_get_key_leader`、`finish_external_get_key_leaders` | +| request plan handle（无 exact-batch shared-op） | 同文件：`external_batch_get_start`、`finish_external_get_start_transfer`、`external_batch_get_transfer` | +| per-key payload transfer submission | `fluxon_rs/fluxon_kv/src/client_kv_api/get.rs`：`batch_get_finish_started` | +| transfer closed-runtime 边界 | `fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs`：`transfer_data_no_copy` | +| external Get 与 reclaim fence | `fluxon_rs/fluxon_kv/src/client_kv_api/reclaim.rs`：`prepare_one` | +| master requester fence 与 TTL | `fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs`：`PreparedGetRequesterTable`、`PreparedGetRequesterLease`、`inflight_gets` | +| master 冲突和 batch handler | `fluxon_rs/fluxon_kv/src/master_kv_router/get.rs`：`handle_get_start`、`handle_get_done`、`handle_batch_get_start` | +| 已有 overlap / ABA unit tests | `external_get_start_batch_tests::overlapping_nonidentical_batches_share_each_key_but_keep_leaders_batched`、`old_singleflight_cleanup_cannot_remove_new_generation`、`prepared_get_singleflight_is_same_owner_only_and_aba_safe` | +| 原设计要求 | `fluxon_doc_cn/design/sglang_fluxon_kv集成设计.md`：`Batch Get 的逐 key 交集、pin 与 singleflight` | + +## 10. 决策记录 + +- **接受** owner 侧 per-key overlap 去重和 master 侧 `(key, requester)` 防御性唯一性。 +- **接受** BatchGetStart / transfer / Done 在数据流上的 cohort 概念。 +- **拒绝** exact-batch shared op 与 per-key shared op 长期并存。 +- **拒绝** 用 bool 和手工 counter 表达跨 `.await` 的资源所有权。 +- **暂缓** `BatchFlight` 和 probe/reserve 协议，直到 P0 / P1 完成且 metrics 指向对应热点。 +- **要求** 在把该路径视为完成前，补齐取消、overlap、Revoke 不确定性、reclaim race、master 收敛测试和资源归零观测。 + +## 11. 2026-07-16 实现跟进 + +本分析提出的 P0/P1 已进入工作区实现： + +- 删除 `ExternalGetStartDedupKey`、`ExternalGetStartSharedOp`、 + `external_get_start_by_key`、request-level waiter/phase/result cache 和 + `external_get_start_transfer` 聚合 task；handle 直接拥有自己的 keys/items plan。 +- `ExternalGetKeyInterest` 以 RAII 表达 prefix decision；未决定的 request future 被取消时， + Drop 自动减少 `undecided`。leader cohort 在第一个 master RPC await 前交给注册后台 task， + 因而请求取消不会丢失 Start/Finish/Revoke executor。 +- prepared slot claim 使用 Drop cleanup guard；明确接受的 slots 才逐项 disarm。external + handle 增加 360 秒 TTL 和统一 Drop 回收，覆盖当前对齐 workload 的 300 秒合法请求窗口。 +- `Revoking` 保存完整 start item；BatchGetRevoke 的 transport、响应长度和 identity 不确定性 + 保留 marker/slot 并重试，只有明确 Revoke 成功才释放 target。transfer/install 失败的 + cleanup 也在独立注册 task 中持有 get ids 和 targets，调用方取消不再丢失清理 ownership。 +- overlap/ABA 测试之外新增了 pending Interest Drop 和真实 task abort 测试；当前仍需用真实 RPC + fault injection 覆盖 Done response 丢失、owner generation 切换和完整 transfer cancellation matrix。 +- BatchGetDone 不再固定重试三次后遗留 pending-visible slot，而是使用同一组 get ids 幂等重试， + 严格校验响应长度和逐项 identity，直到明确终态或 owner shutdown。BatchGetStart 没有响应的 + 未知提交点会将 prepared slots 隔离 65 秒，覆盖 master 60 秒 inflight TTL；已知 get ids 的 + 异常响应则移交独立 Revoke cleanup。 +- owner 周期快照增加 active handle、per-key flight 阶段、undecided/retained interest 和 + Free/Prepared/Pending/Committed local-reserve slot 计数，用于三机 workload 后资源归零验收。 + +这次跟进没有宣称 payload 已成为单次底层 DMA batch。控制面 cohort 仍是一次 +BatchGetStart/Done/Revoke，但 `transfer_data_no_copy` 的 descriptor-first 合并仍属于后续 P3， +必须由 submission metrics 和 benchmark 决定接口形态。 + +2026-07-17 使用宝宝盘 `CARGO_TARGET_DIR` 和单并发构建做了最新工作树回归。 +Clippy 在仅放行仓库既有 `uninit_vec` 的前提下强制 +`-D clippy::await_holding_lock` 通过；master-router、owner-hot、local-reserve、 +member-left 和 `test_memholder_pin` 分别为 `28/28`、`9/9`、`8/8`、`5/5`、`1/1`。 +`fluxon_kv --lib` 全量为 `176 passed, 0 failed`，旧的 `test_memholder_pin` 失败已不再复现。 +上述结果不替代真实 RPC fault injection 和三机生命周期验证；尤其仍需主动丢弃 +Start/Done/Revoke 响应并检查快照临时态回到基线。 diff --git "a/fluxon_doc_cn/design/sglang_fluxon_kv\351\233\206\346\210\220\350\256\276\350\256\241.md" "b/fluxon_doc_cn/design/sglang_fluxon_kv\351\233\206\346\210\220\350\256\276\350\256\241.md" new file mode 100644 index 0000000..8d6fa0f --- /dev/null +++ "b/fluxon_doc_cn/design/sglang_fluxon_kv\351\233\206\346\210\220\350\256\276\350\256\241.md" @@ -0,0 +1,1686 @@ +# SGLang Fluxon KV 集成设计 + +## 背景与目标 + +本文说明开源仓库中 SGLang HiCache 接入 Fluxon KV 的 hostless 实现设计，聚焦接口契约、状态归属和生命周期边界。本文不是通用 `FlatDict` KV API 的完整说明。 + +稳定结论：本集成的目标是把 SGLang HiCache + Mooncake 形态下分裂的 L2/L3 KV Cache 收敛到 Fluxon 统一管理的本机/远端 KV 层。Fluxon 需要同时具备两类能力：一类是 `key -> value` 的分布式缓存、路由、版本和生命周期能力；另一类是 HiCache L2 所依赖的 hostless 数据面能力，也就是让 SGLang native kernel 直接在 GPU KV cache 和 Fluxon 管理的 host value memory 之间搬运数据。 + +关键 insights 先集中列在这里，后文围绕这些结论展开接口、状态和时序设计： + +- **L2/L3 分离是工程解耦选择，但会带来资源冗余和 L2 全局不可见** + - 因果链：HiCache L2 留在推理框架内，Mooncake L3 作为外部 backend，能降低接入复杂度；代价是同一类 KV page 被两套系统分别索引、驻留、驱逐和释放。同一批 page 可能同时驻留在本机 L2 和外部 L3，且其它 worker 或其它节点不能通过统一 route 判断本机 L2 page 是否可复用。 + - 本设计落点：把逻辑 L2/L3 收敛到 Fluxon local side / remote side，用一套 owner、holder、route、commit/release 管理同一批 page。page key 在 Fluxon master/owner 中形成统一路由；只有 commit future 成功后，跨 worker 和跨节点复用才把该 page 视为可见。 +- **同机 worker 缺少共享内存快路径** + - 因果链：多个 SGLang worker 容易各自维护后端 segment、本机缓存副本和独立 pin/release 状态；同机交接也可能绕行后端协议。 + - 本设计落点：同机 worker attach 到同一个 Fluxon owner shared segment，通过 holder 引用管理本机可见 value 的生命周期。 +- **统一 L2/L3 需要分布式 KV 能力** + - 因果链：Mooncake 这类分布式 KV backend 的目标仍然成立：KV Cache 需要利用远端 CPU 内存，尤其是无推理负载机器上的闲置内存，作为 shared backing。只用本进程数组或本机缓存，无法统一索引、放置和复用这些远端闲置内存。 + - 本设计落点：Fluxon 保留 `key -> value` route、inflight 判重、commit future 和 local/remote 放置能力，把远端闲置 CPU 内存纳入 remote side，同时把本机 L2 纳入同一套 owner/holder 生命周期。 +- **保留 L2 低延迟需要 kernel 直连内存** + - 因果链：普通 `put/get(key, bytes)` KV 接口需要在 CPU 侧 materialize payload、查地址、组织拷贝或拆装 bytes；长上下文下如果再按 page 和 layer 展开，会形成 `page_count * layer_count` 级别的 CPU 控制面循环。 + - 本设计落点：Fluxon 只构造 `plan_ptr(value_ptrs)`，SGLang native kernel 直接在 GPU KV cache 和 Fluxon host value memory 之间搬运 bytes。 +- **KV Cache 适合 write-back 最终一致性，最终目标是极致化这条路径** + - 因果链：page value 按不可变对象使用，缓存副本可丢失；写回失败只会降低命中率，不改变推理正确性。因此优化重点应放在缩短 write-back 和 restore 的热路径，而不是把每次缓存写回做成强同步持久写。 + - 本设计落点：native write 完成后异步推进 route commit，`KvFuture` 成功前不作为跨 worker / 跨节点 shared backing；写入侧减少同步等待、CPU 循环和中间 bytes 包装，读取侧先规划 prefix，再只 materialize 可恢复前缀。 +- **需要适配 SGLang KVCache 特化接口** + - 因果链：SGLang KVCache restore 不是普通逐 key `get`，而是面向一批有序 page keys 的 prefix restore。它需要一次性回答“这一批最多恢复到哪里”，并且不能切开一个 radix node 对应的 atomic group。batch 化还能减少逐 key 调用、查表、锁和跨边界调度开销。 + - 本设计落点：Fluxon hostless get 以 batch 为单位计算 `raw_prefix_hit_len`、`transferable_len` 和 `prefix_hit_groups`。`get_start` 合并存在性 / prefix 判断和可恢复前缀的数据拉取启动，返回 handle 和 prefix；`get_transfer` 再消费 handle，把可恢复前缀转换成 readable `plan_ptr(value_ptrs)`。 + +### 关键设计倾向与强制规约 + +本文使用 `MUST`、`MUST NOT`、`SHOULD` 和 `MAY` 区分正确性规约与性能倾向： + +- `MUST / MUST NOT` 是内存、版本、可见性或批量语义的硬不变量；即使某个实现能得到更高 + QPS，违反它也不得验收。 +- `SHOULD` 是默认性能方向；偏离时必须给出可重复的对照实验和没有破坏硬不变量的证据。 +- `MAY` 表示在不改变外部契约时可替换的内部策略。 + +设计取舍的优先级固定为：首先保证 backing/version/holder 正确性，其次保证本地快路径 +不被全局控制面阻塞，再次保证 batch/atomic/TP 语义，最后才是单轮命中率或 QPS。 +关键规约如下： + +| ID | 级别 | 规约 | 验收含义 | +| --- | --- | --- | --- | +| V1 | MUST | native write 和 stream 可见性完成后，Put slot 必须立即进入 `local-read-ready` / `precommit_local_visible_info`。同 owner local read 不等待 master PutDone。 | 在人为延迟 master RPC 时，同 owner Get 仍能命中并取得 holder。 | +| V2 | MUST NOT | `local-read-ready` 不得被解释为 global committed。master route 成功前，其它 owner 不得通过全局 route 观察该 value。 | local 与 remote 可见性分别故障注入，不出现未完成 route 的跨 owner 读。 | +| V3 | MUST | `Committed(route_live=true)` 和 owner-hot admission 必须等 master route/maintenance 成功；这个门禁不得反向阻塞 V1 的 local read。 | hot Moka 中每个条目都可与相同 `put_id + backing identity` 的 live route 对账。 | +| G1 | MUST | required batch `R` 必须与可 pin 的 local-visible `L` 和可共享 inflight Get `F` 完整分流：`LocalJoin=R∩L`、`InflightJoin=R∩F`、`Leader=R-(L∪F)`。每个 key 独立线性化，不要求用一把锁取得整批快照。 | 重叠但不完全相同的 batch 不会对共享 key 重复申请 slot 或传输，同时大 batch 不会阻塞无关 key。 | +| G2 | MUST | 任何 leader GetStart/payload transfer 之前，`LocalJoin` 必须已持有 `MemoryInfo` holder pin，`InflightJoin` 必须已持有 shared-op waiter/pin，`Leader` 必须已原子安装 marker。 | pin 与 hot demotion/reclaim 并发时，kernel 看到的地址不会提前回收或复用。 | +| G3 | MUST NOT | 不得用“snapshot inflight set，稍后批量 insert”实现与非；必须逐 key 使用原子 entry/CAS 确定 leader 和 joiner。 | 并发两批均只有一个 leader，registry 清理不发生 ABA 删新操作。 | +| G4 | MUST | leader 子集的压缩响应必须先按原 index 重组成与 `R` 等长结果，然后才能按原 `atomic_group_lens` 计算 `raw_prefix_hit_len/transferable_len`。 | 不会因压缩 leader 而改变 prefix 顺序，也不恢复半个 radix node。 | +| G5 | MUST | `transferable_len` 之后的 local pin、join waiter 和 leader prepared Get 必须立即释放/revoke；不得保留到 TTL。 | suffix holder、master inflight 和 prepared slots 在每轮结束后可精确对账归零。 | +| G6 | MUST | owner 侧按 key singleflight；master 侧的防御性唯一性粒度为 `(key, requester_owner)`，不是全局 key。 | 同 owner 不重复 materialize；不同 GPU owner 仍可同时拉取同一 key。 | +| G7 | MUST | per-key shared Get 必须至少有 `Starting/Started/Ready/Failed-or-Revoked` 阶段，所有 waiter 观察同一终态并共享同一 canonical `MemoryInfo`。 | 重叠 batch 测试中 master GetStart、DMA 和 committed local route 均只出现一份。 | +| G8 | MUST / SHOULD | per-key singleflight 只是 leader/joiner 的线性化粒度，不得把控制面降级为逐 key RPC。同一 required batch 的 leader 必须压缩成一次 BatchGetStart，GetDone/Revoke 也必须批量收敛（MUST）。payload 当前仍逐 key调用 `transfer_data_no_copy`；应按 peer/transport 汇总 descriptor 后批量提交，或至少在有界 time-window 内并行聚合（SHOULD，尚未实现）。 | Start/Done/Revoke 次数按 batch 而不是 leader key 数增长；另外单独统计 payload submission，未实现 descriptor batch 前不得宣称数据面已完全 batch 化。 | +| G9 | SHOULD | 同时到达的多个 required batches 中，已经判定为 leader 且具有兼容 value geometry/transport 的 keys 应进入短 time-window micro-batch；窗口必须有界，不能为凑 batch 无界延迟 local hit 或 prefix 响应。 | 同时并发压力下 batch size 提升，但 p50/p90 GetStart 不因等待窗口显著恶化。 | +| G10 | MUST NOT | owner Get 规划、hot pin 和 runtime metrics 不得持 process-wide key-control mutex 遍历 batch/cohort；key-control 只能是短 O(1) 的逐 key/shard 临界区，禁止跨 `.await`、RPC、RDMA。规划取消必须用 RAII 清理尚未 BatchGetStart 的 leader marker。 | 大 batch 与无关 key 并发前进；取消任一 await 后没有永久 `Starting` marker、undecided 或 prepared slot。 | +| A1 | MUST | atomic group/TP cohort 可以逐项 local-ready，但只有完整 group 才能对 restore 发布命中；同一 cohort 先全部 promotion，再统一 owner-hot admission。 | 任一 rank/member 失败时不暴露半组 GPU KV 状态，首个成员不会提前 hot eviction。 | +| A2 | MUST | local-first Put 的同一 batch 必须由一个 publish job 持有全部 reservation/pending fence；PutDone 不确定或部分 promotion 时只用同一 identity 整批 roll-forward。声明不完整的 atomic group 禁止退化为逐 key 发布。 | 注入丢响应、取消和局部 promotion 后，所有成员最终一致发布，且不存在半组 hot admission 或提前释放 slot。 | +| S1 | MUST | KV 驱逐和复用单位是单个 slot/key/version；`512 MiB` grant 只是物理容器，不是驱逐、恢复或等待单位。 | 一个 grant 中部分冷 slot 可独立降级、回收并被后续 Put/Get 复用。 | +| S2 | MUST | slot 只有在 `route_live == false && holder_ref_count == 0` 时才能回到 Free；route 释放和 holder 释放的任意先后顺序都不得提前复用。 | 故障注入覆盖 route-first/holder-first，不出现 use-after-free 或重复 free-list entry。 | +| S3 | MUST | owner 容量驱逐与 remote replica 解耦。owner 提交完整、精确的 atomic/TP source cohort；master 只验证并删除该 owner source，不选择 victim，也不要求 CPU route 已存在。最后一份 cache route 允许删除，后续正常 miss/recompute。 | 有 CPU 时只删 GPU source；无 CPU 时 route 消失；remote write 失败只损失后续命中率，不否决 owner 回收。 | +| S4 | MUST | Moka 已发出的 size eviction 必须进入有界、可重试的 exact source-eviction ownership。cohort 元数据暂不完整时保留同一 selected identity 重试，不能静默丢弃、降级成单 key 删除或重新进入 append。 | `size_evictions` 可与 selected/handoff/committed/restored/obsolete/retry 对账；只有物理 Free 或安全恢复 hot 才结清 selected debt。 | +| S5 | MUST NOT | `evict_some`、hot ratio、slot 循环或 grant shrink 不得修改 owner 配置容量。GPU0/GPU1/CPU 实验容量仍为 `128/128/256 GiB`。 | 调优前后 `contribute_to_cluster_pool_size`、Moka max capacity 和 expected-grant 契约可对账。 | +| T1 | MUST | Put/Get/replica/demotion 的结果不确定时必须使用同一 operation ID 重试幂等终态，不得猜测失败并释放可能已发布的 backing。 | 响应丢失/重放后 route、holder、slot 和 quota token 都只提交或归还一次。 | +| T2 | MUST | 相同 `put_id` 的 live-local/GetDone 竞争是可收敛状态，不得作为 `InvalidArgument/Unknown` 永久失败；输家释放新 slot 并复用 canonical holder。不同 `put_id` 才是 stale。 | `cannot replace a live replica` / `could not publish current route` 从可见错误归零，且不会覆盖当前 backing。 | +| P1 | SHOULD | local hit 和 local-ready Put 快路径不同步等待 master/RDMA；跨边界工作按 batch/time-window 合并，避免逐 key `start-transfer-end`。 | 进行 master/RDMA 延迟注入时 local-only 路径延迟和 QPS 不应同比例恶化。 | +| P2 | MUST | 实验必须同时记录 QPS、L1/L2/L3/总命中、local/inflight/leader 分流、CPU source Get、replica/demotion、slot free/used 与业务错误。 | 不得只使用 QPS 或逻辑 admission 声称容量收益、去重或 CPU 已生效。 | +| P3 | MUST | 与 Mooncake 的容量对齐按“进程内 HiCache L2 + 外部 storage”总和计算；GPU 侧 external segment 必须扣除 L2，CPU 无 L2 保留完整 256 GiB。 | 两方总容量均为精确 512 GiB，开跑前 allocated=0，不用重复计入 L2 的基线做对比。 | + +上述规约是后续代码、并发测试和三机压测的验收清单。历史兼容路径可以保留，但不得 +以 fallback 名义跳过 V1/G2/G4/S2/T1 这些正确性门禁。 + +### Mooncake 对齐实验给出的性能约束 + +2026-07-12 的三机对齐实验使用两台 TP=2 GPU 节点和一台纯 CPU 节点，标称容量均为 +GPU0/GPU1/CPU = 128/128/256 GiB，同一批 1152 个 agent 请求中，原版 SGLang + +Mooncake 达到 6.9192 QPS，Fluxon E15 为 3.5134 QPS。Mooncake 没有本设计的 Put +atomic-group admission；它在本轮也没有进入驱逐区。因此该结果证明 group 完整性是 CPU +副本可消费性的必要条件，但不是这轮端到端差距的唯一来源。 + +| 观测 | Fluxon E15 | Mooncake 对齐轮 | 设计约束 | +| --- | ---: | ---: | --- | +| 总命中率 | 44.57% | 93.05% | 先对齐实际 payload 容量和驱逐水位，再比较 group 策略。 | +| cache hit tokens | 12,794,304 | 26,713,408 | Fluxon 多重算 13,919,104 tokens，是吞吐差距主体。 | +| 单 GPU 有效 payload 驱逐点 | 约 57.6 GiB | 本轮约 106 GiB 且未驱逐 | slot 内部碎片和过早水位不能隐藏在相同标称容量后面。 | +| load-back 平均耗时 | 57.1 ms | 5.99 ms | 容量问题消除后，仍需优化 `kernel + layer_first` 读取路径。 | + +E15 的 4,718,592-byte value 被旧 slot class 扩成 8 MiB，payload 利用率 56.25%；再叠加 +旧 Moka 0.8 水位，每个 128 GiB GPU owner 只承载约 57.6 GiB payload 就开始驱逐。 +失败候选立即回插又导致 571 万次请求区间 reclaim candidates 中只有约 4.59 万次成功。 +本设计因此把 exact-fit slot、可配置水位和失败候选退避视为容量控制面的基础正确性，不能 +用新增 atomic group 掩盖这些问题。完整实验设置、计数和日志位置记录在 +`sglang_fluxon_agent_experiments.md`。 + +### 核心实现思路 + +核心实现按上面的 insight 逐条落到接口和状态边界： + +- **对应 L2/L3 分离带来资源冗余和 L2 全局不可见** + - 实现思路：把 L2/L3 从两个后端系统收敛成 Fluxon 的 local side / remote side。local side 承接本机 shared segment、reserve slots 和 holder；remote side 承接跨 owner / 跨机器 KV 数据面。 + - 主要落点：Fluxon master/owner route、`MemHolder`、commit/release 生命周期。 +- **对应同机 worker 缺少共享内存快路径** + - 实现思路：一台机器运行一个 Fluxon owner，同机 SGLang worker 以 external/client attach 到同一个 owner shared segment。worker 只注册本进程 CUDA context 可见的 mapping，底层内存归属和释放由 owner 管理。 + - 主要落点：`wait_local_segments_ready()`、owner shared segment、CUDA host registration。 +- **对应统一 L2/L3 需要分布式 KV 能力** + - 实现思路：继续保留类似 Mooncake 的分布式 KV 目标，用 page key 做全局命名，用 route 和 version 定位本机 owner 或远端 owner，把无推理负载机器上的闲置 CPU 内存作为可放置、可复用的 remote side backing。inflight 判重和 commit future 约束写入发布，避免远端共享副本提前可见。 + - 主要落点：Fluxon master route、local/remote owner 放置、`local_fast_put_start` reservation、`local_fast_put_commit` 发布。 +- **对应保留 L2 低延迟需要 kernel 直连内存** + - 实现思路：在保留分布式 `key -> value` KV 语义的前提下，把普通 `put/get` 的数据面拆成两条 hostless 两阶段路径。Fluxon 不接管 KV page 内部 layout，只提供 key 路由、holder 生命周期和 `plan_ptr(value_ptrs)`。 + - `plan_ptr` 协议：Fluxon 和 SGLang kernel 之间只共享一段短生命周期 plan blob。blob 里保存本次 batch 的 `value_ptrs[]`；Fluxon 负责这些地址的分配、可见性和 holder 生命周期，SGLang kernel 负责按自己的 KV layout 解释并读写这些地址。 + - Put：`put_start -> native write -> put_commit` 构成写入闭环。`put_start` 本质上为本次 batch 分配 Fluxon 管理的 value memory，并返回 `plan_ptr(value_ptrs)`；SGLang native write kernel 直接往这些地址写入 GPU KV bytes；`put_commit` 在 kernel 写完后通知 Fluxon，把这些 slots 发布为可路由的 KV value。 + - Get / restore：`get_start -> get_transfer -> native restore -> release_views` 构成读取闭环。`get_start` 先计算连续安全前缀；`get_transfer` 把可恢复前缀 materialize 成 readable `plan_ptr(value_ptrs)` 并持有 holder；SGLang native restore kernel 直接从这些地址恢复 GPU KV cache，完成后 `release_views` 释放 plan 引用。 + - 主要落点：plan blob ABI、`value_ptrs[]`、`local_fast_put_start/local_fast_put_commit`、`get_start/get_transfer`、SGLang `write_*_to_fluxon_values` / `restore_*_from_fluxon_values`。 +- **对应 KV Cache 适合 write-back 最终一致性** + - 实现思路：KV Cache 的不可变、可丢失语义天然对应 Fluxon 的 master / owner / external 三层架构。master 管全局 route 和最终可见性，owner 管本机共享内存和 holder 生命周期，external 是 SGLang worker 的热路径接入层。 + - Put：native write 完成后先更新 external / owner 本地可见索引并返回 `KvFuture`，后台再推进 master route commit。`KvFuture` 成功前只表示本机 write-back 进入提交流程；响应不确定或可能已发布时必须保留 local-visible/fence 并以同一 identity 整批 roll-forward。只有明确证明 master 从未发布且操作已经终止的 abort 才能清理本地在途状态并退化为后续 cache miss。 + - Get / restore：优先查询 external / owner 本地热路径；本地 miss 时再进入远端 prefetch 和 restore。读取只 materialize 连续安全前缀，失败时 rollback 或按 miss 继续。 + - 主要落点：master route、owner shared segment、external local visible index、`KvFuture`、`get_start/get_transfer/release_views`。 +- **对应适配 SGLang KVCache 特化接口** + - 实现思路：把 restore 查询建模成 SGLang KVCache 专用的 batch prefix planning，而不是逐 key 独立 get。 + - Reject inflight / exist：KV page 按不可变缓存对象写回，重复 page key 不应产生第二份并发写入。`local_fast_put_start` 固定使用 `reject_if_inflight_same_key` 和 `reject_if_exist_same_key`，在分配 writable memory 前拒绝在途或已存在的同 key value，避免重复写回和未提交数据被误发布。 + - Prefetch 对应 `get_start`：`get_start(keys, prefix_best_effort, atomic_group_lens)` 按 key 顺序完成存在性 / prefix 判断。全 owner-local 批次直接返回 `InlineLocal` holder metadata；混合命中或需要远端 materialization 的批次启动 `keys[..transferable_len]` 的后台 prefetch transfer。 + - Batch 和 group 对齐：`get_start` 按 atomic group 边界把 `raw_prefix_hit_len` 收敛成 `transferable_len`；`get_transfer` 只消费这个 handle 对应的 `InlineLocal` 或 `OwnerRpc` 有限分支，并构造 batch 级 readable `plan_ptr(value_ptrs)`，保证 kernel restore 看到的是连续、完整、可回滚的 readable batch plan。 + - 主要落点：`reject_if_inflight_same_key`、`reject_if_exist_same_key`、`GetStartResult`、`atomic_group_lens`、`raw_prefix_hit_len`、`transferable_len`、`prefix_hit_groups`、`get_start/get_transfer`。 + +因此，本设计要统一的是 L2/L3 的对象归属、可见性和生命周期，同时保留 L2 路径的低延迟数据面。单独替换成普通远端 KV 会丢掉 HiCache L2 的 kernel 直连能力；单独保留本地数组或进程内缓存，又无法让 L2 进入全局命中、驱逐和跨节点复用链路。 + +当前调用链主要分为三条主线： + +- 写入侧：`local_fast_put_start -> SGLang native write -> local_fast_put_commit`。 +- 读取侧：`get_start -> get_transfer -> SGLang native restore -> release_views`。 +- 放弃读取侧 restore 时：`cancel_get_transfer` 释放 `get_start` 持有的资源。 + +这个集成把 SGLang 逻辑上的 L2/L3 缓存落到 Fluxon 统一管理的本机/远端 KV 层中。这样可以用同一套 owner、holder、commit/release 语义管理 KV page，减少传统 L2 host cache 与 L3 backend 分属不同系统时产生的同机重复缓存和生命周期割裂。 + +常见部署下，一台机器启动一个 Fluxon owner；同机多个 GPU 对应的多个 SGLang worker 进程通过 external/client attach 到同一个 owner shared segment。相比一个 SGLang 进程一个后端 segment、进程间不共享后端 segment 的形态，Fluxon 把同机 host/shared memory、owner local reserve slots 和 holder 生命周期放到同一个 owner 生命周期模型里。这样多个 SGLang worker 可以通过同一个 owner segment 获得本机快速可见性和受控的本地可写内存供给，减少各进程为了各自安全边界重复持有 KV、固定预留 segment 或独立维护 pin/release 状态带来的浪费。 + +## 范围边界 + +| 范围 | 当前结论 | +| --- | --- | +| SGLang hostless 写入 | 已接入。SGLang 通过 `local_fast_put_start` 取得一批可写 host value 地址，native kernel 写入后再调用 `local_fast_put_commit` 提交。 | +| SGLang hostless 读取 | 使用 `get_start/get_transfer/release_views`。`get_start` 先计算连续可恢复前缀，`get_transfer` 再把可恢复前缀转换成 readable `plan_ptr`。 | +| Fluxon value layout | Fluxon 不理解 KV page 内部布局；只按 `key + value_len` 管理连续字节。 | +| SGLang node 状态 | SGLang 的 `storage_*` 字段是调度层状态，不等同于 Fluxon master route；跨节点复用以 Fluxon commit future 成功为准。 | + +## 总体架构 + +```mermaid +flowchart LR + A["SGLang HiCache radix node"] --> B["HiCacheFluxon backend"] + B --> C["FluxonKVCacheStore"] + C --> D["fluxon_pyo3::KvClient"] + D --> E["Fluxon external / owner"] + E --> F["Fluxon master"] + + B --> G["sgl-kernel kvcacheio"] + D --> H["plan_ptr blob
magic/count/value_ptrs"] + H --> G + G --> I["GPU KV cache"] + E --> J["owner shared segment
local reserve slots / MemHolder"] + J --> H +``` + +这里有两个层次： + +- KV 语义层：SGLang 传入 page key，Fluxon 对外保存 `key -> value`。 +- hostless 数据层：Fluxon 返回 `plan_ptr`，SGLang native kernel 根据 blob 里的 `value_ptrs[]` 直接执行 GPU/host 数据传输。 + +从缓存物理层级看，Fluxon 把传统 L2/L3 逻辑抽象落到 local side / remote side：local side 覆盖本机 GPU KV、owner shared segment、owner local reserve slots 和 `MemHolder`；remote side 覆盖跨 owner 或跨机器的数据面。多个 SGLang worker attach 到同一个 owner segment 时，同机 KV bytes 不需要按 worker 进程重复保存在多个后端 segment 中。 + +`plan_ptr` 只是一轮 backup 或 restore 的短生命周期 carrier。它不能作为 key、缓存地址、跨进程句柄或长期状态保存。 + +## 公共契约 + +本节只列 SGLang HiCache hostless 接入 Fluxon KV 时直接依赖的接口。 + +| 接口 | 层级 | 契约 | +| --- | --- | --- | +| `wait_local_segments_ready()` | Fluxon + SGLang 集成 | 返回当前进程可见的 local segment mapping，供 SGLang 做 CUDA host registration。 | +| `local_fast_put_start(keys, value_len, opts)` | SGLang hostless 写入 | 为一批等长 values 准备可写地址，返回 put `plan_ptr`。 | +| `local_fast_put_commit(plan_ptr)` | SGLang hostless 写入 | 在 SGLang native kernel 写完 `value_ptrs[]` 后消费 put plan，把对应 slots 提交为 Fluxon KV route，返回 `KvFuture`。 | +| `put_abort(plan_ptr)` | SGLang hostless 写入 | 在 commit 前释放 put plan、key reservation 和 local reserve slot lease。 | +| `GetStartResult` | SGLang hostless 读取 | 描述连续命中前缀、可传输长度、atomic group 命中数和第一个 miss 位置。 | +| `GetStartHandle` | SGLang hostless 读取 | 持有一次 get-start 结果和 backend handle；必须被 `get_transfer` 消费或被 `cancel_get_transfer` 取消。 | +| `get_start(keys, prefix_best_effort, atomic_group_lens)` | SGLang hostless 读取 | 按 key 顺序计算连续命中的 prefix，并返回 `GetStartHandle`。 | +| `get_transfer(handle, *, consume_prefix_len=None)` | SGLang hostless 读取 | 消费 handle 的全部可传输前缀，或消费由 `consume_prefix_len` 选定的更短完整 atomic-group 前缀；释放 tail，执行必要 transfer，并返回 readable `plan_ptr`。 | +| `cancel_get_transfer(handle)` | SGLang hostless 读取 | 放弃未 transfer 的 `GetStartHandle`，释放 get-start 期间持有的 owner/external 资源。 | +| `release_views(plan_ptr)` | SGLang hostless 读取 | 释放 get-transfer 产生的 readable plan，丢弃其持有的 holder 引用。 | + +`PutOptionalArgs` 在 SGLang hostless 写入路径中的语义如下： + +| 字段 | SGLang 使用方式 | +| --- | --- | +| `reject_if_inflight_same_key` | 固定开启，避免同一 page key 并发写回造成重复 inflight put。 | +| `reject_if_exist_same_key` | 固定开启，SGLang 把重复 key 当作已写回或冲突重试处理。 | +| `write_through` | 当前配置决定提交策略；调用方显式传入时 Fluxon 必须按字段语义执行。 | +| `make_replica_task` | owner-local write-back 成功后的 batch 级异步副本总开关。 | +| `make_replica_task_mask` | `local_fast_put_start` 的可选逐 key 副本准入结果；长度必须与 `keys` 相同。 | +| `atomic_group_lens` | `local_fast_put_start` 有序 key batch 的原子分组；每项必须大于 0，总和必须严格等于 `keys` 长度。组内副本准入必须一致。 | +| `lease_id` | 当前 SGLang hostless 主线不依赖 lease。 | + +## Key 与组件命名 + +SGLang 传给 Fluxon 的 key 必须先经过 backend namespace 处理： + +```text +storage_key = key_prefix + ":" + logical_key +logical_key = page_hash + optional_component_suffix + config_suffix + optional_extra_backend_tag +``` + +规则： + +- page hash 是 SGLang prefix 复用和 Fluxon KV 存取的共同语义 ID。 +- `PoolName.KV` 使用默认 component；Mamba 等额外 component 通过 suffix 区分。 +- `config_suffix` 编入模型名、TP/PP 等会影响 page layout 的维度，避免不同运行配置复用同一批 physical values。 +- `extra_backend_tag` 用于同一集群内隔离实验或实例，不改变 Fluxon KV 的值格式。 + +Fluxon 只看最终 `storage_key`。page 内部如何拆成 K/V layer、MLA tensor 或 Mamba state，由 SGLang kernel 参数解释。 + +## Segment Registration + +hostless 读写依赖 SGLang 进程可访问的 Fluxon owner segment 已经完成 CUDA host registration。 + +常见部署中，同一台机器上的多个 SGLang worker 连接同一个 Fluxon owner，并映射同一个 owner shared segment。每个 SGLang 进程仍需要在自己的 CUDA context 中完成 host registration；底层内存归属、holder 引用和回收由 owner 统一管理。 + +```mermaid +sequenceDiagram + participant S as SGLang HiCache + participant B as HiCacheFluxon + participant P as Fluxon Python store + participant R as fluxon_pyo3 + participant O as owner segment mapping + participant C as CUDA runtime + + S->>B: register_mem_pool_host / register_mem_host_pool_v2 + B->>P: wait_local_segments_ready() + P->>R: wait_local_segments_ready() + R->>O: wait mapped range + O-->>R: segment_label, write_ptr, read_ptr, len, generation + R-->>P: segment list + P-->>B: dict list + B->>C: cudaHostRegister(write_ptr/read_ptr, len) +``` + +`wait_local_segments_ready()` 返回的 item 至少包含： + +| 字段 | 含义 | +| --- | --- | +| `segment_label` | owner 本地一般为 `cpu:0`；external attach owner 时为 `external_owner:0`。 | +| `write_ptr` | 当前进程可写映射地址。 | +| `read_ptr` | 当前进程可读映射地址，存在时也可注册。 | +| `len` | 映射长度。 | +| `generation` | owner 启动代际，用于拒绝过期 holder 或 mapping。 | +| `node_id` | segment 所属 Fluxon node。 | + +SGLang external-client 模式要求看到 `external_owner:*` segment。注册失败时必须同步报错，不能退回到未注册 host memory 的 direct H2D path。 + +## Plan Blob ABI + +`plan_ptr` 是 Fluxon 返回给 SGLang 的临时句柄，本质上是一段 plan blob 的首地址。SGLang native kernel 通过 `plan_ptr` 找到 blob，再从 blob 里读取本次 batch 对应的 value 地址表。 + +引入 plan blob 的目的不是单纯定义一个新句柄，而是把 KV 数据面的控制面展开从热路径中移走。实际测试中，如果 KV 接口用 Python dict/list 传递 page、layer 和 value 地址信息，或者用 C++ 结构体承载同等信息，调用链仍然会在 Python 层或 C++ 层按 `page_count * layer_count` 做大量遍历、校验和拷贝任务组织；长上下文场景下，这部分 CPU 控制面开销会被显著放大，直接抬高 backup / restore 延迟。 + +因此这里需要一套足够扁平、稳定、可被 native kernel 直接消费的协议：Fluxon 只把本次 batch 已经分配或 materialize 好的 value 起始地址压成 `value_ptrs[]`，通过 `plan_ptr` 交给 SGLang kernel；kernel 再按自己的 KV layout 参数解释这些地址，并直接发起 GPU KV cache 和 Fluxon host value memory 之间的搬运。对 backup 来说，这条路径让 kernel 直接把 GPU KV bytes 写到 CPU host value memory；对 restore 来说，则从这些 host value 地址恢复回 GPU KV cache。这正是开头“hostless 数据面能力”的落点：Fluxon 管地址、路由和生命周期，SGLang kernel 管 KV layout 和真实数据搬运。 + +plan blob 是 Fluxon 在 `local_fast_put_start(...)` 或 `get_transfer(...)` 时创建的一段连续 host memory，格式固定： + +```c +uint64_t magic; // 固定校验值，确认 plan_ptr 指向 Fluxon plan blob +uint64_t count; // value_ptrs 的数量，也就是本次 batch 的 page 数 +uint64_t value_ptrs[count]; // 每个 page 对应的 Fluxon value 起始地址 +``` + +如果 SGLang 一次写入或恢复 10 个 page，Fluxon 会创建一个 blob，并返回一个 `plan_ptr`： + +```text +plan_ptr -> blob 起始地址 + +blob[0] = magic +blob[1] = 10 +blob[2] = value_ptr_0 +blob[3] = value_ptr_1 +... +blob[11] = value_ptr_9 +``` + +`magic` 不是 KV value 地址，只是固定校验值；`value_ptr_0 ... value_ptr_9` 才是 Fluxon 为这些 page 准备的 value 起始地址。它们是当前进程可访问的绝对地址，不是偏移量。 + +`value_len` 不写入 blob。Fluxon 只负责按 `value_len` 分配每个 value 的连续字节区间，并把起始地址放进 `value_ptrs[]`；每个 value 内部如何切成 K/V、layer、MLA 或 Mamba state，由 SGLang 调用 write/restore kernel 时显式传入 layout 参数。 + +`plan_ptr` 只在当前进程、当前 batch 生命周期内有效。`local_fast_put_commit(plan_ptr)`、`put_abort(plan_ptr)` 或 `release_views(plan_ptr)` 后，Fluxon 会清理对应 plan，SGLang 不能继续使用这个 `plan_ptr`。 + +## Backup 时序 + +hostless backup 的核心约束是：Fluxon 负责 GPU KV cache 之下的本机/远端 KV 层的地址分配、route 提交和生命周期管理，但真正的 KV bytes 由 SGLang native kernel 从 GPU KV cache 写入。因此 Fluxon 不能在收到 key 后立刻发布 KV route；它必须先完成 key reservation、put id 分配和 owner-local reserve slot claim，把稳定可写的 `value_ptrs[]` 通过 `plan_ptr` 返回给 SGLang。SGLang native kernel 写完这些地址后，`local_fast_put_commit` 才能把这些 slots 提交为 resident values，并发布 Fluxon KV route。 + +这条 backup 路径按 write-back 最终一致性设计：native write 完成后先进入本地可见和后台 commit 流程，跨 worker / 跨节点可见性以 `KvFuture` 成功为准。`KvFuture` 返回失败时，上层不得把 page 标成全局 `storage_backed`，后续可按 cache miss 处理；但 Fluxon 内部若处于响应不确定或可能已发布状态，仍必须保留 local-visible/fence 并以同一 identity roll-forward，不能把“上层不采用”误解为可以立即释放 backing。 + +### 本地主副本与异步副本控制 + +`write_through` 和 `make_replica_task` 是两个独立契约。`write_through` 选择同步远端放置还是 owner-local write-back；`make_replica_task` 只控制 owner-local commit 成功后是否创建异步副本任务。SGLang 的纯本地模式必须使用 `write_through=false, make_replica_task=false`，不能把关闭副本等同于开启 write-through。 + +| 模式 | `write_through` | `make_replica_task` | 主数据 | 异步副本 | +| --- | --- | --- | --- | --- | +| local-only | `false` | `false` | 当前 GPU owner | 无 | +| local + remote replica | `false` | `true` | 当前 GPU owner | 由 master 的 `replica_task_placement` 选择 | +| write-through | `true` | 不参与异步副本决策 | master 直接选择远端 target | 无额外异步副本任务 | + +`PutOptionalArgs.make_replica_task` 必须贯穿 Python store、PyO3 和 Rust put request。兼容层如果只识别 `write_through` 而丢弃该字段，SGLang 即使记录 `replica_admitted=0`，Rust 仍可能按旧默认创建副本；因此启动检查必须同时确认 Python `PutOptionalArgs` 和 `fluxon_pyo3.KvClient.local_fast_put_start` 的签名都包含 `make_replica_task`。 + +page admission 还要求 `make_replica_task_mask` 贯穿同一条链路。mask 缺省时，Fluxon 把 batch 级 `make_replica_task` 广播到所有 keys，保留 eager-all 和 local-only 的标量契约；mask 存在时，其长度必须严格等于 `keys` 长度，实际逐项决策为 `!write_through && make_replica_task && make_replica_task_mask[i]`。Python store 在调用 PyO3 前检查容器、长度和元素类型，PyO3 在任何 reserve 或 RPC 之前再次检查长度。external 模式把结果写入每个 `ExternalBatchPutStartItemReq.make_replica_task`；owner 模式把整份 mask 保存在 staged plan 中，commit 时按相同 key 顺序写入每个 `OwnerLocalPublishItem.make_replica_task`。 + +Put 与 Get 使用同一种 `atomic_group_lens` 分区语义。SGLang hostless backup 把一个 radix node 的全部 page keys 标成一个 atomic group；副本策略只能整组 admitted 或整组 skipped。`min_replica_pages` 选择完整 group，允许为了满足最小值而超过它；`max_replica_pages_per_batch` 是硬上限，只保留能够完整装入预算的 group，不能截取 group 的前几页。ratio/score 策略先得到 group 级概率和优先级，再展开为组内全 true 或全 false 的 per-key mask。 + +该契约做双层 fail-fast。Python store 要求 group lengths 是严格的 `list[int]`、每项大于 0、总和等于 key 数，并拒绝组内混合 mask；PyO3 在任何 owner slot reserve 或 external RPC 前重复同样校验。未传 `atomic_group_lens` 时，每个 key 视为长度 1 的独立 group，保留已有调用方的逐 key 行为。存在性过滤或冲突重试若只能留下半个 group，SGLang hostless 路径直接放弃该次 backup，不能把剩余 keys 重新声明成一个完整 group。 + +标量降级不能用于 page admission。若一批中只有一页 admitted，把 mask 折叠为 `any(mask)` 会让整个 batch 都创建副本；逻辑 admission 日志与物理 CPU transfer 因而失配。启动自检必须同时确认 Python `PutOptionalArgs` 和 PyO3 签名包含 `make_replica_task_mask`，端到端验收则比较 SGLang admitted page 数、master CPU target 数和 GPU owner 成功 transfer 数，而不能只看 admission 日志。 + +SGLang 的 admitted 只表示该 atomic group 通过客户端准入并为组内 pages 请求创建副本，不表示目标 owner 已经接收整组。当前 replica actor 仍逐 key 完成 transfer/append；相同 key 的并发副本可能在目标侧返回 `KeyBeingWritten`。因此 Put 原子准入消除了策略主动制造的稀疏组，但尚未把远端复制提交升级成组事务。数据面仍需检查“成功 `local_fast_put_start` 行中的 admitted 数 = 成功 transfer 数 + 可逐项对账的目标侧拒绝数”，并继续要求“成功 transfer 数 = `appended=true` 数 = master CPU target 数”。全局 admission 累计计数可能包含随后被 backend 拒绝并重试的尝试，不能直接作为已 staged page 数或物理复制分母。 + +CPU-only 副本模式不新增第二个客户端开关。客户端仍设置 `make_replica_task=true`，master 再通过 `restrict_to_remote_only_node_roles=true` 和 `remote_only_node_roles=["remote_cache"]` 把 replica target 限定到 CPU owner。若 strict 候选为空，任务返回 NoSpace，不回退到其它 GPU owner。 + +端到端验收以实际数据面计数为准：local-only 要求 GPU owner 的 `put_transfer success`、`replica task append done` 和 master `replica_task_target_counts` 全部为 0；CPU-only 副本要求两台 GPU owner 的 transfer/append 总数严格等于 master 的 CPU target 数，并且 GPU target 为 0。仅检查 SGLang 配置日志不足以证明副本策略已经生效。 + +### 弹性本地侧预分配 + +弹性本地侧预分配对应 owner 本地写入预留池。它是在 owner shared segment 中为 SGLang hostless put 预先划分的 writable slots。`local_fast_put_start` 从这些 slots 中为本次 put 分配地址，并把地址写入 `plan_ptr(value_ptrs)` 返回给 SGLang native kernel。此时 slot 只处于 reserved/prepared 状态，还没有绑定为 Fluxon KV 的正式 `key -> value` route。 + +这个 pool 是共享 owner segment 上的弹性本地可写内存供给层。它让多个 SGLang worker 都能快速取得受 owner 生命周期管理的 `value_ptrs[]`，同时避免为每个 worker 固定切出长期独占的后端 segment。reserve slot 不足时可以按需求补充 grant，空闲后再按 cooldown 回收。 + +对象含义： + +| 对象 | 含义 | +| --- | --- | +| grant | owner 侧一次申请的大块本地内存，当前固定为 `512 MiB`。 | +| slot | grant 内按 `slot_size` 切分的小块；一个 slot 承载一个 Fluxon value。 | +| slot lease | `local_fast_put_start` 为本次 batch 临时 claim 到的一组 slots；失败或 abort 时必须释放。 | +| value pointer | slot 的起始地址，会写入 plan blob 的 `value_ptrs[]`，供 SGLang kernel 直接写入。 | +| resident value | `local_fast_put_commit` 后由 slot 构造出的本地可读 value。 | +| route | master/owner 确认后的 key 到 value 位置映射；route 成功后该 value 才是全局可见的 KV replica。 | + +slot 生命周期： + +```text +Free + -> Prepared // local_fast_put_start claim slot + -> PendingLocalVisible // local_fast_put_commit 开始，本地 resident value 已记录为 pending visible + -> Committed // put_done 成功，route 引用该 slot + -> Free // route 和 holder 引用都释放后回收 +``` + +如果 native write 失败，调用方必须执行 `put_abort(plan_ptr)`，Prepared slots 会回到 Free。`local_fast_put_commit` 成功返回后，slot 是否能释放由 route 引用和 holder 引用共同决定；只要 master/owner route 或 `MemHolder` 仍引用该 slot，底层 grant 就不能释放。 + +当前容量策略： + +| 项 | 当前实现 | +| --- | --- | +| grant 物理粒度 | `OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES = 512 * 1024 * 1024` | +| 最小 slot size | `4 KiB` | +| slot size 计算 | `align_up(max(value_len, 4 KiB), 4 KiB)`；按页对齐 exact-fit，不再向上取整到 2 的幂。 | +| slot 上限 | `slot_size <= 512 MiB` | +| refill 触发 | 当前 slot class free slots 不足时登记 pending demand 并唤醒 rebalance actor。 | +| 预期容量预热 | owner 可配置 `owner_local_reserve_expected_capacity: {value_len, payload_capacity_bytes}`。启动时按 grant 内真实可用 slots 计算并申请所需 grants；SGLang 启动脚本必须等预热完成。 | +| expected capacity | 配置预期容量后，`expected_grant_count` 同时是预热目标和物理 grant 硬上限；业务压力只能触发 owner writeback/reclaim，不能把 owner 容量扩到该上限之外。只有未配置 expected capacity 的动态模式可继续按 demand 申请 grant。 | +| 默认等待 | soft wait `10 ms`，hard timeout `30 s`。hard timeout 覆盖一次有界 replica/reclaim 背压事务，不用于掩盖 master 调度饥饿，也不替代 terminal/pending 对账。 | +| shrink 单位 | 整个 grant；不做 live grant compaction。 | +| Moka 计费权重 | committed slot 使用 4 KiB 对齐后的实际 `slot_size`，而不是原始 `value_len`；容量驱逐因此对应真实内存占用。 | +| Moka 容量水位 | owner `replica_writeback_hot_capacity_ratio` 控制 local hot/writeback 窗口；master `replica_cache_capacity_ratio` 只约束 `Allocation && !owner_local_indexed && lease_id.is_none()` 域。两者都不改变 owner segment 或 expected grants。 | +| 常规容量驱逐 | owner hot Moka 的 Size 事件只转移“写回任务 ownership”，不直接释放 backing；master 的环 B Moka Size 事件进入 unindexed Allocation reclaim。两个环分别闭合。 | +| 物理 NoSpace 兜底 | owner pressure actor 根据 free-slot 缺口减去 selection debt，只在这一处调用无谓词 `evict_some`（单次 ≤256 MiB、≥200 ms 间隔）启动 writeback。旧 route 全表扫描已删除；remote write 失败不进入 master drop/reclaim。 | +| 容量不变量 | `evict_some` 只驱逐一部分条目，不修改 Moka `max_capacity`，也不修改 owner 配置容量。 | +| route 维护反压 | route 发布后的 prefix-index/Moka 更新使用既有有界 actor；队列满时 `put_done` 等待，避免 metadata/Moka 水位落后于物理分配。它不负责 owner slot-pressure victim 扫描。 | + +Qwen3-VL-8B、TP=2、page size 64 的单 rank K/V value 是 4,718,592 bytes，正好是 +4 KiB 的整数倍。exact-fit 后一个 512 MiB grant 可以提供 113 个 slots；旧的 +`next_power_of_two` 会把每页扩成 8 MiB，一个 grant 只能提供 64 个 slots。以 128 GiB +owner 为例，旧 `0.8` 水位叠加 56.25% slot payload 利用率后只有约 57.6 GiB 的有效 KV +payload，会在物理 owner 仍有大量空闲时提前驱逐。exact-fit 与 0.95 水位分别解决内部碎片 +和过早驱逐；两项都不修改 owner 的 128 GiB segment 容量。 + +预期容量的 grant 数必须计入每个 grant 尾部不能容纳完整 slot 的碎片，而不能简单用总 payload +除以 512 MiB。计算为： + +```text +slot_size = align_up(value_len, 4 KiB) +slots_per_grant = floor(512 MiB / slot_size) +expected_grants = ceil(payload_capacity_bytes / (slots_per_grant * slot_size)) +``` + +Qwen3-VL 本轮 `value_len=4,718,592`、0.8 payload 目标 109,951,162,777 bytes 对应 +207 grants；可用 23,391 slots / 110,372,585,472 bytes，实际物理预留 +111,132,278,784 bytes。owner 配置容量仍为 128 GiB。配置层会拒绝非正 value/capacity/timeout、 +`hard_timeout <= soft_timeout`、value 大于 512 MiB、预热物理容量超过 owner DRAM，以及在 +master、external client 或 side worker 上设置 owner-only expected capacity。claim 超时统一 +返回 local-reserve timeout 类错误，并用 `stage=claim_turn|refill` 和完整 pool 状态区分等待 +claim 顺序与 refill/reclaim 失败。 + +route 发布与容量记账必须保持有界距离。local-first `BatchPutDone` 可以一次发布上百个 page； +如果每个 page 都独立 spawn prefix/Moka 维护任务，route 已提交和 owner slot 已占用的速度会 +超过 Moka 记账速度。此时即使 segment 仍有大量空闲，master 运行时也可能被任务和逐 page +日志占满，使新的 512 MiB reserve RPC 超过 hard timeout。当前实现用一个有界批处理 actor +统一消费 post-route 事件：一次 prefix-index write lock 覆盖整批，随后逐项推进 Moka；消费 +前校验当前 route 仍匹配 `put_id + owner replica`。容量 512 的队列同时承担反压边界，不能 +改回无界任务生成，也不能用增加 owner 容量代替该控制面约束。 + +底层物理释放收束在 grant 级别。单个 committed slot 只是 grant 内逻辑索引，不直接拥有释放整块 mmap/registered memory 的权力。 + +### Fluxon put_start / put_commit 接口设计 + +Fluxon 的 backup 解法是把普通 KV `put(key, bytes)` 拆成 `put_start -> native write -> put_commit`。这个拆分的核心是接口解耦：Fluxon 先给本次写入分配受 owner 生命周期管理的 value memory，并把地址交给 SGLang；SGLang native/CUDA kernel 负责真正把 GPU KV cache bytes 写入这些地址；kernel 写完后，SGLang 再通知 Fluxon 发布这批 value。这样普通 KV 接口里的 payload materialization、CPU 侧 bytes 拆装和逐 page/layer 拷贝组织都不会进入 backup 热路径。 + +`put_start` 不是一次可见写入，它的核心结果是内存分配和指针返回。Fluxon 为本次 batch claim owner-local reserve slots，生成 `plan_ptr(value_ptrs)`；`value_ptrs[]` 指向 Fluxon owner segment 中已经准备好的 writable value memory，SGLang kernel 可以直接写这些地址。判重、reservation 和 put id 是保护这次内存分配不被重复写入或提前发布的控制面约束；这些 value 在 `put_commit` 前不会作为可读 route 暴露。 + +`put_commit` 是 kernel 写完后的通知和发布点。SGLang 在 CUDA write 完成后调用它，Fluxon 消费 put plan，把对应 slots 转成 resident values，并异步推进 route commit。`KvFuture` 成功前，数据只表示本次 write-back 已进入 Fluxon commit 流程；`KvFuture` 成功后，page 才能成为跨 worker / 跨节点可复用的 shared backing。这样 `put_start -> native write -> put_commit` 就构成了 hostless put 的完整闭环。 + +这套接口拆分直接使用上面的弹性本地侧预分配池：`put_start` 从预分配 reserve slots 中 claim 短生命周期 slot lease，热路径拿到的是已准备好的地址。 + +| 解法层 | 做什么 | 解决的问题 | +| --- | --- | --- | +| `local_fast_put_start(keys, value_len)` | 从 owner-local reserve slots claim 本次 batch 的 slot lease，分配 writable value memory，并返回 `plan_ptr(value_ptrs)`；key reservation 和 put id 保护这次分配 | 在 native write 前只暴露可写地址，不发布可读 route，避免其它 get 读到尚未写完的 value | +| SGLang native write | SGLang native kernel 根据 `plan_ptr(value_ptrs)` 把 GPU KV page 写入 Fluxon 管理的 host value memory | 避免普通 KV `put(bytes)` 路径在 CPU 侧拆装 payload 和逐 page/layer 组织拷贝 | +| `local_fast_put_commit(plan_ptr)` | 消费 put plan，把已写完的 slots 转为 resident values，推进 transfer / route commit，并返回 `KvFuture` | 把本地写完和全局发布分开；只有 future 成功后，SGLang 才能把 node 标记为 `storage_backed` | +| 弹性本地侧预分配 | owner shared segment 中按 slot size 管理 reserve slots；free slots 不足时按 pending demand 补充 grant，空闲后按 cooldown 回收 | 热路径直接 claim 已准备好的可写地址，避免每个 worker 固定独占 segment，也避免每次 backup 临时分配或注册 host memory | + +写入阶段的拆分主要是为了同时满足三个约束： + +- 数据面效率：SGLang 不需要先把 GPU KV page 包装成通用 KV payload 再交给后端，而是直接用 kernel 写入 Fluxon 返回的 value 地址。 +- 可见性安全：`put_start` 阶段只预留地址，不发布 route；避免其它 get 读到尚未写完或尚未 commit 的 value。 +- 容量弹性：本地侧预分配池提供短生命周期 slot lease，而不是为每个 SGLang worker 固定切出长期独占内存；容量不足时由 owner 侧 refill，空闲后按 grant 粒度回收。 + +```mermaid +sequenceDiagram + participant U as UnifiedRadixCache + participant B as HiCacheFluxon + participant F as Fluxon store + participant K as sgl-kernel + participant O as Fluxon owner + participant M as Fluxon master + + U->>B: local_fast_put_start(page_keys, value_len, atomic_group_lens) + B->>F: local_fast_put_start(storage_keys, value_len, opts(mask, groups)) + F->>O: claim local reserve slots / external owner offsets + O-->>F: plan_ptr(value_ptrs) + B-->>U: plan_ptr + U->>K: write_*_to_fluxon_values(plan_ptr, page_indices, layout ptrs) + K-->>U: writes queued on CUDA stream + U->>U: record local_ready_event + U->>B: local_fast_put_commit(plan_ptr) after event ready + B->>F: local_fast_put_commit(plan_ptr) + F->>O: record precommit visible / transfer_end or put_done + O->>M: commit route + F-->>B: KvFuture + U->>U: scheduler poll future +``` + +hostless backup 默认不依赖 put 前 exists 扫描。重复 key 或在途 key 由 `local_fast_put_start` 的 `reject_if_exist_same_key` 和 `reject_if_inflight_same_key` 准入语义处理；SGLang 上层按冲突错误做重试或跳过。 + +`local_fast_put_start(keys, value_len)` 的要求： + +- `keys` 不能为空。 +- `value_len` 必须大于 0，且同一批 keys 共享同一个 value size。 +- `atomic_group_lens` 存在时必须由正整数构成，且总和严格等于 `keys` 长度；同组的 `make_replica_task_mask` 值必须完全一致。 +- SGLang 的 radix node backup 默认使用 `[len(node_page_keys)]`。过滤已有 key 或冲突重试不能切开该组；无法保留完整组时本次 backup 失败关闭。 +- SGLang 必须在 `local_fast_put_commit` 前完成 native write；写入失败时必须调用 `put_abort`。 +- `local_fast_put_commit` 只能调用一次；调用后 plan 从 registry 清理，后续只能等待返回的 `KvFuture`。 + +commit 请求会带上 `len`、`src_offset` 和 target 信息。Fluxon 用这些字段判断本次 value 是否落在当前进程可访问的 owner segment 或 owner-local reserve slot 中；如果需要本地可见索引，`put_done` 会返回 owner 分配的 `local_cache_holder_id`。`MemoryInfo` 和 holder 生命周期在下文说明。 + +## Prefetch 与 Restore 时序 + +hostless restore 的核心约束是：SGLang 只能恢复有序 page keys 的连续前缀，并且不能切开一个 radix node 对应的 atomic group。Fluxon 需要先在本机/远端 KV 层里判断这批 keys 的可恢复边界，再把真正可恢复的部分提前 materialize 到当前进程可读的 holder / memory view 中，最后转换成 SGLang kernel 可读取的 `plan_ptr(value_ptrs)`。 + +因此读取侧分成 prefetch 和 restore 两段。当前实现把 start-to-transfer 计划显式限定为两个分支： + +- Prefetch：`get_start(keys, prefix_best_effort, atomic_group_lens)` 按 key 顺序做 local visible check / owner get start，计算 page 级连续命中前缀 `raw_prefix_hit_len`，再按 `atomic_group_lens` 向下收敛成 `transferable_len`。全 owner-local 批次返回 `InlineLocal { items }`：owner 为每个导出 page 安装独立 external holding，并把 offset、len、holding ID 和 owner generation 随 start response 返回；该分支不创建 transfer registry，也不启动后台 transfer。混合命中、master fallback 或需要远端 materialization 的批次返回 `OwnerRpc`，继续由 shared op 保存后台 prefetch 的 holder / transfer result。 +- Restore：SGLang 拿到 `transferable_len` 后再决定是否分配 GPU KV pages 并继续恢复。`get_transfer(handle, consume_prefix_len=...)` 可以消费不超过 `transferable_len` 的更短完整 atomic-group 前缀；省略该参数时消费全部可传输前缀。`InlineLocal` 分支在 external 进程校验 owner generation 和 mmap 范围后，只为消费前缀构造 holders，tail holding IDs 进入现有 holder-ACK 合批队列；`OwnerRpc` 分支在同一 transfer handler 内 drop tail items，只等待或取得消费前缀的 prefetch 结果。两个分支最终都生成持有 holder 引用的 readable `plan_ptr`，随后 SGLang native kernel 使用 `restore_*_from_fluxon_values(...)` 把 Fluxon value memory 拷回 GPU KV cache。 + +读取阶段拆成 prefetch 和 restore 两步，主要是为了保证： + +- prefix 安全：中间 page miss 时，只恢复连续命中的完整前缀，不构造带洞的 GPU KV 状态。 +- atomic group 安全：`transferable_len` 不会切开 radix node group，避免恢复半个 node。 +- 资源效率：SGLang 在知道可恢复边界后再分配 GPU KV pages。`OwnerRpc` 可以重叠远端 materialization 和 GPU page 分配；`InlineLocal` 省去第二次 owner transfer RPC 和无用后台任务。 +- 生命周期安全：`get_transfer` 返回的 plan 持有 holder 引用，直到 `release_views(plan_ptr)` 后才释放，保证 kernel restore 期间 value 地址稳定。 + +```mermaid +sequenceDiagram + participant U as UnifiedRadixCache + participant B as HiCacheFluxon + participant F as Fluxon store + participant E as external client + participant O as Fluxon owner / external + participant K as sgl-kernel + + U->>B: get_start(page_keys, atomic_group_lens) + B->>F: get_start(storage_keys, prefix_best_effort, atomic_group_lens) + F->>E: batch_get_start(keys) + E->>O: ExternalBatchGetStartReq + O->>O: local visible snapshot / owner get start + O->>O: compute raw_prefix_hit_len and transferable_len + alt fully owner-local + O->>O: allocate external holding IDs and install holdings + O-->>E: handle + InlineLocal(items, generation) + E->>E: cache inline plan by handle + else mixed hit or remote materialization + O->>O: create shared op and prefetch keys[..transferable_len] + O-->>E: handle + OwnerRpc + end + E-->>F: backend handle + raw_prefix_hit_len + F->>F: build GetStartResult(raw_prefix_hit_len, transferable_len, ...) + F-->>B: GetStartHandle + GetStartResult + B-->>U: transferable_len / prefix_hit_groups / first_miss_index + U->>U: allocate GPU KV pages for transferable prefix + alt transferable_len > 0 and caller chooses restore + U->>B: get_transfer(handle, consume_prefix_len) + B->>F: get_transfer(handle, consume_prefix_len) + F->>E: batch_get_transfer(handle, consumed prefix) + alt InlineLocal + E->>E: validate generation and mmap range + E->>E: enqueue tail holding IDs to batched ACK + E->>E: construct consumed-prefix holders + else OwnerRpc + E->>O: ExternalBatchGetTransferReq(consume_prefix_len) + O->>O: validate atomic boundary and drop tail items + O-->>E: consumed-prefix holders / transfer results + end + E-->>F: plan_ptr(value_ptrs, holders kept alive) + F-->>B: plan_ptr + B-->>U: plan_ptr + U->>K: restore_*_from_fluxon_values(plan_ptr, prefix page indices, layout ptrs) + K-->>U: H2D queued on CUDA stream + U->>B: release_views(plan_ptr) after restore finalizer + else caller gives up restore + U->>B: cancel_get_transfer(handle) + B->>F: cancel_get_transfer(handle) + F->>E: cancel_batch_get_start(handle) + alt InlineLocal + E->>O: ExternalBatchGetCancelReq(holding IDs) + O->>O: release exact external holdings + else OwnerRpc + E->>O: ExternalBatchGetCancelReq + O->>O: release shared op / prefetched holders + end + end +``` + +`GetStartResult` 的关键字段如下： + +| 字段 | 含义 | +| --- | --- | +| `raw_prefix_hit_len` | 按 key 顺序连续命中的 page 数，未按 atomic group 收敛。 | +| `transferable_len` | `get_transfer` 最多可消费的 page 数；它不会切开 atomic group。 | +| `prefix_hit_groups` | 完整命中的 atomic group 数。 | +| `first_miss_index` | 第一个 miss page 的 index；全部命中时为 `None`。 | +| `first_miss_group_index` | 第一个 miss 所在 atomic group；全部命中时为 `None`。 | +| `all_hit` | `transferable_len == len(keys)`。 | + +生命周期规则： + +- `get_start` 成功后，调用方必须二选一：`get_transfer(handle)` 或 `cancel_get_transfer(handle)`；放弃 restore 时必须取消 handle，释放可能已经启动的 prefetch 资源。 +- `get_transfer(handle, consume_prefix_len=...)` 的值必须大于 0、不超过 `transferable_len`，并精确落在 atomic-group 边界；参数校验失败时 handle 仍可 cancel 或用合法长度重试。 +- `get_transfer` 成功后，handle 已被消费；tail 资源同时移交给 owner 直接 drop 或 external holder-ACK 合批队列，后续只由 returned `plan_ptr` 和 `release_views(plan_ptr)` 管理消费前缀。 +- `release_views(plan_ptr)` 必须在 native restore 完成后执行，即使 native restore 失败也要释放。 +- `InlineLocal` holding 从 owner 返回 start response 前已经安装；全量 `get_transfer` 把它们移交给 returned plan，部分前缀 `get_transfer` 只移交前缀并将 tail holding IDs 无等待送入 holder-ACK 合批队列，`cancel_get_transfer` 则携带全部 holding IDs 精确释放。owner 重启或 generation 不匹配时不能继续使用旧 inline plan。 +- `get_start` 只命中部分前缀时，SGLang 只能恢复 `transferable_len` 覆盖的完整 atomic groups，不能构造半个 atomic group 的 GPU KV 状态。 +- `get_transfer` 返回 miss / KeyNotFound 时，SGLang 必须放弃本次 restore 并执行 rollback。 +- owner 必须为未被 `get_transfer/cancel_get_transfer` 消费的 handle 设置有界 TTL；TTL 回收与 + 显式 cancel 走同一条 plan Drop 路径，不能永久 pin local holder 或 per-key result。 + +### Batch Get 的逐 key 交集、pin 与 singleflight + +hostless restore 的 batch 通常不会完全相同。多个会话或同一会话的不同轮次会共享 +大量前缀 page keys，但各自的 batch 长度、后缀和 `atomic_group_lens` 不同。因此以 +`keys + atomic_group_lens + prefix_best_effort` 整个结构作为去重 key，只能合并完全相同的 +batch；它不能阻止两个重叠 batch 对同一 page 重复申请 prepared slot、重复发起 Get +和重复 DMA。 + +读路径必须在 owner 侧按单 key 分流，同时保留原 batch 的顺序和 group 边界。对一个有序 +required batch `R`，定义： + +```text +L = 通过 owner reclaim fence 后可以立即取得 holder 的 local-visible keys +F = 当前 owner 上已有可共享 Get 操作的 inflight keys + +LocalJoin = R ∩ L +InflightJoin = R ∩ F +Leader = R - (L ∪ F) +``` + +这里的“逐 key”只指 singleflight registry 的判定和线性化粒度，不指 RPC 粒度。 +owner 逐 key 完成 `R` 的完整分流后，必须把全部 `Leader` 按原 index 压缩成 +`leader_keys[] + prepared_targets[] + original_indices[]`，通过一次 BatchGetStart 发布；全部终态 +使用 BatchGetDone/Revoke 收敛，最后才按 `original_indices[]` scatter 回 `R`。当前 payload +backend 仍对每个 key 调用一次 `transfer_data_no_copy`，因此只能表述为“控制面已 batch、数据面 +descriptor batch 待实现”，不能表述成已按 peer/transport 成批提交。后续应直接按 peer/transport +汇总 descriptor，或使用有界 time-window 聚合兼容传输；不得为了 per-key future 在循环中逐个 +执行完整的 `GetStart -> transfer -> GetDone`。 + +`L` 只包含能在当前 owner generation 内取得稳定 `MemoryInfo` 的 +`precommit_local_visible_info` 和 `get_cached_info`。只有 master route、只有 +`local_snapshot_info`、已进入 reclaim fence，或只存在 hidden pending 但还没有可共享 +future 的 key，都不得当作 local hit。`pending_local_get_info` 属于 `F` 的生命周期， +不应通过把未发布 bytes 提前放入 `L` 来规避竞争。 + +一次 batch 的必须时序如下： + +1. owner 使用 256-shard key-control table 逐 key选择 local / join / leader；每次只持一个 + key 所属 shard 的短临界区，禁止持锁遍历 `R`，也禁止任何锁跨 `.await`。 +2. `LocalJoin` 在离开该 key fence 前 clone 对应的 `Arc`，以 resident holder 引用 pin + 住 backing。即使 route 随后被 hot demotion，该 slot 也要等本次 holder 释放后才能 + 变成 Free。 +3. `InflightJoin` 在离开 registry 前取得逐 key `Interest` guard；该 guard 必须能观察 + 同一个终态，不得另外申请 slot 或发起 DMA。prefix decision 由 guard 的所有权表达， + request future 在任一 `.await` 被取消时，`Drop` 必须自动归还 `undecided`，不得依赖 + `decision_registered: bool` 和尾部手工减计数。 +4. `Leader` 使用逐 key 原子 entry 操作先安装 shared op，再加入本轮压缩后的 + master BatchGetStart。不能先对 inflight set 做一次普通 snapshot，稍后再批量 + insert；否则两个 batch 仍可以同时观察到空集并都成为 leader。 + 若规划在 BatchGetStart handoff 前取消，`ExternalGetPlanningLeadersGuard` 必须把本轮新建且 + 仍为 `Starting` 的 op 发布为安全 miss、按 Arc identity 清 marker 并唤醒 joiner；此时尚无 + prepared target/master operation，所以不得发 Revoke，也不得留下孤儿 marker。 +5. 只有 `LocalJoin` 和 `InflightJoin` 已经取得 pin，且所有 `Leader` 都已发布 + inflight marker 后，才能对 leader 子集发起 master GetStart 和后续 payload transfer。 +6. leader 的压缩响应必须按原 index 放回与 `R` 等长的 result slots，local、joiner + 和 leader 的状态全部拼回后，才可计算 `raw_prefix_hit_len`。不得在压缩后的 + leader 列表上计算 prefix。 +7. `transferable_len` 仍使用原始 `atomic_group_lens` 向下收敛。一个 group 可以同时 + 包含 local、joiner 和 leader，但只有所有成员均可读时才能对上层发布完整命中。 +8. 超出 `transferable_len` 的后缀必须立即释放 local pin、减少 inflight waiter，并对 + 已 start 的 leader 执行 revoke。未发布的 prepared slot 必须返回 Free，不能因为该 + key 不在连续前缀内就保留到 TTL。 +9. 保留前缀的 local holder、joined future 和 leader future 一直存活到 `get_transfer` + 成功移交给 readable plan，或 `cancel_get_transfer` 明确取消。 + +请求级不得再叠加 exact-batch shared-op 状态机。每个 external handle 直接持有自己的 +`BatchPlan(keys + Local/Interest items)`；完全相同和部分重叠的 batch 都自然在 per-key +registry 合流。leader 子集在进入第一个网络 `.await` 前移交给一个注册到 owner task +registry 的 cohort task，该 task 独立完成 `BatchGetStart -> transfer -> BatchGetDone/Revoke`。 +因此 external Start RPC 被取消只会 Drop 本请求的 Interest，不会让已经安装的 leader 失去 +executor。该收缩同时删除 request 级的第二套 phase/Mutex/Notify、waiter_count、完整 key +向量结果缓存和纯聚合后台 task，但不改变公开的 start/transfer/cancel API。 + +prepared local-reserve claim 也必须由 guard 持有。BatchGetStart 明确接受的 target 逐项从 +guard disarm 并转交 per-key flight。已经收到响应、但响应长度或 target identity 异常时， +所有已返回 `get_id` 的成功项必须先移交注册 Revoke cleanup，只有明确未接受的 slots 才能 +立即返回 Free。若 BatchGetStart transport 失败且完全没有响应，owner 不知道 master 是否已经 +接受，也没有 `get_id` 可以 revoke；这类 slots 不得立即复用，当前实现隔离 65 秒，覆盖 master +60 秒 inflight Get TTL 后再尝试释放。该 TTL 隔离是未知提交点的安全兜底，不替代后续带 +request identity 的 terminal query 协议。 + +`Revoking` 状态必须保留完整 `get_id + prepared_target`；RPC 超时、响应丢失或长度/identity +不匹配时保留 marker 和 slot 并重试，只有 master 明确确认 Revoke 后才释放 slot。若 Done +已胜出，则 Revoke 失败方不得释放已经 committed 的 slot。BatchGetDone 同样必须携带原始 +`get_id` 集合幂等重试，并逐项校验响应 identity；不能在固定三次重试后把 +pending-visible slot 遗留为无 owner 状态。owner shutdown 是唯一允许停止该重试的进程级终态。 + +per-key shared op 至少要区分下列阶段： + +| 阶段 | 对 prefix / 资源的含义 | +| --- | --- | +| `Starting` | leader marker 已经对后续 batch 可见，但 master 尚未接受；joiner 只能等待，不能把它计为命中。 | +| `Started` | master 已返回成功的 `put_id/get_id` 和精确 target；可用于计算连续 prefix，但 bytes 还不得暴露给 kernel。 | +| `Finishing` | 至少一个 Interest 保留该 key；唯一 cohort executor 正在 transfer/GetDone，后续请求只能 join。 | +| `Revoking { item }` | 无 Interest 保留该 key；`item` 持有 get_id 和精确 target，直到 Revoke 明确终止并释放 slot。 | +| `Ready` | transfer、幂等 GetDone 和 local promotion 已完成；所有 waiter 共享同一个 `Arc`。 | +| `Failed/Revoked` | 终态错误对所有 waiter 一致可见；leader 只释放一次 prepared slot 和 master inflight。 | + +shared op 的 registry 清理必须同时比对 key 和 op identity/generation，只删除当前这一个 +terminal op。旧 waiter 的延迟 drop 不能删除同 key 随后创建的新 op，这与现有 +`key + Arc::ptr_eq` 的 generation 清理原则相同。同一 required batch 内即使出现 +重复 key，也只能有一个 leader/shared op；原始的多个 index 分别持有结果引用。 + +owner 侧 singleflight 是主去重边界，master 仍需要防御性地保证一个 requester 不发布两个 +local committed targets。master 防重粒度必须是 `(key, requester_owner)`，不是全局 +`key`；GPU0 和 GPU1 同时把同一 key materialize 到各自 owner 是合法的，不应被串行化。 +prepared Get 从 master GetStart 接受到 GetDone/Revoke/TTL 终态期间占有该 +`(key, requester_owner)`。 + +master 已经看到 requester 上有相同版本 live replica，或 GetDone 时另一个操作已经 +发布同版本 route，都属于可收敛竞争，不得使用 `InvalidArgument/Unknown` 当作永久 +失败。master 应返回可区分的 busy/already-exists 终态：owner 释放本次多申请的 +prepared slot，在 reclaim fence 下重新获取 canonical local holder；正在 demotion 的 key 则等待 +route 移除后从 remote replica 重试。相同 `put_id` 的 GetDone 竞争必须幂等收敛到唯一 +local route，输家不覆盖 canonical backing；不同 `put_id` 才是 stale Get，必须 revoke。 + +该路径至少要暴露 `required/local_pinned/inflight_joined/leader_started`、被省略的 +GetStart/DMA 数、suffix local-pin/waiter/revoke 数、join 等待时间以及 conflict retry 数。 +当前 owner 周期快照已暴露 active handles、per-key flights 的 Starting/Finishing/Revoking、 +`undecided/retained` interest，以及 local-reserve 的 Free/Prepared/Pending/Committed slot 数； +workload 结束后这些临时态必须在有界时间内回到基线。 +`prepared local-reserve Get target cannot replace a live replica` 和 +`prepared local-reserve Get target could not publish current route` 应该从可见错误中归零； +仅看最终 QPS 不能证明重叠 batch 已经正确去重。 + +## Fluxon 本地可见索引与生命周期 + +Fluxon client/external 侧会维护当前进程可直接访问的 value 索引，以及 get/put plan 持有的 holder 引用。SGLang hostless 路径主要涉及下面几类状态： + +| 状态 | 创建入口 | 生命周期 | +| --- | --- | --- | +| precommit local visible | `local_fast_put_commit` 开始后，由 owner-local reserve slot 对应的 `MemoryInfo` 记录到 `precommit_local_visible_info` | 表示 value 已在当前进程可读但 global publish 尚未收敛；成功后转为 committed entry。响应不确定/局部失败时保留并 roll-forward，只有证明从未发布的显式 abort 才移除。 | +| committed local visible info | put commit 成功后，由 SGLang 进程内的 Fluxon external/client 将 `MemoryInfo` 记录到 `get_cached_info` | 保存 key、put version、`holder_id`、`offset`、`len` 和 owner node；后续 `get_start/get_transfer` 可以复用这份 `MemoryInfo` 构造 readable plan。 | +| inline external holding | owner 判定整个 transferable batch 都在本地可见后，为每个导出 page 分配 external holding ID 并安装 holding | start response 到达 external 后由 handle 缓存；随后必须移交给 `get_transfer` plan 或由 cancel 携带 holding IDs 释放。holding 绑定 owner generation。 | +| get-transfer holder | `get_transfer` 成功后绑定到 readable plan，并由 plan 持有引用 | `release_views(plan_ptr)` 后释放引用；plan 生命周期内 holder 保证对应 value 不被释放。 | + +Put 的本地可见性和全局发布是两个独立的线性化点，不得把两者合并成 +“等 master 后才能本地读”： + +| 线性化点 | 条件 | 立即允许的操作 | +| --- | --- | --- | +| `local-read-ready` | native kernel 已把 bytes 写入 local-reserve slot，且对应 CUDA stream 可见性已满足；client 将唯一 `MemoryInfo` 安装到 `precommit_local_visible_info` | 同 owner 的 `local_visible_mem_holder(s)` 立即可以命中、pin 并读取该 slot，不等待 master PutDone。 | +| `global-route-ready` | master PutDone 已按 key/version 发布 route，并完成必要的 route maintenance | 该 slot 才进入 `Committed(route_live=true)`，可被其它 owner 通过全局 route 发现，并可进入 owner-hot Moka 参与后续分级。 | + +因此延后的只是 global committed 和 hot admission，不是 local read。之所以不在 +`local-read-ready` 时立即加入 owner-hot Moka，是因为 Size eviction 可能马上启动 +replica/demotion/reclaim；此时 master route 还未发布，会使回收协议无法用同一 +`put_id + backing identity` 证明对象归属。precommit entry 本身持有 resident holder，已经 +足以保证本地读期间地址稳定，不需要依赖 hot Moka。 + +同一 atomic/TP cohort 可以逐项进入 `local-read-ready`，但 batch Get 仍必须在原 +`atomic_group_lens` 边界上判定可恢复前缀；这保留了“本地已写页立即可复用”， +又不会向 SGLang 恢复半个 radix node。master 发布不确定或部分成功时保留 local-visible +entry/fence 并整批 roll-forward；只有明确从未发布的 abort 才能移除 precommit。已经取得的 +local-read holder 仍可安全完成当前读，最后一个 holder 释放后才能回收 slot。 + +收到 `local_cache_holder_id` 后，SGLang 进程内的 Fluxon external/client 使用 `holder_id`、`offset` 和 `len` 构造 `MemoryInfo`，并记录到自身 `get_cached_info`。`MemoryInfo` 记录当前进程访问该 value 所需的地址信息和释放动作；后续 `get_start` 命中 `get_cached_info` 时，可以直接把这份 `MemoryInfo` 纳入本次 get 结果，`get_transfer` 再把这些 value 地址写入 readable plan。底层内存的回收由 owner route、holder 引用和 owner-local reserve grant 生命周期共同约束。 + +内部 `MemoryInfo.holder_id` 与 external holding ID 属于两个生命周期域。resident local-reserve page 的内部 holder ID 可以为 0；每次 external export 仍由 owner 分配非零、递增且独立的 holding ID，同一内部 backing 的并发导出也不能复用 ID。这样 cancel/delete ACK 可以按本次 export 精确对账，多个 resident page 不会因内部 ID 相同而覆盖 holding。external 在构造 inline holder 前还必须验证 response generation 与当前 mmap owner generation 一致，并验证 `offset + len` 位于 mapping 范围内。 + +`precommit_local_visible_info` 覆盖 local-ready 到 global route/promotion 完成之间的窗口。 +put commit 成功后会原子转入 committed index；PutDone 响应不确定或局部失败时 publish job +必须继续持有 precommit/fence 并以同一 identity roll-forward，不能清理可能已被 master 发布的 +backing。只有明确证明 master 从未发布且操作已终止的 abort 路径才可移除 precommit。 + +## Local-side 主动驱逐与安全回收 + +local-side 主动驱逐分成“owner 选择冷 source”和“master/owner 安全删除该 source”两个 +阶段。owner hot Moka 的 Size removal 或 pressure actor 的 `evict_some` 只完成候选选择； +dispatcher 随后提交完整、精确的 atomic/TP source cohort。该容量事务不查询 CPU replica， +也不等待 remote append：有其它 replica 时只删当前 owner source，没有其它 replica 时允许 +删除最后一份 cache route，后续 Get 正常 miss/recompute。proactive remote replica 是独立的 +命中率优化，失败只代表少一份可命中的 CPU copy，不拥有 owner 容量释放的否决权。 + +### 触发与职责边界 + +- owner hot Moka 以 committed slot 的真实 `slot_size` 计费，value 只保存 + `Weak`，不因热度索引额外 pin backing。 +- synchronous eviction callback 只发送 `(key, put_id, Weak, weight)` 到 lossless metadata + channel；callback 不取 key-control、不查 index、不展开 cohort、不 RPC/RDMA。 +- async dispatcher 从当前 local index point-pin `Arc`，在锁外展开完整 atomic/TP cohort， + 构造 `OwnerSourceEvictionMember { key, put_id, backing }`；不使用 cohort-wide 同步锁。 +- `evict_some(requested_weight)` 是 owner slot-pressure actor 的唯一生产调用点；参数是本轮 + 希望选择的字节数，不是新容量。单次 ≤256 MiB、最短间隔 200 ms，调用前后 + `max_capacity`、owner 配置容量和 expected grants 均不变。 +- `selection_debt_bytes` 覆盖已离开 hot window、仍由 source-delete/retry 状态机拥有的字节。 + pressure actor 只选择 `free-slot deficit - selection debt`，避免重复超选。 +- `BatchEvictOwnerSourceReq/Resp` 使用 Msg ID `3069/3070`。RPC 不可达、activity Busy、 + partial overlap 或 actor 暂不可用都按同一 exact identity 有界重试，不重新 append； + physical Free、stale/obsolete 或安全恢复 hot 是 selected debt 的唯一终态。 +- active size class 为 `0` 时是尚未分配的正常 idle，pressure actor 静默等待；当前只支持 + `1` 个 active class；`>1` 时停止选择并限频报错，不能用全局 Moka 在多 class 中猜 victim。 + +### 安全回收时序 + +```mermaid +sequenceDiagram + participant H as owner hot Moka + participant O as GPU owner dispatcher + participant M as Fluxon master + participant R as existing local reader + + H->>O: Size event(key, put_id, Weak, weight) + O->>O: point-pin current source / expand exact cohort + O->>M: BatchEvictOwnerSource(exact cohort) + M->>M: validate caller generation, closure and backing + M->>M: install all master activity fences + M->>O: cohort Prepare + alt reader already holds source + O-->>M: Busy; restore visible index + R-->>O: holder released + M->>M: bounded retry of the same identity + end + O-->>M: all Prepared + M->>O: cohort Commit + O->>O: Prepared -> Releasing under key shard + O->>O: unlock key; release route+holder in one pool update + O->>O: Releasing -> Committed + M->>M: remove exact owner routes; last route may disappear + M->>O: cohort Finalize +``` + +Master 与 owner 两侧分别承担下面的正确性约束： + +| 约束 | 当前实现 | +| --- | --- | +| 并发线性化 | Master 为每个 key 统计 put/get/replica activity。只有三个计数都为零且没有其它 reclaim 时才能原子安装 reclaim fence；安装后新的 put/get/replica 不能取得 activity lease。已经取得 source holder 的读或 replica transfer 必须先排空，但 transfer 是否成功不改变容量事务义务。 | +| 版本隔离 | 每个候选携带 `put_id` 和 reclaim `epoch`。Master 在安装 fence 前、Prepare 后以及删除 route 前重复检查当前 route 仍是同一版本。 | +| backing/slot 防 ABA | committed slot 必须同时匹配 owner、`grant_id`、`slot_index` 和 `slot_size`；旧驱逐请求不能释放新版本或已经复用的 slot。 | +| 现有读者排空 | Owner Prepare 只持该 key 所属的短 sharded key-control fence，先从 local index 隐藏 backing，再要求唯一 holder。若已有 get/plan/transfer pin，则恢复索引并返回 Busy。 | +| 写入冲突隔离 | Owner Prepare 遇到 `local_puts`、`precommit_local_visible_info` 或 `external_pending_puts` 时返回 Busy，不会回收正在写入或尚未发布完成的 value。 | +| 精确物理释放 | Owner Commit 在 key-shard 内只切换 `Prepared -> Releasing`，随即释放 key 锁；之后在一次 slot-pool 临界区同时释放 committed route ref 与 resident holder，再回到 key shard 标记 `Committed`。这避免 key-shard mutex 嵌套 slot-pool mutex，也不暴露两次 pool update 间的中间状态。 | +| route 更新顺序 | Owner Commit 释放物理 backing 后，Master 只删除相同 `put_id` 且 backing identity 相同的 replica；这个短窗口仍受 master reclaim fence 保护，对新请求不可访问。Master 更新 route 后才 Finalize owner fence。 | +| cohort 事务 | owner exact source deletion 和 master allocation capacity reclaim 都以一个 cohort request 进入 lossless queue。全组 master fence 与 Prepare 要么全部取得、要么全部撤销；一旦任一 Commit 已成功，只允许 roll-forward 完成其余成员。 | +| 幂等与失败关闭 | Prepare/Commit/Abort/Finalize 都按 `item + epoch` 判定已应用状态。RPC 不确定时重放同一 identity；已经 Commit 的 backing 不可回滚。未收敛时保留 fence，优先阻塞该 key 而不是暴露已释放内存。 | +| partial/replay 判定 | 全部 source 已不存在是 `Completed`；部分不存在只有在完整 exact identity 已在 reclaim registry 中时才是 `AlreadyInProgress`。无对应在途事务返回 `Stale`，部分 overlap 返回 `RetryableBusy`，不能把残缺 cohort 当作新事务接受。 | +| 失败候选收敛 | 暂时 Busy 的当前版本保持 pending 并退避重试，不立即插回 Moka；stale route 或版本变化直接结束 pending。只有队列关闭等无法继续推进的路径，才在 route 仍匹配 `put_id` 且 owner replica live 时重新插入，避免旧条目覆盖新版本。 | + +从可观察状态看，回收事务允许 Owner Commit 与 Master route 删除之间短暂存在“owner slot 已回收到 allocator、可以复用，但旧 route 对象尚在”的内部状态，但对应 key 的 master reclaim fence 在整个窗口内拒绝新操作；安装 fence 之前已经取得 lease 的 get/put/replica 又必须先结束。因此调用方不会通过该 route 访问已经释放或复用的 backing。 + +### 副本保留语义 + +内存/并发正确性与“驱逐后是否仍有另一份可恢复数据”是两个不同契约： + +- **环 A（owner-indexed source）**：hot removal/`evict_some` 只是 owner 的 victim selection。 + owner 把完整 exact cohort 交给 master 做 source route fence/delete；不要求 CPU 副本。 + 有其它 replica 时只删除该 source，没有其它 replica 时允许最后 route 消失。RPC/Busy 暂时 + 失败时 source 保持可读并按同一 identity 重试；remote write 失败不进入容量失败补偿。 +- **环 B（master-only/unindexed Allocation）**：master 有界 Moka 的原生 Size eviction 可以删除最后 + 一份 cache route，后续正常 miss/recompute；这是 cache 容量语义，与环 A 的结果一致， + 但 victim authority 与 backing 释放位置不同。该分类只取决于 backing/index/lease + 语义，不取决于 CPU/GPU 或 active/remote-only 节点角色；当前默认置换策略为 + Moka TinyLFU，正确性不依赖具体 victim 顺序。 +- 旧 `OwnerReclaimReason::Reserve` route 全表扫描、`reclaim_before_grow` 和 master 侧手动 + `evict_some` 均已删除。NoSpace 只返回明确可重试结果；owner pressure actor 等待真实 Free + slot，不通过增大 grant 数绕过 configured capacity。 +- `atomic_group_lens`/Put group 是 source cohort 描述；master 在 source deletion 当下验证 + 请求集合与当前 atomic/TP closure 完全一致。它不从单 key补全 cohort，也不要求共同 CPU + owner；任一 member stale 时整组不 Commit。 + +### 当前验证证据与覆盖边界 + +2026-07-17 对当前 Fluxon 工作树完成 direct source-delete 本机回归。使用 +`CARGO_TARGET_DIR=/tmp/fluxon_target_direct`、`-j 4` 和 closed-sdk runtime，最新代码 +`cargo test -p fluxon_kv --lib --no-run` 通过；exact source cohort 定向测试 `7/7`，owner +slot/reclaim 定向测试 `19/19`，`fluxon_kv --lib` 全量为 +`171 passed, 0 failed`（193.19s）。覆盖有/无 CPU replica、最后 route 删除、stale backing、 +partial replay-only、cohort fence all-or-nothing、retry exactly-once、`Releasing` 和 +route+holder 单次 pool update。全量测试还修复了 `0 active class` 被误判后每 10ms 报错的 +idle actor 空转。本机回归本身不替代 release 三机部署、真实 RDMA 与 E44 自然冷跑，后者 +按下一段独立验收。 + +上述独立验收已在 2026-07-17 的 E44 r9 完成。三机容量保持 +GPU0/GPU1/CPU=`128/128/256 GiB`，两 GPU 各固定 `232` grants、`26,216` slots； +S96×T24、concurrency 24、session-stream、无预热自然冷跑为 `2304/2304` 成功，QPS +`5.609336`，L1/L2/L3/总命中率 `5.4060%/0%/54.6878%/60.0939%`。CPU owner 被选为 +唯一 proactive replica target（`77,384` 个），并实际为两个 GPU requester 提供 +`183,923` 个 key、`808.26 GiB` Get source 数据。 + +容量闭环终态为：GPU0/GPU1 direct source commit `130,540/145,888`，dispatch failure +均为 0；Free slots `466/488`，Prepared/Pending 均为 0；active flights、retry entries、 +selected 和 selection debt 全为 0；grant 数始终等于 expected `232`。业务错误、refill +timeout、P2P、panic、OOM 和 scheduler exception 均为 0。日志中的 recoverable +`not_ready` 只对应最后一份 route 已按容量策略删除后的正常 miss/recompute；所有真正提交的 +prefetch 均以非零 prefix 成功，没有 TP prepare reject 或 transport error。该结果证明环 A +的 direct source-delete 正确性和容量收敛，不证明 selective CPU replica 已达到 Mooncake 的 +命中率；后者仍是独立性能策略，不能反向成为 Free 的前置条件。 + +三机 SGLang 正式压力还覆盖了 expected-capacity 与常规驱逐的组合。E16a4 使用两台 TP=2 +GPU owner 和一台 CPU-only owner，标称容量保持 128/128/256 GiB；两 GPU 各预热 207 grants， +正式请求期间各按 pending demand 增到 208 grants。1152/1152 请求成功，QPS 4.845，两侧 +Moka 稳定在 23,296 entries / 109,924,319,232 weighted bytes，owner reserve 在继续写入时 +稳定维持约 128--170 free slots；refill timeout、NoSpace、OOM、scheduler fatal 均为 0。 +master 的 replica target 只有 CPU owner，共 19,395 个，GPU 间副本为 0。这证明预热和 +shrink 下限不改变 owner 容量，也不会阻止达到水位后的常规 `evict some` 持续回收。 + +同拓扑 E16b 把水位和 expected payload 从 0.8/102.4 GiB 提到 0.95/121.6 GiB 后，QPS +从 4.845 提高到 5.967，总命中从 73.67% 提高到 90.28%；Mooncake 为 6.919 QPS、 +93.05% 命中。容量对齐后只剩 2.77pp 命中差，但仍有 13.76% QPS 差距。两 GPU 的 +`init_load_back` 平均总时长分别为 80.07/52.49 ms，其中 `restore_sync` 平均 +68.46/42.33 ms、radix eviction 平均 9.66/8.31 ms，而 kernel launch 记账约 0.05 ms。 +因此 expected capacity、exact-fit 和 group 完整性都不能替代数据面优化：需要继续减少 +同步点，分析 GPU0/GPU1 restore sync 不对称，并评估 `page_first_direct` 组织方式。E16b +仍有 26/24 次 TP transferable mismatch；common-prefix retry 是独立的命中修复，不应被当成 +剩余全部性能差距的解释。 + +E16c 的单变量结果进一步确认了这个边界。启用 TP-invariant Put atomic-group admission 和 +TP common-prefix retry 后，总命中从 90.28% 提高到 92.76%，与 Mooncake 的 93.05% 只差 +0.29pp；TP mismatch skip 从 26/24 降到 0。但 QPS 只从 5.967 提高到 6.136，仍比 +Mooncake 低 11.32%。E16c 的 GPU0/GPU1 load-back 平均总时长仍为 80.55/53.84 ms， +`restore_sync` 仍占 69.29/43.51 ms。结论是：atomic group 是 CPU selective replica 可消费性 +与 TP 正确收敛所必需的协议信息，但在命中率已经对齐后，主要性能工作必须转向 restore +kernel 的同步边界、GPU radix eviction 和 page/layer 搬运顺序。 + +后续 E16f 的逐层 restore kernel 把 Fluxon QPS 提到 6.635；E16s 的 deferred +`release_views` 和 E16t 的整批 local-visible snapshot 继续把 QPS 提到 6.808。E16u 再为 +全 owner-local get 增加 `InlineLocal` 分支，两个独立冷栈正式轮分别为 6.754/6.760 QPS， +1152/1152 成功且各自实际备份 3,869,056 tokens。`get_transfer` mean 已从 E16t 的 +4.192/1.957 ms 降到 0.558/0.585 ms，但冷 QPS 没有继续提高；一次复用旧 owner keys 的热轮 +达到 6.894 QPS，却只实际备份 246,976 tokens，不能作为冷 A/B。 + +E16u 冷复测把剩余 radix eviction 进一步定位到同步 write-back。GPU0/GPU1 的 +`init_load_back` mean 为 10.621/8.399 ms，其中 eviction 为 8.859/6.688 ms；阻塞 +`_wait_for_fluxon_hostless_backup` 与 `writing_check(write_back=True)` 的累计 duration 约占 +这些 eviction 累计时间的 95.7%/90.5%。parent-chain backup mean 为 6.508/4.435 ms，主要 +由 stream sync 3.735/2.322 ms 和 `local_fast_put_start` 1.988/1.366 ms 构成。因此当前性能 +边界已经从 get transfer RPC 转到 write-back 时机与批次组织。后续优化需要在不改变 owner +128/128/256 GiB 容量、不降低连续完整前缀命中的前提下，评估有界 proactive write-back 或 +多 radix node 共用 Put plan/一次 stream sync；单纯继续减少 get 控制面 RPC 的收益不足以 +代表冷端到端收益。 + +### Write-through parent 依赖与驱逐时补备份 + +hostless write-through 的 `storage_backed` 也必须满足连续前缀不变量：一个 child 只有在 +parent 已经 remote-backed 后才能成为独立可恢复节点。proactive `write_backup` 遇到 parent +尚在途时会先提交 parent，但不会自动把 child 排队到 parent ACK 之后重试。因此 child 可能 +既没有 `storage_backed`，也没有 `storage_pending`；后续 device eviction 若直接删除该 radix +node，会丢失仍可通过一次补写保留的前缀。 + +驱逐时的收敛顺序应为： + +1. 已经 `storage_backed`：直接走 `_evict_to_fluxon_storage`。 +2. 当前 node 有 `ongoing_fluxon_hostless_backup` 或 pending ACK：只等待已有操作；成功后走 + Fluxon tombstone。 +3. 没有可恢复副本也没有 pending：调用 `write_backup(node, write_back=True)`；该调用递归 + 补齐缺失 parent，随后 `writing_check(write_back=True)` 阻塞到 ACK。 +4. ACK 后再次检查 `_fluxon_hostless_node_storage_ready`。只有确认完整 KV（及需要时 Mamba） + 均 remote-backed 才释放 device backing；提交失败、ACK 失败或仍不可恢复时保留原来的 + fail-closed 删除语义。 + +E16w 只实现第 2 步，但两台 GPU 的新分支触发均为 0，排除了“主要是 ACK 尚在途”的假设。 +E16x 补齐第 3--4 步，正式轮中 GPU0/GPU1 分别产生 118/114 条 TP-rank 同步补备份提交， +失败为 0；其中直接 leaf completion 各 88 条，parent-chain completion 为 30/26 条。 +总命中由 E16w 的 92.38% 提高到 92.73%，增加 99,520 tokens；QPS 仅从 6.8081 降到 +6.7915。每次 load-back 的 radix eviction mean 只增加 0.169/0.127 ms。这说明按需补齐 +parent chain 可以作为正确性和命中保底，不需要把所有 write-through 操作重新变成同步写。 + +该结果也进一步限定 atomic group 的作用边界。Mooncake 没有 Fluxon Put atomic-group +admission，仍达到 93.05% 命中和 6.9192 QPS；E16x 在命中只低 0.323pp 时仍低 1.85% QPS。 +因此 Put atomic group 是 TP 一致准入、CPU selective replica 完整组和“有完整远端副本才 +优先驱逐”的正确性信息，但不是剩余端到端差距的充分解释。Mooncake 每次 restore 平均跨 +3.20/3.26 个 radix node，E16x 为 10.15/10.08 个；Mooncake 的命中主要经过 32 GiB 本地 +HiCache L2，Fluxon hostless 的命中主要表现为 external L3。后续数据面比较必须记录每个 +request 的实际 restore bytes、来源 locality、CUDA event 等待和合并后的 batch 数，不能仅 +比较命中率或 atomic-group 数量。 + +E16x 的分节点指标显示剩余差距集中在 GPU0：Fluxon TTFT/E2E mean 为 1.036/1.629s， +Mooncake 为 0.865/1.522s；GPU1 的 Fluxon 1.719/2.499s 反而略快于 Mooncake 的 +1.728/2.524s。GPU0 同时承载 Fluxon master、owner、router 和 workload，因此还需把 +restore 数据面耗时与 master/owner CPU 竞争分开测量；在证据出来前，不应通过扩大 owner +容量或改变副本语义掩盖该不对称。 + +E16y 对该不对称做了 NUMA 定向实验。两台 GPU 都属于 NUMA node1，但 E16x 的 GPU0/GPU1 +owner backing 分别有 128.14/97.40 GiB 落在远端 NUMA0。把 owner 与 SGLang 绑定到 +NUMA1、把 master/router/workload 绑定到 NUMA0 后，两台 owner 都有约 128.25 GiB 落在 +NUMA1；备份量、命中率和 fallback 数保持相同，QPS 从 6.7915 提高到 6.8520，与 Mooncake +只差 0.97%。因此大块 registered host backing 的 first-touch NUMA 必须成为部署契约，而 +不能依赖进程启动瞬间的随机 CPU。 + +但 E16y 不是最终绑核方案：GPU1 layerwise completion 降低 5.20%，GPU0 却没有改善，且 +两端 `get_start` p50/p90 因 owner 与 SGLang 共用受限 CPU 集而上升。正确的后续拆分是只让 +owner 在 GPU-local NUMA 完成大块 backing first-touch，SGLang/master/router/workload 保持 +正常调度；若需要长期绑核，再为 owner data plane 与 SGLang scheduler 分配不重叠的 core +集合。不能把“NUMA-local memory”与“所有线程绑在同一个 NUMA CPU 集”视为同一要求。 + +E16z 完成该拆分后，在备份量和 fallback 数不变的条件下达到 6.9007 QPS，比 E16x 提高 +1.61%，与 Mooncake 只差 0.27%；完整 1152 请求 wall 仅差 0.45s。由此 owner backing 的 +GPU-local first-touch 应固化为部署步骤，而 SGLang 不应与 owner 共用同一受限 affinity。 +steady-state owner 也不一定要永久绑核：E16z 的 `get_start` p50 仍比 unrestricted E16x 高 +约 0.5--0.6 ms。更精确的生命周期是 owner 在 segment/backing 与 expected-capacity reserve +建立期间绑定 GPU-local NUMA，完成后把全部已有线程 affinity 放宽；first-touch 页面会保持 +原 NUMA placement，控制/RPC 线程则可以重新使用其它 CPU 和 NIC-local core。 + +E16aa 对这一两阶段生命周期做了正式 A/B。放宽后 `get_start` p50 确实从 E16z 的 +2.55/2.60 ms 降到 2.15/2.29 ms，页面仍驻留 NUMA1；但 QPS 为 6.8902，比 E16z 低 +0.15%，没有形成端到端正收益。考虑实现和运维复杂度，当前建议采用 E16z 的简单规则： +GPU owner 整个生命周期绑定 GPU-local NUMA CPU 集，SGLang/master/router/workload 不绑核。 +E16z 6.9007 QPS 与 Mooncake 6.9192 只差 0.27%；若后续追求稳定的亚百分比收益，应先 +增加独立冷轮重复次数和 per-request CUDA/CPU tracing，而不是继续叠加 affinity 状态切换。 + +E16ab 随后按 E16z 配置做了独立全冷栈复测。两个 GPU owner 启动后分别有 +128.151/128.143 GiB 页面位于 NUMA1，容量、副本开关、请求序列和 active source 均未变化。 +本轮 1152/1152 成功、严格 576/576，得到 6.8812 QPS、92.72% 总命中；它比 E16z 低 +0.28%，比 Mooncake 单轮低 0.55%。两次 E16z 形态冷轮均值为 6.8909 QPS，比 Mooncake +低 0.41%。因此原先 0.27% 差值确实落在单轮波动量级，但不能据此宣称 Fluxon 已超过 +Mooncake。两轮的 `get_start`、`init_load_back`、radix eviction、restored-node 数和 +layerwise completion 高度复现，未出现新的系统瓶颈或错误。 + +“Fluxon 有额外优化所以理应直接超过 Mooncake”并不是有效的性能推导。Put atomic group 和 +缺失 parent 补写首先是正确性/可恢复性契约；InlineLocal、batch local-visible、deferred +`release_views` 和 layerwise kernel 则用于抵消 external owner/hostless 路径相对进程内 L2 +的 RPC、holder 和恢复组织成本。Mooncake 对齐轮有 70.55% tokens 命中进程内本地 L2，并把 +512 GiB 五个 segment 全部作为 primary capacity，285.3 GB 数据没有发生容量驱逐；E16z/E16ab +约 62% tokens 经过 external owner，只使用两个 128 GiB GPU owner 保存 primary，256 GiB +CPU owner 为空。因此现有优化是在不同路径上把端到端拉平，并不是在 Mooncake 相同基线上的 +额外加速项。 + +能够形成 Fluxon 独有容量收益的下一闭环仍是“完整 CPU 副本 -> GPU 优先驱逐该完整 group -> +CPU 成为可恢复剩余副本”。Put `atomic_group_lens` 目前只保证准入不主动切组；replica actor +仍逐 key append，owner/Moka 也没有完整 group 的 remote-complete 状态可用于候选排序。只有 +补齐远端整组完成记账、版本失效和 recoverable-first eviction 后，CPU 256 GiB 才会从异步 +复制开销转化为 GPU eviction 后的额外命中。这比继续压缩已经只有约 2--3 ms 的 `get_start` +或 0.7--0.8 ms 的 radix eviction 更可能带来稳定超过 Mooncake 的收益。 + +E16ad 进一步纠正了“配置容量”和“实际参与容量”的表述。E16z/E16ab 虽然启动了 +128/128/256 GiB 三台 owner，但 `replica_task.enabled=false`；SGLang +`external_batch_put_start` 又直接使用 requester-local reserve，因此只有两个 GPU owner 保存 +primary，CPU owner entries 为 0。CPU 的 256 GiB 在这两轮只是注册容量，不能计作实际数据 +容量。Mooncake 对齐轮则确实把五个 segment 都作为 primary 使用，两者的容量利用语义并不 +等价。 + +E16ad 在相同物理容量和 NUMA 配置上启用 CPU-only `kv_score_only` replica，master 严格限制 +remote target role 为 `remote_cache`，GPU 间仍无副本。由于当前 PyO3 契约令 +`write_through || !make_replica_task` 都跳过 replica，本轮 HiCache 使用 write-back。最终 +GPU0/GPU1/CPU 分别保存 121.585/121.584/142.938 GiB，CPU 有 18,296 entries，master +记录 18,336 个 CPU target 和 0 个 GPU target，证明三台 owner 均实际进入数据面。 + +但该轮只有 6.8124 QPS、92.69% 总命中；相对 no-replica 冷复测 E16ab，QPS 低 1.00%， +命中低 0.03pp，同时产生约 0.517 Gbps 聚合 RDMA TX。CPU 没占满 256 GiB 是 kv-score 跳过 +低分 group 的预期结果，更关键的是这 142.938 GiB 仍为重复副本。当前 GPU capacity eviction +没有利用“完整 CPU group 已完成”的状态选择候选，所以 CPU 写入没有转化为更多可消费命中。 + +这里必须区分两个后续实现方向： + +- 若维持“CPU 只存副本”，需要补齐完整组 remote append 完成、版本失效和 + recoverable-first GPU eviction；GPU 本地 copy 被驱逐后，CPU copy 才转化为有效扩容。 +- 若要求像 Mooncake 一样把 128+128+256 GiB 全部作为独立 primary pool，需要实现 + requester-local first、local 压力后 spill 到 CPU 的 remote-primary batch put。当前 + external batch commit 显式拒绝 remote primary，修改 `prefer_local_placement` 不能得到 + 该语义。 + +因此“三 owner 都有占用”已经由 E16ad 验证；“512 GiB 都是独立有效容量”尚未成立，不能把 +副本字节与 primary 字节直接相加。 + +### 低分 KV 提前 CPU write-back 策略 + +SGLang Fluxon HiCache 新增显式准入策略 `kv_score_low_only`。它保留 GPU requester-local +owner 上的 primary，同时为低分 KV atomic group 提前创建 replica task；三机 CPU-only +拓扑由 master 的 `remote_cache` target role 约束副本只落到 CPU owner。策略本身只决定 +replica task 准入，不改变 target role、owner 容量或 GPU 间副本规则。 + +配置示例： + +```json +{ + "replica_task": { + "enabled": true, + "admission": { + "policy": "kv_score_low_only", + "score_threshold": 0.55, + "min_replica_pages": 0, + "max_replica_pages_per_batch": 512 + } + } +} +``` + +该策略使用包含边界的 `score <= score_threshold`。Put 仍按 `atomic_group_lens` 作整体决策， +同一 group 的 mask 全 true 或全 false。group priority 使用 `1 - mean(group scores)`；因此触发 +`min_replica_pages` 补选或 `max_replica_pages_per_batch` 裁剪时，优先选择最低分 group，且不会 +为了填满页数预算切开 group。原 `kv_score_only` 保持 `score >= score_threshold` 和高分优先 +语义不变。 + +物理 storage key 继续区分 TP rank，准入 identity 则保留 TP size 并去掉 TP rank；同一逻辑 +group 在所有 TP ranks 上得到一致的低分准入结果。该实现属于 proactive CPU write-back +admission：CPU append 成功前 GPU primary 不会提前释放，成功后当前 eviction 也尚未按 +remote-complete group 排序。完整的“低分提前写回 -> GPU `evict some` 优先回收已完整备份组 +-> CPU 副本承担后续恢复”闭环仍需 remote-complete 记账、版本失效和 recoverable-first +eviction。 + +E16ae 已用与 E16ad 相同的 128/128/256 GiB、双 TP=2、owner-only NUMA1 和 1152 请求配置 +完成低分策略正式压力。1152/1152 成功，QPS 为 6.8073，总命中 92.72%；E16ad 高分策略为 +6.8124 QPS、92.69% 命中，差异分别为 -0.075% 和 +0.03pp。低分策略把 CPU retained usage +从 142.938 GiB 提高到 243.188 GiB，历史 CPU targets 从 18,336 提高到 42,214，聚合 RDMA +TX 从 0.517 提高到 1.189 Gbps，但没有得到可测的吞吐或命中收益。两 GPU 仍各保存约 +121.570 GiB primary，GPU replica target 为 0；CPU cache 已达到 95% effective capacity。 +这进一步证明 admission 选择本身不能替代 recoverable-first eviction：额外 CPU 副本在当前 +消费路径上仍主要体现为重复容量和写流量。 + +Mooncake 对齐轮也必须按相同抽象层解释。SGLang 启动参数是 +`--hicache-write-policy write_back`，安装版 Mooncake `ReplicateConfig()` 默认 +`replica_num=1`。Mooncake master 管理 120,928 keys，每个对象 2,359,296 bytes；乘积 +285,304,946,688 bytes 与 master allocated bytes、五个 segment allocated bytes 之和完全 +相等。这证明 global store 中每个对象只有一份，五个 segment 共同承担单份对象的 placement。 +它不是“GPU segment 保留 primary，再 write-through 一份 CPU replica”。 + +因此更严格的路径映射是：Mooncake 的 SGLang device/L2/L3 采用 write-back demotion，global +store 的 512 GiB 提供单份对象容量；Fluxon E16ad 则在 GPU owner primary 仍存在时异步写 +CPU replica。二者都可能在层间迁移瞬间短暂同时持有数据，但稳态容量语义不同。若目标是复刻 +Mooncake 同时发挥 Fluxon 的统一管理优势，Fluxon 更合适的实现是 local owner 到 CPU owner +的 single-copy demotion/spill：CPU commit 成功后删除 local copy，而不是长期保留两份。 +recoverable-first replica eviction 可以作为现有协议的过渡实现，但验收最终应统计单份逻辑 +字节、物理总字节和迁移后 source 删除，不能只统计 CPU entries 非零。 + +### Master 两级 Moka 提前写回 + +Master 新增可选的包含式两级 Moka，用 master Moka 的实际容量淘汰序列代替 Put 时的 score +预测。这里的 T1/T2 是 +Fluxon master 的 owner route 元数据层级，不是 SGLang HiCache 的 GPU L1 / host L2： + +| 层级 | 容量 | 淘汰动作 | 是否释放 owner 数据 | +| --- | --- | --- | --- | +| T1 hot tier | `replica_writeback_tier1_capacity_ratio × owner segment` | 批量通知源 owner 发起 remote-only replica task | 否 | +| T2 resident tier | `replica_cache_capacity_ratio × owner segment` | 进入现有 fenced owner reclaim | 是 | + +T1 是 T2 的热数据子集。新 route 同时进入 T2 和 T1，因此启用 T1 不会从 T2 切走容量，也不会 +改变 owner 的 `max_capacity`。经 master get 路径观察到的 T2 命中会把对应 key 重新提升到 +T1。T1 因容量发生 `Size` 淘汰时，master 按源 owner 合并请求，再通过内部 +`BatchEnqueueReplicaTaskReq` 通知 owner。 +Owner 继续复用已有的 `put_append_start -> RDMA transfer -> put_append_done` replica actor；目标仍 +由 `replica_task_placement` 选择。启用两级写回时，配置校验要求 +`restrict_to_remote_only_node_roles=true`，防止写回落到另一个 GPU owner。 + +```mermaid +sequenceDiagram + participant M as Fluxon master T1/T2 + participant G as GPU owner + participant C as remote CPU owner + + G->>M: publish local route + M->>M: insert T2 resident + T1 hot + M->>M: T1 Size eviction + M->>G: BatchEnqueueReplicaTaskReq(key, put_id) + G->>M: put_append_start + M-->>G: reserve CPU target + G->>C: RDMA transfer + G->>M: put_append_done + M->>M: append CPU replica route + Note over M,G: GPU owner source remains in T2 + M->>G: later T2 fenced reclaim +``` + +Master 配置示例： + +```yaml +replica_task_placement: + remote_only_node_roles: ["remote_cache"] + restrict_to_remote_only_node_roles: true +replica_cache_capacity_ratio: 0.95 +replica_writeback_tier1_capacity_ratio: 0.75 +``` + +`replica_writeback_tier1_capacity_ratio` 必须大于 0，且严格小于 +`replica_cache_capacity_ratio`。不配置该字段时保留原单层 resident Moka 行为。以 +128 GiB owner、T2 `0.95`、T1 `0.75` 为例，owner resident 上限仍约为 121.6 GiB；约 25.6 +GiB 的 T1/T2 容量差提供 replica 完成窗口，而不是把 owner 容量降到 96 GiB。 +Master 的周期日志同时报告 `writeback_tier1_triggered`、 +`writeback_tier1_owner_accepted` 和 `writeback_tier1_failed`。其中 accepted 只表示源 owner +已把任务放入 replica actor，CPU route 是否完成仍以 `replica_task_target_counts` 和 +`put_append_done appended=true` 为准。 + +当前实现仍是过渡策略，边界如下： + +- T1 淘汰只启动逐 key remote replica，尚未记录 SGLang radix atomic group，也没有整组完成 + 状态。 +- Owner `InlineLocal` fast path 当前不经过 master get，因此这类本地命中不会刷新 T1；当前顺序 + 是 master 可观察的 insert/get 淘汰序列，不是所有 SGLang 访问的全局 LRU。 +- CPU replica 完成后不会立即删除 GPU owner source；source 只有在后续 T2 reclaim 时释放。 +- 该策略能让写回时机跟随 master 实际热度淘汰，并为异步复制提供有界提前量；它尚未实现 + Mooncake 式 single-copy demotion。 +- 运行该策略时，SGLang Put 侧应关闭 `kv_score_only` / `kv_score_low_only` eager admission， + 避免同一 key 同时由 Put admission 和 T1 淘汰重复请求副本。 + +#### E16ag 三机验收结果 + +首轮 E16af 虽然成功解析 T1/T2 比例，但把 source eligibility 错误绑定到 +`prefill/decode` active-client role。实际 GPU storage owner 的 role 是 `sglang_owner`，因此 +两台 GPU 都没有构造 T1。E16ag 把资格修正为“已注册且非 `remote_cache` 的 owner”，从而 +允许 GPU storage owner 作为写回源，同时继续排除纯 CPU owner。 + +E16ag 在固定 GPU0/GPU1/CPU = 128/128/256 GiB、双 TP=2、96 sessions × 12 turns、 +concurrency 16 的正式三机压力中得到 1152/1152 成功、6.8151 QPS、92.74% 总命中。两台 +GPU 的 resident weighted bytes 均为 121.570 GiB，T2 上限仍为 121.600 GiB；T1 均为 +95.977 GiB，上限为 96 GiB。因此包含式 T1 没有改变 owner 配置容量，也没有把 T2 降到 +T1 容量。 + +CPU 最终持有 13,061 entries / 102.039 GiB，历史完成 13,090 个 replica targets，目标全部是 +`remote_cache` CPU，GPU target 为 0。流水线计数满足： + +```text +T1 triggered 16525 + = pre-dispatch stale 226 + + owner source-missing/version-mismatch 1209 + + owner accepted 15090 + +owner accepted 15090 + = target KeyBeingWritten 2000 + + CPU replica completed 13090 +``` + +这说明策略已经生效，版本检查也正确阻止了过期 source 被复制；同时暴露了两个明确窗口。 +第一，T1 淘汰事件入队后，source 可能先被 T2 回收，现有 holder 获取发生得太晚。第二，owner +接单只表示任务进入 replica actor，目标侧相同 key 仍可能处于写入中。下一版应在 T1 eviction +时为相同 `put_id` 安装短期 source pin，复制完成/失败后释放，并对 `KeyBeingWritten` 做有界、 +版本安全的重试。观测上应把 pre-dispatch stale、source missing 和 busy key 分开计数，不能只 +用 `writeback_tier1_failed` 聚合。 + +性能上，E16ag 相对无 T1 的 E16af 提高 0.75%，但相对 no-replica E16ab 低 0.96%，相对 +Mooncake 低 1.50%。CPU 副本物理存在不等于形成额外可消费命中：当前 source 在 T2 前仍保留， +CPU 字节主要是暂时重复容量；逐 key 成功也不能证明完整 radix group 可恢复。因此 source pin / +busy-key retry 解决复制完成率之后，仍需记录 atomic group remote-complete，并让 T2 +recoverable-first eviction 优先选择已有完整 CPU 副本的 group，才能把 CPU 容量稳定转化为 +命中收益。 + +### Owner-local hot Moka 提前写回 + +当前实现把提前写回的热度观察点下沉到 storage owner。每个 GPU owner 维护一个逐 key 的 +owner-local hot Moka；本地 committed value 进入该 tier，`InlineLocal` 等 owner 本地命中会刷新 +热度。hot tier 发生 `RemovalCause::Size` 时，callback 只把 key/put_id/Weak/weight 交给 +lossless metadata channel；离开 Moka housekeeper 后，异步 dispatcher 才逐 key point-check、 +pin 当前 backing、展开 cohort，并交给 replica actor 写入 remote CPU owner。Master 仍负责 +remote route 成功后的 cohort demotion/reclaim 仲裁。 + +```mermaid +sequenceDiagram + participant S as SGLang / external client + participant G as GPU storage owner + participant H as owner-local hot Moka + participant M as Fluxon master T2 + participant C as remote CPU owner + + S->>G: local commit(key, put_id) + G->>H: insert weak entry(key, put_id, slot_size) + S->>G: InlineLocal hit + G->>H: touch current entry + H->>G: Size eviction(key, put_id) + G->>G: async dispatcher: DashMap point-check + Arc pin + G->>M: put_append_start(key, put_id) + M-->>G: reserve remote-only target + G->>C: RDMA transfer + G->>M: put_append_done + M->>M: append CPU replica route + alt owner-hot demotion intent 且完整 TP cohort 可从同一 CPU owner 恢复 + M->>M: Size-evict exact source cohort + M->>G: fenced Prepare / Commit / Finalize + G-->>M: release current local backing + else append/版本/整组检查失败或 CPU NoSpace + M-->>G: retain GPU source + end + G->>G: release pin and inflight dedup guard +``` + +Owner 配置只增加一个 canonical 参数： + +```yaml +replica_writeback_hot_capacity_ratio: 0.75 +``` + +该比例必须有限且位于 `(0, 1)`，并且只能出现在非零容量 owner 配置中。hot Moka 的 +`max_capacity` 等于该比例乘 owner DRAM；128 GiB owner 配置 `0.75` 时，逻辑 hot 水位为 +96 GiB。这个计算不修改 `contribute_to_cluster_pool_size.dram`、local-reserve grant 数量、owner +segment 或 master T2 的 `replica_cache_capacity_ratio`。Moka value 只保存 +`Weak`；所有 hot entries 因而不会整体 pin 住 owner backing。计费优先使用 +local-reserve 的实际 `slot_size`，普通 allocation 使用 payload length。 + +CPU-only 目标仍由 master 的 placement 契约提供： + +```yaml +replica_task_placement: + remote_only_node_roles: ["remote_cache"] + restrict_to_remote_only_node_roles: true +replica_cache_capacity_ratio: 0.95 +# replica_writeback_tier1_capacity_ratio 不配置 +``` + +第一版实验必须关闭 master T1 和 SGLang Put 侧的 eager replica admission，避免同一个 key 有 +多个独立提前写回触发器。owner hot tier 只决定写回时机；如果 master 没有限定 +`remote_only_node_roles`，目标仍可能由普通 placement 选到其它 owner。 + +并发与失败处理遵守下面的不变量： + +| 场景 | 当前处理 | +| --- | --- | +| hot eviction 与 T2 reclaim 竞争 | owner-hot pin 不取全局 `owner_key_control`。dispatcher 在 DashMap shard read guard 内校验 current identity 并 clone `Arc`；pin 先发生时 reclaim remove 后的唯一强引用检查返回 Busy，Prepare 先发生时 dispatcher 看到 index 缺失且 live Weak，按 `RetryableReclaimFence` 退避，绝不从 Prepared 升级 Weak。 | +| 版本替换、delete、显式失效 | eviction listener 只消费 `RemovalCause::Size`。`Replaced` 和 `Explicit` 不创建 replica；reclaim/delete 按 `put_id` 条件失效 hot entry，不能删除随后安装的新版本。 | +| 旧事件与重复事件 | listener 携带 eviction entry 的 `put_id + Weak identity`；dispatcher 以当前 owner index 再校验。`(key, put_id)` guard 覆盖排队、transfer 和 append-done 全周期。 | +| 控制面阻塞 | Moka callback 只向 lossless 轻量 metadata channel 发送事件，不点查 index、不展开 cohort、不 pin、不等待 RPC/RDMA。dispatcher 完成 point-pin 后再进入容量 128 的 replica task queue；RetryableCold 按 owner committed identity 去重并退避。 | +| append start busy key | owner-hot 任务遇到 `KeyBeingWritten` 时最多尝试 3 次，重试间隔为 25/50 ms。其它错误不做这一层重试；最终失败释放 pin，保留 GPU primary，不会把失败任务解释为 CPU replica 已完成。 | +| 可观察性 | owner 每 30 秒记录 hot capacity/weighted bytes、Size eviction、enqueued/completed/already-satisfied/failed/obsolete、dispatch-failed、inflight 以及 stale/reclaim/duplicate skip 计数。CPU 数据面完成仍以 `put_append_done appended=true` 和 master target 计数为准。 | + +E16ah 的第一版实现修复了 master T1 的两个直接缺口：source 在 eviction callback 内完成 pin， +且 owner `InlineLocal` 命中能参与热度更新；当时仍是逐 key proactive replica write-back，CPU +append 成功后不会直接删除 GPU source。后续版本把 Put atomic group 存入 route，并让 owner +在任一 group member 被 hot Moka 淘汰时 pin、排队完整 group。Master 进一步按 canonical +`__` 后缀验证所有 TP rank 的逻辑 group 边界及同一 CPU owner 上的完整副本。 + +E16ap 在该完整性信息之上增加 single-copy demotion。只有 owner-hot 任务的内部 append 请求 +携带 demotion intent；普通 eager replica 不携带。append 成功或同版本 CPU route 已存在时， +Master 再次验证 source 版本、atomic group、TP cohort 和共同 `remote_cache` owner，然后以 +`RemovalCause::Size` 从 resident Moka 选择该 exact cohort，复用既有两阶段 fenced reclaim。 +检查不通过、CPU NoSpace、append 失败或版本变化时保留 GPU source。该路径不新增用户侧副本 +开关，也不修改 owner segment、local-reserve grant 或 `contribute_to_cluster_pool_size.dram`。 + +#### E16ah 三机验收结果 + +E16ah 在固定 GPU0/GPU1/CPU = 128/128/256 GiB、96 sessions × 12 turns、concurrency 16 +的三机正式压力中得到 1152/1152 成功、6.8174 QPS 和 92.96% 总命中。两个 GPU owner 的 +T2 resident 均为 121.570 GiB，上限仍为 121.600 GiB；两个 owner-local hot Moka 上限均为 +96 GiB。CPU 终态持有 19,538 entries / 152.641 GiB。Master T1 容量为 0，SGLang eager +replica 关闭，19,658 个历史 replica target 全部指向 CPU `remote_cache` owner。 + +两个 GPU 的 `pending_eviction_reclaim_bytes` 终态分别为 4.500/4.406 GiB，且 workload 结束 +超过 15 分钟后仍未下降；E16ab 为 0/0，E16ag 为 0.125/0 GiB。它不改变 128 GiB owner +配置或 121.6 GiB resident T2 上限，但代表已从 Moka 淘汰、尚未完成 owner reclaim 的权重， +会消耗 segment headroom。高峰期 owner 拒绝原因为 active holder，终态 hot inflight 已归零， +后续 retry 却没有再到达 owner phase。当前观测还不能区分 master key-activity fence 泄漏和 +reclaim retry 停滞，因此 local-side 物理回收尚未完全闭环。 + +终态流水线满足： + +```text +44626 Size eviction = 43762 enqueued + 864 duplicate +43762 enqueued = 19658 completed + 24103 already-satisfied + 1 failed +19658 completed = 19658 master CPU replica targets +``` + +stale、reclaim、obsolete、dispatch-failed 和 inflight 均为 0。相较 E16ag，dispatch 前 stale +与 source-gone/version-mismatch 从 226/1,209 降为 0，最终 `KeyBeingWritten` 丢失从 2,000 +降为 1；这验证了 owner callback 内立即 pin 和有界 busy retry。唯一失败是一个 key 在三次 +append-start 尝试后仍处于 `KeyBeingWritten`；transfer 和 append-done failure 均为 0。 + +性能上，E16ah 相对 E16ag QPS 只提高 0.034%，相对 no-replica E16ab 低 0.93%，相对 +Mooncake 低 1.47%。总命中相对 E16ag 提高 0.218 pp，与 Mooncake 只差 0.089 pp，但单轮变化 +不能全部归因于 CPU；当前 get 指标还没有按 source owner 拆分。owner-local hot Moka 已修复 +副本触发和完成窗口，仍需 atomic-group remote-complete 与 recoverable-first T2 eviction 才能 +把包含式 CPU 副本稳定转成独立有效容量和吞吐收益。 + +在启用 group-aware 淘汰前，应先给 pending reclaim 增加 master-activity、owner-holder、 +route-changed、retry-queued 和 retry-completed 分原因计数，并保证 source pin 释放后请求最终 +reclaim 或安全回插。否则 logical T2 命中和 CPU replica 完成率提高时，物理 segment headroom +仍可能被未终结的 reclaim 占用。 + +#### E16ai Local-side reclaim 闭环验收 + +E16ai 在不修改 owner 容量的前提下完成了上述 pending 闭环。修复包含三个关键生命周期约束： + +- completed `get_holding` 只是 holder backing 的观测状态，不再作为 master 侧 reclaim 的 + blanket veto；真实读者由 owner Prepare 对 `MemoryInfo` 强引用进行权威检查； +- put/get/replica activity lease 在 done、revoke、TTL 和响应发送失败等终态显式幂等释放， + `Drop` 只作兜底，Moka retired clone 不再延迟 activity 计数； +- `ReuseReplica` get done 只有在相同版本、相同 Allocation identity 下才把 master route 标成 + `owner_local_indexed`，确保后续进入 owner 两阶段 reclaim，而不是错误的 master-only 删除。 + +Capacity eviction 对 Busy 请求最多执行 8 次有界指数退避；达到上限后，只在 route 仍是同一 +`put_id` 时把当前条目安全回插 Moka。queue 关闭路径也先结束 pending weight，再按同样的 +current-version 条件回插。pending weight 使用 checked subtraction，运行时分别暴露 +master-activity、master-holder-observed、owner-holder、owner-other、route-changed、retry +queued/completed/restored 和 reclaim-completed 计数。 + +固定 GPU0/GPU1/CPU = 128/128/256 GiB、GPU T2 ratio 0.95、hot ratio 0.75、96 sessions × +12 turns、concurrency 16 的三机正式轮中，两个 GPU pending 的最大 30 秒快照为 +1.031/1.477 GiB；workload 结束后的首个快照均为 0，并在超过 12 分钟的静默观察中保持 0。 +GPU T2 resident 仍为 121.570/121.570 GiB，上限仍为 121.600/121.600 GiB。E16ah 长期残留的 +4.500/4.406 GiB 以及两个 GPU 各 12 次 local get-target NoSpace 均消失。 + +```text +reclaim counter GPU0 GPU1 total +master activity deferred 10 7 17 +completed holder observed 5939 5707 11646 +owner real-holder deferred 5092 4978 10070 +route changed / terminal removal 3149 2996 6145 +retry queued 5007 4583 9590 +retry completed 965 868 1833 +bounded-retry restored 95 402 497 +reclaim completed 3149 2996 6145 +terminal pending bytes 0 0 0 +``` + +这些计数是 attempt 计数，不构成逐列相加的 unique-key 等式；唯一终态要求是 pending 请求 +最终 route removal/change 或安全回插，且 pending bytes 归零。本轮没有 pending underflow、 +safe reclaim queue full/closed 或 retry queue closed。静默期 master 仍有 71 个 completed +`get_holding` / 335,020,032 bytes；它们是 client 仍持有的真实 backing，不是 activity lease， +也不再阻塞其它无真实 reader 的 owner reclaim。 + +正式结果为 1152/1152 成功、6.7585 QPS、92.9691% 总命中。CPU 保留 20,717 entries / +161.852 GiB，20,773 个历史副本目标全部指向 CPU；owner-hot 对账严格满足： + +```text +46576 Size eviction = 45204 enqueued + 35 reclaim-race skip + 1337 duplicate +45204 enqueued = 20773 completed + 24431 already-satisfied +20773 completed = 20773 master CPU replica targets +``` + +E16ai 相对 E16ah QPS 低 0.865%，总命中只高 0.008 pp；因此 reclaim 闭环解决了正确性和 +物理 headroom，但没有把包含式 CPU 副本转成性能收益。atomic-group remote-complete 与 +recoverable-first T2 eviction（或 single-copy demotion）仍是下一层容量策略，不能回退上述 +activity、holder、版本和 slot identity 安全约束。 + +#### E16ao 手动聚合回收与 E16ap 完整组降级 + +E16ao 修复了 resident Moka 八个 segment 各自执行硬容量淘汰的问题。resident Moka 只作为 +无界 metadata/LRU index，Master 按 owner 的 0.95 聚合有效容量显式调用 `evict_some_if`；存在 +ready CPU tier 时，GPU ordinary fallback 为 0，只允许完整 TP cohort 且已有 CPU 副本的条目 +进入两阶段 reclaim。固定三机压力得到 1152/1152 成功、6.6634 QPS、92.58% 总命中和 1,240 +次 CPU-source Get。两个 GPU 的 pending reclaim 终态为 0，recoverable selected 合计约 +34.6 GB，CPU 却保留约 200.9 GB；测量窗口只从 CPU 读取 5.85 GB。该结果证明 group、TP +cohort、CPU Get 和物理 reclaim 均已生效，同时表明大部分 CPU 字节仍是包含式重复副本。 + +E16ap 把 owner-hot write-back 的成功终态收敛为 exact-cohort demotion，内部状态流如下： + +| 阶段 | 必须满足的条件 | 失败行为 | +| --- | --- | --- | +| hot 触发 | Moka `RemovalCause::Size`、当前 `put_id` 和 `MemoryInfo` identity 匹配，完整 owner atomic group 可 pin | 不创建任务，GPU source 保持当前状态 | +| remote append | target 被 `remote_only_node_roles=["remote_cache"]` 限定，append 发布同版本 CPU route | NoSpace/transfer/append 失败释放 pin，不请求 source demotion | +| cohort 验证 | 每个 atomic-group member 当前、每个 TP rank 边界一致，所有成员在同一个 CPU owner 上 live | 本次只计 attempt，不选择任何 GPU resident entry | +| exact selection | resident descriptor 仍是预期版本且当前 route 仍可恢复 | 跳过变化条目，不能删除随后写入的新版本 | +| 物理回收 | Moka Size listener 安装 master activity fence，再执行 owner Prepare/Commit/Finalize | Busy 有界重试；耗尽后仅把仍为当前版本的条目安全回插 | + +E16ap 压测把 owner hot ratio 设为 0.95，与 master resident ratio 对齐。128 GiB GPU owner 的 +逻辑触发点为 121.6 GiB；这只改变冷数据开始异步迁移的时机，不改变 128 GiB 物理 segment。 +CPU 未满时，完整冷 cohort 写入 CPU 后释放 GPU copy；CPU 无空间时，任务不降级，GPU 仍可 +使用剩余物理 headroom。Master 每 30 秒额外报告 `owner_hot_demotion_attempts`、 +`owner_hot_demotion_cohorts` 和 `owner_hot_demotion_selected_bytes`，用于区分“复制完成”与 +“已经转化为独立二级容量”。正式三机结果需要同时满足非零 demotion、CPU-source Get、两侧 +pending 归零和 QPS 高于双 GPU no-replica 基线，不能只用 CPU retained bytes 判定策略有效。 + +两 Fluxon 模拟压测的完整记录保存在本工作区的 `remote_cpu.md`。当前一次完整压力运行得到： + +```text +Moka eviction callbacks: 11463 +safe reclaim reclaimed: 11452 +safe reclaim restored: 11 +processed eviction weight: 96158613504 bytes +``` + +`reclaimed + restored` 与 callback 数严格相等；11 个 Busy/未回收条目按当前版本重新插回 Moka。8 个均匀抽样 key 的 payload 长度和首尾字节均校验通过，其中 3 个已确认在 A local backing 被驱逐后通过 B -> A P2P 读取恢复。该次日志中 safe reclaim queue full/closed、NoSpace 返回调用方、panic、OOM 和 batch failure 均为 0。 + +这次稳态压力运行由按真实 slot 权重的常规 Moka capacity eviction 提前释放内存，因此没有实际触发“申请不到 `512 MiB` grant 后调用 `evict_some`”的 NoSpace 兜底。`evict_some` 另有定向 Moka 契约测试，覆盖请求驱逐量大于/小于现有占用、Busy 条目回插后的 LRU 顺序，并断言调用前后 `max_capacity` 不变。Master activity fence 和 Moka 路径已有单元测试；Owner Prepare/Commit/Abort/Finalize 的主要证据目前来自端到端压测，后续仍应补充直接的并发与故障注入单元测试。 + +#### E16au--E16ax 包含式写回与 local IPC 结论 + +E16au 将 owner-hot 的成功终态改为包含式 CPU write-through：CPU route 发布成功后不立即 +demote GPU source，GPU backing 只由独立的 95% resident 压力回收。固定 +GPU0/GPU1/CPU = 128/128/256 GiB 的正式压力得到 1152/1152、6.7446 QPS、92.90% 总命中； +CPU 保留 8,793 entries / 73.761 GB，CPU-source Get 为 4,349 / 20.521 GB，owner-hot +demotion 与终态 pending reclaim 均为 0。包含式语义消除了主动降级争用，但仍未超过双 GPU +baseline 或 Mooncake；提前 CPU copy 大多只是已有 GPU source 的替代恢复源。 + +这轮有 255,972 条 iceoryx2 `FailedToDeliverSignal`。开放层调低日志等级无法控制闭源 +`libfluxon_commu_core.so` 内静态链接的独立 log-level 符号；完全关闭 local IPC 又在同机 +direct P2P 路径触发 608、`KeyBeingWritten` 终态失败和 Prefill OOM。E16ax 只把 +owner/client receiver 改为 WaitSet 后虽完成 1152/1152，QPS 降至 6.6932；warning 仍有 +232,723 条，且 96.85% 移入 SGLang 关键进程。故 local IPC 必须保留，WaitSet 和完全禁用 +两条方案均不能作为当前修复。 + +#### E16ay Get-target 压力与 local reserve 物理容量闭环 + +> 本节记录 E16ay 当时的历史实现，不是当前容量规约。后续修复已经把 configured +> `expected_grant_count` 改为硬上限，并删除 route 全表扫描；当前行为以本章开头规约表和 +> “Local-side 主动驱逐与安全回收”章节为准。 + +E16ay 当时 local reserve 的 expected capacity 是下限而非硬上限。Put 并发需求会使它从配置的 +grant 数继续扩张；后续 slot 回收只让 grant 内部变空，并不会自动把这段 512 MiB 物理 +allocation 归还 master segment allocator。原 shrink actor 还只检查队尾 grant，并要求连续全空 +5 秒。 +另一方面，CPU -> GPU remote Get 的 target 仍是 master 直接从请求 owner 的普通 allocator +分配。因此系统可能同时出现： + +- master Get target 报 `NoSpace/free_capacity=0`； +- owner reserve 内存在大量 free slots； +- `pending_eviction_reclaim_bytes=0`，说明这不是 safe reclaim backlog。 + +E16ax 的实测即为 GPU0/GPU1 18/9 次 Get-target `NoSpace`，而 reserve 最终仍有 1,085/989 个 +空闲 4,718,592-byte slot。E16au 精确重扫也有 10/26 次。该错误会在 atomic-group prefix +计算前截断可传输前缀，使已经存在的 CPU route 不能转成额外命中。 + +修复契约如下： + +| 阶段 | 规则 | +| --- | --- | +| 首次 batch get-start | 保留每个 item 的独立结果，只收集 `API_NO_SPACE` 索引；成功项、KeyNotFound 和其它错误不重发。 | +| 物理让位 | owner 可同步摘除任意完全空闲的多余 grant，不受队尾和 idle cooldown 限制；必须先收到 `release_local_grant` 完成，再重试 Get。 | +| 容量下限 | 归还后 grant 数不得低于 `max(expected_grants, ceil((used_slots + pending_slot_demand) / slots_per_grant))`；Prepared、PendingLocalVisible、Committed 或仍有 holder 的 slot 都不能释放。 | +| 有界重试 | 每轮只替换仍为 NoSpace 的 item，最多 8 轮；不能安全归还 grant 时保留原错误和 prefix fallback。 | +| 容量语义 | 不修改 `contribute_to_cluster_pool_size.dram`、resident/hot Moka capacity 或 expected reserve；这是把已经空闲的物理 grant 及时归还 allocator，不是缩 owner 容量。 | + +该路径与 `evict_some` 的职责不同：本步骤优先回收“已经完全空闲但仍被 reserve 持有”的 +grant，不产生新的 KV 淘汰；若没有这种 grant，现有 safe reclaim/NoSpace 语义继续生效，不能 +为了 Get 强行释放 live backing。新增定向测试覆盖只选 NoSpace item、非队尾 grant 与容量下限。 + +E16ay 的固定 128/128/256 GiB 三机正式验收得到 1152/1152、6.8269 QPS 和 93.02% 总命中。 +GPU0/GPU1 分别通过 2/3 次 pressure retry 恢复 30/54 个 NoSpace item，全部首轮达到 +`no_space_after=0`，没有错误传到 SGLang。CPU-source Get 为 4,942 / +23,319,281,664 bytes；CPU 保留 8,855 entries / 74,281,123,840 bytes；两个 GPU terminal +pending 与 owner-hot demotion 都为 0。配置的 owner 容量、resident/hot 容量和 232-grant +expected 下限没有改变。这证明同步归还空 grant 可以闭合 Get-target 物理 headroom，但 QPS +仍比双 GPU baseline 低 0.79%、比 Mooncake 低 1.33%。 + +#### 单 KV slot 分级循环（取代整 grant Get-target 让位） + +E16ay 的整 grant 让位只是在旧 Get target 仍从 master 普通 allocator 分配时的兼容方案， +不是最终的缓存分级语义。`512 MiB` grant 只是 local-reserve 的物理容器，不应成为 KV +驱逐或恢复的最小单位。多个 grant 中的所有 committed KV slot 已由同一个 owner-hot Moka +按 key、版本和实际 slot weight 统一管理；因此正确闭环是逐 KV slot 循环： + +```text +任意 grant 的 Committed slot + -> 跑出 owner hot window + -> owner 展开完整 exact atomic/TP source cohort + -> BatchEvictOwnerSource + -> master activity fence + owner holder drain + -> owner route 引用与 resident holder 一次 pool update 归还 + -> master 删除 exact source route + -> 该 slot 立即 Free + -> 后续 remote Get 直接 claim 这个 Free slot +``` + +CPU proactive replica 可以在这条路径之前、期间或完全不发生；它只改变删除 source 后还能否 +remote hit，不是释放 slot 的前置条件。 + +这条路径不等待同一 grant 的其它 slot 释放，也不排空或归还整个 grant。普通 idle shrink +仍可在 grant 已经自然全空且高于保留下限时回收物理 allocation，但它不参与 Get 的正确性 +或前进性。 + +Moka size eviction 是冷候选 ownership transfer，不能被当作可丢失通知。如果触发 key 的 +atomic/TP cohort 元数据暂不完整，owner 保留同一个 selected identity，在有界 retry actor 中 +等待当前 route/metadata 可构造完整 cohort；不能退化成逐 key source delete，也不能重新进入 +replica append。RPC/Busy/partial overlap 时 GPU route/holder 保持 live，并且事件必须以 +handoff、committed、restored、obsolete 或 retry 状态可对账，不能让 slot 永久留在 hot cache +之外。 + +remote Get 的 local-side target 使用与 hostless Put 相同的全 pool slot allocator；allocator +可以从任意 grant 取一个匹配 `slot_size` 的 Free slot。没有 Free slot 时，复用现有 Moka +候选与 owner 两阶段 fence，逐 KV reclaim 所需数量的 slot，而不是选择某个 grant drain。 +slot 状态机和失败契约如下： + +| 阶段 | 约束 | +| --- | --- | +| GetStart | client 先 claim `Free -> Prepared`，把精确的 `grant_id/slot_index/slot_size/address` 交给 master；master 校验 owner 和 grant 几何，不能改用另一个 target。 | +| transfer 完成 | `Prepared -> Pending`，唯一 resident `MemoryInfo` 放入 hidden pending-get 表；此时 bytes 不得通过 local-visible index 暴露。 | +| GetDone 成功 | master 以 key/version/owner CAS 发布 committed-slot route，并确保 resident Moka entry 已可见；client 再把同一 `MemoryInfo` 原子移入 committed index、`Pending -> Committed`，并加入 owner-hot Moka。 | +| transfer/GetDone 明确失败 | revoke master inflight，移除 hidden pending；最后一个 `MemoryInfo` holder 归还后单 slot 回到 Free。 | +| GetDone 响应丢失 | 保留 hidden pending，用同一 `get_id` 重试幂等终态；不能猜测失败并释放可能已经被 master 发布的 slot。 | +| 并发或旧版本 | 同 key+owner 只能有一个 live committed route；CAS 输家不覆盖现有 backing。delete/reclaim 按 put id 清理 hidden pending，避免 master 已发布而 client 尚未 promote 的窗口泄漏。 | +| owner source eviction | 完整 exact cohort 可直接请求删除当前 owner source，不依赖 CPU route；最后一份 cache route 也可删除。master fence 与 holder drain 完成前 slot 不得复用。 | + +hostless Put 也必须经过同一发布门禁。`local_fast_put_commit` 先把写完的 slot 置为 +`PendingLocalVisible`，只允许当前 owner 本地读取；master PutDone 已发布 route 并同步完成 +resident Moka admission 后，client 才执行 `PendingLocalVisible -> Committed(route_live=true)`。 +同一 batch 的成功成员要先全部完成 promotion，再统一加入 owner-hot Moka，避免第一个成员刚 +admit 就被容量淘汰，而其 route 或 atomic/TP cohort 仍未完整发布。PutDone 明确失败或响应 +结果不确定时都保留 pending/fence，并使用同一操作身份重试；不能释放可能已经被 master +发布的 backing，也不能把部分 item 的错误当作整批可回滚。 + +发布生命周期必须以 batch/job 为单位持有，而不是每个 key 各自 fire-and-forget。一个 +`OwnerLocalPublishJob` 持有整批 master key reservations；external local-first 路径还持有 +整批 `ExternalPendingPutCtx` 和 owner reclaim fences。PutDone 响应不确定、响应长度/identity +不一致或任一 item 未达终态时，job 保留全部状态并以 25ms～1s 退避重试。全部 master route +成功后才开始 local promotion；若取消或局部错误发生在 promotion 中间，下一轮接受已经 +promoted 的成员并继续其余成员，只向前收敛。所有 promotion 完成后才统一 owner-hot +admission。只要任一 item 声明 atomic group，而 job 未包含连续、同序的完整成员，就禁止 +退化到逐 key PutDone；不得用“部分成功”换取半组可见性。 + +`Committed` slot 的可复用条件仍是 `route_live == false && holder_ref_count == 0`；释放 route +与释放 holder 各执行一次，任意先后顺序都不能提前复用。Get 使用 local-reserve slot 不占旧的 +master 普通-allocation durable quota；旧 quota token 则必须绑定普通 route 生命周期，在 +revoke、失败、替换或删除时恰好归还一次。 + +容量语义保持不变：GPU owner 仍是配置的 `128 GiB`，CPU owner 仍是 `256 GiB`；该改造 +不调整 `contribute_to_cluster_pool_size.dram`、resident/hot Moka capacity 或 expected grant +数量，只让完成 exact source deletion 的单个 KV slot 立即回到同一 size class 的 Free list。 + +#### E16az external local-IPC busy-poll 验收 + +E16ay 的两个 GPU owner 仍分别产生 147,124/150,701 条 iceoryx2 +`FailedToDeliverSignal`。失败的是 owner -> SGLang external 的逐消息 event notification， +subscriber 数据本身仍存在。E16az 保持 E16ay release 与所有容量/策略/工作负载不变，仅令 GPU +owner 和 SGLang external receiver 使用 `iceoryx_external_busy_poll=true`。此模式直接 drain +subscriber，发送端不会创建 notifier；owner/client 内部 receiver 原有的 busy-poll 设置不变。 + +正式轮仍为 1152/1152、576/576，四份相关日志的 `FailedToDeliverSignal` 全部为 0，GPU owner +日志缩到约 132/137 KB。local IPC 和 RDMA 数据路径仍存活，CPU-source Get 为 4,965 / +23,427,809,280 bytes，CPU 保留 8,610 entries / 72,225,914,880 bytes；74 个 Get-target +NoSpace item 全部被 E16ay 路径恢复。terminal pending、surfaced NoSpace、业务 5xx、P2P 608、 +OOM 和 scheduler exception 均为 0。 + +性能为 6.7888 QPS、93.02% 总命中，比 E16ay 低 0.56%、比双 GPU baseline 低 1.34%。因此 +事件告警格式化不是剩余吞吐主瓶颈；持续 external polling 的 CPU contention 可能抵消省下的 +日志工作，且这一小幅差值可能包含单轮波动。设计上可把 external busy-poll 作为避免 notifier +风暴的运行模式，但不能把它视为 CPU replica 转化为吞吐收益的容量策略。下一步仍应直接减少 +包含式重复副本的无效工作，或降低 CPU Get/replica 进入 SGLang 关键路径的等待开销。 + +#### E16ba Hot/resident lead-window 契约 + +包含式 owner-hot write-through 不能通过简单地把 hot capacity 提到 resident capacity 来实现 +“驱逐即 write-back”。hot eviction 只产生异步 CPU append 信号；resident eviction 则要求 +CPU 上已经存在完整、同版本、同 TP cohort 的 durable backing。两者必须保留足够的 lead +window，否则 resident pressure 会先于 remote completion 到达。 + +固定 resident=0.95、reserve=0.90 和 128/128/256 GiB 容量，对 hot=0.92/0.94/0.95 的三次 +正式冷启动扫描结果为： + +| hot | QPS | CPU placement | CPU Gets | recoverable selected | shortfall | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 0.90 E16ay | 6.826884 | 8,860 | 4,942 | 65,729,986,560 B | 0 | +| 0.92 | 6.824876 | 7,413 | 3,214 | 51,091,865,600 B | 0 | +| 0.94 | 6.792892 | 5,893 | 2,108 | 41,800,433,664 B | 0 | +| 0.95 | 6.623042 | 5,341 | 1,496 | 36,394,500,096 B | 32,925,181,564 B | + +四轮 overall hit 都为 93.02%，因此下降不是 cache miss 数量变化。0.92 以几乎相同 QPS +减少 16.3% placement 和 35.0% CPU Get，证明较早触发包含冗余工作；0.94 开始变慢；0.95 +请求 69,316,116,502 bytes recoverable victim 时只有 36,394,500,096 bytes 已 CPU-complete， +QPS 相对 0.90 下降 2.99%。三轮均无 surfaced NoSpace、terminal pending、业务错误或容量变化。 + +因此静态 hot threshold 在该 workload 下必须比 resident 至少提前约 3--5pp。若设计目标是真正 +的 eviction-time write-back，应由 resident 的精确 victim 触发 CPU append，并为该 victim +建立有界 completion fence；或者根据 replica queue depth、传输延迟和 resident growth rate +动态计算 lead。不能让两个独立 Moka 同阈值后依赖调度碰巧完成。 + +#### E16bc--E16bj SGLang 调度容量、replica pipeline 与发布门禁 + +固定 Fluxon GPU0/GPU1/CPU owner 为 128/128/256 GiB 后，端到端吞吐的最大增益并非继续改 +Moka 水位，而是解除 SGLang scheduler 的串行与 token-pool 约束。E16bc 只把每个 TP=2 服务的 +`max_total_tokens` 从 100,000 提到 200,000，QPS 从 E16bb 的 6.773531 提到 8.645063；E16cc +再启用默认 overlap scheduler，正式冷启动达到 10.056034 QPS。该轮 1,152/1,152 成功、两端 +各 576 请求、总命中 92.96%，GPU 初始化后仍各有约 53.51 GiB 可用显存。这里的 +`max_total_tokens` 是 SGLang GPU scheduler/KV pool 参数，不修改 Fluxon owner segment、resident +Moka、hot Moka 或 local-reserve 的容量；配置与结果记录不能把两种“容量”混为一谈。 + +overlap scheduler 允许 HiCache restore/prefill 的调度推进与 GPU decode 工作交叠。关闭 overlap +时，同样的外部命中率并不表示同样的吞吐，因为 storage future、prefill admission 和 decode +会在 scheduler step 边界上串行等待。因此正式性能对比必须同时固定 +`max_total_tokens`、`disable_overlap_schedule`、CUDA graph、并发和 workload round barrier。 + +owner replica actor 原先全局只能有一个 append pipeline 在途，每个不同 key 都依次等待 +append-start、RDMA completion 与 append-done。`test_spec_config.replica_task_max_inflight` 将其改为 +1--64 的有界跨 key 并发，默认 1 保持兼容；全局 semaphore 限制总在途数，每个 key 另有单许可 +semaphore，所以不同 key 可以并行，同 key 的版本/任务仍严格串行。E16bh 在两个 GPU owner 设为 +16，CPU owner 保持 1；该参数改变复制流水线深度，不改变 owner 容量、placement 或原子组完整性。 + +组合发布还必须经过目标机 ABI 预检。E16bf 的 wheel 文件名虽然声明 +`manylinux_2_28_x86_64`，其中新 closed core 实际引用了 `GLIBC_2.39`，三台正式机器因此在 master +import 阶段 fail-fast，未发送任何 workload。发布门禁应对 wheel 内每个 bundled `.so` 检查最高 +GLIBC symbol version，并在三种目标 Python/宿主环境做 import smoke；仅检查 wheel tag 或构建成功 +不足以证明可部署。兼容性回退 E16bd 保留开放层 replica pipeline，并沿用已验证 closed core。 + +E16bj 在真实 manylinux 2.28 builder 中重建 closed core 后，通过了 wheel 内全部 `.so` 的 +symbol-version 扫描和 Python 3.10/3.11/3.12 import smoke。正式三机冷启动保持 +GPU0/GPU1/CPU = 128/128/256 GiB、GPU pipeline=16、CPU-only placement、200k token pool 和 +overlap scheduler，得到 1,152/1,152、576/576、10.045299 QPS 和 92.99% 总命中。CPU 承担 +8,770 个 replica target 与 4,990 次 remote Get；GPU replica target、终态 replica inflight、 +replica failure 和 pending reclaim 均为 0。E16cc/E16bh/E16bj 的最低和平均 QPS 为 +10.020985/10.040772，说明 10-QPS 结果可由可部署组合版本复现。 + +E16be closed sharding 修复要求所有正长度 transfer 的 child 都非空：按 sharding quantum 做 ceil +division并省略空 child；只有显式零长度请求保留一个零长度 child，以维持 completion contract。 +这项约束属于通信正确性和尾延迟卫生，不能用 wheel tag 代替测试，也不能从 E16bj 相对 E16bh +仅 +0.24% 的单轮差值推导出稳定吞吐增益。 + +## SGLang Node Storage 状态 + +SGLang 侧的 node metadata 不等价于 Fluxon master route 状态。当前四个字段建议按下面语义解释： + +| 字段 | true 的含义 | 清理时机 | +| --- | --- | --- | +| `storage_staged` | 该 node 有一批 Fluxon hostless backup 正在 staged 路径中。 | `KvFuture` 完成或失败后清空。 | +| `storage_local_ready` | CUDA write 已完成，SGLang 已调用 `local_fast_put_commit`，但返回的 `KvFuture` 还未 ack；该状态只表示本次 hostless backup 已进入 Fluxon commit 流程，不表示 KV route 已经全局确认。 | async ack 结束后清空。 | +| `storage_pending` | Fluxon `KvFuture` 还没结束。 | future 成功或失败后清空。 | +| `storage_backed` | Fluxon 后台提交成功，KV route 已确认可作为 shared backing。 | 该 node 被删除或失效时清空。 | + +因此，SGLang 可以用 `storage_staged/storage_local_ready` 判断本次 hostless backup 已经推进到本地写入或 commit 阶段；但跨节点复用和长期共享必须等 `KvFuture` 成功，并以 `storage_backed` 为准。 + +TP 场景下，每个 rank 仍有各自的 radix tree 和恢复决策。`get_start` 的结果只描述当前 rank 这批 keys 的可恢复前缀；如果一个 rank miss、另一个 rank hit，上层必须按 SGLang 的 TP restore 约束处理一致性，不能把单 rank 的部分成功当作完整 request 已恢复。 + +Put admission 的随机 identity 必须跨 TP rank 稳定。物理 storage key 保留 `_0_2`、`_1_2` +等 rank 后缀，用于区分不同 rank 的 KV bytes；ratio/score jitter 和 atomic-group hash 则使用 +保留 TP size、去掉 TP rank 的 admission key。同一个逻辑 radix group 因而在所有 TP ranks +上得到同一 admitted/skipped 决策，同时仍生成不同物理 keys。只约束单 rank 内 mask 不切组， +无法满足这个跨 rank 不变量。 + +各 rank 首次 `get_start` 的 `transferable_len` 不一致时，SGLang 先用 collective 取最小值。 +该最小值必须落在共同的 atomic-group 边界；每个 rank 取消原 handle，再以该最小完整-group +前缀重新执行一次 `get_start`，随后再次 collective 校验长度一致。第二次仍不一致才放弃本轮 +restore。不能直接使用较长 rank 的 handle，也不应因为尾部差异把两个 rank 已共同命中的 +前缀全部作废。 + +`get_start` 长度一致只完成了 intent 协商，不能作为 restore 的最终提交点。每个 rank 在 +`get_transfer` 返回可用 plan 后必须再进入一次 TP commit gate；只有所有 rank 都成功，plan +才可进入 ready-prefetch、radix node 和 layerwise restore。任一 rank 超时、P2P 失败或 plan +缺失时，所有 rank 统一取消 handle、释放已成功 rank 的 views，并把本轮 restored-token 数置为 +0，按共同 cache miss 重算。这个门禁必须位于任何 rank 发布 restored prefix 之前，禁止一侧 +restore、另一侧 recompute。异步 CUDA descriptor submission 的错误也必须在消费 prefix 前对 +所有 rank 可见；已经修改临时 node/slot 的实现必须全 rank 回滚后才能继续调度。 + +## 失败处理 + +| 场景 | 必须动作 | +| --- | --- | +| `local_fast_put_start` 后 native write 失败 | 调用 `put_abort(plan_ptr)`，释放 key reservation 和 local reserve slot lease。 | +| `local_fast_put_commit` 返回 future 后后台失败 | SGLang 清理 `storage_staged/storage_pending/storage_local_ready`，必要时删除已 evicted 的 dead leaf。 | +| `get_start` 后放弃 restore | 调用 `cancel_get_transfer(handle)`，释放 get-start 持有的 owner/external 资源。 | +| `get_start` 只命中部分前缀 | 只允许恢复 `transferable_len` 覆盖的完整 atomic groups；后续 page 按 miss 处理。 | +| 任一 TP rank 的 `get_transfer` 失败或 plan 缺失 | TP commit gate 在所有 rank 上拒绝发布；释放成功 rank 的 views，统一按 0-token miss 重算。 | +| `get_transfer` 成功后 native restore 失败 | 先 `release_views(plan_ptr)`，再执行 SGLang rollback；此时 handle 已被消费。 | +| CUDA host registration 失败 | direct path 同步失败，不能降级为未注册 host memory。 | +| `plan_ptr` 类型用错 | `local_fast_put_commit`、`put_abort`、`release_views` 都按 registry entry 类型校验并 fail fast。 | diff --git "a/fluxon_doc_cn/design/teststack_2_Benchmark\347\233\221\346\216\247\345\267\245\345\205\267\351\223\276\344\270\216AIPerf\346\216\245\345\205\245\350\247\204\345\210\222.md" "b/fluxon_doc_cn/design/teststack_2_Benchmark\347\233\221\346\216\247\345\267\245\345\205\267\351\223\276\344\270\216AIPerf\346\216\245\345\205\245\350\247\204\345\210\222.md" new file mode 100644 index 0000000..37cb6c6 --- /dev/null +++ "b/fluxon_doc_cn/design/teststack_2_Benchmark\347\233\221\346\216\247\345\267\245\345\205\267\351\223\276\344\270\216AIPerf\346\216\245\345\205\245\350\247\204\345\210\222.md" @@ -0,0 +1,565 @@ +# Benchmark 监控工具链与 AIPerf 接入规划 + +> 状态：设计阶段 +> 调研基线：2026-07-10 +> AIPerf 基线：[`ActivePeter/aiperf:teleai`](https://github.com/ActivePeter/aiperf/tree/teleai)，提交 [`d72160e20957013d6608afcc88ed24100cb27dc5`](https://github.com/ActivePeter/aiperf/commit/d72160e20957013d6608afcc88ed24100cb27dc5) + +相关文档：[TestStack 架构与 CI 测试流程](./teststack_1_当前架构与CI测试流程.md)、[本地文件日志与 Greptime OTLP 导出链路](./log_1_本地文件日志与Greptime_OTLP导出链路.md)。AIPerf 能力判断依据固定提交下的 [architecture](https://github.com/ActivePeter/aiperf/blob/d72160e20957013d6608afcc88ed24100cb27dc5/docs/architecture.md)、[API endpoints](https://github.com/ActivePeter/aiperf/blob/d72160e20957013d6608afcc88ed24100cb27dc5/docs/reference/api-endpoints.md)、[server metrics](https://github.com/ActivePeter/aiperf/blob/d72160e20957013d6608afcc88ed24100cb27dc5/docs/server-metrics/server-metrics.md) 和 [profile exports](https://github.com/ActivePeter/aiperf/blob/d72160e20957013d6608afcc88ed24100cb27dc5/docs/tutorials/working-with-profile-exports.md)。 + +## 1. 结论 + +Fluxon 当前已经具备 benchmark 编排、终态结果、服务指标、日志和常驻 UI，但这些能力仍分布在不同链路中。现有链路适合 KV、MQ、RPC 和 FS benchmark，缺少生成式 AI 请求的 TTFT、ITL、token throughput、goodput、trace replay 等负载和指标语义。 + +规划采用以下边界： + +- **AIPerf 只负责生成式 AI 负载和请求侧统计**：新增显式的 `AIPERF` test stack mode，首期支持单个 load generator 对 OpenAI-compatible SGLang endpoint 发压。 +- **`test_runner` 继续负责 case 生命周期**：suite 编译、资源准备、启动、超时、取消、终态、历史记录、UI 和对外 API 都由 `test_runner` 管理。 +- **Greptime 继续负责连续时序指标与日志**：Fluxon 服务指标、日志和已有监控链路不迁移到 AIPerf。AIPerf 的请求结果以 run artifact 为事实来源。 +- **AIPerf UI 和 API 不成为新的公共入口**：AIPerf 以 `runtime.ui: none` 运行；其 loopback API 只供同机适配器读取，`test_runner` UI 通过现有 ops 接口展示标准化后的进度和结果。 +- **现有 distributed benchmark coordinator 保持原路径**：KV、MQ、RPC、FS 的 `benchmark_result.json` 语义不由这次接入改写。`AIPERF` 是现有有限模式集合中的新分支。 +- **依赖必须固定到提交和构建产物哈希**：`teleai` 是未保护的可变分支，运行时不能直接按分支头安装。 + +首期不启用 AIPerf sweep、AIPerf multi-run、AIPerf OTel 导出和 GPU telemetry。case 矩阵仍由 TestStack 的 `scene × scale × profile` 决定，避免形成嵌套调度和重复监控通道。 + +## 2. 目标与非目标 + +### 2.1 目标 + +1. 把生成式 AI benchmark 收敛进 `start testbed -> testrunner` 两步模型。 +2. 复用 AIPerf 已有的请求生成、warmup、并发或请求速率控制、请求级记录和聚合统计。 +3. run artifact 和请求链路统一使用 `case_id + run_index`；连续服务指标通过 `cluster/instance + run time window` 与 run 关联。 +4. 让常驻 `test_runner` UI 同时展示实时进度、终态摘要、artifact 和 Greptime 时间窗口。 +5. 保证离线安装、固定依赖、固定随机种子和明确的失败状态。 + +### 2.2 非目标 + +- 不用 AIPerf 替换 `distributed_benchmark_coordinator.py`。 +- 不把 AIPerf FastAPI、TUI 或 plot dashboard 暴露为 TestStack 公共服务。 +- 不在 Fluxon 内重写 TTFT、ITL、token throughput 等 AIPerf 已提供的算法。 +- 不允许 suite 同时传入自由格式 AIPerf YAML、任意 CLI 参数和环境变量覆盖。 +- 首期不支持多 load generator、AIPerf Kubernetes service mode、真实用户 prompt、鉴权 endpoint 和自动性能回归判定。 + +## 3. 当前工具链分析 + +### 3.1 当前数据流 + +```mermaid +flowchart LR + A[ci_test_list.yaml] --> B[test_runner.py] + B --> C[resolved_case.yaml] + B --> D[distributed benchmark coordinator] + D --> E[benchmark nodes] + E --> F[benchmark_result.json] + F --> B + B --> G[summary.yaml + case_runs.yaml] + + H[Fluxon processes] --> I[Prometheus remote write] + I --> J[Greptime] + H --> K[local logs / OTLP logs] + K --> J + + G --> L[test_runner UI] + J --> L + M[ops interfaces] --> L +``` + +当前各层职责如下。 + +| 层 | 当前实现 | 已有能力 | 当前边界 | +| --- | --- | --- | --- | +| suite 与 case 编译 | `test_runner.py` | `scene × scale × profile`、`resolved_case`、run 目录 | 没有生成式 AI workload branch | +| testbed 生命周期 | `start_test_bed.py`、controller、ops | 启动共享服务、apply、进程状态和日志 | 不执行单个 benchmark case | +| benchmark 执行 | coordinator + nodes | KV、MQ、RPC、FS 的分布式同步和聚合 | 指标模型面向 operation、bytes 和内部 phase | +| 终态 | `benchmark_result.json`、`summary.yaml`、`case_runs.yaml` | 强制等待结果、校验节点完成、记录 run outcome | UI 主要展示原始摘要，缺少通用指标视图 | +| 连续监控 | Prometheus-compatible remote write、Greptime | 服务指标、部分 transport 指标、日志 | 与 run artifact 的关联字段尚未统一 | +| 常驻 UI | `test_runner_ui.py` | suite/run 状态、日志、GitOps、ops log | 当前没有 benchmark 曲线和统一结果 API | + +### 3.2 当前链路已经稳定的部分 + +- `case_runs.yaml` 是 suite workdir 内执行状态的单一事实来源。 +- `summary.yaml` 是单次 run 的终态摘要，runner 会在执行前写入可诊断的占位结果。 +- TestStack benchmark 通过 `benchmark_result.json` 完成终态握手。runner 会校验每轮 `completion.status`、节点数量和 operation 数量。 +- UI 是常驻服务，run 结束后仍能读取历史、日志和 ops 状态。 +- Fluxon 指标使用 Prometheus-compatible 协议进入 Greptime，服务日志同时保留本地文件和 OTLP 路径。 + +### 3.3 主要缺口 + +| 缺口 | 影响 | 本设计的处理 | +| --- | --- | --- | +| 缺少生成式 AI 请求语义 | 现有 operation latency 无法表达 TTFT、ITL、token throughput 和多轮会话 | 使用 AIPerf 计算并保留原始 metric tag 与单位 | +| 进度、终态和时序指标分散 | UI 需要分别读状态文件、原始 JSON、日志和 Greptime | 在现有 `/api/run_state` 下增加有界的 `benchmark` 分支 | +| run 关联字段未贯穿 | 很难把请求突刺与服务指标、日志对齐 | 固定 artifact 关联键，并保存 cluster/instance 与 benchmark 时间窗口 | +| 当前 UI 缺少 benchmark 图表 | 只能查看 JSON 和日志 | UI 读取标准化摘要、timeslice 和 server metrics artifact | +| 多节点 percentile 聚合口径有限 | 当前 distributed benchmark 的 p50/p95/p99 是按成功 operation 数加权的节点 percentile，不能作为全量样本 percentile 使用 | AIPerf 路径保留请求级 JSONL，并由 AIPerf 计算请求总体 percentile | +| 工具能力容易重复 | AIPerf 也有 UI、API、server metrics、OTel 和 GPU telemetry | 用明确的所有权表限制每条链路 | + +上表中的 percentile 结论只描述当前 coordinator 的节点聚合实现，不扩展为其他 benchmark 或整个监控系统的结论。 + +## 4. AIPerf `teleai` 分支评估 + +### 4.1 可直接复用的能力 + +| 能力 | `teleai` 快照行为 | Fluxon 用法 | +| --- | --- | --- | +| Python 版本 | `requires-python = ">=3.10,<3.14"` | 满足 Fluxon 的 Python `>=3.10` 要求 | +| 许可证 | Apache-2.0 | 构建产物保留许可证与 attribution | +| endpoint | OpenAI chat/completions 等 | 首期只开放 OpenAI-compatible chat streaming | +| 负载控制 | concurrency、request rate、warmup、ramp、trace replay | 首期开放 fixed concurrency 和 fixed request rate 两个有限分支 | +| 请求指标 | TTFT、ITL、request latency、request/token throughput、error 等 | 直接采用 AIPerf metric tag 和单位，不在适配层重算 | +| artifact | summary JSON/CSV、请求级 JSONL、`inputs.json`、timeslice | summary JSON 和请求级 JSONL 为首期必需产物 | +| server metrics | 自动发现或显式抓取 Prometheus `/metrics`，当前默认每 333 ms 采样 | 首期只抓取 runner 解析出的单个 SGLang metrics endpoint | +| 实时 API | `/api/run`、`/api/progress`、`/api/metrics`、`/api/results`、`/metrics` | 只绑定 loopback，由同机适配器读取 | +| UI | dashboard、simple、none | 固定 `none`，避免出现第二套 UI 归属和入口 | +| OTel | 可流式导出 metrics 和 timing | 首期关闭，待 Greptime OTLP metrics 契约单独评审 | + +### 4.2 依赖快照结论 + +截至 2026-07-10： + +- `teleai` 指向提交 `d72160e20957013d6608afcc88ed24100cb27dc5`。 +- GitHub compare 显示该提交下 [`main...teleai`](https://github.com/ActivePeter/aiperf/compare/main...teleai) 为 `identical`。 +- 仓库中没有可识别的 TeleAI 专属配置、插件或 patch。 +- `teleai` 分支未启用 branch protection。 + +因此，“使用 `teleai` 分支”在工程上必须落实为下面四项记录： + +1. source repo：`https://github.com/ActivePeter/aiperf` +2. requested branch：`teleai` +3. resolved commit：`d72160e20957013d6608afcc88ed24100cb27dc5` +4. wheel 与 wheelhouse manifest 的 SHA-256 + +实现前还需要业务方确认：当前意图是否就是采用这份与 fork `main` 相同的快照；如果预期存在 TeleAI 定制能力，应先明确对应提交或差异清单。 + +### 4.3 不能直接沿用的默认行为 + +- **不能按 branch head 在线安装**：testbed 运行期间不访问 PyPI 或 GitHub。 +- **不能使用环境变量替换普通参数**：suite 编译器生成完整 AIPerf YAML，启动命令只传 `--config`。 +- **不能使用 AIPerf sweep 或 multi-run 展开 case**：TestStack 已拥有 case 空间和 run history。 +- **不能依赖 AIPerf 默认 UI**：非 TTY 与 TTY 下必须都固定为 `none`。 +- **不能依赖 server metrics 自动发现**：显式写入 SGLang `/metrics` URL，并关闭 Kubernetes discovery。 +- **不能直接公开 AIPerf API port**：API 只监听 `127.0.0.1`，端口由 testbed port allocator 分配。 + +## 5. 目标架构 + +### 5.1 核心角色 + +| 角色 | 责任 | 不负责什么 | +| --- | --- | --- | +| `start_test_bed.py` | 准备 SGLang、Greptime、controller、ops 和常驻 runner UI | 不启动单次 AIPerf run | +| `test_runner.py` | 编译 `AIPERF` case、分配资源、驱动 prepare/execute/finalize、写终态 | 不计算 TTFT 或 ITL | +| AIPerf adapter | 在 load generator 上启动 AIPerf、轮询 loopback API、原子写进度、校验 artifact、生成终态 | 不拥有 suite、历史或公共 API | +| AIPerf | 发请求并计算请求侧指标 | 不拥有 testbed 和长期监控 | +| SGLang endpoint | 被测服务，暴露 OpenAI-compatible API 和 `/metrics` | 不决定 benchmark 成败 | +| Greptime | 保存 Fluxon 连续指标与日志 | 首期不保存 AIPerf 请求级记录 | +| `test_runner` UI | 组合 run 状态、AIPerf 摘要、artifact、ops log 和 Greptime 时间窗口 | 不直接连接公开的 AIPerf 服务 | + +实现代码按现有 runner 分层落位： + +| 模块 | 计划改动 | +| --- | --- | +| `test_runner.py` | 增加 `AIPERF` suite schema、case 编译和有界 dispatch，不承载子进程细节 | +| `test_runner_runtime_backend.py` | 增加 `AIPERF` prepare、execute、finalize 和终态读取 | +| `aiperf_adapter.py` | 唯一的直接进程入口，负责 AIPerf 子进程、loopback API、timeout 和结果发布 | +| `aiperf_contract.py` | generated config、progress 和 result 的强类型模型与 validator，不提供第二个执行入口 | +| 现有 runner UI 代码 | 在 `/api/run_state` 与 run 页面增加 `AIPERF` 显式分支 | + +### 5.2 数据流 + +```mermaid +flowchart LR + A[suite] --> B[test_runner compile] + B --> C[resolved_case.yaml] + B --> D[aiperf.generated.yaml] + B --> E[ops apply on load generator] + E --> F[AIPerf adapter] + F --> G[aiperf profile] + G --> H[SGLang OpenAI endpoint] + H --> G + + G --> I[AIPerf artifacts] + G --> J[loopback API] + J --> F + F --> K[aiperf_progress.json] + F --> L[benchmark_result.json] + L --> B + B --> M[summary.yaml + case_runs.yaml] + + H --> N[SGLang /metrics] + N --> G + O[Fluxon services] --> P[Greptime metrics + logs] + + Q[test_runner UI/API] --> M + Q --> K + Q --> I + Q --> P + Q --> R[ops status + logs] +``` + +AIPerf adapter 是单次 case 的直接进程入口，生命周期与该 case 一致。它不作为 testbed 常驻服务运行。 + +### 5.3 所有权与事实来源 + +| 数据 | 事实来源 | 保留周期 | UI 读取方式 | +| --- | --- | --- | --- | +| suite/case 状态 | `case_runs.yaml` | suite workdir 生命周期 | runner 本地读取 | +| run 终态 | `summary.yaml` | run artifact 生命周期 | runner 本地读取 | +| AIPerf 请求摘要 | `aiperf/profile_export_aiperf.json` | run artifact 生命周期 | runner 结果适配器 | +| AIPerf 请求级记录 | `aiperf/profile_export.jsonl` | run artifact 生命周期 | 按需下载或离线分析 | +| AIPerf 实时进度 | `aiperf_progress.json` | case 运行期间，终态后保留最后快照 | 通过 ops 文件读取接口 | +| SGLang run-window 指标 | `aiperf/server_metrics_export.jsonl` | run artifact 生命周期 | runner 图表适配器 | +| Fluxon 连续指标 | Greptime Prometheus-compatible API | Greptime retention | runner monitor 查询接口 | +| 服务与进程日志 | 本地 daily shard、Greptime `fluxon_logs`、run log | 各自既有 retention | 现有 log/ops log API | + +首期不把 AIPerf summary 再复制到 Greptime。这样可以避免 artifact、OTel 和 Prometheus 三条通道同时保存同一组请求指标。若后续需要跨 run 长期聚合，应单独设计一个由 runner 控制的离线导入协议。 + +## 6. 配置与编译模型 + +### 6.1 suite 只增加一个有限模式 + +目标 suite 结构采用 `scene.test_stack.mode: AIPERF`。各层仍保持现有分工： + +| 层 | `AIPERF` 分支内容 | +| --- | --- | +| scene | model、endpoint type、streaming、dataset、warmup、profiling phase、随机种子 | +| scale | load generator target、SGLang endpoint instance、资源规模 | +| profile | AIPerf artifact set、adapter deploy 模板、SGLang runtime 组合 | +| resolved case | 具体 URL、metrics URL、run_dir、loopback API port、完整 source manifest | + +下面是提议中的 suite 片段，只用于说明字段归属，当前代码还不能执行： + +```yaml +scenes: + sglang_chat_concurrency: + test_stack: + mode: AIPERF + aiperf: + model: example-model + endpoint_type: chat + streaming: true + random_seed: 42 + dataset: + type: synthetic + entries: 512 + input_tokens: 1024 + output_tokens: 256 + warmup: + concurrency: 8 + requests: 32 + profiling: + type: concurrency + concurrency: 32 + requests: 512 +``` + +Fluxon 只接受上述有界字段，不接受 `extra_args`、任意 AIPerf config fragment 或环境变量模板。需要新增 workload 能力时，先把它加入明确的 schema 分支和 contract test。 + +### 6.2 生成单一 AIPerf 配置 + +runner 根据 `resolved_case` 生成 `aiperf/aiperf.generated.yaml`。该文件是编译产物，用户不直接维护。Fluxon suite 统一使用 `snake_case`；生成器在 AIPerf 边界按该提交自带 JSON Schema 输出 `camelCase` alias，不在同一份配置中混用两种拼写。 + +```yaml +schemaVersion: "2.0" +randomSeed: 42 + +benchmark: + model: example-model + endpoint: + url: http://sglang-host:30000/v1/chat/completions + type: chat + streaming: true + headers: + X-Fluxon-Case-ID: sglang_chat_concurrency__n1__teleai + X-Fluxon-Run-Index: "1" + dataset: + type: synthetic + entries: 512 + prompts: {isl: 1024, osl: 256} + warmup: + type: concurrency + concurrency: 8 + requests: 32 + profiling: + type: concurrency + concurrency: 32 + requests: 512 + artifacts: + dir: /testbed/run/results/example/run_1/aiperf + summary: [json] + records: [jsonl] + raw: false + trace: false + sliceDuration: 5 + serverMetrics: + enabled: true + urls: + - http://sglang-host:30000/metrics + formats: [json, jsonl] + discovery: + mode: disabled + gpuTelemetry: + enabled: false + runtime: + ui: none + apiHost: 127.0.0.1 + apiPort: 19081 +``` + +启动命令固定为： + +```bash +aiperf profile --config /testbed/run/results/example/run_1/aiperf/aiperf.generated.yaml +``` + +`testbed_aiperf_api_port` 由现有 testbed port allocator 派生并写入 resolved case，不增加用户侧端口配置项。 + +### 6.3 依赖交付 + +构建阶段完成以下动作： + +1. 从固定 commit 构建 AIPerf wheel。 +2. 解析并下载目标平台的完整 dependency wheelhouse。 +3. 生成包含文件名、版本、许可证和 SHA-256 的 manifest。 +4. 把 wheelhouse 作为 profile 所选 artifact set 的一部分发布。 +5. load generator 在隔离 venv 中离线安装，禁止运行时访问 package index。 + +`resolved_case.yaml` 和 `benchmark_result.json` 都要记录 source repo、requested branch、resolved commit、AIPerf package version 与 wheelhouse manifest SHA-256。 + +## 7. 执行与终态契约 + +### 7.1 执行时序 + +```mermaid +sequenceDiagram + participant R as test_runner + participant O as ops/deployer + participant A as AIPerf adapter + participant P as aiperf profile + participant S as SGLang + + R->>O: apply adapter workload + O->>A: start generated config + A->>P: spawn direct process + P->>S: warmup requests + P->>S: profiling requests + loop every 1 second while running + A->>P: GET loopback /api/progress and /api/metrics + A->>A: atomic replace aiperf_progress.json + end + P-->>A: exit code + artifacts + A->>A: validate and write benchmark_result.json + R->>O: read terminal result + R->>R: validate, write summary, finalize apply +``` + +### 7.2 `benchmark_result.json` + +`AIPERF` 使用独立的强类型 payload，runner 根据已编译的 mode 选择专用 validator。建议最小结构如下： + +```json +{ + "schema_version": 1, + "result_kind": "AIPERF", + "case_id": "sglang_chat_concurrency__n1__teleai", + "run_index": 1, + "completion": { + "status": "SUCCESS", + "exit_code": 0, + "error": null + }, + "source": { + "repo": "https://github.com/ActivePeter/aiperf", + "requested_branch": "teleai", + "resolved_commit": "d72160e20957013d6608afcc88ed24100cb27dc5", + "package_version": "", + "wheelhouse_manifest_sha256": "" + }, + "timing": { + "started_at_unix_ns": 0, + "finished_at_unix_ns": 0 + }, + "summary": { + "request_count": 0, + "error_request_count": 0, + "metrics": {} + }, + "artifacts": [] +} +``` + +`metrics` 中保留 AIPerf 的 canonical tag、统计字段和单位。适配器不把 `time_to_first_token` 改名为另一套 Fluxon 私有名称。 + +### 7.3 完成状态 + +完成状态使用有限枚举： + +| 状态 | 含义 | runner outcome | +| --- | --- | --- | +| `SUCCESS` | 进程退出码为 0，必需 artifact 可解析，至少有一个成功请求 | `SUCCESS` | +| `START_FAILED` | venv、配置校验或子进程启动失败 | `FAILED` | +| `PROCESS_FAILED` | AIPerf 非零退出 | `FAILED` | +| `RESULT_INVALID` | summary 缺失、schema 不符或计数不满足不变量 | `FAILED` | +| `TIMEOUT` | 超过 case deadline，adapter 已终止子进程 | `FAILED` | +| `CANCELLED` | runner 或 operator 发出取消 | `FAILED` | + +首期不把普通请求错误率直接转换为 case 失败。请求错误会进入 summary；后续性能门禁必须使用单独、显式、可版本化的 acceptance policy。 + +### 7.4 不变量 + +- `case_id`、`run_index` 必须与 `resolved_case.yaml` 完全一致。 +- `resolved_commit` 和 wheelhouse manifest hash 必须与 staged artifact manifest 一致。 +- `finished_at_unix_ns >= started_at_unix_ns`。 +- `SUCCESS` 要求 AIPerf exit code 为 0、summary JSON 可解析、`request_count > 0`，并且 `request_count - error_request_count > 0`。 +- 所有 artifact path 必须位于当前 run_dir 内，并记录 size 与 SHA-256。 +- adapter 只能用临时文件加原子 rename 发布进度和终态，runner 不读取半写文件。 +- timeout 或 cancel 后必须等待 AIPerf 进程组退出，再写终态。 + +## 8. 监控与 UI 设计 + +### 8.1 指标分域 + +| 域 | 代表指标 | 来源 | 首期展示 | +| --- | --- | --- | --- | +| 生命周期 | status、phase、completed/total requests、elapsed | adapter progress | run 状态卡 | +| 请求体验 | TTFT、ITL、request latency、error rate | AIPerf summary/timeslice | percentile 表与时间曲线 | +| 吞吐 | request throughput、input/output token throughput | AIPerf | 摘要与时间曲线；goodput 随阶段 3 的显式 SLO 配置加入 | +| SUT | queue depth、KV cache usage、running/waiting requests | AIPerf server metrics artifact | 与请求曲线共享时间轴 | +| Fluxon 基础设施 | transport、KV、MQ、FS 和服务指标 | Greptime | 按 run 时间窗口查询 | +| 日志 | adapter、AIPerf、SGLang、Fluxon 服务日志 | ops、本地日志、Greptime | 复用现有日志查看器 | + +请求级 `x_request_id` 和 `x_correlation_id` 只进入 record/log，不作为 Greptime 时序标签，避免高基数。 + +### 8.2 关联字段 + +| 字段 | artifact | 请求 header | Greptime 查询标签 | 说明 | +| --- | --- | --- | --- | --- | +| `case_id` | 必需 | `X-Fluxon-Case-ID` | 不要求 | 逻辑 case，由 runner 关联到查询窗口 | +| `run_index` | 必需 | `X-Fluxon-Run-Index` | 不要求 | 同一 case 的第 N 次运行 | +| `cluster_name` | 必需 | 可选 | 必需 | testbed 集群 | +| `instance_key` | 必需 | 无 | 必需 | SUT 与 Fluxon 服务实例 | +| `model` | 必需 | 请求 body 已包含 | 可选 | 模型身份 | +| `endpoint_type` | 必需 | 无 | 可选 | 首期固定为 `chat` | +| `case_key` | 必需 | 无 | 不使用 | 配置快照 hash，不进入时序标签 | +| source commit/hash | 必需 | 无 | 不使用 | 保存在 provenance 中 | + +长生命周期服务不能随每个 case 动态改写 Prometheus label。runner 应在 resolved case 和终态中保存 `cluster_name`、相关 `instance_key`、`started_at_unix_ns` 与 `finished_at_unix_ns`，再用这组条件查询 Greptime。`case_id` 和 `run_index` 只需要出现在 run artifact、AIPerf 请求 header 及可选的 run marker 日志中。 + +跨 AIPerf 请求时间和 Greptime 时序数据做叠图前，testbed preflight 必须检查 load generator 与 SUT 时钟。首期允许的最大时钟偏差应固定为 1 秒，超过阈值时 case 在启动前失败。 + +### 8.3 runner API 与页面 + +扩展现有 `/api/run_state`，增加一个按 mode 判别的 `benchmark` 对象： + +```json +{ + "benchmark": { + "kind": "AIPERF", + "status": "RUNNING", + "progress": {}, + "summary": null, + "artifacts": [], + "monitor_window": { + "start_unix_ns": 0, + "end_unix_ns": null + } + } +} +``` + +不新增平行的 AIPerf UI service。run 页面增加以下区域： + +1. workload 与 provenance +2. warmup/profiling 进度 +3. TTFT、ITL、latency、throughput 和 error 摘要 +4. timeslice 与 SGLang server metrics 曲线 +5. Greptime 同时间窗口入口 +6. artifact 与 ops log 列表 + +UI 服务重启后，历史页只依赖 run artifact 和 Greptime，不能依赖已经退出的 AIPerf API。 + +## 9. 可复现性、性能边界与安全 + +### 9.1 可复现性 + +- 每个 scene 必须显式给出 `random_seed`、warmup、stop condition、ISL、OSL 和 streaming。 +- endpoint URL 从已解析的 SGLang instance 派生，不允许在 scene 中另写自由 URL。 +- load generator 与 SUT 首期必须位于不同 target，避免 generator CPU、网络和 SUT 资源相互竞争。 +- 记录 AIPerf commit、wheelhouse hash、Python 版本、SGLang image/commit、模型、tokenizer 和完整 generated config。 +- server metrics 始终使用相同的显式采样配置。AIPerf 当前 333 ms 采样行为本身也属于对比条件。 +- raw response 默认不保存；请求级 metrics JSONL 必须保存，以便复核 percentile 和异常请求。 + +### 9.2 性能结论边界 + +AIPerf 给出的请求指标覆盖 load generator 到 SGLang HTTP endpoint 的请求路径。它包含客户端排队、网络和服务端处理的组合影响，不能单独证明 Fluxon 某个内部模块的耗时。 + +SGLang `/metrics` 与 Greptime 指标用于解释同一运行窗口内的服务状态。两者采样频率、时钟和标签不同，叠图只提供时间相关性，不自动证明因果关系。 + +### 9.3 安全与数据保留 + +- AIPerf API 只监听 loopback。 +- generated config、result 和 artifact path 都要经过 run_dir containment 校验。 +- 首期只使用 synthetic dataset，不接入真实用户 prompt。 +- `raw: false`、`trace: false` 为固定默认值。 +- endpoint headers 在 result 和日志中必须经过 AIPerf 现有 redaction，再由 adapter 二次校验。 +- 后续鉴权方案必须使用专门的 secret 注入路径，不能把 token 写入 suite、generated config 或普通环境变量。 + +## 10. 实施阶段 + +### 阶段 0：依赖与样例固化 + +- 确认 `teleai` 快照就是目标版本，或取得明确的 TeleAI patch commit。 +- 构建离线 wheelhouse、license 清单和 SHA-256 manifest。 +- 用固定 SGLang mock/fixture 验证 AIPerf summary、records、server metrics 和 loopback API schema。 +- 保存一份脱敏 golden artifact，作为 adapter contract test 输入。 + +### 阶段 1：终态 MVP + +- 在 suite schema 中加入 `AIPERF` mode 和有限 workload 字段。 +- 增加 case 编译、testbed port 分配、generated YAML 和 isolated venv prepare。 +- 增加 AIPerf adapter 的 direct-process 启动、timeout、cancel、artifact 校验和终态写入。 +- 增加专用 result validator，并把小型摘要写入 `summary.yaml`。 +- UI 先展示终态摘要、artifact 和日志。 + +### 阶段 2：实时监控 + +- adapter 每秒读取 loopback `/api/progress` 和 `/api/metrics`，原子更新进度快照。 +- 扩展现有 `/api/run_state` 和 run 页面，不新增服务。 +- 增加 timeslice 与 SGLang server metrics 图表。 +- 用 run 时间窗口和关联字段查询 Greptime。 + +### 阶段 3：对比与门禁 + +- 由 `test_runner` 管理重复 run 和基线选择，禁止打开 AIPerf multi-run 形成第二层历史。 +- 定义版本化 acceptance policy，明确 metric、stat、方向、阈值和缺失值行为。 +- 对 percentile 门禁优先读取请求级记录的 pooled 统计，不使用节点 percentile 的平均值。 +- 增加跨 run 对比页和机器可读的 evaluation 结果。 + +## 11. 测试计划 + +| 测试 | 执行模型 | 成功条件 | +| --- | --- | --- | +| config compile | contract test | suite 字段唯一映射到 generated YAML，未知字段快速失败 | +| dependency provenance | contract test | commit、package version 和 wheelhouse hash 一致 | +| happy path | 直接启动 adapter 进程 + mock OpenAI endpoint | exit 0、终态 `SUCCESS`、必需 artifact 存在 | +| request failure | 直接启动 adapter 进程 + 失败响应 endpoint | 错误计数保留，至少一个成功请求时完成 | +| invalid result | 直接启动 adapter 进程并注入损坏 artifact | `RESULT_INVALID`，runner case 失败 | +| timeout | 直接启动 adapter 进程 + 挂起 endpoint | 整个进程组退出，终态 `TIMEOUT` | +| cancel | 运行中发送 runner cancel | 终态 `CANCELLED`，没有残留 AIPerf 进程 | +| API isolation | process test | API 只监听 loopback，testbed 外无法访问 | +| UI restart | 启动常驻 UI，完成 run 后重启 UI | 历史摘要和图表仍能从 artifact 恢复 | +| clock preflight | process test | 偏差超过 1 秒时在发压前失败 | +| offline install | clean testbed process | 无外网条件下从 wheelhouse 完成安装和运行 | + +进程生命周期测试按独立脚本或进程直接运行，并显式检查 exit code 与残留进程。不要为了统一外观再套一层 pytest 入口。 + +## 12. 验收条件 + +首期完成必须同时满足： + +1. 用户仍只执行 `start testbed` 和 `testrunner`。 +2. suite 只有一个 `AIPERF` 配置入口，运行时只消费生成的单一 YAML。 +3. testbed 无外网时可以安装并运行固定 AIPerf 快照。 +4. runner 能区分启动失败、进程失败、结果损坏、超时和取消。 +5. UI 不公开 AIPerf API，run 结束和 UI 重启后仍能展示结果。 +6. TTFT、ITL、latency、throughput 和 error 指标可追溯到 AIPerf artifact，不经二次重算。 +7. AIPerf 请求窗口可以和 SGLang server metrics、Fluxon Greptime 指标及日志按 `case_id + run_index + time range` 对齐。 +8. 现有 KV、MQ、RPC、FS benchmark 路径和结果校验保持不变。 diff --git a/fluxon_py/__init__.py b/fluxon_py/__init__.py index 753281c..7c57441 100644 --- a/fluxon_py/__init__.py +++ b/fluxon_py/__init__.py @@ -66,6 +66,9 @@ "KvFuture", "MemHolder", "FluxonMemHolder", + "GpuBufferRegistration", + "GpuDestination", + "GpuGetStartHandle", # Backend management "KvClientType", "new_store", @@ -148,6 +151,9 @@ _LAZY_PYO3 = { "FluxonMemHolder": ("kvclient.fluxon", "FluxonMemHolder"), + "GpuBufferRegistration": ("kvclient.fluxon", "GpuBufferRegistration"), + "GpuDestination": ("kvclient.fluxon", "GpuDestination"), + "GpuGetStartHandle": ("kvclient.fluxon", "GpuGetStartHandle"), } diff --git a/fluxon_py/config.py b/fluxon_py/config.py index 5861f64..7340d5f 100644 --- a/fluxon_py/config.py +++ b/fluxon_py/config.py @@ -66,6 +66,7 @@ def _yaml_template(): protocol_type: # Protocol type (('tcp'|'rdma')) rdma_device_names: # Explicit RDMA devices for protocol config (['{str}'](optional)) pprof_duration_seconds: # Dump pprof flamegraph after N seconds (int(optional)) +replica_writeback_hot_capacity_ratio: 0.75 # Owner-local hot working-set ratio in the open interval zero to one (float(optional)) contribute_to_cluster_pool_size: # Capacity contributed to cluster pool (dict(optional)) dram: 1677721600 # - DRAM contribution (int(multiple of 16777216)) vram: # - VRAM contribution per GPU (dict(dynamic_key)) @@ -82,11 +83,18 @@ def _yaml_template(): prefer_local_placement: false # Prefer placing new KV writes on the requester-local owner when possible (bool(optional)) short_circuit_put_payload_path: false # Keep large put_start allocation but skip payload memcpy + transfer (bool(optional)) skip_put_end_commit: false # Return success after payload transfer without put_done commit; inflight_put TTL cleanup only (bool(optional)) + ssd_read_source_policy: legacy_remote_first # legacy_remote_first|local_ssd_only_first (str(optional)) + owner_local_reserve_soft_wait_timeout_ms: # Local-reserve polling interval, >0 (int(optional)) + owner_local_reserve_hard_timeout_ms: # Local-reserve end-to-end claim timeout, > soft wait (int(optional)) + owner_local_reserve_expected_capacity: # Owner-only local-reserve prewarm target (dict(optional)) + value_len: # Canonical value payload bytes, >0 and <=512 MiB (int) + payload_capacity_bytes: # Expected payload capacity to keep resident, >0 (int) transport_mode: # transfer_only|transfer_with_rpc (str(optional)) tcp_thread_reactor_shard_count: # tcp_thread reactor shard count, 1..16 (int(optional)) tcp_thread_bulk_lane_count: # tcp_thread bulk lane count, 1..8 (int(optional)) tcp_thread_control_lane_count: # tcp_thread control lane count, 1..8 (int(optional)) user_rpc_sync_handler_thread_count: # Owner-dedicated sync user-RPC worker thread count, >0 (int(optional)) + replica_task_max_inflight: # Deprecated compatibility field, still validated as 1..64; direct remote-Put singleflight does not use a global actor queue (int(optional)) require_transfer_rpc_fast_path_ready_timeout_seconds: # Require owner-owner transfer-rpc fast path before owner ready/shared.json publication (int(optional)) rdma_device_names: # Explicit RDMA devices for benchmark/test fast-path fanout (['{str}'](optional)) enable_side_transfer: false # Enable TCP side-transfer fast-path (bool(optional)) @@ -110,6 +118,9 @@ def _yaml_template(): cluster_name: # Cluster name (str) share_mem_path: # Shared bundle path for mmap.file/shared.json/peer metadata (str) large_file_paths: # Owner-mode ordered large-file roots (['{str}'](optional)) + large_limit_size: # Optional per-root SSD capacity in bytes (list(optional)) + ssd_write_rate_limit_bytes_per_sec: # Optional non-queueing SSD write rate (int(optional)) + ssd_write_burst_bytes: # Paired immediate SSD write burst (int(optional)) p2p_listen_port: # P2P QUIC listen port override (int(optional)) redis_compat: # Enable Redis protocol shim (dict(optional)) listen_addr: # TCP listen addr, e.g. "127.0.0.1:16379" (str) @@ -135,11 +146,16 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: "prefer_local_placement", "short_circuit_put_payload_path", "skip_put_end_commit", + "ssd_read_source_policy", + "owner_local_reserve_soft_wait_timeout_ms", + "owner_local_reserve_hard_timeout_ms", + "owner_local_reserve_expected_capacity", "transport_mode", "tcp_thread_reactor_shard_count", "tcp_thread_bulk_lane_count", "tcp_thread_control_lane_count", "user_rpc_sync_handler_thread_count", + "replica_task_max_inflight", "require_transfer_rpc_fast_path_ready_timeout_seconds", "rdma_device_names", "enable_side_transfer", @@ -172,6 +188,21 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: raise ValueError(f"{ctx}.{key} must be a bool") out[key] = value + ssd_read_source_policy = raw.get("ssd_read_source_policy") + if ssd_read_source_policy is not None: + if not isinstance(ssd_read_source_policy, str): + raise ValueError(f"{ctx}.ssd_read_source_policy must be a string") + allowed_ssd_read_source_policies = { + "legacy_remote_first", + "local_ssd_only_first", + } + if ssd_read_source_policy not in allowed_ssd_read_source_policies: + raise ValueError( + f"{ctx}.ssd_read_source_policy must be one of " + f"{sorted(allowed_ssd_read_source_policies)}, got {ssd_read_source_policy!r}" + ) + out["ssd_read_source_policy"] = ssd_read_source_policy + transport_mode = raw.get("transport_mode") transport_mode_was_explicit = transport_mode is not None side_transfer_role_raw = raw.get("side_transfer_role") @@ -234,6 +265,7 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: ("tcp_thread_reactor_shard_count", 1, 16), ("tcp_thread_bulk_lane_count", 1, 8), ("tcp_thread_control_lane_count", 1, 8), + ("replica_task_max_inflight", 1, 64), ): value = raw.get(key) if value is None: @@ -254,6 +286,50 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]: raise ValueError(f"{ctx}.user_rpc_sync_handler_thread_count must be > 0") out["user_rpc_sync_handler_thread_count"] = user_rpc_sync_handler_thread_count + reserve_timeouts: Dict[str, int] = {} + for key in ( + "owner_local_reserve_soft_wait_timeout_ms", + "owner_local_reserve_hard_timeout_ms", + ): + value = raw.get(key) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{ctx}.{key} must be an int") + if value <= 0: + raise ValueError(f"{ctx}.{key} must be > 0") + reserve_timeouts[key] = value + out[key] = value + soft_timeout_ms = reserve_timeouts.get("owner_local_reserve_soft_wait_timeout_ms", 10) + hard_timeout_ms = reserve_timeouts.get("owner_local_reserve_hard_timeout_ms", 10_000) + if hard_timeout_ms <= soft_timeout_ms: + raise ValueError( + f"{ctx}.owner_local_reserve_hard_timeout_ms must be greater than " + f"{ctx}.owner_local_reserve_soft_wait_timeout_ms" + ) + + expected_capacity = raw.get("owner_local_reserve_expected_capacity") + if expected_capacity is not None: + expected_ctx = f"{ctx}.owner_local_reserve_expected_capacity" + if not isinstance(expected_capacity, dict): + raise ValueError(f"{expected_ctx} must be a mapping") + unknown_expected = sorted( + set(expected_capacity.keys()) - {"value_len", "payload_capacity_bytes"} + ) + if unknown_expected: + raise ValueError(f"{expected_ctx} contains unknown keys: {unknown_expected}") + normalized_expected: Dict[str, int] = {} + for key in ("value_len", "payload_capacity_bytes"): + value = expected_capacity.get(key) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{expected_ctx}.{key} must be an int") + if value <= 0: + raise ValueError(f"{expected_ctx}.{key} must be > 0") + normalized_expected[key] = value + if normalized_expected["value_len"] > 512 * 1024 * 1024: + raise ValueError(f"{expected_ctx}.value_len must be <= 536870912") + out["owner_local_reserve_expected_capacity"] = normalized_expected + side_transfer_worker_count = raw.get("side_transfer_worker_count") if side_transfer_worker_count is not None: if isinstance(side_transfer_worker_count, bool) or not isinstance(side_transfer_worker_count, int): @@ -348,6 +424,23 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: raise ValueError("fluxonkv_spec must be a mapping") is_zero_contribution = _is_zero_contribution_fluxonkv_config(cfg) + test_spec_config = cfg.get("test_spec_config") or {} + expected_capacity = test_spec_config.get("owner_local_reserve_expected_capacity") + hot_capacity_ratio = cfg.get("replica_writeback_hot_capacity_ratio") + + if hot_capacity_ratio is not None: + if isinstance(hot_capacity_ratio, bool) or not isinstance( + hot_capacity_ratio, (int, float) + ): + raise ValueError( + "replica_writeback_hot_capacity_ratio must be a number in (0, 1)" + ) + hot_capacity_ratio = float(hot_capacity_ratio) + if not 0.0 < hot_capacity_ratio < 1.0: + raise ValueError( + "replica_writeback_hot_capacity_ratio must be finite and in (0, 1)" + ) + cfg["replica_writeback_hot_capacity_ratio"] = hot_capacity_ratio share_mem_path = spec.get("share_mem_path") if not isinstance(share_mem_path, str) or not share_mem_path.strip(): @@ -360,11 +453,22 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: raise ValueError("fluxonkv_spec.transfer_engine has been removed from Fluxon KV config") if is_zero_contribution: + if hot_capacity_ratio is not None: + raise ValueError( + "replica_writeback_hot_capacity_ratio is only valid on owner configs" + ) + if expected_capacity is not None: + raise ValueError( + "test_spec_config.owner_local_reserve_expected_capacity is only valid on owner configs" + ) forbidden_spec_keys = [ "etcd_addresses", "redis_compat", "sub_cluster", "large_file_paths", + "large_limit_size", + "ssd_write_rate_limit_bytes_per_sec", + "ssd_write_burst_bytes", ] for key in forbidden_spec_keys: if key in spec: @@ -379,6 +483,23 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: if int(contrib["dram"]) == 0: raise ValueError("owner mode requires non-zero contribute_to_cluster_pool_size.dram") + if expected_capacity is not None: + value_len = int(expected_capacity["value_len"]) + payload_capacity_bytes = int(expected_capacity["payload_capacity_bytes"]) + slot_size = max(value_len, 4096) + slot_size = (slot_size + 4095) & ~4095 + slots_per_grant = (512 * 1024 * 1024) // slot_size + value_count = (payload_capacity_bytes + value_len - 1) // value_len + expected_grants = (value_count + slots_per_grant - 1) // slots_per_grant + physical_reserve_bytes = expected_grants * 512 * 1024 * 1024 + owner_dram_bytes = int(contrib["dram"]) + if physical_reserve_bytes > owner_dram_bytes: + raise ValueError( + "test_spec_config.owner_local_reserve_expected_capacity requires " + f"{physical_reserve_bytes} physical bytes across {expected_grants} grants, " + f"exceeding owner dram contribution {owner_dram_bytes}" + ) + if "etcd_addresses" not in spec: raise ValueError("fluxonkv_spec.etcd_addresses is required for owner mode") etcd_addresses = spec.get("etcd_addresses") @@ -404,6 +525,26 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None: f"fluxonkv_spec.large_file_paths[{idx}] must be a non-empty string in owner mode" ) + write_rate = spec.get("ssd_write_rate_limit_bytes_per_sec") + write_burst = spec.get("ssd_write_burst_bytes") + if (write_rate is None) != (write_burst is None): + raise ValueError( + "fluxonkv_spec.ssd_write_rate_limit_bytes_per_sec and " + "ssd_write_burst_bytes must be configured together" + ) + if write_rate is not None: + if ( + isinstance(write_rate, bool) + or not isinstance(write_rate, int) + or write_rate <= 0 + or isinstance(write_burst, bool) + or not isinstance(write_burst, int) + or write_burst <= 0 + ): + raise ValueError("SSD write rate and burst must both be positive integers") + if "large_limit_size" not in spec: + raise ValueError("SSD write admission requires fluxonkv_spec.large_limit_size") + class FluxonKvClientConfig(): """Configuration class for KV Cache stores that reads from YAML config files.""" @@ -434,6 +575,10 @@ def __init__(self, config_dict: Dict[str, Any]): raise ValueError( "exactly one of [mooncake_spec, fluxonkv_spec] is required (and the chosen spec must not be null)" ) + if "replica_writeback_hot_capacity_ratio" in plain and not has_fluxon: + raise ValueError( + "replica_writeback_hot_capacity_ratio requires fluxonkv_spec" + ) pprof_duration_seconds = plain.get("pprof_duration_seconds") if pprof_duration_seconds is None: @@ -750,8 +895,8 @@ def parse_type_recursive(type_str: str) -> Optional[Tuple[str, Dict[str, Any]]]: return parsed_type_name, merged_params return type_name, {"constraint": constraint} - # 6) Primitive types: str, int, bool, None - if type_str in ["str", "int", "bool", "None"]: + # 6) Primitive types: str, int, float, bool, list, None + if type_str in ["str", "int", "float", "bool", "list", "None"]: return type_str, {} debug_print("type_str ", type_str, "not matched to any type") @@ -871,13 +1016,29 @@ def raise_validation_error(msg: str): raise_validation_error(f"Value must be multiple of {multiple}, got {value}") else: return None - + + elif type_name == "float": + if isinstance(value, bool) or not isinstance(value, (int, float)): + if raise_err: + raise_validation_error( + f"Expected float-compatible number, got {type(value).__name__}" + ) + else: + return None + elif type_name == "bool": if not isinstance(value, bool): if raise_err: raise_validation_error(f"Expected bool, got {type(value).__name__}") else: return None + + elif type_name == "list": + if not isinstance(value, list): + if raise_err: + raise_validation_error(f"Expected list, got {type(value).__name__}") + else: + return None elif type_name == "dict": if not isinstance(value, dict): diff --git a/fluxon_py/kvclient/fluxon.py b/fluxon_py/kvclient/fluxon.py index 1325e3d..25f3776 100644 --- a/fluxon_py/kvclient/fluxon.py +++ b/fluxon_py/kvclient/fluxon.py @@ -3,6 +3,7 @@ This module provides a concrete implementation using the PyO3 Rust bindings. """ +from dataclasses import dataclass from typing import Union, Optional, Callable, Any, Dict, List, Tuple import ctypes import os @@ -25,7 +26,7 @@ from .kvclient_interface import KvClient from .kvclient_interface import KvLeaseApi, KvRpcApi, PutOptionalArgs, FlatDict from .backend_fallback_close import unregister_store_from_cleanup -from .kvclient_interface import KvFuture, MemHolder +from .kvclient_interface import GetStartHandle, GetStartResult, KvFuture, MemHolder from .nonzerocopy_encode import ( DLPacked, INTERNAL_DLPACK_META_KEY, @@ -54,6 +55,63 @@ _SIDE_TRANSFER_WORKER_PYTHON_ENV = "FLUXON_KV_SIDE_WORKER_PYTHON" _BLOCKING_PUT_OUTER_TOTAL_LOG_INTERVAL_NS = 10 * 1_000_000_000 +_MIN_EXPLICIT_RPC_TIMEOUT_MS = int(fluxon_pyo3.MIN_EXPLICIT_RPC_TIMEOUT_MS) + + +def _validate_explicit_rpc_timeout_ms(timeout_ms: int) -> None: + if not isinstance(timeout_ms, int): + raise InvalidArgumentError(message=f"timeout_ms must be int; got {type(timeout_ms)}") + if timeout_ms < _MIN_EXPLICIT_RPC_TIMEOUT_MS: + raise InvalidArgumentError( + message=( + f"timeout_ms must be >= {_MIN_EXPLICIT_RPC_TIMEOUT_MS}; " + f"got {timeout_ms}" + ) + ) + + +def _validate_put_atomic_groups( + keys: List[str], + atomic_group_lens: Optional[List[int]], + make_replica_task_mask: Optional[List[bool]], +) -> Optional[List[int]]: + if atomic_group_lens is None: + return None + if not isinstance(atomic_group_lens, list): + raise ValueError( + "atomic_group_lens must be list[int] when provided; " + f"got {type(atomic_group_lens)}" + ) + normalized_group_lens: List[int] = [] + for index, length in enumerate(atomic_group_lens): + if type(length) is not int: + raise ValueError( + "atomic_group_lens items must be int; " + f"index={index} got={type(length)}" + ) + if length <= 0: + raise ValueError( + "atomic_group_lens entries must be > 0; " + f"index={index} got={length}" + ) + normalized_group_lens.append(length) + group_sum = sum(normalized_group_lens) + if group_sum != len(keys): + raise ValueError( + "atomic_group_lens must sum to keys length; " + f"sum={group_sum} keys={len(keys)}" + ) + if make_replica_task_mask is not None: + offset = 0 + for group_index, group_len in enumerate(normalized_group_lens): + group_mask = make_replica_task_mask[offset : offset + group_len] + if any(item != group_mask[0] for item in group_mask[1:]): + raise ValueError( + "make_replica_task_mask must be uniform within each atomic group; " + f"group_index={group_index} offset={offset} len={group_len}" + ) + offset += group_len + return normalized_group_lens def _percentile_nearest_rank_ns(sorted_values: List[int], percentile: int) -> int: @@ -138,11 +196,233 @@ def _resolve_side_transfer_worker_python() -> str: return sys.executable +@dataclass(frozen=True) +class _RegisteredBufferDescriptor: + ptr: int + size: int + device_kind: str = "host" + device_id: str = "" + layout: str = "raw" + metadata: Optional[Dict[str, Any]] = None + + @property + def end(self) -> int: + return self.ptr + self.size + + def contains(self, ptr: int, size: int) -> bool: + req_end = ptr + size + return ptr >= self.ptr and req_end <= self.end + + def as_dict(self) -> Dict[str, Any]: + return { + "ptr": self.ptr, + "size": self.size, + "device_kind": self.device_kind, + "device_id": self.device_id, + "layout": self.layout, + "metadata": dict(self.metadata or {}), + } + + +@dataclass(frozen=True) +class GpuBufferRegistration: + registration_id: int + ptr: int + size: int + device_id: int + + def __post_init__(self) -> None: + if type(self.registration_id) is not int or self.registration_id <= 0: + raise ValueError("registration_id must be a positive int") + if type(self.ptr) is not int or self.ptr <= 0: + raise ValueError("GPU registration ptr must be a positive int") + if type(self.size) is not int or self.size <= 0: + raise ValueError("GPU registration size must be a positive int") + if type(self.device_id) is not int or self.device_id < 0: + raise ValueError("GPU device_id must be a non-negative int") + if self.end > (1 << 64): + raise ValueError("GPU registration range overflows u64") + + @property + def end(self) -> int: + return self.ptr + self.size + + def destination(self, ptr: int, capacity: int) -> "GpuDestination": + destination = GpuDestination( + registration_id=self.registration_id, + ptr=ptr, + capacity=capacity, + ) + if destination.ptr < self.ptr or destination.end > self.end: + raise ValueError( + "GPU destination is outside its registration: " + f"registration=[{self.ptr:#x},{self.end:#x}) " + f"destination=[{destination.ptr:#x},{destination.end:#x})" + ) + return destination + + +@dataclass(frozen=True) +class GpuDestination: + registration_id: int + ptr: int + capacity: int + + def __post_init__(self) -> None: + if type(self.registration_id) is not int or self.registration_id <= 0: + raise ValueError("registration_id must be a positive int") + if type(self.ptr) is not int or self.ptr <= 0: + raise ValueError("GPU destination ptr must be a positive int") + if type(self.capacity) is not int or self.capacity <= 0: + raise ValueError("GPU destination capacity must be a positive int") + if self.end > (1 << 64): + raise ValueError("GPU destination range overflows u64") + + @property + def end(self) -> int: + return self.ptr + self.capacity + + +@dataclass +class GpuGetStartHandle: + """One-shot handle for a background RDMA pull into GPU staging.""" + + keys: Tuple[str, ...] + destinations: Tuple[GpuDestination, ...] + result: GetStartResult + created_at_ns: int + backend_token: int + backend_handle: int + remote_indices: Tuple[int, ...] + closed: bool = False + transfer_wall_us: Optional[int] = None + finish_wait_us: Optional[int] = None + terminal_before_consume: Optional[bool] = None + terminal_to_consume_us: Optional[int] = None + + +@dataclass +class GetPlanHandle: + """Target-free Get plan that must be executed once or cancelled.""" + + keys: Tuple[str, ...] + result: GetStartResult + gpu_result: GetStartResult + created_at_ns: int + backend_token: int + backend_handle: int + gpu_remote_indices: Tuple[int, ...] + closed: bool = False + + def _map_nospace_to_storagefull(err: ApiError) -> ApiError: """Normalize storage-capacity errors without depending on backend internals.""" return err +def _get_start_prefix_hit_groups( + raw_prefix_hit_len: int, + group_lens: Tuple[int, ...], +) -> int: + prefix_hit_groups = 0 + transferable_len = 0 + for group_len in group_lens: + next_transferable_len = transferable_len + group_len + if next_transferable_len > raw_prefix_hit_len: + break + transferable_len = next_transferable_len + prefix_hit_groups += 1 + return prefix_hit_groups + + +def _get_start_group_index_for_key_index( + group_lens: Tuple[int, ...], + key_index: int, +) -> Optional[int]: + cursor = 0 + for group_index, group_len in enumerate(group_lens): + cursor += group_len + if key_index < cursor: + return group_index + return None + + +def _build_get_start_result_from_backend_payload( + payload: Dict[str, Any], + keys: List[str], + prefix_best_effort: bool, + normalized_group_lens: Optional[List[int]], +) -> GetStartResult: + result_keys = tuple(keys) + group_lens = ( + (len(result_keys),) + if normalized_group_lens is None + else tuple(normalized_group_lens) + ) + raw_prefix_hit_len = int(payload["raw_prefix_hit_len"]) + if raw_prefix_hit_len < 0 or raw_prefix_hit_len > len(result_keys): + raise RuntimeError( + "get_start returned invalid raw_prefix_hit_len: " + f"raw_prefix_hit_len={raw_prefix_hit_len} keys={len(result_keys)}" + ) + prefix_hit_groups = _get_start_prefix_hit_groups( + raw_prefix_hit_len, + group_lens, + ) + if not prefix_best_effort and prefix_hit_groups != len(group_lens): + transferable_len = 0 + prefix_hit_groups = 0 + else: + transferable_len = sum(group_lens[:prefix_hit_groups]) + first_miss_index = ( + None if raw_prefix_hit_len == len(result_keys) else raw_prefix_hit_len + ) + first_miss_group_index = ( + None + if first_miss_index is None + else _get_start_group_index_for_key_index(group_lens, first_miss_index) + ) + return GetStartResult( + keys=result_keys, + raw_prefix_hit_len=raw_prefix_hit_len, + transferable_len=transferable_len, + prefix_hit_groups=prefix_hit_groups, + atomic_group_lens=tuple(normalized_group_lens) if normalized_group_lens is not None else None, + prefix_best_effort=prefix_best_effort, + first_miss_index=first_miss_index, + first_miss_group_index=first_miss_group_index, + all_hit=transferable_len == len(result_keys), + ) + + +def _narrow_get_start_result(result: GetStartResult, consume_prefix_len: int) -> GetStartResult: + group_lens = result.atomic_group_lens or (len(result.keys),) + selected_groups: List[int] = [] + cursor = 0 + for group_len in group_lens: + if cursor + group_len > consume_prefix_len: + break + selected_groups.append(int(group_len)) + cursor += int(group_len) + if cursor != consume_prefix_len: + raise ValueError( + "consume_prefix_len must end at an atomic-group boundary: " + f"consume={consume_prefix_len} groups={group_lens}" + ) + keys = result.keys[:consume_prefix_len] + return GetStartResult( + keys=keys, + raw_prefix_hit_len=consume_prefix_len, + transferable_len=consume_prefix_len, + prefix_hit_groups=len(selected_groups), + atomic_group_lens=tuple(selected_groups), + prefix_best_effort=result.prefix_best_effort, + first_miss_index=None, + first_miss_group_index=None, + all_hit=True, + ) + + def _error_to_ret_code(err: ApiError) -> int: if hasattr(err, "code") and callable(err.code): try: @@ -299,7 +579,12 @@ def __init__(self, config: FluxonKvClientConfig): self._client: Optional[fluxon_pyo3.KvClient] = None self._config = config self._init_error: Optional[ApiError] = None + self._registered_buffer_descriptors: List[_RegisteredBufferDescriptor] = [] + self._gpu_buffer_registration: Optional[GpuBufferRegistration] = None cluster_name = config.fluxonkv_spec_cluster_name + self._batch_concurrency = 128 + if self._batch_concurrency <= 0: + raise ValueError("batch_concurrency must be > 0") self._blocking_put_outer_total_log_window = _BlockingPutOuterTotalLogWindow( f"FluxonKVCacheStore[{cluster_name}]" ) @@ -368,11 +653,21 @@ def put( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = ( + bool(opts.make_replica_task) if opts is not None else True + ) inner_res = self._client.put( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, ) if not inner_res.is_ok(): err = inner_res.unwrap_error() @@ -416,11 +711,21 @@ def put_blocking( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = ( + bool(opts.make_replica_task) if opts is not None else True + ) inner_res = self._client.put_blocking( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, ) if not inner_res.is_ok(): return Result.new_error(inner_res.unwrap_error()) @@ -448,6 +753,205 @@ def get_blocking(self, key: str) -> Result[MemHolder, ApiError]: except ApiError as e: return Result.new_error(e) + @staticmethod + def _normalize_batch_result_list(batch_result: Any, expected_len: int, op_name: str) -> List[Any]: + if isinstance(batch_result, Result): + if not batch_result.is_ok(): + raise RuntimeError(f"{op_name} backend error: {batch_result.unwrap_error()}") + batch_result = batch_result.unwrap() + if not isinstance(batch_result, list): + raise RuntimeError(f"{op_name} returned non-list: {type(batch_result)}") + if len(batch_result) != expected_len: + raise RuntimeError( + f"{op_name} returned unexpected length: expected={expected_len} got={len(batch_result)}" + ) + return list(batch_result) + + def batch_put_blocking( + self, + keys: List[str], + values: List[FlatDict], + opts: Optional[PutOptionalArgs] = None, + concurrency: Optional[int] = None, + ) -> List[Result[OkNone, ApiError]]: + if len(keys) != len(values): + raise ValueError("batch_put_blocking requires keys and values to have the same length") + if len(keys) == 0: + return [] + if self._client is None: + err = GeneralError(message="Store not initialized when batch_put_blocking(). Call setup() first.") + return [Result.new_error(err) for _ in keys] + + lease_id: Optional[int] = opts.lease_id if opts is not None else None + reject_if_inflight_same_key = ( + bool(opts.reject_if_inflight_same_key) if opts is not None else False + ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = bool(opts.make_replica_task) if opts is not None else True + + keepalive_groups: List[List[bytes]] = [] + dlpack_groups: List[List[object]] = [] + ptr_groups: List[List[tuple[int, int, int, int, int, Optional[int]]]] = [] + try: + for value in values: + keepalive: List[bytes] = [] + dlpack_capsules: List[object] = [] + ptr_groups.append(build_flat_dict_ptrs(value, keepalive, dlpack_capsules)) + keepalive_groups.append(keepalive) + dlpack_groups.append(dlpack_capsules) + + inner_res = self._client.batch_put_blocking( + keys, + ptr_groups, + lease_id=lease_id, + reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, + concurrency=concurrency if concurrency is not None else self._batch_concurrency, + ) + if not inner_res.is_ok(): + err = inner_res.unwrap_error() + return [Result.new_error(err) for _ in keys] + + batch_results = self._normalize_batch_result_list( + inner_res.unwrap(), len(keys), "batch_put_blocking" + ) + submit_out: List[Result[OkNone, ApiError]] = [] + for idx, item in enumerate(batch_results): + if isinstance(item, Result): + submit_out.append(item) + continue + if item is None: + submit_out.append(Result.new_ok(OkNone())) + continue + if isinstance(item, int) and item == 0: + submit_out.append(Result.new_ok(OkNone())) + continue + if isinstance(item, int): + submit_out.append( + Result.new_error( + GeneralError( + message=( + "batch_put_blocking returned backend code " + f"{item} for key {keys[idx]!r}" + ) + ) + ) + ) + continue + submit_out.append( + Result.new_error( + GeneralError( + message=f"unexpected batch_put result type: {type(item)}" + ) + ) + ) + return submit_out + except ApiError as e: + return [Result.new_error(e) for _ in keys] + except Exception as e: + return [Result.new_error(GeneralError(f"batch_put_blocking failed: {e}")) for _ in keys] + finally: + for keepalive in keepalive_groups: + keepalive.clear() + for dlpack_capsules in dlpack_groups: + dlpack_capsules.clear() + + def batch_get_blocking( + self, + keys: List[str], + concurrency: Optional[int] = None, + ) -> List[Result[Union[Any, MemHolder], ApiError]]: + if len(keys) == 0: + return [] + if self._client is None: + err = GeneralError(message="Store not initialized when batch_get_blocking(). Call setup() first.") + return [Result.new_error(err) for _ in keys] + + try: + inner_res = self._client.batch_get_blocking( + keys, + concurrency=concurrency if concurrency is not None else self._batch_concurrency, + ) + if not inner_res.is_ok(): + err = inner_res.unwrap_error() + return [Result.new_error(err) for _ in keys] + + batch_results = self._normalize_batch_result_list( + inner_res.unwrap(), len(keys), "batch_get_blocking" + ) + out: List[Result[Union[Any, MemHolder], ApiError]] = [] + for idx, item in enumerate(batch_results): + if isinstance(item, Result): + out.append(item) + continue + if item is None: + out.append( + Result.new_error( + GeneralError(message=f"batch_get_blocking returned None for key {keys[idx]!r}") + ) + ) + continue + out.append(Result.new_ok(item)) + return out + except ApiError as e: + return [Result.new_error(e) for _ in keys] + except Exception as e: + return [Result.new_error(GeneralError(f"batch_get_blocking failed: {e}")) for _ in keys] + + def put_payload_from_ptr( + self, + key: str, + payload_ptr: int, + payload_size: int, + opts: Optional[PutOptionalArgs] = None, + ) -> Result[KvFuture, ApiError]: + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when put_payload_from_ptr(). Call setup() first." + ) + ) + + keepalive: List[bytes] = [] + try: + ptrs = _build_payload_field_ptrs(payload_ptr, payload_size, keepalive) + lease_id: Optional[int] = opts.lease_id if opts is not None else None + reject_if_inflight_same_key = ( + bool(opts.reject_if_inflight_same_key) if opts is not None else False + ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = ( + bool(opts.make_replica_task) if opts is not None else True + ) + inner_res = self._client.put( + key, + ptrs, + lease_id=lease_id, + reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, + ) + if not inner_res.is_ok(): + return Result.new_error(_map_nospace_to_storagefull(inner_res.unwrap_error())) + inner_future = inner_res.unwrap() + assert inner_future is not None + outer_future = _FluxonPutFuture(inner_future, keepalive, []) + keepalive = [] + return Result.new_ok(outer_future) + except ApiError as e: + return Result.new_error(e) + finally: + keepalive.clear() + def put_payload_from_ptr_blocking( self, key: str, @@ -469,11 +973,21 @@ def put_payload_from_ptr_blocking( reject_if_inflight_same_key = ( bool(opts.reject_if_inflight_same_key) if opts is not None else False ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = ( + bool(opts.make_replica_task) if opts is not None else True + ) inner_res = self._client.put_blocking( key, ptrs, lease_id=lease_id, reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, ) if not inner_res.is_ok(): return Result.new_error(inner_res.unwrap_error()) @@ -566,17 +1080,466 @@ def get_payload_into_ptr_blocking( del holder return Result.new_ok(payload_size) + def register_buffer( + self, + ptr: int, + size: int, + device_kind: str = "host", + device_id: str = "", + layout: str = "raw", + metadata: Optional[Dict[str, Any]] = None, + ) -> Result[OkNone, ApiError]: + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when register_buffer(). Call setup() first." + ) + ) + if not isinstance(ptr, int): + return Result.new_error( + InvalidArgumentError(message=f"ptr must be int; got {type(ptr)}") + ) + if not isinstance(size, int): + return Result.new_error( + InvalidArgumentError(message=f"size must be int; got {type(size)}") + ) + if ptr < 0: + return Result.new_error( + InvalidArgumentError(message=f"ptr must be >= 0; got {ptr}") + ) + if size < 0: + return Result.new_error( + InvalidArgumentError(message=f"size must be >= 0; got {size}") + ) + if not isinstance(device_kind, str): + return Result.new_error( + InvalidArgumentError( + message=f"device_kind must be str; got {type(device_kind)}" + ) + ) + if device_kind.strip().lower() != "host": + return Result.new_error( + InvalidArgumentError( + message=( + "register_buffer supports host memory only; " + "use register_gpu_buffer for a CUDA staging range" + ) + ) + ) + if not isinstance(device_id, str): + return Result.new_error( + InvalidArgumentError( + message=f"device_id must be str; got {type(device_id)}" + ) + ) + if not isinstance(layout, str): + return Result.new_error( + InvalidArgumentError( + message=f"layout must be str; got {type(layout)}" + ) + ) + if metadata is not None and not isinstance(metadata, dict): + return Result.new_error( + InvalidArgumentError( + message=f"metadata must be dict or None; got {type(metadata)}" + ) + ) + try: + inner_res = self._client.register_buffer(ptr, size) + if not inner_res.is_ok(): + return Result.new_error(inner_res.unwrap_error()) + _ = inner_res.unwrap() + self._registered_buffer_descriptors.append( + _RegisteredBufferDescriptor( + ptr=int(ptr), + size=int(size), + device_kind=device_kind, + device_id=device_id, + layout=layout, + metadata=dict(metadata or {}), + ) + ) + return Result.new_ok(OkNone()) + except ApiError as e: + return Result.new_error(e) + + def register_gpu_buffer( + self, + ptr: int, + size: int, + device_id: int, + ) -> Result[GpuBufferRegistration, ApiError]: + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when register_gpu_buffer(). Call setup() first." + ) + ) + if isinstance(ptr, bool) or not isinstance(ptr, int) or ptr <= 0: + return Result.new_error( + InvalidArgumentError(message=f"GPU ptr must be a positive int; got {ptr!r}") + ) + if isinstance(size, bool) or not isinstance(size, int) or size <= 0: + return Result.new_error( + InvalidArgumentError(message=f"GPU size must be a positive int; got {size!r}") + ) + if ( + isinstance(device_id, bool) + or not isinstance(device_id, int) + or device_id < 0 + ): + return Result.new_error( + InvalidArgumentError( + message=f"GPU device_id must be a non-negative int; got {device_id!r}" + ) + ) + if ptr + size > (1 << 64): + return Result.new_error( + InvalidArgumentError( + message=f"GPU registration range overflows u64: ptr={ptr:#x} size={size}" + ) + ) + if self._gpu_buffer_registration is not None: + return Result.new_error( + InvalidArgumentError( + message=( + "one GPU buffer is already registered: " + f"registration_id={self._gpu_buffer_registration.registration_id}" + ) + ) + ) + try: + inner_res = self._client.register_gpu_buffer(ptr, size, device_id) + if not inner_res.is_ok(): + return Result.new_error(inner_res.unwrap_error()) + payload = inner_res.unwrap() + if not isinstance(payload, dict): + return Result.new_error( + GeneralError( + message=f"register_gpu_buffer returned non-dict payload: {type(payload)}" + ) + ) + registration = GpuBufferRegistration( + registration_id=int(payload["registration_id"]), + ptr=int(payload["ptr"]), + size=int(payload["len"]), + device_id=int(payload["device_id"]), + ) + if ( + registration.registration_id <= 0 + or registration.ptr != ptr + or registration.size != size + or registration.device_id != device_id + ): + return Result.new_error( + GeneralError( + message=( + "register_gpu_buffer returned a mismatched registration: " + f"{registration!r}" + ) + ) + ) + self._gpu_buffer_registration = registration + return Result.new_ok(registration) + except ApiError as e: + return Result.new_error(e) + + def validate_gpu_destination( + self, + destination: GpuDestination, + ) -> Result[OkNone, ApiError]: + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when validate_gpu_destination()." + ) + ) + if not isinstance(destination, GpuDestination): + return Result.new_error( + InvalidArgumentError( + message=( + "validate_gpu_destination requires GpuDestination; " + f"got {type(destination)}" + ) + ) + ) + if destination.capacity <= 0 or destination.ptr <= 0: + return Result.new_error( + InvalidArgumentError( + message=f"invalid GPU destination: {destination!r}" + ) + ) + try: + inner_res = self._client.validate_gpu_destination( + destination.registration_id, + destination.ptr, + destination.capacity, + ) + if not inner_res.is_ok(): + return Result.new_error(inner_res.unwrap_error()) + _ = inner_res.unwrap() + return Result.new_ok(OkNone()) + except ApiError as e: + return Result.new_error(e) + + def unregister_gpu_buffer( + self, + registration: GpuBufferRegistration, + ) -> Result[OkNone, ApiError]: + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when unregister_gpu_buffer()." + ) + ) + if not isinstance(registration, GpuBufferRegistration): + return Result.new_error( + InvalidArgumentError( + message=( + "unregister_gpu_buffer requires GpuBufferRegistration; " + f"got {type(registration)}" + ) + ) + ) + active = self._gpu_buffer_registration + if active != registration: + return Result.new_error( + InvalidArgumentError( + message=( + "GPU registration is not active on this store: " + f"requested={registration!r} active={active!r}" + ) + ) + ) + try: + inner_res = self._client.unregister_gpu_buffer(registration.registration_id) + if not inner_res.is_ok(): + return Result.new_error(inner_res.unwrap_error()) + _ = inner_res.unwrap() + self._gpu_buffer_registration = None + return Result.new_ok(OkNone()) + except ApiError as e: + return Result.new_error(e) + + def batch_put_from( + self, + keys: List[str], + payload_ptrs: List[int], + payload_sizes: List[int], + opts: Optional[PutOptionalArgs] = None, + ) -> List[int]: + if len(keys) != len(payload_ptrs) or len(keys) != len(payload_sizes): + raise ValueError( + "batch_put_from requires keys, payload_ptrs, and payload_sizes to have the same length" + ) + if len(keys) == 0: + return [] + + if self._client is None: + code = _error_to_ret_code( + GeneralError( + message="Store not initialized when batch_put_from(). Call setup() first." + ) + ) + return [code] * len(keys) + + lease_id: Optional[int] = opts.lease_id if opts is not None else None + reject_if_inflight_same_key = ( + bool(opts.reject_if_inflight_same_key) if opts is not None else False + ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = bool(opts.make_replica_task) if opts is not None else True + + try: + inner_res = self._client.batch_put_from( + keys, + payload_ptrs, + payload_sizes, + lease_id=lease_id, + reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, + ) + if inner_res.is_ok(): + submit_results = list(inner_res.unwrap()) + if len(submit_results) != len(keys): + raise RuntimeError( + "batch_put_from returned unexpected result length: " + f"expected={len(keys)} got={len(submit_results)}" + ) + return [int(item) for item in submit_results] + err = inner_res.unwrap_error() + code = _error_to_ret_code(err) + return [code] * len(keys) + except ApiError as e: + code = _error_to_ret_code(e) + return [code] * len(keys) + except Exception as e: + code = _error_to_ret_code(GeneralError(f"batch_put_from failed: {e}")) + return [code] * len(keys) + def get_size(self, key: str) -> Result[int, ApiError]: """Get the size of a stored value (non-blocking).""" return self._client.get_size(key) - def is_exist(self, key: str) -> Result[bool, ApiError]: + def is_exist(self, key: str, allow_local_snapshot: bool = False) -> Result[bool, ApiError]: """Check if a key exists in the store (non-blocking).""" try: - return self._client.is_exist(key) + if self._client is None: + return Result.new_error( + GeneralError( + message="Store not initialized when is_exist(). Call setup() first." + ) + ) + return self._client.is_exist(key, allow_local_snapshot=allow_local_snapshot) except Exception as e: return Result.new_error(GeneralError(f"Existence check failed: {str(e)}")) + def batch_get_into( + self, + keys: List[str], + payload_ptrs: List[int], + payload_capacities: List[int], + ) -> List[int]: + if len(keys) != len(payload_ptrs) or len(keys) != len(payload_capacities): + raise ValueError( + "batch_get_into requires keys, payload_ptrs, and payload_capacities to have the same length" + ) + if len(keys) == 0: + return [] + + if self._client is not None: + inner_res = self._client.batch_get_into(keys, payload_ptrs, payload_capacities) + if inner_res.is_ok(): + return list(inner_res.unwrap()) + err = inner_res.unwrap_error() + return [_error_to_ret_code(err)] * len(keys) + + results: List[int] = [] + for key, ptr, size in zip(keys, payload_ptrs, payload_capacities): + get_result = self.get_payload_into_ptr_blocking(key, ptr, size) + if get_result.is_ok(): + results.append(int(get_result.unwrap())) + else: + results.append(_error_to_ret_code(get_result.unwrap_error())) + return results + + def batch_is_exist( + self, + keys: List[str], + allow_local_snapshot: bool = False, + ) -> List[int]: + if len(keys) == 0: + return [] + + if self._client is None: + code = _error_to_ret_code( + GeneralError( + message="Store not initialized when batch_is_exist(). Call setup() first." + ) + ) + return [code] * len(keys) + + try: + inner_res = self._client.batch_is_exist( + keys, + allow_local_snapshot=allow_local_snapshot, + ) + if not inner_res.is_ok(): + code = _error_to_ret_code(inner_res.unwrap_error()) + return [code] * len(keys) + batch_results = self._normalize_batch_result_list( + inner_res.unwrap(), len(keys), "batch_is_exist" + ) + out: List[int] = [] + for idx, item in enumerate(batch_results): + if not isinstance(item, int): + raise GeneralError( + message=( + f"batch_is_exist returned non-int item for key {keys[idx]!r}: " + f"{type(item)}" + ) + ) + out.append(int(item)) + return out + except ApiError as e: + code = _error_to_ret_code(e) + return [code] * len(keys) + except Exception as e: + code = _error_to_ret_code(GeneralError(f"batch_is_exist failed: {e}")) + return [code] * len(keys) + + def get_meta(self, key: str) -> Result[Dict[str, Any], ApiError]: + """Query key metadata and one live placement without fetching payload bytes.""" + try: + inner = self._client.get_meta(key) + if not inner.is_ok(): + return Result.new_error(inner.unwrap_error()) + meta = inner.unwrap() + assert isinstance(meta, dict), f"get_meta returned non-dict: {type(meta)}" + return Result.new_ok(meta) + except Exception as e: + return Result.new_error(GeneralError(f"GetMeta failed for key '{key}': {str(e)}")) + + def batch_get_meta(self, keys: List[str]) -> List[Dict[str, Any]]: + if len(keys) == 0: + return [] + + if self._client is not None and hasattr(self._client, "batch_get_meta"): + inner_res = self._client.batch_get_meta(keys) + if inner_res.is_ok(): + rows = inner_res.unwrap() + assert isinstance(rows, list), ( + f"batch_get_meta returned non-list: {type(rows)}" + ) + return list(rows) + err = inner_res.unwrap_error() + code = _error_to_ret_code(err) + return [ + { + "exists": False, + "len": 0, + "node_id": "", + "src_addr": 0, + "src_base_addr": 0, + "segment_device_id": "", + "segment_device_desc": "", + "replica_count": 0, + "transport_error": True, + "error_code": -code, + "error_json": str(err), + } + for _ in keys + ] + + rows: List[Dict[str, Any]] = [] + for key in keys: + meta_res = self.get_meta(key) + if meta_res.is_ok(): + rows.append(meta_res.unwrap()) + else: + err = meta_res.unwrap_error() + rows.append( + { + "exists": False, + "len": 0, + "node_id": "", + "src_addr": 0, + "src_base_addr": 0, + "segment_device_id": "", + "segment_device_desc": "", + "replica_count": 0, + "transport_error": True, + "error_code": -_error_to_ret_code(err), + "error_json": str(err), + } + ) + return rows + def count_prefix(self, prefix: str) -> Result[int, ApiError]: """Count number of keys with the given prefix. @@ -604,10 +1567,7 @@ def rpc_call( try: if self._client is None: raise GeneralError(message="Store not initialized when rpc_call(). Call setup() first.") - if not isinstance(timeout_ms, int): - raise InvalidArgumentError(message=f"timeout_ms must be int; got {type(timeout_ms)}") - if timeout_ms < 10_000: - raise InvalidArgumentError(message=f"timeout_ms must be >= 10000; got {timeout_ms}") + _validate_explicit_rpc_timeout_ms(timeout_ms) encoded = encode_flat_kv_dict(payload) if not encoded.is_ok(): @@ -634,10 +1594,7 @@ def rpc_call_bytes( raise GeneralError(message="Store not initialized when rpc_call_bytes(). Call setup() first.") if not isinstance(payload, (bytes, bytearray)): raise InvalidArgumentError(message=f"payload must be bytes; got {type(payload)}") - if not isinstance(timeout_ms, int): - raise InvalidArgumentError(message=f"timeout_ms must be int; got {type(timeout_ms)}") - if timeout_ms < 10_000: - raise InvalidArgumentError(message=f"timeout_ms must be >= 10000; got {timeout_ms}") + _validate_explicit_rpc_timeout_ms(timeout_ms) inner = self._client.rpc_call(node_id, path, bytes(payload), timeout_ms) if not inner.is_ok(): @@ -729,10 +1686,7 @@ def sync_kv_to_file( raise InvalidArgumentError(message=f"file_offset must be int; got {type(file_offset)}") if file_offset < 0: raise InvalidArgumentError(message=f"file_offset must be >= 0; got {file_offset}") - if not isinstance(timeout_ms, int): - raise InvalidArgumentError(message=f"timeout_ms must be int; got {type(timeout_ms)}") - if timeout_ms < 10_000: - raise InvalidArgumentError(message=f"timeout_ms must be >= 10000; got {timeout_ms}") + _validate_explicit_rpc_timeout_ms(timeout_ms) return self._client.sync_kv_to_file( target_instance_key, @@ -776,6 +1730,13 @@ def instance_key(self) -> Result[str, ApiError]: def close(self) -> Result[OkNone, ApiError]: """Close and tear down the store.""" try: + if self._gpu_buffer_registration is not None: + unregister_result = self.unregister_gpu_buffer( + self._gpu_buffer_registration + ) + if not unregister_result.is_ok(): + return Result.new_error(unregister_result.unwrap_error()) + _ = unregister_result.unwrap() # Backend returns a Result; MUST be explicitly consumed to avoid # leaking an unconsumed Result that triggers __del__ assertion. res = self._client.close() @@ -806,6 +1767,757 @@ def get_cluster_name(self) -> str: raise RuntimeError("Store not initialized") return str(self._client.cluster_name()) + def wait_local_segments_ready(self) -> List[dict[str, Any]]: + if self._client is None: + raise RuntimeError( + "Store not initialized when wait_local_segments_ready(). Call setup() first." + ) + inner_res = self._client.wait_local_segments_ready() + if not inner_res.is_ok(): + raise RuntimeError( + f"wait_local_segments_ready backend error: {inner_res.unwrap_error()}" + ) + inner_res = inner_res.unwrap() + if not isinstance(inner_res, list): + raise RuntimeError( + "wait_local_segments_ready must return a list of segment mappings" + ) + out: List[dict[str, Any]] = [] + for item in inner_res: + if not isinstance(item, dict): + raise RuntimeError( + "wait_local_segments_ready segment item must be a dict" + ) + out.append(dict(item)) + return out + + def local_fast_put_start( + self, + keys: List[str], + value_len: int, + opts: Optional[PutOptionalArgs] = None, + ) -> int: + if self._client is None: + raise RuntimeError( + "Store not initialized when local_fast_put_start(). Call setup() first." + ) + if len(keys) == 0: + raise ValueError("local_fast_put_start requires at least one key") + if not isinstance(value_len, int): + raise ValueError(f"value_len must be int; got {type(value_len)}") + if value_len <= 0: + raise ValueError(f"value_len must be > 0; got {value_len}") + reject_if_inflight_same_key = ( + bool(opts.reject_if_inflight_same_key) if opts is not None else False + ) + reject_if_exist_same_key = ( + bool(opts.reject_if_exist_same_key) if opts is not None else False + ) + write_through = bool(opts.write_through) if opts is not None else True + make_replica_task = bool(opts.make_replica_task) if opts is not None else True + make_replica_task_mask: Optional[List[bool]] = None + if opts is not None and opts.make_replica_task_mask is not None: + requested_mask = opts.make_replica_task_mask + if not isinstance(requested_mask, list): + raise ValueError( + "make_replica_task_mask must be list[bool] when provided; " + f"got {type(requested_mask)}" + ) + if len(requested_mask) != len(keys): + raise ValueError( + "make_replica_task_mask length must match keys length; " + f"keys={len(keys)} mask={len(requested_mask)}" + ) + for index, item in enumerate(requested_mask): + if type(item) is not bool: + raise ValueError( + "make_replica_task_mask items must be bool; " + f"index={index} got={type(item)}" + ) + make_replica_task_mask = list(requested_mask) + atomic_group_lens = _validate_put_atomic_groups( + keys, + opts.atomic_group_lens if opts is not None else None, + make_replica_task_mask, + ) + inner_res = self._client.local_fast_put_start( + keys, + value_len, + reject_if_inflight_same_key=reject_if_inflight_same_key, + reject_if_exist_same_key=reject_if_exist_same_key, + write_through=write_through, + make_replica_task=make_replica_task, + make_replica_task_mask=make_replica_task_mask, + atomic_group_lens=atomic_group_lens, + ) + if not inner_res.is_ok(): + err = inner_res.unwrap_error() + if isinstance(err, Exception): + raise err + raise RuntimeError(f"local_fast_put_start backend error: {err}") + plan_ptr = inner_res.unwrap() + if not isinstance(plan_ptr, int) or plan_ptr <= 0: + raise RuntimeError(f"local_fast_put_start returned invalid plan_ptr: {plan_ptr!r}") + return int(plan_ptr) + + def local_fast_put_commit(self, plan_ptr: int) -> KvFuture: + if self._client is None: + raise RuntimeError( + "Store not initialized when local_fast_put_commit(). Call setup() first." + ) + inner_res = self._client.local_fast_put_commit(plan_ptr) + if not inner_res.is_ok(): + raise RuntimeError( + f"local_fast_put_commit backend error: {inner_res.unwrap_error()}" + ) + inner_future = inner_res.unwrap() + if inner_future is None: + raise RuntimeError("local_fast_put_commit returned empty future") + return _FluxonBatchRetCodeFuture(inner_future, [plan_ptr]) + + def put_abort(self, plan_ptr: int) -> None: + if self._client is None: + raise RuntimeError( + "Store not initialized when put_abort(). Call setup() first." + ) + inner_res = self._client.put_abort(plan_ptr) + if not inner_res.is_ok(): + raise RuntimeError(f"put_abort backend error: {inner_res.unwrap_error()}") + _ = inner_res.unwrap() + + def get_views( + self, + keys: List[str], + concurrency: Optional[int] = None, + ) -> int: + if self._client is None: + raise RuntimeError( + "Store not initialized when get_views(). Call setup() first." + ) + if len(keys) == 0: + raise ValueError("get_views requires at least one key") + inner_res = self._client.get_views( + keys, + concurrency=concurrency if concurrency is not None else self._batch_concurrency, + ) + if not inner_res.is_ok(): + raise RuntimeError(f"get_views backend error: {inner_res.unwrap_error()}") + plan_ptr = inner_res.unwrap() + if not isinstance(plan_ptr, int) or plan_ptr <= 0: + raise RuntimeError(f"get_views returned invalid plan_ptr: {plan_ptr!r}") + return int(plan_ptr) + + def release_views(self, plan_ptr: int) -> None: + if self._client is None: + raise RuntimeError( + "Store not initialized when release_views(). Call setup() first." + ) + inner_res = self._client.release_views(plan_ptr) + if not inner_res.is_ok(): + raise RuntimeError(f"release_views backend error: {inner_res.unwrap_error()}") + _ = inner_res.unwrap() + + def get_start( + self, + keys: List[str], + prefix_best_effort: bool = True, + atomic_group_lens: Optional[List[int]] = None, + ) -> GetStartHandle: + if self._client is None: + raise RuntimeError( + "Store not initialized when get_start(). Call setup() first." + ) + if len(keys) == 0: + raise ValueError("get_start requires at least one key") + normalized_group_lens: Optional[List[int]] = None + if atomic_group_lens is not None: + normalized_group_lens = [int(length) for length in atomic_group_lens] + if any(length <= 0 for length in normalized_group_lens): + raise ValueError("get_start atomic_group_lens entries must be > 0") + if sum(normalized_group_lens) != len(keys): + raise ValueError( + "get_start atomic_group_lens must sum to keys length: " + f"sum={sum(normalized_group_lens)} keys={len(keys)}" + ) + + started_at_ns = time.monotonic_ns() + inner_res = self._client.get_start( + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + self._batch_concurrency, + ) + if not inner_res.is_ok(): + raise RuntimeError(f"get_start backend error: {inner_res.unwrap_error()}") + payload = inner_res.unwrap() + if not isinstance(payload, dict): + raise RuntimeError(f"get_start returned non-dict payload: {type(payload)}") + backend_handle = int(payload["handle"]) + result = _build_get_start_result_from_backend_payload( + payload, + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + ) + logging.info( + "FluxonKVCacheStore get_start result: keys=%d raw_prefix_hit_len=%d " + "transferable_len=%d prefix_hit_groups=%d all_hit=%s " + "first_miss_index=%s first_miss_group_index=%s " + "prefix_best_effort=%s duration_ms=%.3f", + len(keys), + result.raw_prefix_hit_len, + result.transferable_len, + result.prefix_hit_groups, + result.all_hit, + result.first_miss_index, + result.first_miss_group_index, + result.prefix_best_effort, + (time.monotonic_ns() - started_at_ns) / 1_000_000.0, + ) + return GetStartHandle( + keys=tuple(keys), + result=result, + created_at_ns=started_at_ns, + backend_token=id(self), + backend_handle=backend_handle, + ) + + def get_plan( + self, + keys: List[str], + prefix_best_effort: bool = True, + atomic_group_lens: Optional[List[int]] = None, + ) -> GetPlanHandle: + if self._client is None: + raise RuntimeError("Store not initialized when get_plan(). Call setup() first.") + if not keys: + raise ValueError("get_plan requires at least one key") + normalized_group_lens: Optional[List[int]] = None + if atomic_group_lens is not None: + normalized_group_lens = [int(length) for length in atomic_group_lens] + if any(length <= 0 for length in normalized_group_lens): + raise ValueError("get_plan atomic_group_lens entries must be > 0") + if sum(normalized_group_lens) != len(keys): + raise ValueError( + "get_plan atomic_group_lens must sum to keys length: " + f"sum={sum(normalized_group_lens)} keys={len(keys)}" + ) + started_at_ns = time.monotonic_ns() + inner_res = self._client.get_plan( + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + ) + if not inner_res.is_ok(): + raise RuntimeError(f"get_plan backend error: {inner_res.unwrap_error()}") + payload = inner_res.unwrap() + if not isinstance(payload, dict): + raise RuntimeError(f"get_plan returned non-dict payload: {type(payload)}") + backend_handle = int(payload["handle"]) + try: + result = _build_get_start_result_from_backend_payload( + payload, + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + ) + gpu_payload = dict(payload) + gpu_payload["raw_prefix_hit_len"] = payload["gpu_raw_prefix_hit_len"] + gpu_result = _build_get_start_result_from_backend_payload( + gpu_payload, + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + ) + gpu_remote_indices = tuple(int(index) for index in payload["gpu_remote_indices"]) + if ( + tuple(sorted(set(gpu_remote_indices))) != gpu_remote_indices + or any( + index < 0 or index >= gpu_result.transferable_len + for index in gpu_remote_indices + ) + ): + raise RuntimeError( + "get_plan returned invalid gpu_remote_indices: " + f"indices={gpu_remote_indices} " + f"transferable={gpu_result.transferable_len}" + ) + except Exception: + cancel_res = self._client.cancel_get_plan(backend_handle) + if not cancel_res.is_ok(): + logging.exception( + "get_plan payload validation failed and cleanup also failed: %s", + cancel_res.unwrap_error(), + ) + else: + _ = cancel_res.unwrap() + raise + return GetPlanHandle( + keys=tuple(keys), + result=result, + gpu_result=gpu_result, + created_at_ns=started_at_ns, + backend_token=id(self), + backend_handle=backend_handle, + gpu_remote_indices=gpu_remote_indices, + ) + + def cancel_get_plan(self, handle: GetPlanHandle) -> None: + if not isinstance(handle, GetPlanHandle): + raise TypeError(f"cancel_get_plan requires GetPlanHandle, got {type(handle)}") + if handle.backend_token != id(self): + raise RuntimeError("cancel_get_plan handle belongs to a different store") + if handle.closed: + return + if self._client is None: + raise RuntimeError("Store not initialized when cancel_get_plan().") + inner_res = self._client.cancel_get_plan(handle.backend_handle) + if not inner_res.is_ok(): + raise RuntimeError(f"cancel_get_plan backend error: {inner_res.unwrap_error()}") + _ = inner_res.unwrap() + handle.closed = True + + def execute_get_plan_cpu( + self, + handle: GetPlanHandle, + *, + consume_prefix_len: int, + concurrency: Optional[int] = None, + ) -> GetStartHandle: + if not isinstance(handle, GetPlanHandle): + raise TypeError( + f"execute_get_plan_cpu requires GetPlanHandle, got {type(handle)}" + ) + if handle.backend_token != id(self) or handle.closed: + raise RuntimeError("execute_get_plan_cpu requires a live plan from this store") + narrowed = _narrow_get_start_result(handle.result, int(consume_prefix_len)) + if consume_prefix_len <= 0 or consume_prefix_len > handle.result.transferable_len: + raise ValueError( + "execute_get_plan_cpu consume prefix is outside the CPU plan: " + f"consume={consume_prefix_len} transferable={handle.result.transferable_len}" + ) + if self._client is None: + raise RuntimeError("Store not initialized when execute_get_plan_cpu().") + inner_res = self._client.execute_get_plan_cpu( + handle.backend_handle, + int(consume_prefix_len), + self._batch_concurrency if concurrency is None else int(concurrency), + ) + if not inner_res.is_ok(): + raise RuntimeError( + f"execute_get_plan_cpu backend error: {inner_res.unwrap_error()}" + ) + _ = inner_res.unwrap() + handle.closed = True + return GetStartHandle( + keys=narrowed.keys, + result=narrowed, + created_at_ns=handle.created_at_ns, + backend_token=id(self), + backend_handle=handle.backend_handle, + ) + + def execute_get_plan_gpu( + self, + handle: GetPlanHandle, + destinations: List[GpuDestination], + *, + consume_prefix_len: int, + concurrency: Optional[int] = None, + ) -> GpuGetStartHandle: + if not isinstance(handle, GetPlanHandle): + raise TypeError( + f"execute_get_plan_gpu requires GetPlanHandle, got {type(handle)}" + ) + if handle.backend_token != id(self) or handle.closed: + raise RuntimeError("execute_get_plan_gpu requires a live plan from this store") + remote_indices = tuple( + index for index in handle.gpu_remote_indices if index < consume_prefix_len + ) + if not remote_indices: + raise ValueError("execute_get_plan_gpu requires at least one remote source") + if len(destinations) != len(remote_indices): + raise ValueError( + "execute_get_plan_gpu requires one exact destination per remote source: " + f"destinations={len(destinations)} remote={len(remote_indices)} " + f"consume={consume_prefix_len}" + ) + narrowed = _narrow_get_start_result(handle.gpu_result, int(consume_prefix_len)) + if consume_prefix_len <= 0 or consume_prefix_len > handle.gpu_result.transferable_len: + raise ValueError( + "execute_get_plan_gpu consume prefix is outside the GPU plan: " + f"consume={consume_prefix_len} transferable={handle.gpu_result.transferable_len}" + ) + active_registration = self._gpu_buffer_registration + if active_registration is None: + raise RuntimeError("execute_get_plan_gpu requires an active GPU registration") + for index, destination in enumerate(destinations): + if not isinstance(destination, GpuDestination): + raise TypeError( + "execute_get_plan_gpu destinations must be GpuDestination: " + f"index={index} got={type(destination)}" + ) + if destination.registration_id != active_registration.registration_id: + raise ValueError( + "execute_get_plan_gpu destination uses a stale registration: " + f"index={index}" + ) + if destination.ptr < active_registration.ptr or destination.end > active_registration.end: + raise ValueError( + "execute_get_plan_gpu destination is outside the active registration: " + f"index={index}" + ) + if self._client is None: + raise RuntimeError("Store not initialized when execute_get_plan_gpu().") + inner_res = self._client.execute_get_plan_gpu( + handle.backend_handle, + [ + (destination.registration_id, destination.ptr, destination.capacity) + for destination in destinations + ], + int(consume_prefix_len), + self._batch_concurrency if concurrency is None else int(concurrency), + ) + if not inner_res.is_ok(): + raise RuntimeError( + f"execute_get_plan_gpu backend error: {inner_res.unwrap_error()}" + ) + _ = inner_res.unwrap() + handle.closed = True + return GpuGetStartHandle( + keys=narrowed.keys, + destinations=tuple(destinations), + result=narrowed, + created_at_ns=handle.created_at_ns, + backend_token=id(self), + backend_handle=handle.backend_handle, + remote_indices=remote_indices, + ) + + def cancel_get_transfer(self, handle: GetStartHandle) -> None: + if not isinstance(handle, GetStartHandle): + raise TypeError(f"cancel_get_transfer requires GetStartHandle, got {type(handle)}") + if handle.backend_token is not None and handle.backend_token != id(self): + raise RuntimeError( + "cancel_get_transfer handle belongs to a different FluxonKVCacheStore" + ) + if handle.closed: + return + if self._client is None: + raise RuntimeError( + "Store not initialized when cancel_get_transfer(). Call setup() first." + ) + inner_res = self._client.cancel_get_transfer(int(handle.backend_handle)) + if not inner_res.is_ok(): + raise RuntimeError( + f"cancel_get_transfer backend error: {inner_res.unwrap_error()}" + ) + _ = inner_res.unwrap() + handle.closed = True + + def get_transfer( + self, + handle: GetStartHandle, + concurrency: Optional[int] = None, + *, + consume_prefix_len: Optional[int] = None, + ) -> int: + if not isinstance(handle, GetStartHandle): + raise TypeError(f"get_transfer requires GetStartHandle, got {type(handle)}") + if handle.backend_token is not None and handle.backend_token != id(self): + raise RuntimeError("get_transfer handle belongs to a different FluxonKVCacheStore") + if handle.closed: + raise RuntimeError("get_transfer handle has been closed") + result = handle.result + if result.transferable_len == 0: + raise RuntimeError( + "get_transfer requires a non-empty transferable prefix: " + f"transferable_len={result.transferable_len} total={len(result.keys)} " + f"raw_prefix_hit_len={result.raw_prefix_hit_len} " + f"first_miss_index={result.first_miss_index} " + f"first_miss_group_index={result.first_miss_group_index}" + ) + if consume_prefix_len is None: + normalized_consume_prefix_len = result.transferable_len + else: + if isinstance(consume_prefix_len, bool) or not isinstance( + consume_prefix_len, int + ): + raise TypeError( + "get_transfer consume_prefix_len must be an int or None, got " + f"{type(consume_prefix_len)}" + ) + normalized_consume_prefix_len = int(consume_prefix_len) + if ( + normalized_consume_prefix_len <= 0 + or normalized_consume_prefix_len > result.transferable_len + ): + raise ValueError( + "get_transfer consume_prefix_len must be within the live " + "transferable prefix: " + f"consume={normalized_consume_prefix_len} " + f"transferable={result.transferable_len}" + ) + group_lens = result.atomic_group_lens or (len(result.keys),) + group_end = 0 + for group_len in group_lens: + group_end += int(group_len) + if group_end >= normalized_consume_prefix_len: + break + if group_end != normalized_consume_prefix_len: + raise ValueError( + "get_transfer consume_prefix_len must end at an atomic-group " + "boundary: " + f"consume={normalized_consume_prefix_len} " + f"atomic_group_lens={group_lens}" + ) + if self._client is None: + raise RuntimeError( + "Store not initialized when get_transfer(). Call setup() first." + ) + _ = concurrency + inner_res = self._client.get_transfer( + handle.backend_handle, normalized_consume_prefix_len + ) + if not inner_res.is_ok(): + handle.closed = True + raise RuntimeError(f"get_transfer backend error: {inner_res.unwrap_error()}") + plan_ptr = inner_res.unwrap() + handle.closed = True + if not isinstance(plan_ptr, int) or plan_ptr <= 0: + raise RuntimeError(f"get_transfer returned invalid plan_ptr: {plan_ptr!r}") + return int(plan_ptr) + + def get_start_gpu( + self, + keys: List[str], + destinations: List[GpuDestination], + prefix_best_effort: bool = True, + atomic_group_lens: Optional[List[int]] = None, + ) -> GpuGetStartHandle: + """Start background remote reads directly into caller-owned GPU staging.""" + if self._client is None: + raise RuntimeError( + "Store not initialized when get_start_gpu(). Call setup() first." + ) + if len(keys) == 0: + raise ValueError("get_start_gpu requires at least one key") + if len(keys) != len(destinations): + raise ValueError( + "get_start_gpu requires one destination per key: " + f"keys={len(keys)} destinations={len(destinations)}" + ) + active_registration = self._gpu_buffer_registration + if active_registration is None: + raise RuntimeError("get_start_gpu requires an active GPU registration") + for index, destination in enumerate(destinations): + if not isinstance(destination, GpuDestination): + raise TypeError( + "get_start_gpu destinations must be GpuDestination: " + f"index={index} got={type(destination)}" + ) + if destination.registration_id != active_registration.registration_id: + raise ValueError( + "get_start_gpu destination uses a stale registration: " + f"index={index} destination_id={destination.registration_id} " + f"active_id={active_registration.registration_id}" + ) + if ( + destination.ptr < active_registration.ptr + or destination.end > active_registration.end + ): + raise ValueError( + "get_start_gpu destination is outside the active registration: " + f"index={index} destination={destination!r} " + f"registration={active_registration!r}" + ) + + normalized_group_lens: Optional[List[int]] = None + if atomic_group_lens is not None: + normalized_group_lens = [int(length) for length in atomic_group_lens] + if any(length <= 0 for length in normalized_group_lens): + raise ValueError("get_start_gpu atomic_group_lens entries must be > 0") + if sum(normalized_group_lens) != len(keys): + raise ValueError( + "get_start_gpu atomic_group_lens must sum to keys length: " + f"sum={sum(normalized_group_lens)} keys={len(keys)}" + ) + + started_at_ns = time.monotonic_ns() + inner_res = self._client.get_start_gpu( + list(keys), + [ + ( + destination.registration_id, + destination.ptr, + destination.capacity, + ) + for destination in destinations + ], + bool(prefix_best_effort), + normalized_group_lens, + self._batch_concurrency, + ) + if not inner_res.is_ok(): + raise RuntimeError(f"get_start_gpu backend error: {inner_res.unwrap_error()}") + payload = inner_res.unwrap() + if not isinstance(payload, dict): + raise RuntimeError( + f"get_start_gpu returned non-dict payload: {type(payload)}" + ) + backend_handle = int(payload["handle"]) + try: + result = _build_get_start_result_from_backend_payload( + payload, + list(keys), + bool(prefix_best_effort), + normalized_group_lens, + ) + except Exception: + cancel_res = self._client.cancel_get_transfer_gpu(backend_handle) + if not cancel_res.is_ok(): + logging.exception( + "get_start_gpu payload validation failed and cleanup also failed: %s", + cancel_res.unwrap_error(), + ) + else: + _ = cancel_res.unwrap() + raise + return GpuGetStartHandle( + keys=tuple(keys), + destinations=tuple(destinations), + result=result, + created_at_ns=started_at_ns, + backend_token=id(self), + backend_handle=backend_handle, + remote_indices=tuple(range(result.transferable_len)), + ) + + def cancel_get_transfer_gpu(self, handle: GpuGetStartHandle) -> None: + if not isinstance(handle, GpuGetStartHandle): + raise TypeError( + f"cancel_get_transfer_gpu requires GpuGetStartHandle, got {type(handle)}" + ) + if handle.backend_token != id(self): + raise RuntimeError( + "cancel_get_transfer_gpu handle belongs to a different FluxonKVCacheStore" + ) + if handle.closed: + return + if self._client is None: + raise RuntimeError( + "Store not initialized when cancel_get_transfer_gpu(). Call setup() first." + ) + inner_res = self._client.cancel_get_transfer_gpu(handle.backend_handle) + handle.closed = True + if not inner_res.is_ok(): + raise RuntimeError( + f"cancel_get_transfer_gpu backend error: {inner_res.unwrap_error()}" + ) + _ = inner_res.unwrap() + + def get_transfer_gpu( + self, + handle: GpuGetStartHandle, + *, + consume_prefix_len: Optional[int] = None, + ) -> int: + """Wait for GPU transfer and return one ordered CPU/GPU source plan.""" + if not isinstance(handle, GpuGetStartHandle): + raise TypeError( + f"get_transfer_gpu requires GpuGetStartHandle, got {type(handle)}" + ) + if handle.backend_token != id(self): + raise RuntimeError( + "get_transfer_gpu handle belongs to a different FluxonKVCacheStore" + ) + if handle.closed: + raise RuntimeError("get_transfer_gpu handle has been closed") + result = handle.result + normalized_consume_prefix_len = ( + result.transferable_len + if consume_prefix_len is None + else consume_prefix_len + ) + if ( + isinstance(normalized_consume_prefix_len, bool) + or not isinstance(normalized_consume_prefix_len, int) + or normalized_consume_prefix_len <= 0 + or normalized_consume_prefix_len > result.transferable_len + ): + raise ValueError( + "get_transfer_gpu consume_prefix_len must be within the live prefix: " + f"consume={normalized_consume_prefix_len!r} " + f"transferable={result.transferable_len}" + ) + group_lens = result.atomic_group_lens or (len(result.keys),) + group_end = 0 + for group_len in group_lens: + group_end += int(group_len) + if group_end >= normalized_consume_prefix_len: + break + if group_end != normalized_consume_prefix_len: + raise ValueError( + "get_transfer_gpu consume_prefix_len must end at an atomic-group boundary: " + f"consume={normalized_consume_prefix_len} groups={group_lens}" + ) + if self._client is None: + raise RuntimeError( + "Store not initialized when get_transfer_gpu(). Call setup() first." + ) + inner_res = self._client.get_transfer_gpu( + handle.backend_handle, + normalized_consume_prefix_len, + ) + handle.closed = True + if not inner_res.is_ok(): + raise RuntimeError( + f"get_transfer_gpu backend error: {inner_res.unwrap_error()}" + ) + payload = inner_res.unwrap() + if not isinstance(payload, dict): + raise RuntimeError( + f"get_transfer_gpu returned non-dict payload: {type(payload)}" + ) + transferred_prefix_len = int(payload["transferred_prefix_len"]) + consumed_prefix_len = int(payload["consumed_prefix_len"]) + if ( + transferred_prefix_len != result.transferable_len + or consumed_prefix_len != normalized_consume_prefix_len + ): + raise RuntimeError( + "get_transfer_gpu returned inconsistent terminal lengths: " + f"transferred={transferred_prefix_len}/{result.transferable_len} " + f"consumed={consumed_prefix_len}/{normalized_consume_prefix_len}" + ) + transfer_wall_us = payload["transfer_wall_us"] + finish_wait_us = payload["finish_wait_us"] + terminal_before_consume = payload["terminal_before_consume"] + terminal_to_consume_us = payload["terminal_to_consume_us"] + for field_name, value in ( + ("transfer_wall_us", transfer_wall_us), + ("finish_wait_us", finish_wait_us), + ("terminal_to_consume_us", terminal_to_consume_us), + ): + if type(value) is not int or value < 0: + raise RuntimeError( + "get_transfer_gpu returned invalid timing: " + f"{field_name}={value!r}" + ) + if type(terminal_before_consume) is not bool: + raise RuntimeError( + "get_transfer_gpu returned invalid terminal_before_consume: " + f"{terminal_before_consume!r}" + ) + handle.transfer_wall_us = transfer_wall_us + handle.finish_wait_us = finish_wait_us + handle.terminal_before_consume = terminal_before_consume + handle.terminal_to_consume_us = terminal_to_consume_us + plan_ptr = int(payload["plan_ptr"]) + if plan_ptr <= 0: + raise RuntimeError(f"get_transfer_gpu returned invalid plan_ptr: {plan_ptr}") + return plan_ptr + def get_etcd_config(self) -> List[str]: if self._client is None: raise RuntimeError("Store not initialized") @@ -889,6 +2601,19 @@ def metrics_snapshot(self) -> MetricSnapshot: return MetricSnapshot(per_segment=normalized) + def observability_snapshot_async(self) -> KvFuture: + """Return a future for Fluxon locality and IO counters.""" + if self._client is None: + raise RuntimeError( + "Store not initialized when observability_snapshot_async(). Call setup() first." + ) + res = self._client.observability_snapshot_async() + if not res.is_ok(): + raise RuntimeError( + f"observability_snapshot_async backend error: {res.unwrap_error()}" + ) + return res.unwrap() + # --- Fluxon-kv lease helpers (synchronous) --- def allocate_lease(self, ttl_seconds: int) -> Result[int, ApiError]: try: @@ -948,7 +2673,7 @@ def __init__(self, inner_future: Any) -> None: self._inner = inner_future def is_waiting(self) -> bool: - return bool(getattr(self._inner, "is_waiting")()) + return bool(self._inner.is_waiting()) def _decode_wait_success( self, @@ -985,7 +2710,7 @@ def __init__(self, inner_future: Any) -> None: self._inner = inner_future def is_waiting(self) -> bool: - return bool(getattr(self._inner, "is_waiting")()) + return bool(self._inner.is_waiting()) def wait(self) -> Result[bytes, ApiError]: res = self._inner.wait() @@ -1022,7 +2747,7 @@ def __del__(self) -> None: self._dlpack_capsules = [] def is_waiting(self) -> bool: - return bool(getattr(self._inner, "is_waiting")()) + return bool(self._inner.is_waiting()) def wait(self) -> Result[Union[Any, MemHolder], ApiError]: from ..api_error import OkNone, Result as PyResult # type: ignore @@ -1035,3 +2760,36 @@ def wait(self) -> Result[Union[Any, MemHolder], ApiError]: _ = res.unwrap() return PyResult.new_ok(OkNone()) # type: ignore + + +class _FluxonBatchRetCodeFuture(KvFuture): + """Future wrapper for batch APIs that resolve to one integer ret-code per key.""" + + def __init__(self, inner_future: Any, keepalive: List[object]) -> None: + self._inner = inner_future + self._keepalive = keepalive + + def __del__(self) -> None: + self._keepalive = [] + + def is_waiting(self) -> bool: + return bool(self._inner.is_waiting()) + + def wait(self) -> Result[List[int], ApiError]: + res = self._inner.wait() + self._keepalive = [] + if not res.is_ok(): + return Result.new_error(res.unwrap_error()) + raw = res.unwrap() + if not isinstance(raw, list): + return Result.new_error( + GeneralError(message=f"batch future returned non-list payload: {type(raw)}") + ) + out: List[int] = [] + for item in raw: + if not isinstance(item, int): + return Result.new_error( + GeneralError(message=f"batch future returned non-int item: {type(item)}") + ) + out.append(int(item)) + return Result.new_ok(out) diff --git a/fluxon_py/kvclient/kvclient_interface.py b/fluxon_py/kvclient/kvclient_interface.py index f50db0f..ead8e4e 100644 --- a/fluxon_py/kvclient/kvclient_interface.py +++ b/fluxon_py/kvclient/kvclient_interface.py @@ -21,6 +21,48 @@ FlatDict = Dict[str, Union[int, float, bool, str, bytes, DLPacked]] +@dataclass(frozen=True) +class GetStartResult: + """ + Result of a group-prefix best-effort get_start(). + + Semantics: + - ``keys`` is the caller-provided ordered page-key sequence. + - ``raw_prefix_hit_len`` is the page-level continuous hit prefix. + - ``transferable_len`` is rounded down to complete atomic groups and is the + maximum prefix that can be consumed by get_transfer(). A caller may + consume a shorter prefix only at an atomic-group boundary. + """ + + keys: Tuple[str, ...] + raw_prefix_hit_len: int + transferable_len: int + prefix_hit_groups: int + atomic_group_lens: Optional[Tuple[int, ...]] + prefix_best_effort: bool + first_miss_index: Optional[int] + first_miss_group_index: Optional[int] + all_hit: bool + + +@dataclass +class GetStartHandle: + """ + Opaque-ish handle returned by get_start() and consumed by get_transfer(). + + Callers must pass this handle to cancel_get_transfer() when abandoning it + without calling get_transfer(). Fluxon backends may keep strong holder + references alive while this handle is live. + """ + + keys: Tuple[str, ...] + result: GetStartResult + created_at_ns: int + backend_token: Optional[int] = None + backend_handle: int = 0 + closed: bool = False + + @dataclass class PutOptionalArgs: """ @@ -29,9 +71,27 @@ class PutOptionalArgs: - lease_id: attach the written key to a lease on commit. - reject_if_inflight_same_key: ask Fluxon to fail-fast when the same key is already being written by another inflight put. + - reject_if_exist_same_key: ask Fluxon to fail-fast when the key already has a + committed live replica. + - write_through: keep synchronous remote-placement semantics when the backend + supports an async write-back path. Defaults to True to match SGLang + HiCache's default write policy. + - make_replica_task: enqueue an asynchronous replica after a local write-back + commit. Set False for a local-only write-back. + - make_replica_task_mask: optional per-key replica admission decisions for + local_fast_put_start(). Its length must match the key batch. The scalar + make_replica_task remains the batch-wide gate. + - atomic_group_lens: optional positive lengths that partition the ordered + local_fast_put_start() key batch into atomic groups. Replica admission + must be uniform within every group. """ lease_id: Optional[int] = None reject_if_inflight_same_key: bool = False + reject_if_exist_same_key: bool = False + write_through: bool = True + make_replica_task: bool = True + make_replica_task_mask: Optional[List[bool]] = None + atomic_group_lens: Optional[List[int]] = None def support_mooncake(self) -> Tuple[bool, List[str]]: """ @@ -48,6 +108,16 @@ def support_mooncake(self) -> Tuple[bool, List[str]]: unsupported.append("lease_id") if self.reject_if_inflight_same_key: unsupported.append("reject_if_inflight_same_key") + if self.reject_if_exist_same_key: + unsupported.append("reject_if_exist_same_key") + if self.write_through: + unsupported.append("write_through") + if not self.make_replica_task: + unsupported.append("make_replica_task") + if self.make_replica_task_mask is not None: + unsupported.append("make_replica_task_mask") + if self.atomic_group_lens is not None: + unsupported.append("atomic_group_lens") return (len(unsupported) == 0, unsupported) @@ -139,7 +209,7 @@ def put_blocking( _ = wait_result.unwrap() return Result.new_ok(OkNone()) - def get_blocking(self, key: str) -> Result["MemHolder", ApiError]: + def get_blocking(self, key: str) -> Result[Union[Any, "MemHolder"], ApiError]: """Synchronously retrieve a value by key. Default implementation delegates to ``get()`` followed by @@ -151,6 +221,103 @@ def get_blocking(self, key: str) -> Result["MemHolder", ApiError]: return Result.new_error(get_result.unwrap_error()) return get_result.unwrap().wait() + def batch_put_blocking( + self, + keys: List[str], + values: List[FlatDict], + opts: Optional[PutOptionalArgs] = None, + concurrency: Optional[int] = None, + ) -> List[Result[OkNone, ApiError]]: + """Synchronously store a batch of key-value pairs.""" + if len(keys) != len(values): + raise ValueError("batch_put_blocking requires keys and values to have the same length") + _ = concurrency + return [self.put_blocking(key, value, opts=opts) for key, value in zip(keys, values)] + + def batch_get_blocking( + self, + keys: List[str], + concurrency: Optional[int] = None, + ) -> List[Result[Union[Any, "MemHolder"], ApiError]]: + """Synchronously retrieve a batch of keys.""" + _ = concurrency + return [self.get_blocking(key) for key in keys] + + def local_fast_put_start( + self, + keys: List[str], + value_len: int, + opts: Optional[PutOptionalArgs] = None, + ) -> int: + _ = keys + _ = value_len + _ = opts + raise NotImplementedError( + "local_fast_put_start is only implemented by backends with native plan_ptr support" + ) + + def local_fast_put_commit(self, plan_ptr: int) -> "KvFuture": + _ = plan_ptr + raise NotImplementedError( + "local_fast_put_commit is only implemented by backends with native plan_ptr support" + ) + + def put_abort(self, plan_ptr: int) -> None: + _ = plan_ptr + raise NotImplementedError( + "put_abort is only implemented by backends with native plan_ptr support" + ) + + def get_views( + self, + keys: List[str], + concurrency: Optional[int] = None, + ) -> int: + _ = keys + _ = concurrency + raise NotImplementedError( + "get_views is only implemented by backends with native plan_ptr support" + ) + + def release_views(self, plan_ptr: int) -> None: + _ = plan_ptr + raise NotImplementedError( + "release_views is only implemented by backends with native plan_ptr support" + ) + + def get_start( + self, + keys: List[str], + prefix_best_effort: bool = True, + atomic_group_lens: Optional[List[int]] = None, + ) -> GetStartHandle: + _ = keys + _ = prefix_best_effort + _ = atomic_group_lens + raise NotImplementedError( + "get_start is only implemented by backends with native prefix get support" + ) + + def get_transfer( + self, + handle: GetStartHandle, + concurrency: Optional[int] = None, + *, + consume_prefix_len: Optional[int] = None, + ) -> int: + _ = handle + _ = concurrency + _ = consume_prefix_len + raise NotImplementedError( + "get_transfer is only implemented by backends with native prefix get support" + ) + + def cancel_get_transfer(self, handle: GetStartHandle) -> None: + _ = handle + raise NotImplementedError( + "cancel_get_transfer is only implemented by backends with native prefix get support" + ) + @abstractmethod def get_size(self, key: str) -> Result[int, ApiError]: """Get the size of a stored value (non-blocking).""" diff --git a/fluxon_py/tests/test_config.py b/fluxon_py/tests/test_config.py index 6de5180..0b7cdc6 100644 --- a/fluxon_py/tests/test_config.py +++ b/fluxon_py/tests/test_config.py @@ -49,11 +49,13 @@ def _build_checks(selected_test_id: Optional[str]) -> List[Tuple[str, Callable[[ ("fluxonkv_owner_requires_sub_cluster", test_fluxonkv_owner_requires_sub_cluster), ("fluxonkv_owner_requires_large_file_paths", test_fluxonkv_owner_requires_large_file_paths), ("fluxonkv_external_forbids_large_file_paths", test_fluxonkv_external_forbids_large_file_paths), + ("fluxonkv_owner_ssd_capacity", test_fluxonkv_owner_ssd_capacity), ("fluxonkv_p2p_relay_removed", test_fluxonkv_p2p_relay_removed), ("fluxon_client_config_yaml_shape", test_fluxon_client_config_yaml_shape), ("fluxonkv_protocol_field", test_fluxonkv_protocol_field), ("fluxonkv_runtime_defaults_are_internal", test_fluxonkv_runtime_defaults_are_internal), ("fluxonkv_removed_rdma_config_keys", test_fluxonkv_removed_rdma_config_keys), + ("fluxonkv_owner_hot_writeback", test_fluxonkv_owner_hot_writeback), ("fluxonkv_test_spec_config", test_fluxonkv_test_spec_config), ("fluxon_pyo3_import_authority", test_fluxon_pyo3_import_authority), ] @@ -333,6 +335,66 @@ def test_fluxonkv_external_forbids_large_file_paths(): print(f"❌ FAIL: test_fluxonkv_external_forbids_large_file_paths - {e}") +def test_fluxonkv_owner_ssd_capacity(): + """Ensure the Python schema forwards owner SSD capacity to Rust unchanged.""" + try: + owner = _owner_fluxonkv_base_config(tag="owner_ssd_capacity") + owner["fluxonkv_spec"]["large_limit_size"] = [67108864] + owner["fluxonkv_spec"]["ssd_write_rate_limit_bytes_per_sec"] = 268435456 + owner["fluxonkv_spec"]["ssd_write_burst_bytes"] = 67108864 + config = FluxonKvClientConfig(owner) + assert config.to_dict()["fluxonkv_spec"]["large_limit_size"] == [67108864] + assert ( + config.to_dict()["fluxonkv_spec"]["ssd_write_rate_limit_bytes_per_sec"] + == 268435456 + ) + assert config.to_dict()["fluxonkv_spec"]["ssd_write_burst_bytes"] == 67108864 + rendered = config.to_fluxon_kv_client_config_yaml_str() + assert "large_limit_size:" in rendered + assert "- 67108864" in rendered + assert "ssd_write_rate_limit_bytes_per_sec: 268435456" in rendered + assert "ssd_write_burst_bytes: 67108864" in rendered + + unpaired = _owner_fluxonkv_base_config(tag="owner_ssd_unpaired_limit") + unpaired["fluxonkv_spec"]["large_limit_size"] = [67108864] + unpaired["fluxonkv_spec"]["ssd_write_rate_limit_bytes_per_sec"] = 1 + try: + FluxonKvClientConfig(unpaired) + print("❌ FAIL: test_fluxonkv_owner_ssd_capacity - unpaired rate should be rejected") + return + except ValueError: + pass + + wrong_shape = _owner_fluxonkv_base_config(tag="owner_ssd_capacity_wrong_shape") + wrong_shape["fluxonkv_spec"]["large_limit_size"] = 67108864 + try: + FluxonKvClientConfig(wrong_shape) + print("❌ FAIL: test_fluxonkv_owner_ssd_capacity - non-list capacity should be rejected") + return + except ValueError: + pass + + external = { + "instance_key": "test_external_ssd_capacity", + "contribute_to_cluster_pool_size": {"dram": 0, "vram": {}}, + "fluxonkv_spec": { + "cluster_name": "test_cluster", + "share_mem_path": "/tmp/kvcache_shared_memory/test", + "large_limit_size": [67108864], + }, + } + try: + FluxonKvClientConfig(external) + print("❌ FAIL: test_fluxonkv_owner_ssd_capacity - external capacity should be rejected") + return + except ValueError: + pass + + print("✅ PASS: test_fluxonkv_owner_ssd_capacity") + except Exception as e: + print(f"❌ FAIL: test_fluxonkv_owner_ssd_capacity - {e}") + + def test_fluxonkv_p2p_relay_removed(): """Ensure removed fluxonkv_spec.p2p_relay is rejected as an unknown key.""" try: @@ -444,6 +506,61 @@ def test_fluxonkv_removed_rdma_config_keys(): print(f"❌ FAIL: test_fluxonkv_removed_rdma_config_keys - {e}") +def test_fluxonkv_owner_hot_writeback(): + """Validate the owner-only hot-tier ratio and preserve it in Rust YAML.""" + try: + owner = _owner_fluxonkv_base_config(tag="owner_hot_writeback") + owner["replica_writeback_hot_capacity_ratio"] = 0.75 + loaded = yaml.safe_load( + FluxonKvClientConfig(owner).to_fluxon_kv_client_config_yaml_str() + ) + assert loaded["replica_writeback_hot_capacity_ratio"] == 0.75 + + for invalid in (0, 1, -0.1, float("nan"), True, "0.75"): + invalid_owner = _owner_fluxonkv_base_config(tag="owner_hot_invalid") + invalid_owner["replica_writeback_hot_capacity_ratio"] = invalid + try: + FluxonKvClientConfig(invalid_owner) + print( + "❌ FAIL: test_fluxonkv_owner_hot_writeback - invalid ratio should be rejected" + ) + return + except ValueError: + pass + + external = { + "instance_key": "external_hot_invalid", + "replica_writeback_hot_capacity_ratio": 0.75, + "fluxonkv_spec": { + "cluster_name": "test_cluster", + "share_mem_path": "/tmp/kvcache_shared_memory/external_hot_invalid", + }, + } + try: + FluxonKvClientConfig(external) + print( + "❌ FAIL: test_fluxonkv_owner_hot_writeback - external ratio should be rejected" + ) + return + except ValueError: + pass + + mooncake = config_dict() + mooncake["replica_writeback_hot_capacity_ratio"] = 0.75 + try: + FluxonKvClientConfig(mooncake) + print( + "❌ FAIL: test_fluxonkv_owner_hot_writeback - Mooncake ratio should be rejected" + ) + return + except ValueError: + pass + + print("✅ PASS: test_fluxonkv_owner_hot_writeback") + except Exception as e: + print(f"❌ FAIL: test_fluxonkv_owner_hot_writeback - {e}") + + def test_fluxonkv_test_spec_config(): """Ensure test_spec_config is accepted, normalized, and serialized.""" try: @@ -473,6 +590,10 @@ def test_fluxonkv_test_spec_config(): rdma_devices["test_spec_config"]["tcp_thread_reactor_shard_count"] = 2 rdma_devices["test_spec_config"]["tcp_thread_bulk_lane_count"] = 4 rdma_devices["test_spec_config"]["tcp_thread_control_lane_count"] = 3 + rdma_devices["test_spec_config"]["replica_task_max_inflight"] = 16 + rdma_devices["test_spec_config"]["ssd_read_source_policy"] = ( + "local_ssd_only_first" + ) rdma_devices["test_spec_config"][ "require_transfer_rpc_fast_path_ready_timeout_seconds" ] = 45 @@ -487,12 +608,76 @@ def test_fluxonkv_test_spec_config(): assert loaded["test_spec_config"]["tcp_thread_reactor_shard_count"] == 2 assert loaded["test_spec_config"]["tcp_thread_bulk_lane_count"] == 4 assert loaded["test_spec_config"]["tcp_thread_control_lane_count"] == 3 + assert loaded["test_spec_config"]["replica_task_max_inflight"] == 16 + assert ( + loaded["test_spec_config"]["ssd_read_source_policy"] + == "local_ssd_only_first" + ) assert ( loaded["test_spec_config"]["require_transfer_rpc_fast_path_ready_timeout_seconds"] == 45 ) assert config.protocol_rdma_device_names == "mlx5_0,mlx5_4" + for removed_policy in ( + "local_ssd_after_memory", + "local_ssd_before_remote_memory", + ): + invalid_policy = copy.deepcopy(rdma_devices) + invalid_policy["test_spec_config"]["ssd_read_source_policy"] = removed_policy + try: + FluxonKvClientConfig(invalid_policy) + raise AssertionError( + f"removed SSD read policy must be rejected: {removed_policy}" + ) + except ValueError: + pass + + expected_capacity = _owner_fluxonkv_base_config(tag="expected_capacity") + expected_capacity["contribute_to_cluster_pool_size"]["dram"] = 137438953472 + expected_capacity["test_spec_config"] = { + "owner_local_reserve_soft_wait_timeout_ms": 10, + "owner_local_reserve_hard_timeout_ms": 10_000, + "owner_local_reserve_expected_capacity": { + "value_len": 4_718_592, + "payload_capacity_bytes": 109_951_162_777, + }, + } + config = FluxonKvClientConfig(expected_capacity) + loaded = yaml.safe_load(config.to_fluxon_kv_client_config_yaml_str()) + assert loaded["test_spec_config"]["owner_local_reserve_soft_wait_timeout_ms"] == 10 + assert loaded["test_spec_config"]["owner_local_reserve_hard_timeout_ms"] == 10_000 + assert loaded["test_spec_config"]["owner_local_reserve_expected_capacity"] == { + "value_len": 4_718_592, + "payload_capacity_bytes": 109_951_162_777, + } + + invalid_reserve_timeout = copy.deepcopy(expected_capacity) + invalid_reserve_timeout["test_spec_config"][ + "owner_local_reserve_hard_timeout_ms" + ] = 10 + try: + FluxonKvClientConfig(invalid_reserve_timeout) + print( + "❌ FAIL: test_fluxonkv_test_spec_config - local-reserve hard timeout should exceed soft timeout" + ) + return + except ValueError: + pass + + invalid_expected_capacity = copy.deepcopy(expected_capacity) + invalid_expected_capacity["test_spec_config"][ + "owner_local_reserve_expected_capacity" + ]["value_len"] = 0 + try: + FluxonKvClientConfig(invalid_expected_capacity) + print( + "❌ FAIL: test_fluxonkv_test_spec_config - expected-capacity value_len=0 should be rejected" + ) + return + except ValueError: + pass + implicit_transport = copy.deepcopy(base) implicit_transport["test_spec_config"] = { "disable_observability": True, diff --git a/fluxon_release/closed_sdk/lib/libfluxon_commu_core.so b/fluxon_release/closed_sdk/lib/libfluxon_commu_core.so index 0d1a1d3..ce9458d 100755 Binary files a/fluxon_release/closed_sdk/lib/libfluxon_commu_core.so and b/fluxon_release/closed_sdk/lib/libfluxon_commu_core.so differ diff --git a/fluxon_rs/Cargo.lock b/fluxon_rs/Cargo.lock index a4b0ecd..59b0cd0 100644 --- a/fluxon_rs/Cargo.lock +++ b/fluxon_rs/Cargo.lock @@ -485,6 +485,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -575,6 +577,12 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" +[[package]] +name = "cmsketch" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7ee2cfacbd29706479902b06d75ad8f1362900836aa32799eabc7e004bfd854" + [[package]] name = "cobs" version = "0.3.0" @@ -615,6 +623,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "cpp_demangle" version = "0.4.5" @@ -960,6 +979,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastant" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" +dependencies = [ + "small_ctor", + "web-time", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -1232,6 +1261,7 @@ dependencies = [ "fluxon_framework_compiled", "fluxon_observability", "fluxon_util", + "foyer", "futures", "hex", "hyper 0.14.32", @@ -1447,6 +1477,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "foreign-types" version = "0.3.2" @@ -1471,6 +1507,130 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "foyer" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0abc0b87814989efa711f9becd9f26969820e2d3905db27d10969c4bd45890" +dependencies = [ + "anyhow", + "equivalent", + "foyer-common", + "foyer-memory", + "foyer-storage", + "foyer-tokio", + "futures-util", + "mea", + "mixtrics", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-common" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3db80d5dece93adb7ad709c84578794724a9cba342a7e566c3551c7ec626789" +dependencies = [ + "anyhow", + "bincode", + "bytes", + "cfg-if", + "foyer-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "serde", + "twox-hash", +] + +[[package]] +name = "foyer-intrusive-collections" +version = "0.10.0-dev" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4fee46bea69e0596130e3210e65d3424e0ac1e6df3bde6636304bdf1ca4a3b" +dependencies = [ + "memoffset", +] + +[[package]] +name = "foyer-memory" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db907f40a527ca2aa2f40a5f68b32ea58aa70f050cd233518e9ffd402cfba6ce" +dependencies = [ + "anyhow", + "bitflags 2.9.1", + "cmsketch", + "equivalent", + "foyer-common", + "foyer-intrusive-collections", + "foyer-tokio", + "futures-util", + "hashbrown 0.16.1", + "itertools 0.14.0", + "mea", + "mixtrics", + "parking_lot", + "paste", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-storage" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1983f1db3d0710e9c9d5fc116d9202dccd41a2d1e032572224f1aff5520aa958" +dependencies = [ + "allocator-api2", + "anyhow", + "bytes", + "core_affinity", + "equivalent", + "fastant", + "foyer-common", + "foyer-memory", + "foyer-tokio", + "fs4", + "futures-core", + "futures-util", + "hashbrown 0.16.1", + "io-uring", + "itertools 0.14.0", + "libc", + "lz4", + "mea", + "parking_lot", + "pin-project", + "rand 0.9.2", + "serde", + "tracing", + "twox-hash", + "zstd", +] + +[[package]] +name = "foyer-tokio" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6577b05a7ffad0db555aedf00bfe52af818220fc4c1c3a7a12520896fc38627" +dependencies = [ + "tokio", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix 1.0.7", + "windows-sys 0.59.0", +] + [[package]] name = "futures" version = "0.3.31" @@ -1618,11 +1778,22 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 5.3.0", "wasi 0.14.2+wasi-0.2.4", "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "gimli" version = "0.31.1" @@ -1711,7 +1882,18 @@ checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" dependencies = [ "allocator-api2", "equivalent", - "foldhash", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", ] [[package]] @@ -2395,6 +2577,17 @@ dependencies = [ "str_stack", ] +[[package]] +name = "io-uring" +version = "0.7.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" +dependencies = [ + "bitflags 2.9.1", + "cfg-if", + "libc", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -2446,12 +2639,40 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.77" @@ -2586,6 +2807,25 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "matchers" version = "0.1.0" @@ -2601,6 +2841,15 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +[[package]] +name = "mea" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2640d335e7273dacdcf51044026139b2e269c3bb0dfc3f8cb3496b85e3f6a42c" +dependencies = [ + "slab", +] + [[package]] name = "memchr" version = "2.7.5" @@ -2668,6 +2917,16 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "mixtrics" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c46b5adfb7a3ae4996d327a5bdc90e78fec025806dd312bdbe6f07a755e0ec9" +dependencies = [ + "itertools 0.15.0", + "parking_lot", +] + [[package]] name = "moka" version = "0.12.11" @@ -2831,6 +3090,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "object" version = "0.36.7" @@ -3498,6 +3767,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.7.3" @@ -4206,9 +4481,15 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "slab" -version = "0.4.10" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04dc19736151f35336d325007ac991178d504a119863a2fcb3758cdb5e52c50d" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "small_ctor" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" [[package]] name = "smallvec" @@ -5080,6 +5361,15 @@ dependencies = [ "utf-8", ] +[[package]] +name = "twox-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" +dependencies = [ + "rand 0.9.2", +] + [[package]] name = "typenum" version = "1.18.0" @@ -5902,3 +6192,31 @@ dependencies = [ "quote", "syn 2.0.104", ] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/fluxon_rs/fluxon_commu/src/facade/p2p.rs b/fluxon_rs/fluxon_commu/src/facade/p2p.rs index 8bcc169..1d4306c 100644 --- a/fluxon_rs/fluxon_commu/src/facade/p2p.rs +++ b/fluxon_rs/fluxon_commu/src/facade/p2p.rs @@ -1297,6 +1297,8 @@ pub mod rpc { timeout: Option, transport_policy: RpcTransportPolicy, ) -> P2PResult> { + validate_explicit_rpc_timeout(timeout)?; + let mut failed_count = 0; let shutdown_poller = p2p.module_view().register_shutdown_poller(); @@ -1371,6 +1373,8 @@ pub mod rpc { timeout: Option, transport_policy: RpcTransportPolicy, ) -> P2PResult> { + validate_explicit_rpc_timeout(timeout)?; + regist_rpc_send::(p2p); { @@ -1381,31 +1385,12 @@ pub mod rpc { .await?; } - let timeout_duration = match timeout { - Some(duration) => { - let min_duration = Duration::from_secs(MIN_EXPLICIT_RPC_TIMEOUT_SECS); - if duration < min_duration { - return Err(P2pError::InvalidRpcTimeout { - timeout_ms: duration.as_secs() * 1_000 - + u64::from(duration.subsec_millis()), - min_timeout_ms: MIN_EXPLICIT_RPC_TIMEOUT_SECS * 1_000, - reason: format!( - "Explicit RPC timeout below {}s is forbidden.", - MIN_EXPLICIT_RPC_TIMEOUT_SECS - ), - }); - } - Some(duration) - } - None => None, - }; - let wire_encode_started_at = Instant::now(); let msg_id = req.msg_id(); let wire_body = req.into_wire_body()?; let wire_encode_us = duration_to_i64_us(wire_encode_started_at.elapsed()); let raw_output = p2p - .call_raw_observed(node, msg_id, wire_body, timeout_duration, transport_policy) + .call_raw_observed(node, msg_id, wire_body, timeout, transport_policy) .await?; let response_local_observe = raw_output.message.local_observe; let resp_decode_started_at = Instant::now(); diff --git a/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs b/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs index 878e5c6..f701659 100644 --- a/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs +++ b/fluxon_rs/fluxon_commu/src/facade/transfer_engine.rs @@ -7,6 +7,9 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use crate::NodeIDString; +pub use fluxon_commu_contract::{ + CLOSED_RUNTIME_DIRECT_FAST_PATH_NOT_READY_MARKER, ClosedRuntimeLocalMemoryKind, +}; use fluxon_commu_contract::{ ClosedRuntimeHandle, ClosedRuntimePeerGen, ClosedRuntimeTransferEngineOpenRuntimeRequest, ClosedRuntimeTransferEngineOpenRuntimeResponse, @@ -74,12 +77,10 @@ impl ClosedLocalSegmentLeaseRegistry { where G: Send + Sync + 'static, { - let boxed = self - .guards - .lock() - .await - .remove(&handle) - .ok_or_else(|| format!("closed sdk local segment lease handle {handle} not found"))?; + let boxed = + self.guards.lock().await.remove(&handle).ok_or_else(|| { + format!("closed sdk local segment lease handle {handle} not found") + })?; boxed.downcast::().map(|guard| *guard).map_err(|_| { format!( "closed sdk local segment lease handle {handle} has unexpected runtime guard type" @@ -379,6 +380,25 @@ impl ClientTransferEngineCore { runtime: R, cpu_mem: &CpuAllocatedMem, ) -> TransferEngineResult<()> + where + R: ClientTransferEngineRuntime, + { + self.register_local_memory( + runtime, + cpu_mem.allocated_addr, + cpu_mem.allocated_size, + ClosedRuntimeLocalMemoryKind::Host, + ) + .await + } + + pub async fn register_local_memory( + &self, + runtime: R, + allocated_addr: u64, + allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, + ) -> TransferEngineResult<()> where R: ClientTransferEngineRuntime, { @@ -386,8 +406,9 @@ impl ClientTransferEngineCore { let handle = self.ensure_closed_runtime_handle(&runtime).await?; transfer_engine_register_local_segment( handle, - cpu_mem.allocated_addr, - cpu_mem.allocated_size, + allocated_addr, + allocated_size, + memory_kind, ) .await .map_err(transfer_engine_closed_sdk_error) @@ -397,13 +418,28 @@ impl ClientTransferEngineCore { pub async fn unregister_local_segment( &self, cpu_mem: &CpuAllocatedMem, + ) -> TransferEngineResult<()> { + self.unregister_local_memory( + cpu_mem.allocated_addr, + cpu_mem.allocated_size, + ClosedRuntimeLocalMemoryKind::Host, + ) + .await + } + + pub async fn unregister_local_memory( + &self, + allocated_addr: u64, + allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, ) -> TransferEngineResult<()> { { if let Some(handle) = self.closed.handle.get().copied() { transfer_engine_unregister_local_segment( handle, - cpu_mem.allocated_addr, - cpu_mem.allocated_size, + allocated_addr, + allocated_size, + memory_kind, ) .await .map_err(transfer_engine_closed_sdk_error) @@ -460,8 +496,9 @@ impl ClientTransferEngineCore { target_addr, len, seg_guard, + false, ) - .await + .await } } @@ -474,6 +511,7 @@ impl ClientTransferEngineCore { target_addr: u64, len: u64, seg_guard: Option, + require_fast_path: bool, ) -> TransferEngineResult where R: ClientTransferEngineRuntime, @@ -482,7 +520,11 @@ impl ClientTransferEngineCore { let initial_local_segment_guard = match seg_guard { Some(guard) => Some(guard), None if runtime.supports_local_segment_transfer() => { - let local_addr = if peer_src_or_target { target_addr } else { src_addr }; + let local_addr = if peer_src_or_target { + target_addr + } else { + src_addr + }; match runtime.ensure_local_segment_guard(local_addr, None).await { Ok(guard) => Some(guard), Err(_) => None, @@ -504,6 +546,7 @@ impl ClientTransferEngineCore { target_addr, len, initial_local_segment_guard_handle, + require_fast_path, ) .await .map_err(transfer_engine_closed_sdk_error); diff --git a/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs b/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs index 6fab54e..7395347 100644 --- a/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs +++ b/fluxon_rs/fluxon_commu_closed_sdk_consumer/src/lib.rs @@ -11,9 +11,9 @@ use fluxon_commu_contract::{ ClosedRuntimeCallRawObservedOutputView, ClosedRuntimeClusterEventStreamItem, ClosedRuntimeClusterManagerCall, ClosedRuntimeClusterManagerResponse, ClosedRuntimeClusterRdmaResolvedConfigStreamItem, ClosedRuntimeDesiredTransferPeer, - ClosedRuntimeDispatchRequestView, - ClosedRuntimeDispatchResponse, ClosedRuntimeDispatchTransportPolicy, ClosedRuntimeError, - ClosedRuntimeHandle, ClosedRuntimeHostCallbackHandle, ClosedRuntimeP2pCall, + ClosedRuntimeDispatchRequestView, ClosedRuntimeDispatchResponse, + ClosedRuntimeDispatchTransportPolicy, ClosedRuntimeError, ClosedRuntimeHandle, + ClosedRuntimeHostCallbackHandle, ClosedRuntimeLocalMemoryKind, ClosedRuntimeP2pCall, ClosedRuntimeP2pCallRawObservedRequestView, ClosedRuntimeP2pResponse, ClosedRuntimeP2pSendResponseRawRequestView, ClosedRuntimePeerGen, ClosedRuntimeRawSlice, ClosedRuntimeRequest, ClosedRuntimeResponse, ClosedRuntimeTransferEngineCall, @@ -26,8 +26,8 @@ use fluxon_commu_contract::{ pub mod rdma_probe; -pub const FLUXON_COMMU_CLOSED_SDK_SCHEMA_VERSION: u32 = 5; -pub const FLUXON_COMMU_CLOSED_ABI_VERSION: u32 = 8; +pub const FLUXON_COMMU_CLOSED_SDK_SCHEMA_VERSION: u32 = 6; +pub const FLUXON_COMMU_CLOSED_ABI_VERSION: u32 = 9; pub const FLUXON_COMMU_CLOSED_HOST_CALLBACKS_ABI_VERSION: u32 = 8; pub const FLUXON_COMMU_CLOSED_RUNTIME_RESULT_OK: i32 = 0; pub const FLUXON_COMMU_CLOSED_RUNTIME_RESULT_ERR: i32 = 1; @@ -491,11 +491,15 @@ impl WireBodyPartsOwner { let (raw_lengths, raw_payload) = match raw_bytes.len() { 0 => (WireBodyRawLengths::Empty, WireBodyRawPayload::Empty), 1 => { - let part = raw_bytes.into_iter().next().expect("single raw part missing"); - let len = - u32::try_from(part.len()).map_err(|_| ClosedSdkConsumerError::RuntimeDecode { + let part = raw_bytes + .into_iter() + .next() + .expect("single raw part missing"); + let len = u32::try_from(part.len()).map_err(|_| { + ClosedSdkConsumerError::RuntimeDecode { detail: format!("wire raw part too large for u32 length: {}", part.len()), - })?; + } + })?; ( WireBodyRawLengths::Single([len]), WireBodyRawPayload::Single(part), @@ -849,8 +853,7 @@ fn decode_call_raw_observed_output_view( return Err(ClosedSdkConsumerError::RuntimeDecode { detail: format!( "closed SDK call_raw_observed serialize_part overflow: serialize_len={} full_len={}", - message_view.body.serialize_part.len, - message_view.body.full_body.len, + message_view.body.serialize_part.len, message_view.body.full_body.len, ), }); } @@ -860,21 +863,19 @@ fn decode_call_raw_observed_output_view( .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { detail: "closed SDK call_raw_observed raw_bytes length overflow".to_string(), })?; - let expected_full_len = - message_view - .body - .serialize_part - .len - .checked_add(raw_total) - .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { - detail: "closed SDK call_raw_observed body length overflow".to_string(), - })?; + let expected_full_len = message_view + .body + .serialize_part + .len + .checked_add(raw_total) + .ok_or_else(|| ClosedSdkConsumerError::RuntimeDecode { + detail: "closed SDK call_raw_observed body length overflow".to_string(), + })?; if expected_full_len != message_view.body.full_body.len { return Err(ClosedSdkConsumerError::RuntimeDecode { detail: format!( "closed SDK call_raw_observed body length mismatch: expected={} full_len={}", - expected_full_len, - message_view.body.full_body.len, + expected_full_len, message_view.body.full_body.len, ), }); } @@ -923,9 +924,7 @@ fn decode_call_raw_observed_output_view( frame_recv_done_ts_us: message_view.local_observe.frame_recv_done_ts_us, dispatch_enqueued_ts_us: message_view.local_observe.dispatch_enqueued_ts_us, dispatch_started_ts_us: message_view.local_observe.dispatch_started_ts_us, - complete_pending_call_ts_us: message_view - .local_observe - .complete_pending_call_ts_us, + complete_pending_call_ts_us: message_view.local_observe.complete_pending_call_ts_us, }, }, observe: fluxon_commu_contract::ClosedRuntimeRpcCallTransportObserveTrace { @@ -1550,8 +1549,8 @@ async fn invoke_completion_async_with_keepalive( ) -> i32, ) -> Result<(i32, Bytes), ClosedSdkConsumerError> { let (sender, receiver) = tokio::sync::oneshot::channel::<(i32, Bytes)>(); - let user_data = Box::into_raw(Box::new(RuntimeCompletionState { sender, keepalive })) - .cast::(); + let user_data = + Box::into_raw(Box::new(RuntimeCompletionState { sender, keepalive })).cast::(); let submit_status = submit(user_data, Some(runtime_completion_callback)); if submit_status != 0 { unsafe { @@ -1867,12 +1866,14 @@ pub async fn transfer_engine_register_local_segment( handle: ClosedRuntimeHandle, allocated_addr: u64, allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, ) -> Result<(), ClosedSdkConsumerError> { match transfer_engine_call( handle, ClosedRuntimeTransferEngineCall::RegisterLocalSegment { allocated_addr, allocated_size, + memory_kind, }, ) .await? @@ -1888,12 +1889,14 @@ pub async fn transfer_engine_unregister_local_segment( handle: ClosedRuntimeHandle, allocated_addr: u64, allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, ) -> Result<(), ClosedSdkConsumerError> { match transfer_engine_call( handle, ClosedRuntimeTransferEngineCall::UnregisterLocalSegment { allocated_addr, allocated_size, + memory_kind, }, ) .await? @@ -1913,6 +1916,7 @@ pub async fn transfer_engine_transfer_data_no_copy( target_addr: u64, len: u64, initial_local_segment_guard_handle: Option, + require_fast_path: bool, ) -> Result { match transfer_engine_call( handle, @@ -1923,6 +1927,7 @@ pub async fn transfer_engine_transfer_data_no_copy( target_addr, len, initial_local_segment_guard_handle, + require_fast_path, }, ) .await? @@ -2082,7 +2087,9 @@ pub async fn p2p_call_raw_observed( ) .await?; match status_code { - FLUXON_COMMU_CLOSED_RUNTIME_RESULT_OK => decode_call_raw_observed_output_view(payload.as_ref()), + FLUXON_COMMU_CLOSED_RUNTIME_RESULT_OK => { + decode_call_raw_observed_output_view(payload.as_ref()) + } FLUXON_COMMU_CLOSED_RUNTIME_RESULT_ERR => { let error = bitcode::decode::(payload.as_ref()).map_err( |decode_error| ClosedSdkConsumerError::RuntimeDecode { diff --git a/fluxon_rs/fluxon_commu_contract/src/closed_runtime.rs b/fluxon_rs/fluxon_commu_contract/src/closed_runtime.rs index 80872e2..17241cf 100644 --- a/fluxon_rs/fluxon_commu_contract/src/closed_runtime.rs +++ b/fluxon_rs/fluxon_commu_contract/src/closed_runtime.rs @@ -744,6 +744,15 @@ pub struct ClosedRuntimeDesiredTransferPeer { pub enable_transfer_segment: bool, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum ClosedRuntimeLocalMemoryKind { + Host, + Gpu { device_id: u32 }, +} + +pub const CLOSED_RUNTIME_DIRECT_FAST_PATH_NOT_READY_MARKER: &str = + "fluxon_direct_fast_path_not_ready"; + #[derive(Debug, Clone, Encode, Decode)] pub enum ClosedRuntimeTransferEngineCall { Init2ForInitDag { @@ -764,10 +773,12 @@ pub enum ClosedRuntimeTransferEngineCall { RegisterLocalSegment { allocated_addr: u64, allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, }, UnregisterLocalSegment { allocated_addr: u64, allocated_size: u64, + memory_kind: ClosedRuntimeLocalMemoryKind, }, TransferDataNoCopy { peer_node: Option, @@ -776,6 +787,7 @@ pub enum ClosedRuntimeTransferEngineCall { target_addr: u64, len: u64, initial_local_segment_guard_handle: Option, + require_fast_path: bool, }, TrySendWireDirect { peer_gen: ClosedRuntimePeerGen, diff --git a/fluxon_rs/fluxon_commu_contract/src/p2p.rs b/fluxon_rs/fluxon_commu_contract/src/p2p.rs index 673692b..c548ee4 100644 --- a/fluxon_rs/fluxon_commu_contract/src/p2p.rs +++ b/fluxon_rs/fluxon_commu_contract/src/p2p.rs @@ -9,15 +9,16 @@ pub mod surface; pub mod wire; pub use rpc::{ - MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, MsgPackSerializePart, RPCReq, RpcCallObserveTrace, - RpcCallObservedOutput, USER_RPC_OBSERVE_TRACE_RAW_BYTES_INDEX, + MIN_EXPLICIT_RPC_TIMEOUT_MS, MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, MsgPackSerializePart, + RPCReq, RpcCallObserveTrace, RpcCallObservedOutput, USER_RPC_OBSERVE_TRACE_RAW_BYTES_INDEX, USER_RPC_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX, USER_RPC_REQ_MSG_ID, USER_RPC_REQUEST_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX, USER_RPC_RESP_MSG_ID, UserRpcBytesAsyncHandler, UserRpcBytesError, UserRpcBytesFuture, UserRpcBytesHandler, UserRpcObserveTrace, UserRpcOwner1ObserveTrace, UserRpcTransportPathKind, current_cross_process_monotonic_us, decode_user_rpc_observe_trace, decode_user_rpc_owner1_observe_trace, encode_user_rpc_observe_trace, - encode_user_rpc_owner1_observe_trace, + encode_user_rpc_owner1_observe_trace, validate_explicit_rpc_timeout, + validate_explicit_rpc_timeout_ms, }; pub use surface::*; pub use wire::*; diff --git a/fluxon_rs/fluxon_commu_contract/src/p2p/rpc.rs b/fluxon_rs/fluxon_commu_contract/src/p2p/rpc.rs index b65ccd9..144034d 100644 --- a/fluxon_rs/fluxon_commu_contract/src/p2p/rpc.rs +++ b/fluxon_rs/fluxon_commu_contract/src/p2p/rpc.rs @@ -17,14 +17,66 @@ use prost::bytes::Bytes as ProstBytes; use std::fmt::Debug; use std::future::Future; use std::pin::Pin; +use std::time::Duration; pub const MIN_EXPLICIT_RPC_TIMEOUT_SECS: u64 = 10; +pub const MIN_EXPLICIT_RPC_TIMEOUT_MS: u64 = MIN_EXPLICIT_RPC_TIMEOUT_SECS * 1_000; pub const USER_RPC_REQ_MSG_ID: u32 = 7001; pub const USER_RPC_RESP_MSG_ID: u32 = 7002; pub const USER_RPC_REQUEST_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX: usize = 1; pub const USER_RPC_OBSERVE_TRACE_RAW_BYTES_INDEX: usize = 1; pub const USER_RPC_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX: usize = 2; +pub fn validate_explicit_rpc_timeout(timeout: Option) -> P2PResult<()> { + let Some(timeout) = timeout else { + return Ok(()); + }; + if timeout < Duration::from_millis(MIN_EXPLICIT_RPC_TIMEOUT_MS) { + return Err(P2pError::InvalidRpcTimeout { + timeout_ms: timeout.as_millis().min(u128::from(u64::MAX)) as u64, + min_timeout_ms: MIN_EXPLICIT_RPC_TIMEOUT_MS, + reason: format!( + "Explicit RPC timeout below {}s is forbidden.", + MIN_EXPLICIT_RPC_TIMEOUT_SECS + ), + }); + } + Ok(()) +} + +pub fn validate_explicit_rpc_timeout_ms(timeout_ms: u64) -> P2PResult<()> { + validate_explicit_rpc_timeout(Some(Duration::from_millis(timeout_ms))) +} + +#[cfg(test)] +mod explicit_rpc_timeout_tests { + use super::*; + + #[test] + fn accepts_default_and_minimum_explicit_timeout() { + assert!(validate_explicit_rpc_timeout(None).is_ok()); + assert!(validate_explicit_rpc_timeout_ms(MIN_EXPLICIT_RPC_TIMEOUT_MS).is_ok()); + } + + #[test] + fn rejects_subminimum_timeout_with_typed_error() { + let timeout_ms = MIN_EXPLICIT_RPC_TIMEOUT_MS - 1; + let err = validate_explicit_rpc_timeout_ms(timeout_ms).unwrap_err(); + assert_eq!(err.code(), 617); + match err { + P2pError::InvalidRpcTimeout { + timeout_ms: actual_timeout_ms, + min_timeout_ms, + .. + } => { + assert_eq!(actual_timeout_ms, timeout_ms); + assert_eq!(min_timeout_ms, MIN_EXPLICIT_RPC_TIMEOUT_MS); + } + other => panic!("expected InvalidRpcTimeout, got {other:?}"), + } + } +} + #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] pub enum UserRpcTransportPathKind { #[default] diff --git a/fluxon_rs/fluxon_fs/src/agent.rs b/fluxon_rs/fluxon_fs/src/agent.rs index a482616..45e15e7 100644 --- a/fluxon_rs/fluxon_fs/src/agent.rs +++ b/fluxon_rs/fluxon_fs/src/agent.rs @@ -1408,13 +1408,24 @@ impl FluxonFsAgent { .id .to_string(); let cache_root_base = if self.kv_framework.is_external_mode() { - self.kv_framework + let shared_file_path = self + .kv_framework .external_client_api_view() .external_client_api() .inner() - .large_file_paths() - .fs_disk_cache_base_dir() - .map_err(|err| format!("invalid external large_file_paths: {}", err))? + .shared_file_path(); + if shared_file_path.is_empty() { + return Err("external shared_file_path is empty".to_string()); + } + let cache_root = Path::new(&shared_file_path).join("fluxon_fs_disk_cache"); + fs::create_dir_all(&cache_root).map_err(|err| { + format!( + "invalid external shared_file_path for fluxon fs disk cache: {} ({})", + cache_root.display(), + err + ) + })?; + cache_root } else { self.kv_framework .client_seg_pool_view() diff --git a/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs b/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs index 1738ade..ca54a71 100644 --- a/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs +++ b/fluxon_rs/fluxon_fs/src/agent_service/transfer_agent.rs @@ -9,28 +9,23 @@ use std::time::{Duration, Instant}; use fluxon_fs_core::config::{ FS_AGENT_TRANSFER_STREAM_CLOSE_RPC_PATH, FS_AGENT_TRANSFER_STREAM_NEXT_RPC_PATH, - FS_AGENT_TRANSFER_STREAM_OPEN_RPC_PATH, - FS_MASTER_TRANSFER_SCHEDULER_HEARTBEAT_RPC_PATH, FS_MASTER_TRANSFER_SCHEDULER_RESULT_RPC_PATH, - FluxonFsTransferBatchCollectInfoWire, FluxonFsTransferBatchKind, - FluxonFsTransferCollectInfoKind, FluxonFsTransferDispositionWire, - FluxonFsTransferFailedFileReasonKindWire, - FluxonFsTransferReadStreamCloseWire, FluxonFsTransferReadStreamNextResultWire, - FluxonFsTransferReadStreamNextWire, FluxonFsTransferReadStreamOpenResultWire, - FluxonFsTransferReadStreamOpenWire, - FluxonFsTransferSkipEntryKind, FluxonFsTransferSkipEntryWire, - FluxonFsTransferManifestEntryWire, FluxonFsTransferManifestWire, - FluxonFsTransferScanMode, - FluxonFsTransferScanEventAckWire, FluxonFsTransferScanEventKindWire, - FluxonFsTransferScanEventWire, FluxonFsTransferScanLaunchResultWire, + FS_AGENT_TRANSFER_STREAM_OPEN_RPC_PATH, FS_MASTER_TRANSFER_SCHEDULER_HEARTBEAT_RPC_PATH, + FS_MASTER_TRANSFER_SCHEDULER_RESULT_RPC_PATH, FluxonFsTransferBatchCollectInfoWire, + FluxonFsTransferBatchKind, FluxonFsTransferCollectInfoKind, FluxonFsTransferDispositionWire, + FluxonFsTransferFailedFileReasonKindWire, FluxonFsTransferManifestEntryWire, + FluxonFsTransferManifestWire, FluxonFsTransferReadStreamCloseWire, + FluxonFsTransferReadStreamNextResultWire, FluxonFsTransferReadStreamNextWire, + FluxonFsTransferReadStreamOpenResultWire, FluxonFsTransferReadStreamOpenWire, FluxonFsTransferScanAssignmentWire, FluxonFsTransferScanBatchWire, - FluxonFsTransferScanChildUnitWire, FluxonFsTransferScanFrontier, + FluxonFsTransferScanChildUnitWire, FluxonFsTransferScanEventAckWire, + FluxonFsTransferScanEventKindWire, FluxonFsTransferScanEventWire, FluxonFsTransferScanFrontier, FluxonFsTransferScanFrontierDirEntry, FluxonFsTransferScanFrontierEntry, - FluxonFsTransferScanResultWire, - FluxonFsTransferSymlinkNoticeEntryWire, FluxonFsTransferWorkerCollectInfoResultWire, - FluxonFsTransferWorkerAssignmentWire, FluxonFsTransferWorkerFileResultWire, - FluxonFsTransferWorkerFailedFileResultWire, - FluxonFsTransferWorkerHeartbeatResultWire, FluxonFsTransferWorkerHeartbeatTelemetryWire, - FluxonFsTransferWorkerHeartbeatWire, + FluxonFsTransferScanLaunchResultWire, FluxonFsTransferScanMode, FluxonFsTransferScanResultWire, + FluxonFsTransferSkipEntryKind, FluxonFsTransferSkipEntryWire, + FluxonFsTransferSymlinkNoticeEntryWire, FluxonFsTransferWorkerAssignmentWire, + FluxonFsTransferWorkerCollectInfoResultWire, FluxonFsTransferWorkerFailedFileResultWire, + FluxonFsTransferWorkerFileResultWire, FluxonFsTransferWorkerHeartbeatResultWire, + FluxonFsTransferWorkerHeartbeatTelemetryWire, FluxonFsTransferWorkerHeartbeatWire, FluxonFsTransferWorkerLaunchResultWire, FluxonFsTransferWorkerResultAckWire, FluxonFsTransferWorkerResultWire, FluxonFsTransferWorkerStopReasonWire, transfer_collect_info_output_relpath, @@ -39,8 +34,8 @@ use fluxon_fs_core::retry::{ BackoffConfig, DEFAULT_WARN_INTERVAL_SECS, WarnConfig, next_backoff, should_warn, }; use fluxon_kv::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError}; -use fluxon_kv::user_api::flat_dict::{FlatDict, FlatValue}; use fluxon_kv::user_api::FluxonUserApi; +use fluxon_kv::user_api::flat_dict::{FlatDict, FlatValue}; use parking_lot::{Condvar, Mutex}; use super::{ @@ -202,16 +197,13 @@ fn transfer_scan_session_state() -> &'static Mutex { TRANSFER_SCAN_SESSION_STATE.get_or_init(|| Mutex::new(TransferScanSessionState::default())) } -fn cleanup_expired_transfer_scan_sessions( - state: &mut TransferScanSessionState, - now_unix_ms: i64, -) { - state - .root_dir_listing_sessions - .retain(|_, session| session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms); - state - .subtree_streaming_sessions - .retain(|_, session| session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms); +fn cleanup_expired_transfer_scan_sessions(state: &mut TransferScanSessionState, now_unix_ms: i64) { + state.root_dir_listing_sessions.retain(|_, session| { + session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms + }); + state.subtree_streaming_sessions.retain(|_, session| { + session.lease_expire_unix_ms <= 0 || session.lease_expire_unix_ms > now_unix_ms + }); } fn same_root_continuation_scan_unit( @@ -301,10 +293,7 @@ fn flush_pending_root_direct_files_batch( return Ok(None); } let batch = build_direct_files_only_batch_from_entries_with_batch_id( - direct_files_only_batch_id_for_partition( - assignment, - session.next_direct_files_batch_index, - ), + direct_files_only_batch_id_for_partition(assignment, session.next_direct_files_batch_index), assignment, assignment.root_relpath.clone(), std::mem::take(&mut session.pending_direct_files), @@ -313,7 +302,8 @@ fn flush_pending_root_direct_files_batch( )?; session.pending_direct_bytes = 0; session.next_direct_files_batch_index = session.next_direct_files_batch_index.saturating_add(1); - session.emitted_direct_files_batch_count = session.emitted_direct_files_batch_count.saturating_add(1); + session.emitted_direct_files_batch_count = + session.emitted_direct_files_batch_count.saturating_add(1); Ok(Some(batch)) } @@ -414,7 +404,8 @@ fn open_transfer_root_dir_listing_session( root_dir_abs: &str, assignment: &FluxonFsTransferScanAssignmentWire, ) -> Result, FlatDict> { - let dir_abs = safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; + let dir_abs = + safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; let read_dir = match retry_after_target_path_chmod( dir_abs.as_path(), "root_read_dir", @@ -458,7 +449,10 @@ fn take_transfer_root_dir_listing_session( let now_unix_ms = chrono::Utc::now().timestamp_millis(); let mut state = transfer_scan_session_state().lock(); cleanup_expired_transfer_scan_sessions(&mut state, now_unix_ms); - if let Some(mut session) = state.root_dir_listing_sessions.remove(assignment.scan_unit_id.as_str()) { + if let Some(mut session) = state + .root_dir_listing_sessions + .remove(assignment.scan_unit_id.as_str()) + { if session.job_id == assignment.job_id && session.scan_epoch == assignment.scan_epoch && session.root_relpath == assignment.root_relpath @@ -507,7 +501,8 @@ fn open_transfer_subtree_streaming_session( if is_relpath_skipped(&assignment.skip_entries, assignment.root_relpath.as_str()) { return Ok(None); } - let dir_abs = safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; + let dir_abs = + safe_join_root(root_dir_abs, assignment.root_relpath.as_str()).map_err(resp_err_kverr)?; let root_md = retry_after_target_path_chmod( Path::new(root_dir_abs), "subtree_stream_root_symlink_metadata", @@ -790,7 +785,8 @@ fn collect_transfer_root_dir_listing_slice( assignment: &FluxonFsTransferScanAssignmentWire, deadline: Option, ) -> Result { - let Some(mut session) = take_transfer_root_dir_listing_session(root_dir_abs, assignment)? else { + let Some(mut session) = take_transfer_root_dir_listing_session(root_dir_abs, assignment)? + else { return Ok(TransferRootDirListingOutcome::Finished( build_finished_empty_transfer_scan_result(assignment), )); @@ -848,7 +844,8 @@ fn collect_transfer_root_dir_listing_slice( }; scanned_entries = scanned_entries.saturating_add(1); let name = ent.file_name().to_string_lossy().to_string(); - let child_relpath = normalize_child_relpath(assignment.root_relpath.as_str(), name.as_str()); + let child_relpath = + normalize_child_relpath(assignment.root_relpath.as_str(), name.as_str()); if is_relpath_skipped(&assignment.skip_entries, child_relpath.as_str()) { continue; } @@ -899,10 +896,12 @@ fn collect_transfer_root_dir_listing_slice( let size = md.len().min(i64::MAX as u64) as i64; session.root_visible_entries = true; session.root_total_bytes = session.root_total_bytes.saturating_add(size); - session.pending_direct_files.push(FluxonFsTransferScanFrontierEntry { - relpath: child_relpath, - size, - }); + session + .pending_direct_files + .push(FluxonFsTransferScanFrontierEntry { + relpath: child_relpath, + size, + }); session.pending_direct_bytes = session.pending_direct_bytes.saturating_add(size); if should_flush_direct_batch( assignment.batch_ready_bytes, @@ -910,7 +909,9 @@ fn collect_transfer_root_dir_listing_slice( session.pending_direct_files.len(), session.pending_direct_empty_dirs.len(), ) { - if let Some(batch) = flush_pending_root_direct_files_batch(assignment, &mut session)? { + if let Some(batch) = + flush_pending_root_direct_files_batch(assignment, &mut session)? + { direct_files_only_batches.push(batch); } } @@ -933,14 +934,18 @@ fn collect_transfer_root_dir_listing_slice( session.pending_direct_files.len(), session.pending_direct_empty_dirs.len(), ) { - if let Some(batch) = flush_pending_root_direct_files_batch(assignment, &mut session)? { + if let Some(batch) = + flush_pending_root_direct_files_batch(assignment, &mut session)? + { direct_files_only_batches.push(batch); } } } else { - session.direct_dirs.push(FluxonFsTransferScanFrontierDirEntry { - relpath: child_relpath, - }); + session + .direct_dirs + .push(FluxonFsTransferScanFrontierDirEntry { + relpath: child_relpath, + }); } } } @@ -1244,12 +1249,14 @@ impl TransferWorkerProgressWindow { fn record_written_bytes_and_maybe_ramp(&self, bytes: i64, now_unix_ms: i64) { let normalized = bytes.max(0); self.window_bytes.fetch_add(normalized, Ordering::SeqCst); - self.total_written_bytes.fetch_add(normalized, Ordering::SeqCst); + self.total_written_bytes + .fetch_add(normalized, Ordering::SeqCst); self.maybe_ramp(now_unix_ms); } fn record_materialized_empty_dir(&self) { - self.total_materialized_empty_dirs.fetch_add(1, Ordering::SeqCst); + self.total_materialized_empty_dirs + .fetch_add(1, Ordering::SeqCst); } fn total_materialized_empty_dirs(&self) -> i64 { @@ -1298,8 +1305,9 @@ impl TransferWorkerProgressWindow { } if previous_goodput > 0 { let delta = current_goodput.saturating_sub(previous_goodput); - let improvement_percent = - delta.saturating_mul(100).saturating_div(previous_goodput.max(1)); + let improvement_percent = delta + .saturating_mul(100) + .saturating_div(previous_goodput.max(1)); if improvement_percent < self.policy.min_improvement_percent { return; } @@ -1335,10 +1343,8 @@ impl TransferWorkerProgressWindow { .saturating_mul(1000) .saturating_div(window_elapsed_ms.max(1)) }; - self.peak_sample_goodput_bytes_per_sec.fetch_max( - window_goodput_bytes_per_sec.max(0), - Ordering::SeqCst, - ); + self.peak_sample_goodput_bytes_per_sec + .fetch_max(window_goodput_bytes_per_sec.max(0), Ordering::SeqCst); Some(TransferWorkerThroughputSample { window_started_unix_ms, window_elapsed_ms, @@ -1448,7 +1454,8 @@ impl TransferReadStreamActorOwned { data: Vec::new(), }); } - self.fill_prefetch_queue().map_err(|err| self.cache_terminal_error(err))?; + self.fill_prefetch_queue() + .map_err(|err| self.cache_terminal_error(err))?; let to_take = std::cmp::min(length as usize, (self.file_size - next_offset) as usize); let buf = self .take_prefetched_bytes(to_take) @@ -1457,7 +1464,8 @@ impl TransferReadStreamActorOwned { self.replay_offset = next_offset; self.replay_data = buf.clone(); self.next_offset = next_offset.saturating_add(buf.len() as i64); - self.fill_prefetch_queue().map_err(|err| self.cache_terminal_error(err))?; + self.fill_prefetch_queue() + .map_err(|err| self.cache_terminal_error(err))?; Ok(FluxonFsTransferReadStreamNextResultWire { stream_missing: false, data: buf, @@ -1598,7 +1606,11 @@ impl TransferReadStreamActorHandle { struct TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { log_context: TransferWorkerLogContext, @@ -1611,7 +1623,11 @@ where impl TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { fn new( @@ -1693,7 +1709,8 @@ where } fn progress_snapshot(&self) -> TransferWorkerProgressSnapshot { - self.progress.snapshot(chrono::Utc::now().timestamp_millis()) + self.progress + .snapshot(chrono::Utc::now().timestamp_millis()) } fn stop(&self) { @@ -1737,7 +1754,8 @@ impl TransferReadStreamRegistryHandle { }); } } - let full_path = safe_join_root(root_dir_abs, open.relpath.as_str()).map_err(resp_err_kverr)?; + let full_path = + safe_join_root(root_dir_abs, open.relpath.as_str()).map_err(resp_err_kverr)?; let file = open_file_with_target_path_chmod_retry(&full_path, "open_stream")?; let md = file.metadata().map_err(resp_err_io)?; let file_size = md.len().min(i64::MAX as u64) as i64; @@ -1764,12 +1782,15 @@ impl TransferReadStreamRegistryHandle { }); } state.streams.insert(stream_id.clone(), entry); - state.dedup_by_worker_file.insert(dedup_key, stream_id.clone()); + state + .dedup_by_worker_file + .insert(dedup_key, stream_id.clone()); drop(state); if let Err(resp) = TransferReadStreamActorHandle::start(stream_id.as_str(), actor) { let mut state = self.state.lock(); state.streams.remove(stream_id.as_str()); - state.dedup_by_worker_file + state + .dedup_by_worker_file .retain(|_, existing_stream_id| existing_stream_id != &stream_id); return Err(resp); } @@ -1811,7 +1832,9 @@ impl TransferReadStreamRegistryHandle { let Some(entry) = state.streams.remove(stream_id) else { return; }; - state.dedup_by_worker_file.retain(|_, existing_stream_id| existing_stream_id != stream_id); + state + .dedup_by_worker_file + .retain(|_, existing_stream_id| existing_stream_id != stream_id); entry.close(); } } @@ -1894,20 +1917,22 @@ fn decode_transfer_stream_open_result_payload( return Err(TransferWorkerRpcFailure::Fatal(resp.clone())); } Ok(FluxonFsTransferReadStreamOpenResultWire { - stream_id: require_str(resp, "stream_id").map_err(resp_err_kverr).map_err( - |err| { + stream_id: require_str(resp, "stream_id") + .map_err(resp_err_kverr) + .map_err(|err| { invalid_transfer_rpc_response(format!( "transfer read stream open response missing stream_id: err={}", transfer_rpc_response_err_text(&err) )) - }, - )?, - size: require_i64(resp, "size").map_err(resp_err_kverr).map_err(|err| { - invalid_transfer_rpc_response(format!( - "transfer read stream open response missing size: err={}", - transfer_rpc_response_err_text(&err) - )) - })?, + })?, + size: require_i64(resp, "size") + .map_err(resp_err_kverr) + .map_err(|err| { + invalid_transfer_rpc_response(format!( + "transfer read stream open response missing size: err={}", + transfer_rpc_response_err_text(&err) + )) + })?, }) } @@ -1980,11 +2005,15 @@ fn is_relpath_skipped(skip_entries: &[FluxonFsTransferSkipEntryWire], relpath: & } fn file_name_from_relpath(relpath: &str) -> Result<&str, FlatDict> { - relpath.rsplit('/').next().filter(|v| !v.is_empty()).ok_or_else(|| { - resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("relpath must contain file name: {}", relpath), - })) - }) + relpath + .rsplit('/') + .next() + .filter(|v| !v.is_empty()) + .ok_or_else(|| { + resp_err_kverr(KvError::Api(ApiError::InvalidArgument { + detail: format!("relpath must contain file name: {}", relpath), + })) + }) } fn transfer_staging_dir_for_file(staging_prefix: &str, relpath: &str) -> String { @@ -2106,8 +2135,12 @@ where match attempt() { Ok(value) => Ok(value), Err(initial_err) if initial_err.kind() == ErrorKind::PermissionDenied => { - let repair_dir = - repair_permission_denied_dir_for_retry(repair_anchor, op, target_path, &initial_err)?; + let repair_dir = repair_permission_denied_dir_for_retry( + repair_anchor, + op, + target_path, + &initial_err, + )?; attempt().map_err(|retry_err| { resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( @@ -2334,10 +2367,11 @@ fn collect_transfer_tree_with_deadline( continue; } }; - out.symlink_notices.push(FluxonFsTransferSymlinkNoticeEntryWire { - relpath: child_rel, - link_target: link_target.to_string_lossy().to_string(), - }); + out.symlink_notices + .push(FluxonFsTransferSymlinkNoticeEntryWire { + relpath: child_rel, + link_target: link_target.to_string_lossy().to_string(), + }); continue; } if md.is_dir() { @@ -2361,7 +2395,8 @@ fn collect_transfer_tree_with_deadline( } out.files.sort_by(|a, b| a.relpath.cmp(&b.relpath)); out.empty_dirs.sort(); - out.symlink_notices.sort_by(|a, b| a.relpath.cmp(&b.relpath)); + out.symlink_notices + .sort_by(|a, b| a.relpath.cmp(&b.relpath)); Ok(out) } @@ -2574,10 +2609,8 @@ fn build_transfer_scan_events_for_result( event_seq_no_start: i64, result: FluxonFsTransferScanResultWire, ) -> (Vec, bool, i64) { - let (child_scan_units, continue_locally) = split_same_root_continuation_from_child_scan_units( - assignment, - result.child_scan_units, - ); + let (child_scan_units, continue_locally) = + split_same_root_continuation_from_child_scan_units(assignment, result.child_scan_units); if continue_locally { let event = build_transfer_scan_event( assignment, @@ -2588,11 +2621,7 @@ fn build_transfer_scan_events_for_result( result.full_dir_batches, String::new(), ); - return ( - vec![event], - true, - event_seq_no_start.saturating_add(1), - ); + return (vec![event], true, event_seq_no_start.saturating_add(1)); } let mut next_event_seq_no = event_seq_no_start; let mut events = Vec::new(); @@ -2620,11 +2649,7 @@ fn build_transfer_scan_events_for_result( Vec::new(), String::new(), )); - ( - events, - false, - next_event_seq_no.saturating_add(1), - ) + (events, false, next_event_seq_no.saturating_add(1)) } fn send_transfer_scan_event_once( @@ -2637,10 +2662,7 @@ fn send_transfer_scan_event_once( } let event_json = serde_json::to_string(event) .map_err(|e| format!("serialize transfer scan event failed: {}", e))?; - let payload = FlatDict::from([( - "scan_event_json".to_string(), - FlatValue::String(event_json), - )]); + let payload = FlatDict::from([("scan_event_json".to_string(), FlatValue::String(event_json))]); let resp = api .rpc_client() .call( @@ -2793,30 +2815,28 @@ fn run_transfer_scan_background_task( } let mut next_event_seq_no = 1_i64; loop { - let result = match build_transfer_scan_result_for_root_dir_abs( - root_dir_abs.as_str(), - &assignment, - ) { - Ok(v) => v, - Err(resp) => { - let failed = build_transfer_scan_event( - &assignment, - next_event_seq_no, - FluxonFsTransferScanEventKindWire::Failed, - Vec::new(), - Vec::new(), - Vec::new(), - transfer_rpc_response_err_text(&resp), - ); - let _ = send_transfer_scan_event_with_retry( - api.as_ref(), - master_id.as_str(), - &mut assignment, - &failed, - ); - break; - } - }; + let result = + match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { + Ok(v) => v, + Err(resp) => { + let failed = build_transfer_scan_event( + &assignment, + next_event_seq_no, + FluxonFsTransferScanEventKindWire::Failed, + Vec::new(), + Vec::new(), + Vec::new(), + transfer_rpc_response_err_text(&resp), + ); + let _ = send_transfer_scan_event_with_retry( + api.as_ref(), + master_id.as_str(), + &mut assignment, + &failed, + ); + break; + } + }; let (events, continue_locally, next_seq_no_after_events) = build_transfer_scan_events_for_result(&assignment, next_event_seq_no, result); next_event_seq_no = next_seq_no_after_events; @@ -2922,17 +2942,14 @@ impl TransferScanRegistryHandle { let assignment2 = assignment.clone(); let thread_name = format!("fluxon_fs_transfer_scan_{}", assignment.scan_task_id); match thread::Builder::new().name(thread_name).spawn(move || { - run_transfer_scan_background_task( - registry, - api2, - master_id2, - exports2, - assignment2, - ); + run_transfer_scan_background_task(registry, api2, master_id2, exports2, assignment2); }) { Ok(_) => Ok(FluxonFsTransferScanLaunchResultWire::started()), Err(err) => { - self.state.lock().tasks.remove(assignment.scan_task_id.as_str()); + self.state + .lock() + .tasks + .remove(assignment.scan_task_id.as_str()); Err(resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( "spawn transfer scan thread failed: scan_task_id={} err={}", @@ -3006,22 +3023,16 @@ impl TransferWorkerRegistryHandle { let master_id2 = master_id.to_string(); let exports2 = exports.clone(); let assignment2 = assignment.clone(); - let thread_name = format!( - "fluxon_fs_transfer_worker_{}", - assignment.worker_task_id - ); + let thread_name = format!("fluxon_fs_transfer_worker_{}", assignment.worker_task_id); match thread::Builder::new().name(thread_name).spawn(move || { - run_transfer_worker_background_task( - registry, - api2, - master_id2, - exports2, - assignment2, - ); + run_transfer_worker_background_task(registry, api2, master_id2, exports2, assignment2); }) { Ok(_) => Ok(FluxonFsTransferWorkerLaunchResultWire::started()), Err(err) => { - self.state.lock().tasks.remove(assignment.worker_task_id.as_str()); + self.state + .lock() + .tasks + .remove(assignment.worker_task_id.as_str()); Err(resp_err_kverr(KvError::Api(ApiError::Unknown { detail: format!( "spawn transfer worker thread failed: worker_task_id={} err={}", @@ -3298,10 +3309,7 @@ fn send_transfer_worker_result_once( detail: format!("serialize transfer worker result failed: {}", e), }))) })?; - let payload = FlatDict::from([( - "result_json".to_string(), - FlatValue::String(result_json), - )]); + let payload = FlatDict::from([("result_json".to_string(), FlatValue::String(result_json))]); let resp = api .rpc_client() .call( @@ -3396,7 +3404,10 @@ fn open_transfer_read_stream_via_rpc_once( "relpath".to_string(), FlatValue::String(file.relpath.clone()), ), - ("initial_offset".to_string(), FlatValue::Int64(initial_offset)), + ( + "initial_offset".to_string(), + FlatValue::Int64(initial_offset), + ), ]); let resp = api .rpc_client() @@ -3518,25 +3529,25 @@ impl TransferWorkerRemoteControl { loop { self.before_heartbeat_retry_attempt()?; let current_materialized_empty_dirs = self.progress.total_materialized_empty_dirs(); - match self - .heartbeat - .ensure_continue( - force, - current_materialized_empty_dirs, - |heartbeat_unix_ms, _heartbeat_detail| { - let progress_snapshot = - self.progress.snapshot(chrono::Utc::now().timestamp_millis()); - let telemetry = - Some(transfer_worker_telemetry_from_progress_snapshot(&progress_snapshot)); - send_transfer_worker_heartbeat_once( - self.api.as_ref(), - self.master_id.as_str(), - &self.assignment, - heartbeat_unix_ms, - telemetry, - ) - }, - ) { + match self.heartbeat.ensure_continue( + force, + current_materialized_empty_dirs, + |heartbeat_unix_ms, _heartbeat_detail| { + let progress_snapshot = self + .progress + .snapshot(chrono::Utc::now().timestamp_millis()); + let telemetry = Some(transfer_worker_telemetry_from_progress_snapshot( + &progress_snapshot, + )); + send_transfer_worker_heartbeat_once( + self.api.as_ref(), + self.master_id.as_str(), + &self.assignment, + heartbeat_unix_ms, + telemetry, + ) + }, + ) { Ok(()) => return Ok(()), Err(TransferWorkerHeartbeatGateError::Terminal(err)) => return Err(err), Err(TransferWorkerHeartbeatGateError::Retryable { @@ -3631,10 +3642,7 @@ impl TransferWorkerRemoteControl { ) } - fn close_stream_with_retry( - &self, - stream_id: &str, - ) -> Result<(), TransferWorkerExecutionError> { + fn close_stream_with_retry(&self, stream_id: &str) -> Result<(), TransferWorkerExecutionError> { let api = self.api.clone(); let assignment = self.assignment.clone(); let op_detail = format!( @@ -3744,9 +3752,9 @@ impl TransferWorkerRemoteControl { if ack.accepted { return Ok(()); } - Err(TransferWorkerExecutionError::Stop(stop_reason_or_superseded( - ack.stop_reason, - ))) + Err(TransferWorkerExecutionError::Stop( + stop_reason_or_superseded(ack.stop_reason), + )) } } @@ -3825,10 +3833,12 @@ impl TransferWorkerHeartbeatGate { mut heartbeat_op: HeartbeatOp, ) -> Result<(), TransferWorkerHeartbeatGateError> where - HeartbeatOp: FnMut( - i64, - &'static str, - ) -> Result, + HeartbeatOp: + FnMut( + i64, + &'static str, + ) + -> Result, { loop { let (heartbeat_unix_ms, heartbeat_detail) = { @@ -3875,15 +3885,13 @@ impl TransferWorkerHeartbeatGate { state.heartbeat_inflight = false; let result = match heartbeat_result { Ok(heartbeat_result) if heartbeat_result.continue_running => { - state.last_heartbeat_completed_unix_ms = - chrono::Utc::now().timestamp_millis(); + state.last_heartbeat_completed_unix_ms = chrono::Utc::now().timestamp_millis(); state.last_heartbeat_materialized_empty_dirs = current_materialized_empty_dirs; state.granted_lease_expire_unix_ms = heartbeat_result.lease_expire_unix_ms; Ok(()) } Ok(heartbeat_result) => { - state.last_heartbeat_completed_unix_ms = - chrono::Utc::now().timestamp_millis(); + state.last_heartbeat_completed_unix_ms = chrono::Utc::now().timestamp_millis(); state.last_heartbeat_materialized_empty_dirs = current_materialized_empty_dirs; let reason = stop_reason_or_superseded(heartbeat_result.stop_reason); state.terminal_state = @@ -3933,7 +3941,8 @@ fn run_transfer_worker_background_task( )); let dedup_expire_unix_ms = match control.ensure_continue(true) { Ok(()) => { - let dst_export_root = match exports.export_root_dir_abs(assignment.dst_export.as_str()) { + let dst_export_root = match exports.export_root_dir_abs(assignment.dst_export.as_str()) + { Ok(v) => v, Err(err) => { tracing::warn!( @@ -3996,7 +4005,9 @@ fn run_transfer_worker_background_task( }, { let control = control.clone(); - move |file, read_offset, length| control.read_chunk_with_retry(file, read_offset, length) + move |file, read_offset, length| { + control.read_chunk_with_retry(file, read_offset, length) + } }, ) { Ok(result) => { @@ -4004,34 +4015,38 @@ fn run_transfer_worker_background_task( if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("before_result_submit", &assignment, &resp); - } - match control.submit_result_with_retry(&result) { - Ok(()) => control.dedup_expire_unix_ms(), - Err(TransferWorkerExecutionError::Stop(reason)) => { - tracing::info!( - "transfer worker result submission stopped: job_id={} batch_id={} worker_id={} worker_task_id={} reason={:?}", - assignment.job_id, - assignment.batch_id, - assignment.worker_id, - assignment.worker_task_id, - reason + log_transfer_worker_cleanup_failure( + "before_result_submit", + &assignment, + &resp, ); - control.dedup_expire_unix_ms() } - Err(TransferWorkerExecutionError::Fatal(resp)) => { - tracing::warn!( - "transfer worker result submission failed: job_id={} batch_id={} worker_id={} worker_task_id={} resp={:?}", - assignment.job_id, - assignment.batch_id, - assignment.worker_id, - assignment.worker_task_id, - resp - ); - control.dedup_expire_unix_ms() + match control.submit_result_with_retry(&result) { + Ok(()) => control.dedup_expire_unix_ms(), + Err(TransferWorkerExecutionError::Stop(reason)) => { + tracing::info!( + "transfer worker result submission stopped: job_id={} batch_id={} worker_id={} worker_task_id={} reason={:?}", + assignment.job_id, + assignment.batch_id, + assignment.worker_id, + assignment.worker_task_id, + reason + ); + control.dedup_expire_unix_ms() + } + Err(TransferWorkerExecutionError::Fatal(resp)) => { + tracing::warn!( + "transfer worker result submission failed: job_id={} batch_id={} worker_id={} worker_task_id={} resp={:?}", + assignment.job_id, + assignment.batch_id, + assignment.worker_id, + assignment.worker_task_id, + resp + ); + control.dedup_expire_unix_ms() + } } } - } Err(TransferWorkerExecutionError::Stop(reason)) => { control.close_all_streams(); if let Err(resp) = @@ -4054,10 +4069,13 @@ fn run_transfer_worker_background_task( if let Err(cleanup_resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("after_fatal", &assignment, &cleanup_resp); + log_transfer_worker_cleanup_failure( + "after_fatal", + &assignment, + &cleanup_resp, + ); } - if let Some((fatal_kind, fatal_message)) = - classify_transfer_worker_fatal(&resp) + if let Some((fatal_kind, fatal_message)) = classify_transfer_worker_fatal(&resp) { match report_transfer_worker_fatal_once( control.api.as_ref(), @@ -4107,15 +4125,25 @@ fn run_transfer_worker_background_task( } } Err(TransferWorkerExecutionError::Stop(reason)) => { - let dst_export_root = exports.export_root_dir_abs(assignment.dst_export.as_str()).ok(); + let dst_export_root = exports + .export_root_dir_abs(assignment.dst_export.as_str()) + .ok(); let dst_root = dst_export_root.and_then(|dst_export_root| { - safe_join_root(dst_export_root.as_str(), assignment.dst_root_relpath.as_str()) - .ok() - .map(PathBuf::from) + safe_join_root( + dst_export_root.as_str(), + assignment.dst_root_relpath.as_str(), + ) + .ok() + .map(PathBuf::from) }); if let Some(dst_root) = dst_root { - if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) { - log_transfer_worker_cleanup_failure("before_execution_stop", &assignment, &resp); + if let Err(resp) = cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment) + { + log_transfer_worker_cleanup_failure( + "before_execution_stop", + &assignment, + &resp, + ); } } tracing::info!( @@ -4129,11 +4157,16 @@ fn run_transfer_worker_background_task( control.dedup_expire_unix_ms() } Err(TransferWorkerExecutionError::Fatal(resp)) => { - let dst_export_root = exports.export_root_dir_abs(assignment.dst_export.as_str()).ok(); + let dst_export_root = exports + .export_root_dir_abs(assignment.dst_export.as_str()) + .ok(); let dst_root = dst_export_root.and_then(|dst_export_root| { - safe_join_root(dst_export_root.as_str(), assignment.dst_root_relpath.as_str()) - .ok() - .map(PathBuf::from) + safe_join_root( + dst_export_root.as_str(), + assignment.dst_root_relpath.as_str(), + ) + .ok() + .map(PathBuf::from) }); if let Some(dst_root) = dst_root { if let Err(cleanup_resp) = @@ -4509,7 +4542,9 @@ fn plan_transfer_subtree_batches( total_bytes: 0, root_is_empty: true, mergeable_empty_dir_count: 1, - mergeable_empty_dir_estimated_bytes: estimate_empty_dir_manifest_entry_bytes(root_relpath), + mergeable_empty_dir_estimated_bytes: estimate_empty_dir_manifest_entry_bytes( + root_relpath, + ), direct_files_only_batches: Vec::new(), full_dir_batches: Vec::new(), child_scan_units: Vec::new(), @@ -4540,8 +4575,7 @@ fn plan_transfer_subtree_batches( let child_empty_dir_count = child_plan.mergeable_empty_dir_count; let child_empty_dir_estimated_bytes = child_plan.mergeable_empty_dir_estimated_bytes; - if mergeable_empty_dir_count - .saturating_add(child_empty_dir_count) + if mergeable_empty_dir_count.saturating_add(child_empty_dir_count) > TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET || mergeable_empty_dir_estimated_bytes .saturating_add(child_empty_dir_estimated_bytes) @@ -4716,7 +4750,11 @@ fn build_root_direct_files_only_batch_from_entries( } fn sort_transfer_scan_batches(batches: &mut [FluxonFsTransferScanBatchWire]) { - batches.sort_by(|a, b| a.root_relpath.cmp(&b.root_relpath).then(a.batch_id.cmp(&b.batch_id))); + batches.sort_by(|a, b| { + a.root_relpath + .cmp(&b.root_relpath) + .then(a.batch_id.cmp(&b.batch_id)) + }); } fn build_full_dir_batch_for_mergeable_subtree( @@ -4750,14 +4788,12 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( root_dir_abs: &str, assignment: &FluxonFsTransferScanAssignmentWire, ) -> Result { - let Some(mut session) = take_transfer_subtree_streaming_session(root_dir_abs, assignment)? else { + let Some(mut session) = take_transfer_subtree_streaming_session(root_dir_abs, assignment)? + else { return Ok(build_finished_empty_subtree_stream_result(assignment)); }; loop { - if session - .dir_stack - .is_empty() - { + if session.dir_stack.is_empty() { let mut full_dir_batches = Vec::new(); if let Some(batch) = flush_pending_subtree_stream_batch(assignment, &mut session)? { full_dir_batches.push(batch); @@ -4776,7 +4812,9 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( finished: true, }); } - if TransferScanDeadline::from_assignment(assignment).is_some_and(|deadline| deadline.reached()) { + if TransferScanDeadline::from_assignment(assignment) + .is_some_and(|deadline| deadline.reached()) + { let mut full_dir_batches = Vec::new(); if let Some(batch) = flush_pending_subtree_stream_batch(assignment, &mut session)? { full_dir_batches.push(batch); @@ -4808,7 +4846,10 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( if should_flush_subtree_stream_batch( assignment.batch_ready_bytes, session.pending_bytes, - session.pending_files.len().saturating_add(session.pending_symlink_notices.len()), + session + .pending_files + .len() + .saturating_add(session.pending_symlink_notices.len()), session.pending_empty_dirs.len(), ) { let batch = flush_pending_subtree_stream_batch(assignment, &mut session)?.unwrap(); @@ -4865,22 +4906,30 @@ fn build_transfer_scan_result_for_subtree_streaming_root_dir_abs( }); } else if md.is_dir() { frame.saw_visible_child = true; - session.dir_stack.push(open_transfer_subtree_streaming_dir_frame( - child_path, - child_relpath, - )?); + session + .dir_stack + .push(open_transfer_subtree_streaming_dir_frame( + child_path, + child_relpath, + )?); } else if md.is_file() { frame.saw_visible_child = true; let size = md.len().min(i64::MAX as u64) as i64; session.pending_bytes = session.pending_bytes.saturating_add(size); session .pending_files - .push(FluxonFsTransferScanFrontierEntry { relpath: child_relpath, size }); + .push(FluxonFsTransferScanFrontierEntry { + relpath: child_relpath, + size, + }); } if should_flush_subtree_stream_batch( assignment.batch_ready_bytes, session.pending_bytes, - session.pending_files.len().saturating_add(session.pending_symlink_notices.len()), + session + .pending_files + .len() + .saturating_add(session.pending_symlink_notices.len()), session.pending_empty_dirs.len(), ) { let batch = flush_pending_subtree_stream_batch(assignment, &mut session)?.unwrap(); @@ -4928,11 +4977,12 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( ); } let deadline = TransferScanDeadline::from_assignment(assignment); - let root_listing = match collect_transfer_root_dir_listing_slice(root_dir_abs, assignment, deadline)? { - TransferRootDirListingOutcome::Complete(v) => v, - TransferRootDirListingOutcome::Finished(result) => return Ok(result), - TransferRootDirListingOutcome::Partial(result) => return Ok(result), - }; + let root_listing = + match collect_transfer_root_dir_listing_slice(root_dir_abs, assignment, deadline)? { + TransferRootDirListingOutcome::Complete(v) => v, + TransferRootDirListingOutcome::Finished(result) => return Ok(result), + TransferRootDirListingOutcome::Partial(result) => return Ok(result), + }; let mut direct_files = root_listing.direct_files; let mut direct_symlink_notices = root_listing.direct_symlink_notices; let mut direct_empty_dirs = root_listing.direct_empty_dirs; @@ -4976,7 +5026,10 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( if (!direct_files.is_empty() || !direct_symlink_notices.is_empty() || !direct_empty_dirs.is_empty()) - && !direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()) + && !direct_files_only_disposition_covers_root( + assignment, + assignment.root_relpath.as_str(), + ) { let mut next_partition_index = root_listing.emitted_direct_files_batch_count; direct_files_only_batches.extend(build_partitioned_root_direct_files_only_batches( @@ -4987,17 +5040,15 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( direct_empty_dirs.clone(), )?); } - child_scan_units.extend( - direct_dirs[delegated_child_scan_unit_count..] - .iter() - .map(|entry| { - new_child_scan_unit( - entry.relpath.clone(), - assignment.generation + 1, - delegated_child_scan_mode(), - ) - }), - ); + child_scan_units.extend(direct_dirs[delegated_child_scan_unit_count..].iter().map( + |entry| { + new_child_scan_unit( + entry.relpath.clone(), + assignment.generation + 1, + delegated_child_scan_mode(), + ) + }, + )); child_scan_units.sort_by(|a, b| a.root_relpath.cmp(&b.root_relpath)); sort_transfer_scan_batches(&mut direct_files_only_batches); return Ok(FluxonFsTransferScanResultWire { @@ -5027,7 +5078,8 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( let mut root_partitioned = root_listing.emitted_direct_files_batch_count > 0 || direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()); let mut mergeable_empty_dir_count = direct_empty_dirs.len(); - let mut mergeable_empty_dir_estimated_bytes = estimate_empty_dir_manifest_bytes(&direct_empty_dirs); + let mut mergeable_empty_dir_estimated_bytes = + estimate_empty_dir_manifest_bytes(&direct_empty_dirs); for child_relpath in direct_dirs.iter().map(|entry| entry.relpath.clone()) { let child_plan = plan_transfer_subtree_batches( root_dir_abs, @@ -5043,8 +5095,7 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( let child_empty_dir_count = child_plan.mergeable_empty_dir_count; let child_empty_dir_estimated_bytes = child_plan.mergeable_empty_dir_estimated_bytes; - if mergeable_empty_dir_count - .saturating_add(child_empty_dir_count) + if mergeable_empty_dir_count.saturating_add(child_empty_dir_count) > TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET || mergeable_empty_dir_estimated_bytes .saturating_add(child_empty_dir_estimated_bytes) @@ -5083,7 +5134,10 @@ pub(crate) fn build_transfer_scan_result_for_root_dir_abs( if (!direct_files.is_empty() || !direct_symlink_notices.is_empty() || !mergeable_empty_child_relpaths.is_empty()) - && !direct_files_only_disposition_covers_root(assignment, assignment.root_relpath.as_str()) + && !direct_files_only_disposition_covers_root( + assignment, + assignment.root_relpath.as_str(), + ) { let mut next_partition_index = root_listing.emitted_direct_files_batch_count; direct_empty_dirs.extend(mergeable_empty_child_relpaths); @@ -5212,10 +5266,11 @@ fn handle_transfer_scan_assignment( assignment.generation, assignment.known_dispositions.len(), ); - let result = match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { - Ok(v) => v, - Err(resp) => return resp, - }; + let result = + match build_transfer_scan_result_for_root_dir_abs(root_dir_abs.as_str(), &assignment) { + Ok(v) => v, + Err(resp) => return resp, + }; encode_transfer_scan_result(&result, "transfer scan result") } @@ -5228,8 +5283,11 @@ fn prepare_transfer_file_streaming( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { let staging_relpath = transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) @@ -5239,9 +5297,12 @@ where .map_err(TransferWorkerExecutionError::fatal)?; ensure_transfer_parent_dirs(dst_root, final_relpath.as_str()) .map_err(TransferWorkerExecutionError::fatal)?; - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_relpath.as_str()) - .map_err(resp_err_kverr) - .map_err(TransferWorkerExecutionError::fatal)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr) + .map_err(TransferWorkerExecutionError::fatal)?; let mut dst_file = open_create_file_with_parent_dir_chmod_retry(&staging_abs) .map_err(TransferWorkerExecutionError::fatal)?; dst_file @@ -5254,14 +5315,14 @@ where let remaining = file.size.saturating_sub(copied); let chunk = coordinator.read_chunk(file, copied, remaining.min(CHUNK_BYTES as i64))?; if chunk.is_empty() { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( - ApiError::InvalidArgument { + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { detail: format!( "transfer worker source ended before expected size: relpath={} expected={} copied={}", file.relpath, file.size, copied ), - }, - )))); + }), + ))); } dst_file .write_all(&chunk) @@ -5271,13 +5332,14 @@ where coordinator.record_written_bytes(chunk.len() as i64); } if copied != file.size { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( - ApiError::InvalidArgument { - detail: format!( - "transfer worker size mismatch before staging completion: relpath={} expected={} actual={}", - file.relpath, file.size, copied - ), - })))); + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "transfer worker size mismatch before staging completion: relpath={} expected={} actual={}", + file.relpath, file.size, copied + ), + }), + ))); } // The staged file is still invisible at this point, so one more checkpoint // keeps supersession able to stop the worker before any later visible @@ -5302,48 +5364,57 @@ fn execute_transfer_single_file( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { coordinator.checkpoint_continue()?; - let prepared = match prepare_transfer_file_streaming(dst_root, staging_prefix, file, coordinator) { - Ok(v) => v, - Err(TransferWorkerExecutionError::Fatal(resp)) => { - if let Some(failed) = classify_transfer_failed_file(file, &resp) { - let staging_relpath = - transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) - .map_err(TransferWorkerExecutionError::fatal)?; - let staging_abs = safe_join_root( - dst_root.to_string_lossy().as_ref(), - staging_relpath.as_str(), - ) - .map_err(resp_err_kverr) - .map_err(TransferWorkerExecutionError::fatal)?; - match fs::remove_file(&staging_abs) { - Ok(()) => {} - Err(err) if err.kind() == ErrorKind::NotFound => {} - Err(err) => return Err(TransferWorkerExecutionError::fatal(resp_err_io(err))), + let prepared = + match prepare_transfer_file_streaming(dst_root, staging_prefix, file, coordinator) { + Ok(v) => v, + Err(TransferWorkerExecutionError::Fatal(resp)) => { + if let Some(failed) = classify_transfer_failed_file(file, &resp) { + let staging_relpath = + transfer_staging_file_relpath(staging_prefix, file.relpath.as_str()) + .map_err(TransferWorkerExecutionError::fatal)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr) + .map_err(TransferWorkerExecutionError::fatal)?; + match fs::remove_file(&staging_abs) { + Ok(()) => {} + Err(err) if err.kind() == ErrorKind::NotFound => {} + Err(err) => { + return Err(TransferWorkerExecutionError::fatal(resp_err_io(err))); + } + } + return Ok(TransferWorkerLaneOutcome::Failed( + TransferWorkerLaneFailedFileResult { result: failed }, + )); } - return Ok(TransferWorkerLaneOutcome::Failed( - TransferWorkerLaneFailedFileResult { result: failed }, - )); + return Err(TransferWorkerExecutionError::Fatal(resp)); } - return Err(TransferWorkerExecutionError::Fatal(resp)); - } - Err(err) => return Err(err), - }; + Err(err) => return Err(err), + }; coordinator.checkpoint_continue()?; - let result = promote_prepared_transfer_file(dst_root, PreparedTransferFile { - staging_relpath: prepared.staging_relpath.clone(), - final_relpath: prepared.final_relpath.clone(), - visible_size: prepared.visible_size, - }) + let result = promote_prepared_transfer_file( + dst_root, + PreparedTransferFile { + staging_relpath: prepared.staging_relpath.clone(), + final_relpath: prepared.final_relpath.clone(), + visible_size: prepared.visible_size, + }, + ) .map_err(TransferWorkerExecutionError::fatal); match result { - Ok(result) => Ok(TransferWorkerLaneOutcome::Visible(TransferWorkerLaneFileResult { - result, - })), + Ok(result) => Ok(TransferWorkerLaneOutcome::Visible( + TransferWorkerLaneFileResult { result }, + )), Err(TransferWorkerExecutionError::Fatal(resp)) => { if let Some(failed) = classify_transfer_failed_file(file, &resp) { let staging_abs = safe_join_root( @@ -5374,8 +5445,11 @@ fn execute_transfer_empty_dir( coordinator: &TransferWorkerCoordinator, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { coordinator.checkpoint_continue()?; @@ -5393,11 +5467,14 @@ fn execute_transfer_worker_assignment_with_policy( read_chunk: ReadChunkFn, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { let policy = policy.normalized(); @@ -5424,11 +5501,14 @@ fn execute_transfer_worker_assignment_with_policy_and_progress Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { create_dir_all_with_parent_dir_chmod_retry(dst_root) @@ -5436,9 +5516,11 @@ where let manifest = FluxonFsTransferManifestWire::decode_from_blob(assignment.manifest_blob.as_slice()) .map_err(|e| { - TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("decode transfer worker manifest failed: {}", e), - }))) + TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api( + ApiError::InvalidArgument { + detail: format!("decode transfer worker manifest failed: {}", e), + }, + ))) })?; if transfer_manifest_is_empty_dirs_only_batch(&manifest, assignment.collect_infos.as_slice()) { // Empty-dir-only batches never generate byte-based ramp-up signals, so @@ -5664,10 +5746,16 @@ fn promote_prepared_transfer_file( dst_root: &PathBuf, file: PreparedTransferFile, ) -> Result { - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), file.staging_relpath.as_str()) - .map_err(resp_err_kverr)?; - let final_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), file.final_relpath.as_str()) - .map_err(resp_err_kverr)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + file.staging_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; + let final_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + file.final_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; rename_with_dst_parent_dir_chmod_retry(&staging_abs, &final_abs)?; Ok(FluxonFsTransferWorkerFileResultWire { relpath: file.final_relpath.clone(), @@ -5694,15 +5782,15 @@ fn prepare_transfer_collect_info_materialization( ), })) })?; - let staging_relpath = transfer_collect_info_staging_relpath( - batch_id, - worker_task_id, - collect_info.collect_kind, - )?; + let staging_relpath = + transfer_collect_info_staging_relpath(batch_id, worker_task_id, collect_info.collect_kind)?; ensure_transfer_parent_dirs(dst_root, staging_relpath.as_str())?; ensure_transfer_parent_dirs(dst_root, output_relpath.as_str())?; - let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_relpath.as_str()) - .map_err(resp_err_kverr)?; + let staging_abs = safe_join_root( + dst_root.to_string_lossy().as_ref(), + staging_relpath.as_str(), + ) + .map_err(resp_err_kverr)?; let mut dst_file = open_create_file_with_parent_dir_chmod_retry(&staging_abs)?; dst_file.set_len(0).map_err(resp_err_io)?; dst_file @@ -5723,14 +5811,15 @@ fn transfer_collect_info_staging_relpath( worker_task_id: &str, collect_kind: FluxonFsTransferCollectInfoKind, ) -> Result { - let output_relpath = transfer_collect_info_output_relpath(batch_id, collect_kind).map_err(|detail| { - resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!( - "build transfer collect info output relpath failed: batch_id={} err={}", - batch_id, detail - ), - })) - })?; + let output_relpath = + transfer_collect_info_output_relpath(batch_id, collect_kind).map_err(|detail| { + resp_err_kverr(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "build transfer collect info output relpath failed: batch_id={} err={}", + batch_id, detail + ), + })) + })?; Ok(format!("{}.{}.fluxon.part", output_relpath, worker_task_id)) } @@ -5750,9 +5839,12 @@ fn prune_empty_parent_dirs(mut current: PathBuf, root: &PathBuf) -> Result<(), F Ok(()) } -fn cleanup_attempt_staging_prefix(dst_root: &PathBuf, staging_prefix: &str) -> Result<(), FlatDict> { - let staging_abs = - safe_join_root(dst_root.to_string_lossy().as_ref(), staging_prefix).map_err(resp_err_kverr)?; +fn cleanup_attempt_staging_prefix( + dst_root: &PathBuf, + staging_prefix: &str, +) -> Result<(), FlatDict> { + let staging_abs = safe_join_root(dst_root.to_string_lossy().as_ref(), staging_prefix) + .map_err(resp_err_kverr)?; match fs::remove_dir_all(&staging_abs) { Ok(()) => {} Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), @@ -5865,7 +5957,8 @@ pub(crate) fn read_transfer_chunk_from_root_dir_abs( return Ok(Vec::new()); } let to_read = std::cmp::min(length, size - offset) as usize; - f.seek(SeekFrom::Start(offset as u64)).map_err(resp_err_io)?; + f.seek(SeekFrom::Start(offset as u64)) + .map_err(resp_err_io)?; let mut buf = vec![0u8; to_read]; f.read_exact(&mut buf).map_err(resp_err_io)?; Ok(buf) @@ -5881,11 +5974,14 @@ pub(crate) fn execute_transfer_worker_assignment( read_chunk: ReadChunkFn, ) -> Result where - ReadChunkFn: - Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError> - + Send - + Sync - + 'static, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError> + + Send + + Sync + + 'static, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError> + Send + Sync + 'static, { execute_transfer_worker_assignment_with_policy( @@ -5943,12 +6039,19 @@ pub(super) fn handle_transfer_read(exports: &AgentExportsHandle, payload: FlatDi Ok(v) => v, Err(e) => return resp_err_kverr(e), }; - let buf = - match read_transfer_chunk_from_root_dir_abs(root_dir_abs.as_str(), relpath.as_str(), offset, length) { - Ok(v) => v, - Err(resp) => return resp, - }; - resp_ok(BTreeMap::from([("data".to_string(), FlatValue::Bytes(buf))])) + let buf = match read_transfer_chunk_from_root_dir_abs( + root_dir_abs.as_str(), + relpath.as_str(), + offset, + length, + ) { + Ok(v) => v, + Err(resp) => return resp, + }; + resp_ok(BTreeMap::from([( + "data".to_string(), + FlatValue::Bytes(buf), + )])) } pub(super) fn handle_transfer_stream_open( @@ -6010,7 +6113,9 @@ pub(super) fn handle_transfer_worker( Err(resp) => return resp, }; match registry.launch_task(api, master_id, exports, assignment) { - Ok(result) => encode_transfer_worker_launch_result(&result, "transfer worker launch result"), + Ok(result) => { + encode_transfer_worker_launch_result(&result, "transfer worker launch result") + } Err(resp) => resp, } } @@ -6086,15 +6191,15 @@ mod tests { .collect() } - fn assert_all_child_scan_units_are_subtree_streaming( - result: &FluxonFsTransferScanResultWire, - ) { - assert!(result - .child_scan_units - .iter() - .all(|child| child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming)); - } - + fn assert_all_child_scan_units_are_subtree_streaming(result: &FluxonFsTransferScanResultWire) { + assert!( + result + .child_scan_units + .iter() + .all(|child| child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming) + ); + } + fn ok_bool(resp: &FlatDict) -> bool { matches!(resp.get("ok"), Some(FlatValue::Bool(true))) } @@ -6106,10 +6211,7 @@ mod tests { panic!("unexpected open result fatal decode error: {:?}", other) } Err(TransferWorkerRpcFailure::Retryable { detail }) => { - panic!( - "unexpected open result retryable decode error: {}", - detail - ) + panic!("unexpected open result retryable decode error: {}", detail) } } } @@ -6121,10 +6223,7 @@ mod tests { panic!("unexpected next result fatal decode error: {:?}", other) } Err(TransferWorkerRpcFailure::Retryable { detail }) => { - panic!( - "unexpected next result retryable decode error: {}", - detail - ) + panic!("unexpected next result retryable decode error: {}", detail) } } } @@ -6140,10 +6239,7 @@ mod tests { .collect() } - fn test_worker_assignment( - relpath: &str, - size: i64, - ) -> FluxonFsTransferWorkerAssignmentWire { + fn test_worker_assignment(relpath: &str, size: i64) -> FluxonFsTransferWorkerAssignmentWire { FluxonFsTransferWorkerAssignmentWire { job_id: "job".to_string(), batch_id: "batch".to_string(), @@ -6158,12 +6254,13 @@ mod tests { root_relpath: ".".to_string(), staging_prefix: ".fluxon.stage/job/batch".to_string(), lease_expire_unix_ms: 0, - manifest_blob: build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { + manifest_blob: build_transfer_manifest_blob( + vec![FluxonFsTransferScanFrontierEntry { relpath: relpath.to_string(), size, - }, - ], Vec::new()) + }], + Vec::new(), + ) .unwrap(), collect_infos: Vec::new(), } @@ -6183,7 +6280,11 @@ mod tests { read_chunk: ReadChunkFn, ) -> TransferWorkerCoordinator where - ReadChunkFn: Fn(&FluxonFsTransferManifestEntryWire, i64, i64) -> Result, TransferWorkerExecutionError>, + ReadChunkFn: Fn( + &FluxonFsTransferManifestEntryWire, + i64, + i64, + ) -> Result, TransferWorkerExecutionError>, CheckpointFn: Fn() -> Result<(), TransferWorkerExecutionError>, { let policy = Arc::new(TransferWorkerLanePolicy::production_default()); @@ -6217,16 +6318,19 @@ mod tests { #[test] fn build_transfer_manifest_blob_round_trips_entries() { - let blob = build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { - relpath: "a".to_string(), - size: 1, - }, - FluxonFsTransferScanFrontierEntry { - relpath: "b/c".to_string(), - size: 2, - }, - ], vec!["empty".to_string()]) + let blob = build_transfer_manifest_blob( + vec![ + FluxonFsTransferScanFrontierEntry { + relpath: "a".to_string(), + size: 1, + }, + FluxonFsTransferScanFrontierEntry { + relpath: "b/c".to_string(), + size: 2, + }, + ], + vec!["empty".to_string()], + ) .unwrap(); let manifest = FluxonFsTransferManifestWire::decode_from_blob(&blob).unwrap(); assert_eq!(manifest.entry_count, 2); @@ -6250,11 +6354,12 @@ mod tests { #[test] fn materialize_transfer_collect_info_writes_task_scoped_staging_then_output_file() { let root = TempDir::new().unwrap(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "root/link-file.bin".to_string(), - link_target: "target/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "root/link-file.bin".to_string(), + link_target: "target/file.bin".to_string(), + }]) + .unwrap(); let prepared = prepare_transfer_collect_info_materialization( &root.path().to_path_buf(), "batch-1", @@ -6383,7 +6488,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(2)), ("length".to_string(), FlatValue::Int64(3)), ]), @@ -6398,7 +6506,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(6)), ("length".to_string(), FlatValue::Int64(1)), ]), @@ -6423,7 +6534,10 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(1)), ("length".to_string(), FlatValue::Int64(3)), ]), @@ -6446,9 +6560,15 @@ mod tests { &exports, FlatDict::from([ ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("offset".to_string(), FlatValue::Int64(0)), - ("length".to_string(), FlatValue::Int64(CHUNK_BYTES as i64 + 1)), + ( + "length".to_string(), + FlatValue::Int64(CHUNK_BYTES as i64 + 1), + ), ]), ); assert!(matches!(resp.get("ok"), Some(FlatValue::Bool(false)))); @@ -6470,7 +6590,10 @@ mod tests { FlatValue::String("task-0".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(0)), ]), ); @@ -6551,7 +6674,10 @@ mod tests { FlatValue::String("task-1".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(0)), ]), ); @@ -6584,7 +6710,10 @@ mod tests { ("length".to_string(), FlatValue::Int64(2)), ]), ); - assert!(matches!(invalid_resp.get("ok"), Some(FlatValue::Bool(false)))); + assert!(matches!( + invalid_resp.get("ok"), + Some(FlatValue::Bool(false)) + )); } #[test] @@ -6603,7 +6732,10 @@ mod tests { FlatValue::String("task-2".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(3)), ]), ); @@ -6647,7 +6779,10 @@ mod tests { FlatValue::String("task-3".to_string()), ), ("export".to_string(), FlatValue::String("src".to_string())), - ("relpath".to_string(), FlatValue::String("f.bin".to_string())), + ( + "relpath".to_string(), + FlatValue::String("f.bin".to_string()), + ), ("initial_offset".to_string(), FlatValue::Int64(3)), ]), ); @@ -6718,7 +6853,10 @@ mod tests { let result = decode_result_json(&resp); assert!(result.finished); assert!(result.direct_files_only_batches.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -6781,7 +6919,8 @@ mod tests { } #[test] - fn handle_transfer_scan_assignment_groups_empty_children_into_direct_batch_without_direct_files() { + fn handle_transfer_scan_assignment_groups_empty_children_into_direct_batch_without_direct_files() + { let root = TempDir::new().unwrap(); write_file(&root, "root/big/data.bin", b"12345"); fs::create_dir_all(root.path().join("root/empty-a")).unwrap(); @@ -6814,9 +6953,10 @@ mod tests { assert_eq!(child_scan_unit_roots(&result), vec!["root/big".to_string()]); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); - let direct_manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let direct_manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert!(direct_manifest.entries.is_empty()); assert_eq!( direct_manifest.empty_dir_relpaths, @@ -6854,16 +6994,17 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.full_dir_batches.len(), 0); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/huge".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/huge".to_string() + ); let manifest = FluxonFsTransferManifestWire::decode_from_blob( &result.direct_files_only_batches[0].manifest_blob, ) .unwrap(); assert!(manifest.entries.is_empty()); assert!(!manifest.empty_dir_relpaths.is_empty()); - assert!( - manifest.empty_dir_relpaths.len() <= TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET - ); + assert!(manifest.empty_dir_relpaths.len() <= TRANSFER_MERGEABLE_EMPTY_DIR_BUDGET); assert!( estimate_empty_dir_manifest_bytes(&manifest.empty_dir_relpaths) <= TRANSFER_MERGEABLE_EMPTY_DIR_ESTIMATED_BYTES_BUDGET @@ -6879,10 +7020,10 @@ mod tests { )) + 1; for idx in 0..child_count { - fs::create_dir_all(root.path().join(format!( - "root/branch-{idx:05}/{}", - "x".repeat(200) - ))) + fs::create_dir_all( + root.path() + .join(format!("root/branch-{idx:05}/{}", "x".repeat(200))), + ) .unwrap(); } let result = build_transfer_scan_result_for_root_dir_abs( @@ -6908,10 +7049,12 @@ mod tests { assert!(!result.finished); assert!(result.direct_files_only_batches.is_empty()); assert!(!result.child_scan_units.is_empty()); - assert!(result - .child_scan_units - .iter() - .any(|child| child.scan_mode == FluxonFsTransferScanMode::FullTree)); + assert!( + result + .child_scan_units + .iter() + .any(|child| child.scan_mode == FluxonFsTransferScanMode::FullTree) + ); assert!(result.child_scan_units.iter().all(|child| { child.scan_mode == FluxonFsTransferScanMode::FullTree || child.scan_mode == FluxonFsTransferScanMode::SubtreeStreaming @@ -6963,10 +7106,16 @@ mod tests { assert!(!continue_locally); assert_eq!(next_event_seq_no, 9); assert_eq!(events.len(), 2); - assert_eq!(events[0].event_kind, FluxonFsTransferScanEventKindWire::Append); + assert_eq!( + events[0].event_kind, + FluxonFsTransferScanEventKindWire::Append + ); assert_eq!(events[0].event_seq_no, 7); assert_eq!(events[0].full_dir_batches.len(), 1); - assert_eq!(events[1].event_kind, FluxonFsTransferScanEventKindWire::Finished); + assert_eq!( + events[1].event_kind, + FluxonFsTransferScanEventKindWire::Finished + ); assert_eq!(events[1].event_seq_no, 8); assert!(events[1].direct_files_only_batches.is_empty()); assert!(events[1].child_scan_units.is_empty()); @@ -6997,17 +7146,21 @@ mod tests { skip_entries: Vec::new(), }; - let first = build_transfer_scan_result_for_root_dir_abs( - root.path().to_str().unwrap(), - &assignment, - ) - .unwrap(); + let first = + build_transfer_scan_result_for_root_dir_abs(root.path().to_str().unwrap(), &assignment) + .unwrap(); assert!(!first.finished); assert!(!first.direct_files_only_batches.is_empty()); assert!(first.full_dir_batches.is_empty()); assert_eq!(first.child_scan_units.len(), 1); - assert_eq!(first.child_scan_units[0].scan_unit_id, assignment.scan_unit_id); - assert_eq!(first.child_scan_units[0].root_relpath, assignment.root_relpath); + assert_eq!( + first.child_scan_units[0].scan_unit_id, + assignment.scan_unit_id + ); + assert_eq!( + first.child_scan_units[0].root_relpath, + assignment.root_relpath + ); assert_eq!(first.child_scan_units[0].generation, assignment.generation); let first_entry_count = first .direct_files_only_batches @@ -7019,7 +7172,10 @@ mod tests { .len() }) .sum::(); - assert_eq!(first_entry_count, TRANSFER_SCAN_ROOT_LISTING_SLICE_ENTRY_LIMIT); + assert_eq!( + first_entry_count, + TRANSFER_SCAN_ROOT_LISTING_SLICE_ENTRY_LIMIT + ); let second_assignment = FluxonFsTransferScanAssignmentWire { scan_task_id: "task-2".to_string(), @@ -7086,7 +7242,8 @@ mod tests { } #[test] - fn build_transfer_scan_result_root_direct_fanout_only_emits_child_scan_units_without_recursing() { + fn build_transfer_scan_result_root_direct_fanout_only_emits_child_scan_units_without_recursing() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/child/payload.bin", b"xyz"); @@ -7114,14 +7271,18 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); assert!(result.full_dir_batches.is_empty()); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree ); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7133,7 +7294,8 @@ mod tests { } #[test] - fn build_transfer_scan_result_directory_direct_fanout_only_emits_child_scan_units_without_recursing() { + fn build_transfer_scan_result_directory_direct_fanout_only_emits_child_scan_units_without_recursing() + { let root = TempDir::new().unwrap(); write_file(&root, "root/child/direct.bin", b"abc"); write_file(&root, "root/child/grand/payload.bin", b"xyz"); @@ -7161,14 +7323,18 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); assert!(result.full_dir_batches.is_empty()); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child/grand".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child/grand".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree ); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7206,10 +7372,14 @@ mod tests { .unwrap(); assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); - assert_eq!(result.direct_files_only_batches[0].root_relpath, "root/child"); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.direct_files_only_batches[0].root_relpath, + "root/child" + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7218,7 +7388,10 @@ mod tests { }] ); assert!(manifest.empty_dir_relpaths.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child/grand".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child/grand".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7253,10 +7426,14 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child".to_string()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child".to_string() + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7330,9 +7507,10 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7407,7 +7585,10 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); assert_eq!(result.child_scan_units.len(), 1); - assert_eq!(result.child_scan_units[0].root_relpath, "root/child-b".to_string()); + assert_eq!( + result.child_scan_units[0].root_relpath, + "root/child-b".to_string() + ); assert_eq!( result.child_scan_units[0].scan_mode, FluxonFsTransferScanMode::FullTree @@ -7447,9 +7628,10 @@ mod tests { assert_eq!(result.direct_files_only_batches.len(), 1); assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7490,9 +7672,10 @@ mod tests { assert!(result.direct_files_only_batches.is_empty()); assert!(result.child_scan_units.is_empty()); assert_eq!(result.full_dir_batches.len(), 1); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.full_dir_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.full_dir_batches[0].manifest_blob, + ) + .unwrap(); assert!(manifest.entries.is_empty()); assert_eq!(manifest.empty_dir_relpaths, vec!["root".to_string()]); } @@ -7533,7 +7716,8 @@ mod tests { } #[test] - fn handle_transfer_scan_assignment_does_not_reaggregate_root_when_descendant_batch_is_durable() { + fn handle_transfer_scan_assignment_does_not_reaggregate_root_when_descendant_batch_is_durable() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/big/data.bin", b"12345"); @@ -7579,21 +7763,29 @@ mod tests { size: 3, }] ); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root")); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root/big")); - assert_eq!(child_scan_unit_roots(&result), vec!["root/small".to_string()]); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root") + ); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root/big") + ); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/small".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } #[test] - fn handle_transfer_scan_assignment_honors_cross_generation_descendant_full_dir_during_restart() { + fn handle_transfer_scan_assignment_honors_cross_generation_descendant_full_dir_during_restart() + { let root = TempDir::new().unwrap(); write_file(&root, "root/direct.bin", b"abc"); write_file(&root, "root/big/data.bin", b"12345"); @@ -7639,21 +7831,29 @@ mod tests { size: 3, }] ); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root")); - assert!(result - .full_dir_batches - .iter() - .all(|batch| batch.root_relpath != "root/big")); - assert_eq!(child_scan_unit_roots(&result), vec!["root/small".to_string()]); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root") + ); + assert!( + result + .full_dir_batches + .iter() + .all(|batch| batch.root_relpath != "root/big") + ); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/small".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } #[test] - fn handle_transfer_scan_assignment_replays_descendant_current_layer_when_only_partial_descendant_direct_files_batch_is_durable() { + fn handle_transfer_scan_assignment_replays_descendant_current_layer_when_only_partial_descendant_direct_files_batch_is_durable() + { let root = TempDir::new().unwrap(); write_file(&root, "root/child/a.bin", b"ab"); write_file(&root, "root/child/b.bin", b"cd"); @@ -7688,10 +7888,14 @@ mod tests { assert!(result.child_scan_units.is_empty()); assert!(result.full_dir_batches.is_empty()); assert_eq!(result.direct_files_only_batches.len(), 1); - assert_eq!(result.direct_files_only_batches[0].root_relpath, "root/child"); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + assert_eq!( + result.direct_files_only_batches[0].root_relpath, + "root/child" + ); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![ @@ -7737,7 +7941,10 @@ mod tests { assert!(result.finished); assert!(result.direct_files_only_batches.is_empty()); - assert_eq!(child_scan_unit_roots(&result), vec!["root/parent".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/parent".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7772,9 +7979,10 @@ mod tests { assert!(result.finished); assert_eq!(result.direct_files_only_batches.len(), 1); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.direct_files_only_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.direct_files_only_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![FluxonFsTransferManifestEntryWire { @@ -7782,7 +7990,10 @@ mod tests { size: 10, }] ); - assert_eq!(child_scan_unit_roots(&result), vec!["root/child".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/child".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7820,7 +8031,10 @@ mod tests { assert!(ok_bool(&resp)); assert!(result.finished); - assert_eq!(child_scan_unit_roots(&result), vec!["root/blocked".to_string()]); + assert_eq!( + child_scan_unit_roots(&result), + vec!["root/blocked".to_string()] + ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); } @@ -7870,7 +8084,11 @@ mod tests { ); assert_eq!( child_scan_unit_roots(&result), - vec!["root/a".to_string(), "root/b".to_string(), "root/c".to_string()] + vec![ + "root/a".to_string(), + "root/b".to_string(), + "root/c".to_string() + ] ); assert_all_child_scan_units_are_subtree_streaming(&result); assert!(result.full_dir_batches.is_empty()); @@ -7905,7 +8123,10 @@ mod tests { ); let result = decode_result_json(&resp); assert_eq!(result.full_dir_batches.len(), 1); - assert_eq!(result.full_dir_batches[0].batch_kind, FluxonFsTransferBatchKind::SubtreeSlice); + assert_eq!( + result.full_dir_batches[0].batch_kind, + FluxonFsTransferBatchKind::SubtreeSlice + ); assert_eq!(result.full_dir_batches[0].collect_infos.len(), 1); assert_eq!( decode_symlink_notice_collect_blob( @@ -7957,11 +8178,15 @@ mod tests { assert!(result.direct_files_only_batches.is_empty()); assert!(result.child_scan_units.is_empty()); assert_eq!(result.full_dir_batches.len(), 1); - assert_eq!(result.full_dir_batches[0].batch_kind, FluxonFsTransferBatchKind::SubtreeSlice); + assert_eq!( + result.full_dir_batches[0].batch_kind, + FluxonFsTransferBatchKind::SubtreeSlice + ); assert_eq!(result.full_dir_batches[0].root_relpath, "root".to_string()); - let manifest = - FluxonFsTransferManifestWire::decode_from_blob(&result.full_dir_batches[0].manifest_blob) - .unwrap(); + let manifest = FluxonFsTransferManifestWire::decode_from_blob( + &result.full_dir_batches[0].manifest_blob, + ) + .unwrap(); assert_eq!( manifest.entries, vec![ @@ -7982,7 +8207,9 @@ mod tests { FluxonFsTransferCollectInfoKind::SymlinkNotice ); let mut notices = decode_symlink_notice_collect_blob( - direct_files_only_batch.collect_infos[0].collect_blob.as_slice() + direct_files_only_batch.collect_infos[0] + .collect_blob + .as_slice(), ); notices.sort_by(|a, b| a.relpath.cmp(&b.relpath)); assert_eq!( @@ -8004,16 +8231,17 @@ mod tests { fn prepare_transfer_file_from_chunks_promotes_staged_file_to_final_path() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"ab".to_vec(), b"cde".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![ + b"ab".to_vec(), + b"cde".to_vec(), + Vec::new(), + ])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let prepared = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8038,32 +8266,31 @@ mod tests { fs::read(root.path().join("dir/file.bin")).unwrap(), b"abcde".to_vec() ); - assert!(!root - .path() - .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part") - .exists()); + assert!( + !root + .path() + .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part") + .exists() + ); } #[test] fn prepare_transfer_file_from_chunks_truncates_existing_staging_file() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let stale_staging = - root.path() - .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part"); + let stale_staging = root + .path() + .join(".fluxon.stage/job/batch/dir/file.bin/file.bin.fluxon.part"); fs::create_dir_all(stale_staging.parent().unwrap()).unwrap(); fs::write(&stale_staging, b"stale-data").unwrap(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let prepared = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8076,23 +8303,23 @@ mod tests { .unwrap(); promote_prepared_transfer_file(&dst_root, prepared).unwrap(); - assert_eq!(fs::read(root.path().join("dir/file.bin")).unwrap(), b"xy".to_vec()); + assert_eq!( + fs::read(root.path().join("dir/file.bin")).unwrap(), + b"xy".to_vec() + ); } #[test] fn prepare_transfer_file_from_chunks_rejects_size_mismatch_and_keeps_staging_file() { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); - let coordinator = test_transfer_coordinator( - || Ok(()), - { - let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); - move |_file, _read_offset, _length| { - let mut chunks = chunks.lock(); - Ok(chunks.remove(0)) - } - }, - ); + let coordinator = test_transfer_coordinator(|| Ok(()), { + let chunks = Arc::new(Mutex::new(vec![b"xy".to_vec(), Vec::new()])); + move |_file, _read_offset, _length| { + let mut chunks = chunks.lock(); + Ok(chunks.remove(0)) + } + }); let err = prepare_transfer_file_streaming( &dst_root, ".fluxon.stage/job/batch", @@ -8275,15 +8502,10 @@ mod tests { let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, _read_offset, _length| Ok(file_bytes.clone()) - }, - ) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, _read_offset, _length| Ok(file_bytes.clone()) + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8303,7 +8525,10 @@ mod tests { create_dir_all_with_parent_dir_chmod_retry(&target).unwrap(); assert!(target.is_dir()); - assert_eq!(fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, 0o777); + assert_eq!( + fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, + 0o777 + ); } #[cfg(unix)] @@ -8318,21 +8543,19 @@ mod tests { let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, _read_offset, _length| Ok(file_bytes.clone()) - }, - ) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, _read_offset, _length| Ok(file_bytes.clone()) + }) .unwrap(); assert_eq!(result.file_results.len(), 1); assert!(dst_root.is_dir()); assert_eq!(fs::read(dst_root.join("dir/file.bin")).unwrap(), file_bytes); - assert_eq!(fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, 0o777); + assert_eq!( + fs::metadata(&locked_parent).unwrap().permissions().mode() & 0o777, + 0o777 + ); } #[test] @@ -8349,47 +8572,48 @@ mod tests { let assignment = assignment.clone(); let heartbeat_attempts = heartbeat_attempts.clone(); move || { - retry_transfer_worker_rpc_with_backoff( - &assignment, - "checkpoint", - "test-checkpoint", - BackoffConfig { - initial_secs: 0, - max_secs: 0, - }, - WarnConfig { - warn_interval_secs: 0, - }, - || { - let attempt = - heartbeat_attempts.fetch_add(1, Ordering::SeqCst) + 1; - if attempt < 3 { - return Err(TransferWorkerRpcFailure::Retryable { - detail: format!( - "transient heartbeat failure attempt={}", - attempt - ), - }); - } - Ok(()) - }, - ) - .map_err(TransferWorkerExecutionError::fatal) - } + retry_transfer_worker_rpc_with_backoff( + &assignment, + "checkpoint", + "test-checkpoint", + BackoffConfig { + initial_secs: 0, + max_secs: 0, + }, + WarnConfig { + warn_interval_secs: 0, + }, + || { + let attempt = heartbeat_attempts.fetch_add(1, Ordering::SeqCst) + 1; + if attempt < 3 { + return Err(TransferWorkerRpcFailure::Retryable { + detail: format!( + "transient heartbeat failure attempt={}", + attempt + ), + }); + } + Ok(()) + }, + ) + .map_err(TransferWorkerExecutionError::fatal) + } }, { let file_bytes = file_bytes.clone(); move |file, read_offset, _length| { - if file.relpath != "dir/file.bin" { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("unexpected file relpath: {}", file.relpath), - })))); - } - if read_offset == 0 { - return Ok(file_bytes.clone()); + if file.relpath != "dir/file.bin" { + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!("unexpected file relpath: {}", file.relpath), + }), + ))); + } + if read_offset == 0 { + return Ok(file_bytes.clone()); + } + Ok(Vec::new()) } - Ok(Vec::new()) - } }, ) .unwrap(); @@ -8409,24 +8633,20 @@ mod tests { let file_bytes = b"payload".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); let read_attempts = Arc::new(AtomicUsize::new(0)); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let assignment = assignment.clone(); - let file_bytes = file_bytes.clone(); - let read_attempts = read_attempts.clone(); - move |file, read_offset, _length| { + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let assignment = assignment.clone(); + let file_bytes = file_bytes.clone(); + let read_attempts = read_attempts.clone(); + move |file, read_offset, _length| { if file.relpath != "dir/file.bin" { - return Err(TransferWorkerExecutionError::fatal(resp_err_kverr(KvError::Api(ApiError::InvalidArgument { - detail: format!("unexpected file relpath: {}", file.relpath), - })))); + return Err(TransferWorkerExecutionError::fatal(resp_err_kverr( + KvError::Api(ApiError::InvalidArgument { + detail: format!("unexpected file relpath: {}", file.relpath), + }), + ))); } - let op_detail = format!( - "test-read relpath={} offset={}", - file.relpath, read_offset - ); + let op_detail = + format!("test-read relpath={} offset={}", file.relpath, read_offset); retry_transfer_worker_rpc_with_backoff( &assignment, "read_chunk", @@ -8443,10 +8663,7 @@ mod tests { let attempt = read_attempts.fetch_add(1, Ordering::SeqCst) + 1; if attempt < 3 { return Err(TransferWorkerRpcFailure::Retryable { - detail: format!( - "transient read failure attempt={}", - attempt - ), + detail: format!("transient read failure attempt={}", attempt), }); } return Ok(file_bytes.clone()); @@ -8456,8 +8673,7 @@ mod tests { ) .map_err(TransferWorkerExecutionError::fatal) } - }, - ) + }) .unwrap(); assert_eq!(read_attempts.load(Ordering::SeqCst), 3); @@ -8481,23 +8697,23 @@ mod tests { { let checkpoint_calls = checkpoint_calls.clone(); move || { - let calls = checkpoint_calls.fetch_add(1, Ordering::SeqCst) + 1; - if calls >= 4 { - return Err(TransferWorkerExecutionError::Stop( - FluxonFsTransferWorkerStopReasonWire::Superseded, - )); + let calls = checkpoint_calls.fetch_add(1, Ordering::SeqCst) + 1; + if calls >= 4 { + return Err(TransferWorkerExecutionError::Stop( + FluxonFsTransferWorkerStopReasonWire::Superseded, + )); + } + Ok(()) } - Ok(()) - } }, { let file_bytes = file_bytes.clone(); move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); + if read_offset == 0 { + return Ok(file_bytes.clone()); + } + Ok(Vec::new()) } - Ok(Vec::new()) - } }, ); assert!(matches!( @@ -8567,12 +8783,7 @@ mod tests { break; } if max_in_flight - .compare_exchange( - observed, - current, - Ordering::SeqCst, - Ordering::SeqCst, - ) + .compare_exchange(observed, current, Ordering::SeqCst, Ordering::SeqCst) .is_ok() { break; @@ -8588,8 +8799,14 @@ mod tests { assert_eq!(result.file_results.len(), 2); assert!(max_in_flight.load(Ordering::SeqCst) >= 2); - assert_eq!(fs::read(root.path().join("dir/a.bin")).unwrap(), b"xxx".to_vec()); - assert_eq!(fs::read(root.path().join("dir/b.bin")).unwrap(), b"xxx".to_vec()); + assert_eq!( + fs::read(root.path().join("dir/a.bin")).unwrap(), + b"xxx".to_vec() + ); + assert_eq!( + fs::read(root.path().join("dir/b.bin")).unwrap(), + b"xxx".to_vec() + ); } #[test] @@ -8597,29 +8814,25 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "dir/link.bin".to_string(), - link_target: "dir/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "dir/link.bin".to_string(), + link_target: "dir/file.bin".to_string(), + }]) + .unwrap(); let assignment = FluxonFsTransferWorkerAssignmentWire { collect_infos: collect_infos.clone(), ..test_worker_assignment("dir/file.bin", file_bytes.len() as i64) }; - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); - } - Ok(Vec::new()) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, read_offset, _length| { + if read_offset == 0 { + return Ok(file_bytes.clone()); } - }, - ) + Ok(Vec::new()) + } + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8633,7 +8846,11 @@ mod tests { "fluxon_collect_info/batches/batch/symlinks.jsonl" ); assert_eq!( - fs::read(root.path().join("fluxon_collect_info/batches/batch/symlinks.jsonl")).unwrap(), + fs::read( + root.path() + .join("fluxon_collect_info/batches/batch/symlinks.jsonl") + ) + .unwrap(), collect_infos[0].collect_blob ); } @@ -8643,16 +8860,19 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let assignment = FluxonFsTransferWorkerAssignmentWire { - manifest_blob: build_transfer_manifest_blob(vec![ - FluxonFsTransferScanFrontierEntry { - relpath: "dir/good.bin".to_string(), - size: 5, - }, - FluxonFsTransferScanFrontierEntry { - relpath: "dir/bad.bin".to_string(), - size: 5, - }, - ], Vec::new()) + manifest_blob: build_transfer_manifest_blob( + vec![ + FluxonFsTransferScanFrontierEntry { + relpath: "dir/good.bin".to_string(), + size: 5, + }, + FluxonFsTransferScanFrontierEntry { + relpath: "dir/bad.bin".to_string(), + size: 5, + }, + ], + Vec::new(), + ) .unwrap(), ..test_worker_assignment("dir/good.bin", 5) }; @@ -8816,20 +9036,16 @@ mod tests { .unwrap(); let progress_heartbeat_count = Arc::new(AtomicUsize::new(0)); - gate.ensure_continue( - false, - TRANSFER_WORKER_HEARTBEAT_EMPTY_DIR_PROGRESS_COUNT, - { - let progress_heartbeat_count = progress_heartbeat_count.clone(); - move |_heartbeat_unix_ms, heartbeat_detail| { - assert_eq!(heartbeat_detail, "empty_dir_progress"); - progress_heartbeat_count.fetch_add(1, Ordering::SeqCst); - Ok(FluxonFsTransferWorkerHeartbeatResultWire::continue_running( - chrono::Utc::now().timestamp_millis() + 60_000, - )) - } - }, - ) + gate.ensure_continue(false, TRANSFER_WORKER_HEARTBEAT_EMPTY_DIR_PROGRESS_COUNT, { + let progress_heartbeat_count = progress_heartbeat_count.clone(); + move |_heartbeat_unix_ms, heartbeat_detail| { + assert_eq!(heartbeat_detail, "empty_dir_progress"); + progress_heartbeat_count.fetch_add(1, Ordering::SeqCst); + Ok(FluxonFsTransferWorkerHeartbeatResultWire::continue_running( + chrono::Utc::now().timestamp_millis() + 60_000, + )) + } + }) .unwrap(); gate.ensure_continue( @@ -8927,20 +9143,15 @@ mod tests { let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); let assignment = test_worker_assignment("dir/file.bin", file_bytes.len() as i64); - let result = execute_transfer_worker_assignment( - &assignment, - &dst_root, - || Ok(()), - { - let file_bytes = file_bytes.clone(); - move |_file, read_offset, _length| { - if read_offset == 0 { - return Ok(file_bytes.clone()); - } - Ok(Vec::new()) + let result = execute_transfer_worker_assignment(&assignment, &dst_root, || Ok(()), { + let file_bytes = file_bytes.clone(); + move |_file, read_offset, _length| { + if read_offset == 0 { + return Ok(file_bytes.clone()); } - }, - ) + Ok(Vec::new()) + } + }) .unwrap(); assert_eq!(result.file_results.len(), 1); @@ -8948,7 +9159,10 @@ mod tests { cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment).unwrap(); - assert_eq!(fs::read(root.path().join("dir/file.bin")).unwrap(), file_bytes); + assert_eq!( + fs::read(root.path().join("dir/file.bin")).unwrap(), + file_bytes + ); assert!(!root.path().join(".fluxon.stage").exists()); } @@ -8957,11 +9171,12 @@ mod tests { let root = TempDir::new().unwrap(); let dst_root = root.path().to_path_buf(); let file_bytes = b"hello".to_vec(); - let collect_infos = build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { - relpath: "root/link-file.bin".to_string(), - link_target: "target/file.bin".to_string(), - }]) - .unwrap(); + let collect_infos = + build_symlink_collect_infos(vec![FluxonFsTransferSymlinkNoticeEntryWire { + relpath: "root/link-file.bin".to_string(), + link_target: "target/file.bin".to_string(), + }]) + .unwrap(); let assignment = FluxonFsTransferWorkerAssignmentWire { collect_infos: collect_infos.clone(), ..test_worker_assignment("dir/file.bin", file_bytes.len() as i64) @@ -8976,9 +9191,11 @@ mod tests { let result = execute_transfer_worker_assignment( &assignment, &dst_root, - || Err(TransferWorkerExecutionError::Stop( - FluxonFsTransferWorkerStopReasonWire::Superseded, - )), + || { + Err(TransferWorkerExecutionError::Stop( + FluxonFsTransferWorkerStopReasonWire::Superseded, + )) + }, { let file_bytes = file_bytes.clone(); move |_file, read_offset, _length| { @@ -8996,11 +9213,20 @@ mod tests { FluxonFsTransferWorkerStopReasonWire::Superseded )) )); - assert!(root.path().join(prepared_collect.staging_relpath.as_str()).exists()); + assert!( + root.path() + .join(prepared_collect.staging_relpath.as_str()) + .exists() + ); cleanup_transfer_worker_attempt_artifacts(&dst_root, &assignment).unwrap(); assert!(!root.path().join(".fluxon.stage").exists()); - assert!(!root.path().join(prepared_collect.staging_relpath.as_str()).exists()); + assert!( + !root + .path() + .join(prepared_collect.staging_relpath.as_str()) + .exists() + ); } } diff --git a/fluxon_rs/fluxon_fs/src/cache_controller.rs b/fluxon_rs/fluxon_fs/src/cache_controller.rs index 8a0845c..13ce5a8 100644 --- a/fluxon_rs/fluxon_fs/src/cache_controller.rs +++ b/fluxon_rs/fluxon_fs/src/cache_controller.rs @@ -429,8 +429,8 @@ fn now_ms() -> i64 { #[cfg(test)] mod tests { use super::*; - use std::sync::mpsc; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::mpsc; use std::sync::{Condvar, Mutex}; use tokio::time::{Duration, sleep}; diff --git a/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs b/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs index 827bb23..0866432 100644 --- a/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs +++ b/fluxon_rs/fluxon_fs_s3_gateway/src/lib.rs @@ -5344,10 +5344,9 @@ mod tests { }; use crate::transfer::encode_transfer_manifest_blob_with_empty_dirs; use fluxon_fs_core::config::{ - FS_CACHE_DEFAULT_WRITE_SESSION_TARGET_INFLIGHT_BYTES_V1, - FS_EXPORT_DEFAULT_INLINE_BYTES_MAX_BYTES_V1, - FS_EXPORT_DEFAULT_METADATA_CACHE_TTL_MS_V1, FLUXON_FS_LOCAL_TRANSFER_CHECK_DST_EXPORT, FLUXON_FS_LOCAL_TRANSFER_CHECK_SRC_EXPORT, + FS_CACHE_DEFAULT_WRITE_SESSION_TARGET_INFLIGHT_BYTES_V1, + FS_EXPORT_DEFAULT_INLINE_BYTES_MAX_BYTES_V1, FS_EXPORT_DEFAULT_METADATA_CACHE_TTL_MS_V1, FluxonFsAccessModel, FluxonFsAccessUser, FluxonFsExport, FluxonFsExportRoutingMode, FluxonFsGlobalConfig, FluxonFsLocalTransferCheckJobSpecWire, FluxonFsRequestIdentity, FluxonFsS3GatewayConfig, FluxonFsS3KvMissPolicy, FluxonFsS3PermissionAccount, diff --git a/fluxon_rs/fluxon_kv/Cargo.toml b/fluxon_rs/fluxon_kv/Cargo.toml index 22ff136..2d0208c 100644 --- a/fluxon_rs/fluxon_kv/Cargo.toml +++ b/fluxon_rs/fluxon_kv/Cargo.toml @@ -75,6 +75,7 @@ bytes = "1" pprof = { version = "0.15", features = ["flamegraph"] } hex = "0.4" sha2 = "0.10" +foyer = { version = "0.22.3", features = ["serde"] } tokio-tungstenite = { version = "0.21", default-features = false, features = ["connect", "handshake"], optional = true } sockudo-ws = { version = "^1.7.4", default-features = false, features = ["tokio-runtime", "fastrand"], optional = true } diff --git a/fluxon_rs/fluxon_kv/framework_init_steps.yaml b/fluxon_rs/fluxon_kv/framework_init_steps.yaml index 923ae30..24ca312 100644 --- a/fluxon_rs/fluxon_kv/framework_init_steps.yaml +++ b/fluxon_rs/fluxon_kv/framework_init_steps.yaml @@ -62,8 +62,8 @@ module_tags: MasterKvRouter: [master] MetricReporter: [master, owner, external] ClientKvApi: [owner] - ClientSegPool: [owner] - ClientTransferEngine: [master, owner] + ClientSegPool: [owner, external] + ClientTransferEngine: [master, owner, external] ExternalClientApi: [external] MasterLeaseManager: [master] diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs index 8e37608..809ca0e 100755 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/delete.rs @@ -16,21 +16,6 @@ use crate::{ }; use limit_thirdparty::tokio; -fn remove_local_cached_info_for_delete( - client_inner: &ClientKvApiInner, - delete_item: &DeleteClientKvMetaCacheItem, -) { - client_inner - .get_cached_info - .remove_if(&delete_item.key, |_, v| { - if v.put_time_ms == delete_item.put_time_ms { - v.put_version <= delete_item.put_version - } else { - v.put_time_ms <= delete_item.put_time_ms - } - }); -} - impl ClientKvApiInner { pub async fn delete(&self, key: &str) -> KvResult<()> { if !self.view.register_shutdown_poller().is_running() { @@ -64,14 +49,12 @@ impl ClientKvApiInner { resp.serialize_part.error_json.clone(), )?; - remove_local_cached_info_for_delete( - self, - &DeleteClientKvMetaCacheItem { - key: key.to_string(), - put_time_ms: resp.serialize_part.deleted_put_time_ms, - put_version: resp.serialize_part.deleted_put_version, - }, - ); + let delete_item = DeleteClientKvMetaCacheItem { + key: key.to_string(), + put_time_ms: resp.serialize_part.deleted_put_time_ms, + put_version: resp.serialize_part.deleted_put_version, + }; + apply_delete_client_kv_meta_cache_item(&self.view, &delete_item).await; Ok(()) } @@ -99,6 +82,92 @@ pub fn spawn_owner_delete_ack_batch( }); } +async fn apply_delete_client_kv_meta_cache_item( + view: &ClientKvApiView, + delete_item: &DeleteClientKvMetaCacheItem, +) { + let client_api = view.client_kv_api(); + let client_inner = client_api.inner(); + + let (cached_info, pending_get) = { + let controls = client_inner.owner_key_control.lock_key(&delete_item.key); + if controls + .get(&delete_item.key) + .is_some_and(|state| state.local_access_fenced()) + { + tracing::debug!( + "skip legacy local-index delete while owner reclaim fence is active: key={}", + delete_item.key + ); + return; + } + let cached_info = client_inner + .get_cached_info + .remove_if(&delete_item.key, |_, v| { + let res = if v.put_time_ms == delete_item.put_time_ms { + v.put_version <= delete_item.put_version + } else { + v.put_time_ms <= delete_item.put_time_ms + }; + if res { + tracing::debug!("do remove local cache for key: {}", delete_item.key); + } else { + tracing::debug!( + "skip remove local cache for key: {}, request ({},{}), local ({},{})", + delete_item.key, + delete_item.put_time_ms, + delete_item.put_version, + v.put_time_ms, + v.put_version + ); + } + res + }) + .map(|(_, cached_info)| cached_info); + let _ = client_inner + .local_snapshot_info + .remove_if(&delete_item.key, |_, snapshot| { + if snapshot.put_time_ms == delete_item.put_time_ms { + snapshot.put_version <= delete_item.put_version + } else { + snapshot.put_time_ms <= delete_item.put_time_ms + } + }); + let pending_get = client_inner + .pending_local_get_info + .remove_if(&delete_item.key, |_, pending| { + if pending.put_id.0 == delete_item.put_time_ms { + pending.put_id.1 <= delete_item.put_version + } else { + pending.put_id.0 <= delete_item.put_time_ms + } + }) + .map(|(_, pending)| pending); + (cached_info, pending_get) + }; + drop(pending_get); + if let Some(cached_info) = cached_info { + client_inner.owner_hot_invalidate_version( + &delete_item.key, + (cached_info.put_time_ms, cached_info.put_version), + ); + client_inner.release_local_reserve_route_for_memory_info(cached_info.mem_holder.as_ref()); + } + + if let Err(err) = client_inner + .external_invalidate_delete + .sender() + .send(delete_item.clone()) + .await + { + tracing::warn!( + "Failed to enqueue external weak-index invalidation for key '{}': {}", + delete_item.key, + err + ); + } +} + /// 批量删除客户端 KV 元数据缓存的处理函数 pub async fn handle_batch_delete_client_kv_meta_cache( view: &ClientKvApiView, @@ -111,9 +180,6 @@ pub async fn handle_batch_delete_client_kv_meta_cache( req.serialize_part.delete_items.len() ); - let client_api = view.client_kv_api(); - let client_inner = client_api.inner(); - let mut deleted_count = 0u32; for delete_item in &req.serialize_part.delete_items { @@ -124,21 +190,8 @@ pub async fn handle_batch_delete_client_kv_meta_cache( delete_item.put_version ); - remove_local_cached_info_for_delete(client_inner, delete_item); + apply_delete_client_kv_meta_cache_item(view, delete_item).await; deleted_count += 1; - - if let Err(err) = client_inner - .external_invalidate_delete - .sender() - .send(delete_item.clone()) - .await - { - tracing::warn!( - "Failed to enqueue external weak-index invalidation for key '{}': {}", - delete_item.key, - err - ); - } } tracing::debug!( diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs index 60986be..c63a298 100755 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/external_api.rs @@ -1,33 +1,2435 @@ +use crate::client_kv_api::ClientKvApi; +use crate::client_kv_api::get::{StartedGetRevokeCleanup, finish_started_get_revoke_cleanup}; use crate::client_kv_api::msg_pack::{ - ExternalDeleteReq, ExternalDeleteResp, ExternalGetReq, ExternalGetResp, ExternalIsExistReq, - ExternalIsExistResp, ExternalPutCommitReq, ExternalPutCommitResp, ExternalPutRevokeReq, - ExternalPutRevokeResp, ExternalPutStartReq, ExternalPutStartResp, ExternalPutTransferEndReq, - ExternalPutTransferEndResp, TestPutPhaseTrace, + ExternalBatchGetCancelPlan, ExternalBatchGetCancelReq, ExternalBatchGetCancelResp, + ExternalBatchGetItemResp, ExternalBatchGetLocalProbeReq, ExternalBatchGetLocalProbeResp, + ExternalBatchGetReq, ExternalBatchGetResp, ExternalBatchGetStartReq, ExternalBatchGetStartResp, + ExternalBatchGetStartTransferPlan, ExternalBatchGetTransferReq, ExternalBatchGetTransferResp, + ExternalBatchIsExistReq, ExternalBatchIsExistResp, ExternalBatchPutCommitItemResp, + ExternalBatchPutCommitReq, ExternalBatchPutCommitResp, ExternalBatchPutStartItemResp, + ExternalBatchPutStartReq, ExternalBatchPutStartResp, ExternalBatchPutTransferEndItemResp, + ExternalBatchPutTransferEndReq, ExternalBatchPutTransferEndResp, ExternalDeleteReq, + ExternalDeleteResp, ExternalExecutePlannedGetReq, ExternalExecutePlannedGetResp, + ExternalGetReq, ExternalGetResp, ExternalIsExistReq, ExternalIsExistResp, + ExternalObservabilitySnapshotReq, ExternalObservabilitySnapshotResp, ExternalPutCommitReq, + ExternalPutCommitResp, ExternalPutRevokeReq, ExternalPutRevokeResp, ExternalPutStartReq, + ExternalPutStartResp, ExternalPutTransferEndReq, ExternalPutTransferEndResp, TestPutPhaseTrace, +}; +use crate::client_kv_api::{ + self, ExternalGetKeyInterest, ExternalGetKeySharedOp, ExternalGetKeySharedPhase, + ExternalGetStartEntry, ExternalGetStartOwnerItem, ExternalGetStartPrefixResult, + ExternalGetStartSharedItemResult, ExternalGetStartTransferOutput, + ExternalLocalFirstPutKeyReservation, ExternalPendingPutCtx, ExternalPutKeyOutcome, + OwnerLocalPublishItem, OwnerLocalPublishJob, OwnerLocalReserveSlotLease, + OwnerLocalReserveSlotRef, }; -use crate::client_kv_api::{ClientKvApi, ExternalHoldingGetInfo, ExternalPendingPutCtx}; use crate::client_seg_pool::{ResolveSideTransferLaneReq, parse_side_transfer_worker_lane_idx}; use crate::cluster_manager::NodeIDString; use crate::cluster_manager::{ META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME, }; -use crate::memholder::ExternalMemHolderInfo; +use crate::master_kv_router::msg_pack::{ + BatchGetStartItemResp, BatchGetStartResp, BatchPutDoneItemReq, BatchPutRevokeItemReq, + GetPreparedLocalReserveTarget, PutAtomicGroup, PutDoneCommittedSlot, + build_put_atomic_group_assignments, +}; use crate::memholder::MemholderManagerTrait; use crate::memholder::NodeHolderKey; +use crate::memholder::{UserMemHolder, UserMemHolderExposeKind}; use crate::p2p::msg_pack::MsgPack; use crate::rpcresp_kvresult_convert::FromError; use crate::rpcresp_kvresult_convert::ToResult; -use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, OK}; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, OK, codes_api}; use async_trait::async_trait; +use futures::StreamExt; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use tracing; -fn duration_to_i64_us(duration: std::time::Duration) -> i64 { - duration.as_micros().min(i64::MAX as u128) as i64 +fn duration_to_i64_us(duration: std::time::Duration) -> i64 { + duration.as_micros().min(i64::MAX as u128) as i64 +} + +const SIDE_TRANSFER_OWNER_RPC_TIMEOUT_SECS: u64 = 30; +const SIDE_TRANSFER_TARGET_RESOLVE_TIMEOUT_SECS: u64 = 10; +// A handle can legitimately wait behind one full SGLang request-timeout +// window before layerwise restore consumes it. Keep a bounded crash cleanup +// lease, but do not expire a live plan earlier than the supported 300-second +// request window used by the aligned agent workload. +const EXTERNAL_GET_START_HANDLE_TTL: Duration = Duration::from_secs(360); +const EXTERNAL_GET_START_HANDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(5); +// A BatchGetStart transport error has an unknown commit point: the master may +// already hold the prepared target even though the owner never received its +// get_id. Keep such slots out of reuse until the master's 60-second inflight +// Get TTL has elapsed, with a small scheduling margin. +const PREPARED_GET_START_UNCERTAIN_QUARANTINE: Duration = Duration::from_secs(65); +// PutStart has the same unknown-commit-point problem as GetStart. Keep the +// owner reclaim fence alive beyond the master's 60-second inflight Put TTL if +// the caller future is cancelled or the RPC returns a transport error. +const EXTERNAL_PUT_START_UNCERTAIN_QUARANTINE: Duration = Duration::from_secs(65); +// Keep control-plane batching, but bound the failure domain of one uncertain +// BatchGetDone. Atomic groups are never split merely to satisfy this soft cap. +const EXTERNAL_GET_FINISH_TARGET_KEYS_PER_BATCH: usize = 128; +const EXTERNAL_GET_FINISH_BATCH_CONCURRENCY: usize = 4; + +/// Resolve the generation carried by a MemberLeft event without guessing. +/// ClusterEvent::MemberLeft contains only a node id, so a currently live +/// generation makes the event ambiguous (normally a delayed leave after a +/// reconnect). Missing previous-generation metadata is ambiguous as well. +pub(crate) fn external_member_left_departed_epoch( + previous_epoch: Option, + current_epoch: Option, +) -> Option { + current_epoch.is_none().then_some(previous_epoch).flatten() +} + +/// MemberLeft is a cold path. It may scan the handle registry, but removal is +/// conditional on both requester identity and membership generation so a +/// collected handle id can never delete a newer generation's entry. +pub(crate) fn cleanup_external_get_start_handles_for_generation( + registry: &dashmap::DashMap, + req_node_id: &str, + requester_node_start_time: i64, +) -> usize { + let handles = registry + .iter() + .filter_map(|entry| { + let value = entry.value(); + (value.req_node_id == req_node_id + && value.requester_node_start_time == Some(requester_node_start_time)) + .then_some(*entry.key()) + }) + .collect::>(); + + handles + .into_iter() + .filter(|handle| { + registry + .remove_if(handle, |_, value| { + value.req_node_id == req_node_id + && value.requester_node_start_time == Some(requester_node_start_time) + }) + .is_some() + }) + .count() +} + +struct ExternalPutStartFenceClaim { + view: Option, + fence: Option>, + uncertain_delay: Duration, + #[cfg(test)] + uncertain_test_sink: + Option>>>>, +} + +impl ExternalPutStartFenceClaim { + fn new( + view: crate::client_kv_api::ClientKvApiView, + fence: Arc, + ) -> Self { + Self { + view: Some(view), + fence: Some(fence), + uncertain_delay: EXTERNAL_PUT_START_UNCERTAIN_QUARANTINE, + #[cfg(test)] + uncertain_test_sink: None, + } + } + + #[cfg(test)] + fn new_for_cancellation_test( + fence: Arc, + sink: Arc>>>, + ) -> Self { + Self { + view: None, + fence: Some(fence), + uncertain_delay: Duration::ZERO, + uncertain_test_sink: Some(sink), + } + } + + fn take_for_pending_context(&mut self) -> Arc { + self.fence + .take() + .expect("external PutStart fence claim is armed") + } + + fn release_after_definite_response(&mut self) { + drop(self.fence.take()); + } +} + +impl Drop for ExternalPutStartFenceClaim { + fn drop(&mut self) { + let Some(fence) = self.fence.take() else { + return; + }; + #[cfg(test)] + if let Some(sink) = self.uncertain_test_sink.take() { + let replaced = sink.lock().replace(fence); + assert!(replaced.is_none(), "cancellation test sink must be empty"); + return; + } + let delay = self.uncertain_delay; + let view = self + .view + .take() + .expect("production uncertain PutStart fence requires owner view"); + let spawn_view = view.clone(); + let worker_view = view; + spawn_view.spawn("external_put_start_uncertain_fence", async move { + let mut shutdown_waiter = worker_view.register_shutdown_waiter(); + tokio::select! { + _ = tokio::time::sleep(delay) => {} + _ = shutdown_waiter.wait() => return, + } + drop(fence); + }); + } +} + +pub(crate) fn spawn_external_get_start_handle_sweeper(view: crate::client_kv_api::ClientKvApiView) { + let spawn_view = view.clone(); + let worker_view = view; + spawn_view.spawn("external_get_start_handle_sweeper", async move { + let mut shutdown_waiter = worker_view.register_shutdown_waiter(); + let mut interval = tokio::time::interval(EXTERNAL_GET_START_HANDLE_SWEEP_INTERVAL); + loop { + tokio::select! { + _ = interval.tick() => { + let expired = worker_view + .client_kv_api() + .inner() + .external_get_start_registry + .iter() + .filter_map(|entry| { + (entry.value().created_at.elapsed() >= EXTERNAL_GET_START_HANDLE_TTL) + .then_some(*entry.key()) + }) + .collect::>(); + if !expired.is_empty() { + let inner = worker_view.client_kv_api().inner(); + let mut removed = 0usize; + for handle in expired { + removed += usize::from(inner.external_get_start_registry.remove(&handle).is_some()); + } + tracing::warn!( + "expired abandoned external Get handles: removed={} ttl_secs={}", + removed, + EXTERNAL_GET_START_HANDLE_TTL.as_secs() + ); + } + } + _ = shutdown_waiter.wait() => break, + } + } + }); +} + +#[derive(Clone)] +struct LocalCommittedCachePublish { + src_offset: u64, + len: u32, +} + +fn local_committed_cache_publish( + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + pending_ctx: &ExternalPendingPutCtx, + req_src_offset: u64, + req_len: u64, +) -> KvResult { + if pending_ctx.src_offset != req_src_offset || pending_ctx.len != req_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} local cache publish request mismatches pending ctx: key={} put_id=({},{}) req_src_offset={} ctx_src_offset={} req_len={} ctx_len={}", + key, + put_id.0, + put_id.1, + req_src_offset, + pending_ctx.src_offset, + req_len, + pending_ctx.len + ), + })); + } + let len = u32::try_from(req_len).map_err(|_| { + KvError::Api(ApiError::Unknown { + detail: format!( + "{op} local cache len does not fit u32: key={} len={}", + key, req_len + ), + }) + })?; + Ok(LocalCommittedCachePublish { + src_offset: req_src_offset, + len, + }) +} + +fn external_pending_put_ctx_missing( + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, +) -> KvError { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} requires a live owner pending Put context: key={} put_id=({},{}) (missing, expired, or already terminal)", + key, put_id.0, put_id.1 + ), + }) +} + +fn require_external_pending_put_ctx( + inner: &client_kv_api::ClientKvApiInner, + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, +) -> KvResult { + require_external_pending_put_ctx_value( + inner + .external_pending_puts + .get(&(key.to_string(), put_id.0, put_id.1)), + op, + key, + put_id, + ) +} + +fn require_external_pending_put_ctx_value( + pending_ctx: Option, + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, +) -> KvResult { + pending_ctx.ok_or_else(|| external_pending_put_ctx_missing(op, key, put_id)) +} + +async fn best_effort_revoke_missing_external_pending_ctx( + inner: &client_kv_api::ClientKvApiInner, + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, +) -> bool { + if let Err(err) = inner.put_revoke(key, put_id).await { + tracing::warn!( + "{} could not clean missing pending Put context at master: key={} put_id=({},{}) err={}", + op, + key, + put_id.0, + put_id.1, + err + ); + false + } else { + true + } +} + +fn external_local_first_error_item(err: &KvError) -> ExternalBatchPutStartItemResp { + ExternalBatchPutStartItemResp { + put_id: None, + ..ExternalBatchPutStartItemResp::from_error(err) + } +} + +async fn commit_external_local_first_pending( + inner: &client_kv_api::ClientKvApiInner, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ctx: &ExternalPendingPutCtx, + req_src_offset: u64, + req_len: u64, + op: &'static str, +) -> KvResult { + if ctx.peer_id.is_some() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} local-first pending ctx must not carry peer_id: key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + })); + } + if ctx.src_offset != req_src_offset || ctx.len != req_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} local-first request mismatches pending ctx: key={} put_id=({},{}) req_src_offset={} ctx_src_offset={} req_len={} ctx_len={}", + key, put_id.0, put_id.1, req_src_offset, ctx.src_offset, req_len, ctx.len + ), + })); + } + let slot_ref = ctx.local_reserve_slot.as_ref().ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} pending ctx is not local-first: key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }) + })?; + let slot_size = ctx.local_reserve_slot_size.ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "{op} local-first pending ctx missing slot_size: key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }) + })?; + let len = u32::try_from(req_len).map_err(|_| { + KvError::Api(ApiError::Unknown { + detail: format!( + "{op} local-first len does not fit u32: key={} len={}", + key, req_len + ), + }) + })?; + let memory_info = inner + .build_local_reserve_resident_memory_info( + key, + slot_ref.ptr, + len, + slot_size, + slot_ref.grant_id, + slot_ref.slot_index, + ) + .await; + inner.install_precommit_local_visible_memory_info(key, memory_info.clone()); + ctx._pending_fence.disarm_local_slot_lease(); + Ok(PutDoneCommittedSlot { + grant_id: slot_ref.grant_id, + slot_index: slot_ref.slot_index, + slot_size, + addr: slot_ref.ptr, + base_addr: slot_ref.base_addr, + len: req_len, + }) +} + +async fn release_external_local_first_pending_slot( + inner: &client_kv_api::ClientKvApiInner, + ctx: &ExternalPendingPutCtx, + op: &'static str, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, +) { + if let Err(err) = ctx._pending_fence.release_local_slot_lease_now(inner).await { + tracing::error!( + "{} could not release local-first pending slot: key={} put_id=({},{}) err={}", + op, + key, + put_id.0, + put_id.1, + err + ); + } +} + +fn spawn_external_local_first_publish( + inner: &client_kv_api::ClientKvApiInner, + items: Vec, + pending_contexts: Vec, +) { + assert_eq!( + items.len(), + pending_contexts.len(), + "external local-first publish job must retain one context per item" + ); + let view = inner.view.clone_view(); + let spawn_view = view.clone(); + spawn_view.spawn("external_local_first_route_publish", async move { + client_kv_api::put::publish_owner_local_job( + view, + OwnerLocalPublishJob { + items, + key_reservation_ids: Vec::new(), + external_pending_contexts: pending_contexts, + }, + ) + .await; + }); +} + +pub(crate) fn normalize_external_get_start_group_lens( + keys_len: usize, + atomic_group_lens: Option>, +) -> KvResult> { + match atomic_group_lens { + Some(group_lens) => { + if group_lens.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_get_start atomic_group_lens must be non-empty" + .to_string(), + })); + } + if group_lens.iter().any(|len| *len == 0) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_get_start atomic_group_lens entries must be > 0" + .to_string(), + })); + } + let group_sum = group_lens.iter().sum::(); + if group_sum != keys_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_get_start atomic_group_lens must sum to keys length: sum={} keys={}", + group_sum, keys_len + ), + })); + } + Ok(group_lens) + } + None => Ok(vec![keys_len]), + } +} + +fn normalize_external_put_start_group_lens( + items_len: usize, + atomic_group_lens: Option>, +) -> KvResult> { + let group_lens = atomic_group_lens.unwrap_or_else(|| vec![1; items_len]); + if items_len != 0 && group_lens.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_put_start atomic_group_lens must be non-empty".to_string(), + })); + } + if group_lens.iter().any(|group_len| *group_len == 0) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_put_start atomic_group_lens entries must be > 0".to_string(), + })); + } + let sum = group_lens + .iter() + .try_fold(0usize, |sum, group_len| sum.checked_add(*group_len)) + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_put_start atomic_group_lens sum overflowed usize" + .to_string(), + }) + })?; + if sum != items_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_put_start atomic_group_lens must sum to items length: sum={} items={}", + sum, items_len + ), + })); + } + Ok(group_lens) +} + +fn external_get_start_error_kind_from_code(error_code: u32, error_json: &str) -> Option { + if error_code == OK || error_code == codes_api::API_KEY_NOT_FOUND { + None + } else if error_json.is_empty() { + Some(format!("error_code:{}", error_code)) + } else { + Some(format!("error_code:{} {}", error_code, error_json)) + } +} + +fn compute_external_get_start_raw_prefix( + item_codes: &[(u32, String)], +) -> (usize, Option, Option) { + let mut prefix_hit_len = 0usize; + for (idx, (error_code, error_json)) in item_codes.iter().enumerate() { + if *error_code == OK { + prefix_hit_len += 1; + continue; + } + return ( + prefix_hit_len, + Some(idx), + external_get_start_error_kind_from_code(*error_code, error_json), + ); + } + (prefix_hit_len, None, None) +} + +pub(crate) fn compute_external_get_start_transfer_prefix( + raw_prefix_hit_len: usize, + group_lens: &[usize], + prefix_best_effort: bool, +) -> usize { + let mut transferable_len = 0usize; + for group_len in group_lens.iter().copied() { + let next = transferable_len + group_len; + if next > raw_prefix_hit_len { + break; + } + transferable_len = next; + } + let requested_len = group_lens.iter().sum::(); + if !prefix_best_effort && transferable_len != requested_len { + 0 + } else { + transferable_len + } +} + +pub(crate) fn validate_external_get_consume_prefix( + consume_prefix_len: usize, + transferable_len: usize, + group_lens: &[usize], +) -> KvResult<()> { + if consume_prefix_len == 0 || consume_prefix_len > transferable_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external get_transfer consume_prefix_len must be within the live prefix: consume={} transferable={}", + consume_prefix_len, transferable_len + ), + })); + } + + let mut group_end = 0usize; + for group_len in group_lens.iter().copied() { + group_end = group_end.checked_add(group_len).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: "external get_transfer atomic-group boundary overflowed usize".to_string(), + }) + })?; + if group_end >= consume_prefix_len { + break; + } + } + if group_end != consume_prefix_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external get_transfer consume_prefix_len splits an atomic group: consume={} group_lens={:?}", + consume_prefix_len, group_lens + ), + })); + } + Ok(()) +} + +fn collect_external_get_start_missing( + keys: &[String], + item_slots: &[Option], +) -> (Vec, Vec) { + assert_eq!( + keys.len(), + item_slots.len(), + "external get_start local snapshot length must match keys" + ); + keys.iter() + .zip(item_slots.iter()) + .enumerate() + .filter_map(|(idx, (key, item))| item.is_none().then(|| (idx, key.clone()))) + .unzip() +} + +fn collect_all_local_external_get_start_infos( + items: &[ExternalGetStartOwnerItem], +) -> Option>> { + items + .iter() + .map(|item| match item { + ExternalGetStartOwnerItem::Local { memory_info } => Some(memory_info.clone()), + ExternalGetStartOwnerItem::Shared { .. } => None, + }) + .collect() +} + +fn prepared_target_from_slot( + slot_size: u64, + slot: &OwnerLocalReserveSlotRef, +) -> GetPreparedLocalReserveTarget { + GetPreparedLocalReserveTarget { + grant_id: slot.grant_id, + slot_index: slot.slot_index, + slot_size, + addr: slot.ptr, + base_addr: slot.base_addr, + } +} + +async fn release_prepared_get_target( + inner: &client_kv_api::ClientKvApiInner, + target: &GetPreparedLocalReserveTarget, +) -> KvResult<()> { + inner + .owner_release_local_reserve_slot_lease(OwnerLocalReserveSlotLease { + value_len: target.slot_size, + slot_size: target.slot_size, + slots: vec![OwnerLocalReserveSlotRef { + grant_id: target.grant_id, + slot_index: target.slot_index, + ptr: target.addr, + base_addr: target.base_addr, + }], + }) + .await +} + +/// Cancellation-safe ownership for slots claimed before BatchGetStart. +/// Accepted slots are disarmed one by one and become owned by their per-key +/// flight. Any slots still present when the future is dropped are returned by +/// a registered owner task, so an RPC cancellation cannot strand Prepared +/// state in the local-reserve pool. +struct PreparedGetSlotClaimGuard { + view: crate::client_kv_api::ClientKvApiView, + lease: Option, + drop_delay: Duration, +} + +impl PreparedGetSlotClaimGuard { + fn new(view: crate::client_kv_api::ClientKvApiView, lease: OwnerLocalReserveSlotLease) -> Self { + Self { + view, + lease: Some(lease), + drop_delay: PREPARED_GET_START_UNCERTAIN_QUARANTINE, + } + } + + fn lease(&self) -> &OwnerLocalReserveSlotLease { + self.lease.as_ref().expect("prepared slot claim is armed") + } + + fn disarm_accepted(&mut self, target: &GetPreparedLocalReserveTarget) { + let lease = self.lease.as_mut().expect("prepared slot claim is armed"); + lease.slots.retain(|slot| { + !(slot.grant_id == target.grant_id + && slot.slot_index == target.slot_index + && slot.ptr == target.addr + && slot.base_addr == target.base_addr) + }); + } + + fn mark_start_response_received(&mut self) { + self.drop_delay = Duration::ZERO; + } + + fn handoff_started( + &mut self, + items: &[crate::master_kv_router::msg_pack::BatchGetStartItemResp], + ) -> Vec { + items + .iter() + .filter(|item| item.error_code == OK) + .map(|item| { + if let Some(target) = item.prepared_target.as_ref() { + self.disarm_accepted(target); + } + StartedGetRevokeCleanup { + get_id: item.get_id, + prepared_target: item.prepared_target.clone(), + } + }) + .collect() + } +} + +impl Drop for PreparedGetSlotClaimGuard { + fn drop(&mut self) { + let Some(lease) = self.lease.take() else { + return; + }; + if lease.slots.is_empty() { + return; + } + let drop_delay = self.drop_delay; + let spawn_view = self.view.clone(); + let worker_view = spawn_view.clone(); + spawn_view.spawn("prepared_get_slot_drop_cleanup", async move { + if !drop_delay.is_zero() { + tracing::warn!( + "quarantining prepared Get slots after uncertain BatchGetStart: slots={} delay_secs={}", + lease.slots.len(), + drop_delay.as_secs() + ); + let mut shutdown_waiter = worker_view.register_shutdown_waiter(); + tokio::select! { + _ = tokio::time::sleep(drop_delay) => {} + _ = shutdown_waiter.wait() => return, + } + } + if let Err(err) = worker_view + .client_kv_api() + .inner() + .owner_release_local_reserve_slot_lease(lease) + .await + { + tracing::error!("prepared Get slot Drop cleanup failed: {}", err); + } + }); + } +} + +async fn batch_get_start_with_local_reserve_targets( + inner: &client_kv_api::ClientKvApiInner, + keys: &[String], +) -> KvResult { + let Some(value_len) = inner + .test_spec_config + .owner_local_reserve_expected_capacity + .as_ref() + .map(|expected| expected.value_len) + else { + return inner.batch_get_start(keys.to_vec()).await; + }; + let lease = inner + .owner_claim_local_reserve_slot_lease(value_len, keys.len()) + .await?; + let mut claim_guard = PreparedGetSlotClaimGuard::new(inner.view.clone_view(), lease); + let prepared_targets = claim_guard + .lease() + .slots + .iter() + .map(|slot| { + Some(prepared_target_from_slot( + claim_guard.lease().slot_size, + slot, + )) + }) + .collect::>(); + let response_result = inner + .batch_get_start_with_prepared_targets(keys.to_vec(), prepared_targets.clone()) + .await; + let response = match response_result { + Ok(response) => { + claim_guard.mark_start_response_received(); + response + } + Err(err) => return Err(err), + }; + if response.items.len() != keys.len() { + let started = claim_guard.handoff_started(&response.items); + if !started.is_empty() { + finish_started_get_revoke_cleanup(inner, started, "BatchGetStart length mismatch") + .await; + } + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_get_start response length mismatch: expected={} got={}", + keys.len(), + response.items.len() + ), + })); + } + + let accepted_exactly = + response + .items + .iter() + .zip(prepared_targets.iter()) + .all(|(item, requested_target)| { + item.error_code != OK || item.prepared_target.as_ref() == requested_target.as_ref() + }); + if !accepted_exactly { + let started = claim_guard.handoff_started(&response.items); + finish_started_get_revoke_cleanup(inner, started, "BatchGetStart target mismatch").await; + return Err(KvError::Api(ApiError::Unknown { + detail: "master did not accept the exact prepared local-reserve Get target".to_string(), + })); + } + for (item, requested_target) in response.items.iter().zip(prepared_targets.iter()) { + if item.error_code == OK { + claim_guard.disarm_accepted( + requested_target + .as_ref() + .expect("external prepared target vector must be dense"), + ); + } + } + Ok(response) +} + +async fn batch_get_bind_with_local_reserve_targets( + inner: &client_kv_api::ClientKvApiInner, + get_ids: &[u64], +) -> KvResult> { + if get_ids.is_empty() { + return Ok(Vec::new()); + } + let Some(value_len) = inner + .test_spec_config + .owner_local_reserve_expected_capacity + .as_ref() + .map(|expected| expected.value_len) + else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "planned external Get requires owner-local reserve capacity".to_string(), + })); + }; + let lease = inner + .owner_claim_local_reserve_slot_lease(value_len, get_ids.len()) + .await?; + let mut claim_guard = PreparedGetSlotClaimGuard::new(inner.view.clone_view(), lease); + let prepared_targets = claim_guard + .lease() + .slots + .iter() + .map(|slot| prepared_target_from_slot(claim_guard.lease().slot_size, slot)) + .collect::>(); + let response = match inner + .batch_get_bind_prepared_targets(get_ids.to_vec(), prepared_targets.clone()) + .await + { + Ok(response) => { + claim_guard.mark_start_response_received(); + response + } + Err(err) => return Err(err), + }; + if response.items.len() != get_ids.len() { + let started = claim_guard.handoff_started(&response.items); + finish_started_get_revoke_cleanup(inner, started, "BatchGetBind length mismatch").await; + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned BatchGetBind response length mismatch: expected={} got={}", + get_ids.len(), + response.items.len() + ), + })); + } + let identities_match = response + .items + .iter() + .zip(get_ids) + .all(|(item, get_id)| item.get_id == *get_id); + let targets_match = + response + .items + .iter() + .zip(prepared_targets.iter()) + .all(|(item, expected)| { + item.error_code != OK || item.prepared_target.as_ref() == Some(expected) + }); + if !identities_match || !targets_match { + let started = claim_guard.handoff_started(&response.items); + finish_started_get_revoke_cleanup(inner, started, "BatchGetBind identity/target mismatch") + .await; + return Err(KvError::Api(ApiError::Unknown { + detail: "master did not replay the exact planned Get identity and prepared target" + .to_string(), + })); + } + for (item, target) in response.items.iter().zip(prepared_targets.iter()) { + if item.error_code == OK { + claim_guard.disarm_accepted(target); + } + } + Ok(response.items) +} + +#[cfg(test)] +mod external_put_pending_tests { + use super::{ExternalPutStartFenceClaim, require_external_pending_put_ctx_value}; + use crate::client_kv_api::{ + ExternalPendingPutFenceGuard, OwnerKeyControlTable, + acquire_external_pending_put_fence_for_key, + }; + use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError}; + use parking_lot::Mutex; + use std::sync::Arc; + + #[test] + fn missing_or_expired_pending_context_is_not_a_commit_capability() { + let result = require_external_pending_put_ctx_value( + None, + "external_put_commit", + "missing-key", + (17, 3), + ); + let Err(KvError::Api(ApiError::InvalidArgument { detail })) = result else { + panic!("missing pending context must be rejected") + }; + assert!(detail.contains("requires a live owner pending Put context")); + assert!(detail.contains("put_id=(17,3)")); + } + + #[test] + fn cancelled_put_start_hands_fence_to_uncertainty_quarantine() { + let controls = Arc::new(OwnerKeyControlTable::default()); + let fence = acquire_external_pending_put_fence_for_key(&controls, "cancelled-start") + .expect("pending fence acquisition must succeed"); + let sink: Arc>>> = + Arc::new(Mutex::new(None)); + + let claim = ExternalPutStartFenceClaim::new_for_cancellation_test(fence, sink.clone()); + drop(claim); + assert_eq!( + controls.lock_key("cancelled-start")["cancelled-start"].external_pending_puts, + 1, + "cancellation must transfer, not drop, the reclaim fence" + ); + + drop(sink.lock().take()); + assert!( + controls + .lock_key("cancelled-start") + .get("cancelled-start") + .is_none() + ); + } +} + +#[cfg(test)] +mod external_get_start_batch_tests { + use super::{ + abandon_unstarted_external_get_key_locked, + cleanup_external_get_start_handles_for_generation, clear_external_get_key_marker_locked, + collect_external_get_start_missing, compute_external_get_start_raw_prefix, + compute_external_get_start_transfer_prefix, decide_external_get_key_item, + external_member_left_departed_epoch, normalize_external_put_start_group_lens, + observe_external_get_consume_phases, partition_external_get_finish_leaders, + register_external_get_key_under_fence, validate_external_get_consume_prefix, + }; + use crate::client_kv_api::{ + ExternalGetKeyInterest, ExternalGetKeySharedOp, ExternalGetKeySharedPhase, + ExternalGetStartEntry, ExternalGetStartOwnerItem, ExternalGetStartSharedItemResult, + OwnerKeyControlState, + }; + use crate::master_kv_router::msg_pack::{ + BatchGetStartItemResp, PutAtomicGroup, PutAtomicGroupMember, + }; + use crate::rpcresp_kvresult_convert::msg_and_error::{OK, codes_api}; + use dashmap::DashMap; + use std::collections::HashMap; + use std::sync::Arc; + use std::time::{Duration, Instant}; + + fn abandoned_handle_entry( + req_node_id: &str, + requester_node_start_time: Option, + ) -> ExternalGetStartEntry { + ExternalGetStartEntry { + req_node_id: req_node_id.to_string(), + requester_node_start_time, + keys: Vec::new(), + items: Vec::new(), + atomic_group_lens: Vec::new(), + created_at: Instant::now(), + } + } + + #[test] + fn member_left_requires_previous_epoch_and_skips_delayed_leave_after_reconnect() { + assert_eq!( + external_member_left_departed_epoch(Some(11), None), + Some(11) + ); + assert_eq!(external_member_left_departed_epoch(None, None), None); + assert_eq!( + external_member_left_departed_epoch(Some(11), Some(12)), + None + ); + assert_eq!(external_member_left_departed_epoch(None, Some(12)), None); + } + + #[test] + fn member_left_removes_only_matching_requester_handle_generation() { + let registry = DashMap::new(); + registry.insert(1, abandoned_handle_entry("external-a", Some(11))); + registry.insert(2, abandoned_handle_entry("external-a", Some(12))); + registry.insert(3, abandoned_handle_entry("external-b", Some(11))); + registry.insert(4, abandoned_handle_entry("external-a", None)); + + assert_eq!( + cleanup_external_get_start_handles_for_generation(®istry, "external-a", 11), + 1 + ); + assert!(!registry.contains_key(&1)); + assert!(registry.contains_key(&2)); + assert!(registry.contains_key(&3)); + assert!(registry.contains_key(&4)); + } + + #[test] + fn all_local_batch_has_no_master_fallback_and_keeps_full_atomic_prefix() { + let keys = vec!["a".to_string(), "b".to_string(), "c".to_string()]; + let local_slots = vec![Some(1_u8), Some(2), Some(3)]; + + let (missing_indices, missing_keys) = + collect_external_get_start_missing(&keys, &local_slots); + assert!(missing_indices.is_empty()); + assert!(missing_keys.is_empty()); + + let item_codes = vec![(OK, String::new()); keys.len()]; + let (raw_prefix, first_miss, first_error) = + compute_external_get_start_raw_prefix(&item_codes); + assert_eq!(raw_prefix, keys.len()); + assert_eq!(first_miss, None); + assert_eq!(first_error, None); + assert_eq!( + compute_external_get_start_transfer_prefix(raw_prefix, &[1, 2], true), + keys.len() + ); + } + + #[test] + fn mixed_batch_falls_back_only_for_missing_keys_and_preserves_group_boundary() { + let keys = vec![ + "a".to_string(), + "b".to_string(), + "c".to_string(), + "d".to_string(), + ]; + let local_slots = vec![Some(1_u8), Some(2), None, Some(4)]; + + let (missing_indices, missing_keys) = + collect_external_get_start_missing(&keys, &local_slots); + assert_eq!(missing_indices, vec![2]); + assert_eq!(missing_keys, vec!["c"]); + + let item_codes = vec![ + (OK, String::new()), + (OK, String::new()), + (codes_api::API_KEY_NOT_FOUND, String::new()), + (OK, String::new()), + ]; + let (raw_prefix, first_miss, first_error) = + compute_external_get_start_raw_prefix(&item_codes); + assert_eq!(raw_prefix, 2); + assert_eq!(first_miss, Some(2)); + assert_eq!(first_error, None); + assert_eq!( + compute_external_get_start_transfer_prefix(raw_prefix, &[2, 2], true), + 2 + ); + assert_eq!( + compute_external_get_start_transfer_prefix(raw_prefix, &[2, 2], false), + 0 + ); + } + + #[test] + fn get_transfer_consume_prefix_must_be_live_and_end_at_group_boundary() { + let group_lens = [2, 2, 3]; + assert!(validate_external_get_consume_prefix(2, 4, &group_lens).is_ok()); + assert!(validate_external_get_consume_prefix(4, 4, &group_lens).is_ok()); + assert!(validate_external_get_consume_prefix(0, 4, &group_lens).is_err()); + assert!(validate_external_get_consume_prefix(1, 4, &group_lens).is_err()); + assert!(validate_external_get_consume_prefix(3, 4, &group_lens).is_err()); + assert!(validate_external_get_consume_prefix(5, 4, &group_lens).is_err()); + } + + #[test] + fn consume_phase_snapshot_distinguishes_ready_data_from_real_wait() { + let finishing = Arc::new(ExternalGetKeySharedOp::new("finishing".to_string())); + finishing.state.lock().phase = ExternalGetKeySharedPhase::Finishing { + item: BatchGetStartItemResp::default(), + }; + let ready = Arc::new(ExternalGetKeySharedOp::new("ready".to_string())); + let observed_at = Instant::now(); + { + let mut state = ready.state.lock(); + state.phase = ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Miss, + }; + state.terminal_at = Some(observed_at - Duration::from_millis(5)); + } + let items = vec![ + ExternalGetStartOwnerItem::Shared { + interest: ExternalGetKeyInterest::new(finishing, false), + }, + ExternalGetStartOwnerItem::Shared { + interest: ExternalGetKeyInterest::new(ready, false), + }, + ]; + + let snapshot = observe_external_get_consume_phases(&items, observed_at); + assert_eq!(snapshot.finishing, 1); + assert_eq!(snapshot.ready, 1); + assert_eq!(snapshot.pending_before_consume(), 1); + assert_eq!(snapshot.terminal_before_consume(), 1); + assert_eq!(snapshot.terminal_age_count, 1); + assert_eq!(snapshot.terminal_age_mean_us(), 5_000); + assert_eq!(snapshot.terminal_age_max_us, 5_000); + } + + #[test] + fn put_groups_default_per_item_and_reject_malformed_partitions() { + assert_eq!( + normalize_external_put_start_group_lens(3, None).unwrap(), + vec![1, 1, 1] + ); + assert_eq!( + normalize_external_put_start_group_lens(3, Some(vec![2, 1])).unwrap(), + vec![2, 1] + ); + assert!(normalize_external_put_start_group_lens(3, Some(vec![1, 1])).is_err()); + assert!(normalize_external_put_start_group_lens(3, Some(vec![1, 0, 2])).is_err()); + } + + #[test] + fn overlapping_nonidentical_batches_share_each_key_but_keep_leaders_batched() { + fn register_batch( + controls: &mut HashMap, + keys: &[&str], + ) -> ( + Vec, + Vec>, + ) { + let mut leaders = Vec::new(); + let items = keys + .iter() + .map(|key| { + let control = controls.entry((*key).to_string()).or_default(); + let interest = + register_external_get_key_under_fence(control, key, &mut leaders); + ExternalGetStartOwnerItem::Shared { interest } + }) + .collect(); + (items, leaders) + } + + let mut controls = HashMap::new(); + let (mut batch_a, leaders_a) = + register_batch(&mut controls, &["a", "shared-b", "shared-c"]); + let (mut batch_b, leaders_b) = + register_batch(&mut controls, &["shared-b", "shared-c", "d"]); + + assert_eq!( + leaders_a.len(), + 3, + "first required batch has one leader subset" + ); + assert_eq!( + leaders_b.len(), + 1, + "second batch starts only its new leader d" + ); + let shared_b_a = match &batch_a[1] { + ExternalGetStartOwnerItem::Shared { interest } => interest.op().clone(), + ExternalGetStartOwnerItem::Local { .. } => unreachable!(), + }; + let shared_b_b = match &batch_b[0] { + ExternalGetStartOwnerItem::Shared { interest } => interest.op().clone(), + ExternalGetStartOwnerItem::Local { .. } => unreachable!(), + }; + assert!(Arc::ptr_eq(&shared_b_a, &shared_b_b)); + + // Batch A keeps only [a,b], while batch B keeps [b,c,d]. Prefix + // decisions are independent, yet each physical key operation is still + // represented by one shared marker. + for (idx, item) in batch_a.iter_mut().enumerate() { + decide_external_get_key_item(item, idx < 2); + } + for item in &mut batch_b { + decide_external_get_key_item(item, true); + } + let state_b = shared_b_a.state.lock(); + assert_eq!(state_b.undecided, 0); + assert_eq!(state_b.retained, 2); + drop(state_b); + let shared_c = controls["shared-c"] + .external_get + .as_ref() + .expect("shared-c marker") + .state + .lock(); + assert_eq!(shared_c.undecided, 0); + assert_eq!(shared_c.retained, 1); + } + + #[test] + fn dropped_pending_interest_retires_decision_without_losing_batch_sharing() { + let key = "cancel-safe-key".to_string(); + let mut controls = HashMap::new(); + let mut first_leaders = Vec::new(); + let mut second_leaders = Vec::new(); + let first = register_external_get_key_under_fence( + controls.entry(key.clone()).or_default(), + &key, + &mut first_leaders, + ); + let mut second = register_external_get_key_under_fence( + controls.entry(key.clone()).or_default(), + &key, + &mut second_leaders, + ); + assert_eq!(first_leaders.len(), 1); + assert!(second_leaders.is_empty()); + assert!(Arc::ptr_eq(first.op(), second.op())); + assert_eq!(first.op().state.lock().undecided, 2); + + drop(first); + assert_eq!(second.op().state.lock().undecided, 1); + second.decide(true); + let state = second.op().state.lock(); + assert_eq!(state.undecided, 0); + assert_eq!(state.retained, 1); + } + + #[limit_thirdparty::tokio::test] + async fn aborting_waiter_future_retires_pending_interest_and_wakes_flight() { + let op = Arc::new(ExternalGetKeySharedOp::new("abort-safe-key".to_string())); + let waiter_op = op.clone(); + let (registered_tx, registered_rx) = ::tokio::sync::oneshot::channel(); + let waiter = ::tokio::spawn(async move { + let _interest = crate::client_kv_api::ExternalGetKeyInterest::new(waiter_op, true); + let _ = registered_tx.send(()); + futures::future::pending::<()>().await; + }); + + registered_rx.await.expect("waiter registered interest"); + assert_eq!(op.state.lock().undecided, 1); + let notified = op.notify.notified(); + futures::pin_mut!(notified); + notified.as_mut().enable(); + + waiter.abort(); + assert!( + waiter + .await + .expect_err("waiter must be aborted") + .is_cancelled() + ); + assert_eq!(op.state.lock().undecided, 0); + ::tokio::time::timeout(std::time::Duration::from_secs(1), notified) + .await + .expect("interest Drop must wake the flight"); + } + + #[test] + fn old_singleflight_cleanup_cannot_remove_new_generation() { + let key = "aba-key".to_string(); + let old = Arc::new(ExternalGetKeySharedOp::new(key.clone())); + let mut controls = HashMap::from([( + key.clone(), + OwnerKeyControlState { + local_puts: 0, + external_pending_puts: 0, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: Some(old.clone()), + local_access_fence: None, + }, + )]); + clear_external_get_key_marker_locked(&mut controls, &old); + assert!(!controls.contains_key(&key)); + + let new = Arc::new(ExternalGetKeySharedOp::new(key.clone())); + controls.entry(key.clone()).or_default().external_get = Some(new.clone()); + clear_external_get_key_marker_locked(&mut controls, &old); + assert!( + controls[&key] + .external_get + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &new)) + ); + } + + #[test] + fn cancelled_planning_leader_becomes_miss_and_clears_only_its_generation() { + let key = "planning-cancel-key".to_string(); + let op = Arc::new(ExternalGetKeySharedOp::new(key.clone())); + let mut controls = HashMap::from([( + key.clone(), + OwnerKeyControlState { + local_puts: 0, + external_pending_puts: 0, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: Some(op.clone()), + local_access_fence: None, + }, + )]); + + assert!(abandon_unstarted_external_get_key_locked( + &mut controls, + &op + )); + assert!(!controls.contains_key(&key)); + assert!(matches!( + &op.state.lock().phase, + crate::client_kv_api::ExternalGetKeySharedPhase::Ready { + result: crate::client_kv_api::ExternalGetStartSharedItemResult::Miss + } + )); + assert!(!abandon_unstarted_external_get_key_locked( + &mut controls, + &op + )); + } + + #[test] + fn done_partition_keeps_atomic_groups_intact_and_bounds_other_batches() { + let group = PutAtomicGroup { + members: (0..3) + .map(|idx| PutAtomicGroupMember { + key: format!("group-{idx}"), + put_id: (7, idx), + }) + .collect(), + }; + let leader = |key: &str, atomic_group: Option| { + ( + Arc::new(ExternalGetKeySharedOp::new(key.to_string())), + BatchGetStartItemResp { + atomic_group, + ..Default::default() + }, + ) + }; + let batches = partition_external_get_finish_leaders( + vec![ + leader("single-a", None), + leader("group-0", Some(group.clone())), + leader("single-b", None), + leader("group-1", Some(group.clone())), + leader("group-2", Some(group.clone())), + leader("single-c", None), + ], + 2, + ); + + let group_batch_count = batches + .iter() + .filter(|batch| batch.iter().any(|(op, _)| op.key.starts_with("group-"))) + .count(); + assert_eq!(group_batch_count, 1, "one atomic group must not be split"); + let group_batch = batches + .iter() + .find(|batch| batch.iter().any(|(op, _)| op.key == "group-0")) + .unwrap(); + assert_eq!( + group_batch + .iter() + .filter(|(op, _)| op.key.starts_with("group-")) + .count(), + 3 + ); + assert!(batches.iter().all(|batch| { + batch.len() <= 2 + || batch + .iter() + .all(|(_, item)| item.atomic_group.as_ref() == Some(&group)) + })); + } +} + +#[derive(Debug, Default)] +struct ExternalGetConsumePhaseSnapshot { + local: usize, + starting: usize, + started: usize, + finishing: usize, + revoking: usize, + ready: usize, + failed: usize, + terminal_age_count: usize, + terminal_age_sum_us: u64, + terminal_age_max_us: u64, +} + +impl ExternalGetConsumePhaseSnapshot { + fn terminal_before_consume(&self) -> usize { + self.local + .saturating_add(self.ready) + .saturating_add(self.failed) + } + + fn pending_before_consume(&self) -> usize { + self.starting + .saturating_add(self.started) + .saturating_add(self.finishing) + .saturating_add(self.revoking) + } + + fn terminal_age_mean_us(&self) -> u64 { + if self.terminal_age_count == 0 { + 0 + } else { + self.terminal_age_sum_us / self.terminal_age_count as u64 + } + } +} + +fn observe_external_get_consume_phases( + items: &[ExternalGetStartOwnerItem], + observed_at: Instant, +) -> ExternalGetConsumePhaseSnapshot { + let mut snapshot = ExternalGetConsumePhaseSnapshot::default(); + for item in items { + let ExternalGetStartOwnerItem::Shared { interest } = item else { + snapshot.local = snapshot.local.saturating_add(1); + continue; + }; + let state = interest.op().state.lock(); + match &state.phase { + ExternalGetKeySharedPhase::Starting => { + snapshot.starting = snapshot.starting.saturating_add(1) + } + ExternalGetKeySharedPhase::Started { .. } => { + snapshot.started = snapshot.started.saturating_add(1) + } + ExternalGetKeySharedPhase::Finishing { .. } => { + snapshot.finishing = snapshot.finishing.saturating_add(1) + } + ExternalGetKeySharedPhase::Revoking { .. } => { + snapshot.revoking = snapshot.revoking.saturating_add(1) + } + ExternalGetKeySharedPhase::Ready { .. } => { + snapshot.ready = snapshot.ready.saturating_add(1) + } + ExternalGetKeySharedPhase::Failed { .. } => { + snapshot.failed = snapshot.failed.saturating_add(1) + } + } + if matches!( + &state.phase, + ExternalGetKeySharedPhase::Ready { .. } | ExternalGetKeySharedPhase::Failed { .. } + ) { + if let Some(age) = state + .terminal_at + .and_then(|terminal_at| observed_at.checked_duration_since(terminal_at)) + { + let age_us = age.as_micros().min(u64::MAX as u128) as u64; + snapshot.terminal_age_count = snapshot.terminal_age_count.saturating_add(1); + snapshot.terminal_age_sum_us = snapshot.terminal_age_sum_us.saturating_add(age_us); + snapshot.terminal_age_max_us = snapshot.terminal_age_max_us.max(age_us); + } + } + } + snapshot +} + +async fn finish_external_get_start_transfer( + view: crate::client_kv_api::ClientKvApiView, + transfer_items: Vec, + _transfer_concurrency: usize, +) -> KvResult { + let client_api = view.client_kv_api(); + let inner = client_api.inner(); + let refcount = inner.get_or_init_all_memholder_refcount(); + let waits = transfer_items.into_iter().map(|item| { + let refcount = refcount.clone(); + async move { + match item { + ExternalGetStartOwnerItem::Local { memory_info } => Ok(Some(( + Arc::new(UserMemHolder::new( + memory_info, + refcount, + UserMemHolderExposeKind::SegPtr, + )), + None, + ))), + ExternalGetStartOwnerItem::Shared { interest } => { + match wait_external_get_key_result(interest.op().clone()).await? { + ExternalGetStartSharedItemResult::Hit { memholder } => { + Ok(Some((memholder, None))) + } + ExternalGetStartSharedItemResult::Miss => Ok(None), + ExternalGetStartSharedItemResult::Error { + error_code, + error_json, + } => Err(KvError::from_json(error_code, &error_json)), + } + } + } + } + }); + Ok(futures::future::join_all(waits).await) +} + +fn external_get_start_error_parts(err: &KvError) -> (u32, String) { + (err.code(), err.to_json()) +} + +async fn wait_external_get_key_not_revoking(op: Arc) { + loop { + let notified = op.notify.notified(); + futures::pin_mut!(notified); + let should_wait = { + let state = op.state.lock(); + if matches!(state.phase, ExternalGetKeySharedPhase::Revoking { .. }) { + notified.as_mut().enable(); + true + } else { + false + } + }; + if !should_wait { + return; + } + notified.await; + } +} + +fn register_external_get_key_under_fence( + control: &mut client_kv_api::OwnerKeyControlState, + key: &str, + leaders: &mut Vec>, +) -> ExternalGetKeyInterest { + if let Some(op) = control.external_get.clone() { + let mut state = op.state.lock(); + let decision_registered = match state.phase { + ExternalGetKeySharedPhase::Starting | ExternalGetKeySharedPhase::Started { .. } => { + state.undecided = state + .undecided + .checked_add(1) + .expect("external Get singleflight undecided overflow"); + true + } + ExternalGetKeySharedPhase::Finishing { .. } + | ExternalGetKeySharedPhase::Ready { .. } + | ExternalGetKeySharedPhase::Failed { .. } => false, + ExternalGetKeySharedPhase::Revoking { .. } => { + unreachable!("revoking markers were checked under the same fence") + } + }; + drop(state); + ExternalGetKeyInterest::new(op, decision_registered) + } else { + let op = Arc::new(ExternalGetKeySharedOp::new(key.to_string())); + control.external_get = Some(op.clone()); + leaders.push(op.clone()); + ExternalGetKeyInterest::new(op, true) + } +} + +/// Owns newly installed `Starting` markers until the complete request batch is +/// ready to hand them to one BatchGetStart worker. If planning is cancelled +/// while waiting for an older Revoke on another key, Drop converts only these +/// never-started operations to a safe miss and removes their markers. No +/// prepared target or master identity exists yet, so there is nothing to +/// revoke and no storage state to guess. +struct ExternalGetPlanningLeadersGuard<'a> { + inner: &'a client_kv_api::ClientKvApiInner, + leaders: Vec>, + handed_off: bool, +} + +impl<'a> ExternalGetPlanningLeadersGuard<'a> { + fn new(inner: &'a client_kv_api::ClientKvApiInner) -> Self { + Self { + inner, + leaders: Vec::new(), + handed_off: false, + } + } + + fn leaders_mut(&mut self) -> &mut Vec> { + &mut self.leaders + } + + fn handoff(mut self) -> Vec> { + self.handed_off = true; + std::mem::take(&mut self.leaders) + } +} + +impl Drop for ExternalGetPlanningLeadersGuard<'_> { + fn drop(&mut self) { + if self.handed_off { + return; + } + for op in &self.leaders { + let notify = { + let mut controls = self.inner.owner_key_control.lock_key(&op.key); + abandon_unstarted_external_get_key_locked(&mut controls, op) + }; + if notify { + self.inner.untrack_external_get_flight(op); + op.notify.notify_waiters(); + } + } + } +} + +async fn plan_external_get_key_items( + inner: &client_kv_api::ClientKvApiInner, + keys: &[String], +) -> ( + Vec, + Vec>, +) { + enum KeyAttempt { + Wait(Arc), + Ready { + item: ExternalGetStartOwnerItem, + hot_touch: Option<( + crate::master_kv_router::put::PutIDForAKey, + Arc, + )>, + }, + } + + let mut items = Vec::with_capacity(keys.len()); + let mut planning_leaders = ExternalGetPlanningLeadersGuard::new(inner); + let mut hot_touches = Vec::new(); + for key in keys { + loop { + // One short per-key shard section atomically chooses local, + // joiner, or leader. No synchronous lock spans the request batch + // or the wait for an older Revoke to finish. + let attempt = { + let mut controls = inner.owner_key_control.lock_key(key); + let revoking = controls + .get(key) + .and_then(|state| state.external_get.clone()) + .filter(|op| { + matches!( + op.state.lock().phase, + ExternalGetKeySharedPhase::Revoking { .. } + ) + }); + if let Some(op) = revoking { + KeyAttempt::Wait(op) + } else { + let fenced = controls + .get(key) + .is_some_and(|state| state.local_access_fenced()); + if !fenced { + if let Some(memory_info) = inner.local_visible_mem_holder_unfenced(key) { + let hot_touch = inner.get_cached_info.get(key).and_then(|cached| { + Arc::ptr_eq(&cached.mem_holder, &memory_info).then_some(( + (cached.put_time_ms, cached.put_version), + memory_info.clone(), + )) + }); + KeyAttempt::Ready { + item: ExternalGetStartOwnerItem::Local { memory_info }, + hot_touch, + } + } else { + let control = controls.entry(key.clone()).or_default(); + let leader_count = planning_leaders.leaders.len(); + let interest = register_external_get_key_under_fence( + control, + key, + planning_leaders.leaders_mut(), + ); + if planning_leaders.leaders.len() != leader_count { + inner.track_external_get_flight(interest.op()); + } + KeyAttempt::Ready { + item: ExternalGetStartOwnerItem::Shared { interest }, + hot_touch: None, + } + } + } else { + let control = controls.entry(key.clone()).or_default(); + let leader_count = planning_leaders.leaders.len(); + let interest = register_external_get_key_under_fence( + control, + key, + planning_leaders.leaders_mut(), + ); + if planning_leaders.leaders.len() != leader_count { + inner.track_external_get_flight(interest.op()); + } + KeyAttempt::Ready { + item: ExternalGetStartOwnerItem::Shared { interest }, + hot_touch: None, + } + } + } + }; + + match attempt { + KeyAttempt::Wait(op) => wait_external_get_key_not_revoking(op).await, + KeyAttempt::Ready { item, hot_touch } => { + if let Some((put_id, memory_info)) = hot_touch { + hot_touches.push((key.clone(), put_id, memory_info)); + } + items.push(item); + break; + } + } + } + } + + for (key, put_id, memory_info) in hot_touches { + inner.owner_hot_touch_or_promote(&key, put_id, &memory_info); + } + (items, planning_leaders.handoff()) +} + +fn clear_external_get_key_marker_locked( + controls: &mut std::collections::HashMap, + op: &Arc, +) { + let remove_control = if let Some(control) = controls.get_mut(&op.key) { + if control + .external_get + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, op)) + { + control.external_get = None; + } + control.is_idle() + } else { + false + }; + if remove_control { + controls.remove(&op.key); + } +} + +fn abandon_unstarted_external_get_key_locked( + controls: &mut std::collections::HashMap, + op: &Arc, +) -> bool { + let mut state = op.state.lock(); + if !matches!(state.phase, ExternalGetKeySharedPhase::Starting) { + return false; + } + state.phase = ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Miss, + }; + state.terminal_at = Some(Instant::now()); + clear_external_get_key_marker_locked(controls, op); + true +} + +fn publish_external_get_key_terminal( + inner: &client_kv_api::ClientKvApiInner, + op: &Arc, + phase: ExternalGetKeySharedPhase, +) { + { + let mut controls = inner.owner_key_control.lock_key(&op.key); + let mut state = op.state.lock(); + state.phase = phase; + state.terminal_at = Some(Instant::now()); + clear_external_get_key_marker_locked(&mut controls, op); + } + inner.untrack_external_get_flight(op); + op.notify.notify_waiters(); +} + +fn publish_external_get_key_failed( + inner: &client_kv_api::ClientKvApiInner, + op: &Arc, + err: &KvError, +) { + let (error_code, error_json) = external_get_start_error_parts(err); + publish_external_get_key_terminal( + inner, + op, + ExternalGetKeySharedPhase::Failed { + error_code, + error_json, + }, + ); +} + +fn publish_external_get_key_started( + op: &Arc, + item: crate::master_kv_router::msg_pack::BatchGetStartItemResp, +) { + { + let mut state = op.state.lock(); + assert!(matches!(state.phase, ExternalGetKeySharedPhase::Starting)); + state.phase = ExternalGetKeySharedPhase::Started { item }; + } + op.notify.notify_waiters(); +} + +fn external_get_key_ready_from_code( + error_code: u32, + error_json: String, +) -> ExternalGetStartSharedItemResult { + if error_code == codes_api::API_KEY_NOT_FOUND { + ExternalGetStartSharedItemResult::Miss + } else { + ExternalGetStartSharedItemResult::Error { + error_code, + error_json, + } + } +} + +async fn wait_external_get_key_start_code( + op: Arc, +) -> KvResult<(u32, String)> { + loop { + let notified = op.notify.notified(); + futures::pin_mut!(notified); + let should_wait = { + let state = op.state.lock(); + match &state.phase { + ExternalGetKeySharedPhase::Starting => { + notified.as_mut().enable(); + true + } + ExternalGetKeySharedPhase::Started { item } + | ExternalGetKeySharedPhase::Finishing { item } => { + return Ok((item.error_code, item.error_json.clone())); + } + ExternalGetKeySharedPhase::Ready { result } => { + return Ok(match result { + ExternalGetStartSharedItemResult::Hit { .. } => (OK, String::new()), + ExternalGetStartSharedItemResult::Miss => { + (codes_api::API_KEY_NOT_FOUND, String::new()) + } + ExternalGetStartSharedItemResult::Error { + error_code, + error_json, + } => (*error_code, error_json.clone()), + }); + } + ExternalGetKeySharedPhase::Failed { + error_code, + error_json, + } => return Err(KvError::from_json(*error_code, error_json)), + ExternalGetKeySharedPhase::Revoking { .. } => { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external Get key entered revoke before prefix decision: key={}", + op.key + ), + })); + } + } + }; + if should_wait { + notified.await; + } + } +} + +async fn wait_external_get_key_result( + op: Arc, +) -> KvResult { + loop { + let notified = op.notify.notified(); + futures::pin_mut!(notified); + let should_wait = { + let state = op.state.lock(); + match &state.phase { + ExternalGetKeySharedPhase::Starting + | ExternalGetKeySharedPhase::Started { .. } + | ExternalGetKeySharedPhase::Finishing { .. } => { + notified.as_mut().enable(); + true + } + ExternalGetKeySharedPhase::Ready { result } => return Ok(result.clone()), + ExternalGetKeySharedPhase::Failed { + error_code, + error_json, + } => return Err(KvError::from_json(*error_code, error_json)), + ExternalGetKeySharedPhase::Revoking { .. } => { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "retained external Get key was unexpectedly revoked: key={}", + op.key + ), + })); + } + } + }; + if should_wait { + notified.await; + } + } +} + +fn decide_external_get_key_item(item: &mut ExternalGetStartOwnerItem, retain: bool) { + if let ExternalGetStartOwnerItem::Shared { interest } = item { + interest.decide(retain); + } +} + +enum ExternalGetKeyLeaderAction { + Finish { + op: Arc, + item: crate::master_kv_router::msg_pack::BatchGetStartItemResp, + }, + Revoke { + op: Arc, + item: crate::master_kv_router::msg_pack::BatchGetStartItemResp, + }, + Terminal, +} + +type ExternalGetFinishLeader = (Arc, BatchGetStartItemResp); + +struct ExternalGetFinishUnit { + group: Option, + leaders: Vec, +} + +/// Partition leaders into independent Done failure domains while preserving +/// caller-declared atomic groups. The target is deliberately soft: one large +/// atomic group stays intact instead of being split across terminal RPCs. +fn partition_external_get_finish_leaders( + leaders: Vec, + target_keys_per_batch: usize, +) -> Vec> { + if leaders.is_empty() { + return Vec::new(); + } + let target_keys_per_batch = target_keys_per_batch.max(1); + let mut units: Vec = Vec::new(); + let mut group_unit_by_anchor = HashMap::<(String, u64, u32), usize>::new(); + + for leader in leaders { + let group = leader.1.atomic_group.as_ref(); + let Some(group) = group else { + units.push(ExternalGetFinishUnit { + group: None, + leaders: vec![leader], + }); + continue; + }; + let Some(anchor) = group.members.first() else { + // Master validation rejects empty groups. Treat malformed legacy + // metadata as a singleton here so partitioning cannot lose work. + units.push(ExternalGetFinishUnit { + group: None, + leaders: vec![leader], + }); + continue; + }; + let signature = (anchor.key.clone(), anchor.put_id.0, anchor.put_id.1); + if let Some(unit_idx) = group_unit_by_anchor.get(&signature).copied() + && units[unit_idx].group.as_ref() == Some(group) + { + units[unit_idx].leaders.push(leader); + continue; + } + let unit_idx = units.len(); + units.push(ExternalGetFinishUnit { + group: Some(group.clone()), + leaders: vec![leader], + }); + group_unit_by_anchor.insert(signature, unit_idx); + } + + let mut batches = Vec::new(); + let mut current = Vec::new(); + for mut unit in units { + if !current.is_empty() + && current.len().saturating_add(unit.leaders.len()) > target_keys_per_batch + { + batches.push(std::mem::take(&mut current)); + } + current.append(&mut unit.leaders); + if current.len() >= target_keys_per_batch { + batches.push(std::mem::take(&mut current)); + } + } + if !current.is_empty() { + batches.push(current); + } + batches +} + +async fn classify_external_get_key_leader( + inner: &client_kv_api::ClientKvApiInner, + op: Arc, +) -> ExternalGetKeyLeaderAction { + loop { + let notified = op.notify.notified(); + futures::pin_mut!(notified); + let wait_for_decisions = { + let state = op.state.lock(); + if matches!( + state.phase, + ExternalGetKeySharedPhase::Starting | ExternalGetKeySharedPhase::Started { .. } + ) && state.undecided != 0 + { + notified.as_mut().enable(); + true + } else { + false + } + }; + if wait_for_decisions { + notified.await; + continue; + } + + let mut notify_terminal = false; + let action = { + // Planner lock order is owner fence -> per-key op. Use the same + // order here so no request can join after undecided reaches zero + // but before the leader closes admission. + let mut controls = inner.owner_key_control.lock_key(&op.key); + let mut state = op.state.lock(); + match &state.phase { + ExternalGetKeySharedPhase::Started { item } if state.undecided == 0 => { + let item = item.clone(); + if item.error_code != OK { + state.phase = ExternalGetKeySharedPhase::Ready { + result: external_get_key_ready_from_code( + item.error_code, + item.error_json, + ), + }; + state.terminal_at = Some(Instant::now()); + clear_external_get_key_marker_locked(&mut controls, &op); + notify_terminal = true; + Some(ExternalGetKeyLeaderAction::Terminal) + } else if state.retained == 0 { + state.phase = ExternalGetKeySharedPhase::Revoking { item: item.clone() }; + Some(ExternalGetKeyLeaderAction::Revoke { + op: op.clone(), + item, + }) + } else { + state.phase = ExternalGetKeySharedPhase::Finishing { item: item.clone() }; + Some(ExternalGetKeyLeaderAction::Finish { + op: op.clone(), + item, + }) + } + } + ExternalGetKeySharedPhase::Starting | ExternalGetKeySharedPhase::Started { .. } => { + None + } + ExternalGetKeySharedPhase::Finishing { .. } + | ExternalGetKeySharedPhase::Revoking { .. } + | ExternalGetKeySharedPhase::Ready { .. } + | ExternalGetKeySharedPhase::Failed { .. } => { + Some(ExternalGetKeyLeaderAction::Terminal) + } + } + }; + if notify_terminal { + inner.untrack_external_get_flight(&op); + op.notify.notify_waiters(); + } + if let Some(action) = action { + return action; + } + tokio::task::yield_now().await; + } } -const SIDE_TRANSFER_OWNER_RPC_TIMEOUT_SECS: u64 = 30; -const SIDE_TRANSFER_TARGET_RESOLVE_TIMEOUT_SECS: u64 = 10; +async fn finish_external_get_key_leaders( + view: crate::client_kv_api::ClientKvApiView, + leaders: Vec>, + transfer_concurrency: usize, +) { + let client_api = view.client_kv_api(); + let inner = client_api.inner(); + let mut finish = Vec::new(); + let mut revoke = Vec::new(); + for op in leaders { + match classify_external_get_key_leader(inner, op).await { + ExternalGetKeyLeaderAction::Finish { op, item } => finish.push((op, item)), + ExternalGetKeyLeaderAction::Revoke { op, item } => revoke.push((op, item)), + ExternalGetKeyLeaderAction::Terminal => {} + } + } + + if !revoke.is_empty() { + let mut attempt = 1u32; + loop { + let get_ids = revoke + .iter() + .map(|(_, item)| item.get_id) + .collect::>(); + let response = inner.batch_get_revoke(get_ids).await; + let resp = match response { + Ok(resp) + if resp.items.len() == revoke.len() + && resp + .items + .iter() + .zip(&revoke) + .all(|(resp, (_, expected))| resp.get_id == expected.get_id) => + { + resp + } + Ok(resp) => { + tracing::warn!( + "BatchGetRevoke response length mismatch; retaining prepared slots and retrying: expected={} got={} attempt={}", + revoke.len(), + resp.items.len(), + attempt + ); + tokio::time::sleep(Duration::from_millis( + (50u64.saturating_mul(1u64 << attempt.min(6))).min(2_000), + )) + .await; + attempt = attempt.saturating_add(1); + continue; + } + Err(err) => { + if matches!(&err, KvError::Api(ApiError::SystemShutdown { .. })) { + tracing::warn!( + "BatchGetRevoke cleanup stopped during owner shutdown: items={}", + revoke.len() + ); + return; + } + tracing::warn!( + "BatchGetRevoke uncertain; retaining operation identity and prepared slots for retry: items={} attempt={} err={}", + revoke.len(), + attempt, + err + ); + tokio::time::sleep(Duration::from_millis( + (50u64.saturating_mul(1u64 << attempt.min(6))).min(2_000), + )) + .await; + attempt = attempt.saturating_add(1); + continue; + } + }; + + for ((op, item), revoke_item) in revoke.drain(..).zip(resp.items) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + revoke_item.error_code, + revoke_item.error_json, + ) { + // Done won the per-get terminal lock (or requester identity + // is invalid). The candidate may be committed, so never + // return its slot from the losing Revoke path. + publish_external_get_key_failed(inner, &op, &err); + continue; + } + if let Some(target) = item.prepared_target.as_ref() { + let mut release_attempt = 1u32; + loop { + match release_prepared_get_target(inner, target).await { + Ok(()) => break, + Err(err) => { + if !inner.view.register_shutdown_poller().is_running() { + return; + } + tracing::error!( + "Revoke confirmed but prepared slot release failed; keeping key fenced for retry: key={} get_id={} attempt={} err={}", + op.key, + item.get_id, + release_attempt, + err + ); + tokio::time::sleep(Duration::from_millis(1000)).await; + release_attempt = release_attempt.saturating_add(1); + } + } + } + } + publish_external_get_key_terminal( + inner, + &op, + ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Miss, + }, + ); + } + break; + } + } + + if finish.is_empty() { + return; + } + let finish_batches = + partition_external_get_finish_leaders(finish, EXTERNAL_GET_FINISH_TARGET_KEYS_PER_BATCH); + let batch_concurrency = EXTERNAL_GET_FINISH_BATCH_CONCURRENCY + .min(finish_batches.len()) + .max(1); + let per_batch_transfer_concurrency = transfer_concurrency + .max(1) + .div_ceil(batch_concurrency) + .max(1); + tracing::debug!( + "external Get singleflight finish partitioned: batches={} batch_concurrency={} transfer_concurrency_per_batch={}", + finish_batches.len(), + batch_concurrency, + per_batch_transfer_concurrency + ); + + let finish_futures = finish_batches.into_iter().map(|batch| { + let batch_view = view.clone(); + async move { + let keys = batch + .iter() + .map(|(op, _)| op.key.clone()) + .collect::>(); + let start_items = batch + .iter() + .map(|(_, item)| item.clone()) + .collect::>(); + let result = batch_view + .client_kv_api() + .inner() + .batch_get_finish_started(keys, start_items, per_batch_transfer_concurrency) + .await; + (batch, result) + } + }); + let mut finish_stream = + futures::stream::iter(finish_futures).buffer_unordered(batch_concurrency); + while let Some((batch, finish_result)) = finish_stream.next().await { + match finish_result { + Ok(results) if results.len() == batch.len() => { + // Publish each confirmed sub-batch immediately. A different + // uncertain Done atomic_batch therefore cannot retain these flights + // or their pending-visible slots. + for ((op, _), result) in batch.into_iter().zip(results) { + let phase = match result { + Ok(Some((memholder, _))) => ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Hit { memholder }, + }, + Ok(None) => ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Miss, + }, + Err(err) => { + let (error_code, error_json) = external_get_start_error_parts(&err); + ExternalGetKeySharedPhase::Ready { + result: ExternalGetStartSharedItemResult::Error { + error_code, + error_json, + }, + } + } + }; + publish_external_get_key_terminal(inner, &op, phase); + } + } + Ok(results) => { + let err = KvError::Api(ApiError::Unknown { + detail: format!( + "singleflight leader finish length mismatch: expected={} got={}", + batch.len(), + results.len() + ), + }); + for (op, _) in batch { + publish_external_get_key_failed(inner, &op, &err); + } + } + Err(err) => { + for (op, _) in batch { + publish_external_get_key_failed(inner, &op, &err); + } + } + } + } +} + +async fn prepare_external_get_batch_plan( + inner: &client_kv_api::ClientKvApiInner, + req: &ExternalBatchGetStartReq, + group_lens: &[usize], +) -> KvResult<( + ExternalGetStartPrefixResult, + Vec, + Vec, +)> { + let (mut items, leaders) = plan_external_get_key_items(inner, &req.keys).await; + if !leaders.is_empty() { + let leader_keys = leaders.iter().map(|op| op.key.clone()).collect::>(); + let leader_count = leaders.len(); + let spawn_view = inner.view.clone_view(); + let worker_view = spawn_view.clone(); + let transfer_concurrency = req.transfer_concurrency; + spawn_view.spawn("external_get_key_singleflight", async move { + let batch_start_started_at = Instant::now(); + let start_result = { + let worker_inner = worker_view.client_kv_api().inner(); + batch_get_start_with_local_reserve_targets(worker_inner, &leader_keys).await + }; + let batch_start_us = duration_to_i64_us(batch_start_started_at.elapsed()); + match start_result { + Ok(start_resp) => { + assert_eq!( + start_resp.items.len(), + leaders.len(), + "validated BatchGetStart response must match leader count" + ); + let start_hits = start_resp + .items + .iter() + .filter(|item| item.error_code == OK) + .count(); + let start_misses = start_resp + .items + .iter() + .filter(|item| item.error_code == codes_api::API_KEY_NOT_FOUND) + .count(); + let start_errors = start_resp + .items + .len() + .saturating_sub(start_hits) + .saturating_sub(start_misses); + tracing::info!( + "external Get leader start lifecycle: leaders={} hits={} misses={} errors={} batch_get_start_us={} outcome=ok", + leader_count, + start_hits, + start_misses, + start_errors, + batch_start_us, + ); + for (op, item) in leaders.iter().zip(start_resp.items) { + publish_external_get_key_started(op, item); + } + finish_external_get_key_leaders(worker_view, leaders, transfer_concurrency) + .await; + } + Err(err) => { + tracing::info!( + "external Get leader start lifecycle: leaders={} hits=0 misses=0 errors={} batch_get_start_us={} outcome=error", + leader_count, + leader_count, + batch_start_us, + ); + let worker_inner = worker_view.client_kv_api().inner(); + for op in &leaders { + publish_external_get_key_failed(worker_inner, op, &err); + } + } + } + }); + } + + let mut item_codes = Vec::with_capacity(items.len()); + for item in &items { + match item { + ExternalGetStartOwnerItem::Local { .. } => item_codes.push((OK, String::new())), + ExternalGetStartOwnerItem::Shared { interest } => { + match wait_external_get_key_start_code(interest.op().clone()).await { + Ok(code) => item_codes.push(code), + Err(err) => return Err(err), + } + } + } + } + + let (raw_prefix_hit_len, first_miss_index, first_error_kind) = + compute_external_get_start_raw_prefix(&item_codes); + let transferable_len = compute_external_get_start_transfer_prefix( + raw_prefix_hit_len, + group_lens, + req.prefix_best_effort, + ); + if let Some(error_kind) = first_error_kind.as_ref() { + tracing::warn!( + "external_batch_get_start prefix stopped on non-key-miss item: req_node_id={} raw_prefix_hit_len={} transferable_len={} first_miss_index={:?} error_kind={}", + req.req_node_id, + raw_prefix_hit_len, + transferable_len, + first_miss_index, + error_kind + ); + } + for (idx, item) in items.iter_mut().enumerate() { + decide_external_get_key_item(item, idx < transferable_len); + } + + let transfer_keys = req.keys[..transferable_len].to_vec(); + let transfer_items = items.into_iter().take(transferable_len).collect::>(); + Ok(( + ExternalGetStartPrefixResult { + raw_prefix_hit_len, + transferable_len, + first_miss_index, + first_error_kind, + }, + transfer_keys, + transfer_items, + )) +} impl ClientKvApi { fn is_side_transfer_worker(&self) -> bool { @@ -253,6 +2655,12 @@ impl ClientKvApi { let commit_req = MsgPack { serialize_part: ExternalPutCommitReq { key: req.key.clone(), + len: req.len, + src_offset: req.src_offset, + remote_target: req + .peer_id + .as_deref() + .is_some_and(|peer| peer != owner_id.as_str()), put_id: Some(put_id), lease_id: req.lease_id, started_time: req.started_time, @@ -298,22 +2706,63 @@ pub trait HandlerForExternalClient { /// Validate external's observed owner start_time (0 means skip validation for legacy callers) fn validate_requester_owner_status_updated(&self, started_time: i64) -> KvResult<()>; async fn external_get(&self, req: ExternalGetReq) -> KvResult; + async fn external_batch_get(&self, req: ExternalBatchGetReq) -> KvResult; + async fn external_batch_get_local_probe( + &self, + req: ExternalBatchGetLocalProbeReq, + ) -> KvResult; + async fn external_batch_get_start( + &self, + req: ExternalBatchGetStartReq, + ) -> KvResult; + async fn external_batch_get_transfer( + &self, + req: ExternalBatchGetTransferReq, + ) -> KvResult; + async fn external_batch_get_cancel( + &self, + req: ExternalBatchGetCancelReq, + ) -> KvResult; + async fn external_execute_planned_get( + &self, + req: ExternalExecutePlannedGetReq, + ) -> KvResult; async fn external_put_start(&self, req: ExternalPutStartReq) -> KvResult; + async fn external_batch_put_start( + &self, + req: ExternalBatchPutStartReq, + ) -> KvResult; // deprecated: transfer merged into external_put_transfer_end async fn external_put_transfer_end( &self, req: ExternalPutTransferEndReq, ) -> KvResult; + async fn external_batch_put_transfer_end( + &self, + req: ExternalBatchPutTransferEndReq, + ) -> KvResult; async fn external_put_commit( &self, req: ExternalPutCommitReq, ) -> KvResult; + async fn external_batch_put_commit( + &self, + req: ExternalBatchPutCommitReq, + ) -> KvResult; async fn external_put_revoke( &self, req: ExternalPutRevokeReq, ) -> KvResult; async fn external_delete(&self, req: ExternalDeleteReq) -> KvResult; async fn external_is_exist(&self, req: ExternalIsExistReq) -> KvResult; + async fn external_batch_is_exist( + &self, + req: ExternalBatchIsExistReq, + ) -> KvResult; + async fn external_observability_snapshot( + &self, + req: ExternalObservabilitySnapshotReq, + ) -> KvResult; } /// Handle external get request @@ -351,27 +2800,703 @@ impl HandlerForExternalClient for ClientKvApi { }); } }; - let memory_info = memholder.memory_info(); - // Build holding info to record that this external client holds the memholder - let client_holding_info = ExternalHoldingGetInfo { - key: req.key.clone(), - req_node_id: req.req_node_id.clone(), - memory_info, + let external_memholder_info = + inner.install_external_get_holding(&req.req_node_id, memholder.memory_info()); + Ok(ExternalGetResp { + external_memholder_info: Some(external_memholder_info), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }) + } + + async fn external_batch_get(&self, req: ExternalBatchGetReq) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_get")); + } + let inner = self.inner(); + + self.validate_requester_owner_status_updated(req.started_time)?; + + let batch_results = inner.batch_get(req.keys, req.transfer_concurrency).await?; + let mut items = Vec::with_capacity(batch_results.len()); + for item_result in batch_results { + match item_result { + Ok(Some((memholder, _))) => { + let external_memholder_info = inner.install_external_get_holding( + &req.req_node_id, + memholder.memory_info(), + ); + items.push(ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some(external_memholder_info), + }); + } + Ok(None) => items.push(ExternalBatchGetItemResp { + error_code: + crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND, + error_json: "Key not found".to_string(), + external_memholder_info: None, + }), + Err(err) => items.push(ExternalBatchGetItemResp { + external_memholder_info: None, + ..ExternalBatchGetItemResp::from_error(&err) + }), + } + } + + Ok(ExternalBatchGetResp { + items, + error_code: OK, + error_json: String::new(), + }) + } + + async fn external_batch_get_local_probe( + &self, + req: ExternalBatchGetLocalProbeReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported( + "external_batch_get_local_probe", + )); + } + if req.keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_get_local_probe requires at least one key".to_string(), + })); + } + if req.plan_handle == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_get_local_probe requires a non-zero plan_handle" + .to_string(), + })); + } + self.validate_requester_owner_status_updated(req.started_time)?; + let inner = self.inner(); + let identity = (req.req_node_id.clone(), req.started_time, req.plan_handle); + let operation_lock = inner + .external_get_local_probe_locks + .get_lock(identity.clone()); + let _operation_guard = operation_lock.lock().await; + if let Some((completed_keys, completed)) = inner + .completed_external_get_local_probes + .get(&identity) + .await + { + if completed_keys != req.keys { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_get_local_probe operation identity was reused with different keys: plan_handle={}", + req.plan_handle + ), + })); + } + return Ok(completed); + } + let mut items = Vec::with_capacity(req.keys.len()); + let mut hot_touches = Vec::new(); + for key in &req.keys { + let (item, hot_touch) = { + let controls = inner.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + (None, None) + } else if let Some(memory_info) = inner.local_visible_mem_holder_unfenced(key) { + let hot_touch = inner.get_cached_info.get(key).and_then(|cached| { + Arc::ptr_eq(&cached.mem_holder, &memory_info).then_some(( + (cached.put_time_ms, cached.put_version), + memory_info.clone(), + )) + }); + // Install the holder under the same per-key fence used for + // local visibility. Reclaim cannot slip between the probe + // decision and the pin that owns this exact backing. + let info = inner.install_external_get_holding(&req.req_node_id, memory_info); + (Some(info), hot_touch) + } else { + (None, None) + } + }; + if let Some((put_id, memory_info)) = hot_touch { + hot_touches.push((key.clone(), put_id, memory_info)); + } + items.push(item); + } + for (key, put_id, memory_info) in hot_touches { + inner.owner_hot_touch_or_promote(&key, put_id, &memory_info); + } + let local_items = items.iter().filter(|item| item.is_some()).count(); + let remote_items = items.len().saturating_sub(local_items); + inner + .planned_get_counters + .local_probe_batches + .fetch_add(1, Ordering::Relaxed); + inner + .planned_get_counters + .local_probe_items + .fetch_add(items.len() as u64, Ordering::Relaxed); + inner + .planned_get_counters + .local_probe_local_items + .fetch_add(local_items as u64, Ordering::Relaxed); + inner + .planned_get_counters + .local_probe_remote_items + .fetch_add(remote_items as u64, Ordering::Relaxed); + tracing::debug!( + requested = req.keys.len(), + local = local_items, + remote = remote_items, + "external Get owner-local probe" + ); + let response = ExternalBatchGetLocalProbeResp { + items, + error_code: OK, + error_json: String::new(), + }; + inner + .completed_external_get_local_probes + .insert(identity, (req.keys, response.clone())) + .await; + Ok(response) + } + + async fn external_batch_get_start( + &self, + req: ExternalBatchGetStartReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_get_start")); + } + if req.keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_get_start requires at least one key".to_string(), + })); + } + let inner = self.inner(); + self.validate_requester_owner_status_updated(req.started_time)?; + let observe_started_at = Instant::now(); + let requested_len = req.keys.len(); + + let group_lens = + normalize_external_get_start_group_lens(req.keys.len(), req.atomic_group_lens.clone())?; + let (prefix, transfer_keys, transfer_items) = + prepare_external_get_batch_plan(inner, &req, &group_lens).await?; + let phase_snapshot = observe_external_get_consume_phases(&transfer_items, Instant::now()); + let handle = inner + .next_external_get_start_handle + .fetch_add(1, Ordering::Relaxed); + let inline_memory_infos = (transfer_items.len() == req.keys.len()) + .then(|| collect_all_local_external_get_start_infos(&transfer_items)) + .flatten(); + let transfer_plan = if let Some(memory_infos) = inline_memory_infos { + assert_eq!( + memory_infos.len(), + req.keys.len(), + "inline local get_start plan must cover the full request" + ); + let items = memory_infos + .into_iter() + .map(|memory_info| ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some( + inner.install_external_get_holding(&req.req_node_id, memory_info), + ), + }) + .collect(); + ExternalBatchGetStartTransferPlan::InlineLocal { items } + } else { + inner.external_get_start_registry.insert( + handle, + ExternalGetStartEntry { + req_node_id: req.req_node_id.clone(), + requester_node_start_time: inner + .view + .cluster_manager() + .get_member_info_cached(&req.req_node_id) + .map(|member| member.node_start_time), + keys: transfer_keys, + items: transfer_items, + atomic_group_lens: group_lens, + created_at: Instant::now(), + }, + ); + ExternalBatchGetStartTransferPlan::OwnerRpc }; + let inline_local = matches!( + &transfer_plan, + ExternalBatchGetStartTransferPlan::InlineLocal { .. } + ); + tracing::info!( + "external Get start lifecycle: handle={} requested={} raw_prefix={} transferable={} inline_local={} local={} starting={} started={} finishing={} revoking={} ready={} failed={} terminal_before_return={} pending_before_return={} terminal_age_mean_us={} terminal_age_max_us={} total_us={}", + handle, + requested_len, + prefix.raw_prefix_hit_len, + prefix.transferable_len, + inline_local, + phase_snapshot.local, + phase_snapshot.starting, + phase_snapshot.started, + phase_snapshot.finishing, + phase_snapshot.revoking, + phase_snapshot.ready, + phase_snapshot.failed, + phase_snapshot.terminal_before_consume(), + phase_snapshot.pending_before_consume(), + phase_snapshot.terminal_age_mean_us(), + phase_snapshot.terminal_age_max_us, + duration_to_i64_us(observe_started_at.elapsed()), + ); + + Ok(ExternalBatchGetStartResp { + error_code: OK, + error_json: String::new(), + handle, + raw_prefix_hit_len: prefix.raw_prefix_hit_len, + transfer_plan, + }) + } + + async fn external_batch_get_transfer( + &self, + req: ExternalBatchGetTransferReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported( + "external_batch_get_transfer", + )); + } + let inner = self.inner(); + self.validate_requester_owner_status_updated(req.started_time)?; + let consume_started_at = Instant::now(); - // Insert/update holding via owned manager - inner.external_get_holding.insert( - NodeHolderKey::new(req.req_node_id.clone(), memholder.holder_id()), - client_holding_info, + let Some((_handle, mut entry)) = inner.external_get_start_registry.remove(&req.handle) + else { + return Err(KvError::Api(ApiError::KeyNotFound { + key: format!("external_get_start_handle:{}", req.handle), + })); + }; + if entry.req_node_id != req.req_node_id { + tracing::warn!( + "external_batch_get_transfer req_node_id mismatch: handle={} expected={} got={}", + req.handle, + entry.req_node_id, + req.req_node_id + ); + } + if let Err(err) = validate_external_get_consume_prefix( + req.consume_prefix_len, + entry.keys.len(), + &entry.atomic_group_lens, + ) { + inner.external_get_start_registry.insert(req.handle, entry); + return Err(err); + } + assert_eq!( + entry.keys.len(), + entry.items.len(), + "external get-start registry keys and items must stay aligned" ); - let external_memholder_info = ExternalMemHolderInfo { - offset: memholder.get_offset(), - len: memholder.get_length() as u32, - holder_id: memholder.holder_id(), + let available_prefix_len = entry.keys.len(); + let handle_age_us = duration_to_i64_us(entry.created_at.elapsed()); + let phase_snapshot = observe_external_get_consume_phases( + &entry.items[..req.consume_prefix_len], + Instant::now(), + ); + let tail_items = entry.items.split_off(req.consume_prefix_len); + entry.keys.truncate(req.consume_prefix_len); + if !tail_items.is_empty() { + tracing::info!( + "external_batch_get_transfer consumed prefix and released owner tail: handle={} available={} consumed={} released_tail={}", + req.handle, + available_prefix_len, + req.consume_prefix_len, + tail_items.len() + ); + } + drop(tail_items); + let keys = entry.keys; + let finish_wait_started_at = Instant::now(); + let transfer_results_result = + finish_external_get_start_transfer(inner.view.clone_view(), entry.items, 0).await; + let finish_wait_us = duration_to_i64_us(finish_wait_started_at.elapsed()); + let transfer_results = match transfer_results_result { + Ok(results) => results, + Err(err) => { + tracing::info!( + "external Get consume lifecycle: handle={} available={} consumed={} released_tail={} handle_age_us={} local_before={} starting_before={} started_before={} finishing_before={} revoking_before={} ready_before={} failed_before={} terminal_before={} pending_before={} terminal_age_mean_us={} terminal_age_max_us={} finish_wait_us={} hits=0 misses=0 errors={} install_us=0 total_us={} outcome=error", + req.handle, + available_prefix_len, + req.consume_prefix_len, + available_prefix_len.saturating_sub(req.consume_prefix_len), + handle_age_us, + phase_snapshot.local, + phase_snapshot.starting, + phase_snapshot.started, + phase_snapshot.finishing, + phase_snapshot.revoking, + phase_snapshot.ready, + phase_snapshot.failed, + phase_snapshot.terminal_before_consume(), + phase_snapshot.pending_before_consume(), + phase_snapshot.terminal_age_mean_us(), + phase_snapshot.terminal_age_max_us, + finish_wait_us, + req.consume_prefix_len, + duration_to_i64_us(consume_started_at.elapsed()), + ); + return Err(err); + } }; - Ok(ExternalGetResp { - external_memholder_info: Some(external_memholder_info), - error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + if transfer_results.len() != keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_get_transfer result length mismatch: expected={} got={}", + keys.len(), + transfer_results.len() + ), + })); + } + let install_started_at = Instant::now(); + let mut items = Vec::with_capacity(transfer_results.len()); + let mut hits = 0usize; + let mut misses = 0usize; + let mut errors = 0usize; + for (key, item_result) in keys.iter().zip(transfer_results) { + match item_result { + Ok(Some((memholder, _))) => { + hits = hits.saturating_add(1); + let external_memholder_info = inner + .install_external_get_holding(&req.req_node_id, memholder.memory_info()); + items.push(ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some(external_memholder_info), + }); + } + Ok(None) => { + misses = misses.saturating_add(1); + items.push(ExternalBatchGetItemResp { + error_code: codes_api::API_KEY_NOT_FOUND, + error_json: format!("Key not found: {}", key), + external_memholder_info: None, + }); + } + Err(err) => { + errors = errors.saturating_add(1); + items.push(ExternalBatchGetItemResp { + external_memholder_info: None, + ..ExternalBatchGetItemResp::from_error(&err) + }); + } + } + } + let install_us = duration_to_i64_us(install_started_at.elapsed()); + tracing::info!( + "external Get consume lifecycle: handle={} available={} consumed={} released_tail={} handle_age_us={} local_before={} starting_before={} started_before={} finishing_before={} revoking_before={} ready_before={} failed_before={} terminal_before={} pending_before={} terminal_age_mean_us={} terminal_age_max_us={} finish_wait_us={} hits={} misses={} errors={} install_us={} total_us={} outcome=ok", + req.handle, + available_prefix_len, + req.consume_prefix_len, + available_prefix_len.saturating_sub(req.consume_prefix_len), + handle_age_us, + phase_snapshot.local, + phase_snapshot.starting, + phase_snapshot.started, + phase_snapshot.finishing, + phase_snapshot.revoking, + phase_snapshot.ready, + phase_snapshot.failed, + phase_snapshot.terminal_before_consume(), + phase_snapshot.pending_before_consume(), + phase_snapshot.terminal_age_mean_us(), + phase_snapshot.terminal_age_max_us, + finish_wait_us, + hits, + misses, + errors, + install_us, + duration_to_i64_us(consume_started_at.elapsed()), + ); + Ok(ExternalBatchGetTransferResp { + items, + error_code: OK, + error_json: String::new(), + }) + } + + async fn external_execute_planned_get( + &self, + req: ExternalExecutePlannedGetReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported( + "external_execute_planned_get", + )); + } + if req.items.is_empty() || req.transfer_concurrency == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_execute_planned_get requires non-empty items and positive concurrency: items={} concurrency={}", + req.items.len(), + req.transfer_concurrency + ), + })); + } + let inner = self.inner(); + self.validate_requester_owner_status_updated(req.started_time)?; + let identity = (req.req_node_id.clone(), req.started_time, req.plan_handle); + let operation_lock = inner + .planned_external_get_execute_locks + .get_lock(identity.clone()); + let _operation_guard = operation_lock.lock().await; + if let Some(completed) = inner + .completed_planned_external_get_executes + .get(&identity) + .await + { + return Ok(completed); + } + + let mut requested_by_key = HashMap::with_capacity(req.items.len()); + for item in &req.items { + if item.key.is_empty() + || requested_by_key + .insert(item.key.clone(), item.get_id) + .is_some() + { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_execute_planned_get requires unique non-empty keys: plan_handle={}", + req.plan_handle + ), + })); + } + } + let keys = req + .items + .iter() + .map(|item| item.key.clone()) + .collect::>(); + let (mut owner_items, leaders) = plan_external_get_key_items(inner, &keys).await; + let local_items = owner_items + .iter() + .filter(|item| matches!(item, ExternalGetStartOwnerItem::Local { .. })) + .count(); + let leader_items = leaders.len(); + let follower_items = owner_items + .len() + .saturating_sub(local_items) + .saturating_sub(leader_items); + inner + .planned_get_counters + .batches + .fetch_add(1, Ordering::Relaxed); + inner + .planned_get_counters + .local_items + .fetch_add(local_items as u64, Ordering::Relaxed); + inner + .planned_get_counters + .leader_items + .fetch_add(leader_items as u64, Ordering::Relaxed); + inner + .planned_get_counters + .follower_items + .fetch_add(follower_items as u64, Ordering::Relaxed); + tracing::info!( + plan_handle = req.plan_handle, + items = owner_items.len(), + local_items, + leader_items, + follower_items, + "planned CPU Get owner classification" + ); + let leader_get_ids = leaders + .iter() + .map(|leader| { + requested_by_key + .get(&leader.key) + .copied() + .expect("planned leader key must come from the request") + }) + .collect::>(); + let leader_id_set = leader_get_ids + .iter() + .copied() + .collect::>(); + let mut cleanup = req + .items + .iter() + .filter(|item| !leader_id_set.contains(&item.get_id)) + .map(|item| StartedGetRevokeCleanup { + get_id: item.get_id, + prepared_target: None, + }) + .collect::>(); + + let bound_items = + match batch_get_bind_with_local_reserve_targets(inner, &leader_get_ids).await { + Ok(items) => items, + Err(err) => { + for leader in &leaders { + publish_external_get_key_failed(inner, leader, &err); + } + cleanup.extend(leader_get_ids.into_iter().map(|get_id| { + StartedGetRevokeCleanup { + get_id, + prepared_target: None, + } + })); + finish_started_get_revoke_cleanup( + inner, + cleanup, + "planned external Get bind failure", + ) + .await; + let response = ExternalExecutePlannedGetResp { + items: Vec::new(), + error_code: err.code(), + error_json: err.to_json(), + }; + inner + .completed_planned_external_get_executes + .insert(identity, response.clone()) + .await; + return Ok(response); + } + }; + assert_eq!(leaders.len(), bound_items.len()); + for (leader, item) in leaders.iter().zip(bound_items) { + if item.error_code != OK { + cleanup.push(StartedGetRevokeCleanup { + get_id: item.get_id, + prepared_target: None, + }); + } + publish_external_get_key_started(leader, item); + } + for item in &mut owner_items { + decide_external_get_key_item(item, true); + } + + // From this point on the owner has published Started for every leader, + // so it must retain lifecycle ownership independently of the inbound + // RPC future. P2P timeout/cancellation drops the handler future; if + // finish lived on that future, the per-key marker would remain in + // Finishing forever and every replay would only become a follower. + // The registered owner task is process-scoped and therefore continues + // both transfer/Done and unused-operation cleanup after caller loss. + let spawn_view = inner.view.clone_view(); + let worker_view = spawn_view.clone(); + let finish_leaders = leaders; + let transfer_concurrency = req.transfer_concurrency; + let plan_handle = req.plan_handle; + spawn_view.spawn( + format!("planned_external_get_finish_{plan_handle}"), + async move { + let finish_view = worker_view.clone(); + let cleanup_view = worker_view; + let finish = finish_external_get_key_leaders( + finish_view, + finish_leaders, + transfer_concurrency, + ); + let cleanup_future = async move { + finish_started_get_revoke_cleanup( + cleanup_view.client_kv_api().inner(), + cleanup, + "planned external Get unused operation", + ) + .await; + }; + tokio::join!(finish, cleanup_future); + }, + ); + + let transfer_results = + finish_external_get_start_transfer(inner.view.clone_view(), owner_items, 0).await?; + if transfer_results.len() != keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned external Get result length mismatch: expected={} got={}", + keys.len(), + transfer_results.len() + ), + })); + } + let mut items = Vec::with_capacity(keys.len()); + for (key, result) in keys.iter().zip(transfer_results) { + match result { + Ok(Some((holder, _))) => items.push(ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some( + inner.install_external_get_holding(&req.req_node_id, holder.memory_info()), + ), + }), + Ok(None) => items.push(ExternalBatchGetItemResp { + error_code: codes_api::API_KEY_NOT_FOUND, + error_json: format!("Key not found: {key}"), + external_memholder_info: None, + }), + Err(err) => items.push(ExternalBatchGetItemResp { + external_memholder_info: None, + ..ExternalBatchGetItemResp::from_error(&err) + }), + } + } + let response = ExternalExecutePlannedGetResp { + items, + error_code: OK, + error_json: String::new(), + }; + inner + .completed_planned_external_get_executes + .insert(identity, response.clone()) + .await; + Ok(response) + } + + async fn external_batch_get_cancel( + &self, + req: ExternalBatchGetCancelReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_get_cancel")); + } + let inner = self.inner(); + self.validate_requester_owner_status_updated(req.started_time)?; + + match req.transfer_plan { + ExternalBatchGetCancelPlan::OwnerRpc => { + if let Some((_handle, entry)) = + inner.external_get_start_registry.remove(&req.handle) + { + if entry.req_node_id != req.req_node_id { + tracing::warn!( + "external_batch_get_cancel req_node_id mismatch: handle={} expected={} got={}", + req.handle, + entry.req_node_id, + req.req_node_id + ); + } + drop(entry); + } + } + ExternalBatchGetCancelPlan::InlineLocal { holder_ids } => { + for holder_id in holder_ids { + let _ = inner + .external_get_holding + .remove(&NodeHolderKey::new(req.req_node_id.clone(), holder_id)); + } + } + } + Ok(ExternalBatchGetCancelResp { + error_code: OK, error_json: String::new(), }) } @@ -384,6 +3509,12 @@ impl HandlerForExternalClient for ClientKvApi { let started_at = Instant::now(); self.validate_requester_owner_status_updated(req.started_time)?; + // Register before the master PutStart RPC. Reclaim Prepare checks this + // counter under the same per-key key-control shard, closing the old gap + // between master admission and insertion into external_pending_puts. + let pending_fence = inner.acquire_external_pending_put_fence(&req.key)?; + let mut pending_fence_claim = + ExternalPutStartFenceClaim::new(inner.view.clone_view(), pending_fence); let put_start_started_at = Instant::now(); let source_node_id = if self.is_side_transfer_worker() { @@ -396,6 +3527,8 @@ impl HandlerForExternalClient for ClientKvApi { &req.key, req.len as u32, req.reject_if_inflight_same_key, + req.reject_if_exist_same_key, + req.make_replica_task, req.preferred_sub_cluster.as_deref(), source_node_id, ) @@ -405,10 +3538,15 @@ impl HandlerForExternalClient for ClientKvApi { e })?; // Ensure master responded OK before using returned addresses - crate::rpcresp_kvresult_convert::try_from_code( + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( put_start_resp.error_code, put_start_resp.error_json.clone(), - )?; + ) { + // An application response proves the master did not admit this Put; + // no uncertainty quarantine is required. + pending_fence_claim.release_after_definite_response(); + return Err(err); + } tracing::debug!( "handle external put start for key: {}, len: {}", req.key, @@ -428,21 +3566,38 @@ impl HandlerForExternalClient for ClientKvApi { // Remote target: external still writes into owner's shared memory (src_offset). src_offset }; - if !is_local_target { - // Stash remote context for transfer_end: peer_id, target_base_addr, target_offset(remote) - let remote_offset = put_start_resp.target_addr - put_start_resp.target_base_addr; - inner.external_pending_puts.insert( - ( - req.key.clone(), - put_start_resp.put_id.0, - put_start_resp.put_id.1, - ), - ExternalPendingPutCtx { - peer_id: put_start_resp.node_id.clone(), - target_base_addr: put_start_resp.target_base_addr, - target_offset: remote_offset, + let remote_offset = put_start_resp.target_addr - put_start_resp.target_base_addr; + let replica_admitted = put_start_resp.replica_target.is_some(); + inner.external_pending_puts.insert( + ( + req.key.clone(), + put_start_resp.put_id.0, + put_start_resp.put_id.1, + ), + ExternalPendingPutCtx { + peer_id: if is_local_target { + None + } else { + Some(put_start_resp.node_id.clone()) }, - ); + src_offset, + target_base_addr: put_start_resp.target_base_addr, + target_offset: if is_local_target { + target_offset + } else { + remote_offset + }, + len: req.len, + make_replica_task: req.make_replica_task, + remote_replica_admitted: replica_admitted, + preferred_sub_cluster: req.preferred_sub_cluster.clone(), + local_reserve_slot: None, + local_reserve_slot_size: None, + atomic_group: None, + _pending_fence: pending_fence_claim.take_for_pending_context(), + }, + ); + if !is_local_target { tracing::debug!( "external_put_start stash remote ctx: key={}, put_id=({},{}) peer_id={}, target_base={:#x}, target_off={:#x}, src_off(staging)={:#x}", req.key, @@ -489,6 +3644,241 @@ impl HandlerForExternalClient for ClientKvApi { }) } + async fn external_batch_put_start( + &self, + req: ExternalBatchPutStartReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_put_start")); + } + let inner = self.inner(); + + self.validate_requester_owner_status_updated(req.started_time)?; + + let atomic_group_lens = normalize_external_put_start_group_lens( + req.items.len(), + req.atomic_group_lens.clone(), + )?; + + let Some(first_item) = req.items.first() else { + return Ok(ExternalBatchPutStartResp { + items: Vec::new(), + error_code: OK, + error_json: String::new(), + }); + }; + let value_len = first_item.len; + if req.items.iter().any(|item| item.len != value_len) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "external_batch_put_start local-first requires uniform item len" + .to_string(), + })); + } + + let pending_fences = loop { + let mut pending_fences = Vec::with_capacity(req.items.len()); + let mut wait_for = None; + let mut wait_for_local_access = None; + for item in &req.items { + match inner.reserve_external_local_first_put_key( + &item.key, + item.reject_if_inflight_same_key, + item.reject_if_exist_same_key, + ) { + Ok(ExternalLocalFirstPutKeyReservation::Leader(pending_fence)) => { + pending_fences.push(pending_fence); + } + Ok(ExternalLocalFirstPutKeyReservation::Wait(op)) => { + if atomic_group_lens.len() != 1 { + let err = KvError::Api(ApiError::KeyBeingWritten { + key: item.key.clone(), + }); + return Ok(ExternalBatchPutStartResp { + items: req + .items + .iter() + .map(|_| external_local_first_error_item(&err)) + .collect(), + error_code: OK, + error_json: String::new(), + }); + } + wait_for = Some((item.key.clone(), op)); + break; + } + Ok(ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(completion)) => { + wait_for_local_access = Some((item.key.clone(), completion)); + break; + } + Err(err) => { + let items = req + .items + .iter() + .map(|_| external_local_first_error_item(&err)) + .collect(); + return Ok(ExternalBatchPutStartResp { + items, + error_code: OK, + error_json: String::new(), + }); + } + } + } + if let Some((fenced_key, mut completion)) = wait_for_local_access { + // A partial atomic_batch must not hold local-put fences while + // the precise source-eviction/reclaim generation drains. The + // watch receiver was subscribed under the per-key lock, so a + // rollback/finalize between this drop and changed().await is + // observed without polling or a lost wakeup. + drop(pending_fences); + let wait_started_at = Instant::now(); + tracing::info!( + fenced_key, + items = req.items.len(), + "external local-first Put waiting for owner source/reclaim fence" + ); + loop { + if *completion.borrow_and_update() { + break; + } + if completion.changed().await.is_err() { + // The state owner disappeared; re-evaluate under the + // key lock instead of treating channel close as success. + break; + } + } + tracing::info!( + fenced_key, + items = req.items.len(), + wait_us = duration_to_i64_us(wait_started_at.elapsed()), + "external local-first Put resumed after owner source/reclaim fence" + ); + continue; + } + let Some((joined_key, op)) = wait_for else { + break pending_fences; + }; + + // Never retain a partial atomic_batch while waiting for another + // key. Two overlapping requests with different key order must not + // each hold one leader fence and wait forever on the other. + drop(pending_fences); + match op.wait().await { + ExternalPutKeyOutcome::Succeeded => { + tracing::info!( + joined_key, + items = req.items.len(), + "external local-first Put reused an inflight leader result" + ); + let err = KvError::Api(ApiError::KeyAlreadyExists { key: joined_key }); + return Ok(ExternalBatchPutStartResp { + items: req + .items + .iter() + .map(|_| external_local_first_error_item(&err)) + .collect(), + error_code: OK, + error_json: String::new(), + }); + } + ExternalPutKeyOutcome::Failed => continue, + ExternalPutKeyOutcome::InFlight => { + unreachable!("shared Put wait must return a terminal outcome") + } + } + }; + + // Finish every fallible, purely logical derivation before claiming + // physical slots. The per-key fences above are RAII-owned, so any + // error or cancellation before cache insertion releases both counters. + let put_ids = req + .items + .iter() + .map(|_| inner.next_external_local_first_put_id()) + .collect::>(); + let keys_and_put_ids = req + .items + .iter() + .map(|item| item.key.clone()) + .zip(put_ids.iter().copied()) + .collect::>(); + let atomic_groups = + build_put_atomic_group_assignments(&keys_and_put_ids, &atomic_group_lens) + .map_err(|detail| KvError::Api(ApiError::InvalidArgument { detail }))?; + + let slot_lease = match inner + .owner_claim_local_reserve_slot_lease(value_len, req.items.len()) + .await + { + Ok(slot_lease) => slot_lease, + Err(err) => return Err(err), + }; + let self_node_id = inner.view.cluster_manager().get_self_info().id.clone(); + let slot_size = slot_lease.slot_size; + for (slot_ref, pending_fence) in slot_lease.slots.iter().zip(&pending_fences) { + pending_fence.attach_local_slot_lease(OwnerLocalReserveSlotLease { + value_len, + slot_size, + slots: vec![slot_ref.clone()], + }); + } + let mut items = Vec::with_capacity(req.items.len()); + for (idx, ((req_item, slot_ref), pending_fence)) in req + .items + .into_iter() + .zip(slot_lease.slots.into_iter()) + .zip(pending_fences.into_iter()) + .enumerate() + { + let put_id = put_ids[idx]; + let src_offset = slot_ref.ptr.saturating_sub(slot_ref.base_addr); + inner.external_pending_puts.insert( + (req_item.key.clone(), put_id.0, put_id.1), + ExternalPendingPutCtx { + peer_id: None, + src_offset, + target_base_addr: slot_ref.base_addr, + target_offset: src_offset, + len: req_item.len, + make_replica_task: req_item.make_replica_task, + remote_replica_admitted: true, + preferred_sub_cluster: req_item.preferred_sub_cluster.clone(), + local_reserve_slot: Some(slot_ref.clone()), + local_reserve_slot_size: Some(slot_size), + atomic_group: atomic_groups[idx].clone(), + _pending_fence: pending_fence, + }, + ); + tracing::debug!( + "external_batch_put_start local-first: key={} put_id=({},{}) node_id={} base={:#x} offset={:#x} len={}", + req_item.key, + put_id.0, + put_id.1, + self_node_id, + slot_ref.base_addr, + src_offset, + req_item.len + ); + items.push(ExternalBatchPutStartItemResp { + error_code: OK, + src_offset, + target_offset: src_offset, + transfer_target_offset: None, + peer_id: None, + src_base_addr: slot_ref.base_addr, + target_base_addr: slot_ref.base_addr, + error_json: String::new(), + put_id: Some(put_id), + }); + } + + Ok(ExternalBatchPutStartResp { + items, + error_code: OK, + error_json: String::new(), + }) + } + /// Handle external transfer+end request - transfer data then commit async fn external_put_transfer_end( &self, @@ -512,83 +3902,57 @@ impl HandlerForExternalClient for ClientKvApi { return Ok(ExternalPutTransferEndResp::from_error(&err)); }; - // Delegate transfer to owner put_transfer (offset-based) then end - // For remote target, ignore caller's target_offset (staging) and use stashed remote ctx. - let (peer_id_for_transfer, target_off_for_transfer, target_base_for_transfer) = - if let Some(peer) = req.peer_id.clone() { - // Remote path requires a stashed context; do not fallback to caller-provided target_offset - match inner - .external_pending_puts - .get(&(req.key.clone(), put_id.0, put_id.1)) - { - Some(ctx) => ( - Some(ctx.peer_id.clone()), - ctx.target_offset, - Some(ctx.target_base_addr), - ), - None => match req.target_base_addr { - Some(target_base_addr) => { - (Some(peer.into()), req.target_offset, Some(target_base_addr)) - } - None => { - let err = KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { - rpc_input_json: format!( - "missing pending remote put ctx and caller target_base_addr; key={}, put_id=({},{}), peer_id={}", - req.key, put_id.0, put_id.1, peer - ), - }, - ); - return Ok(ExternalPutTransferEndResp::from_error(&err)); - } - }, + let pending_ctx = match require_external_pending_put_ctx( + inner, + "external_put_transfer_end", + &req.key, + put_id, + ) { + Ok(ctx) => ctx, + Err(err) => { + best_effort_revoke_missing_external_pending_ctx( + inner, + "external_put_transfer_end", + &req.key, + put_id, + ) + .await; + return Ok(ExternalPutTransferEndResp::from_error(&err)); + } + }; + let admitted_replica_task = + pending_ctx.make_replica_task && pending_ctx.remote_replica_admitted; + let self_node_id = inner.view.cluster_manager().get_self_info().id.clone(); + let req_remote_target = req + .peer_id + .as_deref() + .is_some_and(|peer| peer != self_node_id.as_str()); + let has_remote_target = pending_ctx.peer_id.is_some(); + if req_remote_target != has_remote_target { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_put_transfer_end peer_id mismatches pending ctx: key={} put_id=({},{}) req_remote_target={} ctx_remote_target={}", + req.key, put_id.0, put_id.1, req_remote_target, has_remote_target + ), + }); + match inner.put_revoke(&req.key, put_id).await { + Ok(_) => { + inner + .external_pending_puts + .invalidate(&(req.key.clone(), put_id.0, put_id.1)) } - } else { - (None, req.target_offset, None) - }; - let transfer_peer_id_for_trace = peer_id_for_transfer.as_ref().map(|peer| peer.to_string()); - - tracing::debug!( - "external_put_transfer_end resolved: key={}, put_id=({},{}) src_off={:#x}, tgt_off={:#x}, len={}, peer_id={:?}, target_base={:?}", - req.key, - put_id.0, - put_id.1, - req.src_offset, - target_off_for_transfer, - req.len, - peer_id_for_transfer, - target_base_for_transfer - ); - - let transfer_started_at = Instant::now(); - if let Err(e) = inner - .put_transfer( - &req.key, - put_id, - req.src_offset, - target_off_for_transfer, - req.len, - peer_id_for_transfer, - target_base_for_transfer, - ) - .await - { - tracing::error!("Failed to transfer data: {}", e); - // On transfer failure, revoke the put to release resources - if let Err(revoke_err) = inner.put_revoke(&req.key, put_id).await { - tracing::warn!( - "put_revoke also failed after transfer error: {}", + Err(revoke_err) => tracing::warn!( + "external_put_transfer_end put_revoke failed after peer_id mismatch; pending fence retained: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, revoke_err - ); + ), } - // Cleanup pending ctx on failure - inner - .external_pending_puts - .invalidate(&(req.key.clone(), put_id.0, put_id.1)); - return Ok(crate::rpcresp_kvresult_convert::FromError::from_error(&e)); + return Ok(ExternalPutTransferEndResp::from_error(&err)); } - - let put_transfer_total_us = duration_to_i64_us(transfer_started_at.elapsed()); + let put_transfer_total_us = 0; + let transfer_peer_id_for_trace = None; if inner.skip_put_end_commit_enabled() { inner @@ -613,48 +3977,507 @@ impl HandlerForExternalClient for ClientKvApi { owner_put_transfer_peer_id: transfer_peer_id_for_trace, ..Default::default() }) - } else { - None - }, + } else { + None + }, + }); + } + + let end_started_at = Instant::now(); + let publish_local_cache = true; + let local_cache_publish = match local_committed_cache_publish( + "external_put_transfer_end", + &req.key, + put_id, + &pending_ctx, + req.src_offset, + req.len, + ) { + Ok(publish) => publish, + Err(err) => { + match inner.put_revoke(&req.key, put_id).await { + Ok(_) => inner.external_pending_puts.invalidate(&( + req.key.clone(), + put_id.0, + put_id.1, + )), + Err(revoke_err) => tracing::warn!( + "external_put_transfer_end put_revoke failed after local cache publish precheck error; pending fence retained: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + revoke_err + ), + } + return Ok(crate::rpcresp_kvresult_convert::FromError::from_error(&err)); + } + }; + let put_end = match inner + .put_end_with_local_cache_publish(&req.key, put_id, req.lease_id, publish_local_cache) + .await + { + Ok(end) => end, + Err(e) => { + tracing::error!("Failed to end put operation: {}", e); + tracing::warn!( + "external_put_transfer_end PutDone was not terminal; pending fence retained for retry: key={} put_id=({},{})", + req.key, + put_id.0, + put_id.1 + ); + return Ok(crate::rpcresp_kvresult_convert::FromError::from_error(&e)); + } + }; + let put_end_stats = put_end.stats; + let put_end_total_us = duration_to_i64_us(end_started_at.elapsed()); + + let Some(holder_id) = put_end.local_cache_holder_id else { + let err = KvError::Api(ApiError::Unknown { + detail: format!( + "external_put_transfer_end missing local cache holder after local commit: key={} put_id=({},{})", + req.key, put_id.0, put_id.1 + ), + }); + tracing::warn!( + "external_put_transfer_end retained pending fence after terminal response omitted local holder: key={} put_id=({},{})", + req.key, + put_id.0, + put_id.1 + ); + return Ok(crate::rpcresp_kvresult_convert::FromError::from_error(&err)); + }; + inner + .install_local_committed_memory_info( + &req.key, + put_id, + local_cache_publish.src_offset, + local_cache_publish.len, + holder_id, + ) + .await?; + inner + .external_pending_puts + .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + if pending_ctx.make_replica_task { + inner.start_early_owner_local_ssd_put( + &req.key, + put_id, + u64::from(local_cache_publish.len), + ); + } + if admitted_replica_task { + if let Err(err) = inner + .ensure_remote_put( + &req.key, + put_id, + pending_ctx.preferred_sub_cluster.clone(), + true, + ) + .await + { + tracing::warn!( + "external_put_transfer_end make replica task failed after local commit: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + err + ); + } + } + + Ok(ExternalPutTransferEndResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + test_put_phase_trace: if req.test_observe_put_phases { + Some(TestPutPhaseTrace { + owner_external_put_transfer_end_total_us: duration_to_i64_us( + total_started_at.elapsed(), + ), + owner_put_transfer_total_us: put_transfer_total_us, + owner_put_transfer_peer_id: transfer_peer_id_for_trace, + owner_put_end_total_us: put_end_total_us, + owner_master_put_end_rpc_us: put_end_stats.master_put_end_rpc_us, + owner_master_put_end_server_us: put_end_stats.master_put_end_server_us, + ..Default::default() + }) + } else { + None + }, + }) + } + + async fn external_batch_put_transfer_end( + &self, + req: ExternalBatchPutTransferEndReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported( + "external_batch_put_transfer_end", + )); + } + let inner = self.inner(); + + self.validate_requester_owner_status_updated(req.started_time)?; + + if req.items.is_empty() { + return Ok(ExternalBatchPutTransferEndResp { + items: Vec::new(), + error_code: OK, + error_json: String::new(), + }); + } + + #[derive(Clone)] + struct DonePending { + idx: usize, + key: String, + put_id: crate::master_kv_router::put::PutIDForAKey, + lease_id: Option, + len: u64, + put_locality: Option<(bool, i64)>, + local_cache_publish: LocalCommittedCachePublish, + make_replica_task: bool, + remote_replica_admitted: bool, + // Keep the per-key reclaim fence alive until the master's terminal + // response has been applied to the owner-local index. + _pending_ctx: ExternalPendingPutCtx, + } + + let mut results: Vec> = + (0..req.items.len()).map(|_| None).collect(); + let mut done_pending = Vec::new(); + let mut revoke_pending = Vec::new(); + let mut local_publish_items = Vec::new(); + let mut local_publish_contexts = Vec::new(); + + for (idx, item) in req.items.into_iter().enumerate() { + let Some(put_id) = item.put_id else { + let err = KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { + rpc_input_json: format!( + "missing put_id in external_batch_put_transfer_end; key={}", + item.key + ), + }, + ); + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + continue; + }; + + let pending_ctx = match require_external_pending_put_ctx( + inner, + "external_batch_put_transfer_end", + &item.key, + put_id, + ) { + Ok(ctx) => ctx, + Err(err) => { + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + } + }; + if pending_ctx.local_reserve_slot.is_some() { + if item.peer_id.is_some() { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_put_transfer_end local-first item must be local target: key={} put_id=({},{})", + item.key, put_id.0, put_id.1 + ), + }); + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + release_external_local_first_pending_slot( + inner, + &pending_ctx, + "external_batch_put_transfer_end", + &item.key, + put_id, + ) + .await; + inner + .external_pending_puts + .invalidate(&(item.key.clone(), put_id.0, put_id.1)); + continue; + } + match commit_external_local_first_pending( + inner, + &item.key, + put_id, + &pending_ctx, + item.src_offset, + item.len, + "external_batch_put_transfer_end", + ) + .await + { + Ok(committed_slot) => { + inner.record_put_locality(false, item.len, 0); + inner.external_pending_puts.invalidate(&( + item.key.clone(), + put_id.0, + put_id.1, + )); + local_publish_items.push(OwnerLocalPublishItem { + key: item.key.clone(), + put_id, + value_len: item.len, + lease_id: item.lease_id, + committed_slot, + make_replica_task: pending_ctx.make_replica_task, + remote_replica_admitted: pending_ctx.remote_replica_admitted, + preferred_sub_cluster: pending_ctx.preferred_sub_cluster.clone(), + atomic_group: pending_ctx.atomic_group.clone(), + }); + local_publish_contexts.push(pending_ctx); + results[idx] = Some(ExternalBatchPutTransferEndItemResp { + error_code: OK, + error_json: String::new(), + }); + } + Err(err) => { + release_external_local_first_pending_slot( + inner, + &pending_ctx, + "external_batch_put_transfer_end", + &item.key, + put_id, + ) + .await; + inner.external_pending_puts.invalidate(&( + item.key.clone(), + put_id.0, + put_id.1, + )); + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + } + } + continue; + } + let has_remote_append = pending_ctx.peer_id.is_some(); + if item.peer_id.is_some() || has_remote_append { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_put_transfer_end remote primary target is no longer supported: key={} put_id=({},{}) req_remote_target={} ctx_remote_target={}", + item.key, + put_id.0, + put_id.1, + item.peer_id.is_some(), + has_remote_append + ), + }); + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + } + let local_cache_publish = match local_committed_cache_publish( + "external_batch_put_transfer_end", + &item.key, + put_id, + &pending_ctx, + item.src_offset, + item.len, + ) { + Ok(publish) => publish, + Err(err) => { + results[idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + } + }; + done_pending.push(DonePending { + idx, + key: item.key, + put_id, + lease_id: item.lease_id, + len: item.len, + put_locality: Some((false, 0)), + local_cache_publish, + make_replica_task: pending_ctx.make_replica_task, + remote_replica_admitted: pending_ctx.remote_replica_admitted, + _pending_ctx: pending_ctx, + }); + } + + if !local_publish_items.is_empty() { + spawn_external_local_first_publish(inner, local_publish_items, local_publish_contexts); + } + + if !revoke_pending.is_empty() { + let requested = revoke_pending.clone(); + match inner.batch_put_revoke(revoke_pending).await { + Ok(response) if response.items.len() == requested.len() => { + for (request, response) in requested.into_iter().zip(response.items) { + if response.key == request.key + && response.put_id == request.put_id + && response.error_code == OK + { + inner.external_pending_puts.invalidate(&( + request.key, + request.put_id.0, + request.put_id.1, + )); + } + } + } + Ok(response) => tracing::warn!( + "external_batch_put_transfer_end batch_put_revoke response length mismatch: expected={} got={}", + requested.len(), + response.items.len() + ), + Err(err) => tracing::warn!( + "external_batch_put_transfer_end batch_put_revoke failed after precheck errors; pending fences retained for retry: {}", + err + ), + } + } + + if inner.skip_put_end_commit_enabled() { + for pending in done_pending { + inner.external_pending_puts.invalidate(&( + pending.key.clone(), + pending.put_id.0, + pending.put_id.1, + )); + if let Some((remote, transfer_us)) = pending.put_locality { + inner.record_put_locality(remote, pending.len, transfer_us); + } + results[pending.idx] = Some(ExternalBatchPutTransferEndItemResp { + error_code: OK, + error_json: String::new(), + }); + } + return Ok(ExternalBatchPutTransferEndResp { + items: results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| ExternalBatchPutTransferEndItemResp { + error_code: OK, + error_json: String::new(), + }) + }) + .collect(), + error_code: OK, + error_json: String::new(), }); } - let end_started_at = Instant::now(); - let put_end_stats = match inner.put_end(&req.key, put_id, req.lease_id).await { - Ok(stats) => stats, - Err(e) => { - tracing::error!("Failed to end put operation: {}", e); - inner - .external_pending_puts - .invalidate(&(req.key.clone(), put_id.0, put_id.1)); - return Ok(crate::rpcresp_kvresult_convert::FromError::from_error(&e)); + let done_resp = inner + .batch_put_done( + done_pending + .iter() + .map(|pending| BatchPutDoneItemReq { + key: pending.key.clone(), + put_id: pending.put_id, + lease_id: pending.lease_id, + committed_slot: None, + publish_local_cache: true, + atomic_group: None, + }) + .collect(), + ) + .await?; + if done_resp.items.len() != done_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_put_transfer_end done response length mismatch: expected={} got={}", + done_pending.len(), + done_resp.items.len() + ), + })); + } + + let mut early_ssd_candidates = Vec::new(); + let mut remote_replica_pending = Vec::new(); + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + results[pending.idx] = Some(ExternalBatchPutTransferEndItemResp::from_error(&err)); + continue; } - }; - let put_end_total_us = duration_to_i64_us(end_started_at.elapsed()); + if let Some((remote, transfer_us)) = pending.put_locality { + inner.record_put_locality(remote, pending.len, transfer_us); + } + let Some(holder_id) = done_item.local_cache_holder_id else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_put_transfer_end missing local cache holder after local commit: key={} put_id=({},{})", + pending.key, pending.put_id.0, pending.put_id.1 + ), + })); + }; + inner + .install_local_committed_memory_info( + &pending.key, + pending.put_id, + pending.local_cache_publish.src_offset, + pending.local_cache_publish.len, + holder_id, + ) + .await?; + inner.external_pending_puts.invalidate(&( + pending.key.clone(), + pending.put_id.0, + pending.put_id.1, + )); + if pending.make_replica_task { + early_ssd_candidates.push(( + pending.key.clone(), + pending.put_id, + u64::from(pending.local_cache_publish.len), + )); + } + if pending.make_replica_task && pending.remote_replica_admitted { + remote_replica_pending.push(( + pending.key.clone(), + pending.put_id, + pending._pending_ctx.preferred_sub_cluster.clone(), + )); + } + results[pending.idx] = Some(ExternalBatchPutTransferEndItemResp { + error_code: OK, + error_json: String::new(), + }); + } - // Success; cleanup pending ctx - inner - .external_pending_puts - .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + inner.start_early_owner_local_ssd_puts(early_ssd_candidates); + for (key, put_id, preferred_sub_cluster) in remote_replica_pending { + if let Err(err) = inner + .ensure_remote_put(&key, put_id, preferred_sub_cluster, true) + .await + { + tracing::warn!( + "external_batch_put_transfer_end make replica task failed after local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } - Ok(ExternalPutTransferEndResp { - error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, - error_json: String::new(), - test_put_phase_trace: if req.test_observe_put_phases { - Some(TestPutPhaseTrace { - owner_external_put_transfer_end_total_us: duration_to_i64_us( - total_started_at.elapsed(), - ), - owner_put_transfer_total_us: put_transfer_total_us, - owner_put_transfer_peer_id: transfer_peer_id_for_trace, - owner_put_end_total_us: put_end_total_us, - owner_master_put_end_rpc_us: put_end_stats.master_put_end_rpc_us, - owner_master_put_end_server_us: put_end_stats.master_put_end_server_us, - ..Default::default() + Ok(ExternalBatchPutTransferEndResp { + items: results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| ExternalBatchPutTransferEndItemResp { + error_code: OK, + error_json: String::new(), + }) }) - } else { - None - }, + .collect(), + error_code: OK, + error_json: String::new(), }) } @@ -675,21 +4498,150 @@ impl HandlerForExternalClient for ClientKvApi { ); return Ok(ExternalPutCommitResp::from_error(&err)); }; + let pending_ctx = match require_external_pending_put_ctx( + inner, + "external_put_commit", + &req.key, + put_id, + ) { + Ok(ctx) => ctx, + Err(err) => { + best_effort_revoke_missing_external_pending_ctx( + inner, + "external_put_commit", + &req.key, + put_id, + ) + .await; + return Ok(ExternalPutCommitResp::from_error(&err)); + } + }; + let has_remote_append = pending_ctx.peer_id.is_some(); + if req.remote_target || has_remote_append { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_put_commit remote primary target is no longer supported: key={} put_id=({},{}) req_remote_target={} ctx_remote_target={}", + req.key, put_id.0, put_id.1, req.remote_target, has_remote_append + ), + }); + match inner.put_revoke(&req.key, put_id).await { + Ok(_) => { + inner + .external_pending_puts + .invalidate(&(req.key.clone(), put_id.0, put_id.1)) + } + Err(revoke_err) => tracing::warn!( + "external_put_commit put_revoke failed after remote primary target; pending fence retained: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + revoke_err + ), + } + return Ok(ExternalPutCommitResp::from_error(&err)); + } + let admitted_replica_task = + pending_ctx.make_replica_task && pending_ctx.remote_replica_admitted; let end_started_at = Instant::now(); - let put_end_stats = match inner.put_end(&req.key, put_id, req.lease_id).await { - Ok(stats) => stats, + let local_cache_publish = match local_committed_cache_publish( + "external_put_commit", + &req.key, + put_id, + &pending_ctx, + req.src_offset, + req.len, + ) { + Ok(publish) => publish, + Err(err) => { + match inner.put_revoke(&req.key, put_id).await { + Ok(_) => inner.external_pending_puts.invalidate(&( + req.key.clone(), + put_id.0, + put_id.1, + )), + Err(revoke_err) => tracing::warn!( + "external_put_commit put_revoke failed after local cache publish precheck error; pending fence retained: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + revoke_err + ), + } + return Ok(ExternalPutCommitResp::from_error(&err)); + } + }; + let put_end = match inner + .put_end_with_local_cache_publish(&req.key, put_id, req.lease_id, true) + .await + { + Ok(end) => end, Err(e) => { - inner - .external_pending_puts - .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + tracing::warn!( + "external_put_commit PutDone was not terminal; pending fence retained for retry: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + e + ); return Ok(ExternalPutCommitResp::from_error(&e)); } }; + let put_end_stats = put_end.stats; let put_end_total_us = duration_to_i64_us(end_started_at.elapsed()); + let Some(holder_id) = put_end.local_cache_holder_id else { + let err = KvError::Api(ApiError::Unknown { + detail: format!( + "external_put_commit missing local cache holder after local commit: key={} put_id=({},{})", + req.key, put_id.0, put_id.1 + ), + }); + tracing::warn!( + "external_put_commit retained pending fence after terminal response omitted local holder: key={} put_id=({},{})", + req.key, + put_id.0, + put_id.1 + ); + return Ok(ExternalPutCommitResp::from_error(&err)); + }; + inner + .install_local_committed_memory_info( + &req.key, + put_id, + local_cache_publish.src_offset, + local_cache_publish.len, + holder_id, + ) + .await?; inner .external_pending_puts .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + if pending_ctx.make_replica_task { + inner.start_early_owner_local_ssd_put( + &req.key, + put_id, + u64::from(local_cache_publish.len), + ); + } + if admitted_replica_task { + if let Err(err) = inner + .ensure_remote_put( + &req.key, + put_id, + pending_ctx.preferred_sub_cluster.clone(), + true, + ) + .await + { + tracing::warn!( + "external_put_commit make replica task failed after local commit: key={} put_id=({},{}) err={}", + req.key, + put_id.0, + put_id.1, + err + ); + } + } Ok(ExternalPutCommitResp { error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, @@ -707,6 +4659,362 @@ impl HandlerForExternalClient for ClientKvApi { }) } + async fn external_batch_put_commit( + &self, + req: ExternalBatchPutCommitReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_put_commit")); + } + let inner = self.inner(); + + self.validate_requester_owner_status_updated(req.started_time)?; + + if req.items.is_empty() { + return Ok(ExternalBatchPutCommitResp { + items: Vec::new(), + error_code: OK, + error_json: String::new(), + }); + } + + #[derive(Clone)] + struct DonePending { + idx: usize, + key: String, + put_id: crate::master_kv_router::put::PutIDForAKey, + lease_id: Option, + local_cache_publish: LocalCommittedCachePublish, + make_replica_task: bool, + remote_replica_admitted: bool, + _pending_ctx: ExternalPendingPutCtx, + } + + let mut results: Vec> = + (0..req.items.len()).map(|_| None).collect(); + let mut done_pending = Vec::new(); + let mut revoke_pending = Vec::new(); + let mut local_publish_items = Vec::new(); + let mut local_publish_contexts = Vec::new(); + + for (idx, item) in req.items.into_iter().enumerate() { + let Some(put_id) = item.put_id else { + let err = KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { + rpc_input_json: format!( + "missing put_id in external_batch_put_commit; key={}", + item.key + ), + }, + ); + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + continue; + }; + let pending_ctx = match require_external_pending_put_ctx( + inner, + "external_batch_put_commit", + &item.key, + put_id, + ) { + Ok(ctx) => ctx, + Err(err) => { + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + } + }; + if pending_ctx.local_reserve_slot.is_some() { + if item.remote_target { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_put_commit local-first item must be local target: key={} put_id=({},{})", + item.key, put_id.0, put_id.1 + ), + }); + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + release_external_local_first_pending_slot( + inner, + &pending_ctx, + "external_batch_put_commit", + &item.key, + put_id, + ) + .await; + inner + .external_pending_puts + .invalidate(&(item.key.clone(), put_id.0, put_id.1)); + continue; + } + match commit_external_local_first_pending( + inner, + &item.key, + put_id, + &pending_ctx, + item.src_offset, + item.len, + "external_batch_put_commit", + ) + .await + { + Ok(committed_slot) => { + inner.record_put_locality(false, item.len, 0); + inner.external_pending_puts.invalidate(&( + item.key.clone(), + put_id.0, + put_id.1, + )); + local_publish_items.push(OwnerLocalPublishItem { + key: item.key.clone(), + put_id, + value_len: item.len, + lease_id: item.lease_id, + committed_slot, + make_replica_task: pending_ctx.make_replica_task, + remote_replica_admitted: pending_ctx.remote_replica_admitted, + preferred_sub_cluster: pending_ctx.preferred_sub_cluster.clone(), + atomic_group: pending_ctx.atomic_group.clone(), + }); + local_publish_contexts.push(pending_ctx); + results[idx] = Some(ExternalBatchPutCommitItemResp { + error_code: OK, + error_json: String::new(), + }); + } + Err(err) => { + release_external_local_first_pending_slot( + inner, + &pending_ctx, + "external_batch_put_commit", + &item.key, + put_id, + ) + .await; + inner.external_pending_puts.invalidate(&( + item.key.clone(), + put_id.0, + put_id.1, + )); + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + } + } + continue; + } + let has_remote_append = pending_ctx.peer_id.is_some(); + if item.remote_target || has_remote_append { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external_batch_put_commit remote primary target is no longer supported: key={} put_id=({},{}) req_remote_target={} ctx_remote_target={}", + item.key, put_id.0, put_id.1, item.remote_target, has_remote_append + ), + }); + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + }; + let local_cache_publish = match local_committed_cache_publish( + "external_batch_put_commit", + &item.key, + put_id, + &pending_ctx, + item.src_offset, + item.len, + ) { + Ok(publish) => publish, + Err(err) => { + results[idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + revoke_pending.push(BatchPutRevokeItemReq { + key: item.key.clone(), + put_id, + }); + continue; + } + }; + done_pending.push(DonePending { + idx, + key: item.key.clone(), + put_id, + lease_id: item.lease_id, + local_cache_publish, + make_replica_task: pending_ctx.make_replica_task, + remote_replica_admitted: pending_ctx.remote_replica_admitted, + _pending_ctx: pending_ctx, + }); + } + + if !local_publish_items.is_empty() { + spawn_external_local_first_publish(inner, local_publish_items, local_publish_contexts); + } + + if !revoke_pending.is_empty() { + let requested = revoke_pending.clone(); + match inner.batch_put_revoke(revoke_pending).await { + Ok(response) if response.items.len() == requested.len() => { + for (request, response) in requested.into_iter().zip(response.items) { + if response.key == request.key + && response.put_id == request.put_id + && response.error_code == OK + { + inner.external_pending_puts.invalidate(&( + request.key, + request.put_id.0, + request.put_id.1, + )); + } + } + } + Ok(response) => tracing::warn!( + "external_batch_put_commit batch_put_revoke response length mismatch: expected={} got={}", + requested.len(), + response.items.len() + ), + Err(err) => tracing::warn!( + "external_batch_put_commit batch_put_revoke failed after precheck errors; pending fences retained for retry: {}", + err + ), + } + } + + if inner.skip_put_end_commit_enabled() { + for pending in done_pending { + inner.external_pending_puts.invalidate(&( + pending.key.clone(), + pending.put_id.0, + pending.put_id.1, + )); + results[pending.idx] = Some(ExternalBatchPutCommitItemResp { + error_code: OK, + error_json: String::new(), + }); + } + return Ok(ExternalBatchPutCommitResp { + items: results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| ExternalBatchPutCommitItemResp { + error_code: OK, + error_json: String::new(), + }) + }) + .collect(), + error_code: OK, + error_json: String::new(), + }); + } + + let done_resp = inner + .batch_put_done( + done_pending + .iter() + .map(|pending| BatchPutDoneItemReq { + key: pending.key.clone(), + put_id: pending.put_id, + lease_id: pending.lease_id, + committed_slot: None, + publish_local_cache: true, + atomic_group: None, + }) + .collect(), + ) + .await?; + if done_resp.items.len() != done_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_put_commit done response length mismatch: expected={} got={}", + done_pending.len(), + done_resp.items.len() + ), + })); + } + + let mut early_ssd_candidates = Vec::new(); + let mut remote_replica_pending = Vec::new(); + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + results[pending.idx] = Some(ExternalBatchPutCommitItemResp::from_error(&err)); + continue; + } + let Some(holder_id) = done_item.local_cache_holder_id else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external_batch_put_commit missing local cache holder after local commit: key={} put_id=({},{})", + pending.key, pending.put_id.0, pending.put_id.1 + ), + })); + }; + inner + .install_local_committed_memory_info( + &pending.key, + pending.put_id, + pending.local_cache_publish.src_offset, + pending.local_cache_publish.len, + holder_id, + ) + .await?; + inner.external_pending_puts.invalidate(&( + pending.key.clone(), + pending.put_id.0, + pending.put_id.1, + )); + if pending.make_replica_task { + early_ssd_candidates.push(( + pending.key.clone(), + pending.put_id, + u64::from(pending.local_cache_publish.len), + )); + } + if pending.make_replica_task && pending.remote_replica_admitted { + remote_replica_pending.push(( + pending.key.clone(), + pending.put_id, + pending._pending_ctx.preferred_sub_cluster.clone(), + )); + } + results[pending.idx] = Some(ExternalBatchPutCommitItemResp { + error_code: OK, + error_json: String::new(), + }); + } + + inner.start_early_owner_local_ssd_puts(early_ssd_candidates); + for (key, put_id, preferred_sub_cluster) in remote_replica_pending { + if let Err(err) = inner + .ensure_remote_put(&key, put_id, preferred_sub_cluster, true) + .await + { + tracing::warn!( + "external_batch_put_commit make replica task failed after local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } + + Ok(ExternalBatchPutCommitResp { + items: results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| ExternalBatchPutCommitItemResp { + error_code: OK, + error_json: String::new(), + }) + }) + .collect(), + error_code: OK, + error_json: String::new(), + }) + } + async fn external_put_revoke( &self, req: ExternalPutRevokeReq, @@ -724,14 +5032,52 @@ impl HandlerForExternalClient for ClientKvApi { ); return Ok(ExternalPutRevokeResp::from_error(&err)); }; - inner - .external_pending_puts - .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + let pending_ctx = match require_external_pending_put_ctx( + inner, + "external_put_revoke", + &req.key, + put_id, + ) { + Ok(ctx) => ctx, + Err(err) => { + best_effort_revoke_missing_external_pending_ctx( + inner, + "external_put_revoke", + &req.key, + put_id, + ) + .await; + return Ok(ExternalPutRevokeResp::from_error(&err)); + } + }; + if pending_ctx.local_reserve_slot.is_some() { + return match pending_ctx + ._pending_fence + .release_local_slot_lease_now(inner) + .await + { + Ok(_) => { + inner + .external_pending_puts + .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + Ok(ExternalPutRevokeResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }) + } + Err(e) => Ok(ExternalPutRevokeResp::from_error(&e)), + }; + } match inner.put_revoke(&req.key, put_id).await { - Ok(_) => Ok(ExternalPutRevokeResp { - error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, - error_json: String::new(), - }), + Ok(_) => { + inner + .external_pending_puts + .invalidate(&(req.key.clone(), put_id.0, put_id.1)); + Ok(ExternalPutRevokeResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }) + } Err(e) => Ok(ExternalPutRevokeResp::from_error(&e)), } } @@ -775,4 +5121,46 @@ impl HandlerForExternalClient for ClientKvApi { }), } } + + async fn external_batch_is_exist( + &self, + req: ExternalBatchIsExistReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported("external_batch_is_exist")); + } + let inner = self.inner(); + + self.validate_requester_owner_status_updated(req.started_time)?; + + match inner + .batch_is_exist(req.keys, req.allow_local_snapshot) + .await + { + Ok(exists_list) => Ok(ExternalBatchIsExistResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + exists_list, + error_json: String::new(), + }), + Err(e) => Ok(ExternalBatchIsExistResp { + exists_list: Vec::new(), + ..ExternalBatchIsExistResp::from_error(&e) + }), + } + } + + async fn external_observability_snapshot( + &self, + req: ExternalObservabilitySnapshotReq, + ) -> KvResult { + if self.is_side_transfer_worker() { + return Err(Self::side_transfer_unsupported( + "external_observability_snapshot", + )); + } + self.validate_requester_owner_status_updated(req.started_time)?; + Ok(ExternalObservabilitySnapshotResp::success( + self.inner().locality_snapshot(), + )) + } } diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs index f309dd0..27d57ec 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/get.rs @@ -1,5 +1,7 @@ -use super::ClientKvApiInner; -use crate::client_kv_api::CachedValue; +use super::{ + ClientKvApiInner, ClientKvApiView, KvMetrics, OwnerLocalReserveSlotLease, + OwnerLocalReserveSlotRef, +}; use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::memholder::{MemoryInfo, UserMemHolder, UserMemHolderExposeKind}; // no StageScope; timestamps-based metrics only @@ -12,15 +14,189 @@ use crate::observe_kvope::{ use crate::{ cluster_manager::NodeID, master_kv_router::msg_pack::{ - GetAllocationMode, GetDoneReq, GetDoneResp, GetMetaReq, GetMetaResp, GetRevokeReq, - GetStartReq, GetStartResp, + BatchGetBindItemReq, BatchGetBindReq, BatchGetBindResp, BatchGetDoneReq, BatchGetDoneResp, + BatchGetRevokeReq, BatchGetRevokeResp, BatchGetStartItemResp, BatchGetStartReq, + BatchGetStartResp, BatchIsExistReq, GetAllocationMode, GetBindTarget, GetDoneReq, + GetDoneResp, GetMetaReq, GetMetaResp, GetPreparedLocalReserveTarget, GetRevokeReq, + GetSourceKind, GetStartReq, GetStartResp, }, p2p::msg_pack::MsgPack, rpcresp_kvresult_convert::msg_and_error::codes_api, rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, OK}, }; +use ::tokio::sync::Semaphore; use chrono::Utc; +use futures::stream::{self, StreamExt}; +use limit_thirdparty::tokio; +use std::collections::HashSet; use std::sync::Arc; +use std::sync::OnceLock; +use std::time::Duration; +use std::time::Instant; + +const BATCH_GET_DONE_MAX_INFLIGHT: usize = 4; + +fn batch_get_done_rpc_limiter() -> &'static Semaphore { + static LIMITER: OnceLock = OnceLock::new(); + LIMITER.get_or_init(|| Semaphore::new(BATCH_GET_DONE_MAX_INFLIGHT)) +} + +async fn release_prepared_get_target( + inner: &ClientKvApiInner, + target: &GetPreparedLocalReserveTarget, +) -> KvResult<()> { + inner + .owner_release_local_reserve_slot_lease(OwnerLocalReserveSlotLease { + value_len: target.slot_size, + slot_size: target.slot_size, + slots: vec![OwnerLocalReserveSlotRef { + grant_id: target.grant_id, + slot_index: target.slot_index, + ptr: target.addr, + base_addr: target.base_addr, + }], + }) + .await +} + +fn batch_get_done_response_matches(get_ids: &[u64], response: &BatchGetDoneResp) -> bool { + response.items.len() == get_ids.len() + && response + .items + .iter() + .zip(get_ids) + .all(|(item, expected_get_id)| item.get_id == *expected_get_id) +} + +#[derive(Clone)] +pub(crate) struct StartedGetRevokeCleanup { + pub(crate) get_id: u64, + pub(crate) prepared_target: Option, +} + +async fn run_started_get_revoke_cleanup( + view: ClientKvApiView, + pending: Vec, + context: &'static str, +) { + if pending.is_empty() { + return; + } + let mut attempt = 1u32; + loop { + let get_ids = pending.iter().map(|item| item.get_id).collect::>(); + let response = view.client_kv_api().inner().batch_get_revoke(get_ids).await; + let resp = match response { + Ok(resp) + if resp.items.len() == pending.len() + && resp + .items + .iter() + .zip(&pending) + .all(|(resp, expected)| resp.get_id == expected.get_id) => + { + resp + } + Ok(resp) => { + tracing::warn!( + "{} Revoke response shape/identity mismatch; retaining prepared slots for retry: expected={} got={} attempt={}", + context, + pending.len(), + resp.items.len(), + attempt + ); + tokio::time::sleep(Duration::from_millis( + (50u64.saturating_mul(1u64 << attempt.min(6))).min(2_000), + )) + .await; + attempt = attempt.saturating_add(1); + continue; + } + Err(err) => { + if matches!(&err, KvError::Api(ApiError::SystemShutdown { .. })) { + tracing::warn!( + "{} Revoke cleanup stopped during owner shutdown: items={}", + context, + pending.len() + ); + return; + } + tracing::warn!( + "{} Revoke uncertain; retaining get ids and prepared slots for retry: items={} attempt={} err={}", + context, + pending.len(), + attempt, + err + ); + tokio::time::sleep(Duration::from_millis( + (50u64.saturating_mul(1u64 << attempt.min(6))).min(2_000), + )) + .await; + attempt = attempt.saturating_add(1); + continue; + } + }; + + for (expected, item_resp) in pending.iter().zip(resp.items) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + item_resp.error_code, + item_resp.error_json, + ) { + // A terminal Done may have won. Never release a possibly + // committed slot from the losing Revoke path. + tracing::warn!( + "{} Revoke reached non-releasable terminal: get_id={} err={}", + context, + expected.get_id, + err + ); + continue; + } + let Some(target) = expected.prepared_target.as_ref() else { + continue; + }; + let mut release_attempt = 1u32; + loop { + match release_prepared_get_target(view.client_kv_api().inner(), target).await { + Ok(()) => break, + Err(err) => { + tracing::error!( + "{} Revoke confirmed but prepared slot release failed; retrying: get_id={} attempt={} err={}", + context, + expected.get_id, + release_attempt, + err + ); + tokio::time::sleep(Duration::from_secs(1)).await; + release_attempt = release_attempt.saturating_add(1); + } + } + } + } + return; + } +} + +/// Move cleanup ownership to a registered task before awaiting it. If the +/// caller future is cancelled, the task still drives Revoke to a definite +/// terminal and releases only confirmed-uncommitted prepared slots. +pub(crate) async fn finish_started_get_revoke_cleanup( + inner: &ClientKvApiInner, + pending: Vec, + context: &'static str, +) { + if pending.is_empty() { + return; + } + let (done_tx, done_rx) = ::tokio::sync::oneshot::channel::<()>(); + let spawn_view = inner.view.clone_view(); + let worker_view = spawn_view.clone(); + spawn_view.spawn("started_get_revoke_cleanup", async move { + run_started_get_revoke_cleanup(worker_view, pending, context).await; + let _ = done_tx.send(()); + }); + let _ = done_rx.await; +} #[derive(Debug, Clone)] pub struct RemoteGetInfo { @@ -47,57 +223,588 @@ impl std::fmt::Display for RemoteGetInfo { } } +impl RemoteGetInfo { + pub fn data_len(&self) -> usize { + self.data_len + } + + pub fn is_remote_transfer(&self) -> bool { + self.peer_is_src_or_target + } +} + impl ClientKvApiInner { - /// becaused we cached local kv metadata, so we make `MemHolder` with Arc here - pub async fn get( + pub async fn batch_get_finish_started( &self, - key: &str, - ) -> KvResult, Option)>> { + keys: Vec, + start_items: Vec, + transfer_concurrency: usize, + ) -> KvResult, Option)>>>> { if !self.view.register_shutdown_poller().is_running() { return Err(KvError::Api(ApiError::SystemShutdown { - detail: "ClientKvApi is shutting down; rejecting get".to_string(), + detail: "ClientKvApi is shutting down; rejecting batch_get_finish_started" + .to_string(), + })); + } + if keys.len() != start_items.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch_get_finish_started length mismatch: keys={} start_items={}", + keys.len(), + start_items.len() + ), })); } + if keys.is_empty() { + return Ok(Vec::new()); + } + let lifecycle_started_at = Instant::now(); + let lifecycle_requested_keys = keys.len(); + + #[derive(Clone)] + struct DonePending { + idx: usize, + key: String, + start_item: BatchGetStartItemResp, + peer_is_remote: bool, + transfer_us: i64, + prepared_memory_info: Option>, + } + + let transfer_concurrency = transfer_concurrency.max(1); let metrics = self.metrics_handle(); let client_id = self.client_id_str(); let node_role = self.node_role(); + let self_node_id = self.view.cluster_manager().get_self_info().id.clone(); - if let Some(cached_info) = self.get_cached_info.get(key) { - if let CachedValue::LocalReplica(memory_info) = &cached_info.value { - tracing::debug!( - "cache hit for key: {} with putid({},{}), directly return", + let mut results: Vec< + Option, Option)>>>, + > = (0..keys.len()).map(|_| None).collect(); + let mut done_pending = Vec::new(); + let mut transfer_error_cleanup = Vec::new(); + let mut transfer_futures = Vec::new(); + let mut lifecycle_zero_copy_items = 0usize; + let mut lifecycle_transfer_items = 0usize; + let mut lifecycle_remote_transfer_items = 0usize; + let mut lifecycle_transfer_bytes = 0u64; + let mut lifecycle_remote_transfer_bytes = 0u64; + let mut lifecycle_local_ssd_items = 0usize; + let mut lifecycle_remote_ssd_items = 0usize; + let mut lifecycle_transfer_source_nodes = HashSet::new(); + let mut lifecycle_ssd_source_nodes = HashSet::new(); + + for (idx, (key, start_item)) in keys.into_iter().zip(start_items.into_iter()).enumerate() { + if start_item.error_code == codes_api::API_KEY_NOT_FOUND { + if let Some(target) = start_item.prepared_target.as_ref() { + release_prepared_get_target(self, target).await?; + } + results[idx] = Some(Ok(None)); + continue; + } + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + start_item.error_code, + start_item.error_json.clone(), + ) { + if let Some(target) = start_item.prepared_target.as_ref() { + release_prepared_get_target(self, target).await?; + } + results[idx] = Some(Err(err)); + continue; + } + + let peer_id = if start_item.node_id == self_node_id { + None + } else { + Some(start_item.node_id.clone()) + }; + let peer_is_remote = peer_id.is_some(); + let get_id = start_item.get_id; + let src_addr = start_item.src_addr; + let target_addr = start_item.target_addr; + let len = start_item.len; + + lifecycle_transfer_source_nodes.insert(start_item.node_id.to_string()); + if start_item.source_kind == GetSourceKind::Ssd { + lifecycle_ssd_source_nodes.insert(start_item.node_id.to_string()); + if peer_is_remote { + lifecycle_remote_ssd_items = lifecycle_remote_ssd_items.saturating_add(1); + } else { + lifecycle_local_ssd_items = lifecycle_local_ssd_items.saturating_add(1); + } + } + + if start_item.source_kind == GetSourceKind::Memory + && peer_id.is_none() + && src_addr == target_addr + { + lifecycle_zero_copy_items = lifecycle_zero_copy_items.saturating_add(1); + done_pending.push(DonePending { + idx, key, - cached_info.put_time_ms, - cached_info.put_version - ); - let user_mem_holder = Arc::new(UserMemHolder::new( + start_item, + peer_is_remote, + transfer_us: 0, + prepared_memory_info: None, + }); + continue; + } + + lifecycle_transfer_items = lifecycle_transfer_items.saturating_add(1); + lifecycle_transfer_bytes = lifecycle_transfer_bytes.saturating_add(len); + if peer_is_remote { + lifecycle_remote_transfer_items = lifecycle_remote_transfer_items.saturating_add(1); + lifecycle_remote_transfer_bytes = + lifecycle_remote_transfer_bytes.saturating_add(len); + } + + transfer_futures.push(async move { + let transfer_started_at = Instant::now(); + let transfer_result = if start_item.source_kind == GetSourceKind::Ssd { + self.stage_kv_from_ssd_source( + &start_item.node_id, + &key, + start_item.put_id, + get_id, + src_addr, + len, + target_addr, + len, + ) + .await + } else { + self.view + .client_transfer_engine() + .transfer_data_no_copy(peer_id, true, src_addr, target_addr, len, None) + .await + .map(|_| ()) + .map_err(|err| { + KvError::Api(ApiError::Transfer { + from_addr: src_addr, + to_addr: target_addr, + len, + error: err.to_string(), + }) + }) + }; + let transfer_us = transfer_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + ( + idx, + key, + start_item, + peer_is_remote, + get_id, + transfer_us, + transfer_result, + ) + }); + } + + let lifecycle_plan_us = lifecycle_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + let transfer_wall_started_at = Instant::now(); + let mut lifecycle_transfer_sum_us = 0u64; + let mut lifecycle_transfer_max_us = 0u64; + let mut transfer_stream = + stream::iter(transfer_futures).buffer_unordered(transfer_concurrency); + while let Some(joined) = transfer_stream.next().await { + match joined { + (idx, key, start_item, peer_is_remote, _get_id, transfer_us, Ok(_breakdown)) => { + let transfer_us_u64 = transfer_us.max(0) as u64; + lifecycle_transfer_sum_us = + lifecycle_transfer_sum_us.saturating_add(transfer_us_u64); + lifecycle_transfer_max_us = lifecycle_transfer_max_us.max(transfer_us_u64); + done_pending.push(DonePending { + idx, + key, + start_item, + peer_is_remote, + transfer_us, + prepared_memory_info: None, + }); + } + (idx, _key, start_item, _peer_is_remote, get_id, _transfer_us, Err(err)) => { + results[idx] = Some(Err(err)); + transfer_error_cleanup.push(StartedGetRevokeCleanup { + get_id, + prepared_target: start_item.prepared_target, + }); + } + } + } + let lifecycle_transfer_wall_us = transfer_wall_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + + let transfer_cleanup_started_at = Instant::now(); + finish_started_get_revoke_cleanup( + self, + transfer_error_cleanup, + "batch_get transfer failure", + ) + .await; + let lifecycle_transfer_cleanup_us = transfer_cleanup_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + + let install_started_at = Instant::now(); + let mut ready_done_pending = Vec::with_capacity(done_pending.len()); + let mut install_failed_cleanup = Vec::new(); + for mut pending in done_pending { + let install_result = if let Some(target) = pending.start_item.prepared_target.as_ref() { + match u32::try_from(pending.start_item.len) { + Ok(len) => self + .install_hidden_pending_local_get( + &pending.key, + pending.start_item.get_id, + pending.start_item.put_id, + target.addr, + target.base_addr, + len, + target.slot_size, + target.grant_id, + target.slot_index, + ) + .map(Some), + Err(_) => Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "local-reserve Get value length exceeds u32: key={} len={}", + pending.key, pending.start_item.len + ), + })), + } + } else { + Ok(None) + }; + match install_result { + Ok(memory_info) => { + pending.prepared_memory_info = memory_info; + ready_done_pending.push(pending); + } + Err(err) => { + install_failed_cleanup.push(StartedGetRevokeCleanup { + get_id: pending.start_item.get_id, + prepared_target: pending.start_item.prepared_target.clone(), + }); + results[pending.idx] = Some(Err(err)); + } + } + } + done_pending = ready_done_pending; + + finish_started_get_revoke_cleanup( + self, + install_failed_cleanup, + "batch_get pending install failure", + ) + .await; + let lifecycle_install_us = install_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + + let done_get_ids = done_pending + .iter() + .map(|pending| pending.start_item.get_id) + .collect::>(); + let lifecycle_done_items = done_get_ids.len(); + let done_first_get_id = done_get_ids.first().copied(); + let done_last_get_id = done_get_ids.last().copied(); + let mut done_attempt = 1u32; + let done_started_at = Instant::now(); + let done_resp = loop { + match self.batch_get_done(done_get_ids.clone()).await { + Ok(resp) if batch_get_done_response_matches(&done_get_ids, &resp) => break resp, + Ok(resp) => { + tracing::warn!( + "batch_get_done response shape/identity mismatch; retaining pending-visible slots and retrying the same idempotent get_ids: items={} first_get_id={:?} last_get_id={:?} got_items={} got_first_get_id={:?} got_last_get_id={:?} attempt={}", + done_get_ids.len(), + done_first_get_id, + done_last_get_id, + resp.items.len(), + resp.items.first().map(|item| item.get_id), + resp.items.last().map(|item| item.get_id), + done_attempt + ); + } + Err(err) => { + if matches!(&err, KvError::Api(ApiError::SystemShutdown { .. })) { + return Err(err); + } + tracing::warn!( + "batch_get_done transport uncertain; retaining pending-visible slots and retrying the same idempotent get_ids: items={} first_get_id={:?} last_get_id={:?} attempt={} err={}", + done_get_ids.len(), + done_first_get_id, + done_last_get_id, + done_attempt, + err + ); + } + } + tokio::time::sleep(Duration::from_millis( + (10u64.saturating_mul(1u64 << done_attempt.min(8))).min(2_000), + )) + .await; + done_attempt = done_attempt.saturating_add(1); + }; + let lifecycle_done_us = done_started_at.elapsed().as_micros().min(i64::MAX as u128) as i64; + let publish_started_at = Instant::now(); + let master_node_id: NodeID = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await? + .into(); + let mut local_hot_admissions = Vec::new(); + + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + if let Some(target) = pending.start_item.prepared_target.as_ref() { + self.abort_hidden_pending_local_get(&pending.key, pending.start_item.get_id); + let canonical = self.local_committed_mem_holder_for_put_id( + &pending.key, + pending.start_item.put_id, + ); + if let Err(release_err) = release_prepared_get_target(self, target).await { + results[pending.idx] = Some(Err(release_err)); + continue; + } + // A same-version PutDone/GetDone may have won while this + // transfer was in flight. Converge on the owner's + // canonical local backing instead of turning an already + // available KV page into a prefix miss. + if let Some(memory_info) = canonical { + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info, + self.get_or_init_all_memholder_refcount(), + UserMemHolderExposeKind::SegPtr, + )); + results[pending.idx] = Some(Ok(Some((user_mem_holder, None)))); + continue; + } + } + results[pending.idx] = Some(Err(err)); + continue; + } + let expose_kind = if done_item.allocation_mode == GetAllocationMode::Temporary { + UserMemHolderExposeKind::OwnedCopy + } else { + UserMemHolderExposeKind::SegPtr + }; + let data_len = pending.start_item.len as usize; + metrics.record_l2_hit_locality(pending.peer_is_remote, data_len as u64); + metrics.record_get_io_locality( + pending.peer_is_remote, + data_len as u64, + pending.transfer_us, + ); + let memory_info = if pending.start_item.prepared_target.is_some() { + if done_item.allocation_mode != GetAllocationMode::LocalCommittedSlot { + self.abort_hidden_pending_local_get(&pending.key, pending.start_item.get_id); + results[pending.idx] = Some(Err(KvError::Api(ApiError::Unknown { + detail: format!( + "prepared local-reserve Get completed with unexpected allocation mode: key={} mode={:?}", + pending.key, done_item.allocation_mode + ), + }))); + continue; + } + let memory_info = match self.promote_hidden_pending_local_get( + &pending.key, + pending.start_item.get_id, + pending.start_item.put_id, + ) { + Ok(memory_info) => memory_info, + Err(err) => { + results[pending.idx] = Some(Err(err)); + continue; + } + }; + if let Some(prepared) = pending.prepared_memory_info.as_ref() { + assert!( + Arc::ptr_eq(prepared, &memory_info), + "Get promotion must retain the unique prepared MemoryInfo" + ); + } + local_hot_admissions.push(( + pending.key.clone(), + pending.start_item.put_id, memory_info.clone(), - self.get_or_init_all_memholder_refcount(), - UserMemHolderExposeKind::SegPtr, + pending.start_item.atomic_group.clone(), )); - obe_get_cache_hit( - &metrics, - &client_id, - &node_role, - key, - memory_info.len as u64, + memory_info + } else { + if done_item.allocation_mode == GetAllocationMode::LocalCommittedSlot { + results[pending.idx] = Some(Err(KvError::Api(ApiError::Unknown { + detail: format!( + "master returned local committed-slot mode without a prepared target: key={}", + pending.key + ), + }))); + continue; + } + let offset = pending.start_item.target_addr - pending.start_item.target_base_addr; + let memory_info = Arc::new( + MemoryInfo::new( + offset, + pending.start_item.len as u32, + done_item.holder_id, + pending.key.clone(), + master_node_id.clone(), + self.view.clone(), + ) + .await, ); - return Ok(Some((user_mem_holder, None))); + if done_item.allocation_mode != GetAllocationMode::Temporary + && self.install_get_cached_info_if_unfenced( + &pending.key, + pending.start_item.put_id, + memory_info.clone(), + ) + { + metrics.observe_cache_value_size( + &client_id, + node_role.as_str(), + data_len as u64, + ); + } + memory_info + }; + let get_info = RemoteGetInfo { + get_id: pending.start_item.get_id, + data_len, + src_addr: pending.start_item.src_addr, + target_addr: pending.start_item.target_addr, + node_id: pending.start_item.node_id.clone().into(), + peer_is_src_or_target: pending.peer_is_remote, + }; + if done_item.allocation_mode == GetAllocationMode::LocalCommittedSlot { + metrics.observe_cache_value_size(&client_id, node_role.as_str(), data_len as u64); } + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info, + self.get_or_init_all_memholder_refcount(), + expose_kind, + )); + results[pending.idx] = Some(Ok(Some((user_mem_holder, Some(get_info))))); } - let lock = self.get_remote_kv_lock.get_lock(key.to_owned()); - let _guard = lock.lock().await; + // Publish every local index before Moka admission starts selecting + // individual capacity victims. + let lifecycle_local_hot_admissions = local_hot_admissions.len(); + for (key, put_id, memory_info, _atomic_group) in local_hot_admissions { + self.owner_hot_track_committed(&key, put_id, &memory_info); + } - // Recheck after acquiring the miss lock so concurrent cache-fillers can collapse here - // without forcing every cache hit through the async lock path. - if let Some(cached_info) = self.get_cached_info.get(key) { - if let CachedValue::LocalReplica(memory_info) = &cached_info.value { + let output = results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "batch_get_finish_started result slot was not populated" + .to_string(), + })) + }) + }) + .collect::>(); + let lifecycle_hits = output + .iter() + .filter(|item| matches!(item, Ok(Some(_)))) + .count(); + let lifecycle_misses = output + .iter() + .filter(|item| matches!(item, Ok(None))) + .count(); + let lifecycle_errors = output.iter().filter(|item| matches!(item, Err(_))).count(); + let lifecycle_publish_us = publish_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + let lifecycle_total_us = lifecycle_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + tracing::info!( + "external Get finish lifecycle: requested={} transfer_concurrency={} zero_copy_items={} transfer_items={} remote_transfer_items={} local_ssd_items={} remote_ssd_items={} transfer_source_nodes={} ssd_source_nodes={} transfer_bytes={} remote_transfer_bytes={} plan_us={} transfer_wall_us={} transfer_sum_us={} transfer_max_us={} transfer_cleanup_us={} install_us={} done_items={} done_attempts={} done_us={} local_hot_admissions={} publish_us={} hits={} misses={} errors={} total_us={}", + lifecycle_requested_keys, + transfer_concurrency, + lifecycle_zero_copy_items, + lifecycle_transfer_items, + lifecycle_remote_transfer_items, + lifecycle_local_ssd_items, + lifecycle_remote_ssd_items, + lifecycle_transfer_source_nodes.len(), + lifecycle_ssd_source_nodes.len(), + lifecycle_transfer_bytes, + lifecycle_remote_transfer_bytes, + lifecycle_plan_us, + lifecycle_transfer_wall_us, + lifecycle_transfer_sum_us, + lifecycle_transfer_max_us, + lifecycle_transfer_cleanup_us, + lifecycle_install_us, + lifecycle_done_items, + done_attempt, + lifecycle_done_us, + lifecycle_local_hot_admissions, + lifecycle_publish_us, + lifecycle_hits, + lifecycle_misses, + lifecycle_errors, + lifecycle_total_us, + ); + Ok(output) + } + + pub async fn batch_get( + &self, + keys: Vec, + transfer_concurrency: usize, + ) -> KvResult, Option)>>>> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_get".to_string(), + })); + } + if keys.is_empty() { + return Ok(Vec::new()); + } + + #[derive(Clone)] + struct DonePending { + idx: usize, + key: String, + start_item: crate::master_kv_router::msg_pack::BatchGetStartItemResp, + peer_is_remote: bool, + transfer_us: i64, + } + + let transfer_concurrency = transfer_concurrency.max(1); + let metrics = self.metrics_handle(); + let client_id = self.client_id_str(); + let node_role = self.node_role(); + let self_node_id = self.view.cluster_manager().get_self_info().id.clone(); + + let mut results: Vec< + Option, Option)>>>, + > = (0..keys.len()).map(|_| None).collect(); + let mut missing_indices = Vec::new(); + let mut missing_keys = Vec::new(); + + for (idx, key) in keys.iter().enumerate() { + if let Some(memory_info) = self.local_visible_mem_holder(key) { tracing::debug!( - "cache hit after miss-lock for key: {} with putid({},{}), directly return", - key, - cached_info.put_time_ms, - cached_info.put_version + "batch_get local visible hit for key: {}, directly return", + key ); let user_mem_holder = Arc::new(UserMemHolder::new( memory_info.clone(), @@ -111,10 +818,388 @@ impl ClientKvApiInner { key, memory_info.len as u64, ); - return Ok(Some((user_mem_holder, None))); + metrics.record_get_io_locality(false, memory_info.len as u64, 0); + results[idx] = Some(Ok(Some((user_mem_holder, None)))); + } else { + obe_get_cache_miss(&metrics, &client_id, &node_role, key); + missing_indices.push(idx); + missing_keys.push(key.clone()); } } + if missing_keys.is_empty() { + return Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "batch_get result slot was not populated".to_string(), + })) + }) + }) + .collect()); + } + + let start_resp = self.batch_get_start(missing_keys.clone()).await?; + if start_resp.items.len() != missing_keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_get_start response length mismatch: expected={} got={}", + missing_keys.len(), + start_resp.items.len() + ), + })); + } + + let mut done_pending = Vec::new(); + let mut revoke_get_ids = Vec::new(); + let mut transfer_futures = Vec::new(); + + for ((idx, key), start_item) in missing_indices + .into_iter() + .zip(missing_keys.into_iter()) + .zip(start_resp.items.into_iter()) + { + if start_item.error_code == codes_api::API_KEY_NOT_FOUND { + results[idx] = Some(Ok(None)); + continue; + } + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + start_item.error_code, + start_item.error_json.clone(), + ) { + results[idx] = Some(Err(err)); + continue; + } + + let peer_id = if start_item.node_id == self_node_id { + None + } else { + Some(start_item.node_id.clone()) + }; + let peer_is_remote = peer_id.is_some(); + let get_id = start_item.get_id; + let src_addr = start_item.src_addr; + let target_addr = start_item.target_addr; + let len = start_item.len; + + if start_item.source_kind == GetSourceKind::Memory + && peer_id.is_none() + && src_addr == target_addr + { + done_pending.push(DonePending { + idx, + key, + start_item, + peer_is_remote, + transfer_us: 0, + }); + continue; + } + + transfer_futures.push(async move { + let transfer_started_at = Instant::now(); + let transfer_result = if start_item.source_kind == GetSourceKind::Ssd { + self.stage_kv_from_ssd_source( + &start_item.node_id, + &key, + start_item.put_id, + get_id, + src_addr, + len, + target_addr, + len, + ) + .await + } else { + self.view + .client_transfer_engine() + .transfer_data_no_copy(peer_id, true, src_addr, target_addr, len, None) + .await + .map(|_| ()) + .map_err(|err| { + KvError::Api(ApiError::Transfer { + from_addr: src_addr, + to_addr: target_addr, + len, + error: err.to_string(), + }) + }) + }; + let transfer_us = transfer_started_at + .elapsed() + .as_micros() + .min(i64::MAX as u128) as i64; + ( + idx, + key, + start_item, + peer_is_remote, + get_id, + transfer_us, + transfer_result, + ) + }); + } + + let mut transfer_stream = + stream::iter(transfer_futures).buffer_unordered(transfer_concurrency); + while let Some(joined) = transfer_stream.next().await { + match joined { + (idx, key, start_item, peer_is_remote, _get_id, transfer_us, Ok(_breakdown)) => { + done_pending.push(DonePending { + idx, + key, + start_item, + peer_is_remote, + transfer_us, + }); + } + (idx, _key, _start_item, _peer_is_remote, get_id, _transfer_us, Err(err)) => { + results[idx] = Some(Err(err)); + revoke_get_ids.push(get_id); + } + } + } + + if !revoke_get_ids.is_empty() { + if let Err(err) = self.batch_get_revoke(revoke_get_ids).await { + tracing::warn!("batch_get_revoke failed after transfer errors: {}", err); + } + } + + let done_resp = self + .batch_get_done( + done_pending + .iter() + .map(|pending| pending.start_item.get_id) + .collect(), + ) + .await?; + if done_resp.items.len() != done_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_get_done response length mismatch: expected={} got={}", + done_pending.len(), + done_resp.items.len() + ), + })); + } + let master_node_id: NodeID = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await? + .into(); + + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + results[pending.idx] = Some(Err(err)); + continue; + } + let expose_kind = if done_item.allocation_mode == GetAllocationMode::Temporary { + UserMemHolderExposeKind::OwnedCopy + } else { + UserMemHolderExposeKind::SegPtr + }; + let offset = pending.start_item.target_addr - pending.start_item.target_base_addr; + let data_len = pending.start_item.len as usize; + metrics.record_l2_hit_locality(pending.peer_is_remote, data_len as u64); + metrics.record_get_io_locality( + pending.peer_is_remote, + data_len as u64, + pending.transfer_us, + ); + let memory_info = Arc::new( + MemoryInfo::new( + offset, + pending.start_item.len as u32, + done_item.holder_id, + pending.key.clone(), + master_node_id.clone(), + self.view.clone(), + ) + .await, + ); + let get_info = RemoteGetInfo { + get_id: pending.start_item.get_id, + data_len, + src_addr: pending.start_item.src_addr, + target_addr: pending.start_item.target_addr, + node_id: pending.start_item.node_id.clone().into(), + peer_is_src_or_target: pending.peer_is_remote, + }; + if done_item.allocation_mode != GetAllocationMode::Temporary { + if self.install_get_cached_info_if_unfenced( + &pending.key, + pending.start_item.put_id, + memory_info.clone(), + ) { + metrics.observe_cache_value_size( + &client_id, + node_role.as_str(), + data_len as u64, + ); + } + } + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info, + self.get_or_init_all_memholder_refcount(), + expose_kind, + )); + results[pending.idx] = Some(Ok(Some((user_mem_holder, Some(get_info))))); + } + + Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "batch_get result slot was not populated".to_string(), + })) + }) + }) + .collect()) + } + + pub async fn batch_is_exist( + &self, + keys: Vec, + allow_local_snapshot: bool, + ) -> KvResult> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_is_exist".to_string(), + })); + } + if keys.is_empty() { + return Ok(Vec::new()); + } + + let mut results = vec![false; keys.len()]; + let mut missing_indices = Vec::new(); + let mut missing_keys = Vec::new(); + for (idx, key) in keys.iter().enumerate() { + if allow_local_snapshot && self.has_local_snapshot(key) { + results[idx] = true; + continue; + } + missing_indices.push(idx); + missing_keys.push(key.clone()); + } + if missing_keys.is_empty() { + return Ok(results); + } + + let req = MsgPack { + serialize_part: BatchIsExistReq { keys: missing_keys }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_is_exist + .call(self.view.p2p_module(), master_node_id.into(), req, None, 0) + .await + .map_err(KvError::from)?; + let resp_part = resp.serialize_part; + crate::rpcresp_kvresult_convert::try_from_code( + resp_part.error_code, + resp_part.error_json.clone(), + )?; + if resp_part.exists_list.len() != missing_indices.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_is_exist response length mismatch: expected={} got={}", + missing_indices.len(), + resp_part.exists_list.len() + ), + })); + } + for (idx, exists) in missing_indices + .into_iter() + .zip(resp_part.exists_list.into_iter()) + { + results[idx] = exists; + } + Ok(results) + } + + pub async fn is_exist_with_local_snapshot( + &self, + key: &str, + allow_local_snapshot: bool, + ) -> KvResult { + let mut results = self + .batch_is_exist(vec![key.to_string()], allow_local_snapshot) + .await?; + Ok(results.pop().unwrap_or(false)) + } + + /// becaused we cached local kv metadata, so we make `MemHolder` with Arc here + pub async fn get( + &self, + key: &str, + ) -> KvResult, Option)>> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting get".to_string(), + })); + } + let metrics = self.metrics_handle(); + let client_id = self.client_id_str(); + let node_role = self.node_role(); + + if let Some(memory_info) = self.local_visible_mem_holder(key) { + // exist, directly return + tracing::debug!("local visible cache hit for key: {}, directly return", key); + // Build a fresh UserMemHolder from cached MemoryInfo + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info.clone(), + self.get_or_init_all_memholder_refcount(), + UserMemHolderExposeKind::SegPtr, + )); + obe_get_cache_hit( + &metrics, + &client_id, + &node_role, + key, + memory_info.len as u64, + ); + return Ok(Some((user_mem_holder, None))); + } + + let lock = self.get_remote_kv_lock.get_lock(key.to_owned()); + let _guard = lock.lock().await; + + // Recheck after acquiring the miss lock so concurrent cache-fillers can collapse here + // without forcing every cache hit through the async lock path. + if let Some(memory_info) = self.local_visible_mem_holder(key) { + tracing::debug!( + "local visible cache hit after miss-lock for key: {}, directly return", + key + ); + let user_mem_holder = Arc::new(UserMemHolder::new( + memory_info.clone(), + self.get_or_init_all_memholder_refcount(), + UserMemHolderExposeKind::SegPtr, + )); + obe_get_cache_hit( + &metrics, + &client_id, + &node_role, + key, + memory_info.len as u64, + ); + return Ok(Some((user_mem_holder, None))); + } + obe_get_cache_miss(&metrics, &client_id, &node_role, key); let t1 = Utc::now().timestamp_micros(); let resp = { @@ -177,18 +1262,26 @@ impl ClientKvApiInner { ); } - // transfer data (skip if local and src==target to avoid redundant copy) - if peer_id.is_none() && abs_src == abs_target { + let transfer_result = if resp.source_kind == GetSourceKind::Ssd { + self.stage_kv_from_ssd_source( + &resp.node_id, + key, + put_id, + get_id, + abs_src, + data_len as u64, + abs_target, + data_len as u64, + ) + .await + } else if peer_id.is_none() && abs_src == abs_target { tracing::debug!( "kv get local no-op: src==target {:#x}, len={} (skip transfer)", abs_target, data_len ); + Ok(()) } else { - // tracing::debug!( - // "kv get transfer in transfer engine path from {}", - // peer_id.as_ref().map(|v| &**v).unwrap_or("self") - // ); tracing::debug!( "p2p get transfer: key={}, remote_src={:#x} -> local_target={:#x}, len={}, peer={:?}", key, @@ -197,8 +1290,7 @@ impl ClientKvApiInner { data_len, peer_id ); - if let Err(e) = self - .view + self.view .client_transfer_engine() .transfer_data_no_copy( peer_id.clone(), @@ -209,33 +1301,32 @@ impl ClientKvApiInner { None, ) .await - { - tracing::warn!("transfer data failed: {:?}", e); - - #[cfg(test)] - { - self.test_record.remove_transfering_get(get_id); - } - - obe_get_transfer_error(&metrics, &client_id, &node_role, key, data_len as u64); - self.get_revoke(get_id).await?; - return Err(KvError::Api(ApiError::Transfer { - from_addr: abs_src, - to_addr: abs_target, - len: data_len as u64, - error: e.to_string(), - })); - } else { - tracing::debug!( - "get_transfer success key={}, src_addr={:#x}, target_addr={:#x}, len={}, peer_id={:?}", - key, - abs_src, - abs_target, - data_len, - peer_id - ); - } + .map(|_| ()) + .map_err(|err| { + KvError::Api(ApiError::Transfer { + from_addr: abs_src, + to_addr: abs_target, + len: data_len as u64, + error: err.to_string(), + }) + }) + }; + if let Err(err) = transfer_result { + tracing::warn!("transfer data failed: {:?}", err); + #[cfg(test)] + self.test_record.remove_transfering_get(get_id); + obe_get_transfer_error(&metrics, &client_id, &node_role, key, data_len as u64); + self.get_revoke(get_id).await?; + return Err(err); } + tracing::debug!( + "get_transfer success key={}, src_addr={:#x}, target_addr={:#x}, len={}, peer_id={:?}", + key, + abs_src, + abs_target, + data_len, + peer_id + ); let t3 = Utc::now().timestamp_micros(); obe_get_transfer_success( &metrics, @@ -333,8 +1424,9 @@ impl ClientKvApiInner { }; if done_resp.allocation_mode != GetAllocationMode::Temporary { - self.cache_local_replica_after_get(key, put_id, memory_info.clone()); - metrics.observe_cache_value_size(&client_id, node_role.as_str(), data_len as u64); + if self.install_get_cached_info_if_unfenced(key, put_id, memory_info.clone()) { + metrics.observe_cache_value_size(&client_id, node_role.as_str(), data_len as u64); + } } let user_mem_holder = Arc::new(UserMemHolder::new( memory_info, @@ -347,28 +1439,7 @@ impl ClientKvApiInner { } pub async fn is_exist(&self, key: &str) -> KvResult { - if !self.view.register_shutdown_poller().is_running() { - return Err(KvError::Api(ApiError::SystemShutdown { - detail: "ClientKvApi is shutting down; rejecting is_exist".to_string(), - })); - } - let resp = self.get_meta(key).await?; - if resp.error_code != OK { - // If error code indicates key not found, return false - if resp.error_code - == crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND - { - return Ok(false); - } - // For other errors, propagate the error - crate::rpcresp_kvresult_convert::try_from_code( - resp.error_code, - resp.error_json.clone(), - )?; - unreachable!("error path should have returned above"); - } - - Ok(resp.exists) + self.is_exist_with_local_snapshot(key, false).await } /// Get metadata for a key without transferring data @@ -412,6 +1483,8 @@ impl ClientKvApiInner { let req = MsgPack { serialize_part: GetStartReq { key: key.to_string(), + prepared_target: None, + external_sink_target: None, }, raw_bytes: Vec::new(), }; @@ -433,6 +1506,94 @@ impl ClientKvApiInner { Ok(resp.serialize_part) } + pub async fn batch_get_start(&self, keys: Vec) -> KvResult { + self.batch_get_start_with_prepared_targets(keys, Vec::new()) + .await + } + + pub(crate) async fn batch_get_start_with_prepared_targets( + &self, + keys: Vec, + prepared_targets: Vec>, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_get_start".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchGetStartReq { + keys, + prepared_targets, + external_sink_targets: Vec::new(), + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_get_start + .call(self.view.p2p_module(), master_node_id.into(), req, None, 0) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub(crate) async fn batch_get_bind_prepared_targets( + &self, + get_ids: Vec, + prepared_targets: Vec, + ) -> KvResult { + if get_ids.len() != prepared_targets.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch_get_bind_prepared_targets length mismatch: get_ids={} targets={}", + get_ids.len(), + prepared_targets.len() + ), + })); + } + let items = get_ids + .into_iter() + .zip(prepared_targets) + .map(|(get_id, target)| BatchGetBindItemReq { + get_id, + target: GetBindTarget::PreparedLocalReserve(target), + }) + .collect(); + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_get_bind + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetBindReq { items }, + raw_bytes: Vec::new(), + }, + None, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + /// 撤销 Get 操作，释放已分配的资源 pub async fn get_revoke(&self, get_id: u64) -> KvResult<()> { let req = MsgPack { @@ -457,6 +1618,40 @@ impl ClientKvApiInner { Ok(()) } + pub async fn batch_get_revoke(&self, get_ids: Vec) -> KvResult { + if get_ids.is_empty() { + return Ok(BatchGetRevokeResp { + items: Vec::new(), + error_code: OK, + error_json: String::new(), + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_get_revoke".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchGetRevokeReq { get_ids }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_get_revoke + .call(self.view.p2p_module(), master_node_id.into(), req, None, 0) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + /// 完成 Get 操作，清理资源 pub async fn get_done(&self, get_id: u64) -> KvResult { let req = MsgPack { @@ -480,4 +1675,83 @@ impl ClientKvApiInner { Ok(resp.serialize_part) } + + pub async fn batch_get_done(&self, get_ids: Vec) -> KvResult { + if get_ids.is_empty() { + return Ok(BatchGetDoneResp { + items: Vec::new(), + error_code: OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_get_done".to_string(), + })); + } + // The master performs synchronous Moka policy work before acknowledging + // committed-slot Done. Bound the number of callers entering that path + // so a capacity scan cannot park every master Tokio worker on Moka's + // blocking housekeeper lock. The permit covers one RPC attempt only; + // idempotent retry backoff releases it and lets other atomic_batches converge. + let _done_rpc_permit = batch_get_done_rpc_limiter() + .acquire() + .await + .expect("the process-wide BatchGetDone limiter is never closed"); + let req = MsgPack { + serialize_part: BatchGetDoneReq { get_ids }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_get_done + .call(self.view.p2p_module(), master_node_id.into(), req, None, 2) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } +} + +#[cfg(test)] +mod tests { + use super::batch_get_done_response_matches; + use crate::master_kv_router::msg_pack::{BatchGetDoneItemResp, BatchGetDoneResp}; + + fn response(get_ids: &[u64]) -> BatchGetDoneResp { + BatchGetDoneResp { + items: get_ids + .iter() + .map(|get_id| BatchGetDoneItemResp { + get_id: *get_id, + ..Default::default() + }) + .collect(), + ..Default::default() + } + } + + #[test] + fn batch_get_done_requires_exact_response_identity() { + assert!(batch_get_done_response_matches( + &[11, 22], + &response(&[11, 22]) + )); + assert!(!batch_get_done_response_matches( + &[11, 22], + &response(&[22, 11]) + )); + assert!(!batch_get_done_response_matches( + &[11, 22], + &response(&[11]) + )); + } } diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/local_reserve_rebalance.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/local_reserve_rebalance.rs new file mode 100644 index 0000000..5fec74f --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/local_reserve_rebalance.rs @@ -0,0 +1,864 @@ +use super::{ + ClientKvApiInner, ClientKvApiView, OwnerHotEvictionDispatch, OwnerLocalReserveClassState, + OwnerLocalReservePoolState, +}; +use crate::OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES; +use crate::master_kv_router::msg_pack::ReserveLocalGrantOutcome; +use limit_thirdparty::tokio; +use std::time::{Duration, Instant}; + +const OWNER_LOCAL_RESERVE_REBALANCE_INTERVAL: Duration = Duration::from_millis(200); +const OWNER_LOCAL_RESERVE_SHRINK_IDLE_COOLDOWN: Duration = Duration::from_secs(5); +const OWNER_LOCAL_RESERVE_SHRINK_GROW_COOLDOWN: Duration = Duration::from_secs(1); +const OWNER_LOCAL_RESERVE_DEFAULT_SOFT_WAIT_TIMEOUT: Duration = Duration::from_millis(10); +// AtomicBatch-safe master route deletion plus owner slot release is asynchronous. +// Thirty seconds turns a recoverable pressure burst into a storage exception that +// can wedge SGLang's detokenizer. This remains well below the external request +// timeout while allowing bounded backpressure to finish reclaiming slots. +const OWNER_LOCAL_RESERVE_DEFAULT_HARD_TIMEOUT: Duration = Duration::from_secs(30); +const OWNER_LOCAL_RESERVE_MIN_GRANTS_PER_CLASS: usize = 0; +const OWNER_LOCAL_RESERVE_REPORT_INTERVAL: Duration = Duration::from_secs(30); +const OWNER_LOCAL_RESERVE_FREE_SLOT_HEADROOM_GRANTS: usize = 4; +const OWNER_SLOT_PRESSURE_INTERVAL: Duration = Duration::from_millis(10); +// Source deletion now has an exact, idempotent owner -> master transaction and +// selection debt accounts for candidates until their slots are physically Free. +// Keep pressure retries responsive: a 200 ms gate used to surface directly as +// local_fast_put_start tail latency whenever one bounded kick was insufficient. +const OWNER_SLOT_PRESSURE_MIN_KICK_INTERVAL: Duration = Duration::from_millis(25); +// One 512 MiB grant contains 113 production KV slots. The old 256 MiB cap was +// smaller than a single grant and could never refill the configured four-grant +// high watermark in one pass. Bound a kick at eight grants so it covers the +// four-grant retained headroom plus a normal large claimant, while selection +// debt still prevents repeated over-selection before physical reclaim lands. +const OWNER_SLOT_PRESSURE_MAX_EVICT_BYTES: u64 = OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES * 8; + +#[derive(Debug, Clone, Copy)] +struct ExpectedCapacityLayout { + value_len: u64, + payload_capacity_bytes: u64, + slot_size: u64, + slots_per_grant: u32, + grant_count: usize, +} + +#[derive(Debug, Clone, Copy)] +struct RebalanceClassSnapshot { + slot_size: u64, + slots_per_grant: u32, + used_slots: usize, + free_slots: usize, + grant_count: usize, + pending_slot_demand: usize, + max_observed_claim_slots: usize, + expected_grant_count: usize, +} + +fn owner_local_reserve_target_grant_count( + slots_per_grant: u32, + used_slots: usize, + pending_slot_demand: usize, + expected_grant_count: usize, +) -> usize { + assert!( + slots_per_grant > 0, + "local-reserve grant must contain slots" + ); + let required_slots = used_slots.saturating_add(pending_slot_demand); + let slots_per_grant = slots_per_grant as usize; + let demand_grants = + required_slots / slots_per_grant + usize::from(required_slots % slots_per_grant != 0); + if expected_grant_count == 0 { + demand_grants + } else { + // A configured expected capacity is the owner's physical slot budget, + // not a minimum that demand may grow past. Once all expected grants + // are installed, additional demand must be served by owner-local + // source eviction/reclaim; asking the master for another grant silently + // changes the configured owner capacity and can starve the remote + // allocation domain. + expected_grant_count + } +} + +fn owner_slot_pressure_request_bytes( + snapshot: RebalanceClassSnapshot, + source_eviction_selected_bytes: u64, +) -> u64 { + let low_watermark = (snapshot.slots_per_grant as usize) + .saturating_mul(2) + .max(snapshot.max_observed_claim_slots); + let high_watermark = (snapshot.slots_per_grant as usize) + .saturating_mul(OWNER_LOCAL_RESERVE_FREE_SLOT_HEADROOM_GRANTS); + let projected_eviction_slots = + usize::try_from(source_eviction_selected_bytes / snapshot.slot_size.max(1)) + .unwrap_or(usize::MAX); + let projected_free_slots = snapshot.free_slots.saturating_add(projected_eviction_slots); + let pressure_active = + snapshot.pending_slot_demand > snapshot.free_slots || projected_free_slots < low_watermark; + if !pressure_active { + return 0; + } + let desired_free_slots = snapshot.pending_slot_demand.saturating_add(high_watermark); + let selection_slots = desired_free_slots.saturating_sub(projected_free_slots); + u64::try_from(selection_slots) + .unwrap_or(u64::MAX) + .saturating_mul(snapshot.slot_size) + .min(OWNER_SLOT_PRESSURE_MAX_EVICT_BYTES) +} + +fn owner_slot_pressure_projected_reclaim_bytes(inner: &ClientKvApiInner) -> u64 { + // Candidate debt starts in Moka's synchronous eviction listener. A + // candidate that is stale, pinned, or waiting in the retry queue has + // no installed owner fence and cannot release a slot. Only exact selected + // source debt is valid projected reclaim credit. + inner + .owner_hot_counters + .source_eviction_selected_bytes + .load(std::sync::atomic::Ordering::Acquire) +} + +fn owner_local_reserve_desired_free_slots(snapshot: RebalanceClassSnapshot) -> usize { + let configured_headroom = if snapshot.expected_grant_count == 0 { + 0 + } else { + (snapshot.slots_per_grant as usize) + .saturating_mul(OWNER_LOCAL_RESERVE_FREE_SLOT_HEADROOM_GRANTS) + }; + snapshot + .pending_slot_demand + .saturating_add(configured_headroom) +} + +fn configured_expected_capacity(inner: &ClientKvApiInner) -> Option { + let expected = inner + .test_spec_config + .owner_local_reserve_expected_capacity + .as_ref()?; + let slot_size = crate::owner_local_reserve_slot_size_bytes(expected.value_len) + .expect("owner local-reserve expected capacity must be config-validated"); + let slots_per_grant = crate::owner_local_reserve_slots_per_grant(slot_size) + .expect("validated local-reserve slot size must fit in a grant"); + let grant_count = usize::try_from( + crate::owner_local_reserve_expected_grant_count( + expected.value_len, + expected.payload_capacity_bytes, + ) + .expect("owner local-reserve expected capacity must be config-validated"), + ) + .expect("owner local-reserve expected grant count must fit usize"); + Some(ExpectedCapacityLayout { + value_len: expected.value_len, + payload_capacity_bytes: expected.payload_capacity_bytes, + slot_size, + slots_per_grant, + grant_count, + }) +} + +fn install_expected_capacity_class(inner: &ClientKvApiInner) -> Option { + let expected = configured_expected_capacity(inner)?; + let mut pool = inner.owner_local_reserve_pool.lock(); + let class_state = pool.classes.entry(expected.slot_size).or_insert_with(|| { + OwnerLocalReserveClassState::new(expected.slot_size, expected.slots_per_grant) + }); + assert_eq!(class_state.slots_per_grant, expected.slots_per_grant); + class_state.expected_grant_count = expected.grant_count; + Some(expected) +} + +fn snapshot_rebalance_classes(pool: &OwnerLocalReservePoolState) -> Vec { + pool.classes + .values() + .map(|class_state| RebalanceClassSnapshot { + slot_size: class_state.slot_size, + slots_per_grant: class_state.slots_per_grant, + used_slots: class_state.used_slot_count(), + free_slots: class_state.free_slot_count(), + grant_count: class_state.grant_count(), + pending_slot_demand: class_state.pending_slot_demand, + max_observed_claim_slots: class_state.max_observed_claim_slots, + expected_grant_count: class_state.expected_grant_count, + }) + .collect() +} + +pub(crate) fn owner_local_reserve_timeout_config(inner: &ClientKvApiInner) -> (Duration, Duration) { + let soft_wait_timeout = inner + .test_spec_config + .owner_local_reserve_soft_wait_timeout_ms + .map(Duration::from_millis) + .unwrap_or(OWNER_LOCAL_RESERVE_DEFAULT_SOFT_WAIT_TIMEOUT); + let hard_timeout = inner + .test_spec_config + .owner_local_reserve_hard_timeout_ms + .map(Duration::from_millis) + .unwrap_or(OWNER_LOCAL_RESERVE_DEFAULT_HARD_TIMEOUT); + (soft_wait_timeout, hard_timeout) +} + +async fn try_refill_once(inner: &ClientKvApiInner, snapshot: RebalanceClassSnapshot) { + if snapshot.slot_size > OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES { + tracing::error!( + "owner local reserve refill rejected oversized slot_size={} quantum={} pending_slots={}", + snapshot.slot_size, + OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES, + snapshot.pending_slot_demand + ); + return; + } + let desired_free_slots = owner_local_reserve_desired_free_slots(snapshot); + let target_grants = owner_local_reserve_target_grant_count( + snapshot.slots_per_grant, + snapshot.used_slots, + desired_free_slots, + snapshot.expected_grant_count, + ); + if snapshot.grant_count >= target_grants { + return; + } + let additional_grants = target_grants - snapshot.grant_count; + let batch_started_at = Instant::now(); + let mut added_grants = 0usize; + let mut total_rpc_latency = Duration::ZERO; + let mut max_rpc_latency = Duration::ZERO; + for _ in 0..additional_grants { + let rpc_started_at = Instant::now(); + let outcome = match inner.reserve_local_grant().await { + Ok(outcome) => { + let rpc_latency = rpc_started_at.elapsed(); + total_rpc_latency = total_rpc_latency.saturating_add(rpc_latency); + max_rpc_latency = max_rpc_latency.max(rpc_latency); + outcome + } + Err(err) => { + let rpc_latency = rpc_started_at.elapsed(); + tracing::warn!( + "owner local reserve refill failed: slot_size={} pending_slots={} grants={} target_grants={} expected_grants={} added_grants={} rpc_latency_ms={} err={}", + snapshot.slot_size, + snapshot.pending_slot_demand, + snapshot.grant_count, + target_grants, + snapshot.expected_grant_count, + added_grants, + rpc_latency.as_millis(), + err + ); + return; + } + }; + let (grant_id, base_addr, addr, len) = match outcome { + ReserveLocalGrantOutcome::Granted { + grant_id, + node_id: _, + addr, + base_addr, + len, + } => (grant_id, base_addr, addr, len), + ReserveLocalGrantOutcome::None => { + unreachable!("reserve_local_grant filters empty successful outcomes") + } + }; + if len != OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES { + tracing::warn!( + "owner local reserve refill got unexpected grant size: requested={} got={}", + OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES, + len + ); + if let Err(release_err) = inner.release_local_grant(grant_id).await { + tracing::warn!( + "owner local reserve refill failed to release mismatched grant_id={} err={}", + grant_id, + release_err + ); + } + return; + } + let detached_grant = super::OwnerLocalReserveGrantState::new( + grant_id, + base_addr, + addr, + len, + snapshot.slot_size, + snapshot.slots_per_grant, + ); + { + let mut pool = inner.owner_local_reserve_pool.lock(); + let class_state = pool.classes.entry(snapshot.slot_size).or_insert_with(|| { + OwnerLocalReserveClassState::new(snapshot.slot_size, snapshot.slots_per_grant) + }); + class_state.last_grow_at = Some(Instant::now()); + class_state.install_grant(detached_grant); + } + added_grants += 1; + inner + .owner_local_reserve_rebalance_notify() + .notify_waiters(); + } + if added_grants != 0 { + let average_rpc_latency_ms = total_rpc_latency.as_secs_f64() * 1000.0 / added_grants as f64; + tracing::info!( + "owner local reserve refill batch completed: slot_size={} added_grants={} grants_after={} target_grants={} expected_grants={} pending_slots={} elapsed_ms={} rpc_avg_ms={:.3} rpc_max_ms={}", + snapshot.slot_size, + added_grants, + snapshot.grant_count.saturating_add(added_grants), + target_grants, + snapshot.expected_grant_count, + snapshot.pending_slot_demand, + batch_started_at.elapsed().as_millis(), + average_rpc_latency_ms, + max_rpc_latency.as_millis() + ); + } +} + +fn detach_excess_fully_free_grant( + class_state: &mut OwnerLocalReserveClassState, +) -> Option { + let grow_cooldown_ok = class_state + .last_grow_at + .map(|last_grow_at| last_grow_at.elapsed() >= OWNER_LOCAL_RESERVE_SHRINK_GROW_COOLDOWN) + .unwrap_or(true); + if !grow_cooldown_ok { + return None; + } + let shrink_keep_grants = owner_local_reserve_target_grant_count( + class_state.slots_per_grant, + class_state.used_slot_count(), + class_state.pending_slot_demand, + class_state.expected_grant_count, + ) + .max(OWNER_LOCAL_RESERVE_MIN_GRANTS_PER_CLASS); + if class_state.grant_count() <= shrink_keep_grants { + return None; + } + + let candidate_grant_id = class_state.grants.last().and_then(|grant| { + (grant.is_fully_free() + && grant + .fully_free_since + .map(|since| since.elapsed() >= OWNER_LOCAL_RESERVE_SHRINK_IDLE_COOLDOWN) + .unwrap_or(false)) + .then_some(grant.grant_id) + })?; + class_state.detach_fully_free_grant(candidate_grant_id) +} + +async fn try_shrink_once(inner: &ClientKvApiInner, snapshot: RebalanceClassSnapshot) { + let detached = { + let mut pool = inner.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&snapshot.slot_size) else { + return; + }; + detach_excess_fully_free_grant(class_state) + }; + + let Some(grant) = detached else { + return; + }; + + if let Err(err) = inner.release_local_grant(grant.grant_id).await { + tracing::warn!( + "owner local reserve shrink failed to release grant_id={} err={}", + grant.grant_id, + err + ); + let mut pool = inner.owner_local_reserve_pool.lock(); + let class_state = pool.classes.entry(snapshot.slot_size).or_insert_with(|| { + OwnerLocalReserveClassState::new(snapshot.slot_size, snapshot.slots_per_grant) + }); + class_state.install_grant(grant); + } else { + tracing::info!( + "owner local reserve released excess grant: slot_size={} grant_id={} pending_slots={}", + snapshot.slot_size, + grant.grant_id, + snapshot.pending_slot_demand + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client_kv_api::OwnerHotCacheCounters; + use std::sync::atomic::Ordering; + + const SLOT_SIZE: u64 = 8 * 1024 * 1024; + + #[test] + fn target_grants_cover_used_and_pending_without_double_reserving() { + let slots_per_grant = crate::owner_local_reserve_slots_per_grant(SLOT_SIZE).unwrap(); + assert_eq!( + owner_local_reserve_target_grant_count(slots_per_grant, 128, 1, 0), + 3 + ); + assert_eq!( + owner_local_reserve_target_grant_count(slots_per_grant, 128, 0, 0), + 2 + ); + assert_eq!( + owner_local_reserve_target_grant_count(slots_per_grant, 0, 0, 0), + 0 + ); + } + + #[test] + fn exact_fit_target_accounts_for_unusable_grant_tail() { + const VALUE_LEN: u64 = 4_718_592; + const EXPECTED_PAYLOAD_BYTES: u64 = 109_951_162_777; + let slot_size = crate::owner_local_reserve_slot_size_bytes(VALUE_LEN).unwrap(); + let slots_per_grant = crate::owner_local_reserve_slots_per_grant(slot_size).unwrap(); + let expected_grants = + crate::owner_local_reserve_expected_grant_count(VALUE_LEN, EXPECTED_PAYLOAD_BYTES) + .unwrap(); + assert_eq!(slot_size, VALUE_LEN); + assert_eq!(slots_per_grant, 113); + assert_eq!(expected_grants, 207); + assert_eq!( + owner_local_reserve_target_grant_count(slots_per_grant, 0, 0, expected_grants as usize,), + 207 + ); + } + + #[test] + fn default_hard_timeout_covers_owner_source_eviction_transaction() { + assert_eq!( + OWNER_LOCAL_RESERVE_DEFAULT_HARD_TIMEOUT, + Duration::from_secs(30) + ); + assert!( + OWNER_LOCAL_RESERVE_DEFAULT_HARD_TIMEOUT > OWNER_LOCAL_RESERVE_SHRINK_IDLE_COOLDOWN + ); + } + + #[test] + fn configured_pool_refills_before_pending_claims_consume_headroom() { + let snapshot = RebalanceClassSnapshot { + slot_size: SLOT_SIZE, + slots_per_grant: 64, + used_slots: 64 * 228, + free_slots: 64 * 4, + grant_count: 232, + pending_slot_demand: 0, + max_observed_claim_slots: 0, + expected_grant_count: 232, + }; + assert_eq!(owner_local_reserve_desired_free_slots(snapshot), 64 * 4); + assert_eq!( + owner_local_reserve_target_grant_count( + snapshot.slots_per_grant, + snapshot.used_slots, + owner_local_reserve_desired_free_slots(snapshot), + snapshot.expected_grant_count, + ), + 232 + ); + + let pressured = RebalanceClassSnapshot { + used_slots: snapshot.used_slots + 1, + free_slots: snapshot.free_slots - 1, + ..snapshot + }; + assert_eq!( + owner_local_reserve_target_grant_count( + pressured.slots_per_grant, + pressured.used_slots, + owner_local_reserve_desired_free_slots(pressured), + pressured.expected_grant_count, + ), + 232 + ); + } + + #[test] + fn configured_expected_grants_are_a_hard_upper_bound() { + assert_eq!( + owner_local_reserve_target_grant_count(64, 64 * 300, 64, 232), + 232 + ); + assert_eq!( + owner_local_reserve_target_grant_count(64, 64 * 300, 64, 0), + 301 + ); + } + + #[test] + fn owner_slot_pressure_targets_pending_plus_high_watermark() { + let snapshot = RebalanceClassSnapshot { + slot_size: 1024 * 1024, + slots_per_grant: 4, + used_slots: 4, + free_slots: 1, + grant_count: 1, + pending_slot_demand: 3, + max_observed_claim_slots: 3, + expected_grant_count: 1, + }; + assert_eq!( + owner_slot_pressure_request_bytes(snapshot, 0), + 18 * 1024 * 1024, + "3 pending slots plus 16 slots of retained headroom minus one physical Free" + ); + assert_eq!( + owner_slot_pressure_request_bytes(snapshot, 10 * 1024 * 1024), + 8 * 1024 * 1024, + "exact selected debt contributes projected slots of the same class" + ); + assert_eq!( + owner_slot_pressure_request_bytes( + RebalanceClassSnapshot { + free_slots: 16, + pending_slot_demand: 0, + ..snapshot + }, + 0, + ), + 0, + "pressure stops after the high watermark is physically available" + ); + + let clamped = RebalanceClassSnapshot { + slot_size: 8 * 1024 * 1024, + slots_per_grant: 64, + used_slots: 64, + free_slots: 1, + grant_count: 1, + pending_slot_demand: 1_000, + max_observed_claim_slots: 1_000, + expected_grant_count: 1, + }; + assert_eq!( + owner_slot_pressure_request_bytes(clamped, 0), + OWNER_SLOT_PRESSURE_MAX_EVICT_BYTES + ); + } + + #[test] + fn pre_fence_candidate_debt_is_not_projected_reclaim_credit() { + let snapshot = RebalanceClassSnapshot { + slot_size: 1024 * 1024, + slots_per_grant: 4, + used_slots: 4, + free_slots: 1, + grant_count: 1, + pending_slot_demand: 3, + max_observed_claim_slots: 3, + expected_grant_count: 1, + }; + let counters = OwnerHotCacheCounters::default(); + counters + .selection_debt_bytes + .store(10 * 1024 * 1024, Ordering::Release); + + assert_eq!( + counters + .source_eviction_selected_bytes + .load(Ordering::Acquire), + 0 + ); + assert_eq!( + owner_slot_pressure_request_bytes( + snapshot, + counters + .source_eviction_selected_bytes + .load(Ordering::Acquire), + ), + 18 * 1024 * 1024, + "retry-only candidate debt must not suppress physical victim selection" + ); + + counters.add_source_eviction_selected_bytes(10 * 1024 * 1024); + assert_eq!( + owner_slot_pressure_request_bytes( + snapshot, + counters + .source_eviction_selected_bytes + .load(Ordering::Acquire), + ), + 8 * 1024 * 1024, + "only installed source-selection debt is valid projected reclaim credit" + ); + counters.remove_source_eviction_selected_bytes(10 * 1024 * 1024); + assert_eq!( + counters + .source_eviction_selected_bytes + .load(Ordering::Acquire), + 0 + ); + } + + #[test] + fn production_pressure_fills_high_watermark_in_one_kick() { + const VALUE_LEN: u64 = 4_718_592; + let slots_per_grant = crate::owner_local_reserve_slots_per_grant(VALUE_LEN).unwrap(); + assert_eq!(slots_per_grant, 113); + let snapshot = RebalanceClassSnapshot { + slot_size: VALUE_LEN, + slots_per_grant, + used_slots: 113 * 174 - 8, + free_slots: 8, + grant_count: 174, + pending_slot_demand: 128, + max_observed_claim_slots: 128, + expected_grant_count: 174, + }; + let requested = owner_slot_pressure_request_bytes(snapshot, 0); + assert_eq!( + requested, + u64::from(128_u32 + slots_per_grant * 4 - 8) * VALUE_LEN, + "one pressure kick must cover the claimant and retained four-grant headroom" + ); + assert!(requested < OWNER_SLOT_PRESSURE_MAX_EVICT_BYTES); + } +} + +async fn owner_local_reserve_rebalance_once(view: &ClientKvApiView) { + let inner = view.client_kv_api().inner(); + let snapshots = { + let pool = inner.owner_local_reserve_pool.lock(); + snapshot_rebalance_classes(&pool) + }; + for snapshot in snapshots { + try_refill_once(inner, snapshot).await; + try_shrink_once(inner, snapshot).await; + } +} + +fn report_owner_local_reserve_state(inner: &ClientKvApiInner) { + let snapshots = { + let pool = inner.owner_local_reserve_pool.lock(); + snapshot_rebalance_classes(&pool) + }; + for snapshot in snapshots { + let grant_count = u64::try_from(snapshot.grant_count).unwrap_or(u64::MAX); + let used_slots = u64::try_from(snapshot.used_slots).unwrap_or(u64::MAX); + let free_slots = u64::try_from(snapshot.free_slots).unwrap_or(u64::MAX); + let reserved_slots = grant_count.saturating_mul(u64::from(snapshot.slots_per_grant)); + tracing::info!( + "owner local reserve state: slot_size={} used_slots={} free_slots={} pending_slots={} grants={} expected_grants={} reserved_slots={} used_slot_bytes={} usable_slot_bytes={} physical_reserved_bytes={}", + snapshot.slot_size, + snapshot.used_slots, + snapshot.free_slots, + snapshot.pending_slot_demand, + snapshot.grant_count, + snapshot.expected_grant_count, + reserved_slots, + used_slots.saturating_mul(snapshot.slot_size), + used_slots + .saturating_add(free_slots) + .saturating_mul(snapshot.slot_size), + grant_count.saturating_mul(OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES), + ); + } +} + +pub fn spawn_owner_local_reserve_rebalance_actor(view: ClientKvApiView) { + let view_task = view.clone(); + view.spawn("owner_local_reserve_rebalance_actor", async move { + let expected = install_expected_capacity_class(view_task.client_kv_api().inner()); + if let Some(expected) = expected { + tracing::info!( + "owner local reserve expected capacity configured: value_len={} payload_capacity_bytes={} slot_size={} slots_per_grant={} expected_grants={} physical_reserved_target_bytes={}", + expected.value_len, + expected.payload_capacity_bytes, + expected.slot_size, + expected.slots_per_grant, + expected.grant_count, + u64::try_from(expected.grant_count) + .unwrap_or(u64::MAX) + .saturating_mul(OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES), + ); + } + let shutdown_poller = view_task.register_shutdown_poller(); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + let notify = view_task + .client_kv_api() + .inner() + .owner_local_reserve_rebalance_notify(); + let mut tick = tokio::time::interval(OWNER_LOCAL_RESERVE_REBALANCE_INTERVAL); + let mut last_report_at = Instant::now(); + + loop { + tokio::select! { + biased; + _ = shutdown_waiter.wait() => { + tracing::info!("owner local reserve rebalance actor stopped by shutdown"); + return; + } + _ = tick.tick() => {} + _ = notify.notified() => {} + } + + if !shutdown_poller.is_running() { + tracing::info!("owner local reserve rebalance actor stopped by shutdown"); + return; + } + + owner_local_reserve_rebalance_once(&view_task).await; + if last_report_at.elapsed() >= OWNER_LOCAL_RESERVE_REPORT_INTERVAL { + report_owner_local_reserve_state(view_task.client_kv_api().inner()); + last_report_at = Instant::now(); + } + } + }); +} + +/// Drive owner slot pressure outside RPC handlers and outside Tokio's worker +/// threads. This is the sole production call site for Moka's synchronous +/// `evict_some`: selection is deliberately bounded, and the returned weight +/// means only "source-eviction candidates selected", never "slots reclaimed". +pub fn spawn_owner_slot_pressure_actor(view: ClientKvApiView) { + let view_task = view.clone(); + view.spawn("owner_slot_pressure_actor", async move { + let shutdown_poller = view_task.register_shutdown_poller(); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + let notify = view_task + .client_kv_api() + .inner() + .owner_local_reserve_rebalance_notify(); + let mut tick = tokio::time::interval(OWNER_SLOT_PRESSURE_INTERVAL); + let mut last_kick_at: Option = None; + let mut last_unsupported_class_report_at: Option = None; + + loop { + tokio::select! { + biased; + _ = shutdown_waiter.wait() => { + tracing::info!("owner slot pressure actor stopped by shutdown"); + return; + } + _ = tick.tick() => {} + _ = notify.notified() => {} + } + if !shutdown_poller.is_running() { + return; + } + + let inner = view_task.client_kv_api().inner(); + if last_kick_at + .is_some_and(|last| last.elapsed() < OWNER_SLOT_PRESSURE_MIN_KICK_INTERVAL) + { + continue; + } + let source_eviction_selected_bytes = + owner_slot_pressure_projected_reclaim_bytes(inner); + let requested_bytes = { + let pool = inner.owner_local_reserve_pool.lock(); + let snapshots = snapshot_rebalance_classes(&pool); + match snapshots.as_slice() { + // No local allocation class has been observed yet. This is + // the normal idle state, not a configuration error. + [] => 0, + [snapshot] => owner_slot_pressure_request_bytes( + *snapshot, + source_eviction_selected_bytes, + ), + _ => { + // Selection is intentionally disabled until the victim + // cache is class-indexed. Rate-limit the diagnostic so + // an unsupported configuration cannot become a 100 Hz + // logging/CPU loop. + if last_unsupported_class_report_at.is_none_or(|last| { + last.elapsed() >= OWNER_LOCAL_RESERVE_REPORT_INTERVAL + }) { + tracing::error!( + active_classes = snapshots.len(), + "owner slot pressure requires a class-indexed victim domain when more than one size class is active" + ); + last_unsupported_class_report_at = Some(Instant::now()); + } + 0 + } + } + }; + if requested_bytes == 0 { + continue; + } + let Some(cache) = inner.owner_hot_cache.clone() else { + continue; + }; + if inner + .owner_hot_eviction_tx + .send(OwnerHotEvictionDispatch::BeginPressure { requested_bytes }) + .is_err() + { + tracing::warn!( + requested_bytes, + "owner slot pressure dispatcher is closed before Moka selection" + ); + return; + } + let selected_bytes = match limit_thirdparty::tokio::task::spawn_blocking(move || { + cache.evict_some(requested_bytes) + }) + .await + { + Ok(selected_bytes) => selected_bytes, + Err(err) => { + tracing::warn!( + requested_bytes, + err = ?err, + "owner slot pressure Moka selection task failed" + ); + 0 + } + }; + if inner + .owner_hot_eviction_tx + .send(OwnerHotEvictionDispatch::EndPressure { selected_bytes }) + .is_err() + { + tracing::warn!( + requested_bytes, + selected_bytes, + "owner slot pressure dispatcher closed after Moka selection" + ); + return; + } + last_kick_at = Some(Instant::now()); + tracing::debug!( + requested_bytes, + selected_bytes, + "owner slot pressure selected bounded source-eviction candidates" + ); + } + }); +} + +pub async fn wait_owner_local_reserve_ready( + inner: &ClientKvApiInner, + slot_size: u64, + slots_per_grant: u32, + key_count: usize, + soft_wait_timeout: Duration, + hard_deadline: Instant, +) -> bool { + let notify = inner.owner_local_reserve_rebalance_notify(); + let mut shutdown_waiter = inner.view.register_shutdown_waiter(); + loop { + let notified = notify.notified(); + { + let mut pool = inner.owner_local_reserve_pool.lock(); + let class_state = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)); + if class_state.free_slot_count() >= key_count { + return true; + } + } + + let Some(remaining_hard) = hard_deadline.checked_duration_since(Instant::now()) else { + return false; + }; + let wait_budget = remaining_hard.min(soft_wait_timeout); + + tokio::select! { + _ = shutdown_waiter.wait() => return false, + _ = tokio::time::sleep(wait_budget) => { + inner.owner_local_reserve_rebalance_notify().notify_waiters(); + } + _ = notified => {} + } + } +} diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs index dec19f5..5ecf000 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/mod.rs @@ -1,18 +1,40 @@ use crate::client_kv_api::delete::handle_batch_delete_client_kv_meta_cache; +use crate::client_kv_api::local_reserve_rebalance::{ + spawn_owner_local_reserve_rebalance_actor, spawn_owner_slot_pressure_actor, +}; use crate::client_kv_api::msg_pack::{ - ExternalDeleteAckReq, ExternalDeleteAckResp, ExternalDeleteReq, ExternalDeleteResp, - ExternalGetReq, ExternalGetResp, ExternalIsExistReq, ExternalIsExistResp, ExternalPutCommitReq, - ExternalPutCommitResp, ExternalPutRevokeReq, ExternalPutRevokeResp, ExternalPutStartReq, - ExternalPutStartResp, ExternalPutTransferEndReq, ExternalPutTransferEndResp, SyncKvToFileReq, + ExternalBatchDeleteAckReq, ExternalBatchDeleteAckResp, ExternalBatchGetCancelReq, + ExternalBatchGetCancelResp, ExternalBatchGetLocalProbeReq, ExternalBatchGetLocalProbeResp, + ExternalBatchGetReq, ExternalBatchGetResp, ExternalBatchGetStartReq, ExternalBatchGetStartResp, + ExternalBatchGetTransferReq, ExternalBatchGetTransferResp, ExternalBatchIsExistReq, + ExternalBatchIsExistResp, ExternalBatchPutCommitReq, ExternalBatchPutCommitResp, + ExternalBatchPutStartReq, ExternalBatchPutStartResp, ExternalBatchPutTransferEndReq, + ExternalBatchPutTransferEndResp, ExternalDeleteAckReq, ExternalDeleteAckResp, + ExternalDeleteReq, ExternalDeleteResp, ExternalExecutePlannedGetReq, + ExternalExecutePlannedGetResp, ExternalGetReq, ExternalGetResp, ExternalIsExistReq, + ExternalIsExistResp, ExternalObservabilitySnapshotReq, ExternalObservabilitySnapshotResp, + ExternalPutCommitReq, ExternalPutCommitResp, ExternalPutRevokeReq, ExternalPutRevokeResp, + ExternalPutStartReq, ExternalPutStartResp, ExternalPutTransferEndReq, + ExternalPutTransferEndResp, SsdStageReadReq, SsdStageReadResp, SyncKvToFileReq, SyncKvToFileResp, TestPutPhaseTrace, }; -use crate::cluster_manager::NodeIDString; +use crate::client_kv_api::reclaim::handle_batch_owner_reclaim; +use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; +use crate::cluster_manager::{NodeID, NodeIDString}; use crate::config::TestSpecConfig; +use crate::kv_ssd_storage::{ + KvSsdPersistBatchPermit, KvSsdPersistCopy, KvSsdPersistSource, KvSsdStorage, KvSsdStorageInit, +}; use crate::master_kv_router::msg_pack::{ - BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, DeleteClientKvMetaCacheItem, + BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, BatchEnqueueReplicaTaskReq, + BatchEvictOwnerSourceReq, BatchGetBindReq, BatchGetDoneReq, BatchGetRevokeReq, + BatchGetStartItemResp, BatchGetStartReq, BatchIsExistReq, BatchOwnerReclaimReq, + BatchPreparePutKeysReq, BatchPublishOwnerSsdReq, BatchPutAppendDoneReq, BatchPutAppendStartReq, + BatchPutDoneReq, BatchPutRevokeReq, BatchPutStartReq, BatchReleasePutKeyReservationsReq, + DeleteClientKvMetaCacheItem, GroupedBatchPutDoneReq, }; use crate::master_lease_manager::msg_pack::{AllocateClientLeaseReq, ClientLeaseKeepaliveReq}; -use crate::memholder::{AllMemholderRefCount, MemoryInfo, UserMemHolder}; +use crate::memholder::{AllMemholderRefCount, ExternalMemHolderInfo, MemoryInfo, UserMemHolder}; use crate::memholder::{ EnsureMemholderMgmtDeleteHandle, MemholderManagerTrait, NodeHolderKey, OwnerDeleteAckItem, OwnerDeleteAckMemMgr, OwnerExternalMemMgr, @@ -22,56 +44,472 @@ use crate::{ client_transfer_engine::{ClientTransferEngine, ClientTransferEngineAccessTrait}, cluster_manager::{ClusterEvent, ClusterManager, ClusterManagerAccessTrait}, master_kv_router::msg_pack::{ - DeleteReq, GetDoneReq, GetMetaReq, GetRevokeReq, GetStartReq, PutDoneReq, PutRevokeReq, - PutStartReq, + DeleteReq, GetDoneReq, GetMetaReq, GetRevokeReq, GetStartReq, PutAppendDoneReq, + PutAppendRevokeReq, PutAppendStartReq, PutDoneReq, PutRevokeReq, PutStartReq, + ReleaseLocalGrantReq, ReserveLocalGrantReq, SsdStageBeginReq, SsdStageDoneReq, }, metric_reporter::{MetricReporter, MetricReporterAccessTrait}, - metrics::{MetricsHandle, OperationKind, RequestStage}, + metrics::{KvLocalitySnapshot, MetricsHandle, OperationKind, RequestStage}, p2p::{ msg_pack::{RPCCaller, RPCHandler}, - p2p_module::{P2pModule, P2pModuleAccessTrait}, + p2p_module::{P2pModule, P2pModuleAccessTrait, RpcTransportPolicy}, }, - rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}, + rpcresp_kvresult_convert::msg_and_error::{ApiError, ErrorCode, KvError, KvResult}, }; +use ::tokio::sync::watch; use async_trait::async_trait; -use dashmap::DashMap; +use dashmap::{DashMap, mapref::entry::Entry as DashMapEntry}; use fluxon_framework::{LogicalModule, define_module}; use fluxon_util::map_lock::AMapLock; +use fluxon_util::pin_aware_moka::{PinAwareMoka, PinGuard}; use limit_thirdparty::tokio; +use moka::notification::RemovalCause; use parking_lot::Mutex; +use std::collections::BTreeSet; +use std::collections::{HashMap, HashSet}; +use std::hash::{Hash, Hasher}; use std::sync::Weak; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; use std::sync::{Arc, OnceLock}; -use std::time::Duration; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tracing::warn; +const OWNER_LOCAL_PUBLISH_QUEUE_CAPACITY: usize = 4096; +const OWNER_LOCAL_PUBLISH_MAX_INFLIGHT: usize = 64; +const SSD_STAGE_RPC_TIMEOUT: Duration = Duration::from_secs(300); +const SSD_STAGE_TERMINAL_TTL: Duration = Duration::from_secs(10 * 60); +const SSD_STAGE_DONE_RETRY_INITIAL_BACKOFF: Duration = Duration::from_millis(50); +const SSD_STAGE_DONE_RETRY_MAX_BACKOFF: Duration = Duration::from_secs(5); + +#[derive(Clone, Debug)] +struct CompletedSsdStage { + request: SsdStageReadReq, + response: SsdStageReadResp, +} + +struct SsdStageSharedOp { + request: SsdStageReadReq, + terminal: watch::Sender>, + completed: AtomicBool, +} + +impl SsdStageSharedOp { + fn new(request: SsdStageReadReq) -> Arc { + let (terminal, _receiver) = watch::channel(None); + Arc::new(Self { + request, + terminal, + completed: AtomicBool::new(false), + }) + } + + fn complete(&self, response: SsdStageReadResp) -> bool { + if self + .completed + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return false; + } + self.terminal.send_replace(Some(response)); + true + } + + async fn wait(&self) -> SsdStageReadResp { + let mut terminal = self.terminal.subscribe(); + loop { + if let Some(response) = terminal.borrow_and_update().clone() { + return response; + } + if terminal.changed().await.is_err() { + let err = KvError::Api(ApiError::Unknown { + detail: format!( + "SSD stage singleflight closed without a terminal result: get_id={}", + self.request.get_id + ), + }); + return ssd_stage_error_response(err); + } + } + } +} + +fn ssd_stage_error_response(err: KvError) -> SsdStageReadResp { + SsdStageReadResp { + error_code: err.code(), + error_json: err.to_json(), + } +} + +fn ssd_stage_request_mismatch_response( + expected: &SsdStageReadReq, + actual: &SsdStageReadReq, +) -> SsdStageReadResp { + ssd_stage_error_response(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "SSD stage get_id was reused with a different operation identity: get_id={} expected={:?} actual={:?}", + actual.get_id, expected, actual + ), + })) +} + +#[cfg(test)] +mod ssd_stage_singleflight_tests { + use super::{SsdStageReadReq, SsdStageReadResp, SsdStageSharedOp}; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + use std::sync::Arc; + + fn request(get_id: u64) -> SsdStageReadReq { + SsdStageReadReq { + key: "ssd-key".to_string(), + put_id: (17, 2), + get_id, + stage_addr: 0x1000, + stage_capacity: 8192, + len: 4096, + } + } + + #[limit_thirdparty::tokio::test] + async fn all_ssd_stage_followers_reuse_one_terminal_result() { + let op = SsdStageSharedOp::new(request(91)); + let waiters = futures::future::join_all((0..64).map(|_| { + let waiter = op.clone(); + async move { waiter.wait().await } + })); + let leader = async { + assert!(op.complete(SsdStageReadResp { + error_code: OK, + error_json: String::new(), + })); + assert!( + !op.complete(SsdStageReadResp { + error_code: 1, + error_json: "second terminal".to_string(), + }), + "the source operation must publish only one terminal result" + ); + }; + let (responses, ()) = futures::future::join(waiters, leader).await; + + for response in responses { + assert_eq!(response.error_code, OK); + assert!(response.error_json.is_empty()); + } + assert_eq!(Arc::strong_count(&op), 1); + } +} + /// Information about a memholder held by external client #[derive(Clone)] pub struct ExternalHoldingGetInfo { pub key: String, pub req_node_id: String, + /// Requester membership generation observed when this holding was + /// installed. Unknown generations are never removed by generation-scoped + /// MemberLeft cleanup. + pub requester_node_start_time: Option, pub memory_info: Arc, // The actual memholder being held + _owner_hot_pin: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct OwnerRuntimeObserveSnapshot { + pub ssd_capacity_bytes: u64, + pub ssd_used_bytes: u64, + pub ssd_persist_requests: u64, + pub ssd_persist_successes: u64, + pub ssd_persist_failures: u64, + pub ssd_persist_bytes: u64, + pub ssd_persist_duration_us: u64, + pub ssd_persist_batch_requests: u64, + pub ssd_persist_batch_items: u64, + pub ssd_persist_flush_batches: u64, + pub ssd_persist_busy_batches: u64, + pub ssd_persist_admission_skips: u64, + pub ssd_persist_batch_duration_us: u64, + pub ssd_write_candidate_items: u64, + pub ssd_write_candidate_bytes: u64, + pub ssd_write_admitted_items: u64, + pub ssd_write_admitted_bytes: u64, + pub ssd_write_dropped_items: u64, + pub ssd_write_dropped_bytes: u64, + pub ssd_write_refunded_items: u64, + pub ssd_write_refunded_bytes: u64, + pub ssd_load_requests: u64, + pub ssd_load_successes: u64, + pub ssd_load_misses: u64, + pub ssd_load_failures: u64, + pub ssd_load_bytes: u64, + pub ssd_load_duration_us: u64, + pub ssd_memory_hits: u64, + pub ssd_disk_hits: u64, + pub ssd_outer_hits: u64, + pub ssd_removals: u64, + pub ssd_stage_flights: u64, + pub ssd_stage_terminals: u64, + pub ssd_stage_ready_requests: u64, + pub ssd_stage_ready_successes: u64, + pub ssd_stage_ready_failures: u64, + pub ssd_stage_ready_duration_us: u64, + pub ssd_stage_execute_completions: u64, + pub ssd_stage_terminal_published: u64, + pub ssd_stage_terminal_cache_inserts: u64, + pub ssd_stage_terminal_cache_duration_us: u64, + pub ssd_stage_response_send_attempts: u64, + pub ssd_stage_response_send_successes: u64, + pub ssd_stage_response_send_failures: u64, + pub ssd_stage_response_send_duration_us: u64, + pub ssd_source_ready_wait_requests: u64, + pub ssd_source_ready_wait_successes: u64, + pub ssd_source_ready_wait_failures: u64, + pub ssd_source_ready_wait_duration_us: u64, + pub ssd_target_pull_requests: u64, + pub ssd_target_pull_successes: u64, + pub ssd_target_pull_failures: u64, + pub ssd_target_pull_duration_us: u64, + pub ssd_stage_done_detached: u64, + pub external_get_holding_entries: u64, + pub external_get_holding_bytes: u64, + pub external_get_start_handles: u64, + pub external_get_flights: u64, + pub external_get_flights_starting: u64, + pub external_get_flights_finishing: u64, + pub external_get_flights_revoking: u64, + pub external_get_undecided_interests: u64, + pub external_get_retained_interests: u64, + pub owner_local_probe_batches: u64, + pub owner_local_probe_items: u64, + pub owner_local_probe_local_items: u64, + pub owner_local_probe_remote_items: u64, + pub planned_cpu_get_batches: u64, + pub planned_cpu_get_local_items: u64, + pub planned_cpu_get_leader_items: u64, + pub planned_cpu_get_follower_items: u64, + pub external_pending_put_entries: u64, + pub remote_put_flights_active: u64, + pub remote_put_flight_leaders: u64, + pub remote_put_flight_followers: u64, + pub remote_put_source_unavailable: u64, + pub remote_put_source_fenced: u64, + pub remote_put_source_missing: u64, + pub remote_put_source_version_mismatch: u64, + pub remote_put_transfers: u64, + pub remote_put_published: u64, + pub remote_put_already_satisfied: u64, + pub remote_put_obsolete: u64, + pub remote_put_failed: u64, + pub local_ssd_put_flights_active: u64, + pub local_ssd_put_flight_leaders: u64, + pub local_ssd_put_flight_followers: u64, + pub local_ssd_put_source_unavailable: u64, + pub local_ssd_put_published: u64, + pub local_ssd_put_already_present: u64, + pub local_ssd_put_dropped: u64, + pub local_ssd_put_obsolete: u64, + pub local_ssd_put_failed: u64, + pub local_reserve_slots_free: u64, + pub local_reserve_slots_prepared: u64, + pub local_reserve_slots_pending_visible: u64, + pub local_reserve_slots_committed: u64, + pub hot_cache_capacity_bytes: u64, + pub hot_cache_entries: u64, + pub hot_cache_weighted_bytes: u64, + pub hot_size_evictions: u64, + pub hot_source_evict_handoff_members: u64, + pub hot_source_evict_committed_members: u64, + pub hot_source_evict_restored_members: u64, + pub hot_source_evict_obsolete: u64, + pub hot_source_evict_dispatch_failed: u64, + pub hot_source_eviction_selected: u64, + pub hot_source_evict_retry_entries: u64, + pub hot_source_evict_retry_scheduled: u64, + pub hot_source_evict_retry_emitted: u64, + pub hot_selection_debt_bytes: u64, + pub hot_source_eviction_selected_bytes: u64, + pub hot_eviction_skipped_stale: u64, + pub hot_eviction_skipped_reclaim: u64, + pub hot_eviction_skipped_active_holders: u64, + pub hot_victim_duplicates: u64, + pub hot_victim_invalid_backing: u64, + pub grouped_put_done_batches: u64, + pub grouped_put_done_items: u64, + pub legacy_put_done_batches: u64, + pub legacy_put_done_items: u64, } pub use get::RemoteGetInfo; +pub use put::{OwnerLocalPublishItem, OwnerLocalPublishJob, OwnerReservedPutItem}; pub mod external_api; +mod local_reserve_rebalance; +mod reclaim; pub use external_api::HandlerForExternalClient; pub type TestObservePutPhaseSink = Arc>>; +pub type ExternalGetStartTransferOutput = + Vec, Option)>>>; -#[derive(Debug, Clone)] -enum CachedValue { - MetadataOnly, - LocalReplica(Arc), +pub enum ExternalGetStartOwnerItem { + Local { memory_info: Arc }, + Shared { interest: ExternalGetKeyInterest }, } -impl CachedValue { - fn rank(&self) -> u8 { - match self { - Self::MetadataOnly => 0, - Self::LocalReplica(_) => 1, +/// One request's interest in a per-key Get flight. +/// +/// A pending prefix decision is owned by this guard, not by control flow. If +/// the request future is cancelled at any await point, Drop retires the +/// undecided count so the atomic_batch task can still choose Finish or Revoke. +pub struct ExternalGetKeyInterest { + op: Arc, + decision_pending: bool, +} + +impl ExternalGetKeyInterest { + pub fn new(op: Arc, decision_pending: bool) -> Self { + Self { + op, + decision_pending, + } + } + + pub fn op(&self) -> &Arc { + &self.op + } + + pub fn decide(&mut self, retain: bool) { + if !self.decision_pending { + return; + } + self.decision_pending = false; + + let wake = { + let mut state = self.op.state.lock(); + if state.undecided == 0 { + tracing::error!( + "external Get interest observed undecided underflow: key={}", + self.op.key + ); + return; + } + state.undecided -= 1; + if retain + && matches!( + state.phase, + ExternalGetKeySharedPhase::Starting | ExternalGetKeySharedPhase::Started { .. } + ) + { + state.retained = state + .retained + .checked_add(1) + .expect("external Get singleflight retained overflow"); + } + state.undecided == 0 + }; + if wake { + self.op.notify.notify_waiters(); + } + } +} + +impl Drop for ExternalGetKeyInterest { + fn drop(&mut self) { + self.decide(false); + } +} + +#[derive(Clone)] +pub enum ExternalGetKeySharedPhase { + Starting, + Started { + item: BatchGetStartItemResp, + }, + /// At least one request plan's prefix retained this key. The leader atomic_batch is + /// transferring/completing it and later publishes one canonical result. + Finishing { + item: BatchGetStartItemResp, + }, + /// No request plan's prefix retained this prepared Get. Keep the marker in + /// the owner key fence until BatchGetRevoke and local-slot release finish, + /// so a new overlapping batch cannot race the revoke. + Revoking { + /// Keep the exact master operation and prepared target reachable until + /// Revoke reaches a definite terminal response. + item: BatchGetStartItemResp, + }, + Ready { + result: ExternalGetStartSharedItemResult, + }, + Failed { + error_code: ErrorCode, + error_json: String, + }, +} + +pub struct ExternalGetKeySharedState { + /// Exact-batch operations which joined while the key was Starting/Started + /// and have not yet applied their own atomic-prefix decision. + pub undecided: usize, + /// Number of those operations whose transferable prefix retained the key. + pub retained: usize, + pub phase: ExternalGetKeySharedPhase, + /// Monotonic publication time for Ready/Failed. Observation-only users + /// compare this with the later handle-consume time to distinguish data + /// that was already ready from time actually spent waiting in Transfer. + pub terminal_at: Option, +} + +pub struct ExternalGetKeySharedOp { + pub key: String, + pub state: Mutex, + pub notify: Arc, +} + +impl ExternalGetKeySharedOp { + pub fn new(key: String) -> Self { + Self { + key, + state: Mutex::new(ExternalGetKeySharedState { + undecided: 1, + retained: 0, + phase: ExternalGetKeySharedPhase::Starting, + terminal_at: None, + }), + notify: Arc::new(limit_thirdparty::tokio::sync::Notify::new()), } } } +#[derive(Clone, Debug)] +pub struct ExternalGetStartPrefixResult { + pub raw_prefix_hit_len: usize, + pub transferable_len: usize, + pub first_miss_index: Option, + pub first_error_kind: Option, +} + +#[derive(Clone)] +pub enum ExternalGetStartSharedItemResult { + Hit { + memholder: Arc, + }, + Miss, + Error { + error_code: ErrorCode, + error_json: String, + }, +} + +pub struct ExternalGetStartEntry { + pub req_node_id: String, + /// Requester membership generation observed when this handle was created. + /// `None` is retained only for compatibility with a temporarily incomplete + /// cluster cache; generation cleanup must never guess in that case. + pub requester_node_start_time: Option, + pub keys: Vec, + pub items: Vec, + pub atomic_group_lens: Vec, + pub created_at: Instant, +} + /// Optional arguments for put operations #[derive(Clone, Debug)] pub enum PutOptionalArg { @@ -79,6 +517,10 @@ pub enum PutOptionalArg { LeaseId(u64), /// Ask the master to fail-fast when the same key already has an inflight put. RejectIfInflightSameKey, + /// Ask the master to fail-fast when the key already has a committed live replica. + RejectIfExistSameKey, + /// Disable asynchronous remote replica task after local commit in write-back mode. + SkipMakeReplicaTask, /// Prefer placing the target allocation on a kvclient within this sub_cluster. PreferredSubCluster(String), /// Hidden test-only side-channel for collecting per-put phase timings. @@ -98,6 +540,8 @@ impl PutOptionalArgs { self.0.iter().rev().find_map(|a| match a { PutOptionalArg::LeaseId(id) => Some(*id), PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExistSameKey + | PutOptionalArg::SkipMakeReplicaTask | PutOptionalArg::PreferredSubCluster(_) | PutOptionalArg::TestObservePutPhases(_) => None, }) @@ -109,12 +553,27 @@ impl PutOptionalArgs { .any(|arg| matches!(arg, PutOptionalArg::RejectIfInflightSameKey)) } + pub fn reject_if_exist_same_key(&self) -> bool { + self.0 + .iter() + .any(|arg| matches!(arg, PutOptionalArg::RejectIfExistSameKey)) + } + + pub fn make_replica_task(&self) -> bool { + !self + .0 + .iter() + .any(|arg| matches!(arg, PutOptionalArg::SkipMakeReplicaTask)) + } + /// Get the last provided preferred_sub_cluster if any. pub fn preferred_sub_cluster(&self) -> Option<&str> { self.0.iter().rev().find_map(|a| match a { PutOptionalArg::PreferredSubCluster(sc) => Some(sc.as_str()), PutOptionalArg::LeaseId(_) | PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExistSameKey + | PutOptionalArg::SkipMakeReplicaTask | PutOptionalArg::TestObservePutPhases(_) => None, }) } @@ -124,6 +583,8 @@ impl PutOptionalArgs { PutOptionalArg::TestObservePutPhases(sink) => Some(sink.clone()), PutOptionalArg::LeaseId(_) | PutOptionalArg::RejectIfInflightSameKey + | PutOptionalArg::RejectIfExistSameKey + | PutOptionalArg::SkipMakeReplicaTask | PutOptionalArg::PreferredSubCluster(_) => None, }) } @@ -274,6 +735,7 @@ pub enum KvMetrics { transfer_poll_wait_us: i64, transfer_poll_iters: i64, transfer_used_fast_path: bool, + transfer_used_nixl: bool, transfer_local_noop: bool, transfer_remote_transfer: bool, }, @@ -343,6 +805,161 @@ async fn handle_external_get( } } +async fn handle_external_batch_get( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.keys.len(); + let resp = view + .client_kv_api() + .external_batch_get(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_get error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + let mut r: ExternalBatchGetResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.items = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_batch_get_local_probe( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.keys.len(); + let resp = view + .client_kv_api() + .external_batch_get_local_probe(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_get_local_probe error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + ExternalBatchGetLocalProbeResp { + items: Vec::new(), + error_code: e.code(), + error_json: e.to_json(), + } + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_batch_get_start( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.keys.len(); + let resp = view + .client_kv_api() + .external_batch_get_start(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_get_start error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + crate::rpcresp_kvresult_convert::FromError::from_error(&e) + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_batch_get_transfer( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_handle = req.handle; + let resp = view + .client_kv_api() + .external_batch_get_transfer(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_get_transfer error: {e}; handle={handle}", + handle = dbg_handle + ); + let mut r: ExternalBatchGetTransferResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.items = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_execute_planned_get( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let plan_handle = req.plan_handle; + let item_count = req.items.len(); + let resp = view + .client_kv_api() + .external_execute_planned_get(req) + .await + .unwrap_or_else(|err| { + tracing::error!( + "handle_external_execute_planned_get error: {}; plan_handle={} items={}", + err, + plan_handle, + item_count + ); + ExternalExecutePlannedGetResp { + items: Vec::new(), + error_code: err.code(), + error_json: err.to_json(), + } + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_batch_get_cancel( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_handle = req.handle; + let resp = view + .client_kv_api() + .external_batch_get_cancel(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_get_cancel error: {e}; handle={handle}", + handle = dbg_handle + ); + crate::rpcresp_kvresult_convert::FromError::from_error(&e) + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + async fn handle_external_put_start( view: &ClientKvApiView, msg: &MsgPack, @@ -375,6 +992,33 @@ async fn handle_external_put_start( raw_bytes: Vec::new(), } } + +async fn handle_external_batch_put_start( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.items.len(); + let resp = view + .client_kv_api() + .external_batch_put_start(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_put_start error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + let mut r: ExternalBatchPutStartResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.items = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + async fn handle_external_put_transfer_end( view: &ClientKvApiView, msg: &MsgPack, @@ -401,6 +1045,32 @@ async fn handle_external_put_transfer_end( } } +async fn handle_external_batch_put_transfer_end( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.items.len(); + let resp = view + .client_kv_api() + .external_batch_put_transfer_end(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_put_transfer_end error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + let mut r: ExternalBatchPutTransferEndResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.items = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + async fn handle_external_put_commit( view: &ClientKvApiView, msg: &MsgPack, @@ -426,6 +1096,32 @@ async fn handle_external_put_commit( } } +async fn handle_external_batch_put_commit( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.items.len(); + let resp = view + .client_kv_api() + .external_batch_put_commit(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_put_commit error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + let mut r: ExternalBatchPutCommitResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.items = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + async fn handle_external_put_revoke( view: &ClientKvApiView, msg: &MsgPack, @@ -451,6 +1147,20 @@ async fn handle_external_put_revoke( } } +async fn handle_ssd_stage_read( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + MsgPack { + serialize_part: view + .client_kv_api() + .inner() + .execute_ssd_stage(&msg.serialize_part) + .await, + raw_bytes: Vec::new(), + } +} + async fn handle_external_delete_ack( view: &ClientKvApiView, msg: &MsgPack, @@ -500,6 +1210,89 @@ async fn handle_external_delete_ack( raw_bytes: Vec::new(), } } + +fn release_external_holder_ids_with( + holder_ids: Vec, + mut remove: impl FnMut(u64) -> bool, +) -> (u32, u32) { + let mut seen = HashSet::with_capacity(holder_ids.len()); + let mut released = 0u32; + let mut missing = 0u32; + for holder_id in holder_ids { + if !seen.insert(holder_id) { + continue; + } + if remove(holder_id) { + released = released.saturating_add(1); + } else { + missing = missing.saturating_add(1); + } + } + (released, missing) +} + +async fn handle_external_batch_delete_ack( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let expected = view.cluster_manager().get_self_info().node_start_time; + if req.started_time != 0 && req.started_time != expected { + let err = crate::rpcresp_kvresult_convert::msg_and_error::KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::OwnerStartTimeMismatch { + expected, + got: req.started_time, + }, + ); + return MsgPack { + serialize_part: ExternalBatchDeleteAckResp::from_error(&err), + raw_bytes: Vec::new(), + }; + } + + let inner = view.client_kv_api().inner(); + let external_client_id = req.external_client_id; + let requested_count = req.holder_ids.len(); + let (released_count, missing_count) = + release_external_holder_ids_with(req.holder_ids, |holder_id| { + inner + .external_get_holding + .remove(&NodeHolderKey::new(external_client_id.clone(), holder_id)) + .is_some() + }); + tracing::debug!( + external_client_id, + requested_count, + released_count, + missing_count, + "processed external holder ACK batch" + ); + MsgPack { + serialize_part: ExternalBatchDeleteAckResp { + released_count, + missing_count, + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +#[cfg(test)] +mod external_delete_ack_batch_tests { + use super::release_external_holder_ids_with; + use std::collections::HashSet; + + #[test] + fn duplicate_and_missing_holder_ids_are_idempotent() { + let mut live = HashSet::from([3u64, 5u64]); + let (released, missing) = + release_external_holder_ids_with(vec![3, 3, 4, 5], |holder_id| live.remove(&holder_id)); + assert_eq!((released, missing), (2, 1)); + assert!(live.is_empty()); + } +} + async fn handle_external_delete( view: &ClientKvApiView, msg: &MsgPack, @@ -551,6 +1344,51 @@ async fn handle_external_is_exist( } } +async fn handle_external_batch_is_exist( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let dbg_len = req.keys.len(); + let resp = view + .client_kv_api() + .external_batch_is_exist(req) + .await + .unwrap_or_else(|e| { + tracing::error!( + "handle_external_batch_is_exist error: {e}; batch_len={batch_len}", + batch_len = dbg_len + ); + let mut r: ExternalBatchIsExistResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&e); + r.exists_list = Vec::new(); + r + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + +async fn handle_external_observability_snapshot( + view: &ClientKvApiView, + msg: &MsgPack, +) -> MsgPack { + let req = msg.serialize_part.clone(); + let resp = view + .client_kv_api() + .external_observability_snapshot(req) + .await + .unwrap_or_else(|e| { + tracing::error!("handle_external_observability_snapshot error: {e}"); + crate::rpcresp_kvresult_convert::FromError::from_error(&e) + }); + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + } +} + fn write_all_at(file: &std::fs::File, mut buf: &[u8], mut offset: u64) -> std::io::Result<()> { use std::io::{Error, ErrorKind}; use std::os::unix::fs::FileExt; @@ -729,6 +1567,9 @@ define_module!( #[derive(Clone, Debug)] pub struct ClientKvApiNewArg { pub test_spec_config: TestSpecConfig, + /// Logical hot-tier capacity only. This does not resize the owner segment. + pub owner_hot_cache_capacity_bytes: Option, + pub ssd_storage: Option, } pub struct ClientKvApi(ClientKvApiInner); @@ -737,211 +1578,5067 @@ pub struct ClientKvApi(ClientKvApiInner); pub struct GetCachedInfo { put_time_ms: u64, put_version: u32, - value: CachedValue, + mem_holder: Arc, } -struct ClientKvApiViewHolder { - view: OnceLock, +#[derive(Debug)] +pub struct PrecommitLocalVisibleInfo { + mem_holder: Arc, } -impl ClientKvApiViewHolder { - fn new() -> Self { +#[derive(Debug)] +pub(crate) struct PendingLocalGetInfo { + get_id: u64, + put_id: crate::master_kv_router::put::PutIDForAKey, + mem_holder: Arc, +} + +#[derive(Debug)] +pub(crate) struct LocalSnapshotInfo { + put_time_ms: u64, + put_version: u32, +} + +#[derive(Clone)] +struct OwnerHotCacheEntry { + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: Weak, + weight_bytes: u32, +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub(crate) struct OwnerHotPinAlias { + key: String, + memory_info_ptr: usize, +} + +impl OwnerHotPinAlias { + fn new(memory_info: &Arc) -> Self { Self { - view: OnceLock::new(), + key: memory_info.key.clone(), + memory_info_ptr: Arc::as_ptr(memory_info) as usize, } } +} - fn attach(&self, view: ClientKvApiView) { - // The framework attaches a module's PostView exactly once at the init barrier. - // A second attach indicates a programming error. - self.view - .set(view) - .unwrap_or_else(|_| panic!("ClientKvApi view attached twice")); +type OwnerHotCache = PinAwareMoka; + +#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct OwnerHotReplicaIdentity { + key: String, + put_time_ms: u64, + put_version: u32, +} + +#[derive(Default)] +struct OwnerHotCacheCounters { + size_evictions: AtomicU64, + source_evict_handoff_members: AtomicU64, + source_evict_committed_members: AtomicU64, + source_evict_restored_members: AtomicU64, + source_evict_obsolete: AtomicU64, + source_evict_dispatch_failed: AtomicU64, + source_evict_retry_scheduled: AtomicU64, + source_evict_retry_emitted: AtomicU64, + selection_debt_bytes: Arc, + /// Bytes behind an installed source-selection fence. Unlike candidate + /// debt, every byte here can become a physical Free slot after reclaim. + source_eviction_selected_bytes: AtomicU64, + skipped_stale: AtomicU64, + skipped_reclaim: AtomicU64, + skipped_active_holders: AtomicU64, + victim_duplicates: AtomicU64, + victim_invalid_backing: AtomicU64, + grouped_put_done_batches: AtomicU64, + grouped_put_done_items: AtomicU64, + legacy_put_done_batches: AtomicU64, + legacy_put_done_items: AtomicU64, +} + +#[derive(Default)] +struct OwnerRemotePutCounters { + active: AtomicU64, + leaders: AtomicU64, + followers: AtomicU64, + source_unavailable: AtomicU64, + source_fenced: AtomicU64, + source_missing: AtomicU64, + source_version_mismatch: AtomicU64, + transfers: AtomicU64, + published: AtomicU64, + already_satisfied: AtomicU64, + obsolete: AtomicU64, + failed: AtomicU64, +} + +#[derive(Default)] +struct OwnerLocalSsdPutCounters { + active: AtomicU64, + leaders: AtomicU64, + followers: AtomicU64, + source_unavailable: AtomicU64, + published: AtomicU64, + already_present: AtomicU64, + dropped: AtomicU64, + obsolete: AtomicU64, + failed: AtomicU64, +} + +#[derive(Default)] +struct OwnerPlannedGetCounters { + local_probe_batches: AtomicU64, + local_probe_items: AtomicU64, + local_probe_local_items: AtomicU64, + local_probe_remote_items: AtomicU64, + batches: AtomicU64, + local_items: AtomicU64, + leader_items: AtomicU64, + follower_items: AtomicU64, +} + +#[derive(Default)] +struct OwnerSsdStageCounters { + ready_requests: AtomicU64, + ready_successes: AtomicU64, + ready_failures: AtomicU64, + ready_duration_us: AtomicU64, + execute_completions: AtomicU64, + terminal_published: AtomicU64, + terminal_cache_inserts: AtomicU64, + terminal_cache_duration_us: AtomicU64, + response_send_attempts: AtomicU64, + response_send_successes: AtomicU64, + response_send_failures: AtomicU64, + response_send_duration_us: AtomicU64, + source_ready_wait_requests: AtomicU64, + source_ready_wait_successes: AtomicU64, + source_ready_wait_failures: AtomicU64, + source_ready_wait_duration_us: AtomicU64, + target_pull_requests: AtomicU64, + target_pull_successes: AtomicU64, + target_pull_failures: AtomicU64, + target_pull_duration_us: AtomicU64, + done_detached: AtomicU64, +} + +struct OwnerHotSelectionDebt { + weight_bytes: u64, + outstanding_bytes: Arc, + released: AtomicU32, +} + +impl OwnerHotSelectionDebt { + fn new(weight_bytes: u64, outstanding_bytes: Arc) -> Arc { + outstanding_bytes.fetch_add(weight_bytes, Ordering::AcqRel); + Arc::new(Self { + weight_bytes, + outstanding_bytes, + released: AtomicU32::new(0), + }) } - fn clone_view(&self) -> ClientKvApiView { - self.view.get().unwrap().clone() + fn release(&self) { + if self + .released + .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.outstanding_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.saturating_sub(self.weight_bytes)) + }) + .expect("owner hot selection debt update cannot fail"); + } } } -impl std::ops::Deref for ClientKvApiViewHolder { - type Target = ClientKvApiView; +impl OwnerHotCacheCounters { + fn add_source_eviction_selected_bytes(&self, bytes: u64) { + self.source_eviction_selected_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current.checked_add(bytes) + }) + .expect("owner source-selection byte credit overflowed"); + } - fn deref(&self) -> &Self::Target { - self.view.get().unwrap() + fn remove_source_eviction_selected_bytes(&self, bytes: u64) { + self.source_eviction_selected_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current.checked_sub(bytes) + }) + .expect("owner source-selection byte credit underflowed"); } } -pub struct ClientKvApiInner { - view: ClientKvApiViewHolder, - test_spec_config: TestSpecConfig, - metrics: OnceLock>, +#[derive(Clone)] +pub(crate) struct OwnerHotEvictionEvent { + key: String, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: Weak, + selection_debt: Arc, + /// The event was dispatched from the bounded retry queue. + retry: bool, + /// Exact single-key source-delete transaction retained across RPC retries. + source_eviction_victim: + Option>, + /// Failure count follows the event while it is dispatched, preserving + /// exponential backoff across queue take/reinsert cycles. + retry_failures: u32, +} - /// make sure each remote kv get run in order - pub get_remote_kv_lock: AMapLock, - /// key -> value info on this node - /// we can only remove value if it's put_time_ms and put_version match remote eviction command - get_cached_info: DashMap, +pub(crate) enum OwnerHotEvictionDispatch { + Victim(OwnerHotEvictionEvent), + BeginPressure { requested_bytes: u64 }, + EndPressure { selected_bytes: u64 }, + Flush, +} - /// Shared delete actor input for owner -> external weak-index invalidation. - pub external_invalidate_delete: EnsureMemholderMgmtDeleteHandle, - /// Shared delete actor input for owner -> master delete-ack batching. - pub delete_ack_batch: EnsureMemholderMgmtDeleteHandle, - /// Shared manager for owner -> master delete-ack batching. - pub owner_delete_ack_mgr: OwnerDeleteAckMemMgr, +pub(crate) enum OwnerHotEvictionPreparation { + Ready { + trigger: OwnerHotReplicaIdentity, + source: Arc, + }, + RetryableReclaimFence, + TemporarilyPinned, + Obsolete, +} - // record external_client get_holding info (owned, flattened manager) - pub external_get_holding: OwnerExternalMemMgr, - /// Weak handle to a shared refcount tracker for all UserMemHolder of this client. - /// - /// - A strong `Arc` is given to every `UserMemHolder` created by this client. - /// - When the last `UserMemHolder` is dropped, the strong `Arc` is dropped too, - /// and this weak handle will no longer upgrade, meaning the client can be safely dropped. - /// - Stored as `Weak` in `OnceLock` to avoid cycles and allow lazy initialization. - pub all_memholder_refcount: OnceLock>, - /// External API is implemented directly on ClientKvApi; no handler stored here +pub(crate) enum OwnerHotSelectionFenceOutcome { + Fenced, + Retryable, + TemporarilyPinned, + Obsolete, +} - #[cfg(test)] - test_record: crate::client_kv_api::client_test_record::ClientTestRecord, +struct OwnerHotRetryEntry { + event: OwnerHotEvictionEvent, + failures: u32, + next_attempt_at: Instant, + dispatched: bool, +} - rpc_caller_get_start: RPCCaller, - rpc_caller_get_revoke: RPCCaller, - rpc_caller_get_done: RPCCaller, - rpc_caller_put_start: RPCCaller, - rpc_caller_put_revoke: RPCCaller, - rpc_caller_put_done: RPCCaller, - rpc_caller_delete: RPCCaller, - rpc_caller_batch_delete_ack: RPCCaller, - rpc_caller_get_meta: RPCCaller, - _rpc_caller_allocate_client_lease: RPCCaller, - _rpc_caller_client_lease_keepalive: RPCCaller, - rpc_caller_external_put_commit: RPCCaller, - rpc_caller_external_put_revoke: RPCCaller, - rpc_caller_resolve_side_transfer_lane: RPCCaller, +#[derive(Default)] +struct OwnerHotRetryState { + entries: HashMap, + /// Exactly one deadline for each non-dispatched entry. Unlike a lazy + /// generation heap, rescheduling replaces the old tuple, so both memory + /// and lock-held work are bounded by the live retry set, not its history. + deadlines: BTreeSet<(Instant, OwnerHotReplicaIdentity)>, +} - /// Default lease id recorded for inspection/convenience, but NOT auto-applied. - /// Callers must explicitly pass `Some(lease_id)` to attach a put to a lease. - default_lease_id: parking_lot::RwLock>, - /// External put (remote target) pending context keyed by (key, put_time_ms, put_version). - /// 注意：put_id (time_ms,version) 在不同 key 上并不全局唯一，因此必须携带 key 作为索引的一部分，避免碰撞。 - /// 使用 moka::sync::SegmentedCache 并设置 30 分钟 TTL，避免异常路径未清理导致的泄漏；不设置容量上限，纯 TTL 控制。 - external_pending_puts: moka::sync::SegmentedCache<(String, u64, u32), ExternalPendingPutCtx>, +/// Exactly-once owner-local retry state. It is physically bounded by the +/// owner committed-slot pool: one identity can occupy at most one entry, and +/// obsolete identities are removed when their local version is invalidated. +/// The actor emits only a small due batch and applies exponential backoff. +struct OwnerHotRetryQueue { + state: Mutex, + notify: Arc, + counters: Arc, } -impl ClientKvApiInner { - fn should_replace_cached_info( - current: &GetCachedInfo, - put_time_ms: u64, - put_version: u32, - new_value_rank: u8, - ) -> bool { - if current.put_time_ms != put_time_ms { - return current.put_time_ms <= put_time_ms; +impl OwnerHotRetryQueue { + fn new(counters: Arc) -> Self { + Self { + state: Mutex::new(OwnerHotRetryState::default()), + notify: Arc::new(limit_thirdparty::tokio::sync::Notify::new()), + counters, + } + } + + fn retry_delay(failures: u32) -> Duration { + let shift = failures.saturating_sub(1).min(8); + Duration::from_millis(25u64.saturating_mul(1u64 << shift)).min(Duration::from_secs(5)) + } + + fn schedule(&self, mut event: OwnerHotEvictionEvent, reason: &'static str) { + let identity = OwnerHotReplicaIdentity { + key: event.key.clone(), + put_time_ms: event.put_id.0, + put_version: event.put_id.1, + }; + let now = Instant::now(); + let mut state = self.state.lock(); + let previous_deadline = state + .entries + .get(&identity) + .and_then(|entry| (!entry.dispatched).then_some(entry.next_attempt_at)); + if let Some(previous_deadline) = previous_deadline { + state + .deadlines + .remove(&(previous_deadline, identity.clone())); + } + let entry = state + .entries + .entry(identity.clone()) + .or_insert_with(|| OwnerHotRetryEntry { + event: event.clone(), + failures: event.retry_failures, + next_attempt_at: now, + dispatched: false, + }); + if !Arc::ptr_eq(&entry.event.selection_debt, &event.selection_debt) { + entry.event.selection_debt.release(); } - if current.put_version != put_version { - return current.put_version <= put_version; + if event.source_eviction_victim.is_none() { + event.source_eviction_victim = entry.event.source_eviction_victim.clone(); } - current.value.rank() < new_value_rank + entry.failures = entry.failures.max(event.retry_failures).saturating_add(1); + event.retry = true; + event.retry_failures = entry.failures; + entry.event = event; + entry.next_attempt_at = now + Self::retry_delay(entry.failures); + entry.dispatched = false; + let next_attempt_at = entry.next_attempt_at; + let inserted = state.deadlines.insert((next_attempt_at, identity.clone())); + debug_assert!(inserted, "owner retry deadline must be unique per identity"); + drop(state); + self.counters + .source_evict_retry_scheduled + .fetch_add(1, Ordering::Relaxed); + tracing::debug!( + key = identity.key, + put_time_ms = identity.put_time_ms, + put_version = identity.put_version, + reason, + "owner writeback entered retryable local state" + ); + self.notify.notify_waiters(); } - fn upsert_cached_info(&self, key: &str, candidate: GetCachedInfo) { - match self.get_cached_info.entry(key.to_string()) { - dashmap::mapref::entry::Entry::Occupied(mut entry) => { - if Self::should_replace_cached_info( - entry.get(), - candidate.put_time_ms, - candidate.put_version, - candidate.value.rank(), - ) { - entry.insert(candidate); - } + fn take_due_batch(&self, now: Instant, limit: usize) -> Vec { + let mut state = self.state.lock(); + let mut due = Vec::with_capacity(limit); + while due.len() < limit { + let Some((deadline, identity)) = state.deadlines.iter().next().cloned() else { + break; + }; + if deadline > now { + break; } - dashmap::mapref::entry::Entry::Vacant(entry) => { - entry.insert(candidate); + state.deadlines.remove(&(deadline, identity.clone())); + let Some(entry) = state.entries.get_mut(&identity) else { + debug_assert!(false, "owner retry deadline must reference a live entry"); + continue; + }; + if entry.dispatched || entry.next_attempt_at != deadline { + debug_assert!(false, "owner retry deadline and entry must agree"); + continue; } + due.push(entry.event.clone()); + // Keep the authoritative retry record until the dispatcher has + // atomically pinned the source and installed an inflight guard, + // but emit it only once. A failed dispatcher attempt explicitly + // reschedules it with the next backoff. + entry.dispatched = true; } + due } - pub(crate) fn cache_metadata_only_after_put( - &self, - key: &str, - put_id: crate::master_kv_router::put::PutIDForAKey, - ) { - self.upsert_cached_info( - key, - GetCachedInfo { - put_time_ms: put_id.0, - put_version: put_id.1, - value: CachedValue::MetadataOnly, - }, - ); + fn remove(&self, identity: &OwnerHotReplicaIdentity) { + let entry = { + let mut state = self.state.lock(); + let entry = state.entries.remove(identity); + if let Some(entry) = entry.as_ref() + && !entry.dispatched + { + state + .deadlines + .remove(&(entry.next_attempt_at, identity.clone())); + } + entry + }; + if let Some(entry) = entry { + entry.event.selection_debt.release(); + } } - pub(crate) fn cache_local_replica_after_get( + fn take_for_inflight( &self, - key: &str, - put_id: crate::master_kv_router::put::PutIDForAKey, - memory_info: Arc, - ) { - self.upsert_cached_info( - key, - GetCachedInfo { - put_time_ms: put_id.0, - put_version: put_id.1, - value: CachedValue::LocalReplica(memory_info), - }, - ); + identity: &OwnerHotReplicaIdentity, + ) -> Option { + // The inflight guard takes over the same debt token; do not release it. + let mut state = self.state.lock(); + let entry = state.entries.remove(identity)?; + if !entry.dispatched { + state + .deadlines + .remove(&(entry.next_attempt_at, identity.clone())); + } + Some(entry.event) } - pub(crate) fn short_circuit_put_payload_path_enabled(&self) -> bool { - self.test_spec_config.short_circuit_put_payload_path + fn len(&self) -> usize { + self.state.lock().entries.len() } +} - pub(crate) fn skip_put_end_commit_enabled(&self) -> bool { - self.test_spec_config.skip_put_end_commit - } +pub(crate) struct OwnerPreparedReclaim { + item: crate::master_kv_router::msg_pack::OwnerReclaimItem, + source: OwnerPreparedReclaimSource, + ssd_prepare_lock: Arc>, + ssd_prepare_complete: bool, + ssd_backing: Option, } -#[derive(Debug, Clone)] -pub struct ExternalPendingPutCtx { - pub peer_id: NodeIDString, - pub target_base_addr: u64, - pub target_offset: u64, +pub(crate) enum OwnerPreparedReclaimSource { + /// An owner-indexed source detached from the local Moka while the reclaim fence is active. + Indexed { + cached_info: GetCachedInfo, + local_snapshot: Option, + }, + /// A master-owned allocation that has no owner-local key index. The master route owns the + /// physical allocation through Commit; the owner only reads it under the reclaim fence. + UnindexedAllocation { addr: u64, len: u64 }, } -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct MetricsSet { - pub mean: f64, - pub p99: i64, - pub p95: i64, - pub min: i64, - pub max: i64, - pub timestamps: Vec, +pub(crate) struct OwnerPreparedSsdBacking { + len: u64, + _persist_guard: crate::kv_ssd_storage::KvSsdPersistGuard, } -// Removed StageScope: no longer using stage-scoped gauges; we record -// timestamps (t1..t4) and emit stage success/error directly. +pub(crate) struct OwnerSourceEvictionSelection { + put_id: crate::master_kv_router::put::PutIDForAKey, + cached_info: GetCachedInfo, +} -impl MetricsSet { - /// Convert to Prometheus format string - pub fn to_prometheus_format(&self, metric_name: &str, client_id: &str) -> String { - let mut result = String::new(); +pub(crate) enum OwnerReclaimRecord { + Prepared(OwnerPreparedReclaim), + /// The local index is fenced and the Commit handler owns the detached + /// backing while it updates the slot pool. Keeping this marker in the + /// per-key table lets that O(1) pool update happen without nesting the + /// pool mutex under the key-shard mutex. + Releasing(crate::master_kv_router::msg_pack::OwnerReclaimItem), + Committed(crate::master_kv_router::msg_pack::OwnerReclaimItem), +} - // Traditional aggregated metrics (mean, p99, p95, min, max) - result.push_str(&format!( - "kvcache_{}_mean{{client=\"{}\"}} {}\n", - metric_name, client_id, self.mean - )); +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ExternalPutKeyOutcome { + InFlight, + Succeeded, + Failed, +} - result.push_str(&format!( - "kvcache_{}_p99{{client=\"{}\"}} {}\n", - metric_name, client_id, self.p99 - )); +pub(crate) struct ExternalPutKeySharedOp { + outcome: watch::Sender, +} - result.push_str(&format!( - "kvcache_{}_p95{{client=\"{}\"}} {}\n", - metric_name, client_id, self.p95 - )); +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum OwnerRemotePutOutcome { + InFlight, + Published, + AlreadySatisfied, + Obsolete, + Failed, +} + +/// Shared leader/follower terminal publication used by owner backing writes. +/// Target-specific requests and executors stay outside this type. +struct OwnerTargetPutFlight { + terminal: watch::Sender>, + completed: AtomicBool, +} + +impl OwnerTargetPutFlight { + fn new() -> Self { + let (terminal, _receiver) = watch::channel(None); + Self { + terminal, + completed: AtomicBool::new(false), + } + } + + fn terminal(&self) -> Option { + *self.terminal.borrow() + } + + fn complete(&self, terminal: T) -> bool { + if self + .completed + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return false; + } + self.terminal.send_replace(Some(terminal)); + true + } + + async fn wait(&self) -> Option { + let mut terminal = self.terminal.subscribe(); + loop { + if let Some(terminal) = *terminal.borrow_and_update() { + return Some(terminal); + } + if terminal.changed().await.is_err() { + return None; + } + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct OwnerRemotePutRequest { + pub preferred_sub_cluster: Option, + pub protect_source_on_remote_complete: bool, +} + +/// One owner-initiated remote write for an exact local generation. +/// +/// Every trigger (normal Put, pre-reserved replica, proactive write-back, and +/// tier1) joins this same operation. The trigger is deliberately absent from +/// the identity so policy labels cannot create a second payload transfer. +pub(crate) struct OwnerRemotePutSharedOp { + pub key: String, + pub put_id: crate::master_kv_router::put::PutIDForAKey, + request: Mutex, + flight: OwnerTargetPutFlight, +} + +impl OwnerRemotePutSharedOp { + fn new( + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + preferred_sub_cluster: Option, + protect_source_on_remote_complete: bool, + ) -> Arc { + Arc::new(Self { + key: key.to_string(), + put_id, + request: Mutex::new(OwnerRemotePutRequest { + preferred_sub_cluster, + protect_source_on_remote_complete, + }), + flight: OwnerTargetPutFlight::new(), + }) + } + + fn merge_request( + &self, + preferred_sub_cluster: Option, + protect_source_on_remote_complete: bool, + ) { + let mut request = self.request.lock(); + if request.preferred_sub_cluster.is_none() { + request.preferred_sub_cluster = preferred_sub_cluster; + } + request.protect_source_on_remote_complete |= protect_source_on_remote_complete; + } + + pub(crate) fn request(&self) -> OwnerRemotePutRequest { + self.request.lock().clone() + } + + pub(crate) fn outcome(&self) -> OwnerRemotePutOutcome { + self.flight + .terminal() + .unwrap_or(OwnerRemotePutOutcome::InFlight) + } + + fn complete(&self, outcome: OwnerRemotePutOutcome) -> bool { + debug_assert_ne!(outcome, OwnerRemotePutOutcome::InFlight); + self.flight.complete(outcome) + } + + pub(crate) async fn wait(&self) -> OwnerRemotePutOutcome { + self.flight + .wait() + .await + .unwrap_or(OwnerRemotePutOutcome::Failed) + } +} + +pub(crate) enum OwnerRemotePutReservation { + Leader { + op: Arc, + memory_info: Arc, + }, + Follower(Arc), + SourceUnavailable, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum OwnerLocalSsdPutOutcome { + Published, + AlreadyPresent, + Dropped, + Obsolete, + Failed, +} + +/// One same-owner SSD write for an exact local generation. The common flight +/// owns only terminal publication; admission, copy, durability, and route +/// publication remain in the SSD executor. +pub(crate) struct OwnerLocalSsdPutSharedOp { + pub key: String, + pub put_id: crate::master_kv_router::put::PutIDForAKey, + flight: OwnerTargetPutFlight, +} + +impl OwnerLocalSsdPutSharedOp { + fn new(key: &str, put_id: crate::master_kv_router::put::PutIDForAKey) -> Arc { + Arc::new(Self { + key: key.to_string(), + put_id, + flight: OwnerTargetPutFlight::new(), + }) + } + + fn outcome(&self) -> Option { + self.flight.terminal() + } + + fn complete(&self, outcome: OwnerLocalSsdPutOutcome) -> bool { + self.flight.complete(outcome) + } + + pub(crate) async fn wait(&self) -> OwnerLocalSsdPutOutcome { + self.flight + .wait() + .await + .unwrap_or(OwnerLocalSsdPutOutcome::Failed) + } +} + +pub(crate) enum OwnerLocalSsdPutReservation { + Leader { + op: Arc, + memory_info: Arc, + }, + Follower(Arc), + SourceUnavailable, +} + +impl ExternalPutKeySharedOp { + fn new() -> Arc { + let (outcome, _receiver) = watch::channel(ExternalPutKeyOutcome::InFlight); + Arc::new(Self { outcome }) + } + + fn complete(&self, outcome: ExternalPutKeyOutcome) { + debug_assert_ne!(outcome, ExternalPutKeyOutcome::InFlight); + if *self.outcome.borrow() == ExternalPutKeyOutcome::InFlight { + self.outcome.send_replace(outcome); + } + } + + pub(crate) async fn wait(&self) -> ExternalPutKeyOutcome { + let mut outcome = self.outcome.subscribe(); + loop { + let current = *outcome.borrow_and_update(); + if current != ExternalPutKeyOutcome::InFlight { + return current; + } + if outcome.changed().await.is_err() { + return ExternalPutKeyOutcome::Failed; + } + } + } +} + +pub(crate) enum ExternalLocalFirstPutKeyReservation { + Leader(Arc), + Wait(Arc), + /// A committed owner-local source is between precise source selection and + /// reclaim completion/rollback. The caller owns only a watch receiver: + /// it holds no key fence or physical slot while asynchronously waiting to + /// re-evaluate the complete atomic_batch. + WaitForLocalAccess(watch::Receiver), +} + +#[derive(Default)] +pub(crate) struct OwnerKeyControlState { + local_puts: u32, + /// External Put contexts that may still expose or commit owner-local + /// backing for this key. This counter is maintained by an Arc-backed + /// guard stored in every context, so cache invalidation cannot clear the + /// reclaim fence while a cloned context is still in use. + external_pending_puts: u32, + /// The reject-on-inflight local-first Put leader for this key. Followers + /// subscribe to its terminal result without claiming another slot. + external_put: Option>, + /// Exact-generation owner-side remote Put singleflight. Unlike + /// `external_put`, this remains active after local publication until the + /// remote Start/transfer/Done state machine reaches a terminal outcome. + remote_put: Option>, + /// Exact-generation same-owner SSD write. This is independent from + /// `remote_put`, so both backing writes may run concurrently. + local_ssd_put: Option>, + /// Owner-local pre-Prepare fence installed by the Moka source-eviction + /// dispatcher. The matching committed index is moved into this record, + /// so a new local Get cannot acquire the source between victim selection + /// and the master's reclaim Prepare RPC. + source_eviction_selection: Option, + reclaim: Option, + /// Per-key owner-side Get singleflight marker. It deliberately lives in + /// the same fence as local visibility and reclaim so `R ∩ local`, + /// `R ∩ inflight`, and new leaders are classified atomically. + external_get: Option>, + /// Completion channel for the exact source-selection/reclaim fence. A + /// receiver subscribed under the key-shard lock cannot miss completion, + /// even if the fence clears before the waiter is first polled. + local_access_fence: Option>, +} + +impl OwnerKeyControlState { + fn local_access_fenced(&self) -> bool { + self.source_eviction_selection.is_some() || self.reclaim.is_some() + } + + fn is_idle(&self) -> bool { + self.local_puts == 0 + && self.external_pending_puts == 0 + && self.external_put.is_none() + && self.remote_put.is_none() + && self.local_ssd_put.is_none() + && self.source_eviction_selection.is_none() + && self.reclaim.is_none() + && self.external_get.is_none() + && self.local_access_fence.is_none() + } + + fn begin_local_access_fence(&mut self) { + assert!( + self.local_access_fence.is_none(), + "one key cannot install two owner local-access fence generations" + ); + let (completion, _receiver) = watch::channel(false); + self.local_access_fence = Some(completion); + } + + fn subscribe_local_access_fence(&self) -> watch::Receiver { + assert!( + self.local_access_fenced(), + "local-access waiter requires an active source/reclaim fence" + ); + self.local_access_fence + .as_ref() + .expect("active source/reclaim fence must own a completion channel") + .subscribe() + } + + fn finish_local_access_fence(&mut self) { + assert!( + !self.local_access_fenced(), + "local-access completion cannot publish before the source/reclaim fence clears" + ); + if let Some(completion) = self.local_access_fence.take() { + completion.send_replace(true); + } + } + + fn install_remote_put_leader(&mut self, op: Arc) { + if let Some(displaced) = self.remote_put.replace(op.clone()) { + assert_ne!( + displaced.put_id, op.put_id, + "a matching remote Put generation must join instead of being replaced" + ); + } + } + + fn install_local_ssd_put_leader(&mut self, op: Arc) { + if let Some(displaced) = self.local_ssd_put.replace(op.clone()) { + assert_ne!( + displaced.put_id, op.put_id, + "a matching local SSD Put generation must join instead of being replaced" + ); + } + } +} + +const OWNER_KEY_CONTROL_SHARDS: usize = 256; + +/// Per-key owner fencing without a process-wide mutex. +/// +/// Every operation for one key hashes to the same shard, so local index +/// publication, Get singleflight registration, and reclaim remain linearized. +/// Callers must hold a shard only for one key and must not await, perform RPC, +/// or walk a request batch while holding it. Unrelated keys normally proceed +/// on independent shards; a hash collision only adds a short O(1) critical +/// section and does not change correctness. +pub(crate) struct OwnerKeyControlTable { + shards: Box<[Mutex>]>, +} + +impl Default for OwnerKeyControlTable { + fn default() -> Self { + Self { + shards: (0..OWNER_KEY_CONTROL_SHARDS) + .map(|_| Mutex::new(HashMap::new())) + .collect(), + } + } +} + +impl OwnerKeyControlTable { + fn shard_index(key: &str) -> usize { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + key.hash(&mut hasher); + (hasher.finish() as usize) % OWNER_KEY_CONTROL_SHARDS + } + + pub(crate) fn lock_key( + &self, + key: &str, + ) -> parking_lot::MutexGuard<'_, HashMap> { + self.shards[Self::shard_index(key)].lock() + } +} + +pub(crate) struct ExternalPendingPutFenceGuard { + key: String, + owner_key_control: Arc, + owns_local_put: bool, + local_put_op: Option>, + local_put_succeeded: std::sync::atomic::AtomicBool, + local_slot_cleanup_view: Option, + local_slot_lease: Mutex>, + local_slot_release_failed: std::sync::atomic::AtomicBool, +} + +impl std::fmt::Debug for ExternalPendingPutFenceGuard { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ExternalPendingPutFenceGuard") + .field("key", &self.key) + .field("owns_local_put", &self.owns_local_put) + .finish_non_exhaustive() + } +} + +impl ExternalPendingPutFenceGuard { + pub(crate) fn mark_local_put_succeeded(&self) { + assert!( + self.owns_local_put, + "only a local-first Put can publish a reusable terminal result" + ); + self.local_put_succeeded.store(true, Ordering::Release); + if let Some(op) = self.local_put_op.as_ref() { + op.complete(ExternalPutKeyOutcome::Succeeded); + } + } + + pub(crate) fn attach_local_slot_lease(&self, lease: OwnerLocalReserveSlotLease) { + assert!( + self.owns_local_put, + "only local-first Put owns a slot lease" + ); + assert_eq!( + lease.slots.len(), + 1, + "a pending local-first Put fence must own exactly one slot" + ); + let mut current = self.local_slot_lease.lock(); + assert!( + current.is_none(), + "a pending local-first Put fence cannot replace its slot lease" + ); + *current = Some(lease); + } + + /// Transfer the prepared slot to the precommit/committed MemoryInfo. Once + /// disarmed, dropping the pending context must not return that resident slot + /// to the free list. + pub(crate) fn disarm_local_slot_lease(&self) { + assert!( + self.local_slot_lease.lock().take().is_some(), + "pending local-first Put slot lease is absent while committing" + ); + } + + pub(crate) async fn release_local_slot_lease_now( + &self, + inner: &ClientKvApiInner, + ) -> KvResult<()> { + let lease = self.local_slot_lease.lock().take(); + if let Some(lease) = lease { + if let Err(err) = inner.owner_release_local_reserve_slot_lease(lease).await { + // The lease object has been consumed and the physical state is + // now uncertain. Keep the per-key fence permanently rather + // than allow reclaim to cross a possibly-live prepared slot. + self.local_slot_release_failed + .store(true, Ordering::Release); + return Err(err); + } + } + Ok(()) + } +} + +fn release_external_pending_put_counts( + owner_key_control: &Arc, + key: &str, + owns_local_put: bool, + local_put_op: Option>, + local_put_succeeded: bool, +) { + let mut controls = owner_key_control.lock_key(key); + let remove = { + let state = controls + .get_mut(key) + .expect("external pending Put fence state missing on release"); + state.external_pending_puts = state + .external_pending_puts + .checked_sub(1) + .expect("external pending Put fence counter underflow"); + if owns_local_put { + state.local_puts = state + .local_puts + .checked_sub(1) + .expect("owner local-first Put fence counter underflow"); + if let Some(op) = local_put_op.as_ref() + && state + .external_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, op)) + { + state.external_put = None; + } + } + state.is_idle() + }; + if remove { + controls.remove(key); + } + drop(controls); + if let Some(op) = local_put_op { + op.complete(if local_put_succeeded { + ExternalPutKeyOutcome::Succeeded + } else { + ExternalPutKeyOutcome::Failed + }); + } +} + +fn acquire_external_pending_put_fence_for_key( + owner_key_control: &Arc, + key: &str, +) -> KvResult> { + let mut controls = owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })); + } + let state = controls.entry(key.to_string()).or_default(); + state.external_pending_puts = state + .external_pending_puts + .checked_add(1) + .expect("external pending Put fence counter overflow"); + Ok(Arc::new(ExternalPendingPutFenceGuard { + key: key.to_string(), + owner_key_control: owner_key_control.clone(), + owns_local_put: false, + local_put_op: None, + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + })) +} + +impl Drop for ExternalPendingPutFenceGuard { + fn drop(&mut self) { + let abandoned_slot_lease = self.local_slot_lease.get_mut().take(); + if let Some(lease) = abandoned_slot_lease { + let view = self + .local_slot_cleanup_view + .as_ref() + .expect("local-first Put slot cleanup requires an attached owner view") + .clone(); + let worker_view = view.clone(); + let key = self.key.clone(); + let owner_key_control = self.owner_key_control.clone(); + let owns_local_put = self.owns_local_put; + let local_put_op = self.local_put_op.clone(); + let local_put_succeeded = self.local_put_succeeded.load(Ordering::Acquire); + view.spawn("external_pending_put_slot_drop_cleanup", async move { + if let Err(err) = worker_view + .client_kv_api() + .inner() + .owner_release_local_reserve_slot_lease(lease) + .await + { + tracing::error!("pending local-first Put slot Drop cleanup failed: {}", err); + return; + } + release_external_pending_put_counts( + &owner_key_control, + &key, + owns_local_put, + local_put_op, + local_put_succeeded, + ); + }); + } else if !self.local_slot_release_failed.load(Ordering::Acquire) { + release_external_pending_put_counts( + &self.owner_key_control, + &self.key, + self.owns_local_put, + self.local_put_op.clone(), + self.local_put_succeeded.load(Ordering::Acquire), + ); + } else { + tracing::error!( + "retaining pending Put fence after local slot release failure: key={}", + self.key + ); + } + } +} + +fn allocate_external_holding_id(counter: &AtomicU64) -> u64 { + counter + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + current.checked_add(1) + }) + .expect("external holding id space exhausted") +} + +fn owner_hot_weight_bytes(memory_info: &MemoryInfo) -> u32 { + let bytes = memory_info + .local_reserve_resident_slot_ref() + .map(|(slot_size, _, _)| slot_size) + .unwrap_or(memory_info.len as u64); + u32::try_from(bytes).unwrap_or(u32::MAX) +} + +fn clone_if_owner_hot_entry_matches( + current_put_id: crate::master_kv_router::put::PutIDForAKey, + current: &Arc, + entry_put_id: crate::master_kv_router::put::PutIDForAKey, + entry: &Weak, +) -> Option> { + (current_put_id == entry_put_id && Weak::ptr_eq(entry, &Arc::downgrade(current))) + .then(|| current.clone()) +} + +fn owner_hot_source_has_active_holders(selected_source: &Arc) -> bool { + // A reclaimable committed source has exactly two strong references here: + // one in get_cached_info and one temporary selection pin. Any additional + // reference belongs to an active local reader/transfer. Sending such a + // source to the master would only make Prepare return Busy while its + // selection debt suppresses choosing a different, reclaimable victim. + Arc::strong_count(selected_source) > 2 +} + +enum OwnerHotPinResult { + Pinned(Arc), + ReclaimBusy, + Stale, +} + +fn pin_current_owner_hot_source_from_index( + entry_put_id: crate::master_kv_router::put::PutIDForAKey, + entry: &Weak, + resolve_current: impl FnOnce() -> Option<(crate::master_kv_router::put::PutIDForAKey, Arc)>, +) -> OwnerHotPinResult { + let Some((current_put_id, current)) = resolve_current() else { + // Do not upgrade the listener's Weak after the local index has + // disappeared. A reclaim Prepare may already own the sole Arc and + // Commit relies on that ownership. A still-live Weak therefore means + // "retry after the transient owner transition"; a dead Weak is + // definitively obsolete. + return if entry.strong_count() == 0 { + OwnerHotPinResult::Stale + } else { + OwnerHotPinResult::ReclaimBusy + }; + }; + let Some(pinned) = + clone_if_owner_hot_entry_matches(current_put_id, ¤t, entry_put_id, entry) + else { + return OwnerHotPinResult::Stale; + }; + + // `resolve_current` clones from the DashMap entry while its shard read + // guard is held. Reclaim Prepare cannot remove that entry until the clone + // exists, and its strong-count check will consequently return Busy. This + // gives us the required per-key pin without the global owner-control lock. + Some(pinned).map_or(OwnerHotPinResult::Stale, OwnerHotPinResult::Pinned) +} + +fn pin_current_owner_hot_source( + key: &str, + entry: &OwnerHotCacheEntry, + get_cached_info: &DashMap, + counters: &OwnerHotCacheCounters, +) -> OwnerHotPinResult { + let result = pin_current_owner_hot_source_from_index(entry.put_id, &entry.memory_info, || { + get_cached_info.get(key).map(|cached| { + ( + (cached.put_time_ms, cached.put_version), + cached.mem_holder.clone(), + ) + }) + }); + match result { + OwnerHotPinResult::Pinned(pinned) => OwnerHotPinResult::Pinned(pinned), + OwnerHotPinResult::ReclaimBusy => { + counters.skipped_reclaim.fetch_add(1, Ordering::Relaxed); + OwnerHotPinResult::ReclaimBusy + } + OwnerHotPinResult::Stale => { + counters.skipped_stale.fetch_add(1, Ordering::Relaxed); + OwnerHotPinResult::Stale + } + } +} + +fn build_owner_hot_cache( + capacity_bytes: u64, + counters: Arc, + retry_queue: Arc, + eviction_tx: tokio::sync::ampsc::UnboundedSender, +) -> OwnerHotCache { + assert!( + capacity_bytes > 0, + "owner hot-cache capacity must be positive" + ); + OwnerHotCache::builder(capacity_bytes) + .weigher(|_key: &String, entry: &OwnerHotCacheEntry| entry.weight_bytes) + .eviction_listener(move |key, entry, cause| { + if cause != RemovalCause::Size { + return; + } + counters.size_evictions.fetch_add(1, Ordering::Relaxed); + let identity = OwnerHotReplicaIdentity { + key: (*key).clone(), + put_time_ms: entry.put_id.0, + put_version: entry.put_id.1, + }; + let selection_debt = OwnerHotSelectionDebt::new( + u64::from(entry.weight_bytes), + counters.selection_debt_bytes.clone(), + ); + let event = OwnerHotEvictionEvent { + key: identity.key.clone(), + put_id: entry.put_id, + memory_info: entry.memory_info.clone(), + selection_debt, + retry: false, + source_eviction_victim: None, + retry_failures: 0, + }; + if let Err(err) = eviction_tx.send(OwnerHotEvictionDispatch::Victim(event)) { + let OwnerHotEvictionDispatch::Victim(event) = err.0 else { + unreachable!("the Moka listener only sends victim events") + }; + counters + .source_evict_dispatch_failed + .fetch_add(1, Ordering::Relaxed); + retry_queue.schedule(event, "eviction dispatcher closed"); + } + }) + .build() +} + +struct ClientKvApiViewHolder { + view: OnceLock, +} + +impl ClientKvApiViewHolder { + fn new() -> Self { + Self { + view: OnceLock::new(), + } + } + + fn attach(&self, view: ClientKvApiView) { + // The framework attaches a module's PostView exactly once at the init barrier. + // A second attach indicates a programming error. + self.view + .set(view) + .unwrap_or_else(|_| panic!("ClientKvApi view attached twice")); + } + + fn clone_view(&self) -> ClientKvApiView { + self.view.get().unwrap().clone() + } +} + +fn now_unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u64::MAX as u128) as u64 +} + +impl std::ops::Deref for ClientKvApiViewHolder { + type Target = ClientKvApiView; + + fn deref(&self) -> &Self::Target { + self.view.get().unwrap() + } +} + +pub struct ClientKvApiInner { + view: ClientKvApiViewHolder, + test_spec_config: TestSpecConfig, + ssd_storage: Option>, + metrics: OnceLock>, + + /// make sure each remote kv get run in order + pub get_remote_kv_lock: AMapLock, + /// key -> value info on this node + /// we can only remove value if it's put_time_ms and put_version match remote eviction command + get_cached_info: Arc>, + /// key -> locally readable resident slot before backend put_start/put_done finishes. + precommit_local_visible_info: DashMap, + /// Transferred Get targets awaiting an idempotent master GetDone result. + /// These entries fence reclaim but are never visible to readers. + pending_local_get_info: DashMap, + /// key -> local replica version remembered from local put/get durable-replica success. + /// This authority is positive-only: hit means "can answer exists=true immediately when + /// allow_local_snapshot is enabled"; miss does not imply non-existence. + local_snapshot_info: DashMap, + /// KvOwner-managed resident staging grants for hostless put_start. + owner_local_reserve_pool: Mutex, + /// Serialize claims only within one slot-size class. Tokio's mutex provides FIFO + /// acquisition order for equal-size waiters without making an unrelated class wait + /// behind a pressured class. + owner_local_reserve_claim_locks: DashMap>>, + /// Wake the background reserve actor after demand/free-slot changes. + owner_local_reserve_rebalance_notify: Arc, + /// Owner-local write-back put ids for external local-first path. + external_local_first_put_id_counter: AtomicU32, + /// Correlates idempotent owner source-eviction batches in logs and RPC responses. + next_owner_source_eviction_operation_id: AtomicU64, + /// Sharded per-key gate for local-first puts, local index access, and reclaim fencing. + owner_key_control: Arc, + /// A weak-value admission/recency tier. It never owns resident memory and + /// therefore cannot become a second physical-reclaim authority. + owner_hot_cache: Option, + /// Exact selected identities remain here until their backing is physically + /// Free or the source is restored to owner-hot. + owner_source_eviction_selected: + Arc>>, + owner_hot_counters: Arc, + owner_remote_put_counters: Arc, + owner_local_ssd_put_counters: Arc, + planned_get_counters: OwnerPlannedGetCounters, + ssd_stage_counters: OwnerSsdStageCounters, + owner_hot_retry_queue: Arc, + owner_hot_eviction_tx: tokio::sync::ampsc::UnboundedSender, + owner_hot_eviction_rx: + Mutex>>, + + /// Shared delete actor input for owner -> external weak-index invalidation. + pub external_invalidate_delete: EnsureMemholderMgmtDeleteHandle, + /// Shared delete actor input for owner -> master delete-ack batching. + pub delete_ack_batch: EnsureMemholderMgmtDeleteHandle, + /// Shared manager for owner -> master delete-ack batching. + pub owner_delete_ack_mgr: OwnerDeleteAckMemMgr, + + // record external_client get_holding info (owned, flattened manager) + pub external_get_holding: OwnerExternalMemMgr, + pub external_get_start_registry: DashMap, + /// Metrics-only weak index of active per-key Get flights. Correctness + /// remains in the sharded key-control table; observing metrics must never + /// scan or hold those fences. + external_get_flight_registry: DashMap>, + external_get_local_probe_locks: AMapLock<(String, i64, u64)>, + completed_external_get_local_probes: + moka::future::Cache<(String, i64, u64), (Vec, ExternalBatchGetLocalProbeResp)>, + planned_external_get_execute_locks: AMapLock<(String, i64, u64)>, + completed_planned_external_get_executes: + moka::future::Cache<(String, i64, u64), ExternalExecutePlannedGetResp>, + /// Exact `get_id` SSD source operations. RPC retransmits and concurrent + /// callers join one disk read plus one payload transfer. + ssd_stage_flights: DashMap>, + /// Short-lived terminal replay closes the active-map removal race and + /// prevents a lost response from re-reading or re-transferring the value. + completed_ssd_stages: moka::future::Cache, + next_external_get_start_handle: AtomicU64, + /// External holding identities are independent from upstream and resident holder ids. + next_external_holding_id: AtomicU64, + /// Weak handle to a shared refcount tracker for all UserMemHolder of this client. + /// + /// - A strong `Arc` is given to every `UserMemHolder` created by this client. + /// - When the last `UserMemHolder` is dropped, the strong `Arc` is dropped too, + /// and this weak handle will no longer upgrade, meaning the client can be safely dropped. + /// - Stored as `Weak` in `OnceLock` to avoid cycles and allow lazy initialization. + pub all_memholder_refcount: OnceLock>, + /// External API is implemented directly on ClientKvApi; no handler stored here + + #[cfg(test)] + test_record: crate::client_kv_api::client_test_record::ClientTestRecord, + + rpc_caller_get_start: RPCCaller, + rpc_caller_get_revoke: RPCCaller, + rpc_caller_get_done: RPCCaller, + rpc_caller_batch_get_start: RPCCaller, + rpc_caller_batch_get_bind: RPCCaller, + rpc_caller_batch_get_revoke: RPCCaller, + rpc_caller_batch_get_done: RPCCaller, + rpc_caller_put_start: RPCCaller, + rpc_caller_put_revoke: RPCCaller, + rpc_caller_put_done: RPCCaller, + rpc_caller_batch_put_start: RPCCaller, + rpc_caller_batch_put_revoke: RPCCaller, + rpc_caller_batch_put_done: RPCCaller, + rpc_caller_grouped_batch_put_done: RPCCaller, + rpc_caller_batch_prepare_put_keys: RPCCaller, + rpc_caller_batch_release_put_key_reservations: RPCCaller, + rpc_caller_put_append_start: RPCCaller, + rpc_caller_batch_put_append_start: RPCCaller, + rpc_caller_put_append_revoke: RPCCaller, + rpc_caller_put_append_done: RPCCaller, + rpc_caller_batch_put_append_done: RPCCaller, + rpc_caller_batch_evict_owner_source: RPCCaller, + rpc_caller_batch_publish_owner_ssd: RPCCaller, + rpc_caller_reserve_local_grant: RPCCaller, + rpc_caller_release_local_grant: RPCCaller, + rpc_caller_delete: RPCCaller, + rpc_caller_batch_delete_ack: RPCCaller, + rpc_caller_batch_is_exist: RPCCaller, + rpc_caller_get_meta: RPCCaller, + rpc_caller_allocate_client_lease: RPCCaller, + rpc_caller_client_lease_keepalive: RPCCaller, + rpc_caller_ssd_stage_read: RPCCaller, + rpc_caller_ssd_stage_begin: RPCCaller, + rpc_caller_ssd_stage_done: RPCCaller, + rpc_caller_external_put_commit: RPCCaller, + rpc_caller_external_put_revoke: RPCCaller, + rpc_caller_resolve_side_transfer_lane: RPCCaller, + + /// Default lease id recorded for inspection/convenience, but NOT auto-applied. + /// Callers must explicitly pass `Some(lease_id)` to attach a put to a lease. + default_lease_id: parking_lot::RwLock>, + /// External put (remote target) pending context keyed by (key, put_time_ms, put_version). + /// 注意：put_id (time_ms,version) 在不同 key 上并不全局唯一，因此必须携带 key 作为索引的一部分，避免碰撞。 + /// 使用 moka::sync::SegmentedCache 并设置 30 分钟 TTL，避免异常路径未清理导致的泄漏；不设置容量上限，纯 TTL 控制。 + external_pending_puts: moka::sync::SegmentedCache<(String, u64, u32), ExternalPendingPutCtx>, + owner_local_publish_tx: tokio::sync::ampsc::Sender, + owner_local_publish_rx: Mutex>>, +} + +impl ClientKvApiInner { + fn view(&self) -> &ClientKvApiView { + &self.view + } + + pub(crate) async fn persist_local_kvs_to_ssd( + &self, + sources: &[KvSsdPersistSource], + ) -> KvResult>>> { + let Some(store) = self.ssd_storage.as_ref() else { + return Ok(sources.iter().map(|_| Ok(None)).collect()); + }; + let segment_guard = self.view.client_seg_pool().cpu_mem_read_guard().await?; + for source in sources { + if !segment_guard.contains_rw_or_ro(source.addr, source.len) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "SSD persist source is outside the local segment: key={} put_id=({},{}) addr={:#x} len={}", + source.key, source.put_id.0, source.put_id.1, source.addr, source.len + ), + })); + } + } + let results = store.persist_batch_from_addrs(sources).await; + drop(segment_guard); + Ok(results) + } + + pub(crate) async fn copy_local_kvs_for_ssd( + &self, + sources: &[KvSsdPersistSource], + ) -> KvResult>> { + if self.ssd_storage.is_none() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "local SSD storage is disabled".to_string(), + })); + } + let segment_guard = self.view.client_seg_pool().cpu_mem_read_guard().await?; + for source in sources { + if !segment_guard.contains_rw_or_ro(source.addr, source.len) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "SSD persist source is outside the local segment: key={} put_id=({},{}) addr={:#x} len={}", + source.key, source.put_id.0, source.put_id.1, source.addr, source.len + ), + })); + } + } + let copies = KvSsdStorage::copy_batch_from_addrs(sources); + drop(segment_guard); + Ok(copies) + } + + pub(crate) async fn persist_copied_local_kvs_to_ssd( + &self, + permit: KvSsdPersistBatchPermit, + copies: Vec>, + ) -> KvResult>>> { + let Some(store) = self.ssd_storage.as_ref() else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "local SSD storage is disabled".to_string(), + })); + }; + Ok(store + .persist_batch_from_copies_with_permit(permit, copies) + .await) + } + + pub(crate) fn try_acquire_local_ssd_persist_batch( + &self, + item_count: usize, + ) -> KvResult> { + let Some(store) = self.ssd_storage.as_ref() else { + return Ok(None); + }; + store.try_acquire_persist_batch(item_count) + } + + async fn discard_local_ssd_replica( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) -> bool { + let Some(store) = self.ssd_storage.as_ref() else { + return false; + }; + store.remove_exact(key, put_id).await + } + + async fn begin_ssd_stage(&self, get_id: u64) -> KvResult { + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let response = self + .rpc_caller_ssd_stage_begin + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: SsdStageBeginReq { get_id }, + raw_bytes: Vec::new(), + }, + Some(SSD_STAGE_RPC_TIMEOUT), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + response.serialize_part.error_code, + response.serialize_part.error_json, + )?; + Ok(response.serialize_part.started) + } + + async fn finish_ssd_stage_once(&self, get_id: u64, drop_ssd_source: bool) -> KvResult<()> { + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let response = self + .rpc_caller_ssd_stage_done + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: SsdStageDoneReq { + get_id, + drop_ssd_source, + }, + raw_bytes: Vec::new(), + }, + Some(SSD_STAGE_RPC_TIMEOUT), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + response.serialize_part.error_code, + response.serialize_part.error_json, + ) + } + + async fn finish_ssd_stage_until_acked(&self, get_id: u64, drop_ssd_source: bool) { + let mut attempt = 0_u64; + let mut backoff = SSD_STAGE_DONE_RETRY_INITIAL_BACKOFF; + loop { + attempt = attempt.saturating_add(1); + match self.finish_ssd_stage_once(get_id, drop_ssd_source).await { + Ok(()) => return, + Err(err) => { + if attempt == 1 || attempt.is_power_of_two() { + tracing::warn!( + get_id, + drop_ssd_source, + attempt, + retry_delay_ms = backoff.as_millis(), + error = %err, + "SSD StageDone failed; retaining the source flight and retrying" + ); + } + } + } + tokio::time::sleep(backoff).await; + backoff = backoff + .saturating_mul(2) + .min(SSD_STAGE_DONE_RETRY_MAX_BACKOFF); + } + } + + fn finish_ssd_stage_detached(&self, get_id: u64) { + self.ssd_stage_counters + .done_detached + .fetch_add(1, Ordering::Relaxed); + let spawn_view = self.view.clone_view(); + let task_view = spawn_view.clone(); + spawn_view.spawn("ssd_stage_done_retry", async move { + task_view + .client_kv_api() + .inner() + .finish_ssd_stage_until_acked(get_id, false) + .await; + }); + } + + async fn run_ssd_stage_once(&self, req: &SsdStageReadReq) -> SsdStageReadResp { + self.ssd_stage_counters + .ready_requests + .fetch_add(1, Ordering::Relaxed); + let ready_started_at = Instant::now(); + match self.begin_ssd_stage(req.get_id).await { + Ok(true) => {} + Ok(false) => { + let response = ssd_stage_error_response(KvError::Api(ApiError::InvalidArgument { + detail: format!("SSD stage is not startable: get_id={}", req.get_id), + })); + self.ssd_stage_counters + .ready_failures + .fetch_add(1, Ordering::Relaxed); + self.ssd_stage_counters.ready_duration_us.fetch_add( + u64::try_from(ready_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + return response; + } + Err(err) => { + let response = ssd_stage_error_response(err); + self.ssd_stage_counters + .ready_failures + .fetch_add(1, Ordering::Relaxed); + self.ssd_stage_counters.ready_duration_us.fetch_add( + u64::try_from(ready_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + return response; + } + } + + let load_result = async { + let Some(store) = self.ssd_storage.as_ref() else { + return Err(KvError::Api(ApiError::KeyNotFound { + key: req.key.clone(), + })); + }; + let segment_guard = self.view.client_seg_pool().cpu_mem_read_guard().await?; + if !segment_guard.contains_rw(req.stage_addr, req.stage_capacity) { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "SSD stage is outside the local writable segment: get_id={} addr={:#x} capacity={}", + req.get_id, req.stage_addr, req.stage_capacity + ), + })); + } + store + .load_into_addr( + &req.key, + req.put_id, + req.stage_addr, + req.len, + req.stage_capacity, + ) + .await?; + drop(segment_guard); + Ok(()) + } + .await; + + let response = match load_result { + Ok(()) => SsdStageReadResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }, + Err(load_err) => { + let stale_ssd_source = + matches!(&load_err, KvError::Api(ApiError::KeyNotFound { .. })); + // A failed load never becomes pull-ready. The source owner is + // therefore the only side that can safely close the stage. A + // true miss also removes the exact stale SSD route. + self.finish_ssd_stage_until_acked(req.get_id, stale_ssd_source) + .await; + ssd_stage_error_response(load_err) + } + }; + self.ssd_stage_counters.ready_duration_us.fetch_add( + u64::try_from(ready_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + if response.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK { + self.ssd_stage_counters + .ready_successes + .fetch_add(1, Ordering::Relaxed); + } else { + self.ssd_stage_counters + .ready_failures + .fetch_add(1, Ordering::Relaxed); + } + response + } + + async fn execute_ssd_stage(&self, req: &SsdStageReadReq) -> SsdStageReadResp { + if let Some(completed) = self.completed_ssd_stages.get(&req.get_id).await { + return if completed.request == *req { + completed.response + } else { + ssd_stage_request_mismatch_response(&completed.request, req) + }; + } + + let (op, is_leader) = match self.ssd_stage_flights.entry(req.get_id) { + DashMapEntry::Occupied(entry) => (entry.get().clone(), false), + DashMapEntry::Vacant(entry) => { + let op = SsdStageSharedOp::new(req.clone()); + entry.insert(op.clone()); + (op, true) + } + }; + if op.request != *req { + return ssd_stage_request_mismatch_response(&op.request, req); + } + + if is_leader { + let spawn_view = self.view.clone_view(); + let task_view = spawn_view.clone(); + let task_op = op.clone(); + spawn_view.spawn("ssd_stage_singleflight", async move { + let inner = task_view.client_kv_api().inner(); + let response = inner.run_ssd_stage_once(&task_op.request).await; + inner + .ssd_stage_counters + .execute_completions + .fetch_add(1, Ordering::Relaxed); + assert!( + task_op.complete(response.clone()), + "one SSD stage flight must publish exactly one terminal result" + ); + inner + .ssd_stage_counters + .terminal_published + .fetch_add(1, Ordering::Relaxed); + + // Wake the foreground RPC before maintaining the replay cache. + // The completed flight remains indexed until the cache insert + // finishes, so retransmits cannot start a second disk read. + let cache_started_at = Instant::now(); + inner + .completed_ssd_stages + .insert( + task_op.request.get_id, + CompletedSsdStage { + request: task_op.request.clone(), + response: response.clone(), + }, + ) + .await; + inner + .ssd_stage_counters + .terminal_cache_duration_us + .fetch_add( + u64::try_from(cache_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + inner + .ssd_stage_counters + .terminal_cache_inserts + .fetch_add(1, Ordering::Relaxed); + inner + .ssd_stage_flights + .remove_if(&task_op.request.get_id, |_, current| { + Arc::ptr_eq(current, &task_op) + }); + }); + } + op.wait().await + } + + pub(crate) async fn stage_kv_from_ssd_source( + &self, + source_node_id: &NodeIDString, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + get_id: u64, + stage_addr: u64, + stage_capacity: u64, + target_addr: u64, + len: u64, + ) -> KvResult<()> { + let req = SsdStageReadReq { + key: key.to_string(), + put_id, + get_id, + stage_addr, + stage_capacity, + len, + }; + let self_node_id = self.view.cluster_manager().get_self_info().id.clone(); + self.ssd_stage_counters + .source_ready_wait_requests + .fetch_add(1, Ordering::Relaxed); + let ready_wait_started_at = Instant::now(); + let response_result: KvResult = if source_node_id == &self_node_id { + Ok(self.execute_ssd_stage(&req).await) + } else { + self.rpc_caller_ssd_stage_read + .call_with_transport_policy( + self.view.p2p_module(), + source_node_id.clone().into(), + MsgPack { + serialize_part: req, + raw_bytes: Vec::new(), + }, + Some(SSD_STAGE_RPC_TIMEOUT), + RpcTransportPolicy::ForceTransport, + 1, + ) + .await + .map(|response| response.serialize_part) + .map_err(KvError::from) + }; + self.ssd_stage_counters + .source_ready_wait_duration_us + .fetch_add( + u64::try_from(ready_wait_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + let ready_result = response_result.and_then(|response| { + crate::rpcresp_kvresult_convert::try_from_code(response.error_code, response.error_json) + }); + if ready_result.is_ok() { + self.ssd_stage_counters + .source_ready_wait_successes + .fetch_add(1, Ordering::Relaxed); + } else { + self.ssd_stage_counters + .source_ready_wait_failures + .fetch_add(1, Ordering::Relaxed); + } + ready_result?; + + self.ssd_stage_counters + .target_pull_requests + .fetch_add(1, Ordering::Relaxed); + let pull_started_at = Instant::now(); + let peer_id = (source_node_id != &self_node_id).then(|| source_node_id.clone()); + let transfer_result = self + .view + .client_transfer_engine() + .transfer_data_no_copy(peer_id, true, stage_addr, target_addr, len, None) + .await + .map(|_| ()) + .map_err(|err| { + KvError::Api(ApiError::Transfer { + from_addr: stage_addr, + to_addr: target_addr, + len, + error: err.to_string(), + }) + }); + self.ssd_stage_counters.target_pull_duration_us.fetch_add( + u64::try_from(pull_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + if transfer_result.is_ok() { + self.ssd_stage_counters + .target_pull_successes + .fetch_add(1, Ordering::Relaxed); + } else { + self.ssd_stage_counters + .target_pull_failures + .fetch_add(1, Ordering::Relaxed); + } + + // Payload completion is the ownership hand-off point: the target no + // longer reads the source stage, so master may release it. Do not put + // this control-plane ACK on the foreground Get latency path. + self.finish_ssd_stage_detached(get_id); + transfer_result + } + + pub(crate) fn kv_ssd_storage_usage_snapshot( + &self, + ) -> Option { + self.ssd_storage + .as_ref() + .map(|store| store.usage_snapshot()) + } + + pub(crate) fn track_external_get_flight(&self, op: &Arc) { + self.external_get_flight_registry + .insert(op.key.clone(), Arc::downgrade(op)); + } + + pub(crate) fn untrack_external_get_flight(&self, op: &Arc) { + let weak = Arc::downgrade(op); + self.external_get_flight_registry + .remove_if(&op.key, |_, current| Weak::ptr_eq(current, &weak)); + } + + fn external_get_flight_snapshot(&self) -> Vec> { + let mut ops = Vec::new(); + let mut stale = Vec::new(); + for entry in &self.external_get_flight_registry { + if let Some(op) = entry.value().upgrade() { + ops.push(op); + } else { + stale.push(entry.key().clone()); + } + } + for key in stale { + self.external_get_flight_registry + .remove_if(&key, |_, weak| weak.strong_count() == 0); + } + ops + } + + pub(crate) fn owner_hot_prepare_eviction( + &self, + event: &OwnerHotEvictionEvent, + ) -> OwnerHotEvictionPreparation { + let trigger = OwnerHotReplicaIdentity { + key: event.key.clone(), + put_time_ms: event.put_id.0, + put_version: event.put_id.1, + }; + let cache_entry = OwnerHotCacheEntry { + put_id: event.put_id, + memory_info: event.memory_info.clone(), + weight_bytes: 0, + }; + let memory_info = match pin_current_owner_hot_source( + event.key.as_str(), + &cache_entry, + self.get_cached_info.as_ref(), + self.owner_hot_counters.as_ref(), + ) { + OwnerHotPinResult::Pinned(memory_info) => memory_info, + OwnerHotPinResult::ReclaimBusy => { + return OwnerHotEvictionPreparation::RetryableReclaimFence; + } + OwnerHotPinResult::Stale => return OwnerHotEvictionPreparation::Obsolete, + }; + + if owner_hot_source_has_active_holders(&memory_info) { + OwnerHotEvictionPreparation::TemporarilyPinned + } else { + OwnerHotEvictionPreparation::Ready { + trigger, + source: memory_info, + } + } + } + + pub(crate) fn owner_hot_restore_source_selection( + &self, + identity: &OwnerHotReplicaIdentity, + ) -> bool { + let mut controls = self.owner_key_control.lock_key(&identity.key); + let Some(state) = controls.get_mut(&identity.key) else { + return false; + }; + let matches = state + .source_eviction_selection + .as_ref() + .is_some_and(|selection| { + selection.put_id == (identity.put_time_ms, identity.put_version) + }); + if !matches { + return false; + } + let selection = state + .source_eviction_selection + .take() + .expect("matching owner source selection must exist"); + let replaced = self + .get_cached_info + .insert(identity.key.clone(), selection.cached_info); + assert!( + replaced.is_none(), + "rolling back an owner source selection must restore an empty local index" + ); + state.finish_local_access_fence(); + if state.is_idle() { + controls.remove(&identity.key); + } + true + } + + fn owner_hot_install_source_selection_debt( + &self, + identity: OwnerHotReplicaIdentity, + debt: Arc, + ) -> bool { + match self.owner_source_eviction_selected.entry(identity) { + DashMapEntry::Vacant(entry) => { + entry.insert(debt.clone()); + self.owner_hot_counters + .add_source_eviction_selected_bytes(debt.weight_bytes); + true + } + DashMapEntry::Occupied(_) => false, + } + } + + fn owner_hot_remove_source_selection_debt( + &self, + identity: &OwnerHotReplicaIdentity, + ) -> Option> { + let debt = self + .owner_source_eviction_selected + .remove(identity) + .map(|(_, debt)| debt)?; + self.owner_hot_counters + .remove_source_eviction_selected_bytes(debt.weight_bytes); + Some(debt) + } + + pub(crate) fn owner_hot_install_source_selection_fence( + &self, + identity: &OwnerHotReplicaIdentity, + source: &Arc, + ) -> OwnerHotSelectionFenceOutcome { + let mut controls = self.owner_key_control.lock_key(&identity.key); + let control_busy = controls.get(&identity.key).is_some_and(|state| { + state.local_puts != 0 + || state.external_pending_puts != 0 + || state.remote_put.is_some() + || state.local_ssd_put.is_some() + || state.external_get.is_some() + || state.local_access_fenced() + }); + if control_busy + || self + .precommit_local_visible_info + .contains_key(&identity.key) + || self.pending_local_get_info.contains_key(&identity.key) + { + return OwnerHotSelectionFenceOutcome::Retryable; + } + + let cached_info = self + .get_cached_info + .remove_if(&identity.key, |_, cached| { + (cached.put_time_ms, cached.put_version) + == (identity.put_time_ms, identity.put_version) + && Arc::ptr_eq(&cached.mem_holder, source) + }) + .map(|(_, cached)| cached); + let Some(cached_info) = cached_info else { + return OwnerHotSelectionFenceOutcome::Obsolete; + }; + + let state = controls.entry(identity.key.clone()).or_default(); + assert!(state.source_eviction_selection.is_none()); + assert!(state.reclaim.is_none()); + state.begin_local_access_fence(); + state.source_eviction_selection = Some(OwnerSourceEvictionSelection { + put_id: (identity.put_time_ms, identity.put_version), + cached_info, + }); + drop(controls); + + // The index Arc moved into source_eviction_selection. The temporary + // source Arc is the second expected reference; any extra Arc is an + // active reader or transfer that arrived before the fence was installed. + if owner_hot_source_has_active_holders(source) { + assert!( + self.owner_hot_restore_source_selection(identity), + "a pinned single-key victim must restore its source fence" + ); + OwnerHotSelectionFenceOutcome::TemporarilyPinned + } else { + OwnerHotSelectionFenceOutcome::Fenced + } + } + + fn owner_hot_track_committed( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: &Arc, + ) { + let Some(cache) = self.owner_hot_cache.as_ref() else { + return; + }; + if !self.owner_hot_source_is_current(key, put_id, memory_info) { + return; + } + cache.insert( + key.to_string(), + [OwnerHotPinAlias::new(memory_info)], + OwnerHotCacheEntry { + put_id, + memory_info: Arc::downgrade(memory_info), + weight_bytes: owner_hot_weight_bytes(memory_info.as_ref()), + }, + ); + if !self.owner_hot_source_is_current(key, put_id, memory_info) { + self.owner_hot_invalidate_version(key, put_id); + } + } + + pub(crate) fn owner_hot_admit_published_committed( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) -> bool { + let memory_info = self.get_cached_info.get(key).and_then(|cached| { + ((cached.put_time_ms, cached.put_version) == put_id).then(|| cached.mem_holder.clone()) + }); + let Some(memory_info) = memory_info else { + return false; + }; + self.owner_hot_track_committed(key, put_id, &memory_info); + true + } + + fn owner_hot_touch_or_promote( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: &Arc, + ) { + let Some(cache) = self.owner_hot_cache.as_ref() else { + return; + }; + if cache.get(&key.to_string()).is_some_and(|entry| { + entry.put_id == put_id && Weak::ptr_eq(&entry.memory_info, &Arc::downgrade(memory_info)) + }) { + return; + } + self.owner_hot_track_committed(key, put_id, memory_info); + } + + pub(crate) fn owner_hot_source_is_current( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: &Arc, + ) -> bool { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return false; + } + self.get_cached_info.get(key).is_some_and(|cached| { + cached.put_time_ms == put_id.0 + && cached.put_version == put_id.1 + && Arc::ptr_eq(&cached.mem_holder, memory_info) + }) + } + + pub(crate) fn owner_hot_invalidate_version( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) { + let identity = OwnerHotReplicaIdentity { + key: key.to_string(), + put_time_ms: put_id.0, + put_version: put_id.1, + }; + if let Some(cache) = self.owner_hot_cache.as_ref() { + cache.invalidate_if(&key.to_string(), |entry| entry.put_id == put_id); + } + if let Some(debt) = self.owner_hot_remove_source_selection_debt(&identity) { + debt.release(); + self.owner_hot_counters + .source_evict_committed_members + .fetch_add(1, Ordering::Relaxed); + } + self.owner_hot_retry_queue.remove(&identity); + } + + pub(crate) fn release_local_reserve_route_for_memory_info(&self, memory_info: &MemoryInfo) { + let Some((slot_size, grant_id, slot_index)) = memory_info.local_reserve_resident_slot_ref() + else { + return; + }; + if let Err(err) = + self.owner_release_local_reserve_committed_slot_route(slot_size, grant_id, slot_index) + { + tracing::warn!( + "failed to release local reserve committed slot route: key={} slot_size={} grant_id={} slot_index={} err={}", + memory_info.key, + slot_size, + grant_id, + slot_index, + err + ); + } + } + + pub(crate) fn owner_hot_pin_memory_info( + &self, + memory_info: &Arc, + ) -> Option { + self.owner_hot_cache + .as_ref()? + .try_pin_alias(OwnerHotPinAlias::new(memory_info)) + } + + pub(crate) fn short_circuit_put_payload_path_enabled(&self) -> bool { + self.test_spec_config.short_circuit_put_payload_path + } + + pub(crate) fn skip_put_end_commit_enabled(&self) -> bool { + self.test_spec_config.skip_put_end_commit + } + + pub(crate) fn next_external_local_first_put_id( + &self, + ) -> crate::master_kv_router::put::PutIDForAKey { + ( + now_unix_ms(), + self.external_local_first_put_id_counter + .fetch_add(1, Ordering::Relaxed), + ) + } + + pub fn next_owner_local_first_put_id(&self) -> crate::master_kv_router::put::PutIDForAKey { + self.next_external_local_first_put_id() + } + + pub async fn enqueue_owner_local_publish(&self, job: OwnerLocalPublishJob) -> KvResult<()> { + let key_count = job.items.len(); + let first_key = job + .items + .first() + .map(|item| item.key.as_str()) + .unwrap_or("") + .to_string(); + self.owner_local_publish_tx.try_send(job).map_err(|err| { + KvError::Api(ApiError::Unknown { + detail: format!( + "owner local publish queue is full or closed: first_key={} key_count={} err={}", + first_key, key_count, err + ), + }) + }) + } + + pub(crate) fn reserve_external_local_first_put_key( + &self, + key: &str, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + ) -> KvResult { + let mut controls = self.owner_key_control.lock_key(key); + let reusable_singleflight = reject_if_inflight_same_key && reject_if_exist_same_key; + if let Some(state) = controls.get(key) + && state.local_access_fenced() + { + return if reusable_singleflight { + Ok(ExternalLocalFirstPutKeyReservation::WaitForLocalAccess( + state.subscribe_local_access_fence(), + )) + } else { + Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })) + }; + } + if reject_if_exist_same_key + && (self.precommit_local_visible_info.contains_key(key) + || self.pending_local_get_info.contains_key(key) + || self.get_cached_info.contains_key(key) + || self.local_snapshot_info.contains_key(key)) + { + return Err(KvError::Api(ApiError::KeyAlreadyExists { + key: key.to_string(), + })); + } + let state = controls.entry(key.to_string()).or_default(); + if reject_if_inflight_same_key && state.local_puts > 0 { + return (reusable_singleflight) + .then(|| state.external_put.clone()) + .flatten() + .map_or_else( + || { + Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })) + }, + |op| Ok(ExternalLocalFirstPutKeyReservation::Wait(op)), + ); + } + let local_put_op = reusable_singleflight.then(ExternalPutKeySharedOp::new); + if let Some(op) = local_put_op.as_ref() { + assert!( + state.external_put.replace(op.clone()).is_none(), + "a reject-on-inflight Put leader must own an empty shared-op slot" + ); + } + state.local_puts = state + .local_puts + .checked_add(1) + .expect("owner local-first put counter overflow"); + state.external_pending_puts = state + .external_pending_puts + .checked_add(1) + .expect("external pending Put fence counter overflow"); + Ok(ExternalLocalFirstPutKeyReservation::Leader(Arc::new( + ExternalPendingPutFenceGuard { + key: key.to_string(), + owner_key_control: self.owner_key_control.clone(), + owns_local_put: true, + local_put_op, + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: Some(self.view.clone_view()), + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + }, + ))) + } + + pub(crate) fn acquire_external_pending_put_fence( + &self, + key: &str, + ) -> KvResult> { + acquire_external_pending_put_fence_for_key(&self.owner_key_control, key) + } + + pub(crate) fn remember_local_snapshot( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + tracing::debug!( + "skip local snapshot publication behind owner reclaim fence: key={} put_id=({},{})", + key, + put_id.0, + put_id.1 + ); + return; + } + self.local_snapshot_info.insert( + key.to_string(), + LocalSnapshotInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + }, + ); + } + + pub(crate) fn has_local_snapshot(&self, key: &str) -> bool { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return false; + } + self.precommit_local_visible_info.contains_key(key) + || self.get_cached_info.contains_key(key) + || self.local_snapshot_info.contains_key(key) + } + + pub(crate) fn local_visible_mem_holder(&self, key: &str) -> Option> { + let (memory_info, hot_put_id) = { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return None; + } + let memory_info = self.local_visible_mem_holder_unfenced(key); + let hot_put_id = memory_info.as_ref().and_then(|memory_info| { + self.get_cached_info + .get(key) + .filter(|cached| Arc::ptr_eq(&cached.mem_holder, memory_info)) + .map(|cached| (cached.put_time_ms, cached.put_version)) + }); + (memory_info, hot_put_id) + }; + if let Some(put_id) = hot_put_id { + self.owner_hot_touch_or_promote( + key, + put_id, + memory_info + .as_ref() + .expect("hot touch requires a local memory holder"), + ); + } + memory_info + } + + pub(crate) fn local_committed_mem_holder_for_put_id( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) -> Option> { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return None; + } + self.get_cached_info.get(key).and_then(|info| { + (info.put_time_ms == put_id.0 && info.put_version == put_id.1) + .then(|| info.mem_holder.clone()) + }) + } + + pub(crate) fn begin_owner_remote_put( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + preferred_sub_cluster: Option, + protect_source_on_remote_complete: bool, + ) -> OwnerRemotePutReservation { + let mut controls = self.owner_key_control.lock_key(key); + + if let Some(existing) = controls.get(key).and_then(|state| state.remote_put.clone()) { + if existing.put_id == put_id && existing.outcome() == OwnerRemotePutOutcome::InFlight { + existing.merge_request(preferred_sub_cluster, protect_source_on_remote_complete); + self.owner_remote_put_counters + .followers + .fetch_add(1, Ordering::Relaxed); + return OwnerRemotePutReservation::Follower(existing); + } + if existing.outcome() != OwnerRemotePutOutcome::InFlight { + let state = controls + .get_mut(key) + .expect("terminal remote Put flight control state disappeared"); + if state + .remote_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &existing)) + { + state.remote_put = None; + } + } + // A newer local generation may publish before the old remote task + // observes Obsolete. The new generation may replace the visible + // per-key flight slot after its exact source is verified below; + // the displaced task keeps its own holder and its pointer-checked + // completion cannot clear the replacement. + } + + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + self.owner_remote_put_counters + .source_unavailable + .fetch_add(1, Ordering::Relaxed); + self.owner_remote_put_counters + .source_fenced + .fetch_add(1, Ordering::Relaxed); + return OwnerRemotePutReservation::SourceUnavailable; + } + let memory_info = match self.get_cached_info.get(key) { + Some(info) if (info.put_time_ms, info.put_version) == put_id => info.mem_holder.clone(), + Some(_) => { + self.owner_remote_put_counters + .source_unavailable + .fetch_add(1, Ordering::Relaxed); + self.owner_remote_put_counters + .source_version_mismatch + .fetch_add(1, Ordering::Relaxed); + return OwnerRemotePutReservation::SourceUnavailable; + } + None => { + self.owner_remote_put_counters + .source_unavailable + .fetch_add(1, Ordering::Relaxed); + self.owner_remote_put_counters + .source_missing + .fetch_add(1, Ordering::Relaxed); + return OwnerRemotePutReservation::SourceUnavailable; + } + }; + + let op = OwnerRemotePutSharedOp::new( + key, + put_id, + preferred_sub_cluster, + protect_source_on_remote_complete, + ); + let state = controls.entry(key.to_string()).or_default(); + state.install_remote_put_leader(op.clone()); + self.owner_remote_put_counters + .active + .fetch_add(1, Ordering::Relaxed); + self.owner_remote_put_counters + .leaders + .fetch_add(1, Ordering::Relaxed); + OwnerRemotePutReservation::Leader { op, memory_info } + } + + pub(crate) fn finish_owner_remote_put( + &self, + op: &Arc, + outcome: OwnerRemotePutOutcome, + ) -> bool { + if !op.complete(outcome) { + return false; + } + + let mut controls = self.owner_key_control.lock_key(&op.key); + let remove_control = if let Some(state) = controls.get_mut(&op.key) { + if state + .remote_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, op)) + { + state.remote_put = None; + } + state.is_idle() + } else { + false + }; + if remove_control { + controls.remove(&op.key); + } + drop(controls); + + self.owner_remote_put_counters + .active + .fetch_sub(1, Ordering::Relaxed); + let terminal_counter = match outcome { + OwnerRemotePutOutcome::InFlight => { + unreachable!("remote Put cannot finish with an inflight outcome") + } + OwnerRemotePutOutcome::Published => &self.owner_remote_put_counters.published, + OwnerRemotePutOutcome::AlreadySatisfied => { + &self.owner_remote_put_counters.already_satisfied + } + OwnerRemotePutOutcome::Obsolete => &self.owner_remote_put_counters.obsolete, + OwnerRemotePutOutcome::Failed => &self.owner_remote_put_counters.failed, + }; + terminal_counter.fetch_add(1, Ordering::Relaxed); + true + } + + pub(crate) fn begin_owner_local_ssd_put( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + selected_victim: Option<&crate::master_kv_router::msg_pack::OwnerSourceEvictionVictim>, + ) -> OwnerLocalSsdPutReservation { + if self.ssd_storage.is_none() { + return OwnerLocalSsdPutReservation::SourceUnavailable; + } + let mut controls = self.owner_key_control.lock_key(key); + + if let Some(existing) = controls + .get(key) + .and_then(|state| state.local_ssd_put.clone()) + { + if existing.put_id == put_id && existing.outcome().is_none() { + self.owner_local_ssd_put_counters + .followers + .fetch_add(1, Ordering::Relaxed); + return OwnerLocalSsdPutReservation::Follower(existing); + } + if existing.outcome().is_some() { + let state = controls + .get_mut(key) + .expect("terminal local SSD Put control state disappeared"); + if state + .local_ssd_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &existing)) + { + state.local_ssd_put = None; + } + } + } + + let memory_info = if let Some(victim) = selected_victim { + let selected = controls + .get(key) + .and_then(|state| state.source_eviction_selection.as_ref()) + .filter(|selection| selection.put_id == put_id) + .map(|selection| selection.cached_info.mem_holder.clone()); + selected.filter(|source| { + source.local_reserve_resident_slot_ref().is_some_and( + |(slot_size, grant_id, slot_index)| { + matches!( + &victim.backing, + crate::master_kv_router::msg_pack::OwnerReclaimBacking::CommittedSlot { + grant_id: expected_grant_id, + slot_index: expected_slot_index, + slot_size: expected_slot_size, + } if *expected_grant_id == grant_id + && *expected_slot_index == slot_index + && *expected_slot_size == slot_size + ) + }, + ) + }) + } else { + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + None + } else { + self.get_cached_info.get(key).and_then(|info| { + ((info.put_time_ms, info.put_version) == put_id) + .then(|| info.mem_holder.clone()) + }) + } + }; + let Some(memory_info) = memory_info else { + self.owner_local_ssd_put_counters + .source_unavailable + .fetch_add(1, Ordering::Relaxed); + return OwnerLocalSsdPutReservation::SourceUnavailable; + }; + + let op = OwnerLocalSsdPutSharedOp::new(key, put_id); + controls + .entry(key.to_string()) + .or_default() + .install_local_ssd_put_leader(op.clone()); + self.owner_local_ssd_put_counters + .active + .fetch_add(1, Ordering::Relaxed); + self.owner_local_ssd_put_counters + .leaders + .fetch_add(1, Ordering::Relaxed); + OwnerLocalSsdPutReservation::Leader { op, memory_info } + } + + pub(crate) fn finish_owner_local_ssd_put( + &self, + op: &Arc, + outcome: OwnerLocalSsdPutOutcome, + ) -> bool { + if !op.complete(outcome) { + return false; + } + + let mut controls = self.owner_key_control.lock_key(&op.key); + let remove_control = if let Some(state) = controls.get_mut(&op.key) { + if state + .local_ssd_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, op)) + { + state.local_ssd_put = None; + } + state.is_idle() + } else { + false + }; + if remove_control { + controls.remove(&op.key); + } + drop(controls); + + self.owner_local_ssd_put_counters + .active + .fetch_sub(1, Ordering::Relaxed); + let terminal_counter = match outcome { + OwnerLocalSsdPutOutcome::Published => &self.owner_local_ssd_put_counters.published, + OwnerLocalSsdPutOutcome::AlreadyPresent => { + &self.owner_local_ssd_put_counters.already_present + } + OwnerLocalSsdPutOutcome::Dropped => &self.owner_local_ssd_put_counters.dropped, + OwnerLocalSsdPutOutcome::Obsolete => &self.owner_local_ssd_put_counters.obsolete, + OwnerLocalSsdPutOutcome::Failed => &self.owner_local_ssd_put_counters.failed, + }; + terminal_counter.fetch_add(1, Ordering::Relaxed); + true + } + + pub(crate) fn record_owner_remote_put_transfer(&self) { + self.owner_remote_put_counters + .transfers + .fetch_add(1, Ordering::Relaxed); + } + + fn local_visible_mem_holder_unfenced(&self, key: &str) -> Option> { + if let Some(info) = self.precommit_local_visible_info.get(key) { + return Some(info.mem_holder.clone()); + } + self.get_cached_info + .get(key) + .map(|info| info.mem_holder.clone()) + } + + pub(crate) fn local_visible_mem_holders( + &self, + keys: &[String], + ) -> Vec>> { + // Resolve each page under only its own short sharded fence. The batch + // itself never owns a synchronous lock. A cloned MemoryInfo pins the + // selected backing if reclaim starts after this point. + let resolved = keys + .iter() + .map(|key| { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return None; + } + let memory_info = self.local_visible_mem_holder_unfenced(key)?; + let hot_put_id = self + .get_cached_info + .get(key) + .filter(|cached| Arc::ptr_eq(&cached.mem_holder, &memory_info)) + .map(|cached| (cached.put_time_ms, cached.put_version)); + Some((memory_info, hot_put_id)) + }) + .collect::>(); + resolved + .into_iter() + .zip(keys) + .map(|(resolved, key)| { + let (memory_info, hot_put_id) = resolved?; + if let Some(put_id) = hot_put_id { + self.owner_hot_touch_or_promote(key, put_id, &memory_info); + } + Some(memory_info) + }) + .collect() + } + + pub(crate) fn install_external_get_holding( + &self, + req_node_id: &str, + memory_info: Arc, + ) -> ExternalMemHolderInfo { + let external_holder_id = allocate_external_holding_id(&self.next_external_holding_id); + let key = NodeHolderKey::new(req_node_id.to_string(), external_holder_id); + let external_memholder_info = ExternalMemHolderInfo { + offset: memory_info.offset, + len: memory_info.len, + holder_id: external_holder_id, + }; + let owner_hot_pin = self.owner_hot_pin_memory_info(&memory_info); + let previous = self.external_get_holding.inner().insert( + key, + ExternalHoldingGetInfo { + key: memory_info.key.clone(), + req_node_id: req_node_id.to_string(), + requester_node_start_time: self + .view + .cluster_manager() + .get_member_info_cached(req_node_id) + .map(|member| member.node_start_time), + memory_info, + _owner_hot_pin: owner_hot_pin, + }, + ); + assert!( + previous.is_none(), + "fresh external holding id unexpectedly replaced a live holding" + ); + external_memholder_info + } + + pub async fn build_local_reserve_resident_memory_info( + &self, + key: &str, + addr: u64, + len: u32, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> Arc { + let resident_owner_node_id: NodeID = self.view.cluster_manager().get_self_info().id.into(); + Arc::new( + MemoryInfo::new_local_reserve_resident( + addr, + len, + key.to_string(), + resident_owner_node_id, + self.view.clone(), + slot_size, + grant_id, + slot_index, + ) + .await, + ) + } + + pub(crate) fn install_hidden_pending_local_get( + &self, + key: &str, + get_id: u64, + put_id: crate::master_kv_router::put::PutIDForAKey, + addr: u64, + base_addr: u64, + len: u32, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult> { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + || self.pending_local_get_info.contains_key(key) + { + return Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })); + } + self.owner_mark_local_reserve_slot_pending_visible(slot_size, grant_id, slot_index)?; + self.owner_retain_local_reserve_resident_slot_holder(slot_size, grant_id, slot_index)?; + let resident_owner_node_id: NodeID = self.view.cluster_manager().get_self_info().id.into(); + let memory_info = Arc::new(MemoryInfo::new_local_reserve_resident_with_base( + addr, + base_addr, + len, + key.to_string(), + resident_owner_node_id, + self.view.clone(), + slot_size, + grant_id, + slot_index, + )); + let previous = self.pending_local_get_info.insert( + key.to_string(), + PendingLocalGetInfo { + get_id, + put_id, + mem_holder: memory_info.clone(), + }, + ); + assert!( + previous.is_none(), + "pending local Get must be unique per key" + ); + drop(controls); + Ok(memory_info) + } + + pub(crate) fn abort_hidden_pending_local_get(&self, key: &str, get_id: u64) -> bool { + let _controls = self.owner_key_control.lock_key(key); + self.pending_local_get_info + .remove_if(key, |_, pending| pending.get_id == get_id) + .is_some() + } + + pub(crate) fn promote_hidden_pending_local_get( + &self, + key: &str, + get_id: u64, + put_id: crate::master_kv_router::put::PutIDForAKey, + ) -> KvResult> { + let memory_info = { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })); + } + let Some(pending_memory_info) = + self.pending_local_get_info.get(key).and_then(|pending| { + (pending.get_id == get_id && pending.put_id == put_id) + .then(|| pending.mem_holder.clone()) + }) + else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "hidden pending local Get is absent: key={} get_id={}", + key, get_id + ), + })); + }; + let (slot_size, grant_id, slot_index) = pending_memory_info + .local_reserve_resident_slot_ref() + .expect("pending local Get must carry a local-reserve slot"); + self.owner_promote_local_reserve_pending_slot_to_committed( + slot_size, grant_id, slot_index, + )?; + let removed = self + .pending_local_get_info + .remove_if(key, |_, pending| { + pending.get_id == get_id + && pending.put_id == put_id + && Arc::ptr_eq(&pending.mem_holder, &pending_memory_info) + }) + .is_some(); + if !removed { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "hidden pending local Get changed while promoting: key={} get_id={}", + key, get_id + ), + })); + } + let replaced = self.get_cached_info.insert( + key.to_string(), + GetCachedInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + mem_holder: pending_memory_info.clone(), + }, + ); + if let Some(previous) = replaced { + self.release_local_reserve_route_for_memory_info(previous.mem_holder.as_ref()); + } + self.local_snapshot_info.insert( + key.to_string(), + LocalSnapshotInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + }, + ); + drop(controls); + pending_memory_info + }; + Ok(memory_info) + } + + pub async fn install_local_committed_memory_info( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + offset: u64, + len: u32, + holder_id: u64, + ) -> KvResult<()> { + let master_node_id: NodeID = self.view.cluster_manager().get_self_info().id.into(); + let memory_info = Arc::new( + MemoryInfo::new( + offset, + len, + holder_id, + key.to_string(), + master_node_id, + self.view.clone(), + ) + .await, + ); + { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })); + } + let replaced = self.get_cached_info.insert( + key.to_string(), + GetCachedInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + mem_holder: memory_info.clone(), + }, + ); + if let Some(previous) = replaced { + self.release_local_reserve_route_for_memory_info(previous.mem_holder.as_ref()); + } + self.local_snapshot_info.insert( + key.to_string(), + LocalSnapshotInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + }, + ); + } + self.owner_hot_track_committed(key, put_id, &memory_info); + Ok(()) + } + + pub(crate) fn install_get_cached_info_if_unfenced( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: Arc, + ) -> bool { + { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + tracing::debug!( + "skip get cache publication behind owner reclaim fence: key={} put_id=({},{})", + key, + put_id.0, + put_id.1 + ); + return false; + } + let replaced = self.get_cached_info.insert( + key.to_string(), + GetCachedInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + mem_holder: memory_info.clone(), + }, + ); + if let Some(previous) = replaced { + self.release_local_reserve_route_for_memory_info(previous.mem_holder.as_ref()); + } + self.local_snapshot_info.insert( + key.to_string(), + LocalSnapshotInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + }, + ); + } + self.owner_hot_track_committed(key, put_id, &memory_info); + true + } + + pub fn install_precommit_local_visible_memory_info( + &self, + key: &str, + memory_info: Arc, + ) { + let controls = self.owner_key_control.lock_key(key); + assert!( + !controls + .get(key) + .is_some_and(|state| state.local_access_fenced()), + "precommit local index publication must not cross an owner reclaim fence" + ); + let (slot_size, grant_id, slot_index) = memory_info + .local_reserve_resident_slot_ref() + .expect("resident memory_info must carry local reserve slot ref"); + self.owner_mark_local_reserve_slot_pending_visible(slot_size, grant_id, slot_index) + .expect("marking local reserve slot pending visible must succeed"); + self.owner_retain_local_reserve_resident_slot_holder(slot_size, grant_id, slot_index) + .expect("retaining local reserve resident holder must succeed"); + let replaced = self.precommit_local_visible_info.insert( + key.to_string(), + PrecommitLocalVisibleInfo { + mem_holder: memory_info.clone(), + }, + ); + assert!( + replaced.is_none(), + "precommit local visible cache must not be replaced for the same key" + ); + } + + pub fn remove_precommit_local_reserve_resident_slot_if_same( + &self, + key: &str, + expected_mem_holder: &Arc, + ) -> bool { + let _controls = self.owner_key_control.lock_key(key); + self.precommit_local_visible_info + .remove_if(key, |_, info| { + Arc::ptr_eq(&info.mem_holder, expected_mem_holder) + }) + .is_some() + } + + pub fn precommit_local_visible_memory_info(&self, key: &str) -> Option> { + self.precommit_local_visible_info + .get(key) + .map(|info| info.mem_holder.clone()) + } + + pub(crate) fn committed_local_reserve_slot_is_current( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + expected: &crate::master_kv_router::msg_pack::PutDoneCommittedSlot, + ) -> bool { + self.get_cached_info.get(key).is_some_and(|cached| { + (cached.put_time_ms, cached.put_version) == put_id + && cached.mem_holder.local_reserve_resident_slot_ref() + == Some((expected.slot_size, expected.grant_id, expected.slot_index)) + }) + } + + pub fn promote_precommit_local_reserve_resident_slot_if_same( + &self, + key: &str, + put_id: crate::master_kv_router::put::PutIDForAKey, + memory_info: Arc, + _atomic_group: Option<&crate::master_kv_router::msg_pack::PutAtomicGroup>, + ) -> KvResult<()> { + { + let controls = self.owner_key_control.lock_key(key); + if controls + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + return Err(KvError::Api(ApiError::KeyBeingWritten { + key: key.to_string(), + })); + } + let is_same_pending = self + .precommit_local_visible_info + .get(key) + .map(|info| Arc::ptr_eq(&info.mem_holder, &memory_info)) + .unwrap_or(false); + if !is_same_pending { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "precommit local visible cache missing while promoting key={}", + key + ), + })); + } + let (slot_size, grant_id, slot_index) = memory_info + .local_reserve_resident_slot_ref() + .expect("resident memory_info must carry local reserve slot ref"); + self.owner_promote_local_reserve_pending_slot_to_committed( + slot_size, grant_id, slot_index, + )?; + let removed = self + .precommit_local_visible_info + .remove_if(key, |_, info| Arc::ptr_eq(&info.mem_holder, &memory_info)) + .is_some(); + if !removed { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "precommit local visible cache disappeared while promoting key={}", + key + ), + })); + } + let replaced = self.get_cached_info.insert( + key.to_string(), + GetCachedInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + mem_holder: memory_info.clone(), + }, + ); + if let Some(previous) = replaced { + self.release_local_reserve_route_for_memory_info(previous.mem_holder.as_ref()); + } + self.local_snapshot_info.insert( + key.to_string(), + LocalSnapshotInfo { + put_time_ms: put_id.0, + put_version: put_id.1, + }, + ); + } + Ok(()) + } +} + +#[derive(Debug, Clone)] +pub struct ExternalPendingPutCtx { + pub peer_id: Option, + pub src_offset: u64, + pub target_base_addr: u64, + pub target_offset: u64, + pub len: u64, + /// Original content-selection signal from the caller/adapter. + pub make_replica_task: bool, + /// A remote memory target was pre-reserved, or this local-first path may + /// perform normal append-time remote admission. + pub remote_replica_admitted: bool, + pub preferred_sub_cluster: Option, + pub local_reserve_slot: Option, + pub local_reserve_slot_size: Option, + pub atomic_group: Option, + /// Keep the per-key reclaim fence alive for every cache/user clone of this + /// pending context. The counter is released only by the final Arc drop. + pub(crate) _pending_fence: Arc, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) enum OwnerLocalReserveSlotState { + Free, + Prepared, + PendingLocalVisible { + holder_ref_count: u32, + }, + Committed { + route_live: bool, + holder_ref_count: u32, + }, +} + +#[derive(Debug, Clone)] +pub struct OwnerLocalReserveSlotRef { + pub grant_id: u64, + pub slot_index: u32, + pub ptr: u64, + pub base_addr: u64, +} + +#[derive(Debug, Clone)] +pub struct OwnerLocalReserveSlotLease { + pub value_len: u64, + pub slot_size: u64, + pub slots: Vec, +} + +impl OwnerLocalReserveSlotLease { + pub fn value_ptrs(&self) -> Vec { + self.slots.iter().map(|slot| slot.ptr).collect() + } +} + +#[derive(Debug, Clone)] +pub(crate) struct OwnerLocalReserveGrantState { + pub grant_id: u64, + pub base_addr: u64, + pub addr: u64, + pub len: u64, + pub slot_size: u64, + pub slot_count: u32, + pub slot_states: Vec, + pub free_slots: Vec, + pub fully_free_since: Option, +} + +impl OwnerLocalReserveGrantState { + pub fn new( + grant_id: u64, + base_addr: u64, + addr: u64, + len: u64, + slot_size: u64, + slot_count: u32, + ) -> Self { + let mut free_slots = Vec::with_capacity(slot_count as usize); + for slot_index in (0..slot_count).rev() { + free_slots.push(slot_index); + } + Self { + grant_id, + base_addr, + addr, + len, + slot_size, + slot_count, + slot_states: vec![OwnerLocalReserveSlotState::Free; slot_count as usize], + free_slots, + fully_free_since: Some(Instant::now()), + } + } + + pub fn claim_prepared_slot(&mut self) -> Option { + let slot_index = self.free_slots.pop()?; + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + assert!( + matches!(*state, OwnerLocalReserveSlotState::Free), + "claim_prepared_slot expects a free slot" + ); + *state = OwnerLocalReserveSlotState::Prepared; + self.fully_free_since = None; + Some(OwnerLocalReserveSlotRef { + grant_id: self.grant_id, + slot_index, + ptr: self.addr + self.slot_size * slot_index as u64, + base_addr: self.base_addr, + }) + } + + pub fn mark_prepared_slot_pending_visible(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + assert!( + matches!(*state, OwnerLocalReserveSlotState::Prepared), + "mark_prepared_slot_pending_visible expects a prepared slot" + ); + *state = OwnerLocalReserveSlotState::PendingLocalVisible { + holder_ref_count: 0, + }; + self.fully_free_since = None; + } + + pub fn promote_pending_visible_slot_to_committed(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + let holder_ref_count = match *state { + OwnerLocalReserveSlotState::PendingLocalVisible { holder_ref_count } => { + holder_ref_count + } + _ => { + unreachable!("promote_pending_visible_slot_to_committed expects a pending slot"); + } + }; + *state = OwnerLocalReserveSlotState::Committed { + route_live: true, + holder_ref_count, + }; + self.fully_free_since = None; + } + + pub fn release_prepared_slot(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + assert!( + matches!(*state, OwnerLocalReserveSlotState::Prepared), + "release_prepared_slot expects a prepared slot" + ); + *state = OwnerLocalReserveSlotState::Free; + self.free_slots.push(slot_index); + if self.is_fully_free() { + self.fully_free_since = Some(Instant::now()); + } + } + + pub fn retain_resident_slot_holder(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + match state { + OwnerLocalReserveSlotState::PendingLocalVisible { holder_ref_count } + | OwnerLocalReserveSlotState::Committed { + holder_ref_count, .. + } => { + *holder_ref_count = holder_ref_count + .checked_add(1) + .expect("retain_resident_slot_holder overflow"); + } + _ => { + unreachable!("retain_resident_slot_holder expects a resident slot"); + } + } + } + + pub fn release_resident_slot_holder(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + match state { + OwnerLocalReserveSlotState::PendingLocalVisible { holder_ref_count } => { + assert!( + *holder_ref_count > 0, + "release_resident_slot_holder expects holder_ref_count > 0" + ); + *holder_ref_count -= 1; + if *holder_ref_count == 0 { + *state = OwnerLocalReserveSlotState::Free; + self.free_slots.push(slot_index); + if self.is_fully_free() { + self.fully_free_since = Some(Instant::now()); + } + } + } + OwnerLocalReserveSlotState::Committed { + route_live, + holder_ref_count, + } => { + *holder_ref_count = holder_ref_count + .checked_sub(1) + .expect("release_resident_slot_holder expects holder_ref_count > 0"); + if !*route_live && *holder_ref_count == 0 { + *state = OwnerLocalReserveSlotState::Free; + self.free_slots.push(slot_index); + if self.is_fully_free() { + self.fully_free_since = Some(Instant::now()); + } + } + } + _ => { + unreachable!("release_resident_slot_holder expects a resident slot"); + } + } + } + + pub fn release_committed_slot_route(&mut self, slot_index: u32) { + let state = self + .slot_states + .get_mut(slot_index as usize) + .expect("slot state index out of range"); + match state { + OwnerLocalReserveSlotState::Committed { + route_live, + holder_ref_count, + } => { + assert!( + *route_live, + "release_committed_slot_route expects a live route" + ); + *route_live = false; + if *holder_ref_count == 0 { + *state = OwnerLocalReserveSlotState::Free; + self.free_slots.push(slot_index); + if self.is_fully_free() { + self.fully_free_since = Some(Instant::now()); + } + } + } + _ => { + unreachable!("release_committed_slot_route expects a committed slot"); + } + } + } + + pub fn is_fully_free(&self) -> bool { + self.free_slots.len() == self.slot_count as usize + } + + pub fn used_slot_count(&self) -> usize { + self.slot_count as usize - self.free_slots.len() + } +} + +#[derive(Debug, Clone)] +pub(crate) struct OwnerLocalReserveClassState { + pub slot_size: u64, + pub slots_per_grant: u32, + pub grants: Vec, + /// Stable grant identity -> current index in `grants`. Every vector removal uses + /// swap-remove and repairs the one moved entry, so per-slot state transitions never + /// scan all grants. + grant_indices: HashMap, + /// Dense set of grants which currently contain at least one free slot. This avoids + /// scanning full grants while assembling a lease. + claimable_grant_ids: Vec, + claimable_grant_indices: HashMap, + free_slots: usize, + prepared_slots: usize, + pending_visible_slots: usize, + committed_slots: usize, + pub last_grow_at: Option, + pub pending_slot_demand: usize, + pub max_observed_claim_slots: usize, + pub expected_grant_count: usize, +} + +impl OwnerLocalReserveClassState { + pub fn new(slot_size: u64, slots_per_grant: u32) -> Self { + Self { + slot_size, + slots_per_grant, + grants: Vec::new(), + grant_indices: HashMap::new(), + claimable_grant_ids: Vec::new(), + claimable_grant_indices: HashMap::new(), + free_slots: 0, + prepared_slots: 0, + pending_visible_slots: 0, + committed_slots: 0, + last_grow_at: None, + pending_slot_demand: 0, + max_observed_claim_slots: 0, + expected_grant_count: 0, + } + } + + pub fn free_slot_count(&self) -> usize { + self.free_slots + } + + pub fn used_slot_count(&self) -> usize { + self.prepared_slots + .saturating_add(self.pending_visible_slots) + .saturating_add(self.committed_slots) + } + + pub fn grant_count(&self) -> usize { + self.grants.len() + } + + pub fn prepared_slot_count(&self) -> usize { + self.prepared_slots + } + + pub fn pending_visible_slot_count(&self) -> usize { + self.pending_visible_slots + } + + pub fn committed_slot_count(&self) -> usize { + self.committed_slots + } + + fn add_claimable_grant(&mut self, grant_id: u64) { + if self.claimable_grant_indices.contains_key(&grant_id) { + return; + } + let index = self.claimable_grant_ids.len(); + self.claimable_grant_ids.push(grant_id); + let previous = self.claimable_grant_indices.insert(grant_id, index); + assert!( + previous.is_none(), + "duplicate claimable local-reserve grant" + ); + } + + fn remove_claimable_grant(&mut self, grant_id: u64) { + let Some(index) = self.claimable_grant_indices.remove(&grant_id) else { + return; + }; + let removed = self.claimable_grant_ids.swap_remove(index); + assert_eq!(removed, grant_id, "claimable grant index drift"); + if let Some(moved_grant_id) = self.claimable_grant_ids.get(index).copied() { + let moved_index = self + .claimable_grant_indices + .get_mut(&moved_grant_id) + .expect("moved claimable grant must remain indexed"); + *moved_index = index; + } + } + + pub fn install_grant(&mut self, grant: OwnerLocalReserveGrantState) { + assert_eq!(grant.slot_size, self.slot_size, "grant slot-size drift"); + assert_eq!( + grant.slot_count, self.slots_per_grant, + "grant slot-count drift" + ); + assert!( + !self.grant_indices.contains_key(&grant.grant_id), + "duplicate local-reserve grant id" + ); + + let mut free = 0usize; + let mut prepared = 0usize; + let mut pending_visible = 0usize; + let mut committed = 0usize; + for state in &grant.slot_states { + match state { + OwnerLocalReserveSlotState::Free => free += 1, + OwnerLocalReserveSlotState::Prepared => prepared += 1, + OwnerLocalReserveSlotState::PendingLocalVisible { .. } => pending_visible += 1, + OwnerLocalReserveSlotState::Committed { .. } => committed += 1, + } + } + assert_eq!(free, grant.free_slots.len(), "grant free-slot index drift"); + + let grant_id = grant.grant_id; + let index = self.grants.len(); + self.grants.push(grant); + let previous = self.grant_indices.insert(grant_id, index); + assert!(previous.is_none(), "duplicate local-reserve grant id"); + if free != 0 { + self.add_claimable_grant(grant_id); + } + self.free_slots = self + .free_slots + .checked_add(free) + .expect("free slot overflow"); + self.prepared_slots = self + .prepared_slots + .checked_add(prepared) + .expect("prepared slot overflow"); + self.pending_visible_slots = self + .pending_visible_slots + .checked_add(pending_visible) + .expect("pending-visible slot overflow"); + self.committed_slots = self + .committed_slots + .checked_add(committed) + .expect("committed slot overflow"); + } + + pub fn claim_available(&mut self, max_slots: usize) -> Vec { + let claim_count = self.free_slots.min(max_slots); + let mut slots = Vec::with_capacity(claim_count); + while slots.len() < claim_count { + let grant_id = *self + .claimable_grant_ids + .last() + .expect("free-slot counter requires a claimable grant"); + let grant_index = *self + .grant_indices + .get(&grant_id) + .expect("claimable grant must be installed"); + let (slot, exhausted) = { + let grant = &mut self.grants[grant_index]; + let slot = grant + .claim_prepared_slot() + .expect("claimable grant must contain a free slot"); + (slot, grant.free_slots.is_empty()) + }; + self.free_slots = self + .free_slots + .checked_sub(1) + .expect("free slot counter underflow"); + self.prepared_slots = self + .prepared_slots + .checked_add(1) + .expect("prepared slot overflow"); + if exhausted { + self.remove_claimable_grant(grant_id); + } + slots.push(slot); + } + slots + } + + fn grant_index(&self, grant_id: u64) -> Option { + self.grant_indices.get(&grant_id).copied() + } + + pub fn grant(&self, grant_id: u64) -> Option<&OwnerLocalReserveGrantState> { + self.grant_index(grant_id) + .and_then(|index| self.grants.get(index)) + } + + pub fn release_prepared_slot(&mut self, grant_id: u64, slot_index: u32) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + let was_exhausted = self.grants[grant_index].free_slots.is_empty(); + self.grants[grant_index].release_prepared_slot(slot_index); + self.prepared_slots = self + .prepared_slots + .checked_sub(1) + .expect("prepared slot counter underflow"); + self.free_slots = self.free_slots.checked_add(1).expect("free slot overflow"); + if was_exhausted { + self.add_claimable_grant(grant_id); + } + true + } + + pub fn mark_prepared_slot_pending_visible(&mut self, grant_id: u64, slot_index: u32) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + self.grants[grant_index].mark_prepared_slot_pending_visible(slot_index); + self.prepared_slots = self + .prepared_slots + .checked_sub(1) + .expect("prepared slot counter underflow"); + self.pending_visible_slots = self + .pending_visible_slots + .checked_add(1) + .expect("pending-visible slot overflow"); + true + } + + pub fn promote_pending_visible_slot_to_committed( + &mut self, + grant_id: u64, + slot_index: u32, + ) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + self.grants[grant_index].promote_pending_visible_slot_to_committed(slot_index); + self.pending_visible_slots = self + .pending_visible_slots + .checked_sub(1) + .expect("pending-visible slot counter underflow"); + self.committed_slots = self + .committed_slots + .checked_add(1) + .expect("committed slot overflow"); + true + } + + pub fn retain_resident_slot_holder(&mut self, grant_id: u64, slot_index: u32) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + self.grants[grant_index].retain_resident_slot_holder(slot_index); + true + } + + pub fn release_resident_slot_holder(&mut self, grant_id: u64, slot_index: u32) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + let (was_exhausted, prior_state, became_free) = { + let grant = &mut self.grants[grant_index]; + let was_exhausted = grant.free_slots.is_empty(); + let prior_state = grant + .slot_states + .get(slot_index as usize) + .expect("slot state index out of range") + .clone(); + let free_before = grant.free_slots.len(); + grant.release_resident_slot_holder(slot_index); + ( + was_exhausted, + prior_state, + grant.free_slots.len() != free_before, + ) + }; + if became_free { + match prior_state { + OwnerLocalReserveSlotState::PendingLocalVisible { .. } => { + self.pending_visible_slots = self + .pending_visible_slots + .checked_sub(1) + .expect("pending-visible slot counter underflow"); + } + OwnerLocalReserveSlotState::Committed { .. } => { + self.committed_slots = self + .committed_slots + .checked_sub(1) + .expect("committed slot counter underflow"); + } + _ => unreachable!("resident holder must belong to a resident slot"), + } + self.free_slots = self.free_slots.checked_add(1).expect("free slot overflow"); + if was_exhausted { + self.add_claimable_grant(grant_id); + } + } + true + } + + pub fn release_committed_slot_route(&mut self, grant_id: u64, slot_index: u32) -> bool { + let Some(grant_index) = self.grant_index(grant_id) else { + return false; + }; + let (was_exhausted, became_free) = { + let grant = &mut self.grants[grant_index]; + let was_exhausted = grant.free_slots.is_empty(); + let free_before = grant.free_slots.len(); + grant.release_committed_slot_route(slot_index); + (was_exhausted, grant.free_slots.len() != free_before) + }; + if became_free { + self.committed_slots = self + .committed_slots + .checked_sub(1) + .expect("committed slot counter underflow"); + self.free_slots = self.free_slots.checked_add(1).expect("free slot overflow"); + if was_exhausted { + self.add_claimable_grant(grant_id); + } + } + true + } + + /// Drop the route reference and the resident `MemoryInfo` reference as one + /// slot-pool transaction. This is the normal owner-reclaim transition and + /// avoids exposing an intermediate state across two pool lock acquisitions. + pub fn release_committed_resident_slot(&mut self, grant_id: u64, slot_index: u32) -> bool { + if self.grant_index(grant_id).is_none() { + return false; + } + assert!(self.release_committed_slot_route(grant_id, slot_index)); + assert!(self.release_resident_slot_holder(grant_id, slot_index)); + true + } + + pub fn detach_fully_free_grant( + &mut self, + grant_id: u64, + ) -> Option { + let grant_index = self.grant_indices.remove(&grant_id)?; + assert!( + self.grants[grant_index].is_fully_free(), + "only a fully-free grant may be detached" + ); + self.remove_claimable_grant(grant_id); + let grant = self.grants.swap_remove(grant_index); + assert_eq!(grant.grant_id, grant_id, "grant index drift"); + if let Some(moved_grant) = self.grants.get(grant_index) { + let moved_index = self + .grant_indices + .get_mut(&moved_grant.grant_id) + .expect("moved grant must remain indexed"); + *moved_index = grant_index; + } + self.free_slots = self + .free_slots + .checked_sub(grant.slot_count as usize) + .expect("free slot counter underflow"); + Some(grant) + } + + pub fn take_all_grants(&mut self) -> Vec { + self.grant_indices.clear(); + self.claimable_grant_ids.clear(); + self.claimable_grant_indices.clear(); + self.free_slots = 0; + self.prepared_slots = 0; + self.pending_visible_slots = 0; + self.committed_slots = 0; + std::mem::take(&mut self.grants) + } +} + +#[cfg(test)] +mod owner_reclaim_slot_tests { + use super::{ + ApiError, ClientKvApi, ClientKvApiNewArg, ExternalLocalFirstPutKeyReservation, + ExternalPendingPutCtx, ExternalPendingPutFenceGuard, ExternalPutKeyOutcome, + ExternalPutKeySharedOp, KvError, OwnerHotCacheCounters, OwnerHotCacheEntry, + OwnerHotEvictionDispatch, OwnerHotEvictionEvent, OwnerHotPinAlias, OwnerHotReplicaIdentity, + OwnerHotRetryQueue, OwnerHotSelectionDebt, OwnerHotSelectionFenceOutcome, + OwnerKeyControlState, OwnerKeyControlTable, OwnerLocalReserveGrantState, + OwnerLocalReserveSlotLease, OwnerLocalReserveSlotRef, OwnerLocalReserveSlotState, + OwnerLocalSsdPutOutcome, OwnerLocalSsdPutSharedOp, OwnerReclaimRecord, + OwnerRemotePutOutcome, OwnerRemotePutReservation, OwnerRemotePutSharedOp, + acquire_external_pending_put_fence_for_key, allocate_external_holding_id, + build_owner_hot_cache, clone_if_owner_hot_entry_matches, + owner_hot_source_has_active_holders, pin_current_owner_hot_source_from_index, + }; + use crate::config::TestSpecConfig; + use crate::kv_ssd_storage::{KvSsdStorageInit, KvSsdStorageRootLimit, MIN_CAPACITY_BYTES}; + use crate::master_kv_router::msg_pack::{ + BatchOwnerReclaimReq, OwnerReclaimBacking, OwnerReclaimItem, OwnerReclaimItemState, + OwnerReclaimPhase, OwnerReclaimReason, OwnerSourceEvictionVictim, OwnerSourceSsdPolicy, + }; + use crate::p2p::msg_pack::MsgPack; + use parking_lot::Mutex; + use std::fs; + use std::path::PathBuf; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::{Arc, Weak}; + use std::time::{Duration, Instant}; + + fn pending_put_count(controls: &Arc, key: &str) -> Option { + controls + .lock_key(key) + .get(key) + .map(|state| state.external_pending_puts) + } + + fn controls_with_state(key: &str, state: OwnerKeyControlState) -> Arc { + let controls = Arc::new(OwnerKeyControlTable::default()); + controls.lock_key(key).insert(key.to_string(), state); + controls + } + + #[test] + fn external_pending_put_guard_clone_releases_only_after_last_drop() { + let controls = Arc::new(OwnerKeyControlTable::default()); + let guard = acquire_external_pending_put_fence_for_key(&controls, "clone-key") + .expect("pending fence acquisition must succeed"); + let clone = guard.clone(); + assert_eq!(pending_put_count(&controls, "clone-key"), Some(1)); + + drop(guard); + assert_eq!(pending_put_count(&controls, "clone-key"), Some(1)); + drop(clone); + assert_eq!(pending_put_count(&controls, "clone-key"), None); + } + + #[test] + fn stale_pending_put_guard_drop_cannot_erase_new_generation() { + let controls = Arc::new(OwnerKeyControlTable::default()); + let old = acquire_external_pending_put_fence_for_key(&controls, "aba-key") + .expect("old pending fence acquisition must succeed"); + let new = acquire_external_pending_put_fence_for_key(&controls, "aba-key") + .expect("new pending fence acquisition must succeed"); + assert_eq!(pending_put_count(&controls, "aba-key"), Some(2)); + + drop(old); + assert_eq!(pending_put_count(&controls, "aba-key"), Some(1)); + drop(new); + assert_eq!(pending_put_count(&controls, "aba-key"), None); + } + + #[test] + fn local_first_pending_guard_releases_both_key_counters() { + let controls = controls_with_state( + "local-key", + OwnerKeyControlState { + local_puts: 1, + external_pending_puts: 1, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: None, + local_access_fence: None, + }, + ); + let guard = Arc::new(ExternalPendingPutFenceGuard { + key: "local-key".to_string(), + owner_key_control: controls.clone(), + owns_local_put: true, + local_put_op: None, + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + }); + + drop(guard); + assert!(controls.lock_key("local-key").get("local-key").is_none()); + } + + #[limit_thirdparty::tokio::test] + async fn same_key_put_waiter_reuses_one_leader_terminal_result() { + let key = "put-singleflight"; + let op = ExternalPutKeySharedOp::new(); + let controls = controls_with_state( + key, + OwnerKeyControlState { + local_puts: 1, + external_pending_puts: 1, + external_put: Some(op.clone()), + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: None, + local_access_fence: None, + }, + ); + let leader = Arc::new(ExternalPendingPutFenceGuard { + key: key.to_string(), + owner_key_control: controls.clone(), + owns_local_put: true, + local_put_op: Some(op.clone()), + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + }); + + leader.mark_local_put_succeeded(); + assert_eq!( + ::tokio::time::timeout(Duration::from_secs(1), op.wait()) + .await + .expect("the follower must observe the leader terminal result"), + ExternalPutKeyOutcome::Succeeded + ); + drop(leader); + assert!(controls.lock_key(key).get(key).is_none()); + } + + #[limit_thirdparty::tokio::test] + async fn failed_put_leader_wakes_waiter_only_after_its_fence_is_released() { + let key = "put-singleflight-failed"; + let op = ExternalPutKeySharedOp::new(); + let controls = controls_with_state( + key, + OwnerKeyControlState { + local_puts: 1, + external_pending_puts: 1, + external_put: Some(op.clone()), + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: None, + local_access_fence: None, + }, + ); + let leader = Arc::new(ExternalPendingPutFenceGuard { + key: key.to_string(), + owner_key_control: controls.clone(), + owns_local_put: true, + local_put_op: Some(op.clone()), + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + }); + + drop(leader); + assert_eq!( + ::tokio::time::timeout(Duration::from_secs(1), op.wait()) + .await + .expect("the failed leader must wake its follower"), + ExternalPutKeyOutcome::Failed + ); + assert!(controls.lock_key(key).get(key).is_none()); + } + + #[limit_thirdparty::tokio::test] + async fn every_remote_put_trigger_joins_one_owner_generation_flight() { + let key = "remote-put-singleflight"; + let put_id = (77, 3); + let api = ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"); + let op = OwnerRemotePutSharedOp::new(key, put_id, None, false); + api.inner().owner_key_control.lock_key(key).insert( + key.to_string(), + OwnerKeyControlState { + remote_put: Some(op.clone()), + ..Default::default() + }, + ); + api.inner() + .owner_remote_put_counters + .active + .store(1, Ordering::Relaxed); + api.inner() + .owner_remote_put_counters + .leaders + .store(1, Ordering::Relaxed); + + for follower in 0..64 { + let preferred = (follower == 0).then(|| "tier1".to_string()); + let protect_source = follower == 63; + match api + .inner() + .begin_owner_remote_put(key, put_id, preferred, protect_source) + { + OwnerRemotePutReservation::Follower(joined) => { + assert!(Arc::ptr_eq(&joined, &op)); + } + OwnerRemotePutReservation::Leader { .. } => { + panic!("a follower created a second remote Put leader") + } + OwnerRemotePutReservation::SourceUnavailable => { + panic!("a matching active remote Put flight was not reusable") + } + } + } + + let request = op.request(); + assert_eq!(request.preferred_sub_cluster.as_deref(), Some("tier1")); + assert!(request.protect_source_on_remote_complete); + assert_eq!( + api.inner() + .owner_remote_put_counters + .followers + .load(Ordering::Relaxed), + 64 + ); + assert!(matches!( + api.inner() + .begin_owner_remote_put(key, (78, 0), None, false), + OwnerRemotePutReservation::SourceUnavailable + )); + + assert!( + api.inner() + .finish_owner_remote_put(&op, OwnerRemotePutOutcome::Published) + ); + assert!( + !api.inner() + .finish_owner_remote_put(&op, OwnerRemotePutOutcome::Failed) + ); + assert_eq!( + ::tokio::time::timeout(Duration::from_secs(1), op.wait()) + .await + .expect("remote Put followers must observe the terminal result"), + OwnerRemotePutOutcome::Published + ); + assert!( + api.inner() + .owner_key_control + .lock_key(key) + .get(key) + .is_none() + ); + assert_eq!( + api.inner() + .owner_remote_put_counters + .active + .load(Ordering::Relaxed), + 0 + ); + } + + #[limit_thirdparty::tokio::test] + async fn new_remote_put_generation_replaces_old_without_aba_cleanup() { + let key = "remote-put-generation-aba"; + let api = ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"); + let old = OwnerRemotePutSharedOp::new(key, (80, 0), None, false); + let new = OwnerRemotePutSharedOp::new(key, (81, 0), None, false); + api.inner().owner_key_control.lock_key(key).insert( + key.to_string(), + OwnerKeyControlState { + remote_put: Some(old.clone()), + ..Default::default() + }, + ); + api.inner() + .owner_remote_put_counters + .active + .store(2, Ordering::Relaxed); + api.inner() + .owner_key_control + .lock_key(key) + .get_mut(key) + .expect("old generation control state") + .install_remote_put_leader(new.clone()); + let current = api + .inner() + .owner_key_control + .lock_key(key) + .get(key) + .and_then(|state| state.remote_put.clone()) + .expect("new generation must own the visible flight slot"); + assert!(Arc::ptr_eq(¤t, &new)); + assert_eq!( + api.inner() + .owner_remote_put_counters + .active + .load(Ordering::Relaxed), + 2 + ); + + assert!( + api.inner() + .finish_owner_remote_put(&old, OwnerRemotePutOutcome::Obsolete) + ); + let current = api + .inner() + .owner_key_control + .lock_key(key) + .get(key) + .and_then(|state| state.remote_put.clone()) + .expect("old completion must retain the new generation flight"); + assert!(Arc::ptr_eq(¤t, &new)); + assert_eq!( + api.inner() + .owner_remote_put_counters + .active + .load(Ordering::Relaxed), + 1 + ); + + assert!( + api.inner() + .finish_owner_remote_put(&new, OwnerRemotePutOutcome::Published) + ); + assert!( + api.inner() + .owner_key_control + .lock_key(key) + .get(key) + .is_none() + ); + assert_eq!( + api.inner() + .owner_remote_put_counters + .active + .load(Ordering::Relaxed), + 0 + ); + } + + #[limit_thirdparty::tokio::test] + async fn remote_and_local_ssd_flights_publish_independent_terminals() { + let key = "parallel-backing-flights"; + let put_id = (82, 4); + let api = ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"); + let remote = OwnerRemotePutSharedOp::new(key, put_id, None, false); + let local_ssd = OwnerLocalSsdPutSharedOp::new(key, put_id); + api.inner().owner_key_control.lock_key(key).insert( + key.to_string(), + OwnerKeyControlState { + remote_put: Some(remote.clone()), + local_ssd_put: Some(local_ssd.clone()), + ..Default::default() + }, + ); + api.inner() + .owner_remote_put_counters + .active + .store(1, Ordering::Relaxed); + api.inner() + .owner_local_ssd_put_counters + .active + .store(1, Ordering::Relaxed); + + assert!( + api.inner() + .finish_owner_remote_put(&remote, OwnerRemotePutOutcome::Published) + ); + { + let controls = api.inner().owner_key_control.lock_key(key); + let state = controls.get(key).expect("SSD flight must remain installed"); + assert!(state.remote_put.is_none()); + assert!( + state + .local_ssd_put + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &local_ssd)) + ); + } + assert_eq!( + local_ssd.outcome(), + None, + "remote completion must not publish the SSD terminal" + ); + assert!( + api.inner() + .finish_owner_local_ssd_put(&local_ssd, OwnerLocalSsdPutOutcome::Published) + ); + assert_eq!(local_ssd.wait().await, OwnerLocalSsdPutOutcome::Published); + assert!( + api.inner() + .owner_key_control + .lock_key(key) + .get(key) + .is_none() + ); + } + + #[limit_thirdparty::tokio::test] + async fn early_ssd_byte_rejection_does_not_install_generation_flight() { + let target = std::env::var_os("CARGO_TARGET_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/mnt/nvme0/mjq_build/push_sglang_fluxon_target")); + let root = target.join("kv_ssd_tests").join(format!( + "early-pre-admission-{}-{}", + std::process::id(), + uuid::Uuid::new_v4() + )); + let api = ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: Some(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: Some(1), + write_burst_bytes: Some(4), + }), + }) + .await + .expect("construct SSD-enabled test ClientKvApi"); + let key = "early-pre-admission-drop"; + + api.inner() + .start_early_owner_local_ssd_puts(vec![(key.to_string(), (83, 1), 8)]); + + let usage = api + .inner() + .ssd_storage + .as_ref() + .expect("SSD store") + .usage_snapshot(); + assert_eq!(usage.write_candidate_items, 1); + assert_eq!(usage.write_admitted_items, 0); + assert_eq!(usage.write_dropped_items, 1); + assert_eq!( + api.inner() + .owner_local_ssd_put_counters + .leaders + .load(Ordering::Relaxed), + 0, + "a byte-rejected early candidate must never enter singleflight" + ); + assert!( + api.inner() + .owner_key_control + .lock_key(key) + .get(key) + .is_none(), + "a byte-rejected early candidate must not create per-key state" + ); + + api.inner() + .ssd_storage + .as_ref() + .expect("SSD store") + .close() + .await + .expect("close SSD store"); + drop(api); + fs::remove_dir_all(root).expect("remove SSD test root"); + } + + #[test] + fn failed_local_slot_release_keeps_key_fence_closed() { + let controls = controls_with_state( + "failed-slot", + OwnerKeyControlState { + local_puts: 1, + external_pending_puts: 1, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: None, + local_access_fence: None, + }, + ); + let guard = Arc::new(ExternalPendingPutFenceGuard { + key: "failed-slot".to_string(), + owner_key_control: controls.clone(), + owns_local_put: true, + local_put_op: None, + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(true), + }); + + drop(guard); + let controls = controls.lock_key("failed-slot"); + assert_eq!(controls["failed-slot"].local_puts, 1); + assert_eq!(controls["failed-slot"].external_pending_puts, 1); + } + + #[test] + fn committed_local_first_slot_disarms_drop_cleanup_before_fence_release() { + let controls = controls_with_state( + "committed-slot", + OwnerKeyControlState { + local_puts: 1, + external_pending_puts: 1, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: None, + external_get: None, + local_access_fence: None, + }, + ); + let guard = Arc::new(ExternalPendingPutFenceGuard { + key: "committed-slot".to_string(), + owner_key_control: controls.clone(), + owns_local_put: true, + local_put_op: None, + local_put_succeeded: std::sync::atomic::AtomicBool::new(false), + local_slot_cleanup_view: None, + local_slot_lease: Mutex::new(None), + local_slot_release_failed: std::sync::atomic::AtomicBool::new(false), + }); + guard.attach_local_slot_lease(OwnerLocalReserveSlotLease { + value_len: 8, + slot_size: 8, + slots: vec![OwnerLocalReserveSlotRef { + grant_id: 7, + slot_index: 2, + ptr: 0x1008, + base_addr: 0x1000, + }], + }); + + guard.disarm_local_slot_lease(); + drop(guard); + assert!( + controls + .lock_key("committed-slot") + .get("committed-slot") + .is_none() + ); + } + + #[test] + fn pending_ctx_clone_keeps_fence_after_explicit_cache_invalidation() { + let controls = Arc::new(OwnerKeyControlTable::default()); + let fence = acquire_external_pending_put_fence_for_key(&controls, "cached-key") + .expect("pending fence acquisition must succeed"); + let cache = moka::sync::Cache::new(1); + let identity = ("cached-key".to_string(), 10, 2); + cache.insert( + identity.clone(), + ExternalPendingPutCtx { + peer_id: None, + src_offset: 0, + target_base_addr: 0, + target_offset: 0, + len: 1, + make_replica_task: false, + remote_replica_admitted: false, + preferred_sub_cluster: None, + local_reserve_slot: None, + local_reserve_slot_size: None, + atomic_group: None, + _pending_fence: fence, + }, + ); + let clone = cache.get(&identity).expect("pending ctx must exist"); + cache.invalidate(&identity); + cache.run_pending_tasks(); + assert_eq!(pending_put_count(&controls, "cached-key"), Some(1)); + + drop(clone); + assert_eq!(pending_put_count(&controls, "cached-key"), None); + } + + #[test] + fn external_holding_ids_are_nonzero_and_unique_for_resident_pages() { + let counter = AtomicU64::new(1); + let first = allocate_external_holding_id(&counter); + let second = allocate_external_holding_id(&counter); + assert_eq!(first, 1); + assert_eq!(second, 2); + assert_ne!(first, second); + } + + #[test] + fn hot_source_pin_requires_the_same_version_and_allocation() { + let current = Arc::new(7u64); + let current_weak = Arc::downgrade(¤t); + let other = Arc::new(7u64); + let other_weak = Arc::downgrade(&other); + + let pinned = clone_if_owner_hot_entry_matches((10, 2), ¤t, (10, 2), ¤t_weak) + .expect("matching source should be pinned"); + assert!(Arc::ptr_eq(&pinned, ¤t)); + assert_eq!(Arc::strong_count(¤t), 2); + drop(pinned); + + assert!( + clone_if_owner_hot_entry_matches((10, 2), ¤t, (10, 3), ¤t_weak).is_none() + ); + assert!( + clone_if_owner_hot_entry_matches((10, 2), ¤t, (10, 2), &other_weak).is_none() + ); + } + + #[test] + fn pressure_victim_rejects_an_extra_active_holder() { + let indexed = Arc::new(7u64); + let selected = indexed.clone(); + assert_eq!(Arc::strong_count(&selected), 2); + assert!( + !owner_hot_source_has_active_holders(&selected), + "index plus the temporary selection pin is reclaimable" + ); + + let active_reader = indexed.clone(); + assert_eq!(Arc::strong_count(&selected), 3); + assert!(owner_hot_source_has_active_holders(&selected)); + + drop(active_reader); + assert!(!owner_hot_source_has_active_holders(&selected)); + } + + #[limit_thirdparty::tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn single_key_source_selection_fence_closes_late_get_and_rolls_back() { + use crate::client_kv_api::PutOptionalArgs; + use crate::kvcore_test_lib::{ + integration_test_lock, start_master_and_client, stop_master_and_client, + }; + + let _test_guard = integration_test_lock().await; + let (master, client) = start_master_and_client( + "source_selection_fence_master", + "source_selection_fence_owner", + ) + .await; + let owner_view = client.client_kv_api_view(); + let inner = owner_view.client_kv_api().inner(); + let join_probe_leader = match inner + .reserve_external_local_first_put_key("put-join-probe", true, true) + .expect("the first same-key Put must become leader") + { + ExternalLocalFirstPutKeyReservation::Leader(leader) => leader, + ExternalLocalFirstPutKeyReservation::Wait(_) => { + panic!("the first same-key Put cannot be a follower") + } + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(_) => { + panic!("an unfenced key cannot wait for local access") + } + }; + let join_probe_waiter = match inner + .reserve_external_local_first_put_key("put-join-probe", true, true) + .expect("the second same-key Put must join the leader") + { + ExternalLocalFirstPutKeyReservation::Wait(waiter) => waiter, + ExternalLocalFirstPutKeyReservation::Leader(_) => { + panic!("the second same-key Put must not claim a second leader fence") + } + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(_) => { + panic!("a local-Put follower must wait on the leader result") + } + }; + join_probe_leader.mark_local_put_succeeded(); + assert_eq!( + join_probe_waiter.wait().await, + ExternalPutKeyOutcome::Succeeded + ); + drop(join_probe_leader); + + let key = "selection-single"; + inner + .put(key, &[7u8; 4096], PutOptionalArgs::default()) + .await + .expect("owner put must publish a committed route"); + let (holder, _remote) = inner + .get(key) + .await + .expect("owner get must succeed") + .expect("committed route must be readable"); + drop(holder); + + let cached = inner + .get_cached_info + .get(key) + .expect("committed owner source must be indexed"); + let identity = OwnerHotReplicaIdentity { + key: key.to_string(), + put_time_ms: cached.put_time_ms, + put_version: cached.put_version, + }; + let source = cached.mem_holder.clone(); + drop(cached); + assert!( + !owner_hot_source_has_active_holders(&source), + "the initial victim check must see only index plus selection pins" + ); + + // Reproduce the r11 TOCTOU: a local Get acquires the single victim after + // the dispatcher's first holder check but before its fence is installed. + let late_reader = inner + .local_visible_mem_holder(key) + .expect("late local Get must acquire the source before fencing"); + assert!(matches!( + inner.owner_hot_install_source_selection_fence(&identity, &source), + OwnerHotSelectionFenceOutcome::TemporarilyPinned + )); + assert!(inner.get_cached_info.contains_key(key)); + assert!( + inner + .owner_key_control + .lock_key(key) + .get(key) + .is_none_or(|state| state.source_eviction_selection.is_none()), + "a pinned victim must roll back its source fence" + ); + + drop(late_reader); + assert!(matches!( + inner.owner_hot_install_source_selection_fence(&identity, &source), + OwnerHotSelectionFenceOutcome::Fenced + )); + assert!(!inner.get_cached_info.contains_key(key)); + assert!(inner.local_visible_mem_holder(key).is_none()); + assert!( + acquire_external_pending_put_fence_for_key(&inner.owner_key_control, key).is_err(), + "a new local Put must not cross a source-selection fence" + ); + let rollback_waiter = match inner + .reserve_external_local_first_put_key(key, true, true) + .expect("idempotent local-first Put must asynchronously wait for source selection") + { + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(waiter) => waiter, + ExternalLocalFirstPutKeyReservation::Leader(_) => { + panic!("a source-fenced key cannot claim a local-Put leader") + } + ExternalLocalFirstPutKeyReservation::Wait(_) => { + panic!("a source-fenced key has no local-Put leader to join") + } + }; + let second_rollback_waiter = match inner + .reserve_external_local_first_put_key(key, true, true) + .expect("all idempotent local-first Puts must share the fence completion") + { + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(waiter) => waiter, + _ => panic!("a second source-fence waiter must not claim a Put fence"), + }; + let cancelled_rollback_waiter = match inner + .reserve_external_local_first_put_key(key, true, true) + .expect("a cancellable Put may subscribe to the same fence completion") + { + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(waiter) => waiter, + _ => panic!("a cancellable source-fence waiter must not claim a Put fence"), + }; + // Simulate an external RPC being cancelled before reclaim completes. + // Dropping one receiver must neither retain the physical source nor + // prevent the shared generation from waking the remaining waiters. + drop(cancelled_rollback_waiter); + assert!(matches!( + inner.reserve_external_local_first_put_key(key, true, false), + Err(KvError::Api(ApiError::KeyBeingWritten { .. })) + )); + drop(source); + let master_id = master + .cluster_manager_view() + .cluster_manager() + .get_self_info() + .id; + + let mismatched = OwnerReclaimItem { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version.wrapping_add(1)), + epoch: 90, + backing: OwnerReclaimBacking::Allocation, + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + let mismatch_resp = super::reclaim::handle_batch_owner_reclaim( + &owner_view, + MsgPack { + serialize_part: BatchOwnerReclaimReq { + phase: OwnerReclaimPhase::Prepare, + items: vec![mismatched], + }, + raw_bytes: Vec::new(), + }, + master_id.clone().into(), + ) + .await; + assert_eq!( + mismatch_resp.serialize_part.items[0].state, + OwnerReclaimItemState::Busy + ); + assert!( + inner + .owner_key_control + .lock_key(&identity.key) + .get(&identity.key) + .is_some_and(|state| state.source_eviction_selection.is_some()), + "a mismatched Prepare must not consume the owner selection" + ); + + let matching = OwnerReclaimItem { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version), + epoch: 100, + backing: OwnerReclaimBacking::Allocation, + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + let prepare_resp = super::reclaim::handle_batch_owner_reclaim( + &owner_view, + MsgPack { + serialize_part: BatchOwnerReclaimReq { + phase: OwnerReclaimPhase::Prepare, + items: vec![matching.clone()], + }, + raw_bytes: Vec::new(), + }, + master_id.clone().into(), + ) + .await; + assert!( + prepare_resp.serialize_part.items[0].state == OwnerReclaimItemState::Prepared, + "matching Prepare must promote the single selection into reclaim" + ); + assert!(inner.local_visible_mem_holder(key).is_none()); + + let abort_resp = super::reclaim::handle_batch_owner_reclaim( + &owner_view, + MsgPack { + serialize_part: BatchOwnerReclaimReq { + phase: OwnerReclaimPhase::Abort, + items: vec![matching], + }, + raw_bytes: Vec::new(), + }, + master_id.into(), + ) + .await; + assert!(abort_resp.serialize_part.items[0].state == OwnerReclaimItemState::Aborted); + assert!( + inner.local_visible_mem_holder(key).is_some(), + "Abort must restore the exact committed local index" + ); + for mut waiter in [rollback_waiter, second_rollback_waiter] { + ::tokio::time::timeout(Duration::from_secs(1), async { + loop { + if *waiter.borrow_and_update() { + break; + } + waiter + .changed() + .await + .expect("rollback completion sender must remain live"); + } + }) + .await + .expect("Abort must wake every live source-fence Put waiter"); + } + assert!(matches!( + inner.reserve_external_local_first_put_key(key, true, true), + Err(KvError::Api(ApiError::KeyAlreadyExists { .. })) + )); + + let trigger_weak = { + let cached = inner + .get_cached_info + .get(key) + .expect("Abort-restored trigger must be indexed"); + Arc::downgrade(&cached.mem_holder) + }; + let retry_event = OwnerHotEvictionEvent { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version), + memory_info: trigger_weak, + selection_debt: OwnerHotSelectionDebt::new( + 4096, + inner.owner_hot_counters.selection_debt_bytes.clone(), + ), + retry: true, + source_eviction_victim: None, + retry_failures: 1, + }; + match inner.owner_hot_prepare_eviction(&retry_event) { + super::OwnerHotEvictionPreparation::Ready { trigger, source } => { + assert_eq!(trigger, identity); + assert_eq!(source.key, key); + } + _ => panic!("a current single-key retry must be ready"), + } + retry_event.selection_debt.release(); + + let cached = inner + .get_cached_info + .get(key) + .expect("the aborted source must still be locally indexed"); + let source = cached.mem_holder.clone(); + drop(cached); + assert!(matches!( + inner.owner_hot_install_source_selection_fence(&identity, &source), + OwnerHotSelectionFenceOutcome::Fenced + )); + drop(source); + assert!( + inner.owner_hot_install_source_selection_debt( + identity.clone(), + OwnerHotSelectionDebt::new( + 4096, + inner.owner_hot_counters.selection_debt_bytes.clone(), + ), + ) + ); + let mut direct_delete_waiter = match inner + .reserve_external_local_first_put_key(key, true, true) + .expect("direct-delete source fence must expose an async waiter") + { + ExternalLocalFirstPutKeyReservation::WaitForLocalAccess(waiter) => waiter, + _ => panic!("direct-delete source fence must not admit a local Put"), + }; + super::reclaim::complete_owner_source_eviction( + inner, + &OwnerSourceEvictionVictim { + key: key.to_string(), + put_id: (identity.put_time_ms, identity.put_version), + backing: OwnerReclaimBacking::Allocation, + ssd_backing_len: None, + ssd_policy: OwnerSourceSsdPolicy::Drop, + }, + 101, + ) + .expect("one direct-delete response must release and finalize the local source"); + assert!(inner.local_visible_mem_holder(key).is_none()); + assert!(!inner.owner_source_eviction_selected.contains_key(&identity)); + assert!( + inner.owner_key_control.lock_key(key).get(key).is_none(), + "direct local completion must clear the source fence in one call" + ); + ::tokio::time::timeout(Duration::from_secs(1), async { + loop { + if *direct_delete_waiter.borrow_and_update() { + break; + } + direct_delete_waiter + .changed() + .await + .expect("finalize completion sender must remain live"); + } + }) + .await + .expect("direct-delete Finalize must wake the source-fence Put waiter"); + let post_delete_leader = match inner + .reserve_external_local_first_put_key(key, true, true) + .expect("a physically reclaimed key must be eligible for a new local Put") + { + ExternalLocalFirstPutKeyReservation::Leader(leader) => leader, + _ => panic!("a reclaimed key must re-evaluate to a fresh leader"), + }; + drop(post_delete_leader); + + stop_master_and_client(master, client).await; + } + + #[test] + fn dispatcher_pin_and_reclaim_prepare_are_serialized_by_the_current_index() { + let current = Mutex::new(Some(((10, 2), Arc::new(7u64)))); + let weak = Arc::downgrade(¤t.lock().as_ref().unwrap().1); + + let pinned = match pin_current_owner_hot_source_from_index((10, 2), &weak, || { + current + .lock() + .as_ref() + .map(|(put_id, value)| (*put_id, value.clone())) + }) { + super::OwnerHotPinResult::Pinned(pinned) => pinned, + _ => panic!("current source must pin under the owner fence"), + }; + let prepared_after_pin = current.lock().take().unwrap().1; + assert_eq!(Arc::strong_count(&prepared_after_pin), 2); + drop(pinned); + assert_eq!(Arc::strong_count(&prepared_after_pin), 1); + + // If Prepare wins and moves the sole Arc out of the local index, the + // dispatcher observes an absent current entry. It must not upgrade the + // Weak that now points into Prepared, or Commit's try_unwrap could race + // an unexpected second holder. + let current = Mutex::new(Some(((11, 3), Arc::new(9u64)))); + let weak = Arc::downgrade(¤t.lock().as_ref().unwrap().1); + let prepared_before_pin = current.lock().take().unwrap().1; + assert!(matches!( + pin_current_owner_hot_source_from_index((11, 3), &weak, || { + current + .lock() + .as_ref() + .map(|(put_id, value)| (*put_id, value.clone())) + }), + super::OwnerHotPinResult::ReclaimBusy + )); + assert_eq!(Arc::strong_count(&prepared_before_pin), 1); + + drop(prepared_before_pin); + assert!(matches!( + pin_current_owner_hot_source_from_index((11, 3), &weak, || None), + super::OwnerHotPinResult::Stale + )); + } + + #[test] + fn owner_hot_retry_queue_is_exactly_once_and_keeps_selection_debt() { + let counters = Arc::new(OwnerHotCacheCounters::default()); + let retry_queue = OwnerHotRetryQueue::new(counters.clone()); + let identity = OwnerHotReplicaIdentity { + key: "retry-key".to_string(), + put_time_ms: 12, + put_version: 3, + }; + let debt = OwnerHotSelectionDebt::new(64, counters.selection_debt_bytes.clone()); + let event = OwnerHotEvictionEvent { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version), + memory_info: Weak::new(), + selection_debt: debt, + retry: true, + source_eviction_victim: None, + retry_failures: 0, + }; + retry_queue.schedule(event.clone(), "first failure"); + retry_queue.schedule(event, "duplicate failure"); + assert_eq!(retry_queue.len(), 1); + assert_eq!(counters.selection_debt_bytes.load(Ordering::Acquire), 64); + + let due = retry_queue.take_due_batch(Instant::now() + Duration::from_secs(10), 128); + assert_eq!(due.len(), 1); + assert_eq!(retry_queue.len(), 1); + assert_eq!(counters.selection_debt_bytes.load(Ordering::Acquire), 64); + assert!( + retry_queue + .take_due_batch(Instant::now() + Duration::from_secs(30), 128) + .is_empty(), + "a dispatched retry stays exactly once until dispatcher acknowledgement" + ); + let accepted = retry_queue + .take_for_inflight(&identity) + .expect("dispatcher acknowledgement must take the authoritative event"); + accepted.selection_debt.release(); + assert_eq!(retry_queue.len(), 0); + assert_eq!(counters.selection_debt_bytes.load(Ordering::Acquire), 0); + } + + #[test] + fn owner_hot_retry_queue_deadlines_stay_bounded_under_high_churn() { + const CHURN: usize = 20_000; + + let counters = Arc::new(OwnerHotCacheCounters::default()); + let retry_queue = OwnerHotRetryQueue::new(counters.clone()); + let identity = OwnerHotReplicaIdentity { + key: "retry-churn".to_string(), + put_time_ms: 21, + put_version: 5, + }; + let debt = OwnerHotSelectionDebt::new(64, counters.selection_debt_bytes.clone()); + let event = OwnerHotEvictionEvent { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version), + memory_info: Weak::new(), + selection_debt: debt, + retry: true, + source_eviction_victim: None, + retry_failures: 0, + }; + + for _ in 0..CHURN { + retry_queue.schedule(event.clone(), "repeated failure"); + } + assert_eq!(retry_queue.len(), 1); + assert_eq!(retry_queue.state.lock().deadlines.len(), 1); + + let due = retry_queue.take_due_batch(Instant::now() + Duration::from_secs(10), 1); + assert_eq!(due.len(), 1); + assert_eq!(retry_queue.len(), 1); + assert_eq!(retry_queue.state.lock().deadlines.len(), 0); + + for _ in 0..CHURN { + retry_queue.schedule(due[0].clone(), "repeated dispatcher failure"); + } + assert_eq!(retry_queue.len(), 1); + assert_eq!(retry_queue.state.lock().deadlines.len(), 1); + + let event = retry_queue + .take_for_inflight(&identity) + .expect("live retry must remain available to the dispatcher"); + assert_eq!(retry_queue.len(), 0); + assert_eq!(retry_queue.state.lock().deadlines.len(), 0); + event.selection_debt.release(); + assert_eq!(counters.selection_debt_bytes.load(Ordering::Acquire), 0); + + let remove_identity = OwnerHotReplicaIdentity { + key: "retry-remove-churn".to_string(), + put_time_ms: 22, + put_version: 6, + }; + let remove_debt = OwnerHotSelectionDebt::new(32, counters.selection_debt_bytes.clone()); + let remove_event = OwnerHotEvictionEvent { + key: remove_identity.key.clone(), + put_id: (remove_identity.put_time_ms, remove_identity.put_version), + memory_info: Weak::new(), + selection_debt: remove_debt, + retry: true, + source_eviction_victim: None, + retry_failures: 0, + }; + for _ in 0..CHURN { + retry_queue.schedule(remove_event.clone(), "remove churn"); + } + assert_eq!(retry_queue.len(), 1); + assert_eq!(retry_queue.state.lock().deadlines.len(), 1); + retry_queue.remove(&remove_identity); + assert_eq!(retry_queue.len(), 0); + assert_eq!(retry_queue.state.lock().deadlines.len(), 0); + assert_eq!(counters.selection_debt_bytes.load(Ordering::Acquire), 0); + } + + #[test] + fn hot_cache_only_dispatches_size_removals_and_keeps_capacity() { + let counters = Arc::new(OwnerHotCacheCounters::default()); + let retry_queue = Arc::new(OwnerHotRetryQueue::new(counters.clone())); + let (tx, mut rx) = limit_thirdparty::tokio::sync::ampsc::unbounded_channel(); + let cache = build_owner_hot_cache(10, counters.clone(), retry_queue, tx); + let entry = |put_version| OwnerHotCacheEntry { + put_id: (10, put_version), + memory_info: Weak::new(), + weight_bytes: 6, + }; + let alias = |key: &str, put_version: usize| OwnerHotPinAlias { + key: key.to_string(), + memory_info_ptr: put_version, + }; + + cache.insert("explicit".to_string(), [alias("explicit", 0)], entry(0)); + cache.run_pending_tasks(); + cache.invalidate(&"explicit".to_string()); + cache.run_pending_tasks(); + assert_eq!(counters.size_evictions.load(Ordering::Relaxed), 0); + + cache.insert("size-a".to_string(), [alias("size-a", 1)], entry(1)); + cache.insert("size-b".to_string(), [alias("size-b", 2)], entry(2)); + cache.run_pending_tasks(); + assert!(counters.size_evictions.load(Ordering::Relaxed) >= 1); + assert_eq!(cache.max_capacity(), Some(10)); + let dispatch = rx + .try_recv() + .expect("the Moka listener must emit lightweight metadata without pinning"); + let OwnerHotEvictionDispatch::Victim(event) = dispatch else { + panic!("the Moka listener must emit a victim event") + }; + assert!(event.memory_info.upgrade().is_none()); + } + + #[test] + fn pointwise_batch_visibility_skips_reclaim_fenced_keys() { + let keys = vec![ + "local-a".to_string(), + "fenced".to_string(), + "local-b".to_string(), + ]; + let controls = OwnerKeyControlTable::default(); + controls.lock_key("fenced").insert( + "fenced".to_string(), + OwnerKeyControlState { + local_puts: 0, + external_pending_puts: 0, + external_put: None, + remote_put: None, + local_ssd_put: None, + source_eviction_selection: None, + reclaim: Some(OwnerReclaimRecord::Committed(OwnerReclaimItem { + key: "fenced".to_string(), + ..OwnerReclaimItem::default() + })), + external_get: None, + local_access_fence: Some(::tokio::sync::watch::channel(false).0), + }, + ); + let mut resolved_keys = Vec::new(); + let visible = keys + .iter() + .map(|key| { + let shard = controls.lock_key(key); + if shard + .get(key) + .is_some_and(|state| state.local_access_fenced()) + { + None + } else { + resolved_keys.push(key.to_string()); + Some(key.to_string()) + } + }) + .collect::>(); + + assert_eq!( + visible, + vec![ + Some("local-a".to_string()), + None, + Some("local-b".to_string()) + ] + ); + assert_eq!(resolved_keys, vec!["local-a", "local-b"]); + } + + #[test] + fn sharded_owner_control_does_not_globally_block_unrelated_keys() { + let controls = OwnerKeyControlTable::default(); + let first = "shard-key-a"; + let first_shard = OwnerKeyControlTable::shard_index(first); + let second = (0..10_000) + .map(|idx| format!("shard-key-b-{idx}")) + .find(|key| OwnerKeyControlTable::shard_index(key) != first_shard) + .expect("a key on another owner-control shard must exist"); + + let _first_guard = controls.lock_key(first); + assert!( + controls.shards[OwnerKeyControlTable::shard_index(&second)] + .try_lock() + .is_some(), + "one key fence must not block an unrelated shard" + ); + assert!( + controls.shards[first_shard].try_lock().is_none(), + "the same shard must remain serialized while its guard is held" + ); + } + + #[test] + fn committed_slot_becomes_free_only_after_route_and_holder_are_released() { + let mut grant = OwnerLocalReserveGrantState::new(7, 0, 0, 16, 8, 2); + let slot = grant.claim_prepared_slot().expect("slot should be free"); + grant.mark_prepared_slot_pending_visible(slot.slot_index); + grant.retain_resident_slot_holder(slot.slot_index); + grant.promote_pending_visible_slot_to_committed(slot.slot_index); + + grant.release_committed_slot_route(slot.slot_index); + assert_eq!(grant.free_slots.len(), 1); + + grant.release_resident_slot_holder(slot.slot_index); + assert_eq!(grant.free_slots.len(), 2); + assert!(grant.is_fully_free()); + } + + #[test] + fn owner_reclaim_releases_committed_route_and_resident_holder_as_one_pool_update() { + let mut class = super::OwnerLocalReserveClassState::new(8, 2); + class.install_grant(OwnerLocalReserveGrantState::new(7, 0, 0, 16, 8, 2)); + let slot = class.claim_available(1).pop().expect("slot should be free"); + assert!(class.mark_prepared_slot_pending_visible(slot.grant_id, slot.slot_index)); + assert!(class.retain_resident_slot_holder(slot.grant_id, slot.slot_index)); + assert!(class.promote_pending_visible_slot_to_committed(slot.grant_id, slot.slot_index)); + assert_eq!(class.committed_slot_count(), 1); + assert_eq!(class.free_slot_count(), 1); + + assert!(class.release_committed_resident_slot(slot.grant_id, slot.slot_index)); + assert_eq!(class.committed_slot_count(), 0); + assert_eq!(class.free_slot_count(), 2); + } + + #[test] + fn failed_pending_get_releases_only_its_slot() { + let mut grant = OwnerLocalReserveGrantState::new(7, 0, 0, 24, 8, 3); + let failed = grant.claim_prepared_slot().expect("first slot"); + let unrelated = grant.claim_prepared_slot().expect("second slot"); + grant.mark_prepared_slot_pending_visible(failed.slot_index); + grant.retain_resident_slot_holder(failed.slot_index); + + grant.release_resident_slot_holder(failed.slot_index); + + assert!(matches!( + grant.slot_states[failed.slot_index as usize], + OwnerLocalReserveSlotState::Free + )); + assert!(matches!( + grant.slot_states[unrelated.slot_index as usize], + OwnerLocalReserveSlotState::Prepared + )); + assert_eq!(grant.free_slots.len(), 2); + } +} + +#[derive(Debug, Default)] +pub(crate) struct OwnerLocalReservePoolState { + pub classes: HashMap, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct MetricsSet { + pub mean: f64, + pub p99: i64, + pub p95: i64, + pub min: i64, + pub max: i64, + pub timestamps: Vec, +} + +// Removed StageScope: no longer using stage-scoped gauges; we record +// timestamps (t1..t4) and emit stage success/error directly. + +impl MetricsSet { + /// Convert to Prometheus format string + pub fn to_prometheus_format(&self, metric_name: &str, client_id: &str) -> String { + let mut result = String::new(); + + // Traditional aggregated metrics (mean, p99, p95, min, max) + result.push_str(&format!( + "kvcache_{}_mean{{client=\"{}\"}} {}\n", + metric_name, client_id, self.mean + )); + + result.push_str(&format!( + "kvcache_{}_p99{{client=\"{}\"}} {}\n", + metric_name, client_id, self.p99 + )); + + result.push_str(&format!( + "kvcache_{}_p95{{client=\"{}\"}} {}\n", + metric_name, client_id, self.p95 + )); result.push_str(&format!( "kvcache_{}_min{{client=\"{}\"}} {}\n", @@ -1071,12 +6768,452 @@ impl MetricsSet { } } +fn format_metrics_snapshot_prometheus( + client_id: &str, + timestamp_ms: i64, + metrics: &std::collections::HashMap, +) -> String { + let mut result = String::new(); + + for (metric_name, metric_set) in metrics { + result.push_str(&metric_set.to_prometheus_format(metric_name, client_id)); + result.push_str(&metric_set.to_prometheus_timeline_format(client_id)); + } + + result.push_str(&format!( + "kvcache_metrics_report_timestamp{{client=\"{}\"}} {}\n", + client_id, timestamp_ms + )); + + result +} + impl ClientKvApiInner { + pub(crate) fn owner_local_reserve_claim_lock( + &self, + slot_size: u64, + ) -> Arc> { + self.owner_local_reserve_claim_locks + .entry(slot_size) + .or_insert_with(|| Arc::new(limit_thirdparty::tokio::sync::AMutex::new(()))) + .clone() + } + pub fn get_holding_len(&self) -> usize { self.external_get_holding.total() } + + pub fn runtime_observe_snapshot(&self) -> OwnerRuntimeObserveSnapshot { + let ssd = self.kv_ssd_storage_usage_snapshot().unwrap_or_default(); + let mut external_get_holding_bytes = 0u64; + for entry in self.external_get_holding.inner().iter() { + external_get_holding_bytes = + external_get_holding_bytes.saturating_add(entry.value().memory_info.len as u64); + } + let (hot_cache_capacity_bytes, hot_cache_entries, hot_cache_weighted_bytes) = self + .owner_hot_cache + .as_ref() + .map(|cache| { + ( + cache.max_capacity().unwrap_or(0), + cache.entry_count(), + cache.weighted_size(), + ) + }) + .unwrap_or_default(); + let ( + external_get_flights, + external_get_flights_starting, + external_get_flights_finishing, + external_get_flights_revoking, + external_get_undecided_interests, + external_get_retained_interests, + ) = { + let mut flights = 0u64; + let mut starting = 0u64; + let mut finishing = 0u64; + let mut revoking = 0u64; + let mut undecided = 0u64; + let mut retained = 0u64; + // Metrics use a weak side index and never scan correctness fences. + for op in self.external_get_flight_snapshot() { + flights = flights.saturating_add(1); + let state = op.state.lock(); + undecided = undecided.saturating_add(state.undecided as u64); + retained = retained.saturating_add(state.retained as u64); + match &state.phase { + ExternalGetKeySharedPhase::Starting + | ExternalGetKeySharedPhase::Started { .. } => { + starting = starting.saturating_add(1) + } + ExternalGetKeySharedPhase::Finishing { .. } => { + finishing = finishing.saturating_add(1) + } + ExternalGetKeySharedPhase::Revoking { .. } => { + revoking = revoking.saturating_add(1) + } + ExternalGetKeySharedPhase::Ready { .. } + | ExternalGetKeySharedPhase::Failed { .. } => {} + } + } + (flights, starting, finishing, revoking, undecided, retained) + }; + let ( + local_reserve_slots_free, + local_reserve_slots_prepared, + local_reserve_slots_pending_visible, + local_reserve_slots_committed, + ) = { + let pool = self.owner_local_reserve_pool.lock(); + let mut free = 0u64; + let mut prepared = 0u64; + let mut pending_visible = 0u64; + let mut committed = 0u64; + for class in pool.classes.values() { + free = + free.saturating_add(u64::try_from(class.free_slot_count()).unwrap_or(u64::MAX)); + prepared = prepared + .saturating_add(u64::try_from(class.prepared_slot_count()).unwrap_or(u64::MAX)); + pending_visible = pending_visible.saturating_add( + u64::try_from(class.pending_visible_slot_count()).unwrap_or(u64::MAX), + ); + committed = committed.saturating_add( + u64::try_from(class.committed_slot_count()).unwrap_or(u64::MAX), + ); + } + (free, prepared, pending_visible, committed) + }; + OwnerRuntimeObserveSnapshot { + ssd_capacity_bytes: ssd.capacity_bytes, + ssd_used_bytes: ssd.used_bytes, + ssd_persist_requests: ssd.persist_requests, + ssd_persist_successes: ssd.persist_successes, + ssd_persist_failures: ssd.persist_failures, + ssd_persist_bytes: ssd.persist_bytes, + ssd_persist_duration_us: ssd.persist_duration_us, + ssd_persist_batch_requests: ssd.persist_batch_requests, + ssd_persist_batch_items: ssd.persist_batch_items, + ssd_persist_flush_batches: ssd.persist_flush_batches, + ssd_persist_busy_batches: ssd.persist_busy_batches, + ssd_persist_admission_skips: ssd.persist_admission_skips, + ssd_persist_batch_duration_us: ssd.persist_batch_duration_us, + ssd_write_candidate_items: ssd.write_candidate_items, + ssd_write_candidate_bytes: ssd.write_candidate_bytes, + ssd_write_admitted_items: ssd.write_admitted_items, + ssd_write_admitted_bytes: ssd.write_admitted_bytes, + ssd_write_dropped_items: ssd.write_dropped_items, + ssd_write_dropped_bytes: ssd.write_dropped_bytes, + ssd_write_refunded_items: ssd.write_refunded_items, + ssd_write_refunded_bytes: ssd.write_refunded_bytes, + ssd_load_requests: ssd.load_requests, + ssd_load_successes: ssd.load_successes, + ssd_load_misses: ssd.load_misses, + ssd_load_failures: ssd.load_failures, + ssd_load_bytes: ssd.load_bytes, + ssd_load_duration_us: ssd.load_duration_us, + ssd_memory_hits: ssd.memory_hits, + ssd_disk_hits: ssd.disk_hits, + ssd_outer_hits: ssd.outer_hits, + ssd_removals: ssd.removals, + ssd_stage_flights: self.ssd_stage_flights.len() as u64, + ssd_stage_terminals: self.completed_ssd_stages.entry_count(), + ssd_stage_ready_requests: self + .ssd_stage_counters + .ready_requests + .load(Ordering::Relaxed), + ssd_stage_ready_successes: self + .ssd_stage_counters + .ready_successes + .load(Ordering::Relaxed), + ssd_stage_ready_failures: self + .ssd_stage_counters + .ready_failures + .load(Ordering::Relaxed), + ssd_stage_ready_duration_us: self + .ssd_stage_counters + .ready_duration_us + .load(Ordering::Relaxed), + ssd_stage_execute_completions: self + .ssd_stage_counters + .execute_completions + .load(Ordering::Relaxed), + ssd_stage_terminal_published: self + .ssd_stage_counters + .terminal_published + .load(Ordering::Relaxed), + ssd_stage_terminal_cache_inserts: self + .ssd_stage_counters + .terminal_cache_inserts + .load(Ordering::Relaxed), + ssd_stage_terminal_cache_duration_us: self + .ssd_stage_counters + .terminal_cache_duration_us + .load(Ordering::Relaxed), + ssd_stage_response_send_attempts: self + .ssd_stage_counters + .response_send_attempts + .load(Ordering::Relaxed), + ssd_stage_response_send_successes: self + .ssd_stage_counters + .response_send_successes + .load(Ordering::Relaxed), + ssd_stage_response_send_failures: self + .ssd_stage_counters + .response_send_failures + .load(Ordering::Relaxed), + ssd_stage_response_send_duration_us: self + .ssd_stage_counters + .response_send_duration_us + .load(Ordering::Relaxed), + ssd_source_ready_wait_requests: self + .ssd_stage_counters + .source_ready_wait_requests + .load(Ordering::Relaxed), + ssd_source_ready_wait_successes: self + .ssd_stage_counters + .source_ready_wait_successes + .load(Ordering::Relaxed), + ssd_source_ready_wait_failures: self + .ssd_stage_counters + .source_ready_wait_failures + .load(Ordering::Relaxed), + ssd_source_ready_wait_duration_us: self + .ssd_stage_counters + .source_ready_wait_duration_us + .load(Ordering::Relaxed), + ssd_target_pull_requests: self + .ssd_stage_counters + .target_pull_requests + .load(Ordering::Relaxed), + ssd_target_pull_successes: self + .ssd_stage_counters + .target_pull_successes + .load(Ordering::Relaxed), + ssd_target_pull_failures: self + .ssd_stage_counters + .target_pull_failures + .load(Ordering::Relaxed), + ssd_target_pull_duration_us: self + .ssd_stage_counters + .target_pull_duration_us + .load(Ordering::Relaxed), + ssd_stage_done_detached: self + .ssd_stage_counters + .done_detached + .load(Ordering::Relaxed), + external_get_holding_entries: self.external_get_holding.total() as u64, + external_get_holding_bytes, + external_get_start_handles: self.external_get_start_registry.len() as u64, + external_get_flights, + external_get_flights_starting, + external_get_flights_finishing, + external_get_flights_revoking, + external_get_undecided_interests, + external_get_retained_interests, + owner_local_probe_batches: self + .planned_get_counters + .local_probe_batches + .load(Ordering::Relaxed), + owner_local_probe_items: self + .planned_get_counters + .local_probe_items + .load(Ordering::Relaxed), + owner_local_probe_local_items: self + .planned_get_counters + .local_probe_local_items + .load(Ordering::Relaxed), + owner_local_probe_remote_items: self + .planned_get_counters + .local_probe_remote_items + .load(Ordering::Relaxed), + planned_cpu_get_batches: self.planned_get_counters.batches.load(Ordering::Relaxed), + planned_cpu_get_local_items: self + .planned_get_counters + .local_items + .load(Ordering::Relaxed), + planned_cpu_get_leader_items: self + .planned_get_counters + .leader_items + .load(Ordering::Relaxed), + planned_cpu_get_follower_items: self + .planned_get_counters + .follower_items + .load(Ordering::Relaxed), + external_pending_put_entries: self.external_pending_puts.entry_count(), + remote_put_flights_active: self + .owner_remote_put_counters + .active + .load(Ordering::Relaxed), + remote_put_flight_leaders: self + .owner_remote_put_counters + .leaders + .load(Ordering::Relaxed), + remote_put_flight_followers: self + .owner_remote_put_counters + .followers + .load(Ordering::Relaxed), + remote_put_source_unavailable: self + .owner_remote_put_counters + .source_unavailable + .load(Ordering::Relaxed), + remote_put_source_fenced: self + .owner_remote_put_counters + .source_fenced + .load(Ordering::Relaxed), + remote_put_source_missing: self + .owner_remote_put_counters + .source_missing + .load(Ordering::Relaxed), + remote_put_source_version_mismatch: self + .owner_remote_put_counters + .source_version_mismatch + .load(Ordering::Relaxed), + remote_put_transfers: self + .owner_remote_put_counters + .transfers + .load(Ordering::Relaxed), + remote_put_published: self + .owner_remote_put_counters + .published + .load(Ordering::Relaxed), + remote_put_already_satisfied: self + .owner_remote_put_counters + .already_satisfied + .load(Ordering::Relaxed), + remote_put_obsolete: self + .owner_remote_put_counters + .obsolete + .load(Ordering::Relaxed), + remote_put_failed: self + .owner_remote_put_counters + .failed + .load(Ordering::Relaxed), + local_ssd_put_flights_active: self + .owner_local_ssd_put_counters + .active + .load(Ordering::Relaxed), + local_ssd_put_flight_leaders: self + .owner_local_ssd_put_counters + .leaders + .load(Ordering::Relaxed), + local_ssd_put_flight_followers: self + .owner_local_ssd_put_counters + .followers + .load(Ordering::Relaxed), + local_ssd_put_source_unavailable: self + .owner_local_ssd_put_counters + .source_unavailable + .load(Ordering::Relaxed), + local_ssd_put_published: self + .owner_local_ssd_put_counters + .published + .load(Ordering::Relaxed), + local_ssd_put_already_present: self + .owner_local_ssd_put_counters + .already_present + .load(Ordering::Relaxed), + local_ssd_put_dropped: self + .owner_local_ssd_put_counters + .dropped + .load(Ordering::Relaxed), + local_ssd_put_obsolete: self + .owner_local_ssd_put_counters + .obsolete + .load(Ordering::Relaxed), + local_ssd_put_failed: self + .owner_local_ssd_put_counters + .failed + .load(Ordering::Relaxed), + local_reserve_slots_free, + local_reserve_slots_prepared, + local_reserve_slots_pending_visible, + local_reserve_slots_committed, + hot_cache_capacity_bytes, + hot_cache_entries, + hot_cache_weighted_bytes, + hot_size_evictions: self + .owner_hot_counters + .size_evictions + .load(Ordering::Relaxed), + hot_source_evict_handoff_members: self + .owner_hot_counters + .source_evict_handoff_members + .load(Ordering::Relaxed), + hot_source_evict_committed_members: self + .owner_hot_counters + .source_evict_committed_members + .load(Ordering::Relaxed), + hot_source_evict_restored_members: self + .owner_hot_counters + .source_evict_restored_members + .load(Ordering::Relaxed), + hot_source_evict_obsolete: self + .owner_hot_counters + .source_evict_obsolete + .load(Ordering::Relaxed), + hot_source_evict_dispatch_failed: self + .owner_hot_counters + .source_evict_dispatch_failed + .load(Ordering::Relaxed), + hot_source_eviction_selected: self.owner_source_eviction_selected.len() as u64, + hot_source_evict_retry_entries: self.owner_hot_retry_queue.len() as u64, + hot_source_evict_retry_scheduled: self + .owner_hot_counters + .source_evict_retry_scheduled + .load(Ordering::Relaxed), + hot_source_evict_retry_emitted: self + .owner_hot_counters + .source_evict_retry_emitted + .load(Ordering::Relaxed), + hot_selection_debt_bytes: self + .owner_hot_counters + .selection_debt_bytes + .load(Ordering::Relaxed), + hot_source_eviction_selected_bytes: self + .owner_hot_counters + .source_eviction_selected_bytes + .load(Ordering::Relaxed), + hot_eviction_skipped_stale: self + .owner_hot_counters + .skipped_stale + .load(Ordering::Relaxed), + hot_eviction_skipped_reclaim: self + .owner_hot_counters + .skipped_reclaim + .load(Ordering::Relaxed), + hot_eviction_skipped_active_holders: self + .owner_hot_counters + .skipped_active_holders + .load(Ordering::Relaxed), + hot_victim_duplicates: self + .owner_hot_counters + .victim_duplicates + .load(Ordering::Relaxed), + hot_victim_invalid_backing: self + .owner_hot_counters + .victim_invalid_backing + .load(Ordering::Relaxed), + grouped_put_done_batches: self + .owner_hot_counters + .grouped_put_done_batches + .load(Ordering::Relaxed), + grouped_put_done_items: self + .owner_hot_counters + .grouped_put_done_items + .load(Ordering::Relaxed), + legacy_put_done_batches: self + .owner_hot_counters + .legacy_put_done_batches + .load(Ordering::Relaxed), + legacy_put_done_items: self + .owner_hot_counters + .legacy_put_done_items + .load(Ordering::Relaxed), + } + } + pub fn get_cache_len(&self) -> usize { - self.get_cached_info.len() + self.precommit_local_visible_info.len() + self.get_cached_info.len() } fn metrics_handle(&self) -> Arc { self.metrics @@ -1085,6 +7222,15 @@ impl ClientKvApiInner { .expect("metrics handle not initialized") } + pub fn locality_snapshot(&self) -> KvLocalitySnapshot { + self.metrics_handle().get_locality_snapshot() + } + + pub fn record_put_locality(&self, remote: bool, bytes: u64, transfer_us: i64) { + self.metrics_handle() + .record_put_io_locality(remote, bytes, transfer_us); + } + fn client_id_str(&self) -> String { self.view.cluster_manager().get_self_info().id.to_string() } @@ -1506,26 +7652,285 @@ impl ClientKvApiInner { new_ref } + + pub(crate) fn owner_local_reserve_rebalance_notify( + &self, + ) -> Arc { + self.owner_local_reserve_rebalance_notify.clone() + } + + pub(crate) fn owner_local_reserve_register_pending_demand( + &self, + slot_size: u64, + slots_per_grant: u32, + demand_slots: usize, + ) { + let mut pool = self.owner_local_reserve_pool.lock(); + let class_state = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)); + class_state.pending_slot_demand = + class_state.pending_slot_demand.saturating_add(demand_slots); + class_state.max_observed_claim_slots = + class_state.max_observed_claim_slots.max(demand_slots); + } + + pub(crate) fn owner_local_reserve_consume_pending_demand( + &self, + slot_size: u64, + slots_per_grant: u32, + demand_slots: usize, + ) { + let mut pool = self.owner_local_reserve_pool.lock(); + let class_state = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)); + class_state.pending_slot_demand = + class_state.pending_slot_demand.saturating_sub(demand_slots); + } } impl ClientKvApi { pub fn inner(&self) -> &ClientKvApiInner { &self.0 } + fn spawn_runtime_observe_reporter(&self) { + let view = self.0.view.clone_view(); + let view_task = view.clone(); + view.spawn("client_runtime_observe_reporter", async move { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + tokio::select! { + _ = shutdown_waiter.wait() => break, + _ = interval.tick() => { + let snapshot = view_task.client_kv_api().inner().runtime_observe_snapshot(); + if snapshot.ssd_capacity_bytes > 0 { + tracing::info!( + capacity_bytes = snapshot.ssd_capacity_bytes, + used_bytes = snapshot.ssd_used_bytes, + persist_requests = snapshot.ssd_persist_requests, + persist_successes = snapshot.ssd_persist_successes, + persist_failures = snapshot.ssd_persist_failures, + persist_bytes = snapshot.ssd_persist_bytes, + persist_duration_us = snapshot.ssd_persist_duration_us, + persist_batch_requests = snapshot.ssd_persist_batch_requests, + persist_batch_items = snapshot.ssd_persist_batch_items, + persist_flush_batches = snapshot.ssd_persist_flush_batches, + persist_busy_batches = snapshot.ssd_persist_busy_batches, + persist_admission_skips = snapshot.ssd_persist_admission_skips, + persist_batch_duration_us = snapshot.ssd_persist_batch_duration_us, + write_candidate_items = snapshot.ssd_write_candidate_items, + write_candidate_bytes = snapshot.ssd_write_candidate_bytes, + write_admitted_items = snapshot.ssd_write_admitted_items, + write_admitted_bytes = snapshot.ssd_write_admitted_bytes, + write_dropped_items = snapshot.ssd_write_dropped_items, + write_dropped_bytes = snapshot.ssd_write_dropped_bytes, + write_refunded_items = snapshot.ssd_write_refunded_items, + write_refunded_bytes = snapshot.ssd_write_refunded_bytes, + load_requests = snapshot.ssd_load_requests, + load_successes = snapshot.ssd_load_successes, + load_misses = snapshot.ssd_load_misses, + load_failures = snapshot.ssd_load_failures, + load_bytes = snapshot.ssd_load_bytes, + load_duration_us = snapshot.ssd_load_duration_us, + memory_hits = snapshot.ssd_memory_hits, + disk_hits = snapshot.ssd_disk_hits, + outer_hits = snapshot.ssd_outer_hits, + removals = snapshot.ssd_removals, + active_stage_flights = snapshot.ssd_stage_flights, + retained_stage_terminals = snapshot.ssd_stage_terminals, + "owner KV SSD storage snapshot" + ); + } + let metrics = view_task.metric_reporter().metrics(); + metrics.set_kv_holding_entries( + "owner_external_get_holding", + snapshot.external_get_holding_entries, + ); + metrics.set_kv_holding_bytes( + "owner_external_get_holding", + snapshot.external_get_holding_bytes, + ); + metrics.set_kv_external_pending_put_entries( + snapshot.external_pending_put_entries, + ); + tracing::info!( + active_handles = snapshot.external_get_start_handles, + active_flights = snapshot.external_get_flights, + starting_flights = snapshot.external_get_flights_starting, + finishing_flights = snapshot.external_get_flights_finishing, + revoking_flights = snapshot.external_get_flights_revoking, + undecided_interests = snapshot.external_get_undecided_interests, + retained_interests = snapshot.external_get_retained_interests, + local_probe_batches = snapshot.owner_local_probe_batches, + local_probe_items = snapshot.owner_local_probe_items, + local_probe_local_items = snapshot.owner_local_probe_local_items, + local_probe_remote_items = snapshot.owner_local_probe_remote_items, + planned_cpu_batches = snapshot.planned_cpu_get_batches, + planned_cpu_local_items = snapshot.planned_cpu_get_local_items, + planned_cpu_leader_items = snapshot.planned_cpu_get_leader_items, + planned_cpu_follower_items = snapshot.planned_cpu_get_follower_items, + ssd_stage_ready_requests = snapshot.ssd_stage_ready_requests, + ssd_stage_ready_successes = snapshot.ssd_stage_ready_successes, + ssd_stage_ready_failures = snapshot.ssd_stage_ready_failures, + ssd_stage_ready_duration_us = snapshot.ssd_stage_ready_duration_us, + ssd_stage_execute_completions = snapshot.ssd_stage_execute_completions, + ssd_stage_terminal_published = snapshot.ssd_stage_terminal_published, + ssd_stage_terminal_cache_inserts = snapshot.ssd_stage_terminal_cache_inserts, + ssd_stage_terminal_cache_duration_us = snapshot.ssd_stage_terminal_cache_duration_us, + ssd_stage_response_send_attempts = snapshot.ssd_stage_response_send_attempts, + ssd_stage_response_send_successes = snapshot.ssd_stage_response_send_successes, + ssd_stage_response_send_failures = snapshot.ssd_stage_response_send_failures, + ssd_stage_response_send_duration_us = snapshot.ssd_stage_response_send_duration_us, + ssd_source_ready_wait_requests = snapshot.ssd_source_ready_wait_requests, + ssd_source_ready_wait_successes = snapshot.ssd_source_ready_wait_successes, + ssd_source_ready_wait_failures = snapshot.ssd_source_ready_wait_failures, + ssd_source_ready_wait_duration_us = snapshot.ssd_source_ready_wait_duration_us, + ssd_target_pull_requests = snapshot.ssd_target_pull_requests, + ssd_target_pull_successes = snapshot.ssd_target_pull_successes, + ssd_target_pull_failures = snapshot.ssd_target_pull_failures, + ssd_target_pull_duration_us = snapshot.ssd_target_pull_duration_us, + ssd_stage_done_detached = snapshot.ssd_stage_done_detached, + reserve_free = snapshot.local_reserve_slots_free, + reserve_prepared = snapshot.local_reserve_slots_prepared, + reserve_pending_visible = snapshot.local_reserve_slots_pending_visible, + reserve_committed = snapshot.local_reserve_slots_committed, + "owner Get lifecycle snapshot" + ); + tracing::info!( + active = snapshot.remote_put_flights_active, + leaders = snapshot.remote_put_flight_leaders, + followers = snapshot.remote_put_flight_followers, + source_unavailable = snapshot.remote_put_source_unavailable, + source_fenced = snapshot.remote_put_source_fenced, + source_missing = snapshot.remote_put_source_missing, + source_version_mismatch = snapshot.remote_put_source_version_mismatch, + transfers = snapshot.remote_put_transfers, + published = snapshot.remote_put_published, + already_satisfied = snapshot.remote_put_already_satisfied, + obsolete = snapshot.remote_put_obsolete, + failed = snapshot.remote_put_failed, + "owner unified remote Put flight snapshot" + ); + tracing::info!( + active = snapshot.local_ssd_put_flights_active, + leaders = snapshot.local_ssd_put_flight_leaders, + followers = snapshot.local_ssd_put_flight_followers, + source_unavailable = snapshot.local_ssd_put_source_unavailable, + published = snapshot.local_ssd_put_published, + already_present = snapshot.local_ssd_put_already_present, + dropped = snapshot.local_ssd_put_dropped, + obsolete = snapshot.local_ssd_put_obsolete, + failed = snapshot.local_ssd_put_failed, + "owner local SSD Put flight snapshot" + ); + if snapshot.hot_cache_capacity_bytes > 0 { + tracing::info!( + capacity_bytes = snapshot.hot_cache_capacity_bytes, + entries = snapshot.hot_cache_entries, + weighted_bytes = snapshot.hot_cache_weighted_bytes, + size_evictions = snapshot.hot_size_evictions, + source_evict_handoff_members = snapshot.hot_source_evict_handoff_members, + source_evict_committed_members = snapshot.hot_source_evict_committed_members, + source_evict_restored_members = snapshot.hot_source_evict_restored_members, + source_evict_obsolete = snapshot.hot_source_evict_obsolete, + source_evict_dispatch_failed = snapshot.hot_source_evict_dispatch_failed, + source_eviction_selected = snapshot.hot_source_eviction_selected, + source_evict_retry_entries = snapshot.hot_source_evict_retry_entries, + source_evict_retry_scheduled = snapshot.hot_source_evict_retry_scheduled, + source_evict_retry_emitted = snapshot.hot_source_evict_retry_emitted, + selection_debt_bytes = snapshot.hot_selection_debt_bytes, + source_eviction_selected_bytes = snapshot.hot_source_eviction_selected_bytes, + skipped_stale = snapshot.hot_eviction_skipped_stale, + skipped_reclaim = snapshot.hot_eviction_skipped_reclaim, + skipped_active_holders = snapshot.hot_eviction_skipped_active_holders, + victim_duplicates = snapshot.hot_victim_duplicates, + victim_invalid_backing = snapshot.hot_victim_invalid_backing, + grouped_put_done_batches = snapshot.grouped_put_done_batches, + grouped_put_done_items = snapshot.grouped_put_done_items, + legacy_put_done_batches = snapshot.legacy_put_done_batches, + legacy_put_done_items = snapshot.legacy_put_done_items, + "owner hot source-eviction policy snapshot" + ); + } + } + } + } + }); + } + pub fn attach_view(&self, view: ClientKvApiView) { self.0.view.attach(view); } pub async fn construct(arg: ClientKvApiNewArg) -> Result { tracing::info!("Constructing ClientKvApi in Client mode (PreView)"); + let ClientKvApiNewArg { + test_spec_config, + owner_hot_cache_capacity_bytes, + ssd_storage, + } = arg; + let ssd_storage = match ssd_storage { + Some(init) => Some(Arc::new(KvSsdStorage::new(init).await?)), + None => None, + }; + let (owner_local_publish_tx, owner_local_publish_rx) = + tokio::sync::ampsc::channel(OWNER_LOCAL_PUBLISH_QUEUE_CAPACITY); + // The Moka eviction listener is synchronous and must never block while + // holding Moka's housekeeper lock. Events contain only weak payload + // references and are deduplicated by exact selected identities, so use + // a lossless metadata channel instead of dropping victims when the old + // bounded queue briefly filled under cache pressure. + let (owner_hot_eviction_tx, owner_hot_eviction_rx) = + tokio::sync::ampsc::unbounded_channel(); + let get_cached_info = Arc::new(DashMap::new()); + let owner_key_control = Arc::new(OwnerKeyControlTable::default()); + let owner_source_eviction_selected = Arc::new(DashMap::new()); + let owner_hot_counters = Arc::new(OwnerHotCacheCounters::default()); + let owner_remote_put_counters = Arc::new(OwnerRemotePutCounters::default()); + let owner_local_ssd_put_counters = Arc::new(OwnerLocalSsdPutCounters::default()); + let owner_hot_retry_queue = Arc::new(OwnerHotRetryQueue::new(owner_hot_counters.clone())); + let owner_hot_cache = owner_hot_cache_capacity_bytes.map(|capacity_bytes| { + build_owner_hot_cache( + capacity_bytes, + owner_hot_counters.clone(), + owner_hot_retry_queue.clone(), + owner_hot_eviction_tx.clone(), + ) + }); let inner = ClientKvApiInner { view: ClientKvApiViewHolder::new(), - test_spec_config: arg.test_spec_config, + test_spec_config, + ssd_storage, metrics: OnceLock::new(), all_memholder_refcount: OnceLock::new(), get_remote_kv_lock: AMapLock::new(Duration::from_secs(60)), - get_cached_info: DashMap::new(), + get_cached_info, + precommit_local_visible_info: DashMap::new(), + pending_local_get_info: DashMap::new(), + local_snapshot_info: DashMap::new(), + owner_local_reserve_pool: Mutex::new(OwnerLocalReservePoolState::default()), + owner_local_reserve_claim_locks: DashMap::new(), + owner_local_reserve_rebalance_notify: Arc::new( + limit_thirdparty::tokio::sync::Notify::new(), + ), + external_local_first_put_id_counter: AtomicU32::new(0), + next_owner_source_eviction_operation_id: AtomicU64::new(1), + owner_key_control, + owner_hot_cache, + owner_source_eviction_selected, + owner_hot_counters, + owner_remote_put_counters, + owner_local_ssd_put_counters, + planned_get_counters: OwnerPlannedGetCounters::default(), + ssd_stage_counters: OwnerSsdStageCounters::default(), + owner_hot_retry_queue, + owner_hot_eviction_tx, + owner_hot_eviction_rx: Mutex::new(Some(owner_hot_eviction_rx)), external_invalidate_delete: EnsureMemholderMgmtDeleteHandle::new( OwnerExternalMemMgr::DELETE_SUBMIT_QUEUE_CAPACITY, ), @@ -1534,6 +7939,22 @@ impl ClientKvApi { ), owner_delete_ack_mgr: OwnerDeleteAckMemMgr::default(), external_get_holding: OwnerExternalMemMgr::default(), + external_get_start_registry: DashMap::new(), + external_get_flight_registry: DashMap::new(), + external_get_local_probe_locks: AMapLock::new(Duration::from_secs(120)), + completed_external_get_local_probes: moka::future::Cache::builder() + .time_to_live(Duration::from_secs(120)) + .build(), + planned_external_get_execute_locks: AMapLock::new(Duration::from_secs(120)), + completed_planned_external_get_executes: moka::future::Cache::builder() + .time_to_live(Duration::from_secs(120)) + .build(), + ssd_stage_flights: DashMap::new(), + completed_ssd_stages: moka::future::Cache::builder() + .time_to_live(SSD_STAGE_TERMINAL_TTL) + .build(), + next_external_get_start_handle: AtomicU64::new(1), + next_external_holding_id: AtomicU64::new(1), external_pending_puts: moka::sync::Cache::builder() .time_to_live(Duration::from_secs(30 * 60)) .segments(16) @@ -1543,18 +7964,43 @@ impl ClientKvApi { rpc_caller_get_start: RPCCaller::new(), rpc_caller_get_revoke: RPCCaller::new(), rpc_caller_get_done: RPCCaller::new(), + rpc_caller_batch_get_start: RPCCaller::new(), + rpc_caller_batch_get_bind: RPCCaller::new(), + rpc_caller_batch_get_revoke: RPCCaller::new(), + rpc_caller_batch_get_done: RPCCaller::new(), rpc_caller_put_start: RPCCaller::new(), rpc_caller_put_revoke: RPCCaller::new(), rpc_caller_put_done: RPCCaller::new(), + rpc_caller_batch_put_start: RPCCaller::new(), + rpc_caller_batch_put_revoke: RPCCaller::new(), + rpc_caller_batch_put_done: RPCCaller::new(), + rpc_caller_grouped_batch_put_done: RPCCaller::new(), + rpc_caller_batch_prepare_put_keys: RPCCaller::new(), + rpc_caller_batch_release_put_key_reservations: RPCCaller::new(), + rpc_caller_put_append_start: RPCCaller::new(), + rpc_caller_batch_put_append_start: RPCCaller::new(), + rpc_caller_put_append_revoke: RPCCaller::new(), + rpc_caller_put_append_done: RPCCaller::new(), + rpc_caller_batch_put_append_done: RPCCaller::new(), + rpc_caller_batch_evict_owner_source: RPCCaller::new(), + rpc_caller_batch_publish_owner_ssd: RPCCaller::new(), + rpc_caller_reserve_local_grant: RPCCaller::new(), + rpc_caller_release_local_grant: RPCCaller::new(), rpc_caller_delete: RPCCaller::new(), rpc_caller_batch_delete_ack: RPCCaller::new(), + rpc_caller_batch_is_exist: RPCCaller::new(), rpc_caller_get_meta: RPCCaller::new(), - _rpc_caller_allocate_client_lease: RPCCaller::new(), - _rpc_caller_client_lease_keepalive: RPCCaller::new(), + rpc_caller_allocate_client_lease: RPCCaller::new(), + rpc_caller_client_lease_keepalive: RPCCaller::new(), + rpc_caller_ssd_stage_read: RPCCaller::new(), + rpc_caller_ssd_stage_begin: RPCCaller::new(), + rpc_caller_ssd_stage_done: RPCCaller::new(), rpc_caller_external_put_commit: RPCCaller::new(), rpc_caller_external_put_revoke: RPCCaller::new(), rpc_caller_resolve_side_transfer_lane: RPCCaller::new(), default_lease_id: parking_lot::RwLock::new(None), + owner_local_publish_tx, + owner_local_publish_rx: Mutex::new(Some(owner_local_publish_rx)), }; Ok(Self(inner)) } @@ -1570,14 +8016,83 @@ impl ClientKvApi { inner.rpc_caller_get_start.regist(inner.view.p2p_module()); inner.rpc_caller_get_revoke.regist(inner.view.p2p_module()); inner.rpc_caller_get_done.regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_get_start + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_get_bind + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_get_revoke + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_get_done + .regist(inner.view.p2p_module()); inner.rpc_caller_put_start.regist(inner.view.p2p_module()); inner.rpc_caller_put_revoke.regist(inner.view.p2p_module()); inner.rpc_caller_put_done.regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_put_start + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_put_revoke + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_put_done + .regist(inner.view.p2p_module()); + inner + .rpc_caller_grouped_batch_put_done + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_prepare_put_keys + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_release_put_key_reservations + .regist(inner.view.p2p_module()); + inner + .rpc_caller_put_append_start + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_put_append_start + .regist(inner.view.p2p_module()); + inner + .rpc_caller_put_append_revoke + .regist(inner.view.p2p_module()); + inner + .rpc_caller_put_append_done + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_put_append_done + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_evict_owner_source + .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_publish_owner_ssd + .regist(inner.view.p2p_module()); + inner + .rpc_caller_reserve_local_grant + .regist(inner.view.p2p_module()); + inner + .rpc_caller_release_local_grant + .regist(inner.view.p2p_module()); inner.rpc_caller_delete.regist(inner.view.p2p_module()); inner .rpc_caller_batch_delete_ack .regist(inner.view.p2p_module()); + inner + .rpc_caller_batch_is_exist + .regist(inner.view.p2p_module()); inner.rpc_caller_get_meta.regist(inner.view.p2p_module()); + inner + .rpc_caller_ssd_stage_read + .regist(inner.view.p2p_module()); + inner + .rpc_caller_ssd_stage_begin + .regist(inner.view.p2p_module()); + inner + .rpc_caller_ssd_stage_done + .regist(inner.view.p2p_module()); inner .rpc_caller_external_put_commit .regist(inner.view.p2p_module()); @@ -1592,25 +8107,161 @@ impl ClientKvApi { // Register master-only metric RPC callers crate::metrics::client::init_for_p2p_owner(inner.view.p2p_module()); RPCCaller::::new().regist(inner.view.p2p_module()); + RPCCaller::::new().regist(inner.view.p2p_module()); RPCCaller::::new().regist(inner.view.p2p_module()); + spawn_owner_local_reserve_rebalance_actor(inner.view.clone_view()); + spawn_owner_slot_pressure_actor(inner.view.clone_view()); + external_api::spawn_external_get_start_handle_sweeper(inner.view.clone_view()); + self.spawn_runtime_observe_reporter(); + + let view_ssd = inner.view.clone_view(); + RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { + let view = view_ssd.clone(); + let task_view = view.clone(); + view.spawn("rpc_ssd_stage_read", async move { + let get_id = msg.serialize_part.get_id; + let peer = resp.node_id(); + let task_id = resp.task_id(); + let result = handle_ssd_stage_read(&task_view, &msg).await; + let inner = task_view.client_kv_api().inner(); + inner + .ssd_stage_counters + .response_send_attempts + .fetch_add(1, Ordering::Relaxed); + let send_started_at = Instant::now(); + let send_result = resp + .send_resp_with_transport_policy(result, RpcTransportPolicy::ForceTransport) + .await; + inner + .ssd_stage_counters + .response_send_duration_us + .fetch_add( + u64::try_from(send_started_at.elapsed().as_micros()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + if let Err(err) = send_result { + inner + .ssd_stage_counters + .response_send_failures + .fetch_add(1, Ordering::Relaxed); + tracing::warn!( + get_id, + peer = %peer, + task_id, + error = ?err, + "SSD stage-ready response send failed" + ); + } else { + inner + .ssd_stage_counters + .response_send_successes + .fetch_add(1, Ordering::Relaxed); + } + }); + Ok(()) + }); // External RPC handlers let view_ext = inner.view.clone_view(); RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_get", async move { + view.spawn("rpc_external_get", async move { let result = handle_external_get(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); Ok(()) }); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_get", async move { + let result = handle_external_batch_get(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_get_local_probe", async move { + let result = handle_external_batch_get_local_probe(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_get_start", async move { + let result = handle_external_batch_get_start(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_get_transfer", async move { + let result = handle_external_batch_get_transfer(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_get_cancel", async move { + let result = handle_external_batch_get_cancel(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_execute_planned_get", async move { + let result = handle_external_execute_planned_get(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_put_start", async move { + view.spawn("rpc_external_put_start", async move { let req = msg.serialize_part.clone(); tracing::info!( "rpc_external_put_start received: self={} peer={} task_id={} key={} len={} started_time={}", @@ -1644,13 +8295,27 @@ impl ClientKvApi { Ok(()) }); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_put_start", async move { + let result = handle_external_batch_put_start(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist( inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_put_transfer_end", async move { + view.spawn("rpc_external_put_transfer_end", async move { let result = handle_external_put_transfer_end(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1658,13 +8323,27 @@ impl ClientKvApi { }, ); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_put_transfer_end", async move { + let result = handle_external_batch_put_transfer_end(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist( inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_put_commit", async move { + view.spawn("rpc_external_put_commit", async move { let result = handle_external_put_commit(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1672,13 +8351,27 @@ impl ClientKvApi { }, ); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_put_commit", async move { + let result = handle_external_batch_put_commit(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist( inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_put_revoke", async move { + view.spawn("rpc_external_put_revoke", async move { let result = handle_external_put_revoke(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1690,7 +8383,7 @@ impl ClientKvApi { RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_delete", async move { + view.spawn("rpc_external_delete", async move { let result = handle_external_delete(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1703,7 +8396,7 @@ impl ClientKvApi { move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_is_exist", async move { + view.spawn("rpc_external_is_exist", async move { let result = handle_external_is_exist(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1711,13 +8404,41 @@ impl ClientKvApi { }, ); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_is_exist", async move { + let result = handle_external_batch_is_exist(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_observability_snapshot", async move { + let result = handle_external_observability_snapshot(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + let view_ext = inner.view.clone_view(); RPCHandler::::new().regist( inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_external_delete_ack", async move { + view.spawn("rpc_external_delete_ack", async move { let result = handle_external_delete_ack(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1725,13 +8446,27 @@ impl ClientKvApi { }, ); + let view_ext = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_batch_delete_ack", async move { + let result = handle_external_batch_delete_ack(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + // KV->file sync RPC (bytes field -> file@offset) RPCCaller::::new().regist(inner.view.p2p_module()); let view_ext = inner.view.clone_view(); RPCHandler::::new().regist(inner.view.p2p_module(), move |resp, msg| { let view = view_ext.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_sync_kv_to_file", async move { + view.spawn("rpc_sync_kv_to_file", async move { let result = handle_sync_kv_to_file_client(&view_task, &msg).await; let _ = resp.send_resp(result).await; }); @@ -1739,6 +8474,39 @@ impl ClientKvApi { }); // client rpc handler register + let view = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let view = view.clone(); + let view_task = view.clone(); + view.spawn("rpc_batch_enqueue_replica_tasks", async move { + let ack = put::handle_batch_enqueue_replica_tasks(&view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + warn!("Failed to send BatchEnqueueReplicaTaskResp: {:?}", e); + } + }); + Ok(()) + }, + ); + + let view = inner.view.clone_view(); + RPCHandler::::new().regist( + inner.view.p2p_module(), + move |resp, msg| { + let req_node_id = resp.node_id().clone(); + let view = view.clone(); + let view_task = view.clone(); + view.spawn("rpc_batch_owner_reclaim", async move { + let ack = handle_batch_owner_reclaim(&view_task, msg, req_node_id).await; + if let Err(e) = resp.send_resp(ack).await { + warn!("Failed to send BatchOwnerReclaimResp: {:?}", e); + } + }); + Ok(()) + }, + ); + let view = inner.view.clone_view(); RPCHandler::::new().regist( inner.view.p2p_module(), @@ -1746,7 +8514,7 @@ impl ClientKvApi { let req_node_id = resp.node_id().clone(); let view = view.clone(); let view_task = view.clone(); - let _ = view.spawn("rpc_batch_delete_client_kv_meta_cache", async move { + view.spawn("rpc_batch_delete_client_kv_meta_cache", async move { let ack = handle_batch_delete_client_kv_meta_cache(&view_task, msg, req_node_id) .await; @@ -1773,11 +8541,32 @@ impl ClientKvApi { .expect("delete_ack_batch rx already taken, that's impossible"); delete::spawn_owner_delete_ack_batch(inner.view.clone_view(), delete_ack_batch_rx); - // Spawn cluster listener to clean up get_holding when external_client leaves + if inner.owner_hot_cache.is_some() { + if let Some(owner_hot_eviction_rx) = inner.owner_hot_eviction_rx.lock().take() { + put::spawn_owner_source_eviction_dispatcher( + inner.view.clone_view(), + owner_hot_eviction_rx, + ); + put::spawn_owner_hot_retry_actor(inner.view.clone_view()); + } else { + tracing::warn!("owner_hot_eviction_rx already taken for ClientKvApi"); + } + } + if let Some(owner_local_publish_rx) = inner.owner_local_publish_rx.lock().take() { + put::spawn_owner_local_publish_dispatcher( + inner.view.clone_view(), + owner_local_publish_rx, + OWNER_LOCAL_PUBLISH_MAX_INFLIGHT, + ); + } else { + tracing::warn!("owner_local_publish_rx already taken for ClientKvApi"); + } + + // Spawn cluster listener to retire generation-scoped external requester state. let view = inner.view.clone_view(); let view2 = view.clone(); let view_task = view2.clone(); - let _ = view.spawn("client_cluster_listener", async move { + view.spawn("client_cluster_listener", async move { let mut listen_cluster_event = view_task.cluster_manager().listen(); let mut shutdown_waiter = view_task.register_shutdown_waiter(); @@ -1788,15 +8577,43 @@ impl ClientKvApi { Ok(event) => { match event { ClusterEvent::MemberLeft(node_id) => { - let removed = view_task - .client_kv_api() - .inner() + let departed_epoch = view_task + .cluster_manager() + .get_prev_member_info(&node_id) + .map(|member| member.node_start_time); + let current_epoch = view_task + .cluster_manager() + .get_member_info_cached(&node_id) + .map(|member| member.node_start_time); + let Some(departed_epoch) = external_api::external_member_left_departed_epoch( + departed_epoch, + current_epoch, + ) else { + tracing::debug!( + "Ignoring ambiguous/delayed external MemberLeft: node={} departed_epoch={:?} current_epoch={:?}", + node_id, + departed_epoch, + current_epoch, + ); + continue; + }; + + let inner = view_task.client_kv_api().inner(); + let removed_handles = external_api::cleanup_external_get_start_handles_for_generation( + &inner.external_get_start_registry, + &node_id, + departed_epoch, + ); + let removed_holdings = inner .external_get_holding - .cleanup_node(&node_id); - if removed > 0 { + .cleanup_node_generation(&node_id, departed_epoch); + if removed_handles > 0 || removed_holdings > 0 { tracing::info!( - "Cleaned up get_holding for external_client: {} (removed {} holdings)", - node_id, removed + "Cleaned up departed external requester state: node={} epoch={} handles={} holdings={}", + node_id, + departed_epoch, + removed_handles, + removed_holdings, ); } } @@ -1804,8 +8621,11 @@ impl ClientKvApi { } } Err(e) => { - tracing::error!("Failed to receive cluster event: {}", e); - break; + tracing::warn!( + "Client cluster event receiver error (will resubscribe): {}", + e + ); + listen_cluster_event = view_task.cluster_manager().listen(); } } } @@ -1917,9 +8737,8 @@ impl ClientKvApi { tracing::info!("------------------------------------------------------------"); } - #[cfg(any(test, feature = "test_bins"))] pub fn has_cached_key(&self, key: &str) -> bool { - self.inner().get_cached_info.contains_key(key) + self.inner().has_local_snapshot(key) } // Removed is_client_mode(): ClientKvApi is owner-only and always constructed. @@ -1956,6 +8775,9 @@ impl LogicalModule for ClientKvApi { } async fn shutdown(&self) -> Result<(), Self::Error> { tracing::info!("ClientKvApi shutting down..."); + if let Some(store) = self.0.ssd_storage.as_ref() { + store.close().await?; + } tracing::info!( "ClientKvApi final: holding_len={} , cache_len={}", self.0.get_holding_len(), diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs index 55f0970..4c87394 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/msg_pack.rs @@ -1,6 +1,6 @@ use crate::master_kv_router::put::PutIDForAKey; use crate::p2p::msg_pack::{MsgPackSerializePart, RPCReq}; -use crate::rpcresp_kvresult_convert::msg_and_error::ErrorCode; +use crate::rpcresp_kvresult_convert::msg_and_error::{ErrorCode, OK}; use bitcode::{Decode, Encode}; use crate::memholder::ExternalMemHolderInfo; @@ -59,6 +59,38 @@ impl TestPutPhaseTrace { merge_i64_field!(owner_master_put_end_server_us); } } + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct SsdStageReadReq { + pub key: String, + pub put_id: PutIDForAKey, + pub get_id: u64, + pub stage_addr: u64, + pub stage_capacity: u64, + pub len: u64, +} + +impl MsgPackSerializePart for SsdStageReadReq { + fn msg_id(&self) -> u32 { + 4044 + } +} + +impl RPCReq for SsdStageReadReq { + type Resp = SsdStageReadResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageReadResp { + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for SsdStageReadResp { + fn msg_id(&self) -> u32 { + 4045 + } +} // --- RPC for Physical Node Shared Memory --- #[derive(Default, Debug, Clone, Encode, Decode)] @@ -89,6 +121,625 @@ impl MsgPackSerializePart for ExternalGetResp { } } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetReq { + pub keys: Vec, + pub req_node_id: String, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, + pub transfer_concurrency: usize, +} +impl MsgPackSerializePart for ExternalBatchGetReq { + fn msg_id(&self) -> u32 { + 4020 + } +} +impl RPCReq for ExternalBatchGetReq { + type Resp = ExternalBatchGetResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetItemResp { + pub error_code: ErrorCode, + pub error_json: String, + pub external_memholder_info: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +/// Probe only the requester's share-group owner. A `Some` item is pinned +/// owner-local data; `None` must be resolved through the remote directory. +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetLocalProbeReq { + pub plan_handle: u64, + pub keys: Vec, + pub req_node_id: String, + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalBatchGetLocalProbeReq { + fn msg_id(&self) -> u32 { + 4040 + } +} +impl RPCReq for ExternalBatchGetLocalProbeReq { + type Resp = ExternalBatchGetLocalProbeResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetLocalProbeResp { + pub items: Vec>, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchGetLocalProbeResp { + fn msg_id(&self) -> u32 { + 4041 + } +} + +#[cfg(test)] +mod external_batch_get_local_probe_wire_tests { + use super::{ExternalBatchGetLocalProbeReq, ExternalBatchGetLocalProbeResp}; + use crate::memholder::ExternalMemHolderInfo; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + + #[test] + fn mixed_local_remote_probe_round_trip() { + let request = ExternalBatchGetLocalProbeReq { + plan_handle: 41, + keys: vec!["local".to_string(), "remote".to_string()], + req_node_id: "external-a".to_string(), + started_time: 17, + }; + let decoded: ExternalBatchGetLocalProbeReq = + bitcode::decode(&bitcode::encode(&request)).expect("decode local probe request"); + assert_eq!(decoded.plan_handle, 41); + assert_eq!(decoded.keys, request.keys); + + let response = ExternalBatchGetLocalProbeResp { + items: vec![ + Some(ExternalMemHolderInfo { + offset: 4096, + len: 8192, + holder_id: 23, + }), + None, + ], + error_code: OK, + error_json: String::new(), + }; + let decoded: ExternalBatchGetLocalProbeResp = + bitcode::decode(&bitcode::encode(&response)).expect("decode local probe response"); + assert_eq!(decoded.items.len(), 2); + assert_eq!( + decoded.items[0].as_ref().map(|info| info.holder_id), + Some(23) + ); + assert!(decoded.items[1].is_none()); + } +} +impl MsgPackSerializePart for ExternalBatchGetResp { + fn msg_id(&self) -> u32 { + 4021 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetStartReq { + pub keys: Vec, + pub req_node_id: String, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, + pub prefix_best_effort: bool, + pub atomic_group_lens: Option>, + pub transfer_concurrency: usize, +} +impl MsgPackSerializePart for ExternalBatchGetStartReq { + fn msg_id(&self) -> u32 { + 4030 + } +} +impl RPCReq for ExternalBatchGetStartReq { + type Resp = ExternalBatchGetStartResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetStartResp { + pub error_code: ErrorCode, + pub error_json: String, + pub handle: u64, + pub raw_prefix_hit_len: usize, + pub transfer_plan: ExternalBatchGetStartTransferPlan, +} +impl MsgPackSerializePart for ExternalBatchGetStartResp { + fn msg_id(&self) -> u32 { + 4031 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub enum ExternalBatchGetStartTransferPlan { + #[default] + OwnerRpc, + InlineLocal { + items: Vec, + }, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetTransferReq { + pub handle: u64, + pub req_node_id: String, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, + /// Complete atomic-group prefix selected by the consumer from the live + /// get-start handle. + pub consume_prefix_len: usize, +} +impl MsgPackSerializePart for ExternalBatchGetTransferReq { + fn msg_id(&self) -> u32 { + 4032 + } +} +impl RPCReq for ExternalBatchGetTransferReq { + type Resp = ExternalBatchGetTransferResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetTransferResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchGetTransferResp { + fn msg_id(&self) -> u32 { + 4033 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetCancelReq { + pub handle: u64, + pub req_node_id: String, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, + pub transfer_plan: ExternalBatchGetCancelPlan, +} +impl MsgPackSerializePart for ExternalBatchGetCancelReq { + fn msg_id(&self) -> u32 { + 4034 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub enum ExternalBatchGetCancelPlan { + #[default] + OwnerRpc, + InlineLocal { + holder_ids: Vec, + }, +} +impl RPCReq for ExternalBatchGetCancelReq { + type Resp = ExternalBatchGetCancelResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchGetCancelResp { + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalPlannedGetItem { + pub key: String, + pub get_id: u64, +} + +/// Execute a target-free master plan into the requester's owner-local CPU +/// pool. The stable `(requester generation, plan_handle)` identity makes a +/// lost response replayable without another lookup or transfer. +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalExecutePlannedGetReq { + pub plan_handle: u64, + pub items: Vec, + pub req_node_id: String, + pub started_time: i64, + pub transfer_concurrency: usize, +} +impl MsgPackSerializePart for ExternalExecutePlannedGetReq { + fn msg_id(&self) -> u32 { + 4038 + } +} +impl RPCReq for ExternalExecutePlannedGetReq { + type Resp = ExternalExecutePlannedGetResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalExecutePlannedGetResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[cfg(test)] +mod external_execute_planned_get_wire_tests { + use super::{ + ExternalBatchGetItemResp, ExternalExecutePlannedGetReq, ExternalExecutePlannedGetResp, + ExternalPlannedGetItem, + }; + use crate::memholder::ExternalMemHolderInfo; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + + #[test] + fn operation_identity_and_holder_terminal_round_trip() { + let request = ExternalExecutePlannedGetReq { + plan_handle: 19, + items: vec![ExternalPlannedGetItem { + key: "page-a".to_string(), + get_id: 0, + }], + req_node_id: "external-a".to_string(), + started_time: 23, + transfer_concurrency: 32, + }; + let decoded: ExternalExecutePlannedGetReq = + bitcode::decode(&bitcode::encode(&request)).expect("decode planned execute request"); + assert_eq!(decoded.plan_handle, 19); + assert_eq!(decoded.items[0].get_id, 0); + assert_eq!(decoded.started_time, 23); + + let terminal = ExternalExecutePlannedGetResp { + items: vec![ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some(ExternalMemHolderInfo { + offset: 4096, + len: 8192, + holder_id: 29, + }), + }], + error_code: OK, + error_json: String::new(), + }; + let decoded: ExternalExecutePlannedGetResp = + bitcode::decode(&bitcode::encode(&terminal)).expect("decode planned execute terminal"); + assert_eq!( + decoded.items[0] + .external_memholder_info + .as_ref() + .map(|info| info.holder_id), + Some(29) + ); + } +} +impl MsgPackSerializePart for ExternalExecutePlannedGetResp { + fn msg_id(&self) -> u32 { + 4039 + } +} +impl MsgPackSerializePart for ExternalBatchGetCancelResp { + fn msg_id(&self) -> u32 { + 4035 + } +} + +#[cfg(test)] +mod external_batch_get_plan_wire_tests { + use super::{ + ExternalBatchGetCancelPlan, ExternalBatchGetCancelReq, ExternalBatchGetItemResp, + ExternalBatchGetStartResp, ExternalBatchGetStartTransferPlan, ExternalBatchGetTransferReq, + }; + use crate::memholder::ExternalMemHolderInfo; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + + #[test] + fn inline_start_and_cancel_plans_round_trip() { + let start = ExternalBatchGetStartResp { + error_code: OK, + error_json: String::new(), + handle: 19, + raw_prefix_hit_len: 2, + transfer_plan: ExternalBatchGetStartTransferPlan::InlineLocal { + items: vec![ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some(ExternalMemHolderInfo { + offset: 4096, + len: 8192, + holder_id: 23, + }), + }], + }, + }; + let decoded_start: ExternalBatchGetStartResp = + bitcode::decode(&bitcode::encode(&start)).expect("decode inline start plan"); + assert_eq!(decoded_start.handle, 19); + assert!(matches!( + decoded_start.transfer_plan, + ExternalBatchGetStartTransferPlan::InlineLocal { items } + if items.len() == 1 + && items[0] + .external_memholder_info + .as_ref() + .is_some_and(|info| info.holder_id == 23) + )); + + let transfer = ExternalBatchGetTransferReq { + handle: 19, + req_node_id: "external-a".to_string(), + started_time: 29, + consume_prefix_len: 1, + }; + let decoded_transfer: ExternalBatchGetTransferReq = + bitcode::decode(&bitcode::encode(&transfer)).expect("decode transfer request"); + assert_eq!(decoded_transfer.consume_prefix_len, 1); + + let cancel = ExternalBatchGetCancelReq { + handle: 19, + req_node_id: "external-a".to_string(), + started_time: 29, + transfer_plan: ExternalBatchGetCancelPlan::InlineLocal { + holder_ids: vec![23, 24], + }, + }; + let decoded_cancel: ExternalBatchGetCancelReq = + bitcode::decode(&bitcode::encode(&cancel)).expect("decode inline cancel plan"); + assert!(matches!( + decoded_cancel.transfer_plan, + ExternalBatchGetCancelPlan::InlineLocal { holder_ids } + if holder_ids == vec![23, 24] + )); + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutStartItemReq { + pub key: String, + pub len: u64, + pub reject_if_inflight_same_key: bool, + pub reject_if_exist_same_key: bool, + pub make_replica_task: bool, + pub preferred_sub_cluster: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutStartReq { + pub items: Vec, + /// Positive lengths partitioning `items`; omitted means one group per item. + pub atomic_group_lens: Option>, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalBatchPutStartReq { + fn msg_id(&self) -> u32 { + 4022 + } +} +impl RPCReq for ExternalBatchPutStartReq { + type Resp = ExternalBatchPutStartResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutStartItemResp { + pub error_code: ErrorCode, + pub src_offset: u64, + pub target_offset: u64, + pub transfer_target_offset: Option, + pub peer_id: Option, + pub src_base_addr: u64, + pub target_base_addr: u64, + pub error_json: String, + pub put_id: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutStartResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchPutStartResp { + fn msg_id(&self) -> u32 { + 4023 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutTransferEndItemReq { + pub key: String, + pub len: u64, + pub src_offset: u64, + pub target_offset: u64, + pub peer_id: Option, + pub target_base_addr: Option, + pub put_id: Option, + pub lease_id: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutTransferEndReq { + pub items: Vec, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, + pub transfer_concurrency: usize, +} +impl MsgPackSerializePart for ExternalBatchPutTransferEndReq { + fn msg_id(&self) -> u32 { + 4024 + } +} +impl RPCReq for ExternalBatchPutTransferEndReq { + type Resp = ExternalBatchPutTransferEndResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutTransferEndItemResp { + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutTransferEndResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchPutTransferEndResp { + fn msg_id(&self) -> u32 { + 4025 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutCommitItemReq { + pub key: String, + pub len: u64, + pub src_offset: u64, + pub remote_target: bool, + pub put_id: Option, + pub lease_id: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutCommitReq { + pub items: Vec, + /// Owner node_start_time observed by the caller when request starts + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalBatchPutCommitReq { + fn msg_id(&self) -> u32 { + 4026 + } +} +impl RPCReq for ExternalBatchPutCommitReq { + type Resp = ExternalBatchPutCommitResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutCommitItemResp { + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchPutCommitResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchPutCommitResp { + fn msg_id(&self) -> u32 { + 4027 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalIoLocalitySnapshot { + pub op_count: u64, + pub bytes: u64, + pub transfer_us: u64, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalObservabilitySnapshotReq { + /// Owner node_start_time observed by external when request starts. + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalObservabilitySnapshotReq { + fn msg_id(&self) -> u32 { + 4028 + } +} +impl RPCReq for ExternalObservabilitySnapshotReq { + type Resp = ExternalObservabilitySnapshotResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalObservabilitySnapshotResp { + pub error_code: ErrorCode, + pub error_json: String, + pub l2_local_hit_pages: u64, + pub l2_local_hit_bytes: u64, + pub l2_remote_hit_pages: u64, + pub l2_remote_hit_bytes: u64, + pub put_local: ExternalIoLocalitySnapshot, + pub put_remote: ExternalIoLocalitySnapshot, + pub get_local: ExternalIoLocalitySnapshot, + pub get_remote: ExternalIoLocalitySnapshot, +} +impl ExternalObservabilitySnapshotResp { + pub fn success(snapshot: crate::metrics::KvLocalitySnapshot) -> Self { + Self { + error_code: OK, + error_json: String::new(), + l2_local_hit_pages: snapshot.l2_local_hit_pages, + l2_local_hit_bytes: snapshot.l2_local_hit_bytes, + l2_remote_hit_pages: snapshot.l2_remote_hit_pages, + l2_remote_hit_bytes: snapshot.l2_remote_hit_bytes, + put_local: ExternalIoLocalitySnapshot { + op_count: snapshot.put_local.op_count, + bytes: snapshot.put_local.bytes, + transfer_us: snapshot.put_local.transfer_us, + }, + put_remote: ExternalIoLocalitySnapshot { + op_count: snapshot.put_remote.op_count, + bytes: snapshot.put_remote.bytes, + transfer_us: snapshot.put_remote.transfer_us, + }, + get_local: ExternalIoLocalitySnapshot { + op_count: snapshot.get_local.op_count, + bytes: snapshot.get_local.bytes, + transfer_us: snapshot.get_local.transfer_us, + }, + get_remote: ExternalIoLocalitySnapshot { + op_count: snapshot.get_remote.op_count, + bytes: snapshot.get_remote.bytes, + transfer_us: snapshot.get_remote.transfer_us, + }, + } + } + + pub fn into_snapshot(self) -> crate::metrics::KvLocalitySnapshot { + crate::metrics::KvLocalitySnapshot { + l2_local_hit_pages: self.l2_local_hit_pages, + l2_local_hit_bytes: self.l2_local_hit_bytes, + l2_remote_hit_pages: self.l2_remote_hit_pages, + l2_remote_hit_bytes: self.l2_remote_hit_bytes, + put_local: crate::metrics::KvIoLocalitySnapshot { + op_count: self.put_local.op_count, + bytes: self.put_local.bytes, + transfer_us: self.put_local.transfer_us, + }, + put_remote: crate::metrics::KvIoLocalitySnapshot { + op_count: self.put_remote.op_count, + bytes: self.put_remote.bytes, + transfer_us: self.put_remote.transfer_us, + }, + get_local: crate::metrics::KvIoLocalitySnapshot { + op_count: self.get_local.op_count, + bytes: self.get_local.bytes, + transfer_us: self.get_local.transfer_us, + }, + get_remote: crate::metrics::KvIoLocalitySnapshot { + op_count: self.get_remote.op_count, + bytes: self.get_remote.bytes, + transfer_us: self.get_remote.transfer_us, + }, + } + } +} +impl MsgPackSerializePart for ExternalObservabilitySnapshotResp { + fn msg_id(&self) -> u32 { + 4029 + } +} + // #[derive(Default, Debug, Clone, Encode, Decode)] // pub struct ExternalPutReq { // pub key: String, @@ -114,6 +765,8 @@ pub struct ExternalPutStartReq { pub key: String, pub len: u64, pub reject_if_inflight_same_key: bool, + pub reject_if_exist_same_key: bool, + pub make_replica_task: bool, /// Prefer placing the target allocation on any kvclient within this sub_cluster. pub preferred_sub_cluster: Option, /// Owner node_start_time observed by external when request starts @@ -190,6 +843,9 @@ impl MsgPackSerializePart for ExternalPutTransferEndResp { #[derive(Default, Debug, Clone, Encode, Decode)] pub struct ExternalPutCommitReq { pub key: String, + pub len: u64, + pub src_offset: u64, + pub remote_target: bool, pub put_id: Option, pub lease_id: Option, /// Owner node_start_time observed by the caller when request starts @@ -297,6 +953,34 @@ impl MsgPackSerializePart for ExternalIsExistResp { } } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchIsExistReq { + pub keys: Vec, + pub allow_local_snapshot: bool, + /// Owner node_start_time observed by external when request starts + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalBatchIsExistReq { + fn msg_id(&self) -> u32 { + crate::rpcresp_kvresult_convert::msg_and_error::MsgId::ExternalBatchIsExistReq as u32 + } +} +impl RPCReq for ExternalBatchIsExistReq { + type Resp = ExternalBatchIsExistResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchIsExistResp { + pub error_code: ErrorCode, + pub exists_list: Vec, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchIsExistResp { + fn msg_id(&self) -> u32 { + crate::rpcresp_kvresult_convert::msg_and_error::MsgId::ExternalBatchIsExistResp as u32 + } +} + #[derive(Default, Debug, Clone, Encode, Decode)] pub struct ExternalDeleteAckReq { pub key: String, @@ -325,11 +1009,66 @@ impl MsgPackSerializePart for ExternalDeleteAckResp { } } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchDeleteAckReq { + pub external_client_id: String, + pub holder_ids: Vec, + /// Owner node_start_time observed when these holders were created. + pub started_time: i64, +} +impl MsgPackSerializePart for ExternalBatchDeleteAckReq { + fn msg_id(&self) -> u32 { + 4036 + } +} +impl RPCReq for ExternalBatchDeleteAckReq { + type Resp = ExternalBatchDeleteAckResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalBatchDeleteAckResp { + pub released_count: u32, + pub missing_count: u32, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for ExternalBatchDeleteAckResp { + fn msg_id(&self) -> u32 { + 4037 + } +} + +#[cfg(test)] +mod external_batch_delete_ack_wire_tests { + use super::ExternalBatchDeleteAckReq; + + #[test] + fn compact_holder_id_batch_round_trips() { + let request = ExternalBatchDeleteAckReq { + external_client_id: "external-a".to_string(), + holder_ids: vec![3, 5, 8], + started_time: 17, + }; + let decoded: ExternalBatchDeleteAckReq = + bitcode::decode(&bitcode::encode(&request)).expect("decode holder ACK batch"); + assert_eq!(decoded.external_client_id, "external-a"); + assert_eq!(decoded.holder_ids, vec![3, 5, 8]); + assert_eq!(decoded.started_time, 17); + } +} + // --- RPC: Owner -> External to invalidate weak-index cache for keys --- +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ExternalInvalidateWeakIndexItem { + pub key: String, +} + #[derive(Default, Debug, Clone, Encode, Decode)] pub struct ExternalInvalidateWeakIndexReq { - /// Keys whose weak cache entries should be invalidated on external client + /// Keys whose weak cache entries should be invalidated on external client. + /// Kept for compatibility with older senders; new senders should use `items`. pub keys: Vec, + pub items: Vec, } impl MsgPackSerializePart for ExternalInvalidateWeakIndexReq { fn msg_id(&self) -> u32 { diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs index 374ed71..11a5f5a 100644 --- a/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/put.rs @@ -1,4 +1,12 @@ -use super::ClientKvApiInner; +use super::{ + ClientKvApiInner, OwnerHotEvictionEvent, OwnerHotEvictionPreparation, + OwnerHotSelectionFenceOutcome, OwnerLocalReserveClassState, OwnerLocalReserveGrantState, + OwnerLocalReservePoolState, OwnerLocalReserveSlotLease, OwnerLocalReserveSlotRef, + OwnerLocalSsdPutOutcome, OwnerLocalSsdPutReservation, OwnerLocalSsdPutSharedOp, + OwnerRemotePutOutcome, OwnerRemotePutReservation, OwnerRemotePutSharedOp, + local_reserve_rebalance::{owner_local_reserve_timeout_config, wait_owner_local_reserve_ready}, +}; +use crate::OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES; use crate::cluster_manager::NodeIDString; use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::master_kv_router::put::PutIDForAKey; @@ -6,591 +14,4642 @@ use crate::master_kv_router::put::PutIDForAKey; use crate::memholder::kvclient_encode::{calc_flat_dict_encoded_len, write_flat_dict_ptrs_to_ptr}; use crate::observe_kvope::{ obe_put_start_error_rpc, obe_put_start_error_status, obe_put_start_success, - obe_put_transfer_error, }; use crate::{ - master_kv_router::msg_pack::{PutDoneReq, PutRevokeReq, PutStartReq, PutStartResp}, + client_kv_api::ClientKvApiView, + master_kv_router::msg_pack::{ + BatchEnqueueReplicaTaskReq, BatchEnqueueReplicaTaskResp, BatchEvictOwnerSourceReq, + BatchEvictOwnerSourceResp, BatchPreparePutKeyItemReq, BatchPreparePutKeysReq, + BatchPreparePutKeysResp, BatchPublishOwnerSsdReq, BatchPublishOwnerSsdResp, + BatchPutAppendDoneItemReq, BatchPutAppendDoneReq, BatchPutAppendDoneResp, + BatchPutAppendStartItemReq, BatchPutAppendStartReq, BatchPutAppendStartResp, + BatchPutDoneItemReq, BatchPutDoneItemResp, BatchPutDoneReq, BatchPutDoneResp, + BatchPutRevokeItemReq, BatchPutRevokeReq, BatchPutRevokeResp, BatchPutStartItemReq, + BatchPutStartReq, BatchPutStartResp, BatchReleasePutKeyReservationsReq, + BatchReleasePutKeyReservationsResp, EnqueueReplicaTaskItemResp, GroupedBatchPutDoneItemReq, + GroupedBatchPutDoneReq, GroupedBatchPutDoneResp, OwnerReclaimBacking, + OwnerSourceEvictionOutcome, OwnerSourceEvictionVictim, OwnerSourceEvictionVictimResp, + OwnerSourceSsdPolicy, OwnerSsdPublishItem, OwnerSsdPublishOutcome, PutAppendDoneReq, + PutAppendDoneResp, PutAppendRevokeReq, PutAppendStartOutcome, PutAppendStartReq, + PutAppendStartResp, PutAtomicGroup, PutDoneCommittedSlot, PutDoneReq, PutRevokeReq, + PutStartReq, PutStartResp, ReleaseLocalGrantReq, ReserveLocalGrantOutcome, + ReserveLocalGrantReq, owner_source_eviction_epoch, + }, + memholder::{UserMemHolder, UserMemHolderExposeKind}, p2p::msg_pack::MsgPack, + p2p::p2p_module::RpcTransportPolicy, rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}, }; use chrono::Utc; use fluxon_commu::TransferBreakdown; -use std::time::Instant; +use limit_thirdparty::tokio; +use std::future::Future; +use std::sync::{Arc, atomic::Ordering}; +use std::time::{Duration, Instant}; use tracing::info; fn duration_to_i64_us(duration: std::time::Duration) -> i64 { duration.as_micros().min(i64::MAX as u128) as i64 } -#[derive(Debug, Clone, Copy, Default)] -pub struct PutEndStats { - pub master_put_end_rpc_us: i64, - pub master_put_end_server_us: i64, +fn owner_local_reserve_slot_size(value_len: u64) -> KvResult { + crate::owner_local_reserve_slot_size_bytes(value_len).ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "value_len={} cannot be represented by a resident local-reserve slot no larger than {} bytes", + value_len, OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES + ), + }) + }) } -impl ClientKvApiInner { - async fn put_common( - &self, - key: &str, - payload_len: u64, - len_for_start: u32, - reject_if_inflight_same_key: bool, - preferred_sub_cluster: Option<&str>, - lease_id: Option, - _test_payload_len_u32: u32, - _test_remove_after_fill: bool, - fill_abs_src: F, - dbg_addr_summary: bool, - info_complete_tag: Option<&'static str>, - ) -> KvResult<()> - where - F: FnOnce(u64), - { - let client_id = self.client_id_str(); - let node_role = self.node_role(); - let metrics = self.metrics_handle(); +fn owner_local_reserve_slots_per_grant(slot_size: u64) -> u32 { + crate::owner_local_reserve_slots_per_grant(slot_size) + .expect("validated local-reserve slot size must fit in a grant") +} - let t1 = Utc::now().timestamp_micros(); - let (resp, _rpc_latency) = { - match self - .put_start( - key, - len_for_start, - reject_if_inflight_same_key, - preferred_sub_cluster, - ) - .await - { - Ok(resp) => resp, - Err(err) => { - obe_put_start_error_rpc(&metrics, &client_id, &node_role, key, payload_len); - return Err(err); - } - } - }; - let t2 = Utc::now().timestamp_micros(); - if let Err(e) = - crate::rpcresp_kvresult_convert::try_from_code(resp.error_code, resp.error_json.clone()) - { - obe_put_start_error_status(&metrics, &client_id, &node_role, key, payload_len); - return Err(e); +fn owner_local_reserve_install_grant( + pool: &mut OwnerLocalReservePoolState, + slot_size: u64, + slots_per_grant: u32, + grant: OwnerLocalReserveGrantState, +) { + let class_state = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)); + assert!( + class_state.slot_size == slot_size, + "slot_size drift detected while installing local reserve grant" + ); + assert!( + class_state.slots_per_grant == slots_per_grant, + "slots_per_grant drift detected while installing local reserve grant" + ); + class_state.install_grant(grant); +} + +fn owner_local_reserve_try_claim( + pool: &mut OwnerLocalReservePoolState, + slot_size: u64, + slots_per_grant: u32, + value_len: u64, + key_count: usize, +) -> Option { + let free_slots = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)) + .free_slot_count(); + if free_slots < key_count { + return None; + } + let slots = owner_local_reserve_claim_available(pool, slot_size, slots_per_grant, key_count); + assert_eq!( + slots.len(), + key_count, + "free_slot_count check and claim path diverged" + ); + Some(OwnerLocalReserveSlotLease { + value_len, + slot_size, + slots, + }) +} + +fn owner_local_reserve_claim_available( + pool: &mut OwnerLocalReservePoolState, + slot_size: u64, + slots_per_grant: u32, + max_slots: usize, +) -> Vec { + let class_state = pool + .classes + .entry(slot_size) + .or_insert_with(|| OwnerLocalReserveClassState::new(slot_size, slots_per_grant)); + let claim_count = class_state.free_slot_count().min(max_slots); + let slots = class_state.claim_available(max_slots); + assert_eq!( + slots.len(), + claim_count, + "free_slot_count check and claim path diverged" + ); + slots +} + +struct OwnerLocalReservePendingDemandGuard<'a> { + inner: &'a ClientKvApiInner, + slot_size: u64, + slots_per_grant: u32, + pending_slots: usize, +} + +impl<'a> OwnerLocalReservePendingDemandGuard<'a> { + fn new( + inner: &'a ClientKvApiInner, + slot_size: u64, + slots_per_grant: u32, + pending_slots: usize, + ) -> Self { + inner.owner_local_reserve_register_pending_demand( + slot_size, + slots_per_grant, + pending_slots, + ); + Self { + inner, + slot_size, + slots_per_grant, + pending_slots, } - obe_put_start_success(&metrics, &client_id, &node_role, key, t1, t2); + } - let put_id = resp.put_id; - let peer_id = if &*resp.node_id == &*self.view.cluster_manager().get_self_info().id { - None - } else { - Some(resp.node_id.clone()) - }; - let abs_src = resp.src_addr; - let abs_target = resp.target_addr; + fn consume(&mut self) { + if self.pending_slots == 0 { + return; + } + self.inner.owner_local_reserve_consume_pending_demand( + self.slot_size, + self.slots_per_grant, + self.pending_slots, + ); + self.pending_slots = 0; + self.inner + .owner_local_reserve_rebalance_notify() + .notify_waiters(); + } - #[cfg(test)] - { - self.test_record.add_transfering_put( - key.to_string(), - _test_payload_len_u32, - put_id.0, - put_id.1, - resp.node_id.to_string(), - format!("{:#x}", resp.target_addr), + fn disarm_after_locked_consume(&mut self) { + assert!( + self.pending_slots > 0, + "pending-demand guard was already consumed" + ); + self.pending_slots = 0; + self.inner + .owner_local_reserve_rebalance_notify() + .notify_waiters(); + } +} + +impl Drop for OwnerLocalReservePendingDemandGuard<'_> { + fn drop(&mut self) { + self.consume(); + } +} + +#[cfg(test)] +mod local_reserve_claim_tests { + use super::{ + OwnerLocalReserveGrantState, OwnerLocalReservePoolState, + owner_local_reserve_claim_available, owner_local_reserve_install_grant, + owner_local_reserve_slot_size, owner_local_reserve_slots_per_grant, + }; + use crate::client_kv_api::{ClientKvApi, ClientKvApiNewArg}; + use crate::config::TestSpecConfig; + use std::sync::Arc; + use std::time::Duration; + + #[test] + fn slot_size_uses_page_aligned_exact_fit() { + const SGLANG_KV_PAGE_BYTES: u64 = 4_718_592; + + let slot_size = owner_local_reserve_slot_size(SGLANG_KV_PAGE_BYTES).unwrap(); + assert_eq!(slot_size, SGLANG_KV_PAGE_BYTES); + assert_eq!(owner_local_reserve_slots_per_grant(slot_size), 113); + assert_eq!(owner_local_reserve_slot_size(4_097).unwrap(), 8_192); + } + + #[test] + fn partial_claim_keeps_progress_for_large_waiters() { + let mut pool = OwnerLocalReservePoolState::default(); + owner_local_reserve_install_grant( + &mut pool, + 8, + 4, + OwnerLocalReserveGrantState::new(1, 1000, 1000, 32, 8, 4), + ); + + let first = owner_local_reserve_claim_available(&mut pool, 8, 4, 3); + assert_eq!(first.len(), 3); + assert_eq!(pool.classes.get(&8).unwrap().free_slot_count(), 1); + + let second = owner_local_reserve_claim_available(&mut pool, 8, 4, 3); + assert_eq!(second.len(), 1); + assert_eq!(pool.classes.get(&8).unwrap().free_slot_count(), 0); + } + + #[test] + fn grant_index_and_cached_counters_survive_swap_remove_and_reinstall() { + const SLOT_SIZE: u64 = 8; + const SLOTS_PER_GRANT: u32 = 4; + let mut pool = OwnerLocalReservePoolState::default(); + for grant_id in 1..=3 { + owner_local_reserve_install_grant( + &mut pool, + SLOT_SIZE, + SLOTS_PER_GRANT, + OwnerLocalReserveGrantState::new( + grant_id, + grant_id * 1000, + grant_id * 1000, + SLOT_SIZE * u64::from(SLOTS_PER_GRANT), + SLOT_SIZE, + SLOTS_PER_GRANT, + ), ); } - if self.short_circuit_put_payload_path_enabled() { - #[cfg(test)] - { - if _test_remove_after_fill { - self.test_record - .remove_transfering_put(key.to_string(), put_id); - } - } + let class = pool.classes.get_mut(&SLOT_SIZE).unwrap(); + assert_eq!(class.free_slot_count(), 12); + assert_eq!(class.used_slot_count(), 0); + let detached = class + .detach_fully_free_grant(2) + .expect("middle grant must be indexed"); + assert_eq!(class.free_slot_count(), 8); + assert_eq!(class.grant_count(), 2); - let skipped_breakdown = if peer_id.is_none() && abs_src == abs_target { - TransferBreakdown { - local_noop: true, - ..TransferBreakdown::default() - } - } else { - TransferBreakdown::default() - }; - metrics.pending_put_set_transfer_breakdown( - put_id, - skipped_breakdown.submit_blocking_us, - skipped_breakdown.create_xfer_req_us, - skipped_breakdown.post_xfer_req_us, - skipped_breakdown.poll_wait_us, - skipped_breakdown.poll_iters, - skipped_breakdown.used_fast_path, - skipped_breakdown.local_noop, - skipped_breakdown.remote_transfer, + // Removing the middle Vec entry swap-moves grant 3. A subsequent state + // transition by grant id proves that its repaired index is authoritative. + let claimed = class.claim_available(1); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].grant_id, 3); + assert_eq!(class.free_slot_count(), 7); + assert_eq!(class.prepared_slot_count(), 1); + assert!(class.release_prepared_slot(3, claimed[0].slot_index)); + assert_eq!(class.free_slot_count(), 8); + assert_eq!(class.prepared_slot_count(), 0); + + class.install_grant(detached); + assert_eq!(class.free_slot_count(), 12); + assert_eq!(class.used_slot_count(), 0); + assert_eq!(class.grant_count(), 3); + } + + #[test] + fn committed_slots_are_reclaimed_and_reused_independently_across_grants() { + const SLOT_SIZE: u64 = 8; + const SLOTS_PER_GRANT: u32 = 4; + + let mut pool = OwnerLocalReservePoolState::default(); + for (grant_id, addr) in [(1, 1000), (2, 2000)] { + owner_local_reserve_install_grant( + &mut pool, + SLOT_SIZE, + SLOTS_PER_GRANT, + OwnerLocalReserveGrantState::new( + grant_id, + addr, + addr, + SLOT_SIZE * u64::from(SLOTS_PER_GRANT), + SLOT_SIZE, + SLOTS_PER_GRANT, + ), ); - self.put_end(key, put_id, lease_id).await?; - if let Some(tag) = info_complete_tag { - info!("{tag} complete key={} bytes={}", key, payload_len); - } - return Ok(()); } - fill_abs_src(abs_src); + let initial = owner_local_reserve_claim_available(&mut pool, SLOT_SIZE, SLOTS_PER_GRANT, 8); + assert_eq!(initial.len(), 8); + { + let class = pool.classes.get_mut(&SLOT_SIZE).unwrap(); + for slot in &initial { + assert!(class.mark_prepared_slot_pending_visible(slot.grant_id, slot.slot_index)); + assert!(class.retain_resident_slot_holder(slot.grant_id, slot.slot_index)); + assert!( + class.promote_pending_visible_slot_to_committed(slot.grant_id, slot.slot_index) + ); + } + } - #[cfg(test)] + let victims = [initial[1].clone(), initial[5].clone()]; { - if _test_remove_after_fill { - self.test_record - .remove_transfering_put(key.to_string(), put_id); + let class = pool.classes.get_mut(&SLOT_SIZE).unwrap(); + for victim in &victims { + assert!(class.release_committed_slot_route(victim.grant_id, victim.slot_index)); + assert!(class.release_resident_slot_holder(victim.grant_id, victim.slot_index)); } + assert_eq!(class.free_slot_count(), 2); + assert_eq!(class.grant_count(), 2); + assert!(class.grants.iter().all(|grant| !grant.is_fully_free())); + assert!(class.grants.iter().all(|grant| grant.free_slots.len() == 1)); } - let base_addr = self - .view - .client_seg_pool() - .cpu_mem_read_guard() - .await - .unwrap() - .allocated_addr; - let src_offset = abs_src - base_addr; - let (target_offset, target_base_addr_opt) = match &peer_id { - Some(_) => ( - abs_target - resp.target_base_addr, - Some(resp.target_base_addr), - ), - None => (abs_target - base_addr, None), - }; - if dbg_addr_summary { - tracing::debug!( - "put path addr summary: key={}, put_id=({},{}) local_base={:#x}, abs_src={:#x}, src_off={:#x}, master_target_base={:#x}, abs_target={:#x}, tgt_off={:#x}, peer_id={:?}", - key, - put_id.0, - put_id.1, - base_addr, - abs_src, - src_offset, - target_base_addr_opt.unwrap_or(base_addr), - abs_target, - target_offset, - peer_id - ); + for cycle in 0..10 { + let reused = + owner_local_reserve_claim_available(&mut pool, SLOT_SIZE, SLOTS_PER_GRANT, 2); + assert_eq!(reused.len(), 2, "cycle {cycle} did not reuse both slots"); + let mut reused_grant_ids = reused.iter().map(|slot| slot.grant_id).collect::>(); + reused_grant_ids.sort_unstable(); + assert_eq!(reused_grant_ids, vec![1, 2]); + + let class = pool.classes.get_mut(&SLOT_SIZE).unwrap(); + for slot in reused { + assert!(class.mark_prepared_slot_pending_visible(slot.grant_id, slot.slot_index)); + assert!(class.retain_resident_slot_holder(slot.grant_id, slot.slot_index)); + assert!( + class.promote_pending_visible_slot_to_committed(slot.grant_id, slot.slot_index) + ); + assert!(class.release_committed_slot_route(slot.grant_id, slot.slot_index)); + assert!(class.release_resident_slot_holder(slot.grant_id, slot.slot_index)); + } + assert_eq!(class.free_slot_count(), 2); + assert_eq!(class.grant_count(), 2); + assert!(class.grants.iter().all(|grant| !grant.is_fully_free())); } + } - let transfer_breakdown = match self - .put_transfer( - key, - put_id, - src_offset, - target_offset, - payload_len, - peer_id.clone(), - target_base_addr_opt, - ) + #[limit_thirdparty::tokio::test] + async fn later_waiter_cannot_steal_slot_before_current_claim_turn_completes() { + const SLOT_SIZE: u64 = 4 * 1024; + const SLOTS_PER_GRANT: u32 = 4; + + let api = Arc::new( + ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) .await + .expect("construct test ClientKvApi"), + ); { - Ok(breakdown) => breakdown, - Err(e) => { - self.put_revoke(key, put_id).await?; - obe_put_transfer_error(&metrics, &client_id, &node_role, key, payload_len); - return Err(e); - } + let mut pool = api.inner().owner_local_reserve_pool.lock(); + owner_local_reserve_install_grant( + &mut pool, + SLOT_SIZE, + SLOTS_PER_GRANT, + OwnerLocalReserveGrantState::new( + 1, + 1000, + 1000, + SLOT_SIZE * u64::from(SLOTS_PER_GRANT), + SLOT_SIZE, + SLOTS_PER_GRANT, + ), + ); + } + + // Model a five-slot waiter that has made partial progress while owning the claim turn. + let claim_lock = api.inner().owner_local_reserve_claim_lock(SLOT_SIZE); + let current_claim_turn = claim_lock.lock().await; + let first_partial = { + let mut pool = api.inner().owner_local_reserve_pool.lock(); + owner_local_reserve_claim_available(&mut pool, SLOT_SIZE, SLOTS_PER_GRANT, 3) }; - metrics.pending_put_set_transfer_breakdown( - put_id, - transfer_breakdown.submit_blocking_us, - transfer_breakdown.create_xfer_req_us, - transfer_breakdown.post_xfer_req_us, - transfer_breakdown.poll_wait_us, - transfer_breakdown.poll_iters, - transfer_breakdown.used_fast_path, - transfer_breakdown.local_noop, - transfer_breakdown.remote_transfer, + assert_eq!(first_partial.len(), 3); + + // A later one-slot request must queue even though one slot is currently free. + let later_api = Arc::clone(&api); + let mut later_waiter = tokio::spawn(async move { + later_api + .inner() + .owner_claim_local_reserve_slot_lease(SLOT_SIZE, 1) + .await + }); + assert!( + limit_thirdparty::tokio::time::timeout(Duration::from_millis(25), &mut later_waiter,) + .await + .is_err(), + "later waiter bypassed the active claim turn" + ); + assert_eq!( + api.inner() + .owner_local_reserve_pool + .lock() + .classes + .get(&SLOT_SIZE) + .unwrap() + .free_slot_count(), + 1, + "later waiter stole the current waiter's remaining free slot" ); - if self.skip_put_end_commit_enabled() { - let _ = metrics.pending_put_remove(&put_id); - tracing::warn!( - "skip_put_end_commit test-only fast-path: returning success without put_end; key={} put_id=({},{}) payload_len={}", - key, - put_id.0, - put_id.1, - payload_len + // Refill lets the current waiter reach all five slots before handing off the turn. + let first_remainder = { + let mut pool = api.inner().owner_local_reserve_pool.lock(); + owner_local_reserve_install_grant( + &mut pool, + SLOT_SIZE, + SLOTS_PER_GRANT, + OwnerLocalReserveGrantState::new( + 2, + 2000, + 2000, + SLOT_SIZE * u64::from(SLOTS_PER_GRANT), + SLOT_SIZE, + SLOTS_PER_GRANT, + ), ); - if let Some(tag) = info_complete_tag { - info!( - "{tag} complete_without_put_end key={} bytes={}", - key, payload_len - ); - } - return Ok(()); + owner_local_reserve_claim_available(&mut pool, SLOT_SIZE, SLOTS_PER_GRANT, 2) + }; + assert_eq!(first_partial.len() + first_remainder.len(), 5); + drop(current_claim_turn); + + let later_lease = + limit_thirdparty::tokio::time::timeout(Duration::from_secs(1), later_waiter) + .await + .expect("later waiter did not receive the next claim turn") + .expect("later waiter task panicked") + .expect("later waiter failed to claim a free slot"); + assert_eq!(later_lease.slots.len(), 1); + } + + #[limit_thirdparty::tokio::test] + async fn same_class_claim_waiters_complete_in_fifo_order() { + const SLOT_SIZE: u64 = 4 * 1024; + + let api = Arc::new( + ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"), + ); + let claim_lock = api.inner().owner_local_reserve_claim_lock(SLOT_SIZE); + let claim_turn = claim_lock.lock().await; + let order = Arc::new(std::sync::Mutex::new(Vec::new())); + + let first_queued = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let first_lock = api.inner().owner_local_reserve_claim_lock(SLOT_SIZE); + let first_order = Arc::clone(&order); + let first_queued_task = Arc::clone(&first_queued); + let first = tokio::spawn(async move { + first_queued_task.store(true, std::sync::atomic::Ordering::Release); + let _turn = first_lock.lock_owned().await; + first_order.lock().unwrap().push(1u8); + }); + while !first_queued.load(std::sync::atomic::Ordering::Acquire) { + tokio::task::yield_now().await; } - self.put_end(key, put_id, lease_id).await?; - self.cache_metadata_only_after_put(key, put_id); - if let Some(tag) = info_complete_tag { - info!("{tag} complete key={} bytes={}", key, payload_len); + let second_queued = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let second_lock = api.inner().owner_local_reserve_claim_lock(SLOT_SIZE); + let second_order = Arc::clone(&order); + let second_queued_task = Arc::clone(&second_queued); + let second = tokio::spawn(async move { + second_queued_task.store(true, std::sync::atomic::Ordering::Release); + let _turn = second_lock.lock_owned().await; + second_order.lock().unwrap().push(2u8); + }); + while !second_queued.load(std::sync::atomic::Ordering::Acquire) { + tokio::task::yield_now().await; } - Ok(()) + drop(claim_turn); + + limit_thirdparty::tokio::time::timeout(Duration::from_secs(1), async { + first.await.expect("first same-class waiter panicked"); + second.await.expect("second same-class waiter panicked"); + }) + .await + .expect("same-class FIFO waiters did not complete"); + assert_eq!(*order.lock().unwrap(), vec![1, 2]); } - /// Put a key/value by encoding a flat dict from raw pointers directly into the segment pool. - /// - /// # Safety - /// The caller must guarantee the pointer ranges remain readable for the duration of this async call. - pub async unsafe fn put_flat_dict_ptrs( - &self, - key: &str, - ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, - opts: crate::client_kv_api::PutOptionalArgs, - ) -> KvResult<()> { - let lease_id = opts.lease_id(); - let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); - let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); + #[limit_thirdparty::tokio::test] + async fn pressured_class_does_not_block_an_unrelated_slot_class() { + const BLOCKED_SLOT_SIZE: u64 = 4 * 1024; + const READY_SLOT_SIZE: u64 = 8 * 1024; + let ready_slots_per_grant = owner_local_reserve_slots_per_grant(READY_SLOT_SIZE); - let payload_len = calc_flat_dict_encoded_len(&ptrs)?; - self.put_common( - key, - payload_len, - payload_len as u32, - reject_if_inflight_same_key, - preferred_sub_cluster.as_deref(), - lease_id, - payload_len as u32, - /*test_remove_after_fill=*/ false, - move |abs_src| { - // Fill owner's shared memory at abs_src directly from the raw pointers. - unsafe { - write_flat_dict_ptrs_to_ptr(abs_src as *mut u8, &ptrs); - } - }, - /*dbg_addr_summary=*/ false, - Some("put_flat_dict_ptrs"), + let api = Arc::new( + ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"), + ); + let blocked_lock = api + .inner() + .owner_local_reserve_claim_lock(BLOCKED_SLOT_SIZE); + let _blocked_turn = blocked_lock.lock().await; + { + let mut pool = api.inner().owner_local_reserve_pool.lock(); + owner_local_reserve_install_grant( + &mut pool, + READY_SLOT_SIZE, + ready_slots_per_grant, + OwnerLocalReserveGrantState::new( + 1, + 1000, + 1000, + READY_SLOT_SIZE * u64::from(ready_slots_per_grant), + READY_SLOT_SIZE, + ready_slots_per_grant, + ), + ); + } + + let ready_api = Arc::clone(&api); + let ready_lease = limit_thirdparty::tokio::time::timeout( + Duration::from_secs(1), + ready_api + .inner() + .owner_claim_local_reserve_slot_lease(READY_SLOT_SIZE, 1), ) .await + .expect("an unrelated slot class was head-of-line blocked") + .expect("ready slot class claim failed"); + assert_eq!(ready_lease.slot_size, READY_SLOT_SIZE); + assert_eq!(ready_lease.slots.len(), 1); } - /// Put a key/value with optional args (e.g., lease binding) - pub async fn put( + #[limit_thirdparty::tokio::test] + async fn queued_claims_publish_aggregate_demand_and_cancel_cleanly() { + const SLOT_SIZE: u64 = 4 * 1024; + + let api = Arc::new( + ClientKvApi::construct(ClientKvApiNewArg { + test_spec_config: TestSpecConfig::default(), + owner_hot_cache_capacity_bytes: None, + ssd_storage: None, + }) + .await + .expect("construct test ClientKvApi"), + ); + let claim_lock = api.inner().owner_local_reserve_claim_lock(SLOT_SIZE); + let claim_turn = claim_lock.lock().await; + + let first_api = Arc::clone(&api); + let first = tokio::spawn(async move { + first_api + .inner() + .owner_claim_local_reserve_slot_lease(SLOT_SIZE, 3) + .await + }); + let second_api = Arc::clone(&api); + let second = tokio::spawn(async move { + second_api + .inner() + .owner_claim_local_reserve_slot_lease(SLOT_SIZE, 2) + .await + }); + + limit_thirdparty::tokio::time::timeout(Duration::from_secs(1), async { + loop { + let pending = api + .inner() + .owner_local_reserve_pool + .lock() + .classes + .get(&SLOT_SIZE) + .map(|class| class.pending_slot_demand) + .unwrap_or_default(); + if pending == 5 { + break; + } + limit_thirdparty::tokio::task::yield_now().await; + } + }) + .await + .expect("queued claims did not publish aggregate demand"); + + first.abort(); + second.abort(); + let _ = first.await; + let _ = second.await; + assert_eq!( + api.inner() + .owner_local_reserve_pool + .lock() + .classes + .get(&SLOT_SIZE) + .unwrap() + .pending_slot_demand, + 0, + "cancelled queued claims leaked pending demand" + ); + drop(claim_turn); + } +} + +#[derive(Debug, Clone)] +pub struct OwnerReservedPutItem { + pub key: String, + pub put_id: PutIDForAKey, + pub src_addr: u64, + pub src_base_addr: u64, + pub target_addr: u64, + pub target_base_addr: u64, + pub value_len: u64, + pub lease_id: Option, + pub peer_node_id: Option, + pub remember_local_snapshot: bool, + /// Original caller/content selection, independent of remote capacity. + pub make_replica_task: bool, + /// The master pre-reserved a remote memory target for this Put. + pub remote_replica_admitted: bool, + pub preferred_sub_cluster: Option, +} + +#[derive(Debug, Clone)] +pub struct OwnerLocalPublishItem { + pub key: String, + pub put_id: PutIDForAKey, + pub value_len: u64, + pub lease_id: Option, + pub committed_slot: PutDoneCommittedSlot, + /// Original caller/content selection, independent of remote capacity. + pub make_replica_task: bool, + /// This Put path admitted a remote trigger. Pre-reserved paths set this + /// only when a target exists; owner-local fast Put admits selected items + /// here and performs allocation inside the later append Start. + pub remote_replica_admitted: bool, + pub preferred_sub_cluster: Option, + pub atomic_group: Option, +} + +#[derive(Debug, Clone)] +pub struct OwnerLocalPublishJob { + pub items: Vec, + pub key_reservation_ids: Vec, + /// External local-first requests keep their owner reclaim fences here until + /// the grouped master terminal response and all local promotions complete. + /// Native/Pyo3 jobs use master key reservations instead and leave this empty. + pub external_pending_contexts: Vec, +} + +fn owner_local_publish_starts_remote(item: &OwnerLocalPublishItem) -> bool { + item.make_replica_task && item.remote_replica_admitted +} + +#[derive(Clone)] +struct OwnerLocalSsdPutCandidate { + key: String, + put_id: PutIDForAKey, + /// Exact logical value bytes known by the publishing Put path. Early SSD + /// writes use this to reject over-budget candidates before installing a + /// generation flight or acquiring a source holder. A zero value is kept + /// only for legacy allocation-backed last-backing candidates whose + /// resource admission already happened upstream. + value_len: u64, + selected_victim: Option>, +} + +struct OwnerLocalSsdPutLeader { + op: Arc, + memory_info: Arc, + /// Bytes actually charged by the caller's admission pass. This can be + /// zero for a legacy allocation-backed last-backing candidate, so gate + /// failure must refund this value rather than minting credit from the + /// subsequently discovered source length. + admission_len: u64, +} + +fn owner_reclaim_backing_len(backing: &OwnerReclaimBacking) -> u64 { + match backing { + OwnerReclaimBacking::CommittedSlot { slot_size, .. } => *slot_size, + OwnerReclaimBacking::Allocation | OwnerReclaimBacking::UnindexedAllocation { .. } => 0, + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct PutEndStats { + pub master_put_end_rpc_us: i64, + pub master_put_end_server_us: i64, +} + +pub struct PutEndWithLocalCachePublish { + pub stats: PutEndStats, + pub local_cache_holder_id: Option, +} + +fn owner_local_reserve_timeout_error( + inner: &ClientKvApiInner, + stage: &'static str, + slot_size: u64, + key_count: usize, + soft_wait_timeout: std::time::Duration, + hard_wait_timeout: std::time::Duration, + request_started_at: Instant, +) -> KvError { + let (used_slots, free_slots, pending_slots, grants, expected_grants) = { + let pool = inner.owner_local_reserve_pool.lock(); + pool.classes + .get(&slot_size) + .map(|class_state| { + ( + class_state.used_slot_count(), + class_state.free_slot_count(), + class_state.pending_slot_demand, + class_state.grant_count(), + class_state.expected_grant_count, + ) + }) + .unwrap_or((0, 0, 0, 0, 0)) + }; + KvError::Api(ApiError::Unknown { + detail: format!( + "owner local reserve refill timeout: stage={} slot_size={} key_count={} remaining_slots={} used_slots={} free_slots={} pending_slots={} grants={} expected_grants={} waited_ms={} soft_wait_timeout_ms={} hard_timeout_ms={}", + stage, + slot_size, + key_count, + key_count.saturating_sub(free_slots), + used_slots, + free_slots, + pending_slots, + grants, + expected_grants, + request_started_at.elapsed().as_millis(), + soft_wait_timeout.as_millis(), + hard_wait_timeout.as_millis() + ), + }) +} + +impl ClientKvApiInner { + pub async fn owner_claim_local_reserve_slot_lease( &self, - key: &str, - value: &[u8], - opts: crate::client_kv_api::PutOptionalArgs, + value_len: u64, + key_count: usize, + ) -> KvResult { + let slot_size = owner_local_reserve_slot_size(value_len)?; + let slots_per_grant = owner_local_reserve_slots_per_grant(slot_size); + let (soft_wait_timeout, hard_wait_timeout) = owner_local_reserve_timeout_config(self); + let request_started_at = Instant::now(); + let hard_deadline = request_started_at + .checked_add(hard_wait_timeout) + .ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: "owner local reserve hard timeout overflow".to_string(), + }) + })?; + // Publish demand before waiting for the FIFO turn so one refill can cover all queued + // claimants. The guard also removes demand if the caller is cancelled while queued. + let mut pending_demand = + OwnerLocalReservePendingDemandGuard::new(self, slot_size, slots_per_grant, key_count); + self.owner_local_reserve_rebalance_notify().notify_waiters(); + + let Some(remaining_for_turn) = hard_deadline.checked_duration_since(Instant::now()) else { + return Err(owner_local_reserve_timeout_error( + self, + "claim_turn", + slot_size, + key_count, + soft_wait_timeout, + hard_wait_timeout, + request_started_at, + )); + }; + let claim_lock = self.owner_local_reserve_claim_lock(slot_size); + let _claim_turn = match tokio::time::timeout(remaining_for_turn, claim_lock.lock()).await { + Ok(claim_turn) => claim_turn, + Err(_) => { + return Err(owner_local_reserve_timeout_error( + self, + "claim_turn", + slot_size, + key_count, + soft_wait_timeout, + hard_wait_timeout, + request_started_at, + )); + } + }; + + loop { + let claim = { + let mut pool = self.owner_local_reserve_pool.lock(); + let claim = owner_local_reserve_try_claim( + &mut pool, + slot_size, + slots_per_grant, + value_len, + key_count, + ); + if claim.is_some() { + let class_state = pool + .classes + .get_mut(&slot_size) + .expect("claimed local-reserve class must exist"); + class_state.pending_slot_demand = class_state + .pending_slot_demand + .checked_sub(key_count) + .expect("claimed local-reserve demand underflow"); + } + claim + }; + if let Some(lease) = claim { + pending_demand.disarm_after_locked_consume(); + return Ok(lease); + } + self.owner_local_reserve_rebalance_notify().notify_waiters(); + if !wait_owner_local_reserve_ready( + self, + slot_size, + slots_per_grant, + key_count, + soft_wait_timeout, + hard_deadline, + ) + .await + { + break; + } + } + Err(owner_local_reserve_timeout_error( + self, + "refill", + slot_size, + key_count, + soft_wait_timeout, + hard_wait_timeout, + request_started_at, + )) + } + + pub async fn owner_release_local_reserve_slot_lease( + &self, + lease: OwnerLocalReserveSlotLease, + ) -> KvResult<()> { + { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&lease.slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "resident local reserve class missing while releasing slot lease: slot_size={}", + lease.slot_size + ), + })); + }; + for slot_ref in &lease.slots { + if !class_state.release_prepared_slot(slot_ref.grant_id, slot_ref.slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "resident local reserve grant missing while releasing slot lease: grant_id={}", + slot_ref.grant_id + ), + })); + } + } + } + self.owner_local_reserve_rebalance_notify().notify_waiters(); + Ok(()) + } + + pub fn owner_mark_local_reserve_slot_pending_visible( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while marking pending slot: slot_size={}", + slot_size + ), + })); + }; + if !class_state.mark_prepared_slot_pending_visible(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while marking pending slot: grant_id={}", + grant_id + ), + })); + } + Ok(()) + } + + pub fn owner_promote_local_reserve_pending_slot_to_committed( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while promoting pending slot: slot_size={}", + slot_size + ), + })); + }; + if !class_state.promote_pending_visible_slot_to_committed(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while promoting pending slot: grant_id={}", + grant_id + ), + })); + } + Ok(()) + } + + pub fn owner_retain_local_reserve_resident_slot_holder( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while retaining resident slot holder: slot_size={}", + slot_size + ), + })); + }; + if !class_state.retain_resident_slot_holder(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while retaining resident slot holder: grant_id={}", + grant_id + ), + })); + } + Ok(()) + } + + pub fn owner_release_local_reserve_resident_slot_holder( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while releasing resident slot holder: slot_size={}", + slot_size + ), + })); + }; + if !class_state.release_resident_slot_holder(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while releasing resident slot holder: grant_id={}", + grant_id + ), + })); + } + } + self.owner_local_reserve_rebalance_notify().notify_waiters(); + Ok(()) + } + + pub fn owner_release_local_reserve_committed_slot_route( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while releasing committed slot route: slot_size={}", + slot_size + ), + })); + }; + if !class_state.release_committed_slot_route(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while releasing committed slot route: grant_id={}", + grant_id + ), + })); + } + } + self.owner_local_reserve_rebalance_notify().notify_waiters(); + Ok(()) + } + + pub fn owner_release_local_reserve_committed_resident_slot( + &self, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> KvResult<()> { + { + let mut pool = self.owner_local_reserve_pool.lock(); + let Some(class_state) = pool.classes.get_mut(&slot_size) else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve class missing while reclaiming committed resident slot: slot_size={}", + slot_size + ), + })); + }; + if !class_state.release_committed_resident_slot(grant_id, slot_index) { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "local reserve grant missing while reclaiming committed resident slot: grant_id={}", + grant_id + ), + })); + } + } + self.owner_local_reserve_rebalance_notify().notify_waiters(); + Ok(()) + } + + pub async fn owner_shutdown_local_reserve_pool(&self) -> KvResult<()> { + let grants = { + let mut pool = self.owner_local_reserve_pool.lock(); + let mut detached = Vec::new(); + for (_slot_size, mut class_state) in pool.classes.drain() { + detached.extend(class_state.take_all_grants()); + } + detached + }; + let mut first_err = None; + for grant in grants { + if let Err(err) = self.release_local_grant(grant.grant_id).await { + if first_err.is_none() { + first_err = Some(err); + } else { + tracing::warn!( + "owner_shutdown_local_reserve_pool dropped additional release error after the first one: {}", + err + ); + } + } + } + if let Some(err) = first_err { + return Err(err); + } + Ok(()) + } + + pub async fn owner_batch_put_start_reserved( + &self, + start_items: Vec, + lease_id: Option, + ) -> KvResult> { + if start_items.is_empty() { + return Ok(Vec::new()); + } + if self.short_circuit_put_payload_path_enabled() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: + "owner_batch_put_start_reserved does not support short_circuit_put_payload_path" + .to_string(), + })); + } + let self_node_id = self.view.cluster_manager().get_self_info().id.clone(); + let start_resp = self.batch_put_start(start_items.clone()).await?; + if start_resp.items.len() != start_items.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_put_start response length mismatch: expected={} got={}", + start_items.len(), + start_resp.items.len() + ), + })); + } + + let mut prepared_items = Vec::with_capacity(start_items.len()); + let mut revoke_items = Vec::with_capacity(start_items.len()); + let mut first_error: Option = None; + + for (start_req, start_item) in start_items.into_iter().zip(start_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + start_item.error_code, + start_item.error_json.clone(), + ) { + first_error = Some(err); + break; + } + let peer_node_id = if start_item.node_id == self_node_id { + None + } else { + Some(start_item.node_id.clone()) + }; + let replica_admitted = start_item.replica_target.is_some(); + revoke_items.push(BatchPutRevokeItemReq { + key: start_req.key.clone(), + put_id: start_item.put_id, + }); + prepared_items.push(OwnerReservedPutItem { + key: start_req.key, + put_id: start_item.put_id, + src_addr: start_item.src_addr, + src_base_addr: start_item.src_base_addr, + target_addr: start_item.target_addr, + target_base_addr: start_item.target_base_addr, + value_len: start_req.len, + lease_id, + peer_node_id: peer_node_id.clone(), + remember_local_snapshot: true, + make_replica_task: start_req.make_replica_task, + remote_replica_admitted: replica_admitted, + preferred_sub_cluster: start_req.preferred_sub_cluster, + }); + } + + if let Some(err) = first_error { + if !revoke_items.is_empty() { + if let Err(revoke_err) = self.batch_put_revoke(revoke_items).await { + tracing::warn!( + "owner_batch_put_start_reserved batch_put_revoke failed after partial reserve: {}", + revoke_err + ); + } + } + return Err(err); + } + + Ok(prepared_items) + } + + pub async fn owner_batch_put_commit_reserved( + &self, + items: Vec, + _transfer_concurrency: usize, + ) -> KvResult>> { + if items.is_empty() { + return Ok(Vec::new()); + } + + #[derive(Clone)] + struct DonePending { + idx: usize, + item: OwnerReservedPutItem, + } + + let metrics = self.metrics_handle(); + let mut results: Vec>> = (0..items.len()).map(|_| None).collect(); + let mut done_pending = Vec::with_capacity(items.len()); + + for (idx, item) in items.into_iter().enumerate() { + metrics.record_put_io_locality(false, item.value_len, 0); + done_pending.push(DonePending { idx, item }); + } + + if self.skip_put_end_commit_enabled() { + for pending in done_pending { + results[pending.idx] = Some(Ok(())); + } + return Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "owner_batch_put_commit_reserved result slot was not populated" + .to_string(), + })) + }) + }) + .collect()); + } + + let done_req_items = done_pending + .iter() + .map(|pending| BatchPutDoneItemReq { + key: pending.item.key.clone(), + put_id: pending.item.put_id, + lease_id: pending.item.lease_id, + committed_slot: None, + publish_local_cache: false, + atomic_group: None, + }) + .collect::>(); + let done_resp = self.batch_put_done(done_req_items).await?; + if done_resp.items.len() != done_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_put_done response length mismatch: expected={} got={}", + done_pending.len(), + done_resp.items.len() + ), + })); + } + + let mut early_ssd_candidates = Vec::new(); + let mut remote_replica_pending = Vec::new(); + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + results[pending.idx] = Some(Err(err)); + continue; + } + if pending.item.remember_local_snapshot { + self.remember_local_snapshot(&pending.item.key, pending.item.put_id); + } + if pending.item.make_replica_task { + early_ssd_candidates.push(( + pending.item.key.clone(), + pending.item.put_id, + pending.item.value_len, + )); + } + if pending.item.make_replica_task && pending.item.remote_replica_admitted { + remote_replica_pending.push(( + pending.item.key.clone(), + pending.item.put_id, + pending.item.preferred_sub_cluster.clone(), + )); + } + results[pending.idx] = Some(Ok(())); + } + + self.start_early_owner_local_ssd_puts(early_ssd_candidates); + for (key, put_id, preferred_sub_cluster) in remote_replica_pending { + if let Err(err) = self + .ensure_remote_put(&key, put_id, preferred_sub_cluster, true) + .await + { + tracing::warn!( + "owner_batch_put_commit_reserved make replica task failed after local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } + + Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "owner_batch_put_commit_reserved result slot was not populated" + .to_string(), + })) + }) + }) + .collect()) + } + + pub async fn owner_batch_put_abort_reserved( + &self, + items: Vec, ) -> KvResult<()> { + if items.is_empty() { + return Ok(()); + } + let revoke_items = items + .into_iter() + .map(|item| BatchPutRevokeItemReq { + key: item.key, + put_id: item.put_id, + }) + .collect::>(); + self.batch_put_revoke(revoke_items).await?; + Ok(()) + } + + pub async unsafe fn batch_put_flat_dict_ptrs( + &self, + keys: Vec, + ptrs_groups: Vec)>>, + opts: crate::client_kv_api::PutOptionalArgs, + _transfer_concurrency: usize, + ) -> KvResult>> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_flat_dict_ptrs" + .to_string(), + })); + } + if keys.len() != ptrs_groups.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch_put_flat_dict_ptrs requires keys and ptrs_groups to have the same length: keys={} ptrs_groups={}", + keys.len(), + ptrs_groups.len() + ), + })); + } + if keys.is_empty() { + return Ok(Vec::new()); + } let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); - let payload_len = value.len() as u64; - self.put_common( - key, - payload_len, - value.len() as u32, - reject_if_inflight_same_key, - preferred_sub_cluster.as_deref(), - lease_id, - value.len() as u32, - /*test_remove_after_fill=*/ true, - |abs_src| unsafe { - std::ptr::copy_nonoverlapping(value.as_ptr(), abs_src as *mut u8, value.len()); + + let mut start_items = Vec::with_capacity(keys.len()); + let mut payload_lens = Vec::with_capacity(keys.len()); + for (key, ptrs) in keys.iter().zip(ptrs_groups.iter()) { + let payload_len = calc_flat_dict_encoded_len(ptrs)?; + payload_lens.push(payload_len); + start_items.push(BatchPutStartItemReq { + key: key.clone(), + len: payload_len, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster: preferred_sub_cluster.clone(), + }); + } + + let start_resp = self.batch_put_start(start_items).await?; + if start_resp.items.len() != keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_put_start response length mismatch: expected={} got={}", + keys.len(), + start_resp.items.len() + ), + })); + } + + #[derive(Clone)] + struct DonePending { + idx: usize, + key: String, + put_id: PutIDForAKey, + value_len: u64, + lease_id: Option, + remember_local_snapshot: bool, + make_replica_task: bool, + remote_replica_admitted: bool, + preferred_sub_cluster: Option, + } + + let mut results: Vec>> = (0..keys.len()).map(|_| None).collect(); + let mut done_pending = Vec::new(); + let short_circuit_payload = self.short_circuit_put_payload_path_enabled(); + let skip_put_end_commit = self.skip_put_end_commit_enabled(); + + for (idx, (((key, ptrs), payload_len), start_item)) in keys + .into_iter() + .zip(ptrs_groups.into_iter()) + .zip(payload_lens.into_iter()) + .zip(start_resp.items.into_iter()) + .enumerate() + { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + start_item.error_code, + start_item.error_json.clone(), + ) { + results[idx] = Some(Err(err)); + continue; + } + + let put_id = start_item.put_id; + let remember_local_snapshot = true; + let replica_admitted = start_item.replica_target.is_some(); + + if !short_circuit_payload { + unsafe { + write_flat_dict_ptrs_to_ptr(start_item.src_addr as *mut u8, &ptrs); + } + } + + done_pending.push(DonePending { + idx, + key, + put_id, + value_len: payload_len, + lease_id, + remember_local_snapshot, + make_replica_task, + remote_replica_admitted: replica_admitted, + preferred_sub_cluster: preferred_sub_cluster.clone(), + }); + } + + if skip_put_end_commit { + for pending in done_pending { + results[pending.idx] = Some(Ok(())); + } + return Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "batch_put result slot was not populated".to_string(), + })) + }) + }) + .collect()); + } + + let done_req_items = done_pending + .iter() + .map(|pending| BatchPutDoneItemReq { + key: pending.key.clone(), + put_id: pending.put_id, + lease_id: pending.lease_id, + committed_slot: None, + publish_local_cache: false, + atomic_group: None, + }) + .collect::>(); + let done_resp = self.batch_put_done(done_req_items).await?; + if done_resp.items.len() != done_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "batch_put_done response length mismatch: expected={} got={}", + done_pending.len(), + done_resp.items.len() + ), + })); + } + let mut early_ssd_candidates = Vec::new(); + let mut remote_replica_pending = Vec::new(); + for (pending, done_item) in done_pending.into_iter().zip(done_resp.items.into_iter()) { + if let Err(err) = crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) { + results[pending.idx] = Some(Err(err)); + continue; + } + if pending.remember_local_snapshot { + self.remember_local_snapshot(&pending.key, pending.put_id); + } + if pending.make_replica_task { + early_ssd_candidates.push((pending.key.clone(), pending.put_id, pending.value_len)); + } + if pending.make_replica_task && pending.remote_replica_admitted { + remote_replica_pending.push(( + pending.key.clone(), + pending.put_id, + pending.preferred_sub_cluster.clone(), + )); + } + results[pending.idx] = Some(Ok(())); + } + + self.start_early_owner_local_ssd_puts(early_ssd_candidates); + for (key, put_id, preferred_sub_cluster) in remote_replica_pending { + if let Err(err) = self + .ensure_remote_put(&key, put_id, preferred_sub_cluster, true) + .await + { + tracing::warn!( + "batch make replica task failed after local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } + + Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "batch_put result slot was not populated".to_string(), + })) + }) + }) + .collect()) + } + + fn start_owner_local_ssd_puts( + &self, + candidates: Vec, + resources_pre_admitted: bool, + ) -> Vec>> { + // Early backing is optional work. Apply its independent byte budget + // before touching the per-key generation table or pinning a source. + // Last-backing callers have already performed the same admission while + // deciding which capacity victims may enter durability, so they pass + // resources_pre_admitted=true and must not be charged twice. + let admission = if resources_pre_admitted { + vec![true; candidates.len()] + } else { + let lengths = candidates + .iter() + .map(|candidate| candidate.value_len) + .collect::>(); + self.ssd_storage + .as_ref() + .map(|store| store.admit_owner_write_candidates(&lengths)) + .unwrap_or_else(|| vec![false; candidates.len()]) + }; + let mut operations = std::iter::repeat_with(|| None) + .take(candidates.len()) + .collect::>(); + let mut leaders = Vec::new(); + let mut unused_pre_admission_lengths = Vec::new(); + for (index, (candidate, admitted_by_resource)) in + candidates.into_iter().zip(admission).enumerate() + { + if !admitted_by_resource { + continue; + } + let pre_admitted_len = candidate.value_len; + match self.begin_owner_local_ssd_put( + &candidate.key, + candidate.put_id, + candidate.selected_victim.as_deref(), + ) { + OwnerLocalSsdPutReservation::Leader { op, memory_info } => { + operations[index] = Some(op.clone()); + if pre_admitted_len != 0 && u64::from(memory_info.len) != pre_admitted_len { + tracing::warn!( + key = %candidate.key, + put_id_time = candidate.put_id.0, + put_id_version = candidate.put_id.1, + admitted_len = pre_admitted_len, + source_len = memory_info.len, + "owner local SSD pre-admission length no longer matches exact source" + ); + unused_pre_admission_lengths.push(pre_admitted_len); + self.finish_owner_local_ssd_put(&op, OwnerLocalSsdPutOutcome::Failed); + continue; + } + leaders.push(OwnerLocalSsdPutLeader { + op, + memory_info, + admission_len: pre_admitted_len, + }); + } + OwnerLocalSsdPutReservation::Follower(op) => { + operations[index] = Some(op); + unused_pre_admission_lengths.push(pre_admitted_len); + } + OwnerLocalSsdPutReservation::SourceUnavailable => { + unused_pre_admission_lengths.push(pre_admitted_len); + } + } + } + if let Some(store) = self.ssd_storage.as_ref() { + store.refund_owner_write_admission(&unused_pre_admission_lengths); + } + if leaders.is_empty() { + return operations; + } + + let admitted_lengths = leaders + .iter() + .map(|leader| leader.admission_len) + .collect::>(); + let permit = match self.try_acquire_local_ssd_persist_batch(leaders.len()) { + Ok(Some(permit)) => permit, + Ok(None) => { + if let Some(store) = self.ssd_storage.as_ref() { + store.refund_owner_write_admission(&admitted_lengths); + } + for leader in leaders { + self.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Dropped); + } + return operations; + } + Err(err) => { + if let Some(store) = self.ssd_storage.as_ref() { + store.refund_owner_write_admission(&admitted_lengths); + } + tracing::warn!( + items = leaders.len(), + error = %err, + "owner local SSD durability batch admission failed" + ); + for leader in leaders { + self.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + return operations; + } + }; + + let spawn_view = self.view.clone_view(); + let task_view = spawn_view.clone(); + let _ = spawn_view.spawn("owner_local_ssd_put_batch", async move { + run_owner_local_ssd_put_batch(task_view, leaders, permit).await; + }); + operations + } + + pub(crate) fn start_early_owner_local_ssd_puts( + &self, + candidates: Vec<(String, PutIDForAKey, u64)>, + ) { + let _ = self.start_owner_local_ssd_puts( + candidates + .into_iter() + .map(|(key, put_id, value_len)| OwnerLocalSsdPutCandidate { + key, + put_id, + value_len, + selected_victim: None, + }) + .collect(), + false, + ); + } + + pub(crate) fn start_early_owner_local_ssd_put( + &self, + key: &str, + put_id: PutIDForAKey, + value_len: u64, + ) { + self.start_early_owner_local_ssd_puts(vec![(key.to_string(), put_id, value_len)]); + } + + /// Ensure one remote backing exists for an exact owner-local generation. + /// + /// Normal Put, pre-reserved replica, proactive write-back, and tier1 all + /// enter here. Only the flight leader pins the source and immediately + /// launches Start/transfer/Done in its own async task; there is no actor or + /// intermediate work queue. Followers wait for and reuse that terminal + /// result instead of acquiring their own holder. + pub async fn ensure_remote_put( + &self, + key: &str, + put_id: PutIDForAKey, + preferred_sub_cluster: Option, + protect_source_on_remote_complete: bool, + ) -> KvResult { + let mut follower_takeover_attempted = false; + loop { + match self.begin_owner_remote_put( + key, + put_id, + preferred_sub_cluster.clone(), + protect_source_on_remote_complete, + ) { + OwnerRemotePutReservation::Leader { op, memory_info } => { + let holder = Arc::new(UserMemHolder::new( + memory_info, + self.get_or_init_all_memholder_refcount(), + UserMemHolderExposeKind::SegPtr, + )); + let spawn_view = self.view.clone_view(); + let task_view = spawn_view.clone(); + let _ = spawn_view.spawn("owner_remote_put_leader", async move { + run_owner_remote_put(task_view, op, holder).await; + }); + return Ok(true); + } + OwnerRemotePutReservation::Follower(op) => { + tracing::debug!( + "owner remote Put joined existing flight: key={} put_id=({},{})", + op.key, + op.put_id.0, + op.put_id.1 + ); + match op.wait().await { + OwnerRemotePutOutcome::Published + | OwnerRemotePutOutcome::AlreadySatisfied => return Ok(true), + OwnerRemotePutOutcome::Obsolete => return Ok(false), + OwnerRemotePutOutcome::Failed if !follower_takeover_attempted => { + follower_takeover_attempted = true; + continue; + } + OwnerRemotePutOutcome::Failed => return Ok(false), + OwnerRemotePutOutcome::InFlight => { + unreachable!("remote Put wait must return a terminal outcome") + } + } + } + OwnerRemotePutReservation::SourceUnavailable => { + tracing::warn!( + "owner remote Put source is unavailable, fenced, or version-mismatched: key={} put_id=({},{})", + key, + put_id.0, + put_id.1 + ); + return Ok(false); + } + } + } + } + + async fn put_common( + &self, + key: &str, + payload_len: u64, + len_for_start: u32, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, + preferred_sub_cluster: Option<&str>, + lease_id: Option, + _test_payload_len_u32: u32, + _test_remove_after_fill: bool, + fill_abs_src: F, + dbg_addr_summary: bool, + info_complete_tag: Option<&'static str>, + ) -> KvResult<()> + where + F: FnOnce(u64), + { + let client_id = self.client_id_str(); + let node_role = self.node_role(); + let metrics = self.metrics_handle(); + + let t1 = Utc::now().timestamp_micros(); + let (resp, _rpc_latency) = { + match self + .put_start( + key, + len_for_start, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster, + ) + .await + { + Ok(resp) => resp, + Err(err) => { + obe_put_start_error_rpc(&metrics, &client_id, &node_role, key, payload_len); + return Err(err); + } + } + }; + let t2 = Utc::now().timestamp_micros(); + if let Err(e) = + crate::rpcresp_kvresult_convert::try_from_code(resp.error_code, resp.error_json.clone()) + { + obe_put_start_error_status(&metrics, &client_id, &node_role, key, payload_len); + return Err(e); + } + obe_put_start_success(&metrics, &client_id, &node_role, key, t1, t2); + + let put_id = resp.put_id; + let peer_id = if &*resp.node_id == &*self.view.cluster_manager().get_self_info().id { + None + } else { + Some(resp.node_id.clone()) + }; + let abs_src = resp.src_addr; + let abs_target = resp.target_addr; + let replica_admitted = resp.replica_target.is_some(); + + #[cfg(test)] + { + self.test_record.add_transfering_put( + key.to_string(), + _test_payload_len_u32, + put_id.0, + put_id.1, + resp.node_id.to_string(), + format!("{:#x}", resp.target_addr), + ); + } + + if self.short_circuit_put_payload_path_enabled() { + #[cfg(test)] + { + if _test_remove_after_fill { + self.test_record + .remove_transfering_put(key.to_string(), put_id); + } + } + + let skipped_breakdown = if peer_id.is_none() && abs_src == abs_target { + TransferBreakdown { + local_noop: true, + ..TransferBreakdown::default() + } + } else { + TransferBreakdown::default() + }; + metrics.pending_put_set_transfer_breakdown( + put_id, + skipped_breakdown.submit_blocking_us, + skipped_breakdown.create_xfer_req_us, + skipped_breakdown.post_xfer_req_us, + skipped_breakdown.poll_wait_us, + skipped_breakdown.poll_iters, + skipped_breakdown.used_fast_path, + false, + skipped_breakdown.local_noop, + skipped_breakdown.remote_transfer, + ); + self.put_end(key, put_id, lease_id).await?; + self.remember_local_snapshot(key, put_id); + if make_replica_task { + self.start_early_owner_local_ssd_put(key, put_id, payload_len); + } + if make_replica_task && replica_admitted { + if let Err(err) = self + .ensure_remote_put(key, put_id, preferred_sub_cluster.map(str::to_string), true) + .await + { + tracing::warn!( + "make replica task failed after short-circuit local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } + if let Some(tag) = info_complete_tag { + info!("{tag} complete key={} bytes={}", key, payload_len); + } + return Ok(()); + } + + fill_abs_src(abs_src); + + #[cfg(test)] + { + if _test_remove_after_fill { + self.test_record + .remove_transfering_put(key.to_string(), put_id); + } + } + + let base_addr = self + .view + .client_seg_pool() + .cpu_mem_read_guard() + .await + .unwrap() + .allocated_addr; + if dbg_addr_summary { + tracing::debug!( + "put path addr summary: key={}, put_id=({},{}) local_base={:#x}, abs_src={:#x}, master_target_base={:#x}, abs_target={:#x}, peer_id={:?}", + key, + put_id.0, + put_id.1, + base_addr, + abs_src, + resp.target_base_addr, + abs_target, + peer_id + ); + } + + if self.skip_put_end_commit_enabled() { + let _ = metrics.pending_put_remove(&put_id); + tracing::warn!( + "skip_put_end_commit test-only fast-path: returning success without put_end; key={} put_id=({},{}) payload_len={}", + key, + put_id.0, + put_id.1, + payload_len + ); + if let Some(tag) = info_complete_tag { + info!( + "{tag} complete_without_put_end key={} bytes={}", + key, payload_len + ); + } + return Ok(()); + } + + self.put_end(key, put_id, lease_id).await?; + self.remember_local_snapshot(key, put_id); + if make_replica_task { + self.start_early_owner_local_ssd_put(key, put_id, payload_len); + } + if make_replica_task && replica_admitted { + if let Err(err) = self + .ensure_remote_put(key, put_id, preferred_sub_cluster.map(str::to_string), true) + .await + { + tracing::warn!( + "make replica task failed after local commit: key={} put_id=({},{}) err={}", + key, + put_id.0, + put_id.1, + err + ); + } + } + if let Some(tag) = info_complete_tag { + info!("{tag} complete key={} bytes={}", key, payload_len); + } + Ok(()) + } + + /// Put a key/value by encoding a flat dict from raw pointers directly into the segment pool. + /// + /// # Safety + /// The caller must guarantee the pointer ranges remain readable for the duration of this async call. + pub async unsafe fn put_flat_dict_ptrs( + &self, + key: &str, + ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, + opts: crate::client_kv_api::PutOptionalArgs, + ) -> KvResult<()> { + let lease_id = opts.lease_id(); + let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); + let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); + + let payload_len = calc_flat_dict_encoded_len(&ptrs)?; + self.put_common( + key, + payload_len, + payload_len as u32, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster.as_deref(), + lease_id, + payload_len as u32, + /*test_remove_after_fill=*/ false, + move |abs_src| { + // Fill owner's shared memory at abs_src directly from the raw pointers. + unsafe { + write_flat_dict_ptrs_to_ptr(abs_src as *mut u8, &ptrs); + } + }, + /*dbg_addr_summary=*/ false, + Some("put_flat_dict_ptrs"), + ) + .await + } + + /// Put a key/value with optional args (e.g., lease binding) + pub async fn put( + &self, + key: &str, + value: &[u8], + opts: crate::client_kv_api::PutOptionalArgs, + ) -> KvResult<()> { + let lease_id = opts.lease_id(); + let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); + let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); + let payload_len = value.len() as u64; + self.put_common( + key, + payload_len, + value.len() as u32, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster.as_deref(), + lease_id, + value.len() as u32, + /*test_remove_after_fill=*/ true, + |abs_src| unsafe { + std::ptr::copy_nonoverlapping(value.as_ptr(), abs_src as *mut u8, value.len()); + }, + /*dbg_addr_summary=*/ true, + None, + ) + .await + } + + /// Transfer data by offsets with instrumentation for external/owner callers. + /// Records transfer latency (t2..t3) and emits tsbuckets pulses. + pub async fn put_transfer( + &self, + key: &str, + put_id: PutIDForAKey, + src_offset: u64, + target_offset: u64, + len: u64, + peer_id: Option, + target_base_addr: Option, + ) -> KvResult { + let metrics = self.metrics_handle(); + let client_id = self.client_id_str(); + let node_role = self.node_role(); + + // owner/external inner is stable after construction; base_addr must exist + let base_addr = self + .view + .client_seg_pool() + .cpu_mem_read_guard() + .await + .unwrap() + .allocated_addr; + let abs_src = base_addr + src_offset; + let abs_target = if peer_id.is_some() { + let Some(tb) = target_base_addr else { + // propagate as Unreachable: invalid remote target context from distributed input + let err = crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { + rpc_input_json: format!( + "missing target_base_addr while peer_id present; src_off={:#x}, tgt_off={:#x}", + src_offset, target_offset + ), + }, + ); + return Err(err); + }; + tb + target_offset + } else { + base_addr + target_offset + }; + + // Local placement can resolve to src==target, which means the payload is already in-place. + // Skip the transfer-engine hop for this no-op path to avoid paying an extra fixed cost. + if peer_id.is_none() && abs_src == abs_target { + tracing::debug!( + "put_transfer local no-op: key={}, put_id=({},{}) src==target {:#x}, len={}", + key, + put_id.0, + put_id.1, + abs_target, + len + ); + return Ok(TransferBreakdown { + local_noop: true, + ..TransferBreakdown::default() + }); + } else { + let breakdown = self + .view + .client_transfer_engine() + .transfer_data_no_copy(peer_id.clone(), false, abs_src, abs_target, len, None) + .await?; + tracing::debug!( + "put_transfer breakdown: key={}, put_id=({},{}) fast_path={} nixl={} local_noop={} remote_transfer={} submit_blocking_us={} create_xfer_req_us={} post_xfer_req_us={} poll_wait_us={} poll_iters={}", + key, + put_id.0, + put_id.1, + breakdown.used_fast_path, + false, + breakdown.local_noop, + breakdown.remote_transfer, + breakdown.submit_blocking_us, + breakdown.create_xfer_req_us, + breakdown.post_xfer_req_us, + breakdown.poll_wait_us, + breakdown.poll_iters + ); + tracing::debug!( + "put_transfer success: key={}, put_id=({},{}) src_off={:#x}, tgt_off={:#x}, len={}, peer_id={:?}", + key, + put_id.0, + put_id.1, + src_offset, + target_offset, + len, + peer_id + ); + + // Emit transfer stage success and tsbuckets pulse (computes t2/t3 using pending) + crate::observe_kvope::obe_put_transfer_success( + &metrics, &client_id, &node_role, key, len, put_id, + ); + return Ok(breakdown); + } + #[allow(unreachable_code)] + Ok(TransferBreakdown::default()) + } + + /// 开始 Put 操作，分配存储空间 + pub async fn put_start_with_source_node( + &self, + key: &str, + len: u32, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, + preferred_sub_cluster: Option<&str>, + source_node_id: Option, + ) -> KvResult<(PutStartResp, i64)> { + let req = MsgPack { + serialize_part: PutStartReq { + key: key.to_string(), + len: len as u64, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), + source_node_id, + }, + raw_bytes: Vec::new(), + }; + + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let rpc_started_at = Instant::now(); + let start_rpc_timestamp = Utc::now().timestamp_micros() as i64; + let resp = self + .rpc_caller_put_start + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(|e| KvError::P2p(e))?; + let end_rpc_timestamp = Utc::now().timestamp_micros() as i64; + let ser = resp.serialize_part.clone(); + if crate::rpcresp_kvresult_convert::try_from_code(ser.error_code, ser.error_json.clone()) + .is_ok() + { + let metrics = self.metrics_handle(); + metrics.pending_put_insert( + ser.put_id, + key.to_string(), + len as u64, + start_rpc_timestamp, + end_rpc_timestamp, + ser.server_process_us, + ); + } + let rpc_latency_us = duration_to_i64_us(rpc_started_at.elapsed()); + Ok((ser, rpc_latency_us)) + } + + /// 开始 Put 操作，分配存储空间 + pub async fn put_start( + &self, + key: &str, + len: u32, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, + preferred_sub_cluster: Option<&str>, + ) -> KvResult<(PutStartResp, i64)> { + self.put_start_with_source_node( + key, + len, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster, + None, + ) + .await + } + + pub async fn reserve_local_grant(&self) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting reserve_local_grant".to_string(), + })); + } + let req = MsgPack { + serialize_part: ReserveLocalGrantReq {}, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_reserve_local_grant + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + match resp.serialize_part.outcome { + ReserveLocalGrantOutcome::None => Err(KvError::Api(ApiError::Unknown { + detail: "reserve_local_grant returned success without an outcome".to_string(), + })), + outcome => Ok(outcome), + } + } + + pub async fn release_local_grant(&self, grant_id: u64) -> KvResult<()> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting release_local_grant".to_string(), + })); + } + let req = MsgPack { + serialize_part: ReleaseLocalGrantReq { grant_id }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_release_local_grant + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(()) + } + + /// 撤销 Put 操作，释放已分配的资源 + pub async fn put_revoke(&self, key: &str, put_id: PutIDForAKey) -> KvResult<()> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_revoke".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutRevokeReq { + key: key.to_string(), + put_id, + }, + raw_bytes: Vec::new(), + }; + + // 获取 master 节点 ID + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + + // 调用 RPC + let _resp = self + .rpc_caller_put_revoke + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + // cleanup pending stat if any + let _ = self.metrics_handle().pending_put_remove(&put_id); + Ok(()) + } + + pub async fn put_append_start( + &self, + key: &str, + put_id: PutIDForAKey, + len: u32, + preferred_sub_cluster: Option<&str>, + protect_source_on_remote_complete: bool, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_append_start".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutAppendStartReq { + key: key.to_string(), + put_id, + len: len as u64, + preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), + protect_source_on_remote_complete, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_put_append_start + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_put_append_start( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPutAppendStartResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_append_start" + .to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPutAppendStartReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_put_append_start + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_evict_owner_source( + &self, + victims: Vec, + ) -> KvResult { + if victims.is_empty() { + return Ok(BatchEvictOwnerSourceResp { + operation_id: 0, + victims: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting owner source eviction".to_string(), + })); + } + let operation_id = self + .next_owner_source_eviction_operation_id + .fetch_add(1, Ordering::Relaxed); + let self_info = self.view.cluster_manager().get_self_info(); + let req = MsgPack { + serialize_part: BatchEvictOwnerSourceReq { + operation_id, + owner_node_start_time: self_info.node_start_time, + victims, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_evict_owner_source + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + if resp.serialize_part.operation_id != operation_id { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "owner source-eviction operation id mismatch: requested={} response={}", + operation_id, resp.serialize_part.operation_id + ), + })); + } + Ok(resp.serialize_part) + } + + async fn batch_publish_owner_ssd( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPublishOwnerSsdResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting owner SSD publication".to_string(), + })); + } + let self_info = self.view.cluster_manager().get_self_info(); + let req = MsgPack { + serialize_part: BatchPublishOwnerSsdReq { + owner_node_start_time: self_info.node_start_time, + items, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_publish_owner_ssd + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn put_append_revoke( + &self, + key: &str, + put_id: PutIDForAKey, + operation_id: u64, + ) -> KvResult<()> { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_append_revoke".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutAppendRevokeReq { + key: key.to_string(), + put_id, + operation_id, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let _resp = self + .rpc_caller_put_append_revoke + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + Ok(()) + } + + pub async fn put_append_done( + &self, + key: &str, + put_id: PutIDForAKey, + operation_id: u64, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_append_done".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutAppendDoneReq { + key: key.to_string(), + put_id, + operation_id, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_put_append_done + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_put_append_done( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPutAppendDoneResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_append_done".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPutAppendDoneReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_put_append_done + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_put_start( + &self, + items: Vec, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_start".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPutStartReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_put_start + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_prepare_put_keys( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPreparePutKeysResp { + reservation_ids: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_prepare_put_keys" + .to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPreparePutKeysReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_prepare_put_keys + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_release_put_key_reservations( + &self, + reservation_ids: Vec, + ) -> KvResult { + if reservation_ids.is_empty() { + return Ok(BatchReleasePutKeyReservationsResp { + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: + "ClientKvApi is shutting down; rejecting batch_release_put_key_reservations" + .to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchReleasePutKeyReservationsReq { reservation_ids }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_release_put_key_reservations + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_put_revoke( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPutRevokeResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_revoke".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPutRevokeReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_put_revoke + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 2, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn batch_put_done( + &self, + items: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(BatchPutDoneResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting batch_put_done".to_string(), + })); + } + let req = MsgPack { + serialize_part: BatchPutDoneReq { items }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_batch_put_done + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + pub async fn grouped_batch_put_done( + &self, + items: Vec, + atomic_group_lens: Vec, + ) -> KvResult { + if items.is_empty() { + return Ok(GroupedBatchPutDoneResp { + items: Vec::new(), + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }); + } + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting grouped_batch_put_done" + .to_string(), + })); + } + let req = MsgPack { + serialize_part: GroupedBatchPutDoneReq { + items, + atomic_group_lens, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp = self + .rpc_caller_grouped_batch_put_done + .call_with_transport_policy( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(resp.serialize_part) + } + + /// 完成 Put 操作，提交数据（inner，无监控） + pub async fn put_end_inner( + &self, + key: &str, + put_id: PutIDForAKey, + lease_id: Option, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_end".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutDoneReq { + key: key.to_string(), + put_id, + lease_id, + committed_slot: None, + publish_local_cache: false, + atomic_group: None, + }, + raw_bytes: Vec::new(), + }; + + // 获取 master 节点 ID + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let rpc_started_at = Instant::now(); + + // 调用 RPC + let resp = self + .rpc_caller_put_done + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 0, + ) + .await + .map_err(KvError::from)?; + if let Err(e) = crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + ) { + return Err(e); + } + Ok(PutEndStats { + master_put_end_rpc_us: duration_to_i64_us(rpc_started_at.elapsed()), + master_put_end_server_us: resp.serialize_part.server_process_us, + }) + } + + pub async fn put_end_inner_with_local_cache_publish( + &self, + key: &str, + put_id: PutIDForAKey, + lease_id: Option, + publish_local_cache: bool, + ) -> KvResult { + if !self.view.register_shutdown_poller().is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "ClientKvApi is shutting down; rejecting put_end".to_string(), + })); + } + let req = MsgPack { + serialize_part: PutDoneReq { + key: key.to_string(), + put_id, + lease_id, + committed_slot: None, + publish_local_cache, + atomic_group: None, + }, + raw_bytes: Vec::new(), + }; + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let rpc_started_at = Instant::now(); + let resp = self + .rpc_caller_put_done + .call( + self.view.p2p_module(), + master_node_id.into(), + req, + Some(std::time::Duration::from_secs(60)), + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + Ok(PutEndWithLocalCachePublish { + stats: PutEndStats { + master_put_end_rpc_us: duration_to_i64_us(rpc_started_at.elapsed()), + master_put_end_server_us: resp.serialize_part.server_process_us, + }, + local_cache_holder_id: resp.serialize_part.local_cache_holder_id, + }) + } + + /// 完成 Put 操作，提交数据（带监控）：适配 external 路径，统一聚合 t1..t4 + pub async fn put_end( + &self, + key: &str, + put_id: PutIDForAKey, + lease_id: Option, + ) -> KvResult { + let metrics = self.metrics_handle(); + let client_id = self.client_id_str(); + let node_role = self.node_role(); + + let end_stats = match self.put_end_inner(key, put_id, lease_id).await { + Ok(stats) => stats, + Err(e) => { + // on error, emit end error using pending info if exists, then cleanup + crate::observe_kvope::obe_put_end_error_from_pending( + &metrics, &client_id, &node_role, put_id, + ); + return Err(e); + } + }; + + // record end_handle to pending before aggregation + metrics.pending_put_set_end_handle(put_id, end_stats.master_put_end_server_us); + + // success: aggregate with pending timestamps; this also clears pending + crate::observe_kvope::obe_put_done_success_from_pending( + &metrics, &client_id, &node_role, key, put_id, 0, + ); + Ok(end_stats) + } + + pub async fn put_end_with_local_cache_publish( + &self, + key: &str, + put_id: PutIDForAKey, + lease_id: Option, + publish_local_cache: bool, + ) -> KvResult { + let metrics = self.metrics_handle(); + let client_id = self.client_id_str(); + let node_role = self.node_role(); + + let end = match self + .put_end_inner_with_local_cache_publish(key, put_id, lease_id, publish_local_cache) + .await + { + Ok(end) => end, + Err(e) => { + crate::observe_kvope::obe_put_end_error_from_pending( + &metrics, &client_id, &node_role, put_id, + ); + return Err(e); + } + }; + + metrics.pending_put_set_end_handle(put_id, end.stats.master_put_end_server_us); + crate::observe_kvope::obe_put_done_success_from_pending( + &metrics, &client_id, &node_role, key, put_id, 0, + ); + Ok(end) + } +} + +const OWNER_LOCAL_PUBLISH_RETRY_INITIAL: Duration = Duration::from_millis(25); +const OWNER_LOCAL_PUBLISH_RETRY_MAX: Duration = Duration::from_secs(1); + +#[cfg(test)] +mod owner_hot_replica_policy_tests { + use super::{ + OwnerLocalPublishItem, complete_owner_local_publish_group_lens, + owner_local_publish_atomic_batch_complete, owner_local_publish_starts_remote, + }; + use crate::master_kv_router::msg_pack::{ + PutAtomicGroup, PutAtomicGroupMember, PutDoneCommittedSlot, + }; + + fn publish_item( + key: &str, + put_id: (u64, u32), + atomic_group: Option, + ) -> OwnerLocalPublishItem { + OwnerLocalPublishItem { + key: key.to_string(), + put_id, + value_len: 1, + lease_id: None, + committed_slot: PutDoneCommittedSlot::default(), + make_replica_task: false, + remote_replica_admitted: false, + preferred_sub_cluster: None, + atomic_group, + } + } + + #[test] + fn grouped_publish_partition_requires_a_complete_ordered_group() { + let group = PutAtomicGroup { + members: vec![ + PutAtomicGroupMember { + key: "a".to_string(), + put_id: (1, 0), + }, + PutAtomicGroupMember { + key: "b".to_string(), + put_id: (1, 1), + }, + PutAtomicGroupMember { + key: "c".to_string(), + put_id: (1, 2), + }, + ], + }; + let complete = vec![ + publish_item("a", (1, 0), Some(group.clone())), + publish_item("b", (1, 1), Some(group.clone())), + publish_item("c", (1, 2), Some(group.clone())), + publish_item("single", (2, 0), None), + ]; + assert_eq!( + complete_owner_local_publish_group_lens(&complete), + Some(vec![3, 1]) + ); + + let partial = vec![ + publish_item("a", (1, 0), Some(group.clone())), + publish_item("b", (1, 1), Some(group)), + ]; + assert_eq!(complete_owner_local_publish_group_lens(&partial), None); + } + + #[test] + fn local_ssd_content_selection_does_not_bypass_remote_resource_admission() { + let mut item = publish_item("key", (1, 0), None); + item.make_replica_task = true; + assert!(!owner_local_publish_starts_remote(&item)); + + item.remote_replica_admitted = true; + assert!(owner_local_publish_starts_remote(&item)); + + item.make_replica_task = false; + assert!(!owner_local_publish_starts_remote(&item)); + } + + #[test] + fn hot_admission_waits_for_every_atomic_group_member_to_publish() { + let group = PutAtomicGroup { + members: vec![ + PutAtomicGroupMember { + key: "a".to_string(), + put_id: (1, 0), + }, + PutAtomicGroupMember { + key: "b".to_string(), + put_id: (1, 1), + }, + ], + }; + let items = vec![ + publish_item("a", (1, 0), Some(group.clone())), + publish_item("b", (1, 1), Some(group)), + publish_item("single", (2, 0), None), + ]; + let partial = vec![&items[0]]; + assert!(!owner_local_publish_atomic_batch_complete( + &items[0], &partial + )); + let complete = vec![&items[0], &items[1]]; + assert!(owner_local_publish_atomic_batch_complete( + &items[0], &complete + )); + assert!(owner_local_publish_atomic_batch_complete( + &items[1], &complete + )); + assert!(owner_local_publish_atomic_batch_complete( + &items[2], &partial + )); + } +} + +async fn start_replica_append_with_retry( + inner: &ClientKvApiInner, + op: &OwnerRemotePutSharedOp, + len: u32, +) -> KvResult { + let request = op.request(); + inner + .put_append_start( + &op.key, + op.put_id, + len, + request.preferred_sub_cluster.as_deref(), + request.protect_source_on_remote_complete, + ) + .await +} + +struct RemotePutTarget { + operation_id: u64, + node_id: String, + target_offset: u64, + target_base_addr: u64, + len: u64, +} + +fn owner_source_eviction_identity( + victim: &OwnerSourceEvictionVictim, +) -> super::OwnerHotReplicaIdentity { + super::OwnerHotReplicaIdentity { + key: victim.key.clone(), + put_time_ms: victim.put_id.0, + put_version: victim.put_id.1, + } +} + +fn owner_source_eviction_victim( + identity: &super::OwnerHotReplicaIdentity, + memory_info: &crate::memholder::MemoryInfo, + ssd_enabled: bool, +) -> Option { + let (slot_size, grant_id, slot_index) = memory_info.local_reserve_resident_slot_ref()?; + Some(OwnerSourceEvictionVictim { + key: identity.key.clone(), + put_id: (identity.put_time_ms, identity.put_version), + backing: OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size, + }, + ssd_backing_len: None, + ssd_policy: if ssd_enabled { + OwnerSourceSsdPolicy::SelectLastLive + } else { + OwnerSourceSsdPolicy::Drop + }, + }) +} + +fn finish_owner_source_selection( + inner: &ClientKvApiInner, + victim: &OwnerSourceEvictionVictim, + restore_current: bool, + reason: &str, +) { + let identity = owner_source_eviction_identity(victim); + if let Some(debt) = inner.owner_hot_remove_source_selection_debt(&identity) { + debt.release(); + inner + .owner_hot_counters + .source_evict_restored_members + .fetch_add(1, Ordering::Relaxed); + } + inner.owner_hot_retry_queue.remove(&identity); + if !restore_current { + return; + } + inner.owner_hot_restore_source_selection(&identity); + if inner.owner_hot_admit_published_committed(&victim.key, victim.put_id) { + tracing::warn!( + key = victim.key, + put_time_ms = victim.put_id.0, + put_version = victim.put_id.1, + reason, + "restored current source to owner-hot after source eviction did not enter reclaim" + ); + } +} + +fn schedule_owner_source_eviction_retry( + inner: &ClientKvApiInner, + mut event: OwnerHotEvictionEvent, + victim: Arc, + reason: &'static str, +) { + event.retry = true; + event.source_eviction_victim = Some(victim); + inner.owner_hot_retry_queue.schedule(event, reason); +} + +fn prepare_owner_source_eviction_event( + inner: &ClientKvApiInner, + mut event: OwnerHotEvictionEvent, +) -> Option<(OwnerHotEvictionEvent, Arc)> { + let trigger = super::OwnerHotReplicaIdentity { + key: event.key.clone(), + put_time_ms: event.put_id.0, + put_version: event.put_id.1, + }; + + if event.retry { + let _ = inner.owner_hot_retry_queue.take_for_inflight(&trigger); + } + if let Some(victim) = event.source_eviction_victim.clone() { + if inner + .owner_source_eviction_selected + .contains_key(&owner_source_eviction_identity(&victim)) + { + return Some((event, victim)); + } + // Commit/invalidation may win the race with a due retry. + finish_owner_source_selection(inner, &victim, true, "retry victim lost selected identity"); + event.selection_debt.release(); + return None; + } + + if inner.owner_source_eviction_selected.contains_key(&trigger) { + event.selection_debt.release(); + inner + .owner_hot_counters + .victim_duplicates + .fetch_add(1, Ordering::Relaxed); + return None; + } + + let (resolved_trigger, source) = match inner.owner_hot_prepare_eviction(&event) { + OwnerHotEvictionPreparation::Ready { trigger, source } => (trigger, source), + OwnerHotEvictionPreparation::RetryableReclaimFence => { + inner.owner_hot_retry_queue.schedule( + event, + "owner reclaim fence busy before exact source selection", + ); + return None; + } + OwnerHotEvictionPreparation::TemporarilyPinned => { + // This event was removed from Moka, but its source is still + // serving a local reader. Do not hand it to the master's reclaim + // loop and do not retain projected selection credit. Re-admission + // refreshes the trigger's recency so the next pressure kick can + // choose a different, currently reclaimable victim. + event.selection_debt.release(); + inner + .owner_hot_counters + .skipped_active_holders + .fetch_add(1, Ordering::Relaxed); + let restored = inner.owner_hot_admit_published_committed(&event.key, event.put_id); + tracing::debug!( + key = event.key, + put_time_ms = event.put_id.0, + put_version = event.put_id.1, + restored, + "owner pressure eviction skipped an actively held source" + ); + return None; + } + OwnerHotEvictionPreparation::Obsolete => { + event.selection_debt.release(); + inner + .owner_hot_counters + .source_evict_obsolete + .fetch_add(1, Ordering::Relaxed); + return None; + } + }; + debug_assert_eq!(resolved_trigger, trigger); + + let Some(victim) = + owner_source_eviction_victim(&trigger, source.as_ref(), inner.ssd_storage.is_some()) + else { + event.selection_debt.release(); + inner + .owner_hot_counters + .victim_invalid_backing + .fetch_add(1, Ordering::Relaxed); + tracing::error!( + key = event.key, + put_time_ms = event.put_id.0, + put_version = event.put_id.1, + "owner-hot selected a source without exact CommittedSlot backing" + ); + return None; + }; + let victim = Arc::new(victim); + + match inner.owner_hot_install_source_selection_fence(&trigger, &source) { + OwnerHotSelectionFenceOutcome::Fenced => {} + OwnerHotSelectionFenceOutcome::Retryable => { + inner + .owner_hot_retry_queue + .schedule(event, "owner source selection fence is temporarily busy"); + return None; + } + OwnerHotSelectionFenceOutcome::TemporarilyPinned => { + event.selection_debt.release(); + inner + .owner_hot_counters + .skipped_active_holders + .fetch_add(1, Ordering::Relaxed); + inner.owner_hot_admit_published_committed(&event.key, event.put_id); + return None; + } + OwnerHotSelectionFenceOutcome::Obsolete => { + event.selection_debt.release(); + inner + .owner_hot_counters + .source_evict_obsolete + .fetch_add(1, Ordering::Relaxed); + return None; + } + } + + if inner.owner_source_eviction_selected.contains_key(&trigger) { + inner.owner_hot_restore_source_selection(&trigger); + event.selection_debt.release(); + inner + .owner_hot_counters + .victim_duplicates + .fetch_add(1, Ordering::Relaxed); + return None; + } + + if !inner.owner_hot_install_source_selection_debt(trigger.clone(), event.selection_debt.clone()) + { + inner.owner_hot_restore_source_selection(&trigger); + event.selection_debt.release(); + return None; + } + event.source_eviction_victim = Some(victim.clone()); + Some((event, victim)) +} + +async fn finish_owner_source_eviction_result( + inner: &ClientKvApiInner, + operation_id: u64, + index: usize, + event: OwnerHotEvictionEvent, + victim: Arc, + result: OwnerSourceEvictionVictimResp, +) { + if victim.ssd_policy == OwnerSourceSsdPolicy::Persisted + && !result.ssd_backing_committed + && matches!( + result.outcome, + OwnerSourceEvictionOutcome::Completed + | OwnerSourceEvictionOutcome::Stale + | OwnerSourceEvictionOutcome::RejectedNotEvictable + ) + { + inner + .discard_local_ssd_replica(&victim.key, victim.put_id) + .await; + } + match result.outcome { + OwnerSourceEvictionOutcome::Accepted | OwnerSourceEvictionOutcome::AlreadyInProgress => { + inner + .owner_hot_counters + .source_evict_handoff_members + .fetch_add(1, Ordering::Relaxed); + } + OwnerSourceEvictionOutcome::Completed => { + let epoch = owner_source_eviction_epoch(operation_id, index); + match super::reclaim::complete_owner_source_eviction(inner, &victim, epoch) { + Ok(()) => { + inner + .owner_hot_counters + .source_evict_handoff_members + .fetch_add(1, Ordering::Relaxed); + } + Err(detail) => { + tracing::warn!( + key = victim.key, + put_time_ms = victim.put_id.0, + put_version = victim.put_id.1, + detail, + "master deleted source route but owner slot release must retry" + ); + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "owner slot release after master direct-delete is temporarily busy", + ); + } + } + } + OwnerSourceEvictionOutcome::SsdCandidate => { + tracing::error!( + key = victim.key, + policy = ?victim.ssd_policy, + "master returned an SSD candidate outside the selection phase" + ); + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "unexpected repeated SSD candidate response", + ); + } + OwnerSourceEvictionOutcome::RetryableBusy | OwnerSourceEvictionOutcome::Unspecified => { + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "master source reclaim is temporarily busy", + ); + } + OwnerSourceEvictionOutcome::Stale => { + finish_owner_source_selection( + inner, + &victim, + true, + "master rejected stale source identity", + ); + } + OwnerSourceEvictionOutcome::RejectedNotEvictable => { + tracing::error!( + key = victim.key, + detail = result.detail, + "owner selected a source victim that master declared non-evictable" + ); + finish_owner_source_selection( + inner, + &victim, + true, + "master rejected non-evictable source victim", + ); + } + } +} + +async fn submit_owner_source_eviction_decisions( + inner: &ClientKvApiInner, + prepared: Vec<(OwnerHotEvictionEvent, Arc)>, + rpc_failure_reason: &'static str, + response_length_reason: &'static str, + response_identity_reason: &'static str, +) { + if prepared.is_empty() { + return; + } + let response = inner + .batch_evict_owner_source( + prepared + .iter() + .map(|(_, victim)| victim.as_ref().clone()) + .collect(), + ) + .await; + let Ok(response) = response else { + for (event, victim) in prepared { + schedule_owner_source_eviction_retry(inner, event, victim, rpc_failure_reason); + } + return; + }; + if response.victims.len() != prepared.len() { + for (event, victim) in prepared { + schedule_owner_source_eviction_retry(inner, event, victim, response_length_reason); + } + return; + } + let operation_id = response.operation_id; + for (index, ((event, victim), result)) in prepared.into_iter().zip(response.victims).enumerate() + { + if result.victim_index != u32::try_from(index).unwrap_or(u32::MAX) { + schedule_owner_source_eviction_retry(inner, event, victim, response_identity_reason); + continue; + } + finish_owner_source_eviction_result(inner, operation_id, index, event, victim, result) + .await; + } +} + +async fn reclaim_admission_drops_before_ssd_persist( + drop_reclaim: DropFuture, + persist: PersistFn, +) -> T +where + DropFuture: Future, + PersistFn: FnOnce() -> PersistFuture, + PersistFuture: Future, +{ + drop_reclaim.await; + persist().await +} + +#[cfg(test)] +mod owner_ssd_fast_drop_tests { + use super::reclaim_admission_drops_before_ssd_persist; + use std::sync::{Arc, Mutex}; + + #[limit_thirdparty::tokio::test] + async fn admission_drops_finish_before_ssd_persist_starts() { + let order = Arc::new(Mutex::new(Vec::new())); + let drop_order = order.clone(); + let persist_order = order.clone(); + + let result = reclaim_admission_drops_before_ssd_persist( + async move { + drop_order.lock().unwrap().push("drop-start"); + limit_thirdparty::tokio::task::yield_now().await; + drop_order.lock().unwrap().push("drop-finished"); + }, + move || async move { + let mut order = persist_order.lock().unwrap(); + assert_eq!(order.last(), Some(&"drop-finished")); + order.push("persist-start"); + 17 }, - /*dbg_addr_summary=*/ true, - None, ) - .await + .await; + + assert_eq!(result, 17); + assert_eq!( + order.lock().unwrap().as_slice(), + ["drop-start", "drop-finished", "persist-start"] + ); } +} - /// Transfer data by offsets with instrumentation for external/owner callers. - /// Records transfer latency (t2..t3) and emits tsbuckets pulses. - pub async fn put_transfer( - &self, - key: &str, - put_id: PutIDForAKey, - src_offset: u64, - target_offset: u64, - len: u64, - peer_id: Option, - target_base_addr: Option, - ) -> KvResult { - let metrics = self.metrics_handle(); - let client_id = self.client_id_str(); - let node_role = self.node_role(); +async fn process_owner_source_eviction_events( + view: &ClientKvApiView, + events: Vec, +) { + let inner = view.client_kv_api().inner(); + let prepared = events + .into_iter() + .filter_map(|event| prepare_owner_source_eviction_event(inner, event)) + .collect::>(); + if prepared.is_empty() { + return; + } + // The first pass never touches SSD. Master deletes sources that already + // have another live backing and returns only exact last-copy candidates. + let response = inner + .batch_evict_owner_source( + prepared + .iter() + .map(|(_, victim)| victim.as_ref().clone()) + .collect(), + ) + .await; + let Ok(response) = response else { + for (event, victim) in prepared { + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "owner source-eviction RPC failed", + ); + } + return; + }; + if response.victims.len() != prepared.len() { + for (event, victim) in prepared { + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "owner source-eviction response length mismatch", + ); + } + return; + } - // owner/external inner is stable after construction; base_addr must exist - let base_addr = self - .view - .client_seg_pool() - .cpu_mem_read_guard() - .await - .unwrap() - .allocated_addr; - let abs_src = base_addr + src_offset; - let abs_target = if peer_id.is_some() { - let Some(tb) = target_base_addr else { - // propagate as Unreachable: invalid remote target context from distributed input - let err = crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { - rpc_input_json: format!( - "missing target_base_addr while peer_id present; src_off={:#x}, tgt_off={:#x}", - src_offset, target_offset - ), - }, - ); - return Err(err); - }; - tb + target_offset + let retryable = response + .victims + .iter() + .filter(|result| { + matches!( + result.outcome, + OwnerSourceEvictionOutcome::RetryableBusy | OwnerSourceEvictionOutcome::Unspecified + ) + }) + .count(); + if retryable != 0 + && let Some((index, result)) = response.victims.iter().enumerate().find(|(_, result)| { + matches!( + result.outcome, + OwnerSourceEvictionOutcome::RetryableBusy | OwnerSourceEvictionOutcome::Unspecified + ) + }) + { + let victim = &prepared[index].1; + tracing::warn!( + operation_id = response.operation_id, + victims = response.victims.len(), + retryable, + first_victim_index = result.victim_index, + key = %victim.key, + put_time_ms = victim.put_id.0, + put_version = victim.put_id.1, + detail = %result.detail, + "owner source direct-delete batch has retryable victims" + ); + } + + let first_operation_id = response.operation_id; + let mut candidates = Vec::new(); + for (index, ((event, victim), result)) in prepared + .into_iter() + .zip(response.victims.into_iter()) + .enumerate() + { + if result.victim_index != u32::try_from(index).unwrap_or(u32::MAX) { + schedule_owner_source_eviction_retry( + inner, + event, + victim, + "owner source-eviction response identity mismatch", + ); + continue; + } + if result.outcome == OwnerSourceEvictionOutcome::SsdCandidate + && victim.ssd_policy == OwnerSourceSsdPolicy::SelectLastLive + { + candidates.push((event, victim)); + continue; + } + finish_owner_source_eviction_result( + inner, + first_operation_id, + index, + event, + victim, + result, + ) + .await; + } + + if candidates.is_empty() { + return; + } + let candidate_lengths = candidates + .iter() + .map(|(_, victim)| owner_reclaim_backing_len(&victim.backing)) + .collect::>(); + let admission = inner + .ssd_storage + .as_ref() + .map(|store| store.admit_owner_write_candidates(&candidate_lengths)) + .unwrap_or_else(|| vec![false; candidates.len()]); + let candidate_items = candidates.len(); + let mut dropped = Vec::new(); + let mut admitted_candidates = Vec::new(); + for ((mut event, victim), is_admitted) in candidates.into_iter().zip(admission) { + let mut request = victim.as_ref().clone(); + request.ssd_backing_len = None; + request.ssd_policy = OwnerSourceSsdPolicy::Drop; + let request = Arc::new(request); + event.source_eviction_victim = Some(request.clone()); + if is_admitted { + admitted_candidates.push((event, request)); } else { - base_addr + target_offset + dropped.push((event, request)); + } + } + + // Candidates rejected by the byte budget must release their DRAM slots + // before the admitted durability batch starts. Otherwise one slow SSD + // flush makes every intentional Drop in the same pressure batch wait, + // recreating a reclaim backlog even though no SSD queue exists. + let dropped_items = dropped.len(); + let admitted_items = admitted_candidates.len(); + // Both early Put and last-backing fallback enter the same per-generation + // local_ssd_put flight. The fallback starts only after explicit admission + // Drops have released their DRAM slots. + let operations = reclaim_admission_drops_before_ssd_persist( + submit_owner_source_eviction_decisions( + inner, + dropped, + "owner SSD admission-drop direct-delete RPC failed", + "owner SSD admission-drop response length mismatch", + "owner SSD admission-drop response identity mismatch", + ), + || async { + inner.start_owner_local_ssd_puts( + admitted_candidates + .iter() + .map(|(_, victim)| OwnerLocalSsdPutCandidate { + key: victim.key.clone(), + put_id: victim.put_id, + value_len: owner_reclaim_backing_len(&victim.backing), + selected_victim: Some(victim.clone()), + }) + .collect(), + true, + ) + }, + ) + .await; + let mut admitted_decisions = Vec::with_capacity(admitted_candidates.len()); + let mut persisted_items = 0usize; + for ((mut event, victim), operation) in admitted_candidates.into_iter().zip(operations) { + let outcome = match operation { + Some(operation) => operation.wait().await, + None => OwnerLocalSsdPutOutcome::Failed, }; + if matches!( + outcome, + OwnerLocalSsdPutOutcome::Published | OwnerLocalSsdPutOutcome::AlreadyPresent + ) { + persisted_items += 1; + } + // Publication, when successful, already attached SSD to the exact + // route. Direct-delete now removes only memory. Failed/Drop outcomes + // retain the cache's existing fail-open eviction behavior. + event.source_eviction_victim = Some(victim.clone()); + admitted_decisions.push((event, victim)); + } + tracing::info!( + candidates = candidate_items, + admitted = admitted_items, + persisted = persisted_items, + admission_dropped = dropped_items, + persist_dropped = admitted_items.saturating_sub(persisted_items), + "owner last-backing SSD admission completed without queueing excess victims" + ); - // Local placement can resolve to src==target, which means the payload is already in-place. - // Skip the transfer-engine hop for this no-op path to avoid paying an extra fixed cost. - if peer_id.is_none() && abs_src == abs_target { - tracing::debug!( - "put_transfer local no-op: key={}, put_id=({},{}) src==target {:#x}, len={}", - key, - put_id.0, - put_id.1, - abs_target, - len + submit_owner_source_eviction_decisions( + inner, + admitted_decisions, + "owner SSD persisted decision direct-delete RPC failed", + "owner SSD persisted decision response length mismatch", + "owner SSD persisted decision response identity mismatch", + ) + .await; +} + +pub fn spawn_owner_source_eviction_dispatcher( + view: ClientKvApiView, + mut rx: tokio::sync::ampsc::UnboundedReceiver, +) { + const ORPHAN_MERGE_WINDOW: Duration = Duration::from_millis(2); + + let view_task = view.clone(); + let _ = view.spawn("owner_source_eviction_dispatcher", async move { + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + let mut events = Vec::new(); + let mut pressure_open = false; + loop { + let dispatch = if events.is_empty() || pressure_open { + tokio::select! { + _ = shutdown_waiter.wait() => { + tracing::info!("owner source-eviction dispatcher stopping due to shutdown"); + break; + } + dispatch = rx.recv() => dispatch, + } + } else { + tokio::select! { + _ = shutdown_waiter.wait() => { + tracing::info!("owner source-eviction dispatcher stopping due to shutdown"); + break; + } + dispatch = rx.recv() => dispatch, + _ = tokio::time::sleep(ORPHAN_MERGE_WINDOW) => { + process_owner_source_eviction_events( + &view_task, + std::mem::take(&mut events), + ) + .await; + continue; + } + } + }; + let Some(dispatch) = dispatch else { + if !events.is_empty() { + process_owner_source_eviction_events(&view_task, std::mem::take(&mut events)) + .await; + } + break; + }; + match dispatch { + super::OwnerHotEvictionDispatch::Victim(event) => events.push(event), + super::OwnerHotEvictionDispatch::BeginPressure { requested_bytes } => { + if pressure_open { + tracing::error!( + requested_bytes, + "nested owner pressure selection batch is not allowed" + ); + } + if !events.is_empty() { + process_owner_source_eviction_events( + &view_task, + std::mem::take(&mut events), + ) + .await; + } + pressure_open = true; + } + super::OwnerHotEvictionDispatch::EndPressure { selected_bytes } => { + if !pressure_open { + tracing::warn!( + selected_bytes, + "owner pressure selection ended without a matching begin marker" + ); + } + pressure_open = false; + process_owner_source_eviction_events(&view_task, std::mem::take(&mut events)) + .await; + } + super::OwnerHotEvictionDispatch::Flush => { + if !pressure_open && !events.is_empty() { + process_owner_source_eviction_events( + &view_task, + std::mem::take(&mut events), + ) + .await; + } + } + } + } + }); +} + +pub fn spawn_owner_hot_retry_actor(view: ClientKvApiView) { + const RETRY_EMIT_BATCH: usize = 128; + const RETRY_POLL_INTERVAL: Duration = Duration::from_millis(25); + + let view_task = view.clone(); + let _ = view.spawn("owner_hot_retry_actor", async move { + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + let retry_queue = view_task + .client_kv_api() + .inner() + .owner_hot_retry_queue + .clone(); + let notify = retry_queue.notify.clone(); + let mut tick = tokio::time::interval(RETRY_POLL_INTERVAL); + loop { + tokio::select! { + _ = shutdown_waiter.wait() => return, + _ = tick.tick() => {}, + _ = notify.notified() => {}, + } + let events = retry_queue.take_due_batch(Instant::now(), RETRY_EMIT_BATCH); + if events.is_empty() { + continue; + } + let inner = view_task.client_kv_api().inner(); + for event in events { + if let Err(err) = inner + .owner_hot_eviction_tx + .send(super::OwnerHotEvictionDispatch::Victim(event)) + { + let super::OwnerHotEvictionDispatch::Victim(event) = err.0 else { + unreachable!("the retry actor only sends victim events") + }; + retry_queue.schedule(event, "retry dispatcher closed"); + return; + } + inner + .owner_hot_counters + .source_evict_retry_emitted + .fetch_add(1, Ordering::Relaxed); + } + if inner + .owner_hot_eviction_tx + .send(super::OwnerHotEvictionDispatch::Flush) + .is_err() + { + return; + } + } + }); +} + +async fn run_owner_local_ssd_put_batch( + view: ClientKvApiView, + leaders: Vec, + permit: crate::kv_ssd_storage::KvSsdPersistBatchPermit, +) { + let inner = view.client_kv_api().inner(); + let mut holders = Vec::with_capacity(leaders.len()); + let mut sources = Vec::with_capacity(leaders.len()); + for leader in &leaders { + let holder = Arc::new(UserMemHolder::new( + leader.memory_info.clone(), + inner.get_or_init_all_memholder_refcount(), + UserMemHolderExposeKind::SegPtr, + )); + sources.push(crate::kv_ssd_storage::KvSsdPersistSource { + key: leader.op.key.clone(), + put_id: leader.op.put_id, + addr: holder.memory_info.addr, + len: u64::from(holder.memory_info.len), + }); + holders.push(holder); + } + + let copies = match inner.copy_local_kvs_for_ssd(&sources).await { + Ok(copies) => copies, + Err(err) => { + drop(holders); + tracing::warn!( + items = leaders.len(), + error = %err, + "owner local SSD source copy validation failed" ); - return Ok(TransferBreakdown { - local_noop: true, - ..TransferBreakdown::default() - }); - } else { - let breakdown = self - .view - .client_transfer_engine() - .transfer_data_no_copy(peer_id.clone(), false, abs_src, abs_target, len, None) - .await?; - tracing::debug!( - "put_transfer breakdown: key={}, put_id=({},{}) fast_path={} local_noop={} remote_transfer={} submit_blocking_us={} create_xfer_req_us={} post_xfer_req_us={} poll_wait_us={} poll_iters={}", - key, - put_id.0, - put_id.1, - breakdown.used_fast_path, - breakdown.local_noop, - breakdown.remote_transfer, - breakdown.submit_blocking_us, - breakdown.create_xfer_req_us, - breakdown.post_xfer_req_us, - breakdown.poll_wait_us, - breakdown.poll_iters + for leader in leaders { + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + return; + } + }; + // Durability now reads only owner-owned copies. The source Moka pins may + // be released while this flight remains installed through route publish. + drop(holders); + + let persisted = match inner.persist_copied_local_kvs_to_ssd(permit, copies).await { + Ok(persisted) => persisted, + Err(err) => { + tracing::warn!( + items = leaders.len(), + error = %err, + "owner local SSD durability setup failed" ); - tracing::debug!( - "put_transfer success: key={}, put_id=({},{}) src_off={:#x}, tgt_off={:#x}, len={}, peer_id={:?}", - key, - put_id.0, - put_id.1, - src_offset, - target_offset, - len, - peer_id + for leader in leaders { + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + return; + } + }; + + let mut publish = Vec::new(); + for ((leader, source), outcome) in leaders + .into_iter() + .zip(sources.into_iter()) + .zip(persisted.into_iter()) + { + match outcome { + Ok(Some(guard)) => { + publish.push((leader, source, guard)); + } + Ok(None) => { + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Dropped); + } + Err(err) => { + tracing::warn!( + key = leader.op.key, + put_time_ms = leader.op.put_id.0, + put_version = leader.op.put_id.1, + error = %err, + "owner local SSD persist failed" + ); + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + } + } + if publish.is_empty() { + return; + } + + let response = match inner + .batch_publish_owner_ssd( + publish + .iter() + .map(|(leader, source, _)| OwnerSsdPublishItem { + key: leader.op.key.clone(), + put_id: leader.op.put_id, + len: source.len, + }) + .collect(), + ) + .await + { + Ok(response) => response, + Err(err) => { + tracing::warn!( + items = publish.len(), + error = %err, + "owner local SSD publication was uncertain; retaining durable bytes for replay" + ); + for (leader, _, _guard) in publish { + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + return; + } + }; + if response.items.len() != publish.len() { + tracing::warn!( + requested = publish.len(), + received = response.items.len(), + "owner local SSD publication response length mismatch; retaining bytes for replay" + ); + for (leader, _, _guard) in publish { + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + } + return; + } + + for ((leader, _, _guard), result) in publish.into_iter().zip(response.items) { + if result.key != leader.op.key || result.put_id != leader.op.put_id { + tracing::warn!( + expected_key = leader.op.key, + expected_put_time_ms = leader.op.put_id.0, + expected_put_version = leader.op.put_id.1, + response_key = result.key, + response_put_time_ms = result.put_id.0, + response_put_version = result.put_id.1, + "owner local SSD publication identity mismatch; retaining bytes for replay" ); + inner.finish_owner_local_ssd_put(&leader.op, OwnerLocalSsdPutOutcome::Failed); + continue; + } + let outcome = match result.outcome { + OwnerSsdPublishOutcome::Published => OwnerLocalSsdPutOutcome::Published, + OwnerSsdPublishOutcome::AlreadyPresent => OwnerLocalSsdPutOutcome::AlreadyPresent, + OwnerSsdPublishOutcome::Obsolete => { + inner + .discard_local_ssd_replica(&leader.op.key, leader.op.put_id) + .await; + OwnerLocalSsdPutOutcome::Obsolete + } + OwnerSsdPublishOutcome::Rejected => { + inner + .discard_local_ssd_replica(&leader.op.key, leader.op.put_id) + .await; + OwnerLocalSsdPutOutcome::Failed + } + OwnerSsdPublishOutcome::RetryableBusy | OwnerSsdPublishOutcome::Unspecified => { + OwnerLocalSsdPutOutcome::Failed + } + }; + inner.finish_owner_local_ssd_put(&leader.op, outcome); + } +} - // Emit transfer stage success and tsbuckets pulse (computes t2/t3 using pending) - crate::observe_kvope::obe_put_transfer_success( - &metrics, &client_id, &node_role, key, len, put_id, +async fn fail_owner_remote_put( + view: &ClientKvApiView, + op: Arc, + operation_id: Option, + reason: &'static str, +) { + let inner = view.client_kv_api().inner(); + if let Some(operation_id) = operation_id { + if let Err(revoke_err) = inner + .put_append_revoke(&op.key, op.put_id, operation_id) + .await + { + tracing::warn!( + "owner remote Put revoke failed after {}: key={} put_id=({},{}) operation_id={} err={}", + reason, + op.key, + op.put_id.0, + op.put_id.1, + operation_id, + revoke_err ); - return Ok(breakdown); } - #[allow(unreachable_code)] - Ok(TransferBreakdown::default()) } + inner.finish_owner_remote_put(&op, OwnerRemotePutOutcome::Failed); +} - /// 开始 Put 操作，分配存储空间 - pub async fn put_start_with_source_node( - &self, - key: &str, - len: u32, - reject_if_inflight_same_key: bool, - preferred_sub_cluster: Option<&str>, - source_node_id: Option, - ) -> KvResult<(PutStartResp, i64)> { - let req = MsgPack { - serialize_part: PutStartReq { - key: key.to_string(), - len: len as u64, - reject_if_inflight_same_key, - preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), - source_node_id, - }, - raw_bytes: Vec::new(), - }; +async fn run_owner_remote_put( + view: ClientKvApiView, + op: Arc, + holder: Arc, +) { + let inner = view.client_kv_api().inner(); + let src_offset = holder.memory_info().offset; + let len = holder.get_length() as u64; + let len_u32 = match u32::try_from(len) { + Ok(len_u32) => len_u32, + Err(_) => { + tracing::warn!( + "owner remote Put length does not fit u32: key={} put_id=({},{}) len={}", + op.key, + op.put_id.0, + op.put_id.1, + len + ); + drop(holder); + fail_owner_remote_put(&view, op, None, "length does not fit u32").await; + return; + } + }; - let master_node_id = self - .view - .cluster_manager() - .find_or_wait_master_node() - .await?; - let rpc_started_at = Instant::now(); - let start_rpc_timestamp = Utc::now().timestamp_micros() as i64; - let resp = self - .rpc_caller_put_start - .call( - self.view.p2p_module(), - master_node_id.into(), - req, - Some(std::time::Duration::from_secs(60)), - 2, - ) - .await - .map_err(|e| KvError::P2p(e))?; - let end_rpc_timestamp = Utc::now().timestamp_micros() as i64; - let ser = resp.serialize_part.clone(); - if crate::rpcresp_kvresult_convert::try_from_code(ser.error_code, ser.error_json.clone()) - .is_ok() - { - let metrics = self.metrics_handle(); - metrics.pending_put_insert( - ser.put_id, - key.to_string(), - len as u64, - start_rpc_timestamp, - end_rpc_timestamp, - ser.server_process_us, + let append_start = match start_replica_append_with_retry(inner, &op, len_u32).await { + Ok(resp) => resp, + Err(err) => { + tracing::warn!( + "owner remote Put start failed: key={} put_id=({},{}) err={}", + op.key, + op.put_id.0, + op.put_id.1, + err ); + drop(holder); + fail_owner_remote_put(&view, op, None, "start error").await; + return; + } + }; + match append_start.outcome { + PutAppendStartOutcome::Scheduled => {} + PutAppendStartOutcome::AlreadySatisfied => { + drop(holder); + inner.finish_owner_remote_put(&op, OwnerRemotePutOutcome::AlreadySatisfied); + return; + } + PutAppendStartOutcome::Obsolete => { + drop(holder); + inner.finish_owner_remote_put(&op, OwnerRemotePutOutcome::Obsolete); + return; + } + PutAppendStartOutcome::RetryableNoSpace | PutAppendStartOutcome::Unspecified => { + tracing::debug!( + outcome = ?append_start.outcome, + "owner remote Put deferred: key={} put_id=({},{})", + op.key, + op.put_id.0, + op.put_id.1 + ); + drop(holder); + inner.finish_owner_remote_put(&op, OwnerRemotePutOutcome::Failed); + return; } - let rpc_latency_us = duration_to_i64_us(rpc_started_at.elapsed()); - Ok((ser, rpc_latency_us)) } - /// 开始 Put 操作，分配存储空间 - pub async fn put_start( - &self, - key: &str, - len: u32, - reject_if_inflight_same_key: bool, - preferred_sub_cluster: Option<&str>, - ) -> KvResult<(PutStartResp, i64)> { - self.put_start_with_source_node( - key, + let target = RemotePutTarget { + operation_id: append_start.operation_id, + node_id: append_start.node_id, + target_offset: append_start.target_addr - append_start.target_base_addr, + target_base_addr: append_start.target_base_addr, + len: append_start.len, + }; + if len != target.len { + tracing::warn!( + "owner remote Put length mismatch: key={} put_id=({},{}) src_len={} target_len={}", + op.key, + op.put_id.0, + op.put_id.1, len, - reject_if_inflight_same_key, - preferred_sub_cluster, - None, + target.len + ); + drop(holder); + fail_owner_remote_put(&view, op, Some(target.operation_id), "length mismatch").await; + return; + } + + inner.record_owner_remote_put_transfer(); + if let Err(err) = inner + .put_transfer( + &op.key, + op.put_id, + src_offset, + target.target_offset, + len, + Some(target.node_id), + Some(target.target_base_addr), ) .await + { + tracing::warn!( + "owner remote Put transfer failed: key={} put_id=({},{}) err={}", + op.key, + op.put_id.0, + op.put_id.1, + err + ); + drop(holder); + fail_owner_remote_put(&view, op, Some(target.operation_id), "transfer error").await; + return; } - /// 撤销 Put 操作，释放已分配的资源 - pub async fn put_revoke(&self, key: &str, put_id: PutIDForAKey) -> KvResult<()> { - if !self.view.register_shutdown_poller().is_running() { - return Err(KvError::Api(ApiError::SystemShutdown { - detail: "ClientKvApi is shutting down; rejecting put_revoke".to_string(), - })); + // The payload is no longer read after transfer. Keep the generation + // flight installed through Done so eviction cannot cross the completion + // protocol even though the physical source holder is now releasable. + drop(holder); + match inner + .put_append_done(&op.key, op.put_id, target.operation_id) + .await + { + Ok(resp) => { + let outcome = if resp.appended { + OwnerRemotePutOutcome::Published + } else { + OwnerRemotePutOutcome::AlreadySatisfied + }; + tracing::debug!( + "owner remote Put done: key={} put_id=({},{}) appended={}", + op.key, + op.put_id.0, + op.put_id.1, + resp.appended + ); + inner.finish_owner_remote_put(&op, outcome); } - let req = MsgPack { - serialize_part: PutRevokeReq { - key: key.to_string(), - put_id, - }, - raw_bytes: Vec::new(), - }; - - // 获取 master 节点 ID - let master_node_id = self - .view - .cluster_manager() - .find_or_wait_master_node() - .await?; + Err(err) => { + tracing::warn!( + "owner remote Put done failed: key={} put_id=({},{}) err={}", + op.key, + op.put_id.0, + op.put_id.1, + err + ); + fail_owner_remote_put(&view, op, Some(target.operation_id), "done error").await; + } + } +} - // 调用 RPC - let _resp = self - .rpc_caller_put_revoke - .call( - self.view.p2p_module(), - master_node_id.into(), - req, - Some(std::time::Duration::from_secs(60)), - 2, - ) +pub async fn handle_batch_enqueue_replica_tasks( + view: &ClientKvApiView, + req: MsgPack, +) -> MsgPack { + let inner = view.client_kv_api().inner(); + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let accepted = match inner + .ensure_remote_put(&item.key, item.put_id, None, false) .await - .map_err(KvError::from)?; - // cleanup pending stat if any - let _ = self.metrics_handle().pending_put_remove(&put_id); - Ok(()) + { + Ok(accepted) => accepted, + Err(err) => { + tracing::warn!( + "tier1 write-back enqueue failed: key={} put_id=({},{}) err={}", + item.key, + item.put_id.0, + item.put_id.1, + err + ); + false + } + }; + items.push(EnqueueReplicaTaskItemResp { + key: item.key, + put_id: item.put_id, + accepted, + }); } + MsgPack { + serialize_part: BatchEnqueueReplicaTaskResp { + items, + error_code: crate::rpcresp_kvresult_convert::msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} - /// 完成 Put 操作，提交数据（inner，无监控） - pub async fn put_end_inner( - &self, - key: &str, - put_id: PutIDForAKey, - lease_id: Option, - ) -> KvResult { - if !self.view.register_shutdown_poller().is_running() { - return Err(KvError::Api(ApiError::SystemShutdown { - detail: "ClientKvApi is shutting down; rejecting put_end".to_string(), - })); +pub fn spawn_owner_local_publish_dispatcher( + view: ClientKvApiView, + mut rx: tokio::sync::ampsc::Receiver, + max_inflight: usize, +) { + let view_task = view.clone(); + let _ = view.spawn("owner_local_publish_dispatcher", async move { + let max_inflight = max_inflight.max(1); + let semaphore = Arc::new(::tokio::sync::Semaphore::new(max_inflight)); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + let job = tokio::select! { + _ = shutdown_waiter.wait() => { + tracing::info!("owner_local_publish_dispatcher stopping due to shutdown signal"); + break; + } + job = rx.recv() => { + match job { + Some(job) => job, + None => break, + } + } + }; + let permit = match semaphore.clone().acquire_owned().await { + Ok(permit) => permit, + Err(err) => { + tracing::warn!( + "owner_local_publish_dispatcher semaphore closed; dropping job key_count={} err={}", + job.items.len(), + err + ); + break; + } + }; + let spawn_view = view_task.clone(); + let worker_view = view_task.clone(); + spawn_view.spawn("owner_local_publish_worker", async move { + let _permit = permit; + publish_owner_local_job(worker_view, job).await; + }); } - let req = MsgPack { - serialize_part: PutDoneReq { - key: key.to_string(), - put_id, - lease_id, - }, - raw_bytes: Vec::new(), - }; - - // 获取 master 节点 ID - let master_node_id = self - .view - .cluster_manager() - .find_or_wait_master_node() - .await?; - let rpc_started_at = Instant::now(); + }); +} - // 调用 RPC - let resp = self - .rpc_caller_put_done - .call( - self.view.p2p_module(), - master_node_id.into(), - req, - Some(std::time::Duration::from_secs(60)), - 0, - ) - .await - .map_err(KvError::from)?; - if let Err(e) = crate::rpcresp_kvresult_convert::try_from_code( - resp.serialize_part.error_code, - resp.serialize_part.error_json.clone(), - ) { - return Err(e); +/// Returns the linear group partition only when every multi-key group is present, +/// contiguous, and in the caller-declared member order. A partial batch must use +/// the legacy per-item descriptors so the V2 wire format never changes semantics. +fn complete_owner_local_publish_group_lens(items: &[OwnerLocalPublishItem]) -> Option> { + let mut offset = 0usize; + let mut group_lens = Vec::new(); + while offset < items.len() { + let item = &items[offset]; + let Some(group) = item.atomic_group.as_ref() else { + group_lens.push(1); + offset += 1; + continue; + }; + let group_len = group.members.len(); + if group_len < 2 { + return None; } - Ok(PutEndStats { - master_put_end_rpc_us: duration_to_i64_us(rpc_started_at.elapsed()), - master_put_end_server_us: resp.serialize_part.server_process_us, + let end = offset.checked_add(group_len)?; + let group_items = items.get(offset..end)?; + if group_items + .iter() + .zip(group.members.iter()) + .any(|(group_item, member)| { + group_item.key != member.key + || group_item.put_id != member.put_id + || group_item.atomic_group.as_ref() != Some(group) + }) + { + return None; + } + group_lens.push(group_len); + offset = end; + } + Some(group_lens) +} + +fn owner_local_publish_atomic_batch_complete( + item: &OwnerLocalPublishItem, + promoted_items: &[&OwnerLocalPublishItem], +) -> bool { + item.atomic_group.as_ref().map_or(true, |group| { + group.members.iter().all(|member| { + promoted_items + .iter() + .any(|published| published.key == member.key && published.put_id == member.put_id) }) + }) +} + +pub(crate) async fn publish_owner_local_job(view: ClientKvApiView, job: OwnerLocalPublishJob) { + let inner = view.client_kv_api().inner(); + if job.items.is_empty() { + release_owner_local_publish_reservations(inner, job.key_reservation_ids).await; + return; } - /// 完成 Put 操作，提交数据（带监控）：适配 external 路径，统一聚合 t1..t4 - pub async fn put_end( - &self, - key: &str, - put_id: PutIDForAKey, - lease_id: Option, - ) -> KvResult { - let metrics = self.metrics_handle(); - let client_id = self.client_id_str(); - let node_role = self.node_role(); + let group_lens = complete_owner_local_publish_group_lens(&job.items); + let incomplete_declared_group = + group_lens.is_none() && job.items.iter().any(|item| item.atomic_group.is_some()); + let mut shutdown_waiter = view.register_shutdown_waiter(); + let mut retry_delay = OWNER_LOCAL_PUBLISH_RETRY_INITIAL; + let mut published = false; - let end_stats = match self.put_end_inner(key, put_id, lease_id).await { - Ok(stats) => stats, - Err(e) => { - // on error, emit end error using pending info if exists, then cleanup - crate::observe_kvope::obe_put_end_error_from_pending( - &metrics, &client_id, &node_role, put_id, + loop { + let attempt: Result<(), String> = async { + if incomplete_declared_group { + return Err( + "declared atomic group is incomplete or non-contiguous in publish job" + .to_string(), ); - return Err(e); } - }; + let done_items: Vec = + if let Some(atomic_group_lens) = group_lens.clone() { + inner + .owner_hot_counters + .grouped_put_done_batches + .fetch_add(1, Ordering::Relaxed); + inner + .owner_hot_counters + .grouped_put_done_items + .fetch_add(job.items.len() as u64, Ordering::Relaxed); + let items = job + .items + .iter() + .map(|item| GroupedBatchPutDoneItemReq { + key: item.key.clone(), + put_id: item.put_id, + lease_id: item.lease_id, + committed_slot: Some(item.committed_slot.clone()), + publish_local_cache: false, + }) + .collect::>(); + inner + .grouped_batch_put_done(items, atomic_group_lens) + .await + .map(|resp| resp.items) + .map_err(|err| format!("PutDone RPC uncertain: {err}"))? + } else { + inner + .owner_hot_counters + .legacy_put_done_batches + .fetch_add(1, Ordering::Relaxed); + inner + .owner_hot_counters + .legacy_put_done_items + .fetch_add(job.items.len() as u64, Ordering::Relaxed); + let items = job + .items + .iter() + .map(|item| BatchPutDoneItemReq { + key: item.key.clone(), + put_id: item.put_id, + lease_id: item.lease_id, + committed_slot: Some(item.committed_slot.clone()), + publish_local_cache: false, + atomic_group: item.atomic_group.clone(), + }) + .collect::>(); + inner + .batch_put_done(items) + .await + .map(|resp| resp.items) + .map_err(|err| format!("PutDone RPC uncertain: {err}"))? + }; - // record end_handle to pending before aggregation - metrics.pending_put_set_end_handle(put_id, end_stats.master_put_end_server_us); + if done_items.len() != job.items.len() { + return Err(format!( + "PutDone response length mismatch: expected={} got={}", + job.items.len(), + done_items.len() + )); + } + for (item, done_item) in job.items.iter().zip(done_items.iter()) { + if done_item.key != item.key || done_item.put_id != item.put_id { + return Err(format!( + "PutDone identity mismatch: request=({},({},{}) response=({},({},{}))", + item.key, + item.put_id.0, + item.put_id.1, + done_item.key, + done_item.put_id.0, + done_item.put_id.1, + )); + } + crate::rpcresp_kvresult_convert::try_from_code( + done_item.error_code, + done_item.error_json.clone(), + ) + .map_err(|err| { + format!( + "PutDone item unresolved: key={} put_id=({},{}) err={}", + item.key, item.put_id.0, item.put_id.1, err + ) + })?; + } - // success: aggregate with pending timestamps; this also clears pending - crate::observe_kvope::obe_put_done_success_from_pending( - &metrics, &client_id, &node_role, key, put_id, 0, + // Do not expose only part of an atomic/TP atomic_batch to owner-hot. + // First prove every master route terminal, then roll all local + // precommit slots forward. Replayed attempts accept members that + // were already promoted before a cancellation or partial local + // failure. + for item in &job.items { + if inner.committed_local_reserve_slot_is_current( + &item.key, + item.put_id, + &item.committed_slot, + ) { + continue; + } + let memory_info = inner + .precommit_local_visible_memory_info(&item.key) + .ok_or_else(|| { + format!( + "precommit slot missing before promotion: key={} put_id=({},{})", + item.key, item.put_id.0, item.put_id.1 + ) + })?; + inner + .promote_precommit_local_reserve_resident_slot_if_same( + &item.key, + item.put_id, + memory_info, + item.atomic_group.as_ref(), + ) + .map_err(|err| { + format!( + "local promotion unresolved: key={} put_id=({},{}) err={}", + item.key, item.put_id.0, item.put_id.1, err + ) + })?; + } + Ok(()) + } + .await; + + match attempt { + Ok(()) => { + published = true; + break; + } + Err(reason) => { + tracing::warn!( + "owner local publish retained full atomic_batch for retry: key_count={} external_fences={} retry_ms={} reason={}", + job.items.len(), + job.external_pending_contexts.len(), + retry_delay.as_millis(), + reason, + ); + } + } + + tokio::select! { + _ = tokio::time::sleep(retry_delay) => {} + _ = shutdown_waiter.wait() => break, + } + retry_delay = retry_delay + .saturating_mul(2) + .min(OWNER_LOCAL_PUBLISH_RETRY_MAX); + } + + if published { + let promoted_items = job.items.iter().collect::>(); + for item in &promoted_items { + if !owner_local_publish_atomic_batch_complete(item, &promoted_items) { + tracing::warn!( + "owner local publish skipped hot admission for incomplete atomic group: key={} put_id=({},{})", + item.key, + item.put_id.0, + item.put_id.1 + ); + continue; + } + let _ = inner.owner_hot_admit_published_committed(&item.key, item.put_id); + } + + let _ = inner.start_owner_local_ssd_puts( + promoted_items + .iter() + .filter(|item| item.make_replica_task) + .map(|item| OwnerLocalSsdPutCandidate { + key: item.key.clone(), + put_id: item.put_id, + value_len: item.value_len, + selected_victim: None, + }) + .collect(), + false, + ); + + for item in promoted_items { + if owner_local_publish_starts_remote(item) { + if let Err(err) = inner + .ensure_remote_put( + &item.key, + item.put_id, + item.preferred_sub_cluster.clone(), + true, + ) + .await + { + tracing::warn!( + "owner local publish enqueue replica append failed: key={} put_id=({},{}) err={}", + item.key, + item.put_id.0, + item.put_id.1, + err + ); + } + } + } + + for context in &job.external_pending_contexts { + context._pending_fence.mark_local_put_succeeded(); + } + } + + release_owner_local_publish_reservations(inner, job.key_reservation_ids).await; +} + +async fn release_owner_local_publish_reservations( + inner: &ClientKvApiInner, + key_reservation_ids: Vec, +) { + if let Err(err) = inner + .batch_release_put_key_reservations(key_reservation_ids) + .await + { + tracing::warn!( + "owner local publish key reservation cleanup failed: {}", + err ); - Ok(end_stats) } } diff --git a/fluxon_rs/fluxon_kv/src/client_kv_api/reclaim.rs b/fluxon_rs/fluxon_kv/src/client_kv_api/reclaim.rs new file mode 100644 index 0000000..43a3b47 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/client_kv_api/reclaim.rs @@ -0,0 +1,1008 @@ +use super::{ + ClientKvApiInner, ClientKvApiView, OwnerPreparedReclaim, OwnerPreparedReclaimSource, + OwnerPreparedSsdBacking, OwnerReclaimRecord, +}; +use crate::cluster_manager::{NodeID, NodeRole}; +use crate::master_kv_router::msg_pack::{ + BatchOwnerReclaimReq, BatchOwnerReclaimResp, OwnerReclaimBacking, OwnerReclaimItem, + OwnerReclaimItemResp, OwnerReclaimItemState, OwnerReclaimPhase, OwnerReclaimReason, + OwnerSourceEvictionVictim, +}; +use crate::p2p::msg_pack::MsgPack; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, OK}; +use limit_thirdparty::tokio; +use std::{collections::HashMap, sync::Arc}; + +fn item_resp( + item: &OwnerReclaimItem, + state: OwnerReclaimItemState, + detail: impl Into, +) -> OwnerReclaimItemResp { + OwnerReclaimItemResp { + key: item.key.clone(), + epoch: item.epoch, + state, + ssd_backing_len: None, + detail: detail.into(), + } +} + +fn record_item(record: &OwnerReclaimRecord) -> &OwnerReclaimItem { + match record { + OwnerReclaimRecord::Prepared(prepared) => &prepared.item, + OwnerReclaimRecord::Releasing(item) => item, + OwnerReclaimRecord::Committed(item) => item, + } +} + +fn memory_matches_reclaim_backing( + memory_info: &crate::memholder::MemoryInfo, + backing: &OwnerReclaimBacking, +) -> bool { + match backing { + OwnerReclaimBacking::Allocation => memory_info.local_reserve_resident_slot_ref().is_none(), + OwnerReclaimBacking::UnindexedAllocation { .. } => false, + OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size, + } => memory_info.local_reserve_resident_slot_ref().is_some_and( + |(actual_slot_size, actual_grant_id, actual_slot_index)| { + actual_slot_size == *slot_size + && actual_grant_id == *grant_id + && actual_slot_index == *slot_index + }, + ), + } +} + +fn reclaim_key_control_busy_detail(state: &super::OwnerKeyControlState) -> Option<&'static str> { + if state.local_puts != 0 { + Some("owner local put is inflight") + } else if state.external_pending_puts != 0 { + Some("owner external put context is still pending") + } else if state.remote_put.is_some() { + Some("owner remote put transfer is inflight") + } else if state.local_ssd_put.is_some() { + Some("owner local SSD put is inflight") + } else if state.source_eviction_selection.is_some() { + Some("owner source eviction selection fence is active") + } else if state.external_get.is_some() { + Some("owner external Get is inflight") + } else { + None + } +} + +fn prepare_one(inner: &ClientKvApiInner, item: &OwnerReclaimItem) -> OwnerReclaimItemResp { + let mut controls = inner.owner_key_control.lock_key(&item.key); + if controls + .get(&item.key) + .is_some_and(|state| state.source_eviction_selection.is_some()) + { + let state = controls + .get_mut(&item.key) + .expect("owner source selection control state disappeared"); + let selection_matches = state + .source_eviction_selection + .as_ref() + .is_some_and(|selection| { + selection.put_id == item.put_id + && memory_matches_reclaim_backing( + selection.cached_info.mem_holder.as_ref(), + &item.backing, + ) + }); + if !selection_matches { + return item_resp( + item, + OwnerReclaimItemState::Busy, + "another owner source selection owns the key fence", + ); + } + if state.local_puts != 0 + || state.external_pending_puts != 0 + || state.remote_put.is_some() + || state.local_ssd_put.is_some() + { + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner put crossed the source selection fence", + ); + } + if inner.precommit_local_visible_info.contains_key(&item.key) + || inner.pending_local_get_info.contains_key(&item.key) + { + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner local publication crossed the source selection fence", + ); + } + let selection = state + .source_eviction_selection + .take() + .expect("matching owner source selection must exist"); + if Arc::strong_count(&selection.cached_info.mem_holder) != 1 { + state.source_eviction_selection = Some(selection); + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner local memory still has active holders", + ); + } + let local_snapshot = inner + .local_snapshot_info + .remove_if(&item.key, |_, snapshot| { + snapshot.put_time_ms == item.put_id.0 && snapshot.put_version == item.put_id.1 + }) + .map(|(_, snapshot)| snapshot); + assert!(state.reclaim.is_none()); + assert!( + state.local_access_fence.is_some(), + "source-selection promotion must retain its local-access completion generation" + ); + state.reclaim = Some(OwnerReclaimRecord::Prepared(OwnerPreparedReclaim { + item: item.clone(), + source: OwnerPreparedReclaimSource::Indexed { + cached_info: selection.cached_info, + local_snapshot, + }, + ssd_prepare_lock: Arc::new(tokio::sync::AMutex::new(())), + ssd_prepare_complete: false, + ssd_backing: None, + })); + return item_resp( + item, + OwnerReclaimItemState::Prepared, + "owner source selection promoted to reclaim fence", + ); + } + if let Some(state) = controls.get(&item.key) { + if let Some(detail) = reclaim_key_control_busy_detail(state) { + return item_resp(item, OwnerReclaimItemState::Busy, detail); + } + if let Some(record) = state.reclaim.as_ref() { + if record_item(record) == item { + return item_resp( + item, + match record { + OwnerReclaimRecord::Prepared(_) => OwnerReclaimItemState::Prepared, + OwnerReclaimRecord::Releasing(_) => OwnerReclaimItemState::Busy, + OwnerReclaimRecord::Committed(_) => OwnerReclaimItemState::Committed, + }, + "reclaim phase already applied", + ); + } + return item_resp( + item, + OwnerReclaimItemState::Busy, + "another reclaim epoch owns the key fence", + ); + } + } + if inner.precommit_local_visible_info.contains_key(&item.key) { + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner precommit local index is still visible", + ); + } + if inner.pending_local_get_info.contains_key(&item.key) { + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner local Get commit is pending", + ); + } + if let OwnerReclaimBacking::UnindexedAllocation { + addr, + base_addr, + len, + capacity_bytes, + } = &item.backing + { + if item.reason != OwnerReclaimReason::MasterAllocationCapacity + || *len == 0 + || *capacity_bytes < *len + || *addr < *base_addr + || addr.checked_add(*len).is_none() + { + return item_resp( + item, + OwnerReclaimItemState::Stale, + "invalid unindexed Allocation source identity", + ); + } + if inner.get_cached_info.contains_key(&item.key) + || inner.local_snapshot_info.contains_key(&item.key) + { + return item_resp( + item, + OwnerReclaimItemState::Stale, + "unindexed Allocation unexpectedly has an owner-local index", + ); + } + let state = controls.entry(item.key.clone()).or_default(); + assert!( + state.reclaim.is_none() + && state.local_puts == 0 + && state.external_pending_puts == 0 + && state.external_get.is_none() + ); + state.begin_local_access_fence(); + state.reclaim = Some(OwnerReclaimRecord::Prepared(OwnerPreparedReclaim { + item: item.clone(), + source: OwnerPreparedReclaimSource::UnindexedAllocation { + addr: *addr, + len: *len, + }, + ssd_prepare_lock: Arc::new(tokio::sync::AMutex::new(())), + ssd_prepare_complete: false, + ssd_backing: None, + })); + return item_resp( + item, + OwnerReclaimItemState::Prepared, + "master-owned Allocation source fenced", + ); + } + let Some((_key, cached_info)) = inner.get_cached_info.remove_if(&item.key, |_, cached| { + cached.put_time_ms == item.put_id.0 + && cached.put_version == item.put_id.1 + && memory_matches_reclaim_backing(cached.mem_holder.as_ref(), &item.backing) + }) else { + return item_resp( + item, + OwnerReclaimItemState::Stale, + "matching local backing index is absent", + ); + }; + + // The index entry is now hidden while the same control lock keeps all new local readers out. + // Any reader that cloned the memory just before the fence is visible in the Arc count. + if Arc::strong_count(&cached_info.mem_holder) != 1 { + let replaced = inner.get_cached_info.insert(item.key.clone(), cached_info); + assert!( + replaced.is_none(), + "owner reclaim rollback must restore an empty local index slot" + ); + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner local memory still has active holders", + ); + } + + let local_snapshot = inner + .local_snapshot_info + .remove_if(&item.key, |_, snapshot| { + snapshot.put_time_ms == item.put_id.0 && snapshot.put_version == item.put_id.1 + }) + .map(|(_, snapshot)| snapshot); + let state = controls.entry(item.key.clone()).or_default(); + assert!( + state.reclaim.is_none() + && state.local_puts == 0 + && state.external_pending_puts == 0 + && state.external_get.is_none() + ); + state.begin_local_access_fence(); + state.reclaim = Some(OwnerReclaimRecord::Prepared(OwnerPreparedReclaim { + item: item.clone(), + source: OwnerPreparedReclaimSource::Indexed { + cached_info, + local_snapshot, + }, + ssd_prepare_lock: Arc::new(tokio::sync::AMutex::new(())), + ssd_prepare_complete: false, + ssd_backing: None, + })); + item_resp( + item, + OwnerReclaimItemState::Prepared, + "owner local index fenced", + ) +} + +/// Persist one bounded master-capacity batch while every exact source remains +/// fenced. Per-key async locks make an overlapping RPC replay join the first +/// attempt. The storage layer then admits the whole batch without queueing, +/// inserts every independent generation, and executes one durability barrier. +async fn persist_prepared_reclaim_batch_to_ssd( + inner: &ClientKvApiInner, + items: &[OwnerReclaimItem], +) -> Vec> { + let mut backing_lens = vec![None; items.len()]; + if items.is_empty() || inner.ssd_storage.is_none() { + return backing_lens; + } + + let mut prepare_locks = Vec::new(); + for (index, item) in items.iter().enumerate() { + if item.reason != OwnerReclaimReason::MasterAllocationCapacity { + continue; + } + let controls = inner.owner_key_control.lock_key(&item.key); + let Some(OwnerReclaimRecord::Prepared(prepared)) = controls + .get(&item.key) + .and_then(|state| state.reclaim.as_ref()) + else { + continue; + }; + if prepared.item != *item { + continue; + } + if prepared.ssd_prepare_complete { + backing_lens[index] = prepared.ssd_backing.as_ref().map(|backing| backing.len); + continue; + } + prepare_locks.push((item.key.clone(), index, prepared.ssd_prepare_lock.clone())); + } + prepare_locks.sort_by(|left, right| left.0.cmp(&right.0)); + debug_assert!( + prepare_locks + .windows(2) + .all(|window| window[0].0 != window[1].0) + ); + let mut prepare_guards = Vec::with_capacity(prepare_locks.len()); + for (_, _, lock) in &prepare_locks { + prepare_guards.push(lock.clone().lock_owned().await); + } + + let mut sources = Vec::new(); + let mut source_indices = Vec::new(); + let mut source_holders = Vec::new(); + for (_, index, _) in &prepare_locks { + let item = &items[*index]; + let controls = inner.owner_key_control.lock_key(&item.key); + let Some(OwnerReclaimRecord::Prepared(prepared)) = controls + .get(&item.key) + .and_then(|state| state.reclaim.as_ref()) + else { + continue; + }; + if prepared.item != *item { + continue; + } + if prepared.ssd_prepare_complete { + backing_lens[*index] = prepared.ssd_backing.as_ref().map(|backing| backing.len); + continue; + } + let (addr, len, holder) = match &prepared.source { + OwnerPreparedReclaimSource::Indexed { cached_info, .. } => { + let holder = cached_info.mem_holder.clone(); + (holder.addr, u64::from(holder.len), Some(holder)) + } + OwnerPreparedReclaimSource::UnindexedAllocation { addr, len } => (*addr, *len, None), + }; + sources.push(crate::kv_ssd_storage::KvSsdPersistSource { + key: item.key.clone(), + put_id: item.put_id, + addr, + len, + }); + source_indices.push(*index); + source_holders.push(holder); + } + + let persisted = if sources.is_empty() { + Vec::new() + } else { + match inner.persist_local_kvs_to_ssd(&sources).await { + Ok(results) => results, + Err(err) => { + tracing::warn!( + items = sources.len(), + error = %err, + "master-capacity SSD batch validation failed; continuing DRAM reclaim" + ); + sources + .iter() + .map(|_| None) + .collect::>() + .into_iter() + .map(Ok) + .collect() + } + } + }; + drop(source_holders); + + let mut discard = Vec::new(); + for ((index, source), outcome) in source_indices + .into_iter() + .zip(sources.into_iter()) + .zip(persisted.into_iter()) + { + let item = &items[index]; + let mut persist_guard = match outcome { + Ok(Some(guard)) => Some(guard), + Ok(None) => None, + Err(err) => { + tracing::warn!( + key = item.key, + put_time_ms = item.put_id.0, + put_version = item.put_id.1, + epoch = item.epoch, + len = source.len, + error = %err, + "master-capacity victim SSD write-back failed; continuing DRAM reclaim" + ); + None + } + }; + let attached = { + let mut controls = inner.owner_key_control.lock_key(&item.key); + match controls + .get_mut(&item.key) + .and_then(|state| state.reclaim.as_mut()) + { + Some(OwnerReclaimRecord::Prepared(prepared)) if prepared.item == *item => { + if prepared.ssd_prepare_complete { + Some(prepared.ssd_backing.as_ref().map(|backing| backing.len)) + } else { + prepared.ssd_prepare_complete = true; + if let Some(guard) = persist_guard.take() { + prepared.ssd_backing = Some(OwnerPreparedSsdBacking { + len: source.len, + _persist_guard: guard, + }); + Some(Some(source.len)) + } else { + Some(None) + } + } + } + _ => None, + } + }; + match attached { + Some(len) => backing_lens[index] = len, + None => { + let should_discard = persist_guard.is_some(); + drop(persist_guard); + if should_discard { + discard.push((item.key.clone(), item.put_id)); + } + } + } + } + drop(prepare_guards); + for (key, put_id) in discard { + inner.discard_local_ssd_replica(&key, put_id).await; + } + backing_lens +} + +fn release_prepared_backing_now(inner: &ClientKvApiInner, prepared: OwnerPreparedReclaim) { + match prepared.source { + OwnerPreparedReclaimSource::Indexed { cached_info, .. } => { + let mut memory_info = Arc::try_unwrap(cached_info.mem_holder).unwrap_or_else(|_| { + panic!( + "owner reclaim prepared memory unexpectedly gained a holder: key={} epoch={}", + prepared.item.key, prepared.item.epoch + ) + }); + match &prepared.item.backing { + OwnerReclaimBacking::Allocation => { + assert!( + memory_info.local_reserve_resident_slot_ref().is_none(), + "allocation reclaim must not carry a local-reserve slot" + ); + } + OwnerReclaimBacking::UnindexedAllocation { .. } => { + unreachable!("indexed reclaim source cannot name an unindexed Allocation") + } + OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size, + } => { + let (actual_slot_size, actual_grant_id, actual_slot_index) = memory_info + .take_local_reserve_resident_slot_ref() + .expect("committed-slot reclaim must carry a local-reserve slot"); + assert_eq!(actual_slot_size, *slot_size); + assert_eq!(actual_grant_id, *grant_id); + assert_eq!(actual_slot_index, *slot_index); + + inner + .owner_release_local_reserve_committed_resident_slot( + actual_slot_size, + actual_grant_id, + actual_slot_index, + ) + .expect("owner reclaim committed resident slot release must succeed"); + } + } + drop(memory_info); + } + OwnerPreparedReclaimSource::UnindexedAllocation { .. } => { + assert!(matches!( + prepared.item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + )); + // The master route still owns the Allocation. It is removed only after this Commit + // response, which releases the physical bytes on the master side. + } + } +} + +fn reclaim_release_fence_is_intact( + state: &super::OwnerKeyControlState, + item: &OwnerReclaimItem, +) -> bool { + // Prepare hides the local index before installing the reclaim fence. A + // later external Get may share this key state, but it can only take the + // remote path and therefore does not hold the detached local backing. + matches!( + state.reclaim.as_ref(), + Some(OwnerReclaimRecord::Releasing(releasing)) if releasing == item + ) && state.local_puts == 0 + && state.external_pending_puts == 0 + && state.remote_put.is_none() + && state.local_ssd_put.is_none() + && state.source_eviction_selection.is_none() + && state.local_access_fence.is_some() +} + +fn commit_one(inner: &ClientKvApiInner, item: &OwnerReclaimItem) -> OwnerReclaimItemResp { + let mut controls = inner.owner_key_control.lock_key(&item.key); + let Some(state) = controls.get_mut(&item.key) else { + return item_resp( + item, + OwnerReclaimItemState::Stale, + "owner reclaim fence is absent", + ); + }; + let Some(record) = state.reclaim.take() else { + return item_resp( + item, + OwnerReclaimItemState::Stale, + "owner reclaim fence is absent", + ); + }; + let prepared = match record { + OwnerReclaimRecord::Prepared(prepared) if prepared.item == *item => { + state.reclaim = Some(OwnerReclaimRecord::Releasing(item.clone())); + prepared + } + OwnerReclaimRecord::Releasing(releasing) if releasing == *item => { + state.reclaim = Some(OwnerReclaimRecord::Releasing(releasing)); + return item_resp( + item, + OwnerReclaimItemState::Busy, + "owner reclaim slot release is already in progress", + ); + } + OwnerReclaimRecord::Committed(committed) if committed == *item => { + state.reclaim = Some(OwnerReclaimRecord::Committed(committed)); + return item_resp( + item, + OwnerReclaimItemState::Committed, + "owner reclaim commit already applied", + ); + } + other => { + state.reclaim = Some(other); + return item_resp( + item, + OwnerReclaimItemState::Stale, + "owner reclaim epoch or slot identity changed", + ); + } + }; + + // The Releasing marker keeps local Put/Get out. Drop the key-shard lock + // before touching the slot pool so no synchronous locks are nested. + drop(controls); + release_prepared_backing_now(inner, prepared); + + let mut controls = inner.owner_key_control.lock_key(&item.key); + let state = controls + .get_mut(&item.key) + .expect("owner reclaim releasing fence disappeared"); + assert!( + reclaim_release_fence_is_intact(state, item), + "a local put crossed an owner reclaim fence" + ); + state.reclaim = Some(OwnerReclaimRecord::Committed(item.clone())); + drop(controls); + inner.owner_hot_invalidate_version(&item.key, item.put_id); + item_resp( + item, + OwnerReclaimItemState::Committed, + "owner committed slot released", + ) +} + +#[cfg(test)] +mod tests { + use super::{reclaim_key_control_busy_detail, reclaim_release_fence_is_intact}; + use crate::client_kv_api::{ + ExternalGetKeySharedOp, OwnerKeyControlState, OwnerKeyControlTable, OwnerReclaimRecord, + acquire_external_pending_put_fence_for_key, + }; + use crate::master_kv_router::msg_pack::{ + OwnerReclaimBacking, OwnerReclaimItem, OwnerReclaimReason, + }; + use std::sync::Arc; + + #[test] + fn pending_external_put_rejects_reclaim_prepare_precheck() { + let controls = Arc::new(OwnerKeyControlTable::default()); + let _guard = acquire_external_pending_put_fence_for_key(&controls, "pending-key") + .expect("pending fence acquisition must succeed"); + let controls = controls.lock_key("pending-key"); + assert_eq!( + reclaim_key_control_busy_detail(&controls["pending-key"]), + Some("owner external put context is still pending") + ); + } + + #[test] + fn remote_get_marker_can_overlap_reclaim_commit() { + let item = OwnerReclaimItem { + key: "remote-during-reclaim".to_string(), + put_id: (7, 1), + epoch: 9, + backing: OwnerReclaimBacking::CommittedSlot { + grant_id: 3, + slot_index: 4, + slot_size: 4096, + }, + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + let mut state = OwnerKeyControlState { + external_get: Some(Arc::new(ExternalGetKeySharedOp::new( + "remote-during-reclaim".to_string(), + ))), + ..Default::default() + }; + state.reclaim = Some(OwnerReclaimRecord::Releasing(item.clone())); + state.begin_local_access_fence(); + assert!(reclaim_release_fence_is_intact(&state, &item)); + + let mut local_put_state = OwnerKeyControlState { + local_puts: 1, + external_get: state.external_get, + ..Default::default() + }; + local_put_state.reclaim = Some(OwnerReclaimRecord::Releasing(item.clone())); + local_put_state.begin_local_access_fence(); + assert!(!reclaim_release_fence_is_intact(&local_put_state, &item)); + } +} + +struct AbortOneResult { + response: OwnerReclaimItemResp, + discard_ssd: bool, +} + +fn abort_one_fenced(inner: &ClientKvApiInner, item: &OwnerReclaimItem) -> AbortOneResult { + let mut controls = inner.owner_key_control.lock_key(&item.key); + let Some(state) = controls.get_mut(&item.key) else { + return AbortOneResult { + response: item_resp( + item, + OwnerReclaimItemState::Aborted, + "owner reclaim was already absent", + ), + discard_ssd: false, + }; + }; + let Some(record) = state.reclaim.take() else { + return AbortOneResult { + response: item_resp( + item, + OwnerReclaimItemState::Aborted, + "owner reclaim was already absent", + ), + discard_ssd: false, + }; + }; + match record { + OwnerReclaimRecord::Prepared(prepared) if prepared.item == *item => { + let OwnerPreparedReclaim { + source, + ssd_backing, + .. + } = prepared; + let detail = match source { + OwnerPreparedReclaimSource::Indexed { + cached_info, + local_snapshot, + } => { + let replaced = inner.get_cached_info.insert(item.key.clone(), cached_info); + assert!( + replaced.is_none(), + "owner reclaim abort must restore an empty local index slot" + ); + if let Some(snapshot) = local_snapshot { + let replaced = inner.local_snapshot_info.insert(item.key.clone(), snapshot); + assert!( + replaced.is_none(), + "owner reclaim abort must restore an empty local snapshot slot" + ); + } + "owner local index fence rolled back" + } + OwnerPreparedReclaimSource::UnindexedAllocation { .. } => { + "master-owned Allocation source fence rolled back" + } + }; + state.finish_local_access_fence(); + if state.is_idle() { + controls.remove(&item.key); + } + AbortOneResult { + response: item_resp(item, OwnerReclaimItemState::Aborted, detail), + discard_ssd: ssd_backing.is_some(), + } + } + OwnerReclaimRecord::Releasing(releasing) if releasing == *item => { + state.reclaim = Some(OwnerReclaimRecord::Releasing(releasing)); + AbortOneResult { + response: item_resp( + item, + OwnerReclaimItemState::Busy, + "owner slot release is already in progress and cannot be aborted", + ), + discard_ssd: false, + } + } + OwnerReclaimRecord::Committed(committed) if committed == *item => { + state.reclaim = Some(OwnerReclaimRecord::Committed(committed)); + AbortOneResult { + response: item_resp( + item, + OwnerReclaimItemState::Committed, + "owner slot was already committed and cannot be restored", + ), + discard_ssd: false, + } + } + other => { + state.reclaim = Some(other); + AbortOneResult { + response: item_resp( + item, + OwnerReclaimItemState::Stale, + "owner reclaim epoch or slot identity changed", + ), + discard_ssd: false, + } + } + } +} + +async fn abort_one(inner: &ClientKvApiInner, item: &OwnerReclaimItem) -> OwnerReclaimItemResp { + let outcome = abort_one_fenced(inner, item); + if outcome.discard_ssd { + inner + .discard_local_ssd_replica(&item.key, item.put_id) + .await; + } + outcome.response +} + +fn finalize_one(inner: &ClientKvApiInner, item: &OwnerReclaimItem) -> OwnerReclaimItemResp { + let mut controls = inner.owner_key_control.lock_key(&item.key); + let Some(state) = controls.get_mut(&item.key) else { + return item_resp( + item, + OwnerReclaimItemState::Finalized, + "owner reclaim was already finalized", + ); + }; + match state.reclaim.take() { + Some(OwnerReclaimRecord::Committed(committed)) if committed == *item => { + state.finish_local_access_fence(); + if state.is_idle() { + controls.remove(&item.key); + } + item_resp( + item, + OwnerReclaimItemState::Finalized, + "owner reclaim fence cleared", + ) + } + Some(OwnerReclaimRecord::Releasing(releasing)) if releasing == *item => { + state.reclaim = Some(OwnerReclaimRecord::Releasing(releasing)); + item_resp( + item, + OwnerReclaimItemState::Busy, + "owner slot release is still in progress", + ) + } + Some(other) => { + state.reclaim = Some(other); + item_resp( + item, + OwnerReclaimItemState::Busy, + "owner reclaim is not committed for this epoch", + ) + } + None => item_resp( + item, + OwnerReclaimItemState::Finalized, + "owner reclaim was already finalized", + ), + } +} + +pub(crate) fn complete_owner_source_eviction( + inner: &ClientKvApiInner, + victim: &OwnerSourceEvictionVictim, + epoch: u64, +) -> Result<(), String> { + let item = OwnerReclaimItem { + key: victim.key.clone(), + put_id: victim.put_id, + epoch, + backing: victim.backing.clone(), + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + let prepared = prepare_one(inner, &item); + match prepared.state { + OwnerReclaimItemState::Prepared => {} + OwnerReclaimItemState::Committed => { + let finalized = finalize_one(inner, &item); + return (finalized.state == OwnerReclaimItemState::Finalized) + .then_some(()) + .ok_or(finalized.detail); + } + OwnerReclaimItemState::Finalized => return Ok(()), + _ => return Err(prepared.detail), + } + + let committed = commit_one(inner, &item); + if committed.state != OwnerReclaimItemState::Committed { + return Err(committed.detail); + } + let finalized = finalize_one(inner, &item); + (finalized.state == OwnerReclaimItemState::Finalized) + .then_some(()) + .ok_or(finalized.detail) +} + +pub async fn handle_batch_owner_reclaim( + view: &ClientKvApiView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let requester_is_master = view + .cluster_manager() + .get_member_info_cached(req_node_id.as_ref()) + .is_some_and(|member| matches!(member.node_role(), NodeRole::Master)); + if !requester_is_master { + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch owner reclaim requester is not the current master: requester={}", + req_node_id + ), + }); + return MsgPack { + serialize_part: BatchOwnerReclaimResp { + items: Vec::new(), + error_code: err.code(), + error_json: err.to_json(), + }, + raw_bytes: Vec::new(), + }; + } + let inner = view.client_kv_api().inner(); + let phase = req.serialize_part.phase; + let items = match phase { + OwnerReclaimPhase::Prepare => { + let mut responses = req + .serialize_part + .items + .iter() + .map(|item| prepare_one(inner, item)) + .collect::>(); + let prepared_indices = responses + .iter() + .enumerate() + .filter_map(|(index, response)| { + (response.state == OwnerReclaimItemState::Prepared).then_some(index) + }) + .collect::>(); + for indices in prepared_indices.chunks(crate::kv_ssd_storage::MAX_PERSIST_BATCH_ITEMS) { + let batch = indices + .iter() + .map(|index| req.serialize_part.items[*index].clone()) + .collect::>(); + let persisted = persist_prepared_reclaim_batch_to_ssd(inner, &batch).await; + for (index, ssd_backing_len) in indices.iter().copied().zip(persisted) { + responses[index].ssd_backing_len = ssd_backing_len; + } + } + responses + } + OwnerReclaimPhase::Commit => req + .serialize_part + .items + .iter() + .map(|item| commit_one(inner, item)) + .collect(), + OwnerReclaimPhase::Abort => { + futures::future::join_all( + req.serialize_part + .items + .iter() + .map(|item| abort_one(inner, item)), + ) + .await + } + OwnerReclaimPhase::Finalize => req + .serialize_part + .items + .iter() + .map(|item| finalize_one(inner, item)) + .collect(), + }; + let prepared = items + .iter() + .filter(|item| item.state == OwnerReclaimItemState::Prepared) + .count(); + let committed = items + .iter() + .filter(|item| item.state == OwnerReclaimItemState::Committed) + .count(); + let finalized = items + .iter() + .filter(|item| item.state == OwnerReclaimItemState::Finalized) + .count(); + let ssd_prepared = items + .iter() + .filter(|item| item.ssd_backing_len.is_some()) + .count(); + let ssd_prepared_bytes = items + .iter() + .filter_map(|item| item.ssd_backing_len) + .fold(0_u64, u64::saturating_add); + let busy_or_stale = items + .iter() + .filter(|item| { + matches!( + item.state, + OwnerReclaimItemState::Busy | OwnerReclaimItemState::Stale + ) + }) + .count(); + let mut rejection_reason_counts = HashMap::::new(); + for item in &items { + if matches!( + item.state, + OwnerReclaimItemState::Busy | OwnerReclaimItemState::Stale + ) { + *rejection_reason_counts + .entry(item.detail.clone()) + .or_default() += 1; + } + } + let mut rejection_reason_counts = rejection_reason_counts.into_iter().collect::>(); + rejection_reason_counts.sort_by(|a, b| a.0.cmp(&b.0)); + tracing::info!( + "owner reclaim phase completed: phase={:?} items={} prepared={} committed={} finalized={} ssd_prepared={} ssd_prepared_bytes={} busy_or_stale={} rejection_reasons={:?}", + phase, + items.len(), + prepared, + committed, + finalized, + ssd_prepared, + ssd_prepared_bytes, + busy_or_stale, + rejection_reason_counts + ); + MsgPack { + serialize_part: BatchOwnerReclaimResp { + items, + error_code: OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} diff --git a/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs b/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs index 1aa6954..8dd61aa 100644 --- a/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs +++ b/fluxon_rs/fluxon_kv/src/client_seg_pool/mod.rs @@ -1,7 +1,7 @@ use crate::ClientKvApiAccessTrait; use crate::client_kv_api::ClientKvApi; use crate::client_transfer_engine::{ClientTransferEngine, ClientTransferEngineAccessTrait}; -use crate::cluster_manager::ClusterManagerAccessTrait; +use crate::cluster_manager::{ClusterManagerAccessTrait, NodeRole}; use crate::config::ContributeToClusterPoolSize; use crate::master_seg_manager::msg_pack::SegmentDeviceMemInfo; use crate::p2p::p2p_module::P2pModule; @@ -237,10 +237,7 @@ impl ClientSegPool { std::path::Path::new(share_mem_path).join(SIDE_TRANSFER_PEERS_DIRNAME) } - pub fn side_transfer_peer_file_path( - share_mem_path: &str, - side_id: &str, - ) -> std::path::PathBuf { + pub fn side_transfer_peer_file_path(share_mem_path: &str, side_id: &str) -> std::path::PathBuf { Self::side_transfer_peers_dir(share_mem_path).join(format!("{side_id}.json")) } @@ -399,17 +396,13 @@ impl ClientSegPool { crate::rpcresp_kvresult_convert::msg_and_error::SharedMemError::MappingFailed { path: String::new(), len: map_len as u64, - detail: "share_mem_path is empty; explicit configuration required" - .to_string(), + detail: "share_mem_path is empty; explicit configuration required".to_string(), }, )); } let base_path = &share_mem_path; - tracing::info!( - "Using share_mem_path: {} for memory-mapped file", - base_path - ); + tracing::info!("Using share_mem_path: {} for memory-mapped file", base_path); std::fs::create_dir_all(base_path).map_err(|e| { KvError::SharedMem( crate::rpcresp_kvresult_convert::msg_and_error::SharedMemError::MappingFailed { @@ -975,15 +968,24 @@ impl LogicalModule for ClientSegPool { async fn shutdown(&self) -> Result<(), Self::Error> { let _ = self.remove_side_transfer_peer().await; - loop { - if self.inner().view().client_kv_api().can_be_dropped() { - tracing::info!("ClientSegPool can be dropped"); - break; + if !matches!( + self.inner() + .view() + .cluster_manager() + .get_self_info() + .node_role(), + NodeRole::External + ) { + loop { + if self.inner().view().client_kv_api().can_be_dropped() { + tracing::info!("ClientSegPool can be dropped"); + break; + } + tracing::info!( + "ClientSegPool waiting ClientKvApi can not be dropped , will try again after 3s (some user memholder may still be in use)" + ); + tokio::time::sleep(Duration::from_secs(3)).await; } - tracing::info!( - "ClientSegPool waiting ClientKvApi can not be dropped , will try again after 3s (some user memholder may still be in use)" - ); - tokio::time::sleep(Duration::from_secs(3)).await; } if self.0.cpu_allocated_mem.read().await.is_some() { @@ -996,6 +998,15 @@ impl LogicalModule for ClientSegPool { impl ClientSegPool { pub async fn init2_for_init_dag(&self) -> KvResult<()> { let inner = &self.0; + if matches!( + inner.view().cluster_manager().get_self_info().node_role(), + NodeRole::External + ) { + // External clients contribute no cache segment. They construct this + // module only to provide lifetime guards for an explicitly + // registered GPU staging range. + return Ok(()); + } // English note: // - Invariant: register inbound RPC handlers before any awaited etcd operations that diff --git a/fluxon_rs/fluxon_kv/src/client_transfer_engine/mod.rs b/fluxon_rs/fluxon_kv/src/client_transfer_engine/mod.rs index 75f82f5..829b57e 100644 --- a/fluxon_rs/fluxon_kv/src/client_transfer_engine/mod.rs +++ b/fluxon_rs/fluxon_kv/src/client_transfer_engine/mod.rs @@ -21,7 +21,10 @@ use crate::{P2pModuleAccessTrait, cluster_manager::ClusterManagerAccessTrait}; use async_trait::async_trait; use fluxon_commu::ClosedRuntimeHandle; use fluxon_commu::p2p::PeerGen; -use fluxon_commu::transfer_engine::AttachedTransferEngine; +use fluxon_commu::transfer_engine::{ + AttachedTransferEngine, CLOSED_RUNTIME_DIRECT_FAST_PATH_NOT_READY_MARKER, + ClosedRuntimeLocalMemoryKind, +}; use fluxon_commu::{ ClientTransferEngineClusterRuntime, ClientTransferEngineCore, ClientTransferEngineRuntime, CpuAllocatedMem, TransferBreakdown, @@ -29,7 +32,9 @@ use fluxon_commu::{ use fluxon_framework::{LogicalModule, define_module}; use std::future::Future; use std::pin::Pin; -use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{Duration, Instant}; pub use fluxon_commu::{ClientTransferEngineNewArg, ClientTransferEngineRuntimeConfig}; @@ -43,15 +48,228 @@ define_module!( (client_seg_pool, ClientSegPool) ); +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GpuMemoryRegistration { + pub registration_id: u64, + pub addr: u64, + pub len: u64, + pub device_id: u32, +} + +impl GpuMemoryRegistration { + fn contains(&self, addr: u64, len: u64) -> bool { + if len == 0 { + return addr >= self.addr + && self + .addr + .checked_add(self.len) + .is_some_and(|end| addr <= end); + } + let Some(end) = addr.checked_add(len) else { + return false; + }; + let Some(registration_end) = self.addr.checked_add(self.len) else { + return false; + }; + addr >= self.addr && end <= registration_end + } +} + +#[derive(Clone, Debug)] +pub struct GpuMemoryGuard { + registration: Arc, +} + +impl GpuMemoryGuard { + pub fn registration(&self) -> &GpuMemoryRegistration { + self.registration.as_ref() + } + + fn contains(&self, addr: u64, len: u64) -> bool { + self.registration.contains(addr, len) + } +} + +pub enum ClientTransferMemoryGuard { + Cpu(ClientCpuMemReadGuard), + Gpu(GpuMemoryGuard), +} + +#[derive(Debug)] +enum GpuMemoryRegistryState { + Empty, + Registering(GpuMemoryRegistration), + Active(Arc), + Unregistering(GpuMemoryRegistration), +} + +#[derive(Debug)] +struct GpuMemoryRegistry { + next_registration_id: AtomicU64, + state: Mutex, +} + +impl GpuMemoryRegistry { + fn new() -> Self { + Self { + next_registration_id: AtomicU64::new(1), + state: Mutex::new(GpuMemoryRegistryState::Empty), + } + } + + fn begin_register( + &self, + addr: u64, + len: u64, + device_id: u32, + ) -> Result { + if addr == 0 { + return Err("GPU registration address must be non-zero".to_string()); + } + if len == 0 { + return Err("GPU registration length must be non-zero".to_string()); + } + if addr.checked_add(len).is_none() { + return Err(format!( + "GPU registration range overflows: addr={:#x} len={}", + addr, len + )); + } + let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed); + let registration = GpuMemoryRegistration { + registration_id, + addr, + len, + device_id, + }; + let mut state = self.state.lock().expect("gpu memory registry poisoned"); + match &*state { + GpuMemoryRegistryState::Empty => { + *state = GpuMemoryRegistryState::Registering(registration.clone()); + Ok(registration) + } + current => Err(format!( + "one GPU registration is already active or changing: {current:?}" + )), + } + } + + fn finish_register(&self, registration_id: u64, success: bool) { + let mut state = self.state.lock().expect("gpu memory registry poisoned"); + let previous = std::mem::replace(&mut *state, GpuMemoryRegistryState::Empty); + match previous { + GpuMemoryRegistryState::Registering(registration) + if registration.registration_id == registration_id => + { + if success { + *state = GpuMemoryRegistryState::Active(Arc::new(registration)); + } + } + other => panic!( + "GPU registration completion does not match registry state: id={} state={other:?}", + registration_id + ), + } + } + + fn begin_unregister(&self, registration_id: u64) -> Result { + let mut state = self.state.lock().expect("gpu memory registry poisoned"); + let previous = std::mem::replace(&mut *state, GpuMemoryRegistryState::Empty); + match previous { + GpuMemoryRegistryState::Active(registration) => { + if registration.registration_id != registration_id { + let actual = registration.registration_id; + *state = GpuMemoryRegistryState::Active(registration); + return Err(format!( + "GPU registration id mismatch: expected={} got={}", + actual, registration_id + )); + } + let guard_count = Arc::strong_count(®istration).saturating_sub(1); + if guard_count != 0 { + *state = GpuMemoryRegistryState::Active(registration); + return Err(format!( + "GPU registration is busy: registration_id={} active_guards={}", + registration_id, guard_count + )); + } + let descriptor = registration.as_ref().clone(); + *state = GpuMemoryRegistryState::Unregistering(descriptor.clone()); + Ok(descriptor) + } + other => { + *state = other; + Err(format!( + "GPU registration is not active: registration_id={}", + registration_id + )) + } + } + } + + fn finish_unregister(&self, registration_id: u64, success: bool) { + let mut state = self.state.lock().expect("gpu memory registry poisoned"); + let previous = std::mem::replace(&mut *state, GpuMemoryRegistryState::Empty); + match previous { + GpuMemoryRegistryState::Unregistering(registration) + if registration.registration_id == registration_id => + { + if !success { + *state = GpuMemoryRegistryState::Active(Arc::new(registration)); + } + } + other => panic!( + "GPU unregistration completion does not match registry state: id={} state={other:?}", + registration_id + ), + } + } + + fn guard_for_range(&self, addr: u64, len: u64) -> Option { + let state = self.state.lock().expect("gpu memory registry poisoned"); + match &*state { + GpuMemoryRegistryState::Active(registration) if registration.contains(addr, len) => { + Some(GpuMemoryGuard { + registration: registration.clone(), + }) + } + _ => None, + } + } + + fn validate_destination( + &self, + registration_id: u64, + addr: u64, + capacity: u64, + ) -> Result { + let guard = self.guard_for_range(addr, capacity).ok_or_else(|| { + format!( + "GPU destination is outside the active registration: registration_id={} addr={:#x} capacity={}", + registration_id, addr, capacity + ) + })?; + if guard.registration().registration_id != registration_id { + return Err(format!( + "GPU destination registration id mismatch: expected={} got={}", + guard.registration().registration_id, + registration_id + )); + } + Ok(guard) + } +} + #[derive(Clone)] struct ClientTransferRuntimeAdapter { view: ClientTransferEngineView, + gpu_memory_registry: Arc, } impl ClientTransferRuntimeAdapter { fn local_segment_transfer_enabled(&self) -> bool { let self_info = self.view.cluster_manager().get_self_info(); - matches!(self_info.node_role(), NodeRole::Client) + matches!(self_info.node_role(), NodeRole::Client | NodeRole::External) || self_info .metadata .get("side_transfer_worker") @@ -156,7 +374,7 @@ impl ClientTransferEngineClusterRuntime for ClientTransferRuntimeAdapter { #[async_trait] impl ClientTransferEngineRuntime for ClientTransferRuntimeAdapter { - type LocalSegmentGuard = ClientCpuMemReadGuard; + type LocalSegmentGuard = ClientTransferMemoryGuard; fn supports_local_segment_transfer(&self) -> bool { self.local_segment_transfer_enabled() @@ -181,12 +399,40 @@ impl ClientTransferEngineRuntime for ClientTransferRuntimeAdapter { async fn ensure_local_segment_guard( &self, local_addr: u64, - seg_guard: Option, - ) -> Result { + seg_guard: Option, + ) -> Result { if !self.local_segment_transfer_enabled() { return Err("local segment transfer is not supported on this node role".to_string()); } - p2p_transfer_rpc::ensure_local_segment_guard(&self.view, local_addr, seg_guard).await + if let Some(guard) = seg_guard { + return match guard { + ClientTransferMemoryGuard::Cpu(cpu_guard) => { + p2p_transfer_rpc::ensure_local_segment_guard( + &self.view, + local_addr, + Some(cpu_guard), + ) + .await + .map(ClientTransferMemoryGuard::Cpu) + } + ClientTransferMemoryGuard::Gpu(gpu_guard) => { + if gpu_guard.contains(local_addr, 1) { + Ok(ClientTransferMemoryGuard::Gpu(gpu_guard)) + } else { + Err(format!( + "GPU guard does not cover local address: addr={:#x}", + local_addr + )) + } + } + }; + } + if let Some(gpu_guard) = self.gpu_memory_registry.guard_for_range(local_addr, 1) { + return Ok(ClientTransferMemoryGuard::Gpu(gpu_guard)); + } + p2p_transfer_rpc::ensure_local_segment_guard(&self.view, local_addr, None) + .await + .map(ClientTransferMemoryGuard::Cpu) } fn register_p2p_transfer_rpc(&self) { @@ -231,20 +477,27 @@ impl ClientTransferEngineRuntime for ClientTransferRuntimeAdapter { remote_src: u64, local_target: u64, len: u64, - seg_guard: ClientCpuMemReadGuard, + seg_guard: ClientTransferMemoryGuard, ) -> Result<(), String> { if !self.local_segment_transfer_enabled() { return Err("p2p raw-memory read is not supported on this node role".to_string()); } - p2p_transfer_rpc::p2p_read_to_local( - &self.view, - peer, - remote_src, - local_target, - len, - seg_guard, - ) - .await + match seg_guard { + ClientTransferMemoryGuard::Cpu(cpu_guard) => { + p2p_transfer_rpc::p2p_read_to_local( + &self.view, + peer, + remote_src, + local_target, + len, + cpu_guard, + ) + .await + } + ClientTransferMemoryGuard::Gpu(_) => Err( + "GPU destination requires the RDMA fast path; P2P fallback is disabled".to_string(), + ), + } } async fn p2p_write_from_local( @@ -254,21 +507,28 @@ impl ClientTransferEngineRuntime for ClientTransferRuntimeAdapter { remote_target: u64, len: u64, copy_from: Option>, - seg_guard: ClientCpuMemReadGuard, + seg_guard: ClientTransferMemoryGuard, ) -> Result<(), String> { if !self.local_segment_transfer_enabled() { return Err("p2p raw-memory write is not supported on this node role".to_string()); } - p2p_transfer_rpc::p2p_write_from_local( - &self.view, - peer, - local_src, - remote_target, - len, - copy_from, - seg_guard, - ) - .await + match seg_guard { + ClientTransferMemoryGuard::Cpu(cpu_guard) => { + p2p_transfer_rpc::p2p_write_from_local( + &self.view, + peer, + local_src, + remote_target, + len, + copy_from, + cpu_guard, + ) + .await + } + ClientTransferMemoryGuard::Gpu(_) => { + Err("GPU source requires the RDMA fast path; P2P fallback is disabled".to_string()) + } + } } fn try_record_local_ipc_bytes_for_owner_topology( @@ -308,12 +568,14 @@ impl ClientTransferEngineRuntime for ClientTransferRuntimeAdapter { pub struct ClientTransferEngine { view: OnceLock, core: ClientTransferEngineCore, + gpu_memory_registry: Arc, } impl ClientTransferEngine { fn runtime(&self) -> ClientTransferRuntimeAdapter { ClientTransferRuntimeAdapter { view: self.view.get().unwrap().clone(), + gpu_memory_registry: self.gpu_memory_registry.clone(), } } @@ -333,6 +595,7 @@ impl ClientTransferEngine { Ok(Self { view: OnceLock::new(), core, + gpu_memory_registry: Arc::new(GpuMemoryRegistry::new()), }) } @@ -369,6 +632,81 @@ impl ClientTransferEngine { .map_err(KvError::from) } + pub async fn register_gpu_memory( + &self, + addr: u64, + len: u64, + device_id: u32, + ) -> KvResult { + let registration = self + .gpu_memory_registry + .begin_register(addr, len, device_id) + .map_err(|detail| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail, + }, + ) + })?; + let result = self + .core + .register_local_memory( + self.runtime(), + addr, + len, + ClosedRuntimeLocalMemoryKind::Gpu { device_id }, + ) + .await + .map_err(KvError::from); + self.gpu_memory_registry + .finish_register(registration.registration_id, result.is_ok()); + result.map(|()| registration) + } + + pub async fn unregister_gpu_memory(&self, registration_id: u64) -> KvResult<()> { + let registration = self + .gpu_memory_registry + .begin_unregister(registration_id) + .map_err(|detail| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail, + }, + ) + })?; + let result = self + .core + .unregister_local_memory( + registration.addr, + registration.len, + ClosedRuntimeLocalMemoryKind::Gpu { + device_id: registration.device_id, + }, + ) + .await + .map_err(KvError::from); + self.gpu_memory_registry + .finish_unregister(registration_id, result.is_ok()); + result + } + + pub fn validate_gpu_destination( + &self, + registration_id: u64, + addr: u64, + capacity: u64, + ) -> KvResult { + self.gpu_memory_registry + .validate_destination(registration_id, addr, capacity) + .map_err(|detail| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail, + }, + ) + }) + } + pub async fn write_data( &self, data: Pin<&[u8]>, @@ -386,7 +724,7 @@ impl ClientTransferEngine { target_addr, peer_id, do_copy, - seg_guard, + seg_guard.map(ClientTransferMemoryGuard::Cpu), ) .await .map_err(KvError::from) @@ -409,11 +747,87 @@ impl ClientTransferEngine { src_addr, target_addr, len, - seg_guard, + seg_guard.map(ClientTransferMemoryGuard::Cpu), + false, ) .await .map_err(KvError::from) } + + /// Pull remote bytes into an explicitly validated caller-owned GPU range. + /// + /// Keeping this separate from the CPU entry point prevents a CUDA virtual + /// address from silently falling through CPU/P2P guard discovery. The + /// guard is moved into the transfer engine and retains the exact + /// registration generation until the backend releases the local segment + /// lease. + pub async fn transfer_data_no_copy_to_gpu( + &self, + peer_node: NodeIDString, + src_addr: u64, + target_addr: u64, + len: u64, + gpu_guard: GpuMemoryGuard, + ) -> KvResult { + if len == 0 { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: "GPU transfer length must be non-zero".to_string(), + }, + )); + } + if !gpu_guard.contains(target_addr, len) { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GPU transfer exceeds destination guard: target={:#x} len={} registration_id={}", + target_addr, + len, + gpu_guard.registration().registration_id + ), + }, + )); + } + const DIRECT_READY_TIMEOUT: Duration = Duration::from_secs(5); + const DIRECT_RETRY_DELAY: Duration = Duration::from_millis(10); + + let started = Instant::now(); + let mut attempt = 0_u32; + loop { + attempt = attempt.saturating_add(1); + let result = self + .core + .transfer_data_no_copy( + self.runtime(), + Some(peer_node.clone()), + true, + src_addr, + target_addr, + len, + Some(ClientTransferMemoryGuard::Gpu(gpu_guard.clone())), + true, + ) + .await; + match result { + Ok(breakdown) => return Ok(breakdown), + Err(err) + if err + .to_string() + .contains(CLOSED_RUNTIME_DIRECT_FAST_PATH_NOT_READY_MARKER) + && started.elapsed() < DIRECT_READY_TIMEOUT => + { + tracing::debug!( + peer = %peer_node, + attempt, + elapsed_ms = started.elapsed().as_millis(), + "GPU direct transfer is waiting for the RDMA peer fast path" + ); + limit_thirdparty::tokio::time::sleep(DIRECT_RETRY_DELAY).await; + } + Err(err) => return Err(KvError::from(err)), + } + } + } } #[async_trait] @@ -435,3 +849,101 @@ impl LogicalModule for ClientTransferEngine { Ok(()) } } + +#[cfg(test)] +mod gpu_memory_registry_tests { + use super::GpuMemoryRegistry; + + fn active_registry() -> (GpuMemoryRegistry, u64) { + let registry = GpuMemoryRegistry::new(); + let registration = registry + .begin_register(0x1000, 0x1000, 2) + .expect("registration must start"); + let registration_id = registration.registration_id; + registry.finish_register(registration_id, true); + (registry, registration_id) + } + + #[test] + fn validates_exact_registration_generation_and_range() { + let (registry, registration_id) = active_registry(); + + let guard = registry + .validate_destination(registration_id, 0x1400, 0x400) + .expect("subrange must validate"); + assert_eq!(guard.registration().registration_id, registration_id); + assert_eq!(guard.registration().device_id, 2); + + assert!( + registry + .validate_destination(registration_id + 1, 0x1400, 0x400) + .unwrap_err() + .contains("registration id mismatch") + ); + assert!( + registry + .validate_destination(registration_id, 0x1f00, 0x200) + .unwrap_err() + .contains("outside the active registration") + ); + } + + #[test] + fn rejects_invalid_registration_geometry() { + let registry = GpuMemoryRegistry::new(); + assert!(registry.begin_register(0, 1, 0).is_err()); + assert!(registry.begin_register(1, 0, 0).is_err()); + assert!(registry.begin_register(u64::MAX, 2, 0).is_err()); + } + + #[test] + fn unregister_waits_for_destination_guards() { + let (registry, registration_id) = active_registry(); + let guard = registry + .validate_destination(registration_id, 0x1000, 0x1000) + .expect("full range must validate"); + + assert!( + registry + .begin_unregister(registration_id) + .unwrap_err() + .contains("GPU registration is busy") + ); + drop(guard); + + let descriptor = registry + .begin_unregister(registration_id) + .expect("unregister must start after guards drop"); + assert_eq!(descriptor.registration_id, registration_id); + registry.finish_unregister(registration_id, true); + assert!( + registry + .validate_destination(registration_id, 0x1000, 1) + .is_err() + ); + } + + #[test] + fn failed_backend_transitions_restore_retryable_state() { + let registry = GpuMemoryRegistry::new(); + let first = registry + .begin_register(0x1000, 0x1000, 0) + .expect("first registration must start"); + registry.finish_register(first.registration_id, false); + + let second = registry + .begin_register(0x3000, 0x1000, 1) + .expect("failed registration must return to empty"); + registry.finish_register(second.registration_id, true); + let descriptor = registry + .begin_unregister(second.registration_id) + .expect("unregister must start"); + registry.finish_unregister(descriptor.registration_id, false); + + assert!( + registry + .validate_destination(second.registration_id, 0x3000, 1) + .is_ok() + ); + } +} diff --git a/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs b/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs index 97a5758..b563d5a 100644 --- a/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs +++ b/fluxon_rs/fluxon_kv/src/cluster_manager/mod.rs @@ -5,6 +5,8 @@ pub use fluxon_commu::{ META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME, }; +pub(crate) const META_KEY_KV_SSD_STORAGE: &str = "kv_ssd_storage"; + pub mod app_logic_ext; #[cfg(test)] diff --git a/fluxon_rs/fluxon_kv/src/config.rs b/fluxon_rs/fluxon_kv/src/config.rs index f9c7691..3634c63 100644 --- a/fluxon_rs/fluxon_kv/src/config.rs +++ b/fluxon_rs/fluxon_kv/src/config.rs @@ -97,12 +97,38 @@ pub enum SideTransferRole { Worker, } +/// Experimental source ordering for an external planned Get when SSD routes +/// are present. The default preserves the sealed r89 behavior so a new +/// release does not silently change an existing baseline. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum SsdReadSourcePolicy { + /// Any memory source wins over any SSD source; within a tier a remote + /// source wins over the requester-local owner. This is the r89 order. + #[default] + LegacyRemoteFirst, + /// Treat SSD as a requester-local tier, never as a cross-node source. + /// The eligible order is local memory, local SSD, then remote memory. + /// A remote owner's SSD route is deliberately excluded rather than used + /// as a fallback. + LocalSsdOnlyFirst, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct OwnerLocalReserveExpectedCapacity { + pub value_len: u64, + pub payload_capacity_bytes: u64, +} + const TEST_SPEC_TCP_THREAD_REACTOR_SHARD_COUNT_MIN: u8 = 1; const TEST_SPEC_TCP_THREAD_REACTOR_SHARD_COUNT_MAX: u8 = 16; const TEST_SPEC_TCP_THREAD_BULK_LANE_COUNT_MIN: u8 = 1; const TEST_SPEC_TCP_THREAD_BULK_LANE_COUNT_MAX: u8 = 8; const TEST_SPEC_TCP_THREAD_CONTROL_LANE_COUNT_MIN: u8 = 1; const TEST_SPEC_TCP_THREAD_CONTROL_LANE_COUNT_MAX: u8 = 8; +const TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MIN: u16 = 1; +const TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MAX: u16 = 64; fn default_iceoryx_owner_client_busy_poll() -> bool { true @@ -133,6 +159,14 @@ pub struct TestSpecConfig { pub short_circuit_put_payload_path: bool, #[serde(default)] pub skip_put_end_commit: bool, + #[serde(default)] + pub ssd_read_source_policy: SsdReadSourcePolicy, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub owner_local_reserve_soft_wait_timeout_ms: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub owner_local_reserve_hard_timeout_ms: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub owner_local_reserve_expected_capacity: Option, #[serde(skip_serializing_if = "Option::is_none")] pub transport_mode: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -143,6 +177,12 @@ pub struct TestSpecConfig { pub tcp_thread_control_lane_count: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub user_rpc_sync_handler_thread_count: Option, + /// Deprecated compatibility field. Remote Put now uses direct per-generation singleflight: + /// the elected leader starts immediately and followers await its terminal result, without a + /// replica dispatcher or a global inflight queue. The value is still parsed and validated so + /// existing configurations remain loadable, but it no longer throttles remote Put execution. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replica_task_max_inflight: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub rdma_device_names: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -171,11 +211,16 @@ impl Default for TestSpecConfig { prefer_local_placement: false, short_circuit_put_payload_path: false, skip_put_end_commit: false, + ssd_read_source_policy: SsdReadSourcePolicy::default(), + owner_local_reserve_soft_wait_timeout_ms: None, + owner_local_reserve_hard_timeout_ms: None, + owner_local_reserve_expected_capacity: None, transport_mode: None, tcp_thread_reactor_shard_count: None, tcp_thread_bulk_lane_count: None, tcp_thread_control_lane_count: None, user_rpc_sync_handler_thread_count: None, + replica_task_max_inflight: None, rdma_device_names: None, require_transfer_rpc_fast_path_ready_timeout_seconds: None, enable_side_transfer: false, @@ -297,6 +342,74 @@ fn validate_required_transfer_rpc_fast_path_ready_timeout( Ok(()) } +fn validate_owner_local_reserve_timeouts(test_spec_config: &TestSpecConfig) -> KvResult<()> { + let soft_wait_timeout_ms = test_spec_config + .owner_local_reserve_soft_wait_timeout_ms + .unwrap_or(10); + let hard_timeout_ms = test_spec_config + .owner_local_reserve_hard_timeout_ms + .unwrap_or(10_000); + if soft_wait_timeout_ms == 0 { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_soft_wait_timeout_ms must be > 0" + .to_string(), + } + .into_kverror()); + } + if hard_timeout_ms == 0 { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_hard_timeout_ms must be > 0".to_string(), + } + .into_kverror()); + } + if hard_timeout_ms <= soft_wait_timeout_ms { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_hard_timeout_ms must be greater than owner_local_reserve_soft_wait_timeout_ms".to_string(), + } + .into_kverror()); + } + + Ok(()) +} + +fn validate_owner_local_reserve_expected_capacity( + test_spec_config: &TestSpecConfig, +) -> KvResult<()> { + let Some(expected) = &test_spec_config.owner_local_reserve_expected_capacity else { + return Ok(()); + }; + if expected.value_len == 0 { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_expected_capacity.value_len must be > 0" + .to_string(), + } + .into_kverror()); + } + if expected.payload_capacity_bytes == 0 { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_expected_capacity.payload_capacity_bytes must be > 0" + .to_string(), + } + .into_kverror()); + } + if crate::owner_local_reserve_expected_grant_count( + expected.value_len, + expected.payload_capacity_bytes, + ) + .is_none() + { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "test_spec_config.owner_local_reserve_expected_capacity.value_len={} cannot be represented by a local-reserve slot no larger than {} bytes", + expected.value_len, + crate::OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES + ), + } + .into_kverror()); + } + Ok(()) +} + fn apply_test_spec_rdma_device_names_to_protocol( mut protocol: ProtocolConfig, normalized_rdma_device_names: Option<&Vec>, @@ -353,6 +466,21 @@ fn validate_test_spec_tcp_thread_tuning(test_spec_config: &TestSpecConfig) -> Kv .into_kverror()); } } + if let Some(value) = test_spec_config.replica_task_max_inflight { + if !(TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MIN..=TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MAX) + .contains(&value) + { + return Err(ConfigError::InvalidTestConfig { + detail: format!( + "test_spec_config.replica_task_max_inflight must be in [{}, {}], got {}", + TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MIN, + TEST_SPEC_REPLICA_TASK_MAX_INFLIGHT_MAX, + value + ), + } + .into_kverror()); + } + } Ok(()) } @@ -444,6 +572,198 @@ pub const DEFAULT_OTLP_LOG_TABLE_NAME: &str = "fluxon_logs"; pub const DEFAULT_OTLP_LOG_FLUSH_INTERVAL_MS: u64 = 2000; pub const DEFAULT_OTLP_LOG_MAX_BATCH_LINES: usize = 2000; pub const DEFAULT_OTLP_LOG_MAX_QUEUE_LINES: usize = 20000; +pub const DEFAULT_REPLICA_CACHE_CAPACITY_RATIO: f64 = 0.95; + +fn default_replica_cache_capacity_ratio() -> f64 { + DEFAULT_REPLICA_CACHE_CAPACITY_RATIO +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ReplicaTaskPlacementPolicyKind { + LocalFirst, + Random, + QueueAware, + WeightedRoleAware, + BoundedRoleQueueAware, + PressureRoleQueueAware, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct ReplicaTaskPlacementConfigYaml { + #[serde(default = "default_replica_task_placement_policy")] + pub policy: ReplicaTaskPlacementPolicyKind, + #[serde(default = "default_replica_task_active_node_roles")] + pub active_node_roles: Vec, + #[serde(default = "default_replica_task_remote_only_node_roles")] + pub remote_only_node_roles: Vec, + #[serde(default = "default_replica_task_restrict_to_remote_only_node_roles")] + pub restrict_to_remote_only_node_roles: bool, + #[serde(default = "default_replica_task_remote_only_shard_weight")] + pub remote_only_shard_weight: f64, + #[serde(default = "default_replica_task_role_queue_window_ms")] + pub role_queue_window_ms: f64, + #[serde(default = "default_replica_task_role_pressure_gap_ms")] + pub role_pressure_gap_ms: f64, + #[serde(default = "default_replica_task_role_fabric_guard_ms")] + pub role_fabric_guard_ms: f64, + #[serde(default = "default_replica_task_role_max_shard_imbalance")] + pub role_max_shard_imbalance: f64, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ReplicaTaskPlacementConfig { + pub policy: ReplicaTaskPlacementPolicyKind, + pub active_node_roles: Vec, + pub remote_only_node_roles: Vec, + pub restrict_to_remote_only_node_roles: bool, + pub remote_only_shard_weight: f64, + pub role_queue_window_ms: f64, + pub role_pressure_gap_ms: f64, + pub role_fabric_guard_ms: f64, + pub role_max_shard_imbalance: f64, +} + +fn default_replica_task_placement_policy() -> ReplicaTaskPlacementPolicyKind { + ReplicaTaskPlacementPolicyKind::LocalFirst +} + +fn default_replica_task_active_node_roles() -> Vec { + vec!["prefill".to_string(), "decode".to_string()] +} + +fn default_replica_task_remote_only_node_roles() -> Vec { + vec!["remote_cache".to_string()] +} + +fn default_replica_task_restrict_to_remote_only_node_roles() -> bool { + true +} + +fn default_replica_task_remote_only_shard_weight() -> f64 { + 1.02 +} + +fn default_replica_task_role_queue_window_ms() -> f64 { + 1.0 +} + +fn default_replica_task_role_pressure_gap_ms() -> f64 { + 0.5 +} + +fn default_replica_task_role_fabric_guard_ms() -> f64 { + 50.0 +} + +fn default_replica_task_role_max_shard_imbalance() -> f64 { + 1.30 +} + +impl Default for ReplicaTaskPlacementConfigYaml { + fn default() -> Self { + Self { + policy: default_replica_task_placement_policy(), + active_node_roles: default_replica_task_active_node_roles(), + remote_only_node_roles: default_replica_task_remote_only_node_roles(), + restrict_to_remote_only_node_roles: + default_replica_task_restrict_to_remote_only_node_roles(), + remote_only_shard_weight: default_replica_task_remote_only_shard_weight(), + role_queue_window_ms: default_replica_task_role_queue_window_ms(), + role_pressure_gap_ms: default_replica_task_role_pressure_gap_ms(), + role_fabric_guard_ms: default_replica_task_role_fabric_guard_ms(), + role_max_shard_imbalance: default_replica_task_role_max_shard_imbalance(), + } + } +} + +impl Default for ReplicaTaskPlacementConfig { + fn default() -> Self { + ReplicaTaskPlacementConfigYaml::default() + .verify() + .expect("default replica_task_placement config must be valid") + } +} + +fn normalize_role_list(raw_roles: Vec, field_name: &str) -> KvResult> { + let mut deduped = std::collections::BTreeSet::new(); + for (idx, raw) in raw_roles.into_iter().enumerate() { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field_name}[{idx}] must be a non-empty string"), + } + .into_kverror()); + } + deduped.insert(trimmed.to_string()); + } + if deduped.is_empty() { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field_name} must contain at least one role"), + } + .into_kverror()); + } + Ok(deduped.into_iter().collect()) +} + +fn verify_finite_nonnegative_f64(value: f64, field_name: &str) -> KvResult { + if !value.is_finite() || value < 0.0 { + return Err(ConfigError::InvalidClientConfig { + detail: format!("{field_name} must be a finite non-negative number"), + } + .into_kverror()); + } + Ok(value) +} + +impl ReplicaTaskPlacementConfigYaml { + pub fn verify(self) -> KvResult { + let remote_only_shard_weight = self.remote_only_shard_weight; + if !remote_only_shard_weight.is_finite() || remote_only_shard_weight <= 0.0 { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_task_placement.remote_only_shard_weight must be a finite positive number" + .to_string(), + } + .into_kverror()); + } + let role_max_shard_imbalance = self.role_max_shard_imbalance; + if !role_max_shard_imbalance.is_finite() || role_max_shard_imbalance < 1.0 { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_task_placement.role_max_shard_imbalance must be finite and >= 1.0" + .to_string(), + } + .into_kverror()); + } + + Ok(ReplicaTaskPlacementConfig { + policy: self.policy, + active_node_roles: normalize_role_list( + self.active_node_roles, + "replica_task_placement.active_node_roles", + )?, + remote_only_node_roles: normalize_role_list( + self.remote_only_node_roles, + "replica_task_placement.remote_only_node_roles", + )?, + restrict_to_remote_only_node_roles: self.restrict_to_remote_only_node_roles, + remote_only_shard_weight, + role_queue_window_ms: verify_finite_nonnegative_f64( + self.role_queue_window_ms, + "replica_task_placement.role_queue_window_ms", + )?, + role_pressure_gap_ms: verify_finite_nonnegative_f64( + self.role_pressure_gap_ms, + "replica_task_placement.role_pressure_gap_ms", + )?, + role_fabric_guard_ms: verify_finite_nonnegative_f64( + self.role_fabric_guard_ms, + "replica_task_placement.role_fabric_guard_ms", + )?, + role_max_shard_imbalance, + }) + } +} fn verify_otlp_log_api(cfg: &mut GreptimeOtlpLogConfigYaml) -> KvResult { let endpoint = cfg.otlp_endpoint.trim(); @@ -540,6 +860,14 @@ pub struct MasterConfigYaml { pub log_dir: String, #[serde(skip_serializing_if = "Option::is_none")] pub master_ui: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replica_task_placement: Option, + #[serde(default = "default_replica_cache_capacity_ratio")] + pub replica_cache_capacity_ratio: f64, + /// Optional inclusive hot-tier capacity. Entries evicted from this tier remain in the + /// resident replica cache and are proactively copied to a remote-only owner. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replica_writeback_tier1_capacity_ratio: Option, #[serde(default)] pub test_spec_config: TestSpecConfig, } @@ -559,6 +887,9 @@ pub struct MasterConfig { pub pprof_duration_seconds: Option, pub log_dir: String, pub master_ui: Option, + pub replica_task_placement: ReplicaTaskPlacementConfig, + pub replica_cache_capacity_ratio: f64, + pub replica_writeback_tier1_capacity_ratio: Option, pub test_spec_config: TestSpecConfig, } @@ -580,6 +911,19 @@ pub struct FluxonKvSpecYaml { pub share_mem_path: String, #[serde(skip_serializing_if = "Option::is_none")] pub large_file_paths: Option, + /// Optional per-owner SSD capacity. When present, it must contain one + /// value matching the single local `large_file_paths` root. + #[serde(skip_serializing_if = "Option::is_none")] + pub large_limit_size: Option>>, + /// Optional sustained SSD write admission rate. This is a load-shedding + /// limit, not a queue: owner-local victims beyond the available budget are + /// reclaimed without an SSD copy. + #[serde(skip_serializing_if = "Option::is_none")] + pub ssd_write_rate_limit_bytes_per_sec: Option, + /// Maximum immediately admissible SSD bytes. It must be paired with + /// `ssd_write_rate_limit_bytes_per_sec`. + #[serde(skip_serializing_if = "Option::is_none")] + pub ssd_write_burst_bytes: Option, #[serde(skip_serializing_if = "Option::is_none")] pub p2p_listen_port: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -592,6 +936,13 @@ pub struct FluxonKvSpecYaml { #[serde(transparent)] pub struct LargeFilePathsYaml(pub Vec); +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct KvSsdStorageConfig { + pub limit_bytes: u64, + pub write_rate_limit_bytes_per_sec: Option, + pub write_burst_bytes: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct RedisCompatConfigYaml { @@ -609,6 +960,10 @@ pub struct ClientConfigYaml { pub contribute_to_cluster_pool_size: Option, #[serde(skip_serializing_if = "Option::is_none")] pub pprof_duration_seconds: Option, + /// Optional owner-local hot working-set threshold. Size eviction from this + /// logical tier schedules an asynchronous remote replica write-back. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replica_writeback_hot_capacity_ratio: Option, pub fluxonkv_spec: FluxonKvSpecYaml, #[serde(default)] pub test_spec_config: TestSpecConfig, @@ -714,6 +1069,14 @@ impl LargeFilePaths { "fluxon fs disk cache", ) } + + pub fn kv_ssd_storage_root(&self, cluster_name: &str, instance_key: &str) -> KvResult { + let relative_dir = PathBuf::from(format!( + "{cluster_name}_cluster_kv_ssd_storage/{}", + crate::kv_ssd_storage::safe_path_component(instance_key) + )); + self.resolve_preferred_root_subdir(&relative_dir, "kv ssd storage") + } } /// KV client backend types supported by the system @@ -731,10 +1094,12 @@ pub struct ClientConfig { pub contribute_to_cluster_pool_size: ContributeToClusterPoolSize, pub protocol: ProtocolConfig, pub pprof_duration_seconds: Option, + pub replica_writeback_hot_capacity_ratio: Option, pub redis_compat_listen_addr: Option, pub fluxonkv_spec: FluxonKvSpec, - pub share_mem_path: String, // Mandatory shared bundle path + pub share_mem_path: String, // Mandatory shared bundle path pub large_file_paths: LargeFilePaths, // Mandatory large-file roots for logs and caches + pub ssd_storage: Option, pub test_spec_config: TestSpecConfig, } @@ -913,6 +1278,8 @@ impl ClientConfigYaml { )?; materialize_default_test_spec_transport_mode(&mut test_spec_config); validate_required_transfer_rpc_fast_path_ready_timeout(&test_spec_config)?; + validate_owner_local_reserve_timeouts(&test_spec_config)?; + validate_owner_local_reserve_expected_capacity(&test_spec_config)?; validate_test_spec_tcp_thread_tuning(&test_spec_config)?; // Role selection contract: @@ -975,6 +1342,58 @@ impl ClientConfigYaml { } } + if let Some(hot_ratio) = self.replica_writeback_hot_capacity_ratio { + if !hot_ratio.is_finite() || hot_ratio <= 0.0 || hot_ratio >= 1.0 { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_writeback_hot_capacity_ratio must be finite and in (0, 1)" + .to_string(), + } + .into_kverror()); + } + if is_external { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_writeback_hot_capacity_ratio is only valid on owner configs" + .to_string(), + } + .into_kverror()); + } + } + + if let Some(expected) = &test_spec_config.owner_local_reserve_expected_capacity { + if is_external { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_expected_capacity is only valid on owner configs" + .to_string(), + } + .into_kverror()); + } + let expected_grants = crate::owner_local_reserve_expected_grant_count( + expected.value_len, + expected.payload_capacity_bytes, + ) + .expect("owner local-reserve expected capacity was validated"); + let physical_reserve_bytes = expected_grants + .checked_mul(crate::OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES) + .ok_or_else(|| { + ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_expected_capacity physical reserve size overflows u64" + .to_string(), + } + .into_kverror() + })?; + if physical_reserve_bytes > contribute_to_cluster_pool_size.dram { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "test_spec_config.owner_local_reserve_expected_capacity requires {} physical bytes across {} grants, exceeding owner dram contribution {}", + physical_reserve_bytes, + expected_grants, + contribute_to_cluster_pool_size.dram + ), + } + .into_kverror()); + } + } + let is_side_transfer_worker = matches!( test_spec_config.side_transfer_role, Some(SideTransferRole::Worker) @@ -1028,6 +1447,13 @@ impl ClientConfigYaml { } .into_kverror()); } + if self.fluxonkv_spec.large_limit_size.is_some() { + return Err(ConfigError::InvalidClientConfig { + detail: "fluxonkv_spec.large_limit_size is forbidden in zero-contribution mode" + .to_string(), + } + .into_kverror()); + } } // Preserve historical behavior for configs that omit `protocol`, but allow @@ -1170,13 +1596,15 @@ impl ClientConfigYaml { } else { let Some(large_file_paths_yaml) = self.fluxonkv_spec.large_file_paths.as_ref() else { return Err(ConfigError::InvalidClientConfig { - detail: "fluxonkv_spec.large_file_paths is required for owner mode" - .to_string(), + detail: "fluxonkv_spec.large_file_paths is required for owner mode".to_string(), } .into_kverror()); }; LargeFilePaths { - paths: verify_non_empty_root_path_list(&large_file_paths_yaml.0, "large_file_paths")?, + paths: verify_non_empty_root_path_list( + &large_file_paths_yaml.0, + "large_file_paths", + )?, } }; @@ -1204,6 +1632,81 @@ impl ClientConfigYaml { } }; + let ssd_write_rate_limit_bytes_per_sec = + self.fluxonkv_spec.ssd_write_rate_limit_bytes_per_sec; + let ssd_write_burst_bytes = self.fluxonkv_spec.ssd_write_burst_bytes; + let write_limit = match (ssd_write_rate_limit_bytes_per_sec, ssd_write_burst_bytes) { + (None, None) => None, + (Some(rate), Some(burst)) if rate > 0 && burst > 0 => Some((rate, burst)), + (Some(_), Some(_)) => { + return Err(ConfigError::InvalidClientConfig { + detail: "fluxonkv_spec SSD write rate and burst must both be positive" + .to_string(), + } + .into_kverror()); + } + _ => { + return Err(ConfigError::InvalidClientConfig { + detail: "fluxonkv_spec.ssd_write_rate_limit_bytes_per_sec and ssd_write_burst_bytes must be configured together" + .to_string(), + } + .into_kverror()); + } + }; + + let ssd_storage = if is_external { + if write_limit.is_some() { + return Err(ConfigError::InvalidClientConfig { + detail: "SSD write admission is only valid on owner configs".to_string(), + } + .into_kverror()); + } + None + } else { + match std::mem::take(&mut self.fluxonkv_spec.large_limit_size) { + None | Some(YamlNullable::Null) => { + if write_limit.is_some() { + return Err(ConfigError::InvalidClientConfig { + detail: "SSD write admission requires fluxonkv_spec.large_limit_size" + .to_string(), + } + .into_kverror()); + } + None + } + Some(YamlNullable::Value(limits)) => { + if large_file_paths.paths.len() != 1 || limits.len() != 1 { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "SSD-enabled owners require exactly one local large_file_paths root and one large_limit_size value: roots={} limits={}", + large_file_paths.paths.len(), + limits.len() + ), + } + .into_kverror()); + } + let limit_bytes = limits[0]; + if limit_bytes < crate::kv_ssd_storage::MIN_CAPACITY_BYTES { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "fluxonkv_spec.large_limit_size[0] must be at least {} bytes", + crate::kv_ssd_storage::MIN_CAPACITY_BYTES + ), + } + .into_kverror()); + } + let (write_rate_limit_bytes_per_sec, write_burst_bytes) = write_limit + .map(|(rate, burst)| (Some(rate), Some(burst))) + .unwrap_or((None, None)); + Some(KvSsdStorageConfig { + limit_bytes, + write_rate_limit_bytes_per_sec, + write_burst_bytes, + }) + } + } + }; + Ok(ClientConfig { cluster_name: fluxonkv_spec.cluster_name.clone(), etcd_addresses_raw, @@ -1211,10 +1714,12 @@ impl ClientConfigYaml { contribute_to_cluster_pool_size, protocol, pprof_duration_seconds, + replica_writeback_hot_capacity_ratio: self.replica_writeback_hot_capacity_ratio, redis_compat_listen_addr, fluxonkv_spec, share_mem_path, large_file_paths, + ssd_storage, test_spec_config, }) } @@ -1490,6 +1995,41 @@ impl MasterConfigYaml { None => None, }; + let replica_task_placement = match self.replica_task_placement { + Some(cfg) => cfg.verify()?, + None => ReplicaTaskPlacementConfig::default(), + }; + if !self.replica_cache_capacity_ratio.is_finite() + || self.replica_cache_capacity_ratio <= 0.0 + || self.replica_cache_capacity_ratio > 1.0 + { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_cache_capacity_ratio must be finite and in (0, 1]".to_string(), + } + .into_kverror()); + } + if let Some(tier1_ratio) = self.replica_writeback_tier1_capacity_ratio { + if !tier1_ratio.is_finite() + || tier1_ratio <= 0.0 + || tier1_ratio >= self.replica_cache_capacity_ratio + { + return Err(ConfigError::InvalidClientConfig { + detail: format!( + "replica_writeback_tier1_capacity_ratio must be finite, positive, and smaller than replica_cache_capacity_ratio ({})", + self.replica_cache_capacity_ratio + ), + } + .into_kverror()); + } + if !replica_task_placement.restrict_to_remote_only_node_roles { + return Err(ConfigError::InvalidClientConfig { + detail: "replica_writeback_tier1_capacity_ratio requires replica_task_placement.restrict_to_remote_only_node_roles=true" + .to_string(), + } + .into_kverror()); + } + } + let mut test_spec_config = self.test_spec_config; let transport_mode_was_explicit = test_spec_config.transport_mode.is_some(); let normalized_rdma_device_names = normalize_test_spec_rdma_device_names( @@ -1498,6 +2038,18 @@ impl MasterConfigYaml { )?; materialize_default_test_spec_transport_mode(&mut test_spec_config); validate_required_transfer_rpc_fast_path_ready_timeout(&test_spec_config)?; + validate_owner_local_reserve_timeouts(&test_spec_config)?; + validate_owner_local_reserve_expected_capacity(&test_spec_config)?; + if test_spec_config + .owner_local_reserve_expected_capacity + .is_some() + { + return Err(ConfigError::InvalidClientConfig { + detail: "test_spec_config.owner_local_reserve_expected_capacity is only valid on owner configs" + .to_string(), + } + .into_kverror()); + } validate_test_spec_tcp_thread_tuning(&test_spec_config)?; let protocol = apply_test_spec_rdma_device_names_to_protocol( self.protocol.unwrap_or(ProtocolConfig { @@ -1525,6 +2077,9 @@ impl MasterConfigYaml { network, log_dir: self.log_dir, master_ui, + replica_task_placement, + replica_cache_capacity_ratio: self.replica_cache_capacity_ratio, + replica_writeback_tier1_capacity_ratio: self.replica_writeback_tier1_capacity_ratio, test_spec_config, }) } @@ -1616,6 +2171,195 @@ fluxonkv_spec: assert!(verified.fluxonkv_spec.enable_transfer_rpc_fast_path); } + #[test] + fn owner_replica_writeback_hot_capacity_ratio_is_owner_only_and_bounded() { + let owner = r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +replica_writeback_hot_capacity_ratio: RATIO +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner_hot + large_file_paths: [/tmp/test_owner_hot_large] + sub_cluster: rack-a +"#; + let verified = ClientConfigYaml::from_str(&owner.replace("RATIO", "0.75")) + .unwrap() + .verify() + .unwrap(); + assert_eq!(verified.replica_writeback_hot_capacity_ratio, Some(0.75)); + + for invalid in ["0", "1", "-0.1", ".nan"] { + let err = ClientConfigYaml::from_str(&owner.replace("RATIO", invalid)) + .unwrap() + .verify() + .unwrap_err(); + assert!( + err.to_string() + .contains("replica_writeback_hot_capacity_ratio") + ); + } + + let external = r#" +instance_key: test_external +replica_writeback_hot_capacity_ratio: 0.75 +fluxonkv_spec: + cluster_name: test_cluster + share_mem_path: /tmp/test_external_hot +"#; + let err = ClientConfigYaml::from_str(external) + .unwrap() + .verify() + .unwrap_err(); + assert!( + err.to_string() + .contains("replica_writeback_hot_capacity_ratio is only valid on owner configs") + ); + } + + #[test] + fn owner_local_reserve_expected_capacity_accepts_exact_fit_payload_target() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 137438953472 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner_expected + large_file_paths: [/tmp/test_owner_expected_large] + sub_cluster: rack-a +test_spec_config: + owner_local_reserve_expected_capacity: + value_len: 4718592 + payload_capacity_bytes: 109951162777 +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + let expected = verified + .test_spec_config + .owner_local_reserve_expected_capacity + .unwrap(); + assert_eq!(expected.value_len, 4_718_592); + assert_eq!(expected.payload_capacity_bytes, 109_951_162_777); + assert_eq!( + crate::owner_local_reserve_expected_grant_count( + expected.value_len, + expected.payload_capacity_bytes, + ), + Some(207) + ); + } + + #[test] + fn owner_local_reserve_expected_capacity_rejects_zero_value_len() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 1073741824 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner_expected_zero + large_file_paths: [/tmp/test_owner_expected_zero_large] + sub_cluster: rack-a +test_spec_config: + owner_local_reserve_expected_capacity: + value_len: 0 + payload_capacity_bytes: 4096 +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + assert!(err.to_string().contains( + "test_spec_config.owner_local_reserve_expected_capacity.value_len must be > 0" + )); + } + + #[test] + fn owner_local_reserve_expected_capacity_rejects_physical_target_over_owner_capacity() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /tmp/test_owner_expected_too_large + large_file_paths: [/tmp/test_owner_expected_too_large_large] + sub_cluster: rack-a +test_spec_config: + owner_local_reserve_expected_capacity: + value_len: 4096 + payload_capacity_bytes: 4096 +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + assert!( + err.to_string() + .contains("exceeding owner dram contribution") + ); + } + + #[test] + fn owner_local_reserve_hard_timeout_is_validated_without_soft_override() { + let mut cfg = TestSpecConfig { + owner_local_reserve_hard_timeout_ms: Some(0), + ..Default::default() + }; + assert!( + validate_owner_local_reserve_timeouts(&cfg) + .unwrap_err() + .to_string() + .contains("owner_local_reserve_hard_timeout_ms must be > 0") + ); + + cfg.owner_local_reserve_hard_timeout_ms = Some(5); + assert!( + validate_owner_local_reserve_timeouts(&cfg) + .unwrap_err() + .to_string() + .contains("must be greater than owner_local_reserve_soft_wait_timeout_ms") + ); + + cfg.owner_local_reserve_hard_timeout_ms = Some(20); + validate_owner_local_reserve_timeouts(&cfg).unwrap(); + } + + #[test] + fn replica_task_max_inflight_is_strictly_bounded() { + for value in [1, 16, 64] { + let cfg = TestSpecConfig { + replica_task_max_inflight: Some(value), + ..Default::default() + }; + validate_test_spec_tcp_thread_tuning(&cfg).unwrap(); + } + for value in [0, 65] { + let cfg = TestSpecConfig { + replica_task_max_inflight: Some(value), + ..Default::default() + }; + assert!( + validate_test_spec_tcp_thread_tuning(&cfg) + .unwrap_err() + .to_string() + .contains("replica_task_max_inflight must be in [1, 64]") + ); + } + } + #[test] fn client_config_zero_contribution_allows_owner_bootstrapped_large_file_paths() { let cfg = ClientConfigYaml::from_str( @@ -1647,7 +2391,117 @@ fluxonkv_spec: .unwrap(); let err = cfg.verify().unwrap_err(); let text = format!("{err}"); - assert!(text.contains("fluxonkv_spec.large_file_paths is forbidden in zero-contribution mode")); + assert!( + text.contains("fluxonkv_spec.large_file_paths is forbidden in zero-contribution mode") + ); + } + + #[test] + fn owner_ssd_config_requires_one_local_root_and_one_capacity() { + let valid = ClientConfigYaml::from_str( + r#" +instance_key: test_ssd_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /mnt/nvme0/mjq_build/fluxon_config_ssd_share + large_file_paths: [/mnt/nvme0/mjq_build/fluxon_config_ssd_local] + large_limit_size: [67108864] + ssd_write_rate_limit_bytes_per_sec: 268435456 + ssd_write_burst_bytes: 67108864 + sub_cluster: rack-a +"#, + ) + .unwrap() + .verify() + .unwrap(); + assert_eq!( + valid.ssd_storage.as_ref().map(|ssd| ssd.limit_bytes), + Some(crate::kv_ssd_storage::MIN_CAPACITY_BYTES) + ); + assert_eq!( + valid + .ssd_storage + .as_ref() + .and_then(|ssd| ssd.write_rate_limit_bytes_per_sec), + Some(268435456) + ); + assert_eq!( + valid + .ssd_storage + .as_ref() + .and_then(|ssd| ssd.write_burst_bytes), + Some(67108864) + ); + + let multiple_roots = ClientConfigYaml::from_str( + r#" +instance_key: test_ssd_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /mnt/nvme0/mjq_build/fluxon_config_ssd_share + large_file_paths: + - /mnt/nvme0/mjq_build/fluxon_config_ssd_local_a + - /mnt/nvme0/mjq_build/fluxon_config_ssd_local_b + large_limit_size: [67108864] + sub_cluster: rack-a +"#, + ) + .unwrap() + .verify() + .unwrap_err(); + assert!( + multiple_roots + .to_string() + .contains("SSD-enabled owners require exactly one local large_file_paths root") + ); + + let too_small = ClientConfigYaml::from_str( + r#" +instance_key: test_ssd_owner +contribute_to_cluster_pool_size: + dram: 16777216 + vram: {} +fluxonkv_spec: + etcd_addresses: ["127.0.0.1:2379"] + cluster_name: test_cluster + share_mem_path: /mnt/nvme0/mjq_build/fluxon_config_ssd_share + large_file_paths: [/mnt/nvme0/mjq_build/fluxon_config_ssd_local] + large_limit_size: [67108863] + sub_cluster: rack-a +"#, + ) + .unwrap() + .verify() + .unwrap_err(); + assert!(too_small.to_string().contains("must be at least 67108864")); + } + + #[test] + fn zero_contribution_rejects_owner_ssd_capacity() { + let cfg = ClientConfigYaml::from_str( + r#" +instance_key: test_external +fluxonkv_spec: + cluster_name: test_cluster + share_mem_path: /mnt/nvme0/mjq_build/fluxon_config_ssd_external + large_limit_size: [67108864] +"#, + ) + .unwrap(); + assert!( + cfg.verify() + .unwrap_err() + .to_string() + .contains("large_limit_size is forbidden in zero-contribution mode") + ); } #[test] @@ -1667,7 +2521,9 @@ fluxonkv_spec: let logs_dir = large_file_paths.kv_logs_dir("test_cluster").unwrap(); assert_eq!( logs_dir, - first_root.join("child").join("test_cluster_cluster_kv_logs") + first_root + .join("child") + .join("test_cluster_cluster_kv_logs") ); assert!(logs_dir.exists()); @@ -2199,6 +3055,149 @@ test_spec_config: assert!(verified.enable_transfer_rpc_fast_path); } + #[test] + fn master_config_accepts_replica_task_placement_policy() { + let cfg = MasterConfigYaml::from_str( + r#" +instance_key: test_master +cluster_name: test_cluster +port: 18080 +etcd_endpoints: ["127.0.0.1:2379"] +network: + subnet_whitelist: ["127.0.0.0/8"] +monitoring: + prometheus_base_url: "http://127.0.0.1:4000/v1/prometheus" +log_dir: /tmp/test_master_logs +replica_task_placement: + policy: bounded_role_queue_aware + active_node_roles: [" decode ", "prefill", "decode"] + remote_only_node_roles: ["remote_cache", " mem_only "] + restrict_to_remote_only_node_roles: true + remote_only_shard_weight: 1.08 + role_queue_window_ms: 2.5 + role_pressure_gap_ms: 0.75 + role_fabric_guard_ms: 40.0 + role_max_shard_imbalance: 1.5 +"#, + ) + .unwrap(); + let verified = cfg.verify().unwrap(); + assert_eq!( + verified.replica_task_placement.policy, + ReplicaTaskPlacementPolicyKind::BoundedRoleQueueAware + ); + assert_eq!( + verified.replica_task_placement.active_node_roles, + vec!["decode".to_string(), "prefill".to_string()] + ); + assert_eq!( + verified.replica_task_placement.remote_only_node_roles, + vec!["mem_only".to_string(), "remote_cache".to_string()] + ); + assert!( + verified + .replica_task_placement + .restrict_to_remote_only_node_roles + ); + assert_eq!(verified.replica_task_placement.role_queue_window_ms, 2.5); + assert_eq!( + verified.replica_cache_capacity_ratio, + DEFAULT_REPLICA_CACHE_CAPACITY_RATIO + ); + } + + #[test] + fn master_config_validates_replica_cache_capacity_ratio() { + let base = r#" +instance_key: test_master +cluster_name: test_cluster +port: 18080 +etcd_endpoints: ["127.0.0.1:2379"] +network: + subnet_whitelist: ["127.0.0.0/8"] +monitoring: + prometheus_base_url: "http://127.0.0.1:4000/v1/prometheus" +log_dir: /tmp/test_master_logs +replica_cache_capacity_ratio: RATIO +"#; + + let valid = MasterConfigYaml::from_str(&base.replace("RATIO", "0.9")) + .unwrap() + .verify() + .unwrap(); + assert_eq!(valid.replica_cache_capacity_ratio, 0.9); + + for invalid in ["0", "-0.1", "1.01", ".nan"] { + let cfg = MasterConfigYaml::from_str(&base.replace("RATIO", invalid)).unwrap(); + let err = cfg.verify().unwrap_err(); + assert!(format!("{err}").contains("replica_cache_capacity_ratio")); + } + } + + #[test] + fn master_config_validates_replica_writeback_tier1_capacity_ratio() { + let base = r#" +instance_key: test_master +cluster_name: test_cluster +port: 18080 +etcd_endpoints: ["127.0.0.1:2379"] +network: + subnet_whitelist: ["127.0.0.0/8"] +monitoring: + prometheus_base_url: "http://127.0.0.1:4000/v1/prometheus" +log_dir: /tmp/test_master_logs +replica_cache_capacity_ratio: 0.95 +replica_writeback_tier1_capacity_ratio: RATIO +"#; + + let valid = MasterConfigYaml::from_str(&base.replace("RATIO", "0.75")) + .unwrap() + .verify() + .unwrap(); + assert_eq!(valid.replica_writeback_tier1_capacity_ratio, Some(0.75)); + + for invalid in ["0", "-0.1", "0.95", "0.96", ".nan"] { + let cfg = MasterConfigYaml::from_str(&base.replace("RATIO", invalid)).unwrap(); + let err = cfg.verify().unwrap_err(); + assert!(format!("{err}").contains("replica_writeback_tier1_capacity_ratio")); + } + + let unrestricted = base + .replace("RATIO", "0.75") + .replace( + "replica_cache_capacity_ratio: 0.95", + "replica_cache_capacity_ratio: 0.95\nreplica_task_placement:\n restrict_to_remote_only_node_roles: false", + ); + let err = MasterConfigYaml::from_str(&unrestricted) + .unwrap() + .verify() + .unwrap_err(); + assert!(format!("{err}").contains("restrict_to_remote_only_node_roles=true")); + } + + #[test] + fn master_config_rejects_invalid_replica_task_placement_weight() { + let cfg = MasterConfigYaml::from_str( + r#" +instance_key: test_master +cluster_name: test_cluster +port: 18080 +etcd_endpoints: ["127.0.0.1:2379"] +network: + subnet_whitelist: ["127.0.0.0/8"] +monitoring: + prometheus_base_url: "http://127.0.0.1:4000/v1/prometheus" +log_dir: /tmp/test_master_logs +replica_task_placement: + policy: weighted_role_aware + remote_only_shard_weight: 0.0 +"#, + ) + .unwrap(); + let err = cfg.verify().unwrap_err(); + assert!(format!("{err}").contains("remote_only_shard_weight")); + } + #[test] fn master_config_accepts_missing_port_for_auto_discovery() { let cfg = MasterConfigYaml::from_str( diff --git a/fluxon_rs/fluxon_kv/src/external_client_api/delete_ack_batch.rs b/fluxon_rs/fluxon_kv/src/external_client_api/delete_ack_batch.rs new file mode 100644 index 0000000..6ed9528 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/external_client_api/delete_ack_batch.rs @@ -0,0 +1,366 @@ +use super::ExternalClientApiView; +use ::tokio::sync::mpsc; +use limit_thirdparty::tokio; +use parking_lot::Mutex; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +pub(crate) const EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS: usize = 1024; +const EXTERNAL_DELETE_ACK_BATCH_MERGE_WINDOW: Duration = Duration::from_millis(1); +const EXTERNAL_DELETE_ACK_BATCH_RPC_TIMEOUT: Duration = Duration::from_secs(5); +const EXTERNAL_DELETE_ACK_BATCH_LOG_EVERY: u64 = 512; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ExternalDeleteAckItem { + pub external_client_id: String, + pub holder_id: u64, + pub owner_start_time: i64, +} + +#[derive(Debug, Eq, PartialEq)] +struct ExternalDeleteAckBatch { + external_client_id: String, + owner_start_time: i64, + holder_ids: Vec, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct ExternalDeleteAckBatchSnapshot { + pub enqueued_items: u64, + pub enqueue_failures: u64, + pub rpc_batches: u64, + pub rpc_items: u64, + pub released_items: u64, + pub missing_items: u64, + pub generation_mismatch_items: u64, + pub rpc_failures: u64, + pub max_batch_items: u64, +} + +#[derive(Default)] +struct ExternalDeleteAckBatchCounters { + enqueued_items: AtomicU64, + enqueue_failures: AtomicU64, + rpc_batches: AtomicU64, + rpc_items: AtomicU64, + released_items: AtomicU64, + missing_items: AtomicU64, + generation_mismatch_items: AtomicU64, + rpc_failures: AtomicU64, + max_batch_items: AtomicU64, +} + +impl ExternalDeleteAckBatchCounters { + fn snapshot(&self) -> ExternalDeleteAckBatchSnapshot { + ExternalDeleteAckBatchSnapshot { + enqueued_items: self.enqueued_items.load(Ordering::Relaxed), + enqueue_failures: self.enqueue_failures.load(Ordering::Relaxed), + rpc_batches: self.rpc_batches.load(Ordering::Relaxed), + rpc_items: self.rpc_items.load(Ordering::Relaxed), + released_items: self.released_items.load(Ordering::Relaxed), + missing_items: self.missing_items.load(Ordering::Relaxed), + generation_mismatch_items: self.generation_mismatch_items.load(Ordering::Relaxed), + rpc_failures: self.rpc_failures.load(Ordering::Relaxed), + max_batch_items: self.max_batch_items.load(Ordering::Relaxed), + } + } +} + +pub(crate) struct ExternalDeleteAckBatchHandle { + tx: mpsc::UnboundedSender, + rx: Mutex>>, + counters: ExternalDeleteAckBatchCounters, +} + +impl ExternalDeleteAckBatchHandle { + pub(crate) fn new() -> Self { + let (tx, rx) = mpsc::unbounded_channel(); + Self { + tx, + rx: Mutex::new(Some(rx)), + counters: ExternalDeleteAckBatchCounters::default(), + } + } + + pub(crate) fn enqueue(&self, item: ExternalDeleteAckItem) -> Result<(), String> { + match self.tx.send(item) { + Ok(()) => { + self.counters.enqueued_items.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + Err(err) => { + self.counters + .enqueue_failures + .fetch_add(1, Ordering::Relaxed); + Err(format!( + "external holder ACK batch worker is unavailable for holder_id={}", + err.0.holder_id + )) + } + } + } + + pub(crate) fn take_rx(&self) -> Option> { + self.rx.lock().take() + } + + pub(crate) fn snapshot(&self) -> ExternalDeleteAckBatchSnapshot { + self.counters.snapshot() + } + + fn record_batch(&self, item_count: usize) -> u64 { + let item_count = u64::try_from(item_count).unwrap_or(u64::MAX); + self.counters + .max_batch_items + .fetch_max(item_count, Ordering::Relaxed); + self.counters + .rpc_items + .fetch_add(item_count, Ordering::Relaxed); + self.counters.rpc_batches.fetch_add(1, Ordering::Relaxed) + 1 + } + + fn record_result(&self, result: super::ExternalDeleteAckBatchSendResult) { + match result { + super::ExternalDeleteAckBatchSendResult::Applied { released, missing } => { + self.counters + .released_items + .fetch_add(u64::from(released), Ordering::Relaxed); + self.counters + .missing_items + .fetch_add(u64::from(missing), Ordering::Relaxed); + } + super::ExternalDeleteAckBatchSendResult::OwnerGenerationChanged { items } => { + self.counters + .generation_mismatch_items + .fetch_add(items, Ordering::Relaxed); + } + } + } + + fn record_rpc_failure(&self) { + self.counters.rpc_failures.fetch_add(1, Ordering::Relaxed); + } +} + +fn build_external_delete_ack_batches( + items: Vec, +) -> Vec { + let mut groups: BTreeMap<(String, i64), BTreeSet> = BTreeMap::new(); + for item in items { + groups + .entry((item.external_client_id, item.owner_start_time)) + .or_default() + .insert(item.holder_id); + } + let mut batches = Vec::new(); + for ((external_client_id, owner_start_time), holder_ids) in groups { + let holder_ids: Vec<_> = holder_ids.into_iter().collect(); + for chunk in holder_ids.chunks(EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS) { + batches.push(ExternalDeleteAckBatch { + external_client_id: external_client_id.clone(), + owner_start_time, + holder_ids: chunk.to_vec(), + }); + } + } + batches +} + +fn log_external_delete_ack_batch_snapshot(view: &ExternalClientApiView, reason: &'static str) { + let snapshot = view + .external_client_api() + .inner() + .external_delete_ack_batch_snapshot(); + tracing::info!( + reason, + enqueued_items = snapshot.enqueued_items, + enqueue_failures = snapshot.enqueue_failures, + rpc_batches = snapshot.rpc_batches, + rpc_items = snapshot.rpc_items, + released_items = snapshot.released_items, + missing_items = snapshot.missing_items, + generation_mismatch_items = snapshot.generation_mismatch_items, + rpc_failures = snapshot.rpc_failures, + max_batch_items = snapshot.max_batch_items, + "external holder ACK batch snapshot" + ); +} + +pub(crate) fn spawn_external_delete_ack_batch( + view: ExternalClientApiView, + mut rx: mpsc::UnboundedReceiver, +) { + let spawn_view = view.clone(); + let worker_view = view.clone(); + spawn_view.spawn("external_delete_ack_batch", async move { + tracing::info!( + max_items = EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS, + merge_window_us = EXTERNAL_DELETE_ACK_BATCH_MERGE_WINDOW.as_micros(), + "external holder ACK batch worker started" + ); + let mut shutdown_waiter = worker_view.register_shutdown_waiter(); + loop { + let first = tokio::select! { + biased; + _ = shutdown_waiter.wait() => { + log_external_delete_ack_batch_snapshot(&worker_view, "shutdown"); + return; + } + item = rx.recv() => { + let Some(item) = item else { + log_external_delete_ack_batch_snapshot(&worker_view, "channel_closed"); + return; + }; + item + } + }; + + let mut pending = Vec::with_capacity(EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS); + pending.push(first); + let merge_window = tokio::time::sleep(EXTERNAL_DELETE_ACK_BATCH_MERGE_WINDOW); + tokio::pin!(merge_window); + let mut shutting_down = false; + while pending.len() < EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS { + tokio::select! { + biased; + _ = shutdown_waiter.wait() => { + shutting_down = true; + break; + } + item = rx.recv() => { + match item { + Some(item) => pending.push(item), + None => break, + } + } + _ = &mut merge_window => break, + } + } + if shutting_down { + tracing::info!( + dropped_pending_items = pending.len(), + "external holder ACK batch worker skipped pending items during shutdown" + ); + log_external_delete_ack_batch_snapshot(&worker_view, "shutdown_with_pending"); + return; + } + + for batch in build_external_delete_ack_batches(pending) { + let item_count = batch.holder_ids.len(); + let batch_number = worker_view + .external_client_api() + .inner() + .external_delete_ack_batch + .record_batch(item_count); + let send_result = tokio::time::timeout( + EXTERNAL_DELETE_ACK_BATCH_RPC_TIMEOUT, + worker_view + .external_client_api() + .inner() + .send_external_delete_ack_batch( + &batch.external_client_id, + batch.owner_start_time, + batch.holder_ids, + ), + ) + .await; + match send_result { + Ok(Ok(result)) => worker_view + .external_client_api() + .inner() + .external_delete_ack_batch + .record_result(result), + Ok(Err(err)) => { + worker_view + .external_client_api() + .inner() + .external_delete_ack_batch + .record_rpc_failure(); + tracing::warn!( + external_client_id = %batch.external_client_id, + owner_start_time = batch.owner_start_time, + items = item_count, + error = %err, + "external holder ACK batch RPC failed" + ); + } + Err(_) => { + worker_view + .external_client_api() + .inner() + .external_delete_ack_batch + .record_rpc_failure(); + tracing::warn!( + external_client_id = %batch.external_client_id, + owner_start_time = batch.owner_start_time, + items = item_count, + timeout_secs = EXTERNAL_DELETE_ACK_BATCH_RPC_TIMEOUT.as_secs(), + "external holder ACK batch RPC timed out" + ); + } + } + if batch_number % EXTERNAL_DELETE_ACK_BATCH_LOG_EVERY == 0 { + log_external_delete_ack_batch_snapshot(&worker_view, "periodic"); + } + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::{ + EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS, ExternalDeleteAckItem, + build_external_delete_ack_batches, + }; + + fn item(client: &str, generation: i64, holder_id: u64) -> ExternalDeleteAckItem { + ExternalDeleteAckItem { + external_client_id: client.to_string(), + holder_id, + owner_start_time: generation, + } + } + + #[test] + fn grouping_is_generation_safe_and_deduplicates_holder_ids() { + let batches = build_external_delete_ack_batches(vec![ + item("client-a", 11, 3), + item("client-a", 11, 3), + item("client-a", 11, 2), + item("client-a", 12, 4), + item("client-b", 11, 5), + ]); + assert_eq!(batches.len(), 3); + assert_eq!(batches[0].external_client_id, "client-a"); + assert_eq!(batches[0].owner_start_time, 11); + assert_eq!(batches[0].holder_ids, vec![2, 3]); + assert_eq!(batches[1].owner_start_time, 12); + assert_eq!(batches[1].holder_ids, vec![4]); + assert_eq!(batches[2].external_client_id, "client-b"); + assert_eq!(batches[2].holder_ids, vec![5]); + } + + #[test] + fn batches_never_exceed_wire_batch_limit() { + let item_count = EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS * 2 + 1; + let items = (0..item_count) + .map(|holder_id| item("client-a", 11, holder_id as u64)) + .collect(); + let batches = build_external_delete_ack_batches(items); + assert_eq!(batches.len(), 3); + assert!( + batches + .iter() + .all(|batch| batch.holder_ids.len() <= EXTERNAL_DELETE_ACK_BATCH_MAX_ITEMS) + ); + assert_eq!( + batches + .iter() + .map(|batch| batch.holder_ids.len()) + .sum::(), + item_count + ); + } +} diff --git a/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs b/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs index da701cd..e1d4162 100644 --- a/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs +++ b/fluxon_rs/fluxon_kv/src/external_client_api/external_client_test.rs @@ -3,7 +3,8 @@ use std::collections::HashMap; use crate::cluster_manager::NodeID; use crate::config::{ ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, LargeFilePaths, MasterConfig, - MonitoringConfig, ProtocolConfig, ProtocolType, TestSpecConfig, TransferEngineType, + MonitoringConfig, ProtocolConfig, ProtocolType, ReplicaTaskPlacementConfig, TestSpecConfig, + TransferEngineType, }; use crate::master_kv_router::MasterKvRouterView; use crate::{ConfigArg, run_client, run_master}; @@ -46,6 +47,9 @@ fn new_master_config( log_dir, pprof_duration_seconds: None, master_ui: None, + replica_task_placement: ReplicaTaskPlacementConfig::default(), + replica_cache_capacity_ratio: crate::config::DEFAULT_REPLICA_CACHE_CAPACITY_RATIO, + replica_writeback_tier1_capacity_ratio: None, test_spec_config: TestSpecConfig::default(), } } @@ -76,6 +80,7 @@ fn new_client_config( rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: vec![etcd.to_string()], @@ -89,6 +94,7 @@ fn new_client_config( large_file_paths: LargeFilePaths { paths: vec![format!("{}_large", shm_path)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -119,6 +125,7 @@ fn new_zero_contribution_client_config( rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: Vec::new(), @@ -130,6 +137,7 @@ fn new_zero_contribution_client_config( }, share_mem_path: shm_path.to_string(), large_file_paths: LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } diff --git a/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs b/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs index 9cb291f..d984e05 100644 --- a/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs +++ b/fluxon_rs/fluxon_kv/src/external_client_api/mod.rs @@ -1,19 +1,41 @@ +use crate::ClientTransferEngineAccessTrait; use crate::SharedJsonMeta; +use crate::client_kv_api::external_api::{ + compute_external_get_start_transfer_prefix, normalize_external_get_start_group_lens, + validate_external_get_consume_prefix, +}; use crate::client_kv_api::msg_pack::{ - ExternalInvalidateWeakIndexReq, ExternalInvalidateWeakIndexResp, + ExternalBatchDeleteAckReq, ExternalExecutePlannedGetReq, ExternalExecutePlannedGetResp, + ExternalInvalidateWeakIndexItem, ExternalInvalidateWeakIndexReq, + ExternalInvalidateWeakIndexResp, ExternalPlannedGetItem, }; use crate::client_seg_pool::{ClientSegPool, SideTransferPeerFileMeta}; +use crate::client_transfer_engine::{ClientTransferEngine, GpuMemoryGuard}; use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::cluster_manager::{ META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME, }; +use crate::master_kv_router::msg_pack::{ + BatchGetBindItemReq, BatchGetBindReq, BatchGetBindResp, BatchGetDoneReq, BatchGetDoneResp, + BatchGetPlanItemResp, BatchGetPlanReq, BatchGetPlanResp, BatchGetRevokeReq, BatchGetRevokeResp, + BatchGetStartItemResp, BatchGetStartReq, BatchGetStartResp, GetAllocationMode, GetBindTarget, + GetExternalSinkTarget, +}; use crate::rpcresp_kvresult_convert::ToResult; use crate::{ client_kv_api::msg_pack::{ - ExternalDeleteAckReq, ExternalDeleteReq, ExternalGetReq, ExternalIsExistReq, - ExternalPutCommitReq, ExternalPutCommitResp, ExternalPutStartReq, ExternalPutStartResp, - ExternalPutTransferEndReq, ExternalPutTransferEndResp, SyncKvToFileReq, SyncKvToFileResp, - TestPutPhaseTrace, + ExternalBatchGetCancelPlan, ExternalBatchGetCancelReq, ExternalBatchGetItemResp, + ExternalBatchGetLocalProbeReq, ExternalBatchGetLocalProbeResp, ExternalBatchGetReq, + ExternalBatchGetStartReq, ExternalBatchGetStartResp, ExternalBatchGetStartTransferPlan, + ExternalBatchGetTransferReq, ExternalBatchGetTransferResp, ExternalBatchIsExistReq, + ExternalBatchPutCommitItemReq, ExternalBatchPutCommitReq, ExternalBatchPutCommitResp, + ExternalBatchPutStartItemReq, ExternalBatchPutStartReq, ExternalBatchPutStartResp, + ExternalBatchPutTransferEndItemReq, ExternalBatchPutTransferEndReq, + ExternalBatchPutTransferEndResp, ExternalDeleteAckReq, ExternalDeleteReq, ExternalGetReq, + ExternalIsExistReq, ExternalObservabilitySnapshotReq, ExternalPutCommitReq, + ExternalPutCommitResp, ExternalPutRevokeReq, ExternalPutRevokeResp, ExternalPutStartReq, + ExternalPutStartResp, ExternalPutTransferEndReq, ExternalPutTransferEndResp, + SyncKvToFileReq, SyncKvToFileResp, TestPutPhaseTrace, }, cluster_manager::{ ClusterManager, ClusterManagerAccessTrait, IpcBandwidthAttributorHandle, NodeRole, @@ -22,10 +44,11 @@ use crate::{ memholder::ExternalMemHolder, p2p::{ msg_pack::{MsgPack, RPCCaller, RPCHandler}, - p2p_module::{P2pModule, P2pModuleAccessTrait}, + p2p_module::{P2pModule, P2pModuleAccessTrait, RpcTransportPolicy}, }, - rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, SharedMemError}, + rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, OK, SharedMemError}, }; +use ::tokio::sync::watch; use async_trait::async_trait; use core::panic; use dashmap::DashMap; @@ -33,6 +56,7 @@ use fluxon_commu::ShareGroupOwnerRef; use fluxon_framework::{LogicalModule, define_module}; use fluxon_observability::kv_metrics_actor::{ObserveComponent, ObserveDirection}; use fluxon_util::semaphore_map::SemaphoreMap; +use futures::{StreamExt, stream}; use libc::{MAP_SHARED, PROT_READ, PROT_WRITE, mmap}; use limit_thirdparty::tokio; use limit_thirdparty::tokio::sync::{ARwLock, Notify}; @@ -43,7 +67,7 @@ use std::{ // path::PathBuf, // 不再使用PathBuf sync::{ Arc, OnceLock, Weak, - atomic::{AtomicUsize, Ordering}, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, }, time::{Duration, Instant}, }; @@ -54,6 +78,12 @@ pub mod external_client_test; type SharedMetaSignature = fluxon_util::fs_watch::FileSignature; +mod delete_ack_batch; +pub(crate) use delete_ack_batch::{ + ExternalDeleteAckBatchHandle, ExternalDeleteAckBatchSnapshot, ExternalDeleteAckItem, + spawn_external_delete_ack_batch, +}; + // External->Owner staged put consists of multiple potentially slow components: // - ExternalPutStartReq triggers owner->master PutStart RPC (60s timeout). // - ExternalPutTransferEndReq executes transfer (can be slow) and then owner->master PutEnd RPC (60s timeout). @@ -62,9 +92,713 @@ const EXTERNAL_PUT_START_RPC_TIMEOUT_SECS: u64 = 30; const EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS: u64 = 30; const EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS: usize = 3; const EXTERNAL_PUT_TRACE_LOG_WINDOW_SECS: u64 = 10; -const EXTERNAL_INIT_CONTROL_PLANE_READY_TIMEOUT_SECS: u64 = 30; -const EXTERNAL_INIT_CONTROL_PLANE_READY_POLL_MS: u64 = 100; -const EXTERNAL_INIT_CONTROL_PLANE_READY_CONSECUTIVE_SUCCESSES: usize = 2; +const EXTERNAL_OWNER_INTRA_RPC_READY_TIMEOUT_SECS: u64 = 30; +// This is a foreground scheduler wait, not the owner operation lifetime. +// Owner finish is cancellation-safe and the uncertain replay below keeps its +// longer timeout. Fail the foreground request at the P2P minimum so SGLang +// can fall back to compute instead of parking a TP scheduler for 300 seconds. +const EXTERNAL_PLANNED_CPU_GET_FOREGROUND_RPC_TIMEOUT_SECS: u64 = + crate::p2p::msg_pack::MIN_EXPLICIT_RPC_TIMEOUT_SECS; +const EXTERNAL_PLANNED_CPU_GET_REPLAY_RPC_TIMEOUT_SECS: u64 = 300; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ExternalDeleteAckBatchSendResult { + Applied { released: u32, missing: u32 }, + OwnerGenerationChanged { items: u64 }, +} + +#[derive(Debug, Clone)] +pub struct ExternalClientGetStartResp { + pub handle: u64, + pub raw_prefix_hit_len: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalGpuDestination { + pub registration_id: u64, + pub addr: u64, + pub capacity: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalGpuGetStartResp { + pub handle: u64, + pub raw_prefix_hit_len: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalGetPlanResp { + pub handle: u64, + pub raw_prefix_hit_len: usize, + /// Prefix that can be executed by the mixed GPU path. CPU-backed sources + /// inside this prefix remain holder/H2D sources; only indices listed in + /// `gpu_remote_indices` consume GPU destinations. + pub gpu_raw_prefix_hit_len: usize, + /// Original key positions that can bind remote GPU destinations. Local + /// DRAM, requester-local SSD, and other CPU-only positions remain + /// holder/H2D sources and are absent from this vector. + pub gpu_remote_indices: Vec, +} + +pub struct ExternalGpuGetTransferResp { + pub transferred_prefix_len: usize, + pub consumed_prefix_len: usize, + pub value_ptrs: Vec, + pub local_holders: Vec>, + /// Wall time from publishing the live GPU Get handle until the transfer, + /// cleanup, and master Done path reached a terminal state. + pub transfer_wall_us: i64, + /// Time spent by the consuming call waiting for that terminal state. + pub finish_wait_us: i64, + /// Whether the terminal state was already available when consumption began. + pub terminal_before_consume: bool, + /// Ready-but-unconsumed residence when the terminal preceded consumption. + pub terminal_to_consume_us: i64, +} + +fn external_gpu_transfer_plan_geometry_is_valid( + item: &BatchGetStartItemResp, + destination: &ExternalGpuDestination, + registered_generation: u64, +) -> bool { + item.len != 0 + && item.target_addr == destination.addr + && item.target_base_addr == destination.addr + && item.len <= destination.capacity + && item.prepared_target.is_none() + && registered_generation == destination.registration_id +} + +fn external_get_plan_raw_prefixes(items: &[BatchGetPlanItemResp]) -> (usize, usize) { + external_get_plan_raw_prefixes_from_statuses( + items + .iter() + .map(|item| (item.error_code == OK, item.gpu_direct_eligible)), + ) +} + +fn external_get_plan_raw_prefixes_from_statuses( + statuses: impl IntoIterator, +) -> (usize, usize) { + let mut cpu_prefix = 0usize; + for (hit, _gpu_eligible) in statuses { + if !hit { + break; + } + cpu_prefix += 1; + } + // The GPU execution path is mixed: CPU-only sources are materialized as + // holders while later eligible remote-memory sources still bind GPU + // destinations. Therefore one CPU-only hit no longer truncates the plan. + (cpu_prefix, cpu_prefix) +} + +#[derive(Clone, Debug)] +enum ExternalGpuGetTerminal { + Completed { + planned_cpu_items: Vec, + planned_cpu_owner_start_time: Option, + }, + Revoked { + transfer_error: Option, + }, + Failed { + detail: String, + }, +} + +#[derive(Clone, Debug)] +struct ExternalGpuGetTerminalEvent { + outcome: ExternalGpuGetTerminal, + terminal_at: Instant, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct ExternalGpuGetConsumeTiming { + transfer_wall_us: i64, + finish_wait_us: i64, + terminal_before_consume: bool, + terminal_to_consume_us: i64, +} + +fn observe_external_gpu_get_consume_timing( + transfer_started_at: Instant, + terminal_at: Instant, + consume_started_at: Instant, + finish_wait: Duration, +) -> ExternalGpuGetConsumeTiming { + let terminal_to_consume = consume_started_at.checked_duration_since(terminal_at); + ExternalGpuGetConsumeTiming { + transfer_wall_us: duration_to_i64_us( + terminal_at + .checked_duration_since(transfer_started_at) + .unwrap_or_default(), + ), + finish_wait_us: duration_to_i64_us(finish_wait), + terminal_before_consume: terminal_to_consume.is_some(), + terminal_to_consume_us: duration_to_i64_us(terminal_to_consume.unwrap_or_default()), + } +} + +struct PendingExternalGpuGet { + transferable_len: usize, + atomic_group_lens: Vec, + value_ptrs: Vec, + local_holders: Vec<(usize, Arc)>, + planned_cpu_sources: Vec<(usize, String)>, + cancel_requested: Arc, + transfer_started_at: Instant, + terminal_rx: watch::Receiver>, +} + +enum PendingExternalGetPlanItem { + Local { + holder: Arc, + }, + Remote { + key: String, + plan: BatchGetPlanItemResp, + }, +} + +struct PendingExternalGetPlan { + items: Vec, + transferable_len: usize, + gpu_transferable_len: usize, + gpu_remote_indices: Vec, + atomic_group_lens: Vec, +} + +#[derive(Clone, Debug)] +enum ExternalPlannedCpuGetTerminal { + Completed { + items: Vec, + owner_start_time: i64, + }, + Revoked, + Failed { + detail: String, + }, +} + +struct PendingExternalPlannedCpuGet { + sources: Vec, + transferable_len: usize, + atomic_group_lens: Vec, + cancel_requested: Arc, + terminal_rx: watch::Receiver>, +} + +enum PendingExternalCpuSource { + Local { holder: Arc }, + Remote { key: String }, +} + +struct PendingExternalGetStart { + keys: Vec, + transferable_len: usize, + atomic_group_lens: Vec, + first_miss_index: Option, +} + +/// Keeps ownership in a pending registry across cancellation points. A +/// future that is dropped while waiting automatically restores the exact +/// entry, so a later transfer/cancel call can still drive its terminal +/// cleanup. `take()` disarms the guard once no further await can lose the +/// entry. +struct PendingRegistryEntryGuard<'a, T> { + registry: &'a DashMap, + handle: u64, + entry: Option, +} + +impl<'a, T> PendingRegistryEntryGuard<'a, T> { + fn new(registry: &'a DashMap, handle: u64, entry: T) -> Self { + Self { + registry, + handle, + entry: Some(entry), + } + } + + fn entry(&self) -> &T { + self.entry + .as_ref() + .expect("pending registry guard must be armed") + } + + fn take(mut self) -> T { + self.entry + .take() + .expect("pending registry guard must be armed") + } +} + +impl Drop for PendingRegistryEntryGuard<'_, T> { + fn drop(&mut self) { + if let Some(entry) = self.entry.take() { + self.registry.insert(self.handle, entry); + } + } +} + +#[derive(Clone)] +struct PendingInlineExternalGetStart { + keys: Vec, + items: Vec, + owner_start_time: i64, +} + +fn validate_inline_external_get_start_plan( + keys_len: usize, + items: &[ExternalBatchGetItemResp], +) -> KvResult<()> { + if items.len() != keys_len { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "inline external get_start plan length mismatch: expected={} got={}", + keys_len, + items.len() + ), + })); + } + for (idx, item) in items.iter().enumerate() { + if item.error_code != OK || item.external_memholder_info.is_none() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "inline external get_start plan item must be a hit: index={} error_code={} has_memholder={}", + idx, + item.error_code, + item.external_memholder_info.is_some() + ), + })); + } + } + Ok(()) +} + +fn validate_inline_external_get_owner_generation( + plan_owner_start_time: i64, + current_owner_start_time: i64, +) -> KvResult<()> { + if plan_owner_start_time == current_owner_start_time { + return Ok(()); + } + Err(KvError::Api(ApiError::OwnerStartTimeMismatch { + expected: current_owner_start_time, + got: plan_owner_start_time, + })) +} + +#[allow(clippy::too_many_arguments)] +fn validate_external_local_holder_geometry( + index: usize, + holder_owner_start_time: i64, + holder_offset: u64, + holder_len: u32, + holder_addr: u64, + current_owner_start_time: i64, + base_ptr: u64, + mapped_len: u64, +) -> KvResult<()> { + validate_inline_external_get_owner_generation( + holder_owner_start_time, + current_owner_start_time, + )?; + let end = holder_offset + .checked_add(u64::from(holder_len)) + .ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get local holder range overflow: index={} offset={} len={}", + index, holder_offset, holder_len + ), + }) + })?; + let pointer = base_ptr.checked_add(holder_offset).ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get local holder pointer overflow: index={} base={:#x} offset={}", + index, base_ptr, holder_offset + ), + }) + })?; + if end > mapped_len || pointer != holder_addr { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get local holder no longer matches owner mapping: index={} end={} mapped_len={} pointer={:#x} expected={:#x}", + index, end, mapped_len, pointer, holder_addr + ), + })); + } + Ok(()) +} + +fn validate_external_local_holder_mapping( + index: usize, + holder: &ExternalMemHolder, + current_owner_start_time: i64, + base_ptr: u64, + mapped_len: u64, +) -> KvResult<()> { + validate_external_local_holder_geometry( + index, + holder.owner_start_time, + holder.offset, + holder.len, + holder.addr, + current_owner_start_time, + base_ptr, + mapped_len, + ) +} + +fn validate_external_local_holders_mapping( + holders: &[(usize, Arc)], + current_owner_start_time: i64, + base_ptr: u64, + mapped_len: u64, +) -> KvResult<()> { + holders.iter().try_for_each(|(index, holder)| { + validate_external_local_holder_mapping( + *index, + holder, + current_owner_start_time, + base_ptr, + mapped_len, + ) + }) +} + +fn inline_external_get_tail_holder_ids( + items: &[ExternalBatchGetItemResp], + consume_prefix_len: usize, +) -> KvResult> { + if consume_prefix_len == 0 || consume_prefix_len > items.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "inline get_transfer consume prefix is out of range: consume={} items={}", + consume_prefix_len, + items.len() + ), + })); + } + items[consume_prefix_len..] + .iter() + .enumerate() + .map(|(tail_idx, item)| { + item.external_memholder_info + .as_ref() + .map(|info| info.holder_id) + .ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: format!( + "inline get_transfer tail item has no holder: index={}", + consume_prefix_len + tail_idx + ), + }) + }) + }) + .collect() +} + +#[cfg(test)] +mod inline_external_get_start_tests { + use super::{ + EXTERNAL_PLANNED_CPU_GET_FOREGROUND_RPC_TIMEOUT_SECS, + EXTERNAL_PLANNED_CPU_GET_REPLAY_RPC_TIMEOUT_SECS, ExternalGpuDestination, + PendingRegistryEntryGuard, external_get_plan_raw_prefixes, + external_get_plan_raw_prefixes_from_statuses, external_gpu_transfer_plan_geometry_is_valid, + inline_external_get_tail_holder_ids, observe_external_gpu_get_consume_timing, + validate_external_local_holder_geometry, validate_inline_external_get_owner_generation, + validate_inline_external_get_start_plan, validate_mixed_planned_cpu_terminal, + }; + use crate::client_kv_api::msg_pack::ExternalBatchGetItemResp; + use crate::master_kv_router::msg_pack::{BatchGetPlanItemResp, BatchGetStartItemResp}; + use crate::memholder::ExternalMemHolderInfo; + use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, OK}; + use dashmap::DashMap; + use std::time::{Duration, Instant}; + + fn inline_hit(holder_id: u64) -> ExternalBatchGetItemResp { + ExternalBatchGetItemResp { + error_code: OK, + error_json: String::new(), + external_memholder_info: Some(ExternalMemHolderInfo { + offset: holder_id * 4096, + len: 4096, + holder_id, + }), + } + } + + #[test] + fn inline_plan_requires_one_hit_per_requested_key() { + let items = vec![inline_hit(1), inline_hit(2)]; + assert!(validate_inline_external_get_start_plan(2, &items).is_ok()); + assert!(validate_inline_external_get_start_plan(3, &items).is_err()); + + let mut missing = items; + missing[1].external_memholder_info = None; + assert!(validate_inline_external_get_start_plan(2, &missing).is_err()); + } + + #[test] + fn planned_cpu_get_foreground_timeout_is_bounded_below_replay_cleanup() { + assert_eq!( + EXTERNAL_PLANNED_CPU_GET_FOREGROUND_RPC_TIMEOUT_SECS, + crate::p2p::msg_pack::MIN_EXPLICIT_RPC_TIMEOUT_SECS, + ); + assert!( + EXTERNAL_PLANNED_CPU_GET_FOREGROUND_RPC_TIMEOUT_SECS + < EXTERNAL_PLANNED_CPU_GET_REPLAY_RPC_TIMEOUT_SECS + ); + } + + #[test] + fn inline_plan_rejects_a_stale_owner_generation() { + assert!(validate_inline_external_get_owner_generation(17, 17).is_ok()); + let err = validate_inline_external_get_owner_generation(17, 18) + .expect_err("stale inline plan must not expose an old mapping"); + assert!(matches!( + err, + KvError::Api(ApiError::OwnerStartTimeMismatch { + expected: 18, + got: 17 + }) + )); + } + + #[test] + fn mixed_source_local_holder_requires_the_exact_live_mapping() { + assert!( + validate_external_local_holder_geometry( + 2, 17, 0x1000, 0x1000, 0x11_000, 17, 0x10_000, 0x4000, + ) + .is_ok() + ); + + let stale = validate_external_local_holder_geometry( + 2, 16, 0x1000, 0x1000, 0x11_000, 17, 0x10_000, 0x4000, + ) + .expect_err("an old owner generation must be rejected before GPU Bind"); + assert!(matches!( + stale, + KvError::Api(ApiError::OwnerStartTimeMismatch { + expected: 17, + got: 16 + }) + )); + + assert!( + validate_external_local_holder_geometry( + 2, 17, 0x3800, 0x1000, 0x13_800, 17, 0x10_000, 0x4000, + ) + .is_err() + ); + assert!( + validate_external_local_holder_geometry( + 2, 17, 0x1000, 0x1000, 0x21_000, 17, 0x10_000, 0x4000, + ) + .is_err() + ); + } + + #[test] + fn mixed_planned_cpu_terminal_requires_exact_owner_mapping() { + let items = vec![inline_hit(1), inline_hit(2)]; + assert!( + validate_mixed_planned_cpu_terminal(&items, 2, Some(17), 17, 0x10_000, 0x10_000,) + .is_ok() + ); + assert!( + validate_mixed_planned_cpu_terminal(&items, 1, Some(17), 17, 0x10_000, 0x10_000) + .is_err() + ); + assert!( + validate_mixed_planned_cpu_terminal(&items, 2, None, 17, 0x10_000, 0x10_000).is_err() + ); + assert!( + validate_mixed_planned_cpu_terminal(&items, 2, Some(16), 17, 0x10_000, 0x10_000,) + .is_err() + ); + assert!( + validate_mixed_planned_cpu_terminal(&items, 2, Some(17), 17, 0x10_000, 4096).is_err() + ); + } + + #[test] + fn inline_partial_consume_returns_only_tail_holder_ids() { + let items = vec![inline_hit(11), inline_hit(12), inline_hit(13)]; + assert_eq!( + inline_external_get_tail_holder_ids(&items, 2).unwrap(), + vec![13] + ); + assert!( + inline_external_get_tail_holder_ids(&items, 3) + .unwrap() + .is_empty() + ); + assert!(inline_external_get_tail_holder_ids(&items, 0).is_err()); + assert!(inline_external_get_tail_holder_ids(&items, 4).is_err()); + } + + #[test] + fn gpu_transfer_plan_accepts_zero_as_the_first_master_get_id() { + let destination = ExternalGpuDestination { + registration_id: 7, + addr: 0x1000, + capacity: 4096, + }; + let item = BatchGetStartItemResp { + get_id: 0, + target_addr: destination.addr, + target_base_addr: destination.addr, + len: destination.capacity, + ..Default::default() + }; + + assert!(external_gpu_transfer_plan_geometry_is_valid( + &item, + &destination, + destination.registration_id, + )); + } + + #[test] + fn mixed_gpu_prefix_keeps_cpu_only_sources_and_later_gpu_sources() { + let items = vec![ + BatchGetPlanItemResp { + error_code: OK, + gpu_direct_eligible: true, + ..Default::default() + }, + BatchGetPlanItemResp { + error_code: OK, + gpu_direct_eligible: false, + ..Default::default() + }, + BatchGetPlanItemResp { + error_code: OK, + gpu_direct_eligible: true, + ..Default::default() + }, + ]; + assert_eq!(external_get_plan_raw_prefixes(&items), (3, 3)); + } + + #[test] + fn owner_local_positions_do_not_consume_remote_gpu_destinations() { + // local, remote-GPU, local stays a three-page GPU-capable source + // prefix, while only the middle page needs a GPU destination. + assert_eq!( + external_get_plan_raw_prefixes_from_statuses([ + (true, true), + (true, true), + (true, true), + ]), + (3, 3) + ); + // A CPU-only source is materialized through the existing planned CPU + // holder path and does not hide a later GPU-eligible source. + assert_eq!( + external_get_plan_raw_prefixes_from_statuses([ + (true, true), + (true, false), + (true, true), + ]), + (3, 3) + ); + assert_eq!( + external_get_plan_raw_prefixes_from_statuses([ + (true, true), + (false, false), + (true, true), + ]), + (1, 1) + ); + } + + #[test] + fn gpu_get_timing_separates_ready_residence_from_real_wait() { + let transfer_started_at = Instant::now(); + let terminal_at = transfer_started_at + Duration::from_millis(20); + let consume_after_terminal = terminal_at + Duration::from_millis(30); + let ready = observe_external_gpu_get_consume_timing( + transfer_started_at, + terminal_at, + consume_after_terminal, + Duration::from_micros(7), + ); + assert_eq!(ready.transfer_wall_us, 20_000); + assert!(ready.terminal_before_consume); + assert_eq!(ready.terminal_to_consume_us, 30_000); + assert_eq!(ready.finish_wait_us, 7); + + let consume_before_terminal = transfer_started_at + Duration::from_millis(5); + let waiting = observe_external_gpu_get_consume_timing( + transfer_started_at, + terminal_at, + consume_before_terminal, + Duration::from_millis(15), + ); + assert_eq!(waiting.transfer_wall_us, 20_000); + assert!(!waiting.terminal_before_consume); + assert_eq!(waiting.terminal_to_consume_us, 0); + assert_eq!(waiting.finish_wait_us, 15_000); + } + + #[test] + fn pending_registry_guard_reinserts_on_drop_and_disarms_on_take() { + let registry = DashMap::new(); + let handle = 17; + + { + let guard = PendingRegistryEntryGuard::new(®istry, handle, "pending-a"); + assert_eq!(guard.entry(), &"pending-a"); + assert!(!registry.contains_key(&handle)); + } + assert_eq!( + registry.remove(&handle).map(|(_, value)| value), + Some("pending-a") + ); + + let guard = PendingRegistryEntryGuard::new(®istry, handle, "pending-b"); + assert_eq!(guard.take(), "pending-b"); + assert!(!registry.contains_key(&handle)); + } + + #[limit_thirdparty::tokio::test] + async fn aborted_terminal_waiter_restores_the_pending_entry() { + let registry = std::sync::Arc::new(DashMap::new()); + let handle = 23; + registry.insert(handle, "terminal-pending"); + let task_registry = registry.clone(); + let (armed_tx, armed_rx) = ::tokio::sync::oneshot::channel(); + let waiter = ::tokio::spawn(async move { + let (_, entry) = task_registry + .remove(&handle) + .expect("test pending entry must exist"); + let _guard = PendingRegistryEntryGuard::new(&task_registry, handle, entry); + let _ = armed_tx.send(()); + futures::future::pending::<()>().await; + }); + + armed_rx.await.expect("waiter armed its guard"); + assert!(!registry.contains_key(&handle)); + waiter.abort(); + assert!( + waiter + .await + .expect_err("waiter must be aborted") + .is_cancelled() + ); + assert_eq!( + registry.remove(&handle).map(|(_, value)| value), + Some("terminal-pending") + ); + } +} fn duration_to_i64_us(duration: std::time::Duration) -> i64 { duration.as_micros().min(i64::MAX as u128) as i64 @@ -195,9 +929,9 @@ struct SharedMemoryPtr { /// Length of the mapping in bytes len: u64, /// Base directory of the shared-memory bundle (used to locate shared.json/mmap.file) - _path: String, + path: String, /// Handle to the mmap backing file. Keeping the FD open is harmless and simplifies lifecycle. - _file: File, + file: File, /// Metadata signature read from shared.json for change detection. memory_signature: SharedMetaSignature, } @@ -218,8 +952,8 @@ impl SharedMemoryPtr { ptr_rw, ptr_ro, len, - _path: path, - _file: file, + path, + file, memory_signature, } } @@ -245,14 +979,15 @@ define_module!( ExternalClientApi, (external_client_api, ExternalClientApi), (p2p, P2pModule), - (cluster_manager, ClusterManager) + (cluster_manager, ClusterManager), + (client_transfer_engine, ClientTransferEngine) ); /// External Client configuration parameters #[derive(Clone, Debug)] pub struct ExternalClientApiNewArg { - pub share_mem_path: String, - pub large_file_paths: crate::config::LargeFilePaths, + pub shared_memory_path: String, + pub shared_file_path: String, pub expected_cluster_name: String, pub expected_protocol_version: String, pub enable_side_transfer: bool, @@ -310,1239 +1045,4579 @@ pub struct ExternalInner { initial_sub_cluster: OnceLock>, expected_cluster_name: String, expected_protocol_version: String, - external_share_mem_path: String, - external_large_file_paths: crate::config::LargeFilePaths, - _enable_side_transfer: bool, + external_shared_memory_path: String, + external_shared_file_path: String, + enable_side_transfer: bool, short_circuit_put_payload_path: bool, side_rr_next: AtomicUsize, side_transfer_put_bindings: moka::sync::SegmentedCache<(u64, u32), (String, u16)>, rpc_caller_external_get: RPCCaller, + rpc_caller_external_batch_get: RPCCaller, + rpc_caller_external_batch_get_local_probe: RPCCaller, + rpc_caller_external_batch_get_start: RPCCaller, + rpc_caller_external_batch_get_transfer: RPCCaller, + rpc_caller_external_batch_get_cancel: RPCCaller, + rpc_caller_master_batch_get_start: RPCCaller, + rpc_caller_master_batch_get_plan: RPCCaller, + rpc_caller_master_batch_get_bind: RPCCaller, + rpc_caller_master_batch_get_done: RPCCaller, + rpc_caller_master_batch_get_revoke: RPCCaller, + rpc_caller_external_execute_planned_get: RPCCaller, rpc_caller_external_put_commit: RPCCaller, + rpc_caller_external_batch_put_commit: RPCCaller, rpc_caller_external_put_start: RPCCaller, + rpc_caller_external_batch_put_start: RPCCaller, rpc_caller_external_put_transfer_end: RPCCaller, + rpc_caller_external_batch_put_transfer_end: RPCCaller, rpc_caller_external_delete: RPCCaller, rpc_caller_external_is_exist: RPCCaller, + rpc_caller_external_batch_is_exist: RPCCaller, + rpc_caller_external_observability_snapshot: RPCCaller, rpc_caller_external_delete_ack: RPCCaller, + rpc_caller_external_batch_delete_ack: RPCCaller, + rpc_caller_external_put_revoke: RPCCaller, /// Lease RPC callers for external mode - _rpc_caller_allocate_client_lease: RPCCaller, - _rpc_caller_client_lease_keepalive: RPCCaller, + rpc_caller_allocate_client_lease: RPCCaller, + rpc_caller_client_lease_keepalive: RPCCaller, /// key -> Weak index (dashmap-based) key_weak_memholder_index: DashMap>, + pending_external_get_start: DashMap, + /// Fully local plans consumed without a follow-up owner transfer RPC. + pending_inline_external_get_start: DashMap, + next_gpu_get_handle: AtomicU64, + pending_external_get_plan: DashMap, + pending_external_gpu_get: DashMap, + pending_external_planned_cpu_get: DashMap, /// per-key semaphore (permits=1) to ensure single inflight per key inflight1_per_key: SemaphoreMap, put_trace_log_window: Mutex, + pub(crate) external_delete_ack_batch: ExternalDeleteAckBatchHandle, } pub struct ExternalClientApi(ExternalInner); -impl ExternalClientApi { - /// Access inner external-only API. Safe to unwrap in external role. - pub fn inner(&self) -> &ExternalInner { - &self.0 +async fn wait_external_gpu_get_terminal( + mut terminal_rx: watch::Receiver>, +) -> KvResult { + loop { + if let Some(terminal) = terminal_rx.borrow().clone() { + return Ok(terminal); + } + if terminal_rx.changed().await.is_err() { + return Err(KvError::Api(ApiError::Unknown { + detail: "GPU Get transfer task ended without publishing a terminal state" + .to_string(), + })); + } } +} - pub fn attach_view(&self, view: ExternalClientApiView) { - // This module is constructed only for the external variant; view attachment is - // therefore an invariant. - self.inner().view.attach(view); +async fn wait_external_planned_cpu_get_terminal( + mut terminal_rx: watch::Receiver>, +) -> KvResult { + loop { + if let Some(terminal) = terminal_rx.borrow().clone() { + return Ok(terminal); + } + if terminal_rx.changed().await.is_err() { + return Err(KvError::Api(ApiError::Unknown { + detail: "planned CPU Get task ended without publishing a terminal state" + .to_string(), + })); + } } +} - pub async fn construct(arg: ExternalClientApiNewArg) -> Result { - tracing::info!( - "Constructing ExternalClientApi in ExternalClient mode (PreView): shm_dir={}", - arg.share_mem_path - ); - - Ok(Self(ExternalInner { - view: ExternalClientApiViewHolder::new(), - current_owner: ARwLock::new(None), - owner_remap_notify: Arc::new(Notify::new()), - wait_owner_gate: ARwLock::new(()), - initial_sub_cluster: OnceLock::new(), - expected_cluster_name: arg.expected_cluster_name, - expected_protocol_version: arg.expected_protocol_version, - external_share_mem_path: arg.share_mem_path, - external_large_file_paths: arg.large_file_paths, - _enable_side_transfer: arg.enable_side_transfer, - short_circuit_put_payload_path: arg.short_circuit_put_payload_path, - side_rr_next: AtomicUsize::new(0), - side_transfer_put_bindings: moka::sync::Cache::builder() - .time_to_live(Duration::from_secs(10 * 60)) - .segments(16) - .build(), - rpc_caller_external_get: RPCCaller::::new(), - rpc_caller_external_put_commit: RPCCaller::::new(), - rpc_caller_external_put_start: RPCCaller::::new(), - rpc_caller_external_put_transfer_end: RPCCaller::::new(), - rpc_caller_external_delete: RPCCaller::::new(), - rpc_caller_external_is_exist: RPCCaller::::new(), - rpc_caller_external_delete_ack: RPCCaller::::new(), - _rpc_caller_allocate_client_lease: RPCCaller::::new(), - _rpc_caller_client_lease_keepalive: RPCCaller::::new(), - key_weak_memholder_index: DashMap::new(), - inflight1_per_key: SemaphoreMap::new(1, std::time::Duration::from_secs(120)), - put_trace_log_window: Mutex::new(ExternalPutTraceLogWindow::new()), - })) +async fn run_planned_get_revoke_cleanup( + view: ExternalClientApiView, + get_ids: Vec, + context: &'static str, +) -> KvResult<()> { + if get_ids.is_empty() { + return Ok(()); } - - pub async fn init2_prepare(&self) -> Result<(), KvError> { - // Prepare external client api initialization without waiting for owner readiness. - // - // All owner readiness (shared.json + mmap.file + membership observation) is handled by - // the init resource hook `owner_shared_mem_bundle_ready`. - Ok(()) + let mut attempt = 1u32; + let mut shutdown = view.register_shutdown_waiter(); + loop { + match view + .external_client_api() + .inner() + .master_batch_gpu_get_revoke(get_ids.clone()) + .await + { + Ok(()) => return Ok(()), + Err(err) if matches!(&err, KvError::Api(ApiError::SystemShutdown { .. })) => { + return Err(err); + } + Err(err) => { + tracing::warn!( + "{} planned Get Revoke uncertain; retaining cleanup ownership: items={} attempt={} err={}", + context, + get_ids.len(), + attempt, + err + ); + } + } + attempt = attempt.saturating_add(1); + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis( + (50u64.saturating_mul(1u64 << attempt.min(6))).min(2_000), + )) => {} + _ = shutdown.wait() => { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: format!( + "{} planned Get Revoke cleanup stopped during shutdown", + context + ), + })); + } + } } +} - pub(crate) async fn wait_owner_shared_mem_bundle_ready_for_init_resource( - &self, - ) -> Result<(), KvError> { - let ext = &self.0; +fn spawn_planned_get_revoke_cleanup( + view: ExternalClientApiView, + get_ids: Vec, + context: &'static str, +) -> ::tokio::sync::oneshot::Receiver> { + let (done_tx, done_rx) = ::tokio::sync::oneshot::channel(); + let spawn_view = view.clone(); + let worker_view = view; + spawn_view.spawn("planned_get_revoke_cleanup", async move { + let result = run_planned_get_revoke_cleanup(worker_view, get_ids, context).await; + let _ = done_tx.send(result); + }); + done_rx +} - if ext.current_owner.read().await.is_none() { - // Initial attach: accept the current shared.json without requiring a post-wait write_ts. - let wait_start_ts = i64::MIN; - let OwnerRestartPayload { meta, signature } = task_wait_owner_restart( - ext.view.clone_view(), - ext.external_share_mem_path.clone(), - None, - wait_start_ts, - None, - ext.expected_cluster_name.clone(), - ext.expected_protocol_version.clone(), - ) - .await?; +async fn finish_planned_get_revoke_cleanup( + view: ExternalClientApiView, + get_ids: Vec, + context: &'static str, +) -> KvResult<()> { + spawn_planned_get_revoke_cleanup(view, get_ids, context) + .await + .map_err(|_| { + KvError::Api(ApiError::Unknown { + detail: format!( + "{} planned Get Revoke task ended without publishing a terminal", + context + ), + }) + })? +} - let shared_memory_ptr = ExternalInner::init_shared_memory_from_meta( - &ext.external_share_mem_path, - &meta, - signature, - )?; +/// Owns master plan identities until a durable local pending entry or a +/// registered cleanup task takes over. This closes the cancellation window +/// around late Bind RPCs without adding a normal-path RPC. +struct PlannedGetRevokeGuard { + view: ExternalClientApiView, + get_ids: Option>, + context: &'static str, +} - ext.initial_sub_cluster - .set(meta.sub_cluster.clone()) - .unwrap(); - *ext.current_owner.write().await = Some(CurrentOwner { - node_id: meta.owner_id.clone(), - owner_start_time: meta.node_start_time, - shared_memory: shared_memory_ptr, - }); - ext.owner_remap_notify.notify_waiters(); +impl PlannedGetRevokeGuard { + fn new(view: ExternalClientApiView, get_ids: Vec, context: &'static str) -> Self { + Self { + view, + get_ids: Some(get_ids), + context, } + } - // Make the resource include the cluster membership observation as well. - self.init3_wait_owner_present().await?; - Ok(()) + fn disarm(&mut self) { + self.get_ids = None; } +} - pub async fn init2_after_owner_shared_mem_bundle_ready(&self) -> Result<(), KvError> { - let ext = &self.0; +impl Drop for PlannedGetRevokeGuard { + fn drop(&mut self) { + let Some(get_ids) = self.get_ids.take() else { + return; + }; + if get_ids.is_empty() { + return; + } + drop(spawn_planned_get_revoke_cleanup( + self.view.clone(), + get_ids, + self.context, + )); + } +} - let owner_id = ext.shared_storage_node_id().await.expect( - "ExternalClientApi expects current_owner to be Some after owner_shared_mem_bundle_ready", - ); +fn release_planned_cpu_response_holders( + inner: &ExternalInner, + response: &ExternalExecutePlannedGetResp, + owner_start_time: i64, +) { + release_planned_cpu_item_holders(inner, &response.items, owner_start_time); +} - // English note: - // Register inbound RPC handlers before any awaited etcd operations that publish or mutate - // member metadata. Otherwise, other nodes can observe this member and send RPCs while the - // handler set is still incomplete, leading to transient "No handler found" drops. - // - // Owner binding (current_owner) is already established by the init resource - // `owner_shared_mem_bundle_ready`, so handler registration is safe here. - ext.rpc_caller_external_get.regist(ext.view.p2p_module()); - ext.rpc_caller_external_put_commit - .regist(ext.view.p2p_module()); - ext.rpc_caller_external_put_start - .regist(ext.view.p2p_module()); - ext.rpc_caller_external_put_transfer_end - .regist(ext.view.p2p_module()); - ext.rpc_caller_external_delete.regist(ext.view.p2p_module()); - ext.rpc_caller_external_is_exist - .regist(ext.view.p2p_module()); - ext.rpc_caller_external_delete_ack - .regist(ext.view.p2p_module()); - crate::key_prefix::init_for_p2p_owner(ext.view.p2p_module()); - crate::kvlease::init_for_p2p_owner(ext.view.p2p_module()); - crate::metrics::client::init_for_p2p_owner(ext.view.p2p_module()); +fn release_planned_cpu_item_holders( + inner: &ExternalInner, + items: &[ExternalBatchGetItemResp], + owner_start_time: i64, +) { + let external_client_id = inner.view.cluster_manager().get_self_info().id; + for holder_id in items.iter().filter_map(|item| { + item.external_memholder_info + .as_ref() + .map(|info| info.holder_id) + }) { + if let Err(err) = inner.enqueue_external_delete_ack( + external_client_id.clone(), + holder_id, + owner_start_time, + ) { + tracing::warn!( + "planned CPU Get could not enqueue unused holder release: holder_id={} err={}", + holder_id, + err + ); + } + } +} - let view_ext = ext.view.clone_view(); - RPCHandler::::new().regist( - ext.view.p2p_module(), - move |resp, msg| { - let view = view_ext.clone(); - let view_task = view.clone(); - let _ = view.spawn("rpc_external_invalidate_weak_index", async move { - let result = handle_external_invalidate_weak_index(&view_task, &msg).await; - let _ = resp.send_resp(result).await; - }); - Ok(()) - }, +fn release_optional_planned_cpu_item_holders( + inner: &ExternalInner, + items: &[ExternalBatchGetItemResp], + owner_start_time: Option, +) { + if items.is_empty() { + return; + } + let Some(owner_start_time) = owner_start_time else { + tracing::error!( + items = items.len(), + "mixed Get terminal lost the owner generation needed to release CPU holders" ); + return; + }; + release_planned_cpu_item_holders(inner, items, owner_start_time); +} - RPCCaller::::new().regist(ext.view.p2p_module()); - let view_ext = ext.view.clone_view(); - RPCHandler::::new().regist(ext.view.p2p_module(), move |resp, msg| { - let view = view_ext.clone(); - let view_task = view.clone(); - let _ = view.spawn("rpc_sync_kv_to_file", async move { - let result = handle_sync_kv_to_file_external(&view_task, &msg).await; - let _ = resp.send_resp(result).await; - }); - Ok(()) - }); - tracing::info!("ExternalClientApi RPC callers registered"); +fn validate_mixed_planned_cpu_terminal( + items: &[ExternalBatchGetItemResp], + expected_items: usize, + owner_start_time: Option, + current_owner_start_time: i64, + base_ptr: u64, + mapped_len: u64, +) -> KvResult { + if items.len() != expected_items { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get planned CPU terminal length mismatch: expected={} got={}", + expected_items, + items.len() + ), + })); + } + let owner_start_time = owner_start_time.ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: "mixed Get planned CPU terminal omitted its owner generation".to_string(), + }) + })?; + validate_inline_external_get_owner_generation(owner_start_time, current_owner_start_time)?; + for (index, item) in items.iter().enumerate() { + if item.error_code != OK { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get planned CPU item failed: index={} error_code={} error_json={}", + index, item.error_code, item.error_json + ), + })); + } + let Some(info) = item.external_memholder_info.as_ref() else { + return Err(KvError::Api(ApiError::Unknown { + detail: format!("mixed Get planned CPU item has no holder: index={index}"), + })); + }; + let end = info + .offset + .checked_add(u64::from(info.len)) + .ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get planned CPU holder range overflow: index={} offset={} len={}", + index, info.offset, info.len + ), + }) + })?; + if end > mapped_len || base_ptr.checked_add(info.offset).is_none() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get planned CPU holder is outside owner mapping: index={} end={} mapped_len={} base={:#x} offset={}", + index, end, mapped_len, base_ptr, info.offset + ), + })); + } + } + Ok(owner_start_time) +} - ext.view - .cluster_manager() - .set_self_share_group_binding(ShareGroupOwnerRef { - owner_id: owner_id.clone(), - owner_start_time: ext.current_owner_start_time().await, - }) - .await?; - ext.view - .cluster_manager() - .set_self_sub_cluster(ext.initial_sub_cluster.get().unwrap().clone()) - .await - .map_err(KvError::from)?; +fn spawn_uncertain_planned_cpu_get_cleanup( + view: ExternalClientApiView, + owner: String, + request: MsgPack, + owner_start_time: i64, +) { + let spawn_view = view.clone(); + let task_view = view.clone(); + spawn_view.spawn("uncertain_planned_cpu_get_cleanup", async move { + let mut shutdown = task_view.register_shutdown_waiter(); + loop { + let inner = task_view.external_client_api().inner(); + if inner.current_owner_start_time().await != owner_start_time { + return; + } + let attempt = inner + .rpc_caller_external_execute_planned_get + .call( + inner.view.p2p_module(), + owner.clone().into(), + request.clone(), + Some(Duration::from_secs( + EXTERNAL_PLANNED_CPU_GET_REPLAY_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await; + if let Ok(response) = attempt { + release_planned_cpu_response_holders( + inner, + &response.serialize_part, + owner_start_time, + ); + return; + } + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(1)) => {} + _ = shutdown.wait() => return, + } + } + }); +} +async fn run_external_planned_cpu_get( + view: ExternalClientApiView, + plan_handle: u64, + plan_items: Vec<(String, u64)>, + skipped_get_ids: Vec, + transfer_concurrency: usize, + cancel_requested: Arc, +) -> ExternalPlannedCpuGetTerminal { + let inner = view.external_client_api().inner(); + let mut all_plan_get_ids = skipped_get_ids.clone(); + all_plan_get_ids.extend(plan_items.iter().map(|(_, get_id)| *get_id)); + if let Err(err) = inner.master_batch_gpu_get_revoke(skipped_get_ids).await { + let cleanup = finish_planned_get_revoke_cleanup( + view.clone(), + all_plan_get_ids, + "planned CPU tail failure", + ) + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + return ExternalPlannedCpuGetTerminal::Failed { + detail: format!( + "planned CPU Get could not revoke its unconsumed tail: {err}; cleanup_error={cleanup:?}" + ), + }; + } + if cancel_requested.load(Ordering::Acquire) { + let get_ids = plan_items.iter().map(|(_, get_id)| *get_id).collect(); + return match finish_planned_get_revoke_cleanup( + view.clone(), + get_ids, + "planned CPU pre-owner cancel", + ) + .await { - let view = ext.view.clone_view(); - let view_task = view.clone(); - let _ = view.spawn("external_owner_remap_actor", async move { - let shutdown_poller = view_task.register_shutdown_poller(); - let mut cluster_rx = view_task.cluster_manager().listen(); - let mut tick = tokio::time::interval(Duration::from_millis(200)); + Ok(()) => ExternalPlannedCpuGetTerminal::Revoked, + Err(err) => ExternalPlannedCpuGetTerminal::Failed { + detail: format!("planned CPU Get cancel cleanup failed: {err}"), + }, + }; + } - loop { - if !shutdown_poller.is_running() { - tracing::info!("external owner remap actor stopped by shutdown"); - break; - } + if plan_items.is_empty() { + return ExternalPlannedCpuGetTerminal::Completed { + items: Vec::new(), + owner_start_time: inner.current_owner_start_time().await, + }; + } - let Some(view_guard) = view_task.try_upgrade() else { - tracing::info!( - "external owner remap actor stopped because view was dropped" - ); - break; - }; - let _keep_view_alive = view_guard; + let Some(owner) = inner.shared_storage_node_id().await else { + let get_ids = plan_items.iter().map(|(_, get_id)| *get_id).collect(); + let cleanup = + finish_planned_get_revoke_cleanup(view.clone(), get_ids, "planned CPU missing owner") + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + return ExternalPlannedCpuGetTerminal::Failed { + detail: format!( + "planned CPU Get has no current share-group owner; cleanup_error={cleanup:?}" + ), + }; + }; + let owner_start_time = inner.current_owner_start_time().await; + let external_client_id = inner.view.cluster_manager().get_self_info().id; + let request = MsgPack { + serialize_part: ExternalExecutePlannedGetReq { + plan_handle, + items: plan_items + .iter() + .map(|(key, get_id)| ExternalPlannedGetItem { + key: key.clone(), + get_id: *get_id, + }) + .collect(), + req_node_id: external_client_id, + started_time: owner_start_time, + transfer_concurrency, + }, + raw_bytes: Vec::new(), + }; + let response = match inner + .rpc_caller_external_execute_planned_get + .call( + inner.view.p2p_module(), + owner.clone().into(), + request.clone(), + Some(Duration::from_secs( + EXTERNAL_PLANNED_CPU_GET_FOREGROUND_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + { + Ok(response) => response.serialize_part, + Err(err) => { + spawn_uncertain_planned_cpu_get_cleanup(view.clone(), owner, request, owner_start_time); + return ExternalPlannedCpuGetTerminal::Failed { + detail: format!( + "planned CPU Get owner RPC failed; replay cleanup continues in background: error={}", + KvError::from(err) + ), + }; + } + }; + if response.error_code != OK || response.items.len() != plan_items.len() { + release_planned_cpu_response_holders(inner, &response, owner_start_time); + return ExternalPlannedCpuGetTerminal::Failed { + detail: format!( + "planned CPU Get owner response failed or changed shape: error_code={} expected={} got={} error_json={}", + response.error_code, + plan_items.len(), + response.items.len(), + response.error_json + ), + }; + } + if let Some((index, item)) = response + .items + .iter() + .enumerate() + .find(|(_, item)| item.error_code != OK || item.external_memholder_info.is_none()) + { + release_planned_cpu_response_holders(inner, &response, owner_start_time); + return ExternalPlannedCpuGetTerminal::Failed { + detail: format!( + "planned CPU Get item failed: index={} error_code={} error_json={}", + index, item.error_code, item.error_json + ), + }; + } + if cancel_requested.load(Ordering::Acquire) { + release_planned_cpu_response_holders(inner, &response, owner_start_time); + return ExternalPlannedCpuGetTerminal::Revoked; + } + ExternalPlannedCpuGetTerminal::Completed { + items: response.items, + owner_start_time, + } +} - if let Err(err) = view_task - .external_client_api() - .inner() - .try_background_owner_remap_once() - .await - { - tracing::warn!("external owner remap actor probe failed: {}", err); - } +async fn run_external_gpu_get_transfer( + view: ExternalClientApiView, + transfer_items: Vec<(BatchGetStartItemResp, GpuMemoryGuard)>, + skipped_get_ids: Vec, + transfer_concurrency: usize, + cancel_requested: Arc, +) -> ExternalGpuGetTerminal { + let transfer_get_ids = transfer_items + .iter() + .map(|(item, _)| item.get_id) + .collect::>(); + let mut all_get_ids = skipped_get_ids.clone(); + all_get_ids.extend(transfer_get_ids.iter().copied()); + + let skipped_revoke_error = view + .external_client_api() + .inner() + .master_batch_gpu_get_revoke(skipped_get_ids) + .await + .err() + .map(|err| err.to_string()); - tokio::select! { - _ = tick.tick() => {} - recv = cluster_rx.recv() => { - if recv.is_err() { - sleep(Duration::from_millis(200)).await; - cluster_rx = view_task.cluster_manager().listen(); - } - } - } - } - }); - } + if transfer_items.is_empty() { + return match skipped_revoke_error { + Some(detail) => ExternalGpuGetTerminal::Failed { + detail: format!("GPU Get could not revoke non-transferable starts: {detail}"), + }, + None => ExternalGpuGetTerminal::Revoked { + transfer_error: None, + }, + }; + } - // Attribute local IPC bandwidth to the owner daemon (machine-level view). - // - // Causal chain: - // - External<->external traffic can use the local IPC tier (iceoryx2) when both are in the - // same share-group (same owner_id + local_ipc_root). - // - Topology aggregates bandwidth at the owner/machine level, so local IPC bytes must be - // charged to the owner, otherwise the UI under-reports throughput. - // - We keep the P2P hot path allocation-free by recording bytes into atomics, and flush - // them periodically via a background task. - { - let cm = ext.view.cluster_manager(); - let handle = IpcBandwidthAttributorHandle::new(); - cm.attach_ipc_bandwidth_attributor_handle(handle.clone()); - if let Some(observe) = cm.observe_handle().cloned() { - let self_member_id = cm.self_member_id().to_string(); - let owner_role = NodeRole::Client.to_string(); - let owner_id_for_task = owner_id.clone(); - let view_task = ext.view.clone_view(); - let view_task2 = view_task.clone(); - let _ = view_task.spawn("ipc_bandwidth_attributor", async move { - let mut shutdown_waiter = view_task2.register_shutdown_waiter(); - let mut interval = tokio::time::interval(Duration::from_secs( - crate::metric_reporter::METRICS_FLUSH_INTERVAL_SECS, - )); - - loop { - tokio::select! { - _ = interval.tick() => { - let tx_bytes = handle.take_tx_bytes(); - if tx_bytes > 0 { - observe.try_record_peer_network_bytes_override( - ObserveComponent::LocalIpc, - owner_id_for_task.as_str(), - owner_role.as_str(), - self_member_id.as_str(), - ObserveDirection::Tx, - tx_bytes, - ); - } - let rx_bytes = handle.take_rx_bytes(); - if rx_bytes > 0 { - observe.try_record_peer_network_bytes_override( - ObserveComponent::LocalIpc, - owner_id_for_task.as_str(), - owner_role.as_str(), - self_member_id.as_str(), - ObserveDirection::Rx, - rx_bytes, - ); - } - } - _ = shutdown_waiter.wait() => { - break; - } - } - } - }); - } else { - tracing::info!( - "ExternalClientApi local IPC bandwidth attribution disabled: ObserveHandle not attached" - ); + let transfer_futures = transfer_items.into_iter().map(|(item, gpu_guard)| { + let transfer_view = view.clone(); + let cancel_requested = cancel_requested.clone(); + async move { + if cancel_requested.load(Ordering::Acquire) { + return None; } + transfer_view + .client_transfer_engine() + .transfer_data_no_copy_to_gpu( + item.node_id.clone(), + item.src_addr, + item.target_addr, + item.len, + gpu_guard, + ) + .await + .err() + .map(|err| { + format!( + "GPU Get transfer failed: get_id={} source={} src={:#x} target={:#x} len={} error={}", + item.get_id, + item.node_id, + item.src_addr, + item.target_addr, + item.len, + err + ) + }) + } + }); + let mut transfer_stream = + stream::iter(transfer_futures).buffer_unordered(transfer_concurrency.max(1)); + let mut transfer_error = None; + while let Some(item_error) = transfer_stream.next().await { + if transfer_error.is_none() { + transfer_error = item_error; } - Ok(()) } - pub async fn init3_wait_owner_present(&self) -> Result<(), KvError> { - let ext = &self.0; - let owner_id = ext - .shared_storage_node_id() - .await - .expect("external role expects current_owner to be Some after init2"); - let owner_start_time = ext.current_owner_start_time().await; - let cm = ext.view.cluster_manager(); - if cm - .get_member_info_cached(&owner_id) - .map(|member| member.node_start_time == owner_start_time) - .unwrap_or(false) + let cancelled = cancel_requested.load(Ordering::Acquire); + if cancelled || transfer_error.is_some() || skipped_revoke_error.is_some() { + if let Err(err) = view + .external_client_api() + .inner() + .master_batch_gpu_get_revoke(all_get_ids) + .await { - return Ok(()); + return ExternalGpuGetTerminal::Failed { + detail: format!( + "GPU Get cleanup failed after transfer/cancel: transfer_error={:?} skipped_revoke_error={:?} revoke_error={}", + transfer_error, skipped_revoke_error, err + ), + }; } - - tracing::info!( - "External init: waiting for owner generation to join (owner_id={} owner_start_time={})", - owner_id, - owner_start_time - ); - let mut rx = cm.listen(); - loop { - if cm - .get_member_info_cached(&owner_id) - .map(|member| member.node_start_time == owner_start_time) - .unwrap_or(false) - { - tracing::info!( - "External init: owner generation observed (owner_id={} owner_start_time={})", - owner_id, - owner_start_time - ); - return Ok(()); - } - match rx.recv().await { - Ok(_ev) => { - // Yield once to allow watcher to update member cache after emitting an event. - limit_thirdparty::tokio::task::yield_now().await; - } - Err(e) => { - return Err(KvError::Api(ApiError::Unknown { - detail: format!( - "cluster event channel closed while waiting for owner generation (owner_id={} owner_start_time={}): {}", - owner_id, owner_start_time, e - ), - })); - } - } + if let Some(detail) = skipped_revoke_error { + return ExternalGpuGetTerminal::Failed { + detail: format!( + "GPU Get initially failed to revoke non-transferable starts: {detail}" + ), + }; } + return ExternalGpuGetTerminal::Revoked { transfer_error }; } -} -impl ExternalInner { - async fn wait_initial_control_plane_ready(&self, phase: &'static str) -> KvResult<()> { - let owner_id = self.shared_storage_node_id().await.expect( - "ExternalClientApi expects current_owner to be Some before control-plane readiness wait", - ); - self.wait_peer_send_ready(owner_id, "owner", phase).await?; + match view + .external_client_api() + .inner() + .master_batch_gpu_get_done(transfer_get_ids) + .await + { + Ok(()) => ExternalGpuGetTerminal::Completed { + planned_cpu_items: Vec::new(), + planned_cpu_owner_start_time: None, + }, + Err(err) => ExternalGpuGetTerminal::Failed { + detail: format!("GPU Get BatchDone failed: {err}"), + }, + } +} - let master_node_id = self - .view - .cluster_manager() - .find_or_wait_master_node() - .await?; - self.wait_peer_send_ready(master_node_id, "master", phase) - .await +async fn run_external_gpu_get_transfer_timed( + view: ExternalClientApiView, + transfer_items: Vec<(BatchGetStartItemResp, GpuMemoryGuard)>, + skipped_get_ids: Vec, + transfer_concurrency: usize, + cancel_requested: Arc, +) -> ExternalGpuGetTerminalEvent { + let outcome = run_external_gpu_get_transfer( + view, + transfer_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested, + ) + .await; + ExternalGpuGetTerminalEvent { + outcome, + terminal_at: Instant::now(), } +} - async fn wait_peer_send_ready( - &self, - logical_target: String, - target_role: &'static str, - phase: &'static str, - ) -> KvResult<()> { - let deadline = - Instant::now() + Duration::from_secs(EXTERNAL_INIT_CONTROL_PLANE_READY_TIMEOUT_SECS); - let shutdown_poller = self.view.register_shutdown_poller(); - let mut attempts = 0u64; - let mut consecutive_ready = 0usize; - let mut last_transient_err = None; +async fn run_external_mixed_gpu_get_transfer_timed( + view: ExternalClientApiView, + plan_handle: u64, + gpu_transfer_items: Vec<(BatchGetStartItemResp, GpuMemoryGuard)>, + planned_cpu_items: Vec<(String, u64)>, + skipped_get_ids: Vec, + transfer_concurrency: usize, + cancel_requested: Arc, +) -> ExternalGpuGetTerminalEvent { + if planned_cpu_items.is_empty() { + return run_external_gpu_get_transfer_timed( + view, + gpu_transfer_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested, + ) + .await; + } - loop { - if !shutdown_poller.is_running() { - return Err(KvError::Api(ApiError::SystemShutdown { + // The GPU branch owns tail Revoke. The planned CPU branch receives an + // empty tail so every master operation identity is finalized exactly + // once while both source classes still execute concurrently. + let gpu_future = run_external_gpu_get_transfer( + view.clone(), + gpu_transfer_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested.clone(), + ); + let cpu_future = run_external_planned_cpu_get( + view.clone(), + plan_handle, + planned_cpu_items, + Vec::new(), + transfer_concurrency, + cancel_requested, + ); + let (gpu_terminal, cpu_terminal) = futures::future::join(gpu_future, cpu_future).await; + + let outcome = match (gpu_terminal, cpu_terminal) { + ( + ExternalGpuGetTerminal::Completed { .. }, + ExternalPlannedCpuGetTerminal::Completed { + items, + owner_start_time, + }, + ) => ExternalGpuGetTerminal::Completed { + planned_cpu_items: items, + planned_cpu_owner_start_time: Some(owner_start_time), + }, + (gpu_terminal, cpu_terminal) => { + if let ExternalPlannedCpuGetTerminal::Completed { + items, + owner_start_time, + } = &cpu_terminal + { + release_planned_cpu_item_holders( + view.external_client_api().inner(), + items, + *owner_start_time, + ); + } + match (&gpu_terminal, &cpu_terminal) { + ( + ExternalGpuGetTerminal::Failed { detail: gpu_detail }, + ExternalPlannedCpuGetTerminal::Failed { detail: cpu_detail }, + ) => ExternalGpuGetTerminal::Failed { detail: format!( - "external control-plane wait stopped during {phase}: target_role={target_role} target={logical_target}" + "mixed Get GPU and CPU branches failed: gpu={gpu_detail}; cpu={cpu_detail}" ), - })); - } - - let readiness = self - .view - .p2p_module() - .ensure_peer_send_ready(&logical_target.clone().into()) - .await; - match readiness { - Ok(()) => { - consecutive_ready += 1; - if consecutive_ready >= EXTERNAL_INIT_CONTROL_PLANE_READY_CONSECUTIVE_SUCCESSES - { - if attempts > 0 { - tracing::info!( - "external control-plane route ready: phase={} target_role={} target={} attempts={}", - phase, - target_role, - logical_target, - attempts + 1, - ); - } - return Ok(()); + }, + (ExternalGpuGetTerminal::Failed { detail }, _) => ExternalGpuGetTerminal::Failed { + detail: format!("mixed Get GPU branch failed: {detail}"), + }, + (_, ExternalPlannedCpuGetTerminal::Failed { detail }) => { + ExternalGpuGetTerminal::Failed { + detail: format!("mixed Get CPU branch failed: {detail}"), } } - Err(err) - if matches!( - err, - crate::p2p::P2PError::NoConnectionReady { .. } - | crate::p2p::P2PError::NodeNotFound { .. } - | crate::p2p::P2PError::NodeNotConnected { .. } - | crate::p2p::P2PError::NodePortNotReady { .. } - | crate::p2p::P2PError::ConnectionError { .. } - | crate::p2p::P2PError::SendFailed { .. } - | crate::p2p::P2PError::Iceoryx2TransportNotStarted {} - ) => - { - consecutive_ready = 0; - last_transient_err = Some(err.to_string()); + (ExternalGpuGetTerminal::Revoked { transfer_error }, _) => { + ExternalGpuGetTerminal::Revoked { + transfer_error: transfer_error.clone(), + } } - Err(err) => return Err(KvError::from(err)), - } - - attempts += 1; - if attempts == 1 || attempts % 20 == 0 { - tracing::info!( - "waiting for external control-plane route: phase={} target_role={} target={} attempts={} last_transient_err={:?}", - phase, - target_role, - logical_target, - attempts, - last_transient_err, - ); + (_, ExternalPlannedCpuGetTerminal::Revoked) => ExternalGpuGetTerminal::Revoked { + transfer_error: Some("mixed Get CPU branch was revoked".to_string()), + }, + _ => unreachable!("mixed Get non-completed branches must fail or revoke"), } + } + }; + ExternalGpuGetTerminalEvent { + outcome, + terminal_at: Instant::now(), + } +} - if Instant::now() >= deadline { - return Err(KvError::Api(ApiError::Unknown { - detail: format!( - "timed out waiting for external control-plane route during {phase}: target_role={target_role} target={logical_target} attempts={attempts} last_transient_err={last_transient_err:?}" - ), - })); - } +impl ExternalClientApi { + /// Access inner external-only API. Safe to unwrap in external role. + pub fn inner(&self) -> &ExternalInner { + &self.0 + } - sleep(Duration::from_millis( - EXTERNAL_INIT_CONTROL_PLANE_READY_POLL_MS, - )) - .await; - } + pub fn attach_view(&self, view: ExternalClientApiView) { + // This module is constructed only for the external variant; view attachment is + // therefore an invariant. + self.inner().view.attach(view); } - fn maybe_log_external_put_trace_window(&self, sample: &TestPutPhaseTrace) { - let maybe_window = { - let mut guard = self.put_trace_log_window.lock(); - guard.push_and_maybe_take(sample) - }; - let Some((elapsed, samples)) = maybe_window else { - return; - }; - if samples.is_empty() { - return; - } - let summary = summarize_external_put_trace_window(&samples); - if summary.is_empty() { - return; - } + pub async fn construct(arg: ExternalClientApiNewArg) -> Result { tracing::info!( - "external_put_trace_window samples={} window_s={:.1} {}", - samples.len(), - elapsed.as_secs_f64(), - summary + "Constructing ExternalClientApi in ExternalClient mode (PreView): shm_dir={}", + arg.shared_memory_path ); - } - async fn current_owner_start_time(&self) -> i64 { - let g = self.current_owner.read().await; - g.as_ref().map(|o| o.owner_start_time).unwrap_or_default() + Ok(Self(ExternalInner { + view: ExternalClientApiViewHolder::new(), + current_owner: ARwLock::new(None), + owner_remap_notify: Arc::new(Notify::new()), + wait_owner_gate: ARwLock::new(()), + initial_sub_cluster: OnceLock::new(), + expected_cluster_name: arg.expected_cluster_name, + expected_protocol_version: arg.expected_protocol_version, + external_shared_memory_path: arg.shared_memory_path, + external_shared_file_path: arg.shared_file_path, + enable_side_transfer: arg.enable_side_transfer, + short_circuit_put_payload_path: arg.short_circuit_put_payload_path, + side_rr_next: AtomicUsize::new(0), + side_transfer_put_bindings: moka::sync::Cache::builder() + .time_to_live(Duration::from_secs(10 * 60)) + .segments(16) + .build(), + rpc_caller_external_get: RPCCaller::::new(), + rpc_caller_external_batch_get: RPCCaller::::new(), + rpc_caller_external_batch_get_local_probe: + RPCCaller::::new(), + rpc_caller_external_batch_get_start: RPCCaller::::new(), + rpc_caller_external_batch_get_transfer: RPCCaller::::new(), + rpc_caller_external_batch_get_cancel: RPCCaller::::new(), + rpc_caller_master_batch_get_start: RPCCaller::::new(), + rpc_caller_master_batch_get_plan: RPCCaller::::new(), + rpc_caller_master_batch_get_bind: RPCCaller::::new(), + rpc_caller_master_batch_get_done: RPCCaller::::new(), + rpc_caller_master_batch_get_revoke: RPCCaller::::new(), + rpc_caller_external_execute_planned_get: RPCCaller::::new( + ), + rpc_caller_external_put_commit: RPCCaller::::new(), + rpc_caller_external_batch_put_commit: RPCCaller::::new(), + rpc_caller_external_put_start: RPCCaller::::new(), + rpc_caller_external_batch_put_start: RPCCaller::::new(), + rpc_caller_external_put_transfer_end: RPCCaller::::new(), + rpc_caller_external_batch_put_transfer_end: + RPCCaller::::new(), + rpc_caller_external_delete: RPCCaller::::new(), + rpc_caller_external_is_exist: RPCCaller::::new(), + rpc_caller_external_batch_is_exist: RPCCaller::::new(), + rpc_caller_external_observability_snapshot: + RPCCaller::::new(), + rpc_caller_external_delete_ack: RPCCaller::::new(), + rpc_caller_external_batch_delete_ack: RPCCaller::::new(), + rpc_caller_external_put_revoke: RPCCaller::::new(), + rpc_caller_allocate_client_lease: RPCCaller::::new(), + rpc_caller_client_lease_keepalive: RPCCaller::::new(), + key_weak_memholder_index: DashMap::new(), + pending_external_get_start: DashMap::new(), + pending_inline_external_get_start: DashMap::new(), + next_gpu_get_handle: AtomicU64::new(1), + pending_external_get_plan: DashMap::new(), + pending_external_gpu_get: DashMap::new(), + pending_external_planned_cpu_get: DashMap::new(), + inflight1_per_key: SemaphoreMap::new(1, std::time::Duration::from_secs(120)), + put_trace_log_window: Mutex::new(ExternalPutTraceLogWindow::new()), + external_delete_ack_batch: ExternalDeleteAckBatchHandle::new(), + })) } - async fn current_owner_snapshot(&self) -> Option<(String, i64, SharedMetaSignature)> { - let guard = self.current_owner.read().await; - let owner = guard.as_ref()?; - Some(( - owner.node_id.clone(), - owner.owner_start_time, - owner.shared_memory.memory_signature().clone(), - )) + pub async fn init2_prepare(&self) -> Result<(), KvError> { + // Prepare external client api initialization without waiting for owner readiness. + // + // All owner readiness (shared.json + mmap.file + membership observation) is handled by + // the init resource hook `owner_shared_mem_bundle_ready`. + Ok(()) } - async fn current_owner_base_if_advanced( + pub(crate) async fn wait_owner_shared_mem_bundle_ready_for_init_resource( &self, - prev_owner_start_time: i64, - ) -> Option<(i64, usize)> { - let guard = self.current_owner.read().await; - let owner = guard.as_ref()?; - if owner.owner_start_time == prev_owner_start_time { - return None; - } - Some(( - owner.owner_start_time, - owner.shared_memory.as_ptr() as usize, - )) - } - - async fn owner_generation_changed_in_cluster(&self, prev_owner_start_time: i64) -> bool { - let Some(owner_id) = self.shared_storage_node_id().await else { - return false; - }; - self.view - .cluster_manager() - .get_member_info_cached(&owner_id) - .is_some_and(|member| member.node_start_time != prev_owner_start_time) - } + ) -> Result<(), KvError> { + let ext = &self.0; - async fn try_background_owner_remap_once(&self) -> KvResult { - let Some((owner_id, owner_start_time, current_signature)) = - self.current_owner_snapshot().await - else { - return Ok(false); - }; + if ext.current_owner.read().await.is_none() { + // Initial attach: accept the current shared.json without requiring a post-wait write_ts. + let wait_start_ts = i64::MIN; + let OwnerRestartPayload { meta, signature } = task_wait_owner_restart( + ext.view.clone_view(), + ext.external_shared_memory_path.clone(), + ext.external_shared_file_path.clone(), + None, + wait_start_ts, + None, + ext.expected_cluster_name.clone(), + ext.expected_protocol_version.clone(), + ) + .await?; - let share_mem_path = self.share_mem_path(); - let shared_meta_path = format!("{}/shared.json", share_mem_path); - let probe = probe_owner_restart_payload( - &self.view.clone_view(), - &share_mem_path, - &shared_meta_path, - Some(¤t_signature), - i64::MIN, - Some(owner_id.as_str()), - &self.expected_cluster_name, - &self.expected_protocol_version, - ) - .await?; + let shared_memory_ptr = ExternalInner::init_shared_memory_from_meta( + &ext.external_shared_memory_path, + &meta, + signature, + )?; - let OwnerRestartProbe::Ready(payload) = probe else { - return Ok(false); - }; - if payload.meta.node_start_time == owner_start_time - && payload.signature == current_signature - { - return Ok(false); + ext.initial_sub_cluster + .set(meta.sub_cluster.clone()) + .unwrap(); + *ext.current_owner.write().await = Some(CurrentOwner { + node_id: meta.owner_id.clone(), + owner_start_time: meta.node_start_time, + shared_memory: shared_memory_ptr, + }); + ext.owner_remap_notify.notify_waiters(); } - self.finish_owner_recover(&share_mem_path, payload) - .await?; - Ok(true) + // Make the resource include the cluster membership observation as well. + self.init3_wait_owner_present().await?; + Ok(()) } - /// Try to get a live ExternalMemHolder from weak index. - async fn try_get_from_weak_cache(&self, key: &str) -> Option> { - if let Some(w_ref) = self.key_weak_memholder_index.get(key) { - let w = w_ref.value().clone(); - drop(w_ref); - if let Some(h) = w.upgrade() { - // Ensure holder belongs to current owner generation - if h.owner_start_time == self.current_owner_start_time().await { - return Some(h); - } else { - // Stale generation; remove and fall through - let _ = self.key_weak_memholder_index.remove(key); - } - } else { - // Dead weak; remove to keep cache clean - let _ = self.key_weak_memholder_index.remove(key); - } - } - None - } - // Removed trivial helper: inline-match OwnerStartTimeMismatch directly where needed. - /// 获取共享内存基址（以 usize 表示的地址）；未就绪时返回 NotConfigured - async fn base_ptr(&self) -> KvResult { - let lock = self.current_owner.read().await; - if let Some(o) = lock.as_ref() { - return Ok(o.shared_memory.as_ptr() as usize); - } - Err(KvError::SharedMem(SharedMemError::NotConfigured { - node_id: self.shared_storage_node_id().await, - detail: Some("Shared memory not ready".to_string()), - })) - } + pub async fn init2_after_owner_shared_mem_bundle_ready(&self) -> Result<(), KvError> { + let ext = &self.0; - async fn base_ptr_ro(&self) -> KvResult { - let lock = self.current_owner.read().await; - if let Some(o) = lock.as_ref() { - return Ok(o.shared_memory.as_ptr_ro() as usize); - } - Err(KvError::SharedMem(SharedMemError::NotConfigured { - node_id: self.shared_storage_node_id().await, - detail: Some("Shared memory not ready".to_string()), - })) - } + let owner_id = ext.shared_storage_node_id().await.expect( + "ExternalClientApi expects current_owner to be Some after owner_shared_mem_bundle_ready", + ); - async fn ensure_owner_ready(&self, prev_owner_start_time: &mut i64) -> KvResult { - match self.base_ptr().await { - Ok(addr) => Ok(addr), - Err(_) => { - let path = self.share_mem_path(); - let (st, addr) = self - .wait_owner_recover_only(&path, *prev_owner_start_time) - .await?; - *prev_owner_start_time = st; - Ok(addr) - } - } - } + // English note: + // Register inbound RPC handlers before any awaited etcd operations that publish or mutate + // member metadata. Otherwise, other nodes can observe this member and send RPCs while the + // handler set is still incomplete, leading to transient "No handler found" drops. + // + // Owner binding (current_owner) is already established by the init resource + // `owner_shared_mem_bundle_ready`, so handler registration is safe here. + ext.rpc_caller_external_get.regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_get + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_get_local_probe + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_get_start + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_get_transfer + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_get_cancel + .regist(ext.view.p2p_module()); + ext.rpc_caller_master_batch_get_start + .regist(ext.view.p2p_module()); + ext.rpc_caller_master_batch_get_plan + .regist(ext.view.p2p_module()); + ext.rpc_caller_master_batch_get_bind + .regist(ext.view.p2p_module()); + ext.rpc_caller_master_batch_get_done + .regist(ext.view.p2p_module()); + ext.rpc_caller_master_batch_get_revoke + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_execute_planned_get + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_put_commit + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_put_commit + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_put_start + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_put_start + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_put_transfer_end + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_put_transfer_end + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_delete.regist(ext.view.p2p_module()); + ext.rpc_caller_external_is_exist + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_is_exist + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_observability_snapshot + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_delete_ack + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_batch_delete_ack + .regist(ext.view.p2p_module()); + ext.rpc_caller_external_put_revoke + .regist(ext.view.p2p_module()); + crate::key_prefix::init_for_p2p_owner(ext.view.p2p_module()); + crate::kvlease::init_for_p2p_owner(ext.view.p2p_module()); + crate::metrics::client::init_for_p2p_owner(ext.view.p2p_module()); - /// Note: ExternalInner is only constructed in ExternalClient role. + let external_delete_ack_rx = ext + .external_delete_ack_batch + .take_rx() + .expect("external holder ACK batch worker initialized twice"); + spawn_external_delete_ack_batch(ext.view.clone_view(), external_delete_ack_rx); - async fn finish_owner_recover( - &self, - share_mem_path: &str, - payload: OwnerRestartPayload, - ) -> KvResult<(i64, usize)> { - self.remap_shared_memory_with_payload(share_mem_path, &payload) - .await?; - self.view + let view_ext = ext.view.clone_view(); + RPCHandler::::new().regist( + ext.view.p2p_module(), + move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_external_invalidate_weak_index", async move { + let result = handle_external_invalidate_weak_index(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }, + ); + + RPCCaller::::new().regist(ext.view.p2p_module()); + let view_ext = ext.view.clone_view(); + RPCHandler::::new().regist(ext.view.p2p_module(), move |resp, msg| { + let view = view_ext.clone(); + let view_task = view.clone(); + view.spawn("rpc_sync_kv_to_file", async move { + let result = handle_sync_kv_to_file_external(&view_task, &msg).await; + let _ = resp.send_resp(result).await; + }); + Ok(()) + }); + tracing::info!("ExternalClientApi RPC callers registered"); + + ext.view .cluster_manager() .set_self_share_group_binding(ShareGroupOwnerRef { - owner_id: payload.meta.owner_id.clone(), - owner_start_time: payload.meta.node_start_time, + owner_id: owner_id.clone(), + owner_start_time: ext.current_owner_start_time().await, }) .await?; - self.view + ext.view .cluster_manager() - .set_self_sub_cluster(payload.meta.sub_cluster.clone()) + .set_self_sub_cluster(ext.initial_sub_cluster.get().unwrap().clone()) .await .map_err(KvError::from)?; - self.wait_initial_control_plane_ready("owner_recover") - .await?; - let base_addr = self.base_ptr().await?; - Ok((self.current_owner_start_time().await, base_addr)) - } - - async fn wait_owner_recover_only( - &self, - share_mem_path: &str, - prev_owner_start_time: i64, - ) -> KvResult<(i64, usize)> { - self.wait_owner_recover(share_mem_path, prev_owner_start_time) - .await - } - - async fn recover_after_owner_start_time_mismatch( - &self, - prev_owner_start_time: &mut i64, - ) -> KvResult { - let path = self.share_mem_path(); - let (st, addr) = self - .wait_owner_recover_only(&path, *prev_owner_start_time) + // Publishing the share-group binding changes the desired local-owner route from the + // pre-binding direct lane to intra-machine-only. Do not announce external init complete + // until that topology transition has converged for the exact owner generation. + self.wait_current_owner_intra_rpc_ready_after_binding() .await?; - *prev_owner_start_time = st; - Ok(addr) - } - async fn recover_after_p2p_error(&self, prev_owner_start_time: &mut i64) -> KvResult { - if !self - .owner_generation_changed_in_cluster(*prev_owner_start_time) - .await { - return match self.base_ptr().await { - Ok(addr) => Ok(addr), - Err(_) => { - let path = self.share_mem_path(); - let (st, addr) = self - .wait_owner_recover_only(&path, *prev_owner_start_time) - .await?; - *prev_owner_start_time = st; - Ok(addr) - } - }; - } + let view = ext.view.clone_view(); + let view_task = view.clone(); + view.spawn("external_owner_remap_actor", async move { + let shutdown_poller = view_task.register_shutdown_poller(); + let mut cluster_rx = view_task.cluster_manager().listen(); + let mut tick = tokio::time::interval(Duration::from_millis(200)); - let path = self.share_mem_path(); - let (st, addr) = self - .wait_owner_recover_only(&path, *prev_owner_start_time) - .await?; - *prev_owner_start_time = st; - Ok(addr) - } + loop { + if !shutdown_poller.is_running() { + tracing::info!("external owner remap actor stopped by shutdown"); + break; + } - /// Wait for owner recovery until shared memory has been remapped and `owner_start_time` - /// has advanced. - async fn wait_owner_recover( - &self, - _share_mem_path: &str, - prev_owner_start_time: i64, - ) -> KvResult<(i64, usize)> { - if let Some(res) = self - .current_owner_base_if_advanced(prev_owner_start_time) - .await + if let Err(err) = view_task + .external_client_api() + .inner() + .try_background_owner_remap_once() + .await + { + tracing::warn!("external owner remap actor probe failed: {}", err); + } + + tokio::select! { + _ = tick.tick() => {} + recv = cluster_rx.recv() => { + if recv.is_err() { + sleep(Duration::from_millis(200)).await; + cluster_rx = view_task.cluster_manager().listen(); + } + } + } + } + }); + } + + // Attribute local IPC bandwidth to the owner daemon (machine-level view). + // + // Causal chain: + // - External<->external traffic can use the local IPC tier (iceoryx2) when both are in the + // same share-group (same owner_id + local_ipc_root). + // - Topology aggregates bandwidth at the owner/machine level, so local IPC bytes must be + // charged to the owner, otherwise the UI under-reports throughput. + // - We keep the P2P hot path allocation-free by recording bytes into atomics, and flush + // them periodically via a background task. { - return Ok(res); + let cm = ext.view.cluster_manager(); + let handle = IpcBandwidthAttributorHandle::new(); + cm.attach_ipc_bandwidth_attributor_handle(handle.clone()); + if let Some(observe) = cm.observe_handle().cloned() { + let self_member_id = cm.self_member_id().to_string(); + let owner_role = NodeRole::Client.to_string(); + let owner_id_for_task = owner_id.clone(); + let view_task = ext.view.clone_view(); + let view_task2 = view_task.clone(); + view_task.spawn("ipc_bandwidth_attributor", async move { + let mut shutdown_waiter = view_task2.register_shutdown_waiter(); + let mut interval = tokio::time::interval(Duration::from_secs( + crate::metric_reporter::METRICS_FLUSH_INTERVAL_SECS, + )); + + loop { + tokio::select! { + _ = interval.tick() => { + let tx_bytes = handle.take_tx_bytes(); + if tx_bytes > 0 { + observe.try_record_peer_network_bytes_override( + ObserveComponent::LocalIpc, + owner_id_for_task.as_str(), + owner_role.as_str(), + self_member_id.as_str(), + ObserveDirection::Tx, + tx_bytes, + ); + } + let rx_bytes = handle.take_rx_bytes(); + if rx_bytes > 0 { + observe.try_record_peer_network_bytes_override( + ObserveComponent::LocalIpc, + owner_id_for_task.as_str(), + owner_role.as_str(), + self_member_id.as_str(), + ObserveDirection::Rx, + rx_bytes, + ); + } + } + _ = shutdown_waiter.wait() => { + break; + } + } + } + }); + } else { + tracing::info!( + "ExternalClientApi local IPC bandwidth attribution disabled: ObserveHandle not attached" + ); + } } + Ok(()) + } + pub async fn init3_wait_owner_present(&self) -> Result<(), KvError> { + let ext = &self.0; + let owner_id = ext + .shared_storage_node_id() + .await + .expect("external role expects current_owner to be Some after init2"); + let owner_start_time = ext.current_owner_start_time().await; - let _wait_guard = self.wait_owner_gate.write().await; - let shutdown_poller = self.view.register_shutdown_poller(); - let mut waited_ticks = 0u64; + let cm = ext.view.cluster_manager(); + if cm + .get_member_info_cached(&owner_id) + .map(|member| member.node_start_time == owner_start_time) + .unwrap_or(false) + { + return Ok(()); + } + tracing::info!( + "External init: waiting for owner generation to join (owner_id={} owner_start_time={})", + owner_id, + owner_start_time + ); + let mut rx = cm.listen(); loop { - if let Some(res) = self - .current_owner_base_if_advanced(prev_owner_start_time) - .await + if cm + .get_member_info_cached(&owner_id) + .map(|member| member.node_start_time == owner_start_time) + .unwrap_or(false) { - return Ok(res); + tracing::info!( + "External init: owner generation observed (owner_id={} owner_start_time={})", + owner_id, + owner_start_time + ); + return Ok(()); } - if !shutdown_poller.is_running() { - return Err(KvError::Api(ApiError::SystemShutdown { - detail: "Owner recovery wait aborted due to shutdown".to_string(), - })); + match rx.recv().await { + Ok(_ev) => { + // Yield once to allow watcher to update member cache after emitting an event. + limit_thirdparty::tokio::task::yield_now().await; + } + Err(e) => { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "cluster event channel closed while waiting for owner generation (owner_id={} owner_start_time={}): {}", + owner_id, owner_start_time, e + ), + })); + } } + } + } - let notified = self.owner_remap_notify.notified(); - if let Some(res) = self - .current_owner_base_if_advanced(prev_owner_start_time) - .await + async fn wait_current_owner_intra_rpc_ready_after_binding(&self) -> Result<(), KvError> { + let ext = &self.0; + let owner_id = ext + .shared_storage_node_id() + .await + .expect("external role expects current_owner to be Some after init2"); + let owner_node_id = owner_id.clone().into(); + let owner_start_time = ext.current_owner_start_time().await; + let expected_binding = ShareGroupOwnerRef { + owner_id: owner_id.clone(), + owner_start_time, + }; + let started_at = Instant::now(); + let timeout = Duration::from_secs(EXTERNAL_OWNER_INTRA_RPC_READY_TIMEOUT_SECS); + + tracing::info!( + owner_id = %owner_id, + owner_start_time, + "External init: waiting for current owner intra-machine RPC route after share-group binding" + ); + loop { + let snapshot = ext.view.p2p_module().tier_snapshot(); + let self_binding_ready = snapshot.share_group_owner(&snapshot.self_peer_gen.peer_id) + == Some(&expected_binding); + if self_binding_ready + && let Some(peer_gen) = snapshot.peer_gen(&owner_node_id) + && peer_gen.node_start_time == owner_start_time + && snapshot.is_send_ready_intra_effective(&peer_gen) { - return Ok(res); - } - tokio::select! { - _ = notified => {} - _ = sleep(Duration::from_millis(200)) => {} + tracing::info!( + owner_id = %owner_id, + owner_start_time, + elapsed_ms = started_at.elapsed().as_millis(), + "External init: current owner intra-machine RPC route ready after share-group binding" + ); + return Ok(()); } - waited_ticks += 1; - if waited_ticks % 25 == 0 { - tracing::warn!( - "[wait_owner_remap] waiting for owner remap... ({}s)", - waited_ticks / 5 - ); + if started_at.elapsed() >= timeout { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "timed out waiting for current owner intra-machine RPC route after share-group binding: owner_id={} owner_start_time={} timeout_s={} self_binding={:?} peer={:?}", + owner_id, + owner_start_time, + EXTERNAL_OWNER_INTRA_RPC_READY_TIMEOUT_SECS, + snapshot.share_group_owner(&snapshot.self_peer_gen.peer_id), + snapshot.peers.get(owner_id.as_str()), + ), + })); } + sleep(Duration::from_millis(20)).await; } } +} - /// Read shared.json to get shared memory metadata - fn read_shared_json(shared_meta_path: &str) -> KvResult { - let mut file = File::open(shared_meta_path).map_err(|e| { - KvError::SharedMem(SharedMemError::MetaDataLoadError { - path: shared_meta_path.to_string(), - detail: format!("Failed to open shared.json: {}", e), - }) - })?; - let mut buf = String::new(); - use std::io::Read as _; - file.read_to_string(&mut buf).map_err(|e| { - KvError::SharedMem(SharedMemError::MetaDataLoadError { - path: shared_meta_path.to_string(), - detail: format!("Failed to read shared.json: {}", e), - }) - })?; - let meta: SharedJsonMeta = serde_json::from_str(&buf).map_err(|e| { - KvError::SharedMem(SharedMemError::MetaDataLoadError { - path: shared_meta_path.to_string(), - detail: format!("Failed to parse shared.json: {}", e), - }) - })?; +impl ExternalInner { + fn maybe_log_external_put_trace_window(&self, sample: &TestPutPhaseTrace) { + let maybe_window = { + let mut guard = self.put_trace_log_window.lock(); + guard.push_and_maybe_take(sample) + }; + let Some((elapsed, samples)) = maybe_window else { + return; + }; + if samples.is_empty() { + return; + } + let summary = summarize_external_put_trace_window(&samples); + if summary.is_empty() { + return; + } + tracing::info!( + "external_put_trace_window samples={} window_s={:.1} {}", + samples.len(), + elapsed.as_secs_f64(), + summary + ); + } - Ok(meta) + pub async fn current_owner_start_time(&self) -> i64 { + let g = self.current_owner.read().await; + g.as_ref().map(|o| o.owner_start_time).unwrap_or_default() } - fn get_shared_meta_signature(shared_meta_path: &str) -> KvResult { - fluxon_util::fs_watch::get_file_signature(shared_meta_path).map_err(KvError::from) + pub async fn wait_current_owner_mapped_range(&self) -> KvResult<(String, i64, u64, u64, u64)> { + let mut prev_owner_start_time = i64::MIN; + let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; + let guard = self.current_owner.read().await; + let owner = guard.as_ref().ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared memory not ready".to_string()), + }) + })?; + Ok(( + owner.node_id.clone(), + owner.owner_start_time, + owner.shared_memory.as_ptr() as u64, + owner.shared_memory.as_ptr_ro() as u64, + owner.shared_memory.len(), + )) } - async fn remap_shared_memory_with_payload( - &self, - share_mem_path: &str, - payload: &OwnerRestartPayload, - ) -> KvResult<()> { - let shared_memory = Self::init_shared_memory_from_meta( - share_mem_path, - &payload.meta, - payload.signature.clone(), - )?; - let len = shared_memory.len(); - let mut lock = self.current_owner.write().await; - if let Some(owner) = lock.as_mut() { - owner.shared_memory = shared_memory; - owner.owner_start_time = payload.meta.node_start_time; - owner.node_id = payload.meta.owner_id.clone(); - } else { - // If no owner set yet, set node_id from shared.json - *lock = Some(CurrentOwner { - node_id: payload.meta.owner_id.clone(), - owner_start_time: payload.meta.node_start_time, - shared_memory, - }); - } - tracing::info!( - "[wait_owner_client_recover] Ownerclient recovered, mmap remapped: len={}", - len - ); - self.key_weak_memholder_index.clear(); - self.owner_remap_notify.notify_waiters(); - Ok(()) + async fn current_owner_snapshot(&self) -> Option<(String, i64, SharedMetaSignature)> { + let guard = self.current_owner.read().await; + let owner = guard.as_ref()?; + Some(( + owner.node_id.clone(), + owner.owner_start_time, + owner.shared_memory.memory_signature().clone(), + )) } - /// Initialize shared memory mapping using file path directly - fn init_shared_memory( - mmap_file_path: &str, - len: u64, - memory_signature: SharedMetaSignature, - ) -> KvResult> { - use std::fs::OpenOptions; - use std::os::unix::io::AsRawFd; - - tracing::info!( - "Initializing shared memory mapping: file={}, len={}", - mmap_file_path, - len - ); + async fn current_owner_base_if_advanced( + &self, + prev_owner_start_time: i64, + ) -> Option<(i64, usize)> { + let guard = self.current_owner.read().await; + let owner = guard.as_ref()?; + if owner.owner_start_time == prev_owner_start_time { + return None; + } + Some(( + owner.owner_start_time, + owner.shared_memory.as_ptr() as usize, + )) + } - let file = OpenOptions::new() - .read(true) - .write(true) - .open(mmap_file_path) - .map_err(|e| { - KvError::SharedMem(SharedMemError::MappingFailed { - path: mmap_file_path.to_string(), - len, - detail: format!("Failed to open shared memory file: {}", e), - }) - })?; + async fn owner_generation_changed_in_cluster(&self, prev_owner_start_time: i64) -> bool { + let Some(owner_id) = self.shared_storage_node_id().await else { + return false; + }; + self.view + .cluster_manager() + .get_member_info_cached(&owner_id) + .is_some_and(|member| member.node_start_time != prev_owner_start_time) + } - let fd = file.as_raw_fd(); - tracing::debug!("Opened shared memory file: fd={}", fd); + async fn try_background_owner_remap_once(&self) -> KvResult { + let Some((owner_id, owner_start_time, current_signature)) = + self.current_owner_snapshot().await + else { + return Ok(false); + }; - unsafe { - let addr_rw = mmap( - std::ptr::null_mut(), - len as usize, - PROT_READ | PROT_WRITE, - MAP_SHARED, - fd, - 0, - ); + let shared_memory_path = self.shared_memory_path(); + let shared_file_path = self.shared_file_path(); + let shared_meta_path = format!("{}/shared.json", shared_file_path); + let probe = probe_owner_restart_payload( + &self.view.clone_view(), + &shared_memory_path, + &shared_file_path, + &shared_meta_path, + Some(¤t_signature), + i64::MIN, + Some(owner_id.as_str()), + &self.expected_cluster_name, + &self.expected_protocol_version, + ) + .await?; - if addr_rw == libc::MAP_FAILED { - return Err(KvError::SharedMem(SharedMemError::MappingFailed { - path: mmap_file_path.to_string(), - len, - detail: "mmap failed".to_string(), - })); - } + let OwnerRestartProbe::Ready(payload) = probe else { + return Ok(false); + }; + if payload.meta.node_start_time == owner_start_time + && payload.signature == current_signature + { + return Ok(false); + } - let addr_ro = mmap( - std::ptr::null_mut(), - len as usize, - PROT_READ, - MAP_SHARED, - fd, - 0, - ); + self.finish_owner_recover(&shared_memory_path, payload) + .await?; + Ok(true) + } - if addr_ro == libc::MAP_FAILED { - libc::munmap(addr_rw, len as usize); - return Err(KvError::SharedMem(SharedMemError::MappingFailed { - path: mmap_file_path.to_string(), - len, - detail: "mmap (read-only) failed".to_string(), - })); + /// Try to get a live ExternalMemHolder from weak index. + async fn try_get_from_weak_cache(&self, key: &str) -> Option> { + if let Some(w_ref) = self.key_weak_memholder_index.get(key) { + let w = w_ref.value().clone(); + drop(w_ref); + if let Some(h) = w.upgrade() { + // Ensure holder belongs to current owner generation + if h.owner_start_time == self.current_owner_start_time().await { + return Some(h); + } else { + // Stale generation; remove and fall through + let _ = self.key_weak_memholder_index.remove(key); + } + } else { + // Dead weak; remove to keep cache clean + let _ = self.key_weak_memholder_index.remove(key); } - - tracing::info!( - "Successfully mapped shared memory: file={}, len={}, addr={:?}", - mmap_file_path, - len, - addr_rw - ); - // Store the directory path (shared memory base path), not the mmap file path. - // Many recovery routines expect a directory path to locate memory.file and mmap.file. - let dir_path = std::path::Path::new(mmap_file_path) - .parent() - .map(|p| p.to_string_lossy().to_string()) - .unwrap_or_else(|| String::new()); - - Ok(Arc::new(SharedMemoryPtr::new( - addr_rw as *mut u8, - addr_ro as *mut u8, - len, - dir_path, - file, - memory_signature, - ))) } + None } - fn init_shared_memory_from_meta( - share_mem_path: &str, - meta: &SharedJsonMeta, - memory_signature: SharedMetaSignature, - ) -> KvResult> { - let mmap_file_path = format!("{}/mmap.file", share_mem_path); - Self::init_shared_memory(&mmap_file_path, meta.segment_len, memory_signature) + async fn try_get_local_complete_holder(&self, key: &str) -> Option> { + if let Some(holder) = self.try_get_from_weak_cache(key).await { + return Some(holder); + } + None } - /// Get the shared storage node ID this client connects to - pub async fn shared_storage_node_id(&self) -> Option { - let g = self.current_owner.read().await; - g.as_ref().map(|o| o.node_id.clone()) + // Removed trivial helper: inline-match OwnerStartTimeMismatch directly where needed. + /// 获取共享内存基址（以 usize 表示的地址）；未就绪时返回 NotConfigured + pub async fn base_ptr(&self) -> KvResult { + let lock = self.current_owner.read().await; + if let Some(o) = lock.as_ref() { + return Ok(o.shared_memory.as_ptr() as usize); + } + Err(KvError::SharedMem(SharedMemError::NotConfigured { + node_id: self.shared_storage_node_id().await, + detail: Some("Shared memory not ready".to_string()), + })) } - /// Get the configured shared-memory base path (external mode). - /// Non-external modes return empty string. - pub fn share_mem_path(&self) -> String { - self.external_share_mem_path.clone() + async fn base_ptr_ro(&self) -> KvResult { + let lock = self.current_owner.read().await; + if let Some(o) = lock.as_ref() { + return Ok(o.shared_memory.as_ptr_ro() as usize); + } + Err(KvError::SharedMem(SharedMemError::NotConfigured { + node_id: self.shared_storage_node_id().await, + detail: Some("Shared memory not ready".to_string()), + })) } - pub fn large_file_paths(&self) -> &crate::config::LargeFilePaths { - &self.external_large_file_paths + async fn ensure_owner_ready(&self, prev_owner_start_time: &mut i64) -> KvResult { + match self.base_ptr().await { + Ok(addr) => Ok(addr), + Err(_) => { + let path = self.shared_memory_path(); + let (st, addr) = self + .wait_owner_recover_only(&path, *prev_owner_start_time) + .await?; + *prev_owner_start_time = st; + Ok(addr) + } + } } - fn should_fallback_side_p2p_error(err: &crate::p2p::P2PError) -> bool { - matches!( - err, - crate::p2p::P2PError::NoConnectionReady { .. } - | crate::p2p::P2PError::NodeNotFound { .. } - | crate::p2p::P2PError::NodeNotConnected { .. } - | crate::p2p::P2PError::NodePortNotReady { .. } - | crate::p2p::P2PError::ConnectionError { .. } - | crate::p2p::P2PError::SendFailed { .. } - | crate::p2p::P2PError::StartServerError { .. } - | crate::p2p::P2PError::Iceoryx2TransportNotStarted {} - ) - } + /// Note: ExternalInner is only constructed in ExternalClient role. - fn read_side_transfer_peer(path: &std::path::Path) -> KvResult { - let buf = std::fs::read_to_string(path).map_err(|e| { - KvError::SharedMem(SharedMemError::MetaDataLoadError { - path: path.to_string_lossy().to_string(), - detail: format!("Failed to read side-transfer peer file: {}", e), + async fn finish_owner_recover( + &self, + shared_memory_path: &str, + payload: OwnerRestartPayload, + ) -> KvResult<(i64, usize)> { + self.remap_shared_memory_with_payload(shared_memory_path, &payload) + .await?; + self.view + .cluster_manager() + .set_self_share_group_binding(ShareGroupOwnerRef { + owner_id: payload.meta.owner_id.clone(), + owner_start_time: payload.meta.node_start_time, }) - })?; - serde_json::from_str(&buf).map_err(|e| { - KvError::SharedMem(SharedMemError::MetaDataLoadError { - path: path.to_string_lossy().to_string(), - detail: format!("Failed to parse side-transfer peer file: {}", e), - }) - }) + .await?; + self.view + .cluster_manager() + .set_self_sub_cluster(payload.meta.sub_cluster.clone()) + .await + .map_err(KvError::from)?; + let base_addr = self.base_ptr().await?; + Ok((self.current_owner_start_time().await, base_addr)) } - async fn pick_side_transfer_peer(&self, put_id: Option<(u64, u32)>) -> Option<(String, u16)> { - // External attach auto-detects owner side workers from the shared-memory peer files. - // Owner-side config still controls whether workers exist; external callers should not - // require an extra enable flag once the owner has published ready lanes. - let owner_id = self.shared_storage_node_id().await?; - let owner_start_time = self.current_owner_start_time().await; - let peers_dir = ClientSegPool::side_transfer_peers_dir(&self.external_share_mem_path); - let entries = std::fs::read_dir(&peers_dir).ok()?; - let mut ready = Vec::new(); - for entry in entries.flatten() { - let path = entry.path(); - if path.extension().and_then(|s| s.to_str()) != Some("json") { - continue; - } - let Ok(meta) = Self::read_side_transfer_peer(&path) else { - continue; - }; - if meta.owner_id != owner_id || meta.owner_start_time != owner_start_time { - continue; - } - let Some(member) = self - .view - .cluster_manager() - .get_member_info_cached(&meta.side_id) - else { - continue; + async fn wait_owner_recover_only( + &self, + shared_memory_path: &str, + prev_owner_start_time: i64, + ) -> KvResult<(i64, usize)> { + self.wait_owner_recover(shared_memory_path, prev_owner_start_time) + .await + } + + async fn recover_after_owner_start_time_mismatch( + &self, + prev_owner_start_time: &mut i64, + ) -> KvResult { + let path = self.shared_memory_path(); + let (st, addr) = self + .wait_owner_recover_only(&path, *prev_owner_start_time) + .await?; + *prev_owner_start_time = st; + Ok(addr) + } + + async fn recover_after_p2p_error(&self, prev_owner_start_time: &mut i64) -> KvResult { + if !self + .owner_generation_changed_in_cluster(*prev_owner_start_time) + .await + { + return match self.base_ptr().await { + Ok(addr) => Ok(addr), + Err(_) => { + let path = self.shared_memory_path(); + let (st, addr) = self + .wait_owner_recover_only(&path, *prev_owner_start_time) + .await?; + *prev_owner_start_time = st; + Ok(addr) + } }; - if member - .metadata - .get("side_transfer_worker") - .is_some_and(|v| v == "true") - == false + } + + let path = self.shared_memory_path(); + let (st, addr) = self + .wait_owner_recover_only(&path, *prev_owner_start_time) + .await?; + *prev_owner_start_time = st; + Ok(addr) + } + + /// Wait for owner recovery until shared memory has been remapped and `owner_start_time` + /// has advanced. + async fn wait_owner_recover( + &self, + _shared_memory_path: &str, + prev_owner_start_time: i64, + ) -> KvResult<(i64, usize)> { + if let Some(res) = self + .current_owner_base_if_advanced(prev_owner_start_time) + .await + { + return Ok(res); + } + + let _wait_guard = self.wait_owner_gate.write().await; + let shutdown_poller = self.view.register_shutdown_poller(); + let mut waited_ticks = 0u64; + + loop { + if let Some(res) = self + .current_owner_base_if_advanced(prev_owner_start_time) + .await { - continue; + return Ok(res); } - if member - .metadata - .get(META_KEY_SHARED_STORAGE_NODE_ID) - .is_some_and(|v| v == &owner_id) - == false - { - continue; + if !shutdown_poller.is_running() { + return Err(KvError::Api(ApiError::SystemShutdown { + detail: "Owner recovery wait aborted due to shutdown".to_string(), + })); } - if member - .metadata - .get(META_KEY_SHARED_STORAGE_NODE_START_TIME) - .and_then(|v| v.parse::().ok()) - != Some(owner_start_time) + + let notified = self.owner_remap_notify.notified(); + if let Some(res) = self + .current_owner_base_if_advanced(prev_owner_start_time) + .await { - continue; + return Ok(res); } - let Some(lane_idx) = meta.worker_idx() else { - continue; - }; - ready.push((lane_idx, meta.side_id)); - } + tokio::select! { + _ = notified => {} + _ = sleep(Duration::from_millis(200)) => {} + } + waited_ticks += 1; - if ready.is_empty() { - return None; - } - ready.sort_by(|lhs, rhs| lhs.cmp(rhs)); - if let Some(put_id) = put_id { - let lane_space = ready - .iter() - .map(|(lane_idx, _)| usize::from(*lane_idx)) - .max() - .map(|max_lane_idx| max_lane_idx + 1)?; - let desired_lane = stable_side_transfer_lane_for_put(put_id, lane_space)?; - let selected = ready - .iter() - .find(|(lane_idx, _)| *lane_idx == desired_lane) - .cloned(); - if selected.is_none() { + if waited_ticks % 25 == 0 { tracing::warn!( - "side-transfer desired lane not ready locally; falling back to owner: desired_lane={} lane_space={} owner_id={}", - desired_lane, - lane_space, - owner_id + "[wait_owner_remap] waiting for owner remap... ({}s)", + waited_ticks / 5 ); } - return selected.map(|(lane_idx, side_id)| (side_id, lane_idx)); } - - let idx = self.side_rr_next.fetch_add(1, Ordering::Relaxed); - let ready_len = ready.len(); - ready - .into_iter() - .nth(idx % ready_len) - .map(|(lane_idx, side_id)| (side_id, lane_idx)) } - fn remember_side_transfer_binding( - &self, - put_id: Option<(u64, u32)>, - binding: Option<(String, u16)>, - ) { - if let (Some(put_id), Some(binding)) = (put_id, binding) { - self.side_transfer_put_bindings.insert(put_id, binding); - } + /// Read shared.json to get shared memory metadata + fn read_shared_json(shared_meta_path: &str) -> KvResult { + let mut file = File::open(shared_meta_path).map_err(|e| { + KvError::SharedMem(SharedMemError::MetaDataLoadError { + path: shared_meta_path.to_string(), + detail: format!("Failed to open shared.json: {}", e), + }) + })?; + let mut buf = String::new(); + use std::io::Read as _; + file.read_to_string(&mut buf).map_err(|e| { + KvError::SharedMem(SharedMemError::MetaDataLoadError { + path: shared_meta_path.to_string(), + detail: format!("Failed to read shared.json: {}", e), + }) + })?; + let meta: SharedJsonMeta = serde_json::from_str(&buf).map_err(|e| { + KvError::SharedMem(SharedMemError::MetaDataLoadError { + path: shared_meta_path.to_string(), + detail: format!("Failed to parse shared.json: {}", e), + }) + })?; + + Ok(meta) } - fn bound_side_transfer_peer(&self, put_id: Option<(u64, u32)>) -> Option<(String, u16)> { - put_id.and_then(|put_id| self.side_transfer_put_bindings.get(&put_id)) + fn get_shared_meta_signature(shared_meta_path: &str) -> KvResult { + fluxon_util::fs_watch::get_file_signature(shared_meta_path).map_err(KvError::from) } - fn clear_side_transfer_binding(&self, put_id: Option<(u64, u32)>) { - if let Some(put_id) = put_id { - self.side_transfer_put_bindings.invalidate(&put_id); + async fn remap_shared_memory_with_payload( + &self, + shared_memory_path: &str, + payload: &OwnerRestartPayload, + ) -> KvResult<()> { + let shared_memory = Self::init_shared_memory_from_meta( + shared_memory_path, + &payload.meta, + payload.signature.clone(), + )?; + let len = shared_memory.len(); + let mut lock = self.current_owner.write().await; + if let Some(owner) = lock.as_mut() { + owner.shared_memory = shared_memory; + owner.owner_start_time = payload.meta.node_start_time; + owner.node_id = payload.meta.owner_id.clone(); + } else { + // If no owner set yet, set node_id from shared.json + *lock = Some(CurrentOwner { + node_id: payload.meta.owner_id.clone(), + owner_start_time: payload.meta.node_start_time, + shared_memory, + }); } + tracing::info!( + "[wait_owner_client_recover] Ownerclient recovered, mmap remapped: len={}", + len + ); + self.key_weak_memholder_index.clear(); + self.owner_remap_notify.notify_waiters(); + Ok(()) } - fn short_circuit_put_payload_path_enabled(&self) -> bool { - self.short_circuit_put_payload_path - } + /// Initialize shared memory mapping using file path directly + fn init_shared_memory( + mmap_file_path: &str, + len: u64, + memory_signature: SharedMetaSignature, + ) -> KvResult> { + use std::fs::OpenOptions; + use std::os::unix::io::AsRawFd; - async fn call_put_start_with_side_fallback( - &self, - owner_id: String, - req: MsgPack, - ) -> KvResult<(MsgPack, Option<(String, u16)>)> { - if let Some((side_id, lane_idx)) = self.pick_side_transfer_peer(None).await { - match self - .rpc_caller_external_put_start - .call( + tracing::info!( + "Initializing shared memory mapping: file={}, len={}", + mmap_file_path, + len + ); + + let file = OpenOptions::new() + .read(true) + .write(true) + .open(mmap_file_path) + .map_err(|e| { + KvError::SharedMem(SharedMemError::MappingFailed { + path: mmap_file_path.to_string(), + len, + detail: format!("Failed to open shared memory file: {}", e), + }) + })?; + + let fd = file.as_raw_fd(); + tracing::debug!("Opened shared memory file: fd={}", fd); + + unsafe { + let addr_rw = mmap( + std::ptr::null_mut(), + len as usize, + PROT_READ | PROT_WRITE, + MAP_SHARED, + fd, + 0, + ); + + if addr_rw == libc::MAP_FAILED { + return Err(KvError::SharedMem(SharedMemError::MappingFailed { + path: mmap_file_path.to_string(), + len, + detail: "mmap failed".to_string(), + })); + } + + let addr_ro = mmap( + std::ptr::null_mut(), + len as usize, + PROT_READ, + MAP_SHARED, + fd, + 0, + ); + + if addr_ro == libc::MAP_FAILED { + libc::munmap(addr_rw, len as usize); + return Err(KvError::SharedMem(SharedMemError::MappingFailed { + path: mmap_file_path.to_string(), + len, + detail: "mmap (read-only) failed".to_string(), + })); + } + + tracing::info!( + "Successfully mapped shared memory: file={}, len={}, addr={:?}", + mmap_file_path, + len, + addr_rw + ); + // Store the directory path (shared memory base path), not the mmap file path. + // Many recovery routines expect a directory path to locate memory.file and mmap.file. + let dir_path = std::path::Path::new(mmap_file_path) + .parent() + .map(|p| p.to_string_lossy().to_string()) + .unwrap_or_else(|| String::new()); + + Ok(Arc::new(SharedMemoryPtr::new( + addr_rw as *mut u8, + addr_ro as *mut u8, + len, + dir_path, + file, + memory_signature, + ))) + } + } + + fn init_shared_memory_from_meta( + shared_memory_path: &str, + meta: &SharedJsonMeta, + memory_signature: SharedMetaSignature, + ) -> KvResult> { + let mmap_file_path = format!("{}/mmap.file", shared_memory_path); + Self::init_shared_memory(&mmap_file_path, meta.segment_len, memory_signature) + } + /// Get the shared storage node ID this client connects to + pub async fn shared_storage_node_id(&self) -> Option { + let g = self.current_owner.read().await; + g.as_ref().map(|o| o.node_id.clone()) + } + + /// Get the configured shared-memory base path (external mode). + /// Non-external modes return empty string. + pub fn shared_memory_path(&self) -> String { + self.external_shared_memory_path.clone() + } + + /// Get the configured shared-file base path (external mode). + /// Non-external modes return empty string. + pub fn shared_file_path(&self) -> String { + self.external_shared_file_path.clone() + } + + fn should_fallback_side_p2p_error(err: &crate::p2p::P2PError) -> bool { + matches!( + err, + crate::p2p::P2PError::NoConnectionReady { .. } + | crate::p2p::P2PError::NodeNotFound { .. } + | crate::p2p::P2PError::NodeNotConnected { .. } + | crate::p2p::P2PError::NodePortNotReady { .. } + | crate::p2p::P2PError::ConnectionError { .. } + | crate::p2p::P2PError::SendFailed { .. } + | crate::p2p::P2PError::StartServerError { .. } + | crate::p2p::P2PError::Iceoryx2TransportNotStarted {} + ) + } + + fn read_side_transfer_peer(path: &std::path::Path) -> KvResult { + let buf = std::fs::read_to_string(path).map_err(|e| { + KvError::SharedMem(SharedMemError::MetaDataLoadError { + path: path.to_string_lossy().to_string(), + detail: format!("Failed to read side-transfer peer file: {}", e), + }) + })?; + serde_json::from_str(&buf).map_err(|e| { + KvError::SharedMem(SharedMemError::MetaDataLoadError { + path: path.to_string_lossy().to_string(), + detail: format!("Failed to parse side-transfer peer file: {}", e), + }) + }) + } + + async fn pick_side_transfer_peer(&self, put_id: Option<(u64, u32)>) -> Option<(String, u16)> { + // External attach auto-detects owner side workers from the shared-memory peer files. + // Owner-side config still controls whether workers exist; external callers should not + // require an extra enable flag once the owner has published ready lanes. + let owner_id = self.shared_storage_node_id().await?; + let owner_start_time = self.current_owner_start_time().await; + let peers_dir = ClientSegPool::side_transfer_peers_dir(&self.external_shared_file_path); + let entries = std::fs::read_dir(&peers_dir).ok()?; + let mut ready = Vec::new(); + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().and_then(|s| s.to_str()) != Some("json") { + continue; + } + let Ok(meta) = Self::read_side_transfer_peer(&path) else { + continue; + }; + if meta.owner_id != owner_id || meta.owner_start_time != owner_start_time { + continue; + } + let Some(member) = self + .view + .cluster_manager() + .get_member_info_cached(&meta.side_id) + else { + continue; + }; + if member + .metadata + .get("side_transfer_worker") + .is_some_and(|v| v == "true") + == false + { + continue; + } + if member + .metadata + .get(META_KEY_SHARED_STORAGE_NODE_ID) + .is_some_and(|v| v == &owner_id) + == false + { + continue; + } + if member + .metadata + .get(META_KEY_SHARED_STORAGE_NODE_START_TIME) + .and_then(|v| v.parse::().ok()) + != Some(owner_start_time) + { + continue; + } + let Some(lane_idx) = meta.worker_idx() else { + continue; + }; + ready.push((lane_idx, meta.side_id)); + } + + if ready.is_empty() { + return None; + } + ready.sort_by(|lhs, rhs| lhs.cmp(rhs)); + if let Some(put_id) = put_id { + let lane_space = ready + .iter() + .map(|(lane_idx, _)| usize::from(*lane_idx)) + .max() + .map(|max_lane_idx| max_lane_idx + 1)?; + let desired_lane = stable_side_transfer_lane_for_put(put_id, lane_space)?; + let selected = ready + .iter() + .find(|(lane_idx, _)| *lane_idx == desired_lane) + .cloned(); + if selected.is_none() { + tracing::warn!( + "side-transfer desired lane not ready locally; falling back to owner: desired_lane={} lane_space={} owner_id={}", + desired_lane, + lane_space, + owner_id + ); + } + return selected.map(|(lane_idx, side_id)| (side_id, lane_idx)); + } + + let idx = self.side_rr_next.fetch_add(1, Ordering::Relaxed); + let ready_len = ready.len(); + ready + .into_iter() + .nth(idx % ready_len) + .map(|(lane_idx, side_id)| (side_id, lane_idx)) + } + + fn remember_side_transfer_binding( + &self, + put_id: Option<(u64, u32)>, + binding: Option<(String, u16)>, + ) { + if let (Some(put_id), Some(binding)) = (put_id, binding) { + self.side_transfer_put_bindings.insert(put_id, binding); + } + } + + fn bound_side_transfer_peer(&self, put_id: Option<(u64, u32)>) -> Option<(String, u16)> { + put_id.and_then(|put_id| self.side_transfer_put_bindings.get(&put_id)) + } + + fn clear_side_transfer_binding(&self, put_id: Option<(u64, u32)>) { + if let Some(put_id) = put_id { + self.side_transfer_put_bindings.invalidate(&put_id); + } + } + + pub fn short_circuit_put_payload_path_enabled(&self) -> bool { + self.short_circuit_put_payload_path + } + + pub async fn external_batch_put_start_rpc( + &self, + req: ExternalBatchPutStartReq, + ) -> KvResult { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = self + .rpc_caller_external_batch_put_start + .call_with_transport_policy( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: req, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK { + return Err(KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + )); + } + Ok(resp.serialize_part) + } + + pub async fn external_batch_put_transfer_end_rpc( + &self, + req: ExternalBatchPutTransferEndReq, + ) -> KvResult { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = self + .rpc_caller_external_batch_put_transfer_end + .call_with_transport_policy( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: req, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK { + return Err(KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + )); + } + Ok(resp.serialize_part) + } + + pub async fn external_batch_put_commit_rpc( + &self, + req: ExternalBatchPutCommitReq, + ) -> KvResult { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = self + .rpc_caller_external_batch_put_commit + .call_with_transport_policy( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: req, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK { + return Err(KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + )); + } + Ok(resp.serialize_part) + } + + pub async fn external_put_revoke_rpc( + &self, + req: ExternalPutRevokeReq, + ) -> KvResult { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = self + .rpc_caller_external_put_revoke + .call( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: req, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + .map_err(KvError::from)?; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK { + return Err(KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + )); + } + Ok(resp.serialize_part) + } + + async fn call_put_start_with_side_fallback( + &self, + owner_id: String, + req: MsgPack, + ) -> KvResult<(MsgPack, Option<(String, u16)>)> { + if let Some((side_id, lane_idx)) = self.pick_side_transfer_peer(None).await { + match self + .rpc_caller_external_put_start + .call( + self.view.p2p_module(), + side_id.clone().into(), + req.clone(), + Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), + 0, + ) + .await + { + Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), + Err(err) if Self::should_fallback_side_p2p_error(&err) => { + tracing::warn!( + "side-transfer peer unavailable for put_start; falling back to owner: side={} lane={} owner={} err={}", + side_id, + lane_idx, + owner_id, + err + ); + } + Err(err) => return Err(KvError::from(err)), + } + } + + self.rpc_caller_external_put_start + .call( + self.view.p2p_module(), + owner_id.into(), + req, + Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), + 0, + ) + .await + .map(|resp| (resp, None)) + .map_err(KvError::from) + } + + async fn call_put_commit( + &self, + owner_id: String, + req: MsgPack, + ) -> KvResult> { + self.rpc_caller_external_put_commit + .call( + self.view.p2p_module(), + owner_id.into(), + req, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + .map_err(KvError::from) + } + + async fn call_put_transfer_end_with_side_fallback( + &self, + owner_id: String, + req: MsgPack, + ) -> KvResult<(MsgPack, Option<(String, u16)>)> { + let mut attempted_side = None; + if let Some((side_id, lane_idx)) = self.bound_side_transfer_peer(req.serialize_part.put_id) + { + attempted_side = Some((side_id.clone(), lane_idx)); + match self + .rpc_caller_external_put_transfer_end + .call( + self.view.p2p_module(), + side_id.clone().into(), + req.clone(), + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + { + Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), + Err(err) if Self::should_fallback_side_p2p_error(&err) => { + tracing::warn!( + "bound side-transfer peer unavailable for put_transfer_end; retrying alternate path: side={} lane={} owner={} err={}", + side_id, + lane_idx, + owner_id, + err + ); + } + Err(err) => return Err(KvError::from(err)), + } + } + + if let Some((side_id, lane_idx)) = self + .pick_side_transfer_peer(req.serialize_part.put_id) + .await + { + if attempted_side.as_ref() != Some(&(side_id.clone(), lane_idx)) { + match self + .rpc_caller_external_put_transfer_end + .call( + self.view.p2p_module(), + side_id.clone().into(), + req.clone(), + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + { + Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), + Err(err) if Self::should_fallback_side_p2p_error(&err) => { + tracing::warn!( + "side-transfer peer unavailable; falling back to owner: side={} lane={} owner={} err={}", + side_id, + lane_idx, + owner_id, + err + ); + } + Err(err) => return Err(KvError::from(err)), + } + } + } + + self.rpc_caller_external_put_transfer_end + .call( + self.view.p2p_module(), + owner_id.into(), + req, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + 0, + ) + .await + .map(|resp| (resp, None)) + .map_err(KvError::from) + } + + /// Check a batch of keys in the external storage (loop+wait). + pub async fn batch_is_exist( + &self, + keys: Vec, + allow_local_snapshot: bool, + ) -> KvResult> { + tracing::debug!( + "External batch_is_exist request: batch_len={}, allow_local_snapshot={}", + keys.len(), + allow_local_snapshot + ); + if keys.is_empty() { + return Ok(Vec::new()); + } + + let mut prev_owner_start_time = self.current_owner_start_time().await; + let mut recover_attempts = 0usize; + if self.base_ptr().await.is_err() { + let path = self.shared_memory_path(); + tracing::info!( + "ExternalClientApi.batch_is_exist waiting for owner at: {}", + path + ); + let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; + } + + loop { + let mut results = vec![false; keys.len()]; + let mut missing_indices = Vec::new(); + let mut missing_keys = Vec::new(); + for (idx, key) in keys.iter().enumerate() { + if allow_local_snapshot && self.try_get_local_complete_holder(key).await.is_some() { + results[idx] = true; + continue; + } + missing_indices.push(idx); + missing_keys.push(key.clone()); + } + if missing_keys.is_empty() { + return Ok(results); + } + + let req = MsgPack { + serialize_part: ExternalBatchIsExistReq { + keys: missing_keys.clone(), + allow_local_snapshot, + started_time: self.current_owner_start_time().await, + }, + raw_bytes: Vec::new(), + }; + + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = match self + .rpc_caller_external_batch_is_exist + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await + { + Ok(resp) => resp, + Err(e) => { + let err = KvError::from(e); + if matches!(&err, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "batch_is_exist: transient P2P error; retrying after owner-state recovery check: batch_len={}, attempt={}/{}, err={}", + keys.len(), + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + err + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + return Err(err); + } + }; + + match resp.serialize_part.to_result() { + Ok(exists_list) => { + if exists_list.len() != missing_indices.len() { + break Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_is_exist response length mismatch: expected={} got={}", + missing_indices.len(), + exists_list.len() + ), + })); + } + for (idx, exists) in + missing_indices.iter().copied().zip(exists_list.into_iter()) + { + results[idx] = exists; + } + break Ok(results); + } + Err(e) => { + if matches!(&e, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { + tracing::warn!( + "batch_is_exist: OwnerStartTimeMismatch; remapping and retrying" + ); + let _ = self + .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) + .await?; + continue; + } + if matches!(&e, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "batch_is_exist: transient P2P error; retrying after owner-state recovery check: batch_len={}, attempt={}/{}, err={}", + keys.len(), + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + e + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + tracing::warn!( + "External batch_is_exist failed for batch_len {}: {}", + keys.len(), + e + ); + break Err(e); + } + } + } + } + + pub async fn observability_snapshot(&self) -> KvResult { + let mut prev_owner_start_time = self.current_owner_start_time().await; + let mut recover_attempts = 0usize; + if self.base_ptr_ro().await.is_err() { + let path = self.shared_memory_path(); + tracing::info!( + "ExternalClientApi.observability_snapshot waiting for owner at: {}", + path + ); + let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; + } + + loop { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let req = MsgPack { + serialize_part: ExternalObservabilitySnapshotReq { + started_time: self.current_owner_start_time().await, + }, + raw_bytes: Vec::new(), + }; + let resp = match self + .rpc_caller_external_observability_snapshot + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await + { + Ok(resp) => resp, + Err(e) => { + let err = KvError::from(e); + if matches!(&err, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "observability_snapshot: transient P2P error; retrying after owner-state recovery check: attempt={}/{}, err={}", + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + err + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + return Err(err); + } + }; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK + { + let err = KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + ); + if matches!(&err, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { + tracing::warn!( + "observability_snapshot: OwnerStartTimeMismatch; remapping and retrying" + ); + let _ = self + .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) + .await?; + continue; + } + if matches!(&err, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "observability_snapshot: transient P2P error in response; retrying after owner-state recovery check: attempt={}/{}, err={}", + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + err + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + return Err(err); + } + return Ok(resp.serialize_part.into_snapshot()); + } + } + + async fn master_batch_gpu_get_revoke(&self, get_ids: Vec) -> KvResult<()> { + if get_ids.is_empty() { + return Ok(()); + } + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let expected_get_ids = get_ids.clone(); + let resp: MsgPack = self + .rpc_caller_master_batch_get_revoke + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetRevokeReq { get_ids }, + raw_bytes: Vec::new(), + }, + None, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + if resp.serialize_part.items.len() != expected_get_ids.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "GPU Get BatchRevoke response length mismatch: expected={} got={}", + expected_get_ids.len(), + resp.serialize_part.items.len() + ), + })); + } + for (expected_get_id, item) in expected_get_ids + .into_iter() + .zip(resp.serialize_part.items.into_iter()) + { + if item.get_id != expected_get_id { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "GPU Get BatchRevoke response identity mismatch: expected={} got={}", + expected_get_id, item.get_id + ), + })); + } + crate::rpcresp_kvresult_convert::try_from_code(item.error_code, item.error_json)?; + } + Ok(()) + } + + async fn master_batch_gpu_get_done(&self, get_ids: Vec) -> KvResult<()> { + if get_ids.is_empty() { + return Ok(()); + } + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let expected_get_ids = get_ids.clone(); + let resp: MsgPack = self + .rpc_caller_master_batch_get_done + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetDoneReq { get_ids }, + raw_bytes: Vec::new(), + }, + None, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + if resp.serialize_part.items.len() != expected_get_ids.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "GPU Get BatchDone response length mismatch: expected={} got={}", + expected_get_ids.len(), + resp.serialize_part.items.len() + ), + })); + } + for (expected_get_id, item) in expected_get_ids + .into_iter() + .zip(resp.serialize_part.items.into_iter()) + { + if item.get_id != expected_get_id { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "GPU Get BatchDone response identity mismatch: expected={} got={}", + expected_get_id, item.get_id + ), + })); + } + crate::rpcresp_kvresult_convert::try_from_code(item.error_code, item.error_json)?; + if item.holder_id != 0 || item.allocation_mode != GetAllocationMode::ExternalSink { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "GPU Get BatchDone returned a cache-owned target: get_id={} holder_id={} allocation_mode={:?}", + item.get_id, item.holder_id, item.allocation_mode + ), + })); + } + } + Ok(()) + } + + async fn master_batch_get_bind_external( + &self, + plan_items: &[BatchGetPlanItemResp], + destinations: &[ExternalGpuDestination], + requester_node_start_time: i64, + ) -> KvResult> { + if plan_items.len() != destinations.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "external GetBind plan/destination length mismatch: plan={} destinations={}", + plan_items.len(), + destinations.len() + ), + })); + } + let items = plan_items + .iter() + .zip(destinations) + .map(|(plan, destination)| BatchGetBindItemReq { + get_id: plan.get_id, + target: GetBindTarget::ExternalSink(GetExternalSinkTarget { + addr: destination.addr, + capacity: destination.capacity, + registration_id: destination.registration_id, + requester_node_start_time, + }), + }) + .collect(); + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let response: MsgPack = self + .rpc_caller_master_batch_get_bind + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetBindReq { items }, + raw_bytes: Vec::new(), + }, + None, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + response.serialize_part.error_code, + response.serialize_part.error_json.clone(), + )?; + if response.serialize_part.items.len() != plan_items.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external GetBind response length mismatch: expected={} got={}", + plan_items.len(), + response.serialize_part.items.len() + ), + })); + } + for (plan, bound) in plan_items.iter().zip(&response.serialize_part.items) { + if bound.get_id != plan.get_id { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external GetBind response identity mismatch: expected={} got={}", + plan.get_id, bound.get_id + ), + })); + } + crate::rpcresp_kvresult_convert::try_from_code( + bound.error_code, + bound.error_json.clone(), + )?; + } + Ok(response.serialize_part.items) + } + + async fn probe_owner_local_gets( + &self, + plan_handle: u64, + keys: &[String], + ) -> KvResult>>> { + let mut previous_owner_start_time = self.current_owner_start_time().await; + let mut recover_attempts = 0usize; + loop { + let (owner, owner_start_time, _, base_ptr, mapped_len) = + self.wait_current_owner_mapped_range().await?; + let request = MsgPack { + serialize_part: ExternalBatchGetLocalProbeReq { + plan_handle, + keys: keys.to_vec(), + req_node_id: self.view.cluster_manager().get_self_info().id.clone(), + started_time: owner_start_time, + }, + raw_bytes: Vec::new(), + }; + let response: MsgPack = match self + .rpc_caller_external_batch_get_local_probe + .call(self.view.p2p_module(), owner.into(), request, None, 0) + .await + { + Ok(response) => response, + Err(error) => { + let error = KvError::from(error); + if matches!(&error, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + let _ = self + .recover_after_p2p_error(&mut previous_owner_start_time) + .await?; + continue; + } + return Err(error); + } + }; + if response.serialize_part.error_code != OK { + let error = KvError::from_json( + response.serialize_part.error_code, + &response.serialize_part.error_json, + ); + if matches!( + &error, + KvError::Api(ApiError::OwnerStartTimeMismatch { .. }) + ) { + let _ = self + .recover_after_owner_start_time_mismatch(&mut previous_owner_start_time) + .await?; + continue; + } + return Err(error); + } + let infos = response.serialize_part.items; + let release_infos = |infos: &[Option]| { + let external_client_id = self.view.cluster_manager().get_self_info().id.clone(); + for info in infos.iter().flatten() { + if let Err(error) = self.enqueue_external_delete_ack( + external_client_id.clone(), + info.holder_id, + owner_start_time, + ) { + tracing::warn!( + plan_handle, + holder_id = info.holder_id, + %error, + "owner-local Get probe cleanup enqueue failed" + ); + } + } + }; + if infos.len() != keys.len() { + release_infos(&infos); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "owner-local Get probe length mismatch: expected={} got={}", + keys.len(), + infos.len() + ), + })); + } + let mut holders = Vec::with_capacity(keys.len()); + for (index, (key, info)) in keys.iter().zip(&infos).enumerate() { + let Some(info) = info else { + holders.push(None); + continue; + }; + let Some(end) = info.offset.checked_add(u64::from(info.len)) else { + release_infos(&infos); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "owner-local Get probe range overflow: index={} offset={} len={}", + index, info.offset, info.len + ), + })); + }; + if end > mapped_len { + release_infos(&infos); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "owner-local Get probe exceeds owner mapping: index={} end={} mapped_len={}", + index, end, mapped_len + ), + })); + } + let Some(pointer) = base_ptr.checked_add(info.offset) else { + release_infos(&infos); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "owner-local Get probe pointer overflow: index={} base={:#x} offset={}", + index, base_ptr, info.offset + ), + })); + }; + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + pointer, + info.len, + info.holder_id, + key.clone(), + self.view.cluster_manager().get_self_info().id.clone(), + self.view.clone(), + owner_start_time, + )); + self.key_weak_memholder_index + .insert(key.clone(), Arc::downgrade(&holder)); + holders.push(Some(holder)); + } + return Ok(holders); + } + } + + pub async fn get_plan( + &self, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + ) -> KvResult { + if keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "get_plan requires at least one key".to_string(), + })); + } + let group_lens = normalize_external_get_start_group_lens(keys.len(), atomic_group_lens)?; + let self_info = self.view.cluster_manager().get_self_info(); + if self_info.node_role() != NodeRole::External { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "get_plan is supported only in external-client mode".to_string(), + })); + } + let handle = self.next_gpu_get_handle.fetch_add(1, Ordering::Relaxed); + if handle == 0 { + return Err(KvError::Api(ApiError::Unknown { + detail: "get_plan local handle space exhausted".to_string(), + })); + } + + // Resolve and pin owner-local pages before consulting the cluster + // directory. Only the remaining positions are remote plan work. + let local_holders = self.probe_owner_local_gets(handle, &keys).await?; + let remote_keys = keys + .iter() + .zip(&local_holders) + .filter_map(|(key, local)| local.is_none().then_some(key.clone())) + .collect::>(); + let plan_items = if remote_keys.is_empty() { + Vec::new() + } else { + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let response: MsgPack = self + .rpc_caller_master_batch_get_plan + .call( self.view.p2p_module(), - side_id.clone().into(), - req.clone(), - Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetPlanReq { + keys: remote_keys.clone(), + }, + raw_bytes: Vec::new(), + }, + None, 0, ) .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + response.serialize_part.error_code, + response.serialize_part.error_json.clone(), + )?; + response.serialize_part.items + }; + let started_get_ids = plan_items + .iter() + .filter(|item| item.error_code == OK) + .map(|item| item.get_id) + .collect::>(); + let mut cleanup_guard = PlannedGetRevokeGuard::new( + self.view.clone_view(), + started_get_ids.clone(), + "get_plan abandoned", + ); + if plan_items.len() != remote_keys.len() { + let cleanup = finish_planned_get_revoke_cleanup( + self.view.clone_view(), + started_get_ids, + "get_plan shape failure", + ) + .await + .err() + .map(|err| err.to_string()); + if cleanup.is_none() { + cleanup_guard.disarm(); + } + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_plan response length mismatch: expected={} got={} cleanup_error={cleanup:?}", + remote_keys.len(), + plan_items.len() + ), + })); + } + for item in &plan_items { + if item.error_code != OK + && item.error_code + != crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND + { + let error = crate::rpcresp_kvresult_convert::try_from_code( + item.error_code, + item.error_json.clone(), + ) + .expect_err("non-OK GetPlan item must decode as an error"); + let cleanup = finish_planned_get_revoke_cleanup( + self.view.clone_view(), + started_get_ids, + "get_plan item failure", + ) + .await + .err() + .map(|err| err.to_string()); + if cleanup.is_none() { + cleanup_guard.disarm(); + } + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_plan item failed: error={} cleanup_error={cleanup:?}", + error + ), + })); + } + } + let mut remote_items = plan_items.into_iter(); + let mut combined_items = Vec::with_capacity(keys.len()); + for (key, local) in keys.iter().cloned().zip(local_holders) { + if let Some(holder) = local { + combined_items.push(PendingExternalGetPlanItem::Local { holder }); + } else { + let plan = remote_items + .next() + .expect("validated remote plan shape must cover every remote key"); + combined_items.push(PendingExternalGetPlanItem::Remote { key, plan }); + } + } + assert!(remote_items.next().is_none()); + let (raw_prefix_hit_len, gpu_raw_prefix_hit_len) = + external_get_plan_raw_prefixes_from_statuses(combined_items.iter().map( + |item| match item { + PendingExternalGetPlanItem::Local { .. } => (true, true), + PendingExternalGetPlanItem::Remote { plan, .. } => { + (plan.error_code == OK, plan.gpu_direct_eligible) + } + }, + )); + let transferable_len = compute_external_get_start_transfer_prefix( + raw_prefix_hit_len, + &group_lens, + prefix_best_effort, + ); + let gpu_transferable_len = compute_external_get_start_transfer_prefix( + gpu_raw_prefix_hit_len, + &group_lens, + prefix_best_effort, + ); + let kept_ids = combined_items[..transferable_len] + .iter() + .filter_map(|item| match item { + PendingExternalGetPlanItem::Remote { plan, .. } => Some(plan.get_id), + PendingExternalGetPlanItem::Local { .. } => None, + }) + .collect::>(); + let skipped_get_ids = started_get_ids + .iter() + .copied() + .filter(|get_id| !kept_ids.contains(get_id)) + .collect::>(); + if let Err(err) = self.master_batch_gpu_get_revoke(skipped_get_ids).await { + return Err(err); + } + combined_items.truncate(transferable_len); + let gpu_remote_indices = combined_items + .iter() + .take(gpu_transferable_len) + .enumerate() + .filter_map(|(index, item)| match item { + PendingExternalGetPlanItem::Remote { plan, .. } if plan.gpu_direct_eligible => { + Some(index) + } + _ => None, + }) + .collect::>(); + self.pending_external_get_plan.insert( + handle, + PendingExternalGetPlan { + items: combined_items, + transferable_len, + gpu_transferable_len, + gpu_remote_indices: gpu_remote_indices.clone(), + atomic_group_lens: group_lens, + }, + ); + cleanup_guard.disarm(); + Ok(ExternalGetPlanResp { + handle, + raw_prefix_hit_len, + gpu_raw_prefix_hit_len, + gpu_remote_indices, + }) + } + + pub async fn cancel_get_plan(&self, handle: u64) -> KvResult<()> { + let Some((_handle, plan)) = self.pending_external_get_plan.remove(&handle) else { + return Ok(()); + }; + finish_planned_get_revoke_cleanup( + self.view.clone_view(), + plan.items + .into_iter() + .filter_map(|item| match item { + PendingExternalGetPlanItem::Remote { plan, .. } => Some(plan.get_id), + PendingExternalGetPlanItem::Local { .. } => None, + }) + .collect(), + "cancel_get_plan", + ) + .await + } + + pub async fn execute_get_plan_gpu( + &self, + handle: u64, + destinations: Vec, + consume_prefix_len: usize, + transfer_concurrency: usize, + ) -> KvResult<()> { + let Some((_handle, plan)) = self.pending_external_get_plan.remove(&handle) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("execute_get_plan_gpu requires a live plan: {handle}"), + })); + }; + let expected_remote_destinations = plan + .gpu_remote_indices + .iter() + .take_while(|index| **index < consume_prefix_len) + .count(); + if transfer_concurrency == 0 + || expected_remote_destinations == 0 + || destinations.len() != expected_remote_destinations + || validate_external_get_consume_prefix( + consume_prefix_len, + plan.gpu_transferable_len, + &plan.atomic_group_lens, + ) + .is_err() + { + self.pending_external_get_plan.insert(handle, plan); + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "execute_get_plan_gpu invalid consume/remote-destinations/concurrency: consume={} destinations={} expected_remote={} concurrency={}", + consume_prefix_len, + destinations.len(), + expected_remote_destinations, + transfer_concurrency + ), + })); + } + let self_info = self.view.cluster_manager().get_self_info(); + let mut guards = Vec::with_capacity(destinations.len()); + for destination in &destinations { + match self.view.client_transfer_engine().validate_gpu_destination( + destination.registration_id, + destination.addr, + destination.capacity, + ) { + Ok(guard) => guards.push(guard), + Err(err) => { + self.pending_external_get_plan.insert(handle, plan); + return Err(err); + } + } + } + let all_get_ids = plan + .items + .iter() + .filter_map(|item| match item { + PendingExternalGetPlanItem::Remote { plan, .. } => Some(plan.get_id), + PendingExternalGetPlanItem::Local { .. } => None, + }) + .collect::>(); + let mut cleanup_guard = PlannedGetRevokeGuard::new( + self.view.clone_view(), + all_get_ids.clone(), + "execute_get_plan_gpu abandoned", + ); + if plan + .items + .iter() + .any(|item| matches!(item, PendingExternalGetPlanItem::Local { .. })) + { + let (_, current_owner_start_time, _, base_ptr, mapped_len) = + self.wait_current_owner_mapped_range().await?; + for (index, item) in plan.items.iter().enumerate() { + if let PendingExternalGetPlanItem::Local { holder } = item { + validate_external_local_holder_mapping( + index, + holder, + current_owner_start_time, + base_ptr, + mapped_len, + )?; + } + } + } + let PendingExternalGetPlan { + mut items, + atomic_group_lens, + .. + } = plan; + let tail_items = items.split_off(consume_prefix_len); + let skipped_get_ids = tail_items + .into_iter() + .filter_map(|item| match item { + PendingExternalGetPlanItem::Remote { plan, .. } => Some(plan.get_id), + PendingExternalGetPlanItem::Local { .. } => None, + }) + .collect::>(); + let mut remote_plans = Vec::with_capacity(destinations.len()); + let mut planned_cpu_items = Vec::new(); + let mut planned_cpu_sources = Vec::new(); + let mut local_holders = Vec::new(); + let mut value_ptrs = Vec::with_capacity(consume_prefix_len); + let mut destination_index = 0usize; + for (source_index, item) in items.into_iter().enumerate() { + match item { + PendingExternalGetPlanItem::Local { holder, .. } => { + // The exact owner mapping was validated immediately above. + // Keep the saved address opaque until SGLang submits the + // restore; constructing a slice here would dereference a + // mapping that could become stale after an owner restart. + value_ptrs.push(holder.addr); + local_holders.push((source_index, holder)); + } + PendingExternalGetPlanItem::Remote { key, plan } => { + if plan.gpu_direct_eligible { + let destination = &destinations[destination_index]; + value_ptrs.push(destination.addr); + remote_plans.push(plan); + destination_index += 1; + } else { + // Filled with the owner-mapped holder pointer after the + // planned CPU branch reaches its terminal. + value_ptrs.push(0); + planned_cpu_sources.push((source_index, key.clone())); + planned_cpu_items.push((key, plan.get_id)); + } + } + } + } + assert_eq!(destination_index, destinations.len()); + let bound_items = match self + .master_batch_get_bind_external(&remote_plans, &destinations, self_info.node_start_time) + .await + { + Ok(items) => items, + Err(err) => { + let cleanup = finish_planned_get_revoke_cleanup( + self.view.clone_view(), + all_get_ids, + "execute_get_plan_gpu bind failure", + ) + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + if cleanup.is_none() { + cleanup_guard.disarm(); + } + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "execute_get_plan_gpu bind failed: error={} cleanup_error={cleanup:?}", + err + ), + })); + } + }; + let mut transfer_items = Vec::with_capacity(destinations.len()); + for (index, (((planned, bound), destination), guard)) in remote_plans + .iter() + .zip(bound_items) + .zip(destinations.iter()) + .zip(guards) + .enumerate() + { + if bound.node_id != planned.node_id + || bound.src_addr != planned.src_addr + || bound.src_base_addr != planned.src_base_addr + || bound.len != planned.len + || !external_gpu_transfer_plan_geometry_is_valid( + &bound, + destination, + guard.registration().registration_id, + ) + || bound.node_id == self_info.id + { + let cleanup = finish_planned_get_revoke_cleanup( + self.view.clone_view(), + all_get_ids, + "execute_get_plan_gpu changed bind", + ) + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + if cleanup.is_none() { + cleanup_guard.disarm(); + } + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "execute_get_plan_gpu received a changed bind plan at index={index}; cleanup_error={cleanup:?}" + ), + })); + } + transfer_items.push((bound, guard)); + } + let cancel_requested = Arc::new(AtomicBool::new(false)); + let (terminal_tx, terminal_rx) = watch::channel(None); + let transfer_started_at = Instant::now(); + self.pending_external_gpu_get.insert( + handle, + PendingExternalGpuGet { + transferable_len: consume_prefix_len, + atomic_group_lens, + value_ptrs, + local_holders, + planned_cpu_sources, + cancel_requested: cancel_requested.clone(), + transfer_started_at, + terminal_rx, + }, + ); + let view = self.view.clone_view(); + let task_view = view.clone(); + view.spawn(format!("external_gpu_get_execute_{handle}"), async move { + let terminal = run_external_mixed_gpu_get_transfer_timed( + task_view, + handle, + transfer_items, + planned_cpu_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested, + ) + .await; + let _ = terminal_tx.send(Some(terminal)); + }); + cleanup_guard.disarm(); + Ok(()) + } + + pub async fn execute_get_plan_cpu( + &self, + handle: u64, + consume_prefix_len: usize, + transfer_concurrency: usize, + ) -> KvResult<()> { + let Some((_handle, plan)) = self.pending_external_get_plan.remove(&handle) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("execute_get_plan_cpu requires a live plan: {handle}"), + })); + }; + if transfer_concurrency == 0 + || validate_external_get_consume_prefix( + consume_prefix_len, + plan.transferable_len, + &plan.atomic_group_lens, + ) + .is_err() + { + self.pending_external_get_plan.insert(handle, plan); + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "execute_get_plan_cpu invalid consume/concurrency: consume={} concurrency={}", + consume_prefix_len, transfer_concurrency + ), + })); + } + let PendingExternalGetPlan { + mut items, + atomic_group_lens, + .. + } = plan; + let tail_items = items.split_off(consume_prefix_len); + let skipped_get_ids = tail_items + .into_iter() + .filter_map(|item| match item { + PendingExternalGetPlanItem::Remote { plan, .. } => Some(plan.get_id), + PendingExternalGetPlanItem::Local { .. } => None, + }) + .collect(); + let mut plan_items = Vec::new(); + let mut sources = Vec::with_capacity(consume_prefix_len); + for item in items { + match item { + PendingExternalGetPlanItem::Local { holder, .. } => { + sources.push(PendingExternalCpuSource::Local { holder }); + } + PendingExternalGetPlanItem::Remote { key, plan } => { + plan_items.push((key.clone(), plan.get_id)); + sources.push(PendingExternalCpuSource::Remote { key }); + } + } + } + let cancel_requested = Arc::new(AtomicBool::new(false)); + let (terminal_tx, terminal_rx) = watch::channel(None); + self.pending_external_planned_cpu_get.insert( + handle, + PendingExternalPlannedCpuGet { + sources, + transferable_len: consume_prefix_len, + atomic_group_lens, + cancel_requested: cancel_requested.clone(), + terminal_rx, + }, + ); + let view = self.view.clone_view(); + let task_view = view.clone(); + view.spawn(format!("external_cpu_get_execute_{handle}"), async move { + let terminal = run_external_planned_cpu_get( + task_view, + handle, + plan_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested, + ) + .await; + let _ = terminal_tx.send(Some(terminal)); + }); + Ok(()) + } + + pub async fn get_start_gpu( + &self, + keys: Vec, + destinations: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + transfer_concurrency: usize, + ) -> KvResult { + if keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "get_start_gpu requires at least one key".to_string(), + })); + } + if keys.len() != destinations.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "get_start_gpu keys/destinations length mismatch: keys={} destinations={}", + keys.len(), + destinations.len() + ), + })); + } + if transfer_concurrency == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "get_start_gpu transfer_concurrency must be > 0".to_string(), + })); + } + let group_lens = normalize_external_get_start_group_lens(keys.len(), atomic_group_lens)?; + let self_info = self.view.cluster_manager().get_self_info(); + if self_info.node_role() != NodeRole::External { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "get_start_gpu requires an external node membership".to_string(), + })); + } + + let mut destination_guards = Vec::with_capacity(destinations.len()); + let mut external_sink_targets = Vec::with_capacity(destinations.len()); + for destination in &destinations { + let guard = self + .view + .client_transfer_engine() + .validate_gpu_destination( + destination.registration_id, + destination.addr, + destination.capacity, + )?; + destination_guards.push(guard); + external_sink_targets.push(Some(GetExternalSinkTarget { + addr: destination.addr, + capacity: destination.capacity, + registration_id: destination.registration_id, + requester_node_start_time: self_info.node_start_time, + })); + } + + let master_node_id = self + .view + .cluster_manager() + .find_or_wait_master_node() + .await?; + let resp: MsgPack = self + .rpc_caller_master_batch_get_start + .call( + self.view.p2p_module(), + master_node_id.into(), + MsgPack { + serialize_part: BatchGetStartReq { + keys: keys.clone(), + prepared_targets: Vec::new(), + external_sink_targets, + }, + raw_bytes: Vec::new(), + }, + None, + 0, + ) + .await + .map_err(KvError::from)?; + crate::rpcresp_kvresult_convert::try_from_code( + resp.serialize_part.error_code, + resp.serialize_part.error_json.clone(), + )?; + + let start_items = resp.serialize_part.items; + let started_get_ids = start_items + .iter() + .filter(|item| item.error_code == OK) + .map(|item| item.get_id) + .collect::>(); + if start_items.len() != keys.len() { + let cleanup = self + .master_batch_gpu_get_revoke(started_get_ids) + .await + .err() + .map(|err| err.to_string()); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_start_gpu response length mismatch: expected={} got={} cleanup_error={:?}", + keys.len(), + start_items.len(), + cleanup + ), + })); + } + + for item in &start_items { + if item.error_code != OK + && item.error_code + != crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND + { + let err = crate::rpcresp_kvresult_convert::try_from_code( + item.error_code, + item.error_json.clone(), + ) + .expect_err("non-OK GPU GetStart item must decode as an error"); + let cleanup = self + .master_batch_gpu_get_revoke(started_get_ids) + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_start_gpu item failed: error={} cleanup_error={:?}", + err, cleanup + ), + })); + } + } + + for (idx, ((item, destination), guard)) in start_items + .iter() + .zip(destinations.iter()) + .zip(destination_guards.iter()) + .enumerate() + { + if item.error_code != OK { + continue; + } + let geometry_is_valid = external_gpu_transfer_plan_geometry_is_valid( + item, + destination, + guard.registration().registration_id, + ); + if !geometry_is_valid || item.node_id == self_info.id { + let cleanup = self + .master_batch_gpu_get_revoke(started_get_ids) + .await + .err() + .map(|cleanup_err| cleanup_err.to_string()); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_start_gpu invalid master transfer plan: index={} get_id={} source={} target={:#x} base={:#x} len={} destination={:?} cleanup_error={:?}", + idx, + item.get_id, + item.node_id, + item.target_addr, + item.target_base_addr, + item.len, + destination, + cleanup + ), + })); + } + } + + let raw_prefix_hit_len = start_items + .iter() + .take_while(|item| item.error_code == OK) + .count(); + let transferable_len = compute_external_get_start_transfer_prefix( + raw_prefix_hit_len, + &group_lens, + prefix_best_effort, + ); + let mut transfer_items = Vec::with_capacity(transferable_len); + let mut skipped_get_ids = Vec::new(); + for (idx, (item, guard)) in start_items + .into_iter() + .zip(destination_guards.into_iter()) + .enumerate() + { + if item.error_code != OK { + continue; + } + if idx < transferable_len { + transfer_items.push((item, guard)); + } else { + skipped_get_ids.push(item.get_id); + } + } + + let handle = self.next_gpu_get_handle.fetch_add(1, Ordering::Relaxed); + if handle == 0 { + let mut all_get_ids = skipped_get_ids.clone(); + all_get_ids.extend(transfer_items.iter().map(|(item, _)| item.get_id)); + let _ = self.master_batch_gpu_get_revoke(all_get_ids).await; + return Err(KvError::Api(ApiError::Unknown { + detail: "get_start_gpu local handle space exhausted".to_string(), + })); + } + let cancel_requested = Arc::new(AtomicBool::new(false)); + let (terminal_tx, terminal_rx) = watch::channel(None); + let transfer_started_at = Instant::now(); + self.pending_external_gpu_get.insert( + handle, + PendingExternalGpuGet { + transferable_len, + atomic_group_lens: group_lens, + value_ptrs: destinations + .into_iter() + .take(transferable_len) + .map(|destination| destination.addr) + .collect(), + local_holders: Vec::new(), + planned_cpu_sources: Vec::new(), + cancel_requested: cancel_requested.clone(), + transfer_started_at, + terminal_rx, + }, + ); + let view = self.view.clone_view(); + let view_task = view.clone(); + view.spawn(format!("external_gpu_get_transfer_{handle}"), async move { + let terminal = run_external_gpu_get_transfer_timed( + view_task, + transfer_items, + skipped_get_ids, + transfer_concurrency, + cancel_requested, + ) + .await; + let _ = terminal_tx.send(Some(terminal)); + }); + Ok(ExternalGpuGetStartResp { + handle, + raw_prefix_hit_len, + }) + } + + pub async fn get_transfer_gpu( + &self, + handle: u64, + consume_prefix_len: Option, + ) -> KvResult { + let consume_started_at = Instant::now(); + let Some((_, pending)) = self.pending_external_gpu_get.remove(&handle) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("get_transfer_gpu requires a live handle: {handle}"), + })); + }; + let pending_guard = + PendingRegistryEntryGuard::new(&self.pending_external_gpu_get, handle, pending); + let consumed_prefix_len = + consume_prefix_len.unwrap_or(pending_guard.entry().transferable_len); + if let Err(err) = validate_external_get_consume_prefix( + consumed_prefix_len, + pending_guard.entry().transferable_len, + &pending_guard.entry().atomic_group_lens, + ) { + return Err(err); + } + let finish_wait_started_at = Instant::now(); + let terminal_event = + match wait_external_gpu_get_terminal(pending_guard.entry().terminal_rx.clone()).await { + Ok(terminal) => terminal, + Err(err) => { + let _ = pending_guard.take(); + return Err(err); + } + }; + let finish_wait = finish_wait_started_at.elapsed(); + let mut pending = pending_guard.take(); + let timing = observe_external_gpu_get_consume_timing( + pending.transfer_started_at, + terminal_event.terminal_at, + consume_started_at, + finish_wait, + ); + let outcome = match &terminal_event.outcome { + ExternalGpuGetTerminal::Completed { .. } => "completed", + ExternalGpuGetTerminal::Revoked { .. } => "revoked", + ExternalGpuGetTerminal::Failed { .. } => "failed", + }; + let local_source_count = pending.local_holders.len(); + let planned_cpu_source_count = pending.planned_cpu_sources.len(); + let gpu_direct_source_count = pending + .transferable_len + .saturating_sub(local_source_count) + .saturating_sub(planned_cpu_source_count); + tracing::info!( + "external GPU Get consume lifecycle: handle={} transferred={} consumed={} local_sources={} planned_cpu_sources={} gpu_direct_sources={} outcome={} transfer_wall_us={} terminal_before_consume={} terminal_to_consume_us={} finish_wait_us={}", + handle, + pending.transferable_len, + consumed_prefix_len, + local_source_count, + planned_cpu_source_count, + gpu_direct_source_count, + outcome, + timing.transfer_wall_us, + timing.terminal_before_consume, + timing.terminal_to_consume_us, + timing.finish_wait_us, + ); + match terminal_event.outcome { + ExternalGpuGetTerminal::Completed { + planned_cpu_items, + planned_cpu_owner_start_time, + } => { + assert_eq!(pending.value_ptrs.len(), pending.transferable_len); + let planned_cpu_sources = std::mem::take(&mut pending.planned_cpu_sources); + let cpu_terminal_shape_is_valid = planned_cpu_sources.len() + == planned_cpu_items.len() + && planned_cpu_sources.iter().enumerate().all( + |(source_order, (source_index, _))| { + *source_index < pending.value_ptrs.len() + && pending.value_ptrs[*source_index] == 0 + && source_order.checked_sub(1).is_none_or(|previous_order| { + planned_cpu_sources[previous_order].0 < *source_index + }) + }, + ); + if !cpu_terminal_shape_is_valid { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get CPU source/terminal shape mismatch: sources={} items={} values={}", + planned_cpu_sources.len(), + planned_cpu_items.len(), + pending.value_ptrs.len() + ), + })); + } + + let needs_owner_mapping = + !pending.local_holders.is_empty() || !planned_cpu_sources.is_empty(); + let owner_mapping = if needs_owner_mapping { + match self.wait_current_owner_mapped_range().await { + Ok(mapping) => Some(mapping), + Err(err) => { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(err); + } + } + } else { + None + }; + if let Some((_, current_owner_start_time, _, base_ptr, mapped_len)) = owner_mapping + { + if let Err(err) = validate_external_local_holders_mapping( + &pending.local_holders, + current_owner_start_time, + base_ptr, + mapped_len, + ) { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(err); + } + if !planned_cpu_sources.is_empty() { + let owner_start_time = match validate_mixed_planned_cpu_terminal( + &planned_cpu_items, + planned_cpu_sources.len(), + planned_cpu_owner_start_time, + current_owner_start_time, + base_ptr, + mapped_len, + ) { + Ok(owner_start_time) => owner_start_time, + Err(err) => { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(err); + } + }; + for ((source_index, _), item) in + planned_cpu_sources.iter().zip(&planned_cpu_items) + { + if *source_index >= consumed_prefix_len { + let info = item + .external_memholder_info + .as_ref() + .expect("validated mixed CPU item must have a holder"); + if let Err(detail) = self.enqueue_external_delete_ack( + self.view.cluster_manager().get_self_info().id.clone(), + info.holder_id, + owner_start_time, + ) { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(KvError::Api(ApiError::Unknown { detail })); + } + } + } + let external_client_id = + self.view.cluster_manager().get_self_info().id.clone(); + for ((source_index, key), item) in planned_cpu_sources + .into_iter() + .zip(planned_cpu_items.iter()) + .take_while(|((source_index, _), _)| { + *source_index < consumed_prefix_len + }) + { + let info = item + .external_memholder_info + .as_ref() + .expect("validated mixed CPU item must have a holder"); + let holder_ptr = base_ptr + .checked_add(info.offset) + .expect("validated mixed CPU holder pointer must not overflow"); + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + holder_ptr, + info.len, + info.holder_id, + key.clone(), + external_client_id.clone(), + self.view.clone(), + owner_start_time, + )); + pending.value_ptrs[source_index] = holder_ptr; + self.key_weak_memholder_index + .insert(key, Arc::downgrade(&holder)); + pending.local_holders.push((source_index, holder)); + } + } + } + if let Some(index) = pending.value_ptrs[..consumed_prefix_len] + .iter() + .position(|pointer| *pointer == 0) + { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "mixed Get left a consumed source pointer unresolved: index={index}" + ), + })); + } + let bandwidth_handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("GPU get_transfer expects an IPC bandwidth handle"); + let local_holders = pending + .local_holders + .into_iter() + .filter_map(|(index, holder)| { + (index < consumed_prefix_len).then(|| { + bandwidth_handle.record_tx_bytes(u64::from(holder.len)); + holder + }) + }) + .collect(); + Ok(ExternalGpuGetTransferResp { + transferred_prefix_len: pending.transferable_len, + consumed_prefix_len, + value_ptrs: pending.value_ptrs[..consumed_prefix_len].to_vec(), + local_holders, + transfer_wall_us: timing.transfer_wall_us, + finish_wait_us: timing.finish_wait_us, + terminal_before_consume: timing.terminal_before_consume, + terminal_to_consume_us: timing.terminal_to_consume_us, + }) + } + ExternalGpuGetTerminal::Revoked { transfer_error } => { + Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_transfer_gpu was revoked: handle={} transfer_error={:?}", + handle, transfer_error + ), + })) + } + ExternalGpuGetTerminal::Failed { detail } => { + Err(KvError::Api(ApiError::Unknown { detail })) + } + } + } + + pub async fn cancel_get_transfer_gpu(&self, handle: u64) -> KvResult<()> { + let Some((_, pending)) = self.pending_external_gpu_get.remove(&handle) else { + return Ok(()); + }; + pending.cancel_requested.store(true, Ordering::Release); + match wait_external_gpu_get_terminal(pending.terminal_rx) + .await? + .outcome + { + ExternalGpuGetTerminal::Completed { + planned_cpu_items, + planned_cpu_owner_start_time, + } => { + release_optional_planned_cpu_item_holders( + self, + &planned_cpu_items, + planned_cpu_owner_start_time, + ); + Ok(()) + } + ExternalGpuGetTerminal::Revoked { .. } => Ok(()), + ExternalGpuGetTerminal::Failed { detail } => { + Err(KvError::Api(ApiError::Unknown { detail })) + } + } + } + + pub async fn batch_get_start( + &self, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + transfer_concurrency: usize, + ) -> KvResult { + tracing::debug!( + "External batch_get_start request: batch_len={}, prefix_best_effort={}, transfer_concurrency={}", + keys.len(), + prefix_best_effort, + transfer_concurrency + ); + if keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "ExternalClientApi.batch_get_start requires at least one key".to_string(), + })); + } + + let mut prev_owner_start_time = self.current_owner_start_time().await; + let mut recover_attempts = 0usize; + if self.base_ptr_ro().await.is_err() { + let path = self.shared_memory_path(); + tracing::info!( + "ExternalClientApi.batch_get_start waiting for owner at: {}", + path + ); + let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; + } + + loop { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let started_time = self.current_owner_start_time().await; + let req = MsgPack { + serialize_part: ExternalBatchGetStartReq { + keys: keys.clone(), + req_node_id: self.view.cluster_manager().get_self_info().id.clone(), + started_time, + prefix_best_effort, + atomic_group_lens: atomic_group_lens.clone(), + transfer_concurrency, + }, + raw_bytes: Vec::new(), + }; + let resp = match self + .rpc_caller_external_batch_get_start + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await { - Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), - Err(err) if Self::should_fallback_side_p2p_error(&err) => { + Ok(resp) => resp, + Err(e) => { + let err = KvError::from(e); + if matches!(&err, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "batch_get_start: transient P2P error; retrying after owner-state recovery check: batch_len={}, attempt={}/{}, err={}", + keys.len(), + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + err + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + return Err(err); + } + }; + if resp.serialize_part.error_code != OK { + let err = KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + ); + if matches!(&err, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { tracing::warn!( - "side-transfer peer unavailable for put_start; falling back to owner: side={} lane={} owner={} err={}", - side_id, - lane_idx, - owner_id, - err + "batch_get_start: OwnerStartTimeMismatch; remapping and retrying" ); + let _ = self + .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) + .await?; + continue; } - Err(err) => return Err(KvError::from(err)), + return Err(err); + } + let response = resp.serialize_part; + let _ = self + .pending_inline_external_get_start + .remove(&response.handle); + if let ExternalBatchGetStartTransferPlan::InlineLocal { items } = + &response.transfer_plan + { + validate_inline_external_get_start_plan(keys.len(), items)?; + self.pending_inline_external_get_start.insert( + response.handle, + PendingInlineExternalGetStart { + keys: keys.clone(), + items: items.clone(), + owner_start_time: started_time, + }, + ); + } + return Ok(response); + } + } + + pub async fn batch_get_transfer( + &self, + handle: u64, + keys: Vec, + consume_prefix_len: usize, + ) -> KvResult>>>> { + tracing::debug!( + "External batch_get_transfer request: handle={}, batch_len={}, consume_prefix_len={}", + handle, + keys.len(), + consume_prefix_len + ); + if consume_prefix_len == 0 || keys.len() != consume_prefix_len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch_get_transfer keys must equal the non-empty consumed prefix: keys={} consume_prefix_len={}", + keys.len(), + consume_prefix_len + ), + })); + } + if let Some((_handle, inline_plan)) = self.pending_inline_external_get_start.remove(&handle) + { + if consume_prefix_len > inline_plan.keys.len() + || keys.as_slice() != &inline_plan.keys[..consume_prefix_len] + { + let expected_keys = inline_plan.keys.clone(); + self.pending_inline_external_get_start + .insert(handle, inline_plan); + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "inline external batch_get_transfer prefix mismatch: handle={} consume_prefix_len={} available_keys={:?} got={:?}", + handle, consume_prefix_len, expected_keys, keys + ), + })); + } + + let owner_snapshot = match self.wait_current_owner_mapped_range().await { + Ok(snapshot) => snapshot, + Err(err) => { + self.pending_inline_external_get_start + .insert(handle, inline_plan); + return Err(err); + } + }; + let (_, current_owner_start_time, _, base_ptr_ro, mapped_len) = owner_snapshot; + if let Err(err) = validate_inline_external_get_owner_generation( + inline_plan.owner_start_time, + current_owner_start_time, + ) { + return Err(err); + } + if let Err(err) = + validate_inline_external_get_start_plan(inline_plan.keys.len(), &inline_plan.items) + { + self.pending_inline_external_get_start + .insert(handle, inline_plan); + return Err(err); + } + let range_validation = inline_plan.items.iter().enumerate().try_for_each( + |(idx, item)| -> KvResult<()> { + let info = item + .external_memholder_info + .as_ref() + .expect("inline plan was validated above"); + let end = info.offset.checked_add(u64::from(info.len)).ok_or_else(|| { + KvError::Api(ApiError::Unknown { + detail: format!( + "inline external get_start item range overflow: index={} offset={} len={}", + idx, info.offset, info.len + ), + }) + })?; + if end > mapped_len { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "inline external get_start item exceeds owner mapping: index={} end={} mapped_len={}", + idx, end, mapped_len + ), + })); + } + Ok(()) + }, + ); + if let Err(err) = range_validation { + self.pending_inline_external_get_start + .insert(handle, inline_plan); + return Err(err); + } + + let tail_holder_ids = + inline_external_get_tail_holder_ids(&inline_plan.items, consume_prefix_len) + .expect("validated inline plan must yield tail holder ids"); + if !tail_holder_ids.is_empty() { + let external_client_id = self.view.cluster_manager().get_self_info().id.clone(); + let mut enqueue_failures = 0usize; + for holder_id in tail_holder_ids.iter().copied() { + if let Err(err) = self.enqueue_external_delete_ack( + external_client_id.clone(), + holder_id, + inline_plan.owner_start_time, + ) { + enqueue_failures += 1; + tracing::warn!( + "External inline get_transfer could not enqueue tail holder release: handle={} holder_id={} error={}", + handle, + holder_id, + err + ); + } + } + tracing::info!( + "External inline get_transfer enqueued tail release: handle={}, consumed={}, released_tail={}, enqueue_failures={}", + handle, + consume_prefix_len, + tail_holder_ids.len(), + enqueue_failures + ); + } + + let bandwidth_handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("ExternalClientApi.batch_get_transfer expects IpcBandwidthAttributor handle to be attached"); + let external_client_id = self.view.cluster_manager().get_self_info().id.clone(); + let mut results = Vec::with_capacity(keys.len()); + for (key, item) in keys + .into_iter() + .zip(inline_plan.items.into_iter().take(consume_prefix_len)) + { + let info = item + .external_memholder_info + .expect("inline plan was validated above"); + if info.len > 0 { + bandwidth_handle.record_tx_bytes(info.len as u64); + } + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + base_ptr_ro + info.offset, + info.len, + info.holder_id, + key.clone(), + external_client_id.clone(), + self.view.clone(), + inline_plan.owner_start_time, + )); + self.key_weak_memholder_index + .insert(key, Arc::downgrade(&holder)); + results.push(Ok(Some(holder))); + } + return Ok(results); + } + if keys.is_empty() { + return Ok(Vec::new()); + } + + let mut prev_owner_start_time = self.current_owner_start_time().await; + if self.base_ptr_ro().await.is_err() { + let path = self.shared_memory_path(); + tracing::info!( + "ExternalClientApi.batch_get_transfer waiting for owner at: {}", + path + ); + let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; + } + + let started_time = self.current_owner_start_time().await; + let base_ptr = self.base_ptr_ro().await.expect( + "ExternalClientApi.batch_get_transfer requires shared memory to be ready after ensure_owner_ready", + ) as u64; + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let req = MsgPack { + serialize_part: ExternalBatchGetTransferReq { + handle, + req_node_id: self.view.cluster_manager().get_self_info().id.clone(), + started_time, + consume_prefix_len, + }, + raw_bytes: Vec::new(), + }; + let resp: MsgPack = self + .rpc_caller_external_batch_get_transfer + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await + .map_err(KvError::from)?; + if resp.serialize_part.error_code != OK { + return Err(KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + )); + } + if resp.serialize_part.items.len() != keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_get_transfer response length mismatch: expected={} got={}", + keys.len(), + resp.serialize_part.items.len() + ), + })); + } + + let bandwidth_handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("ExternalClientApi.batch_get_transfer expects IpcBandwidthAttributor handle to be attached"); + let mut results = Vec::with_capacity(keys.len()); + for (key, item) in keys.into_iter().zip(resp.serialize_part.items.into_iter()) { + if item.error_code == OK { + match item.external_memholder_info { + Some(info) => { + if info.len > 0 { + bandwidth_handle.record_tx_bytes(info.len as u64); + } + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + base_ptr + info.offset, + info.len, + info.holder_id, + key.clone(), + self.view.cluster_manager().get_self_info().id.clone(), + self.view.clone(), + started_time, + )); + self.key_weak_memholder_index + .insert(key, Arc::downgrade(&holder)); + results.push(Ok(Some(holder))); + } + None => results.push(Ok(None)), + } + continue; + } + if item.error_code + == crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND + { + results.push(Ok(None)); + continue; } + results.push(Err(KvError::from_json(item.error_code, &item.error_json))); + } + Ok(results) + } + + async fn send_batch_get_cancel_plan( + &self, + handle: u64, + started_time: i64, + transfer_plan: ExternalBatchGetCancelPlan, + ) -> KvResult<()> { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let req = MsgPack { + serialize_part: ExternalBatchGetCancelReq { + handle, + req_node_id: self.view.cluster_manager().get_self_info().id.clone(), + started_time, + transfer_plan, + }, + raw_bytes: Vec::new(), + }; + let resp = match self + .rpc_caller_external_batch_get_cancel + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await + { + Ok(resp) => resp, + Err(err) => return Err(KvError::from(err)), + }; + if resp.serialize_part.error_code == OK { + return Ok(()); } + let err = KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + ); + if matches!(&err, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { + tracing::info!( + "send_batch_get_cancel_plan: owner start_time mismatch; owner restarted, treating handle as gone" + ); + return Ok(()); + } + Err(err) + } - self.rpc_caller_external_put_start - .call( - self.view.p2p_module(), - owner_id.into(), - req, - Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), - 0, - ) - .await - .map(|resp| (resp, None)) - .map_err(KvError::from) + pub async fn cancel_batch_get_start(&self, handle: u64) -> KvResult<()> { + tracing::debug!("External cancel_batch_get_start request: handle={}", handle); + let inline_plan = self + .pending_inline_external_get_start + .remove(&handle) + .map(|(_handle, plan)| plan); + let current_owner_start_time = self.current_owner_start_time().await; + if inline_plan + .as_ref() + .is_some_and(|plan| plan.owner_start_time != current_owner_start_time) + { + tracing::info!( + "cancel_batch_get_start: inline plan owner generation is stale; treating holdings as gone" + ); + return Ok(()); + } + let started_time = inline_plan + .as_ref() + .map(|plan| plan.owner_start_time) + .unwrap_or(current_owner_start_time); + let transfer_plan = match inline_plan.as_ref() { + Some(plan) => ExternalBatchGetCancelPlan::InlineLocal { + holder_ids: plan + .items + .iter() + .filter_map(|item| { + item.external_memholder_info + .as_ref() + .map(|info| info.holder_id) + }) + .collect(), + }, + None => ExternalBatchGetCancelPlan::OwnerRpc, + }; + let cancel_result = self + .send_batch_get_cancel_plan(handle, started_time, transfer_plan) + .await; + if cancel_result.is_err() { + if let Some(plan) = inline_plan { + self.pending_inline_external_get_start.insert(handle, plan); + } + } + cancel_result } - async fn call_put_commit( + pub async fn get_start( &self, - owner_id: String, - req: MsgPack, - ) -> KvResult> { - self.rpc_caller_external_put_commit - .call( - self.view.p2p_module(), - owner_id.into(), - req, - Some(Duration::from_secs( - EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, - )), - 0, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + transfer_concurrency: usize, + ) -> KvResult { + if keys.is_empty() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "ExternalClientApi.get_start requires at least one key".to_string(), + })); + } + let started = self + .batch_get_start( + keys.clone(), + prefix_best_effort, + atomic_group_lens.clone(), + transfer_concurrency, ) - .await - .map_err(KvError::from) + .await?; + if started.raw_prefix_hit_len > keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "ExternalClientApi.get_start prefix response out of range: raw_prefix_hit_len={} keys={}", + started.raw_prefix_hit_len, + keys.len() + ), + })); + } + let group_lens = normalize_external_get_start_group_lens(keys.len(), atomic_group_lens)?; + let transferable_len = compute_external_get_start_transfer_prefix( + started.raw_prefix_hit_len, + &group_lens, + prefix_best_effort, + ); + let first_miss_index = if started.raw_prefix_hit_len < keys.len() { + Some(started.raw_prefix_hit_len) + } else { + None + }; + + self.pending_external_get_start.insert( + started.handle, + PendingExternalGetStart { + keys: keys.clone(), + transferable_len, + atomic_group_lens: group_lens, + first_miss_index, + }, + ); + + Ok(ExternalClientGetStartResp { + handle: started.handle, + raw_prefix_hit_len: started.raw_prefix_hit_len, + }) } - async fn call_put_transfer_end_with_side_fallback( + pub async fn get_transfer( &self, - owner_id: String, - req: MsgPack, - ) -> KvResult<(MsgPack, Option<(String, u16)>)> { - let mut attempted_side = None; - if let Some((side_id, lane_idx)) = self.bound_side_transfer_peer(req.serialize_part.put_id) - { - attempted_side = Some((side_id.clone(), lane_idx)); - match self - .rpc_caller_external_put_transfer_end - .call( - self.view.p2p_module(), - side_id.clone().into(), - req.clone(), - Some(Duration::from_secs( - EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, - )), - 0, - ) - .await + handle: u64, + consume_prefix_len: Option, + ) -> KvResult>>>> { + if let Some((_handle, pending)) = self.pending_external_planned_cpu_get.remove(&handle) { + let pending_guard = PendingRegistryEntryGuard::new( + &self.pending_external_planned_cpu_get, + handle, + pending, + ); + let consumed_prefix_len = + consume_prefix_len.unwrap_or(pending_guard.entry().transferable_len); + if let Err(err) = validate_external_get_consume_prefix( + consumed_prefix_len, + pending_guard.entry().transferable_len, + &pending_guard.entry().atomic_group_lens, + ) { + return Err(err); + } + let terminal = match wait_external_planned_cpu_get_terminal( + pending_guard.entry().terminal_rx.clone(), + ) + .await { - Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), - Err(err) if Self::should_fallback_side_p2p_error(&err) => { - tracing::warn!( - "bound side-transfer peer unavailable for put_transfer_end; retrying alternate path: side={} lane={} owner={} err={}", - side_id, - lane_idx, - owner_id, - err - ); + Ok(terminal) => terminal, + Err(err) => { + let _ = pending_guard.take(); + return Err(err); } - Err(err) => return Err(KvError::from(err)), + }; + let ExternalPlannedCpuGetTerminal::Completed { + items, + owner_start_time, + } = terminal + else { + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: match terminal { + ExternalPlannedCpuGetTerminal::Revoked => { + format!("planned CPU Get was revoked: handle={handle}") + } + ExternalPlannedCpuGetTerminal::Failed { detail } => detail, + ExternalPlannedCpuGetTerminal::Completed { .. } => unreachable!(), + }, + })); + }; + let response = ExternalExecutePlannedGetResp { + items: items.clone(), + error_code: OK, + error_json: String::new(), + }; + let expected_remote_items = pending_guard + .entry() + .sources + .iter() + .filter(|source| matches!(source, PendingExternalCpuSource::Remote { .. })) + .count(); + if items.len() != expected_remote_items { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned CPU Get terminal remote length mismatch: expected={} got={}", + expected_remote_items, + items.len() + ), + })); } - } - - if let Some((side_id, lane_idx)) = self - .pick_side_transfer_peer(req.serialize_part.put_id) - .await - { - if attempted_side.as_ref() != Some(&(side_id.clone(), lane_idx)) { - match self - .rpc_caller_external_put_transfer_end - .call( - self.view.p2p_module(), - side_id.clone().into(), - req.clone(), - Some(Duration::from_secs( - EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, - )), - 0, - ) - .await - { - Ok(resp) => return Ok((resp, Some((side_id, lane_idx)))), - Err(err) if Self::should_fallback_side_p2p_error(&err) => { - tracing::warn!( - "side-transfer peer unavailable; falling back to owner: side={} lane={} owner={} err={}", - side_id, - lane_idx, - owner_id, - err - ); + let owner_snapshot = match self.wait_current_owner_mapped_range().await { + Ok(snapshot) => snapshot, + Err(err) => { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(err); + } + }; + let (_, current_owner_start_time, _, base_ptr, mapped_len) = owner_snapshot; + if let Err(err) = validate_inline_external_get_owner_generation( + owner_start_time, + current_owner_start_time, + ) { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(err); + } + if let Some(local) = + pending_guard + .entry() + .sources + .iter() + .find_map(|source| match source { + PendingExternalCpuSource::Local { holder } + if holder.owner_start_time != current_owner_start_time => + { + Some(holder.owner_start_time) + } + _ => None, + }) + { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::OwnerStartTimeMismatch { + expected: current_owner_start_time, + got: local, + })); + } + for (index, item) in items.iter().enumerate() { + let Some(info) = item.external_memholder_info.as_ref() else { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned CPU Get terminal remote item has no holder: index={index}" + ), + })); + }; + let Some(end) = info.offset.checked_add(u64::from(info.len)) else { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned CPU Get holder range overflow: index={} offset={} len={}", + index, info.offset, info.len + ), + })); + }; + if end > mapped_len { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned CPU Get holder exceeds owner mapping: index={} end={} mapped_len={}", + index, end, mapped_len + ), + })); + } + if base_ptr.checked_add(info.offset).is_none() { + release_planned_cpu_response_holders(self, &response, owner_start_time); + let _ = pending_guard.take(); + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "planned CPU Get holder pointer overflow: index={} base={:#x} offset={}", + index, base_ptr, info.offset + ), + })); + } + } + let mut pending = pending_guard.take(); + let consumed_remote_items = pending.sources[..consumed_prefix_len] + .iter() + .filter(|source| matches!(source, PendingExternalCpuSource::Remote { .. })) + .count(); + for item in items.iter().skip(consumed_remote_items) { + let info = item + .external_memholder_info + .as_ref() + .expect("validated planned CPU remote item must have a holder"); + if let Err(detail) = self.enqueue_external_delete_ack( + self.view.cluster_manager().get_self_info().id.clone(), + info.holder_id, + owner_start_time, + ) { + release_planned_cpu_response_holders(self, &response, owner_start_time); + return Err(KvError::Api(ApiError::Unknown { detail })); + } + } + pending.sources.truncate(consumed_prefix_len); + let bandwidth_handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("planned CPU get_transfer expects an IPC bandwidth handle"); + let external_client_id = self.view.cluster_manager().get_self_info().id.clone(); + let mut remote_items = items.into_iter(); + let mut results = Vec::with_capacity(consumed_prefix_len); + for source in pending.sources { + match source { + PendingExternalCpuSource::Local { holder } => { + bandwidth_handle.record_tx_bytes(u64::from(holder.len)); + results.push(Ok(Some(holder))); + } + PendingExternalCpuSource::Remote { key } => { + let item = remote_items + .next() + .expect("validated remote terminal must match source positions"); + let info = item + .external_memholder_info + .expect("validated planned CPU remote item must have a holder"); + let holder_ptr = base_ptr + .checked_add(info.offset) + .expect("validated planned CPU holder pointer must not overflow"); + bandwidth_handle.record_tx_bytes(u64::from(info.len)); + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + holder_ptr, + info.len, + info.holder_id, + key.clone(), + external_client_id.clone(), + self.view.clone(), + owner_start_time, + )); + self.key_weak_memholder_index + .insert(key, Arc::downgrade(&holder)); + results.push(Ok(Some(holder))); } - Err(err) => return Err(KvError::from(err)), } } + return Ok(results); + } + let Some((_handle, entry)) = self.pending_external_get_start.remove(&handle) else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("get_transfer requires a live get-start handle: {}", handle), + })); + }; + if entry.transferable_len == 0 { + let _ = self.cancel_batch_get_start(handle).await; + let key = entry + .first_miss_index + .and_then(|idx| entry.keys.get(idx).cloned()) + .unwrap_or_else(|| format!("external_get_start_handle:{}", handle)); + return Err(KvError::Api(ApiError::KeyNotFound { key })); } + if entry.transferable_len > entry.keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "get_transfer stored prefix out of range: transferable_len={} keys={}", + entry.transferable_len, + entry.keys.len() + ), + })); + } + let consume_prefix_len = consume_prefix_len.unwrap_or(entry.transferable_len); + if let Err(err) = validate_external_get_consume_prefix( + consume_prefix_len, + entry.transferable_len, + &entry.atomic_group_lens, + ) { + self.pending_external_get_start.insert(handle, entry); + return Err(err); + } + self.batch_get_transfer( + handle, + entry.keys[..consume_prefix_len].to_vec(), + consume_prefix_len, + ) + .await + } - self.rpc_caller_external_put_transfer_end - .call( - self.view.p2p_module(), - owner_id.into(), - req, - Some(Duration::from_secs( - EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, - )), - 0, - ) - .await - .map(|resp| (resp, None)) - .map_err(KvError::from) + pub async fn cancel_get_transfer(&self, handle: u64) -> KvResult<()> { + if self.pending_external_get_plan.contains_key(&handle) { + return self.cancel_get_plan(handle).await; + } + if let Some((_handle, pending)) = self.pending_external_planned_cpu_get.remove(&handle) { + pending.cancel_requested.store(true, Ordering::Release); + return match wait_external_planned_cpu_get_terminal(pending.terminal_rx).await? { + ExternalPlannedCpuGetTerminal::Revoked => Ok(()), + ExternalPlannedCpuGetTerminal::Completed { + items, + owner_start_time, + } => { + release_planned_cpu_response_holders( + self, + &ExternalExecutePlannedGetResp { + items, + error_code: OK, + error_json: String::new(), + }, + owner_start_time, + ); + Ok(()) + } + ExternalPlannedCpuGetTerminal::Failed { detail } => { + Err(KvError::Api(ApiError::Unknown { detail })) + } + }; + } + let removed = self.pending_external_get_start.remove(&handle); + if removed.is_none() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "cancel_get_transfer requires a live get-start handle: {}", + handle + ), + })); + } + self.cancel_batch_get_start(handle).await } - /// Check if a key exists in the external storage (loop+wait) - pub async fn is_exist(&self, key: &str) -> KvResult { - tracing::debug!("External is_exist request for key: {}", key); + pub async fn batch_get( + &self, + keys: Vec, + transfer_concurrency: usize, + ) -> KvResult>>>> { + tracing::debug!( + "External batch_get request: batch_len={}, transfer_concurrency={}", + keys.len(), + transfer_concurrency + ); + if keys.is_empty() { + return Ok(Vec::new()); + } + let mut prev_owner_start_time = self.current_owner_start_time().await; let mut recover_attempts = 0usize; - if self.base_ptr().await.is_err() { - let path = self.share_mem_path(); - tracing::info!("ExternalClientApi.is_exist waiting for owner at: {}", path); + if self.base_ptr_ro().await.is_err() { + let path = self.shared_memory_path(); + tracing::info!("ExternalClientApi.batch_get waiting for owner at: {}", path); let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; } + let mut results: Vec>>>> = + (0..keys.len()).map(|_| None).collect(); + let mut missing_indices = Vec::new(); + let mut missing_keys = Vec::new(); + for (idx, key) in keys.iter().enumerate() { + if let Some(holder) = self.try_get_from_weak_cache(key).await { + results[idx] = Some(Ok(Some(holder))); + continue; + } + missing_indices.push(idx); + missing_keys.push(key.clone()); + } + if missing_keys.is_empty() { + return Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "external batch_get result slot was not populated".to_string(), + })) + }) + }) + .collect()); + } + loop { + let started_time = self.current_owner_start_time().await; + let base_ptr = self.base_ptr_ro().await.expect( + "ExternalClientApi.batch_get requires shared memory to be ready after ensure_owner_ready", + ) as u64; let req = MsgPack { - serialize_part: ExternalIsExistReq { - key: key.to_string(), - started_time: self.current_owner_start_time().await, + serialize_part: ExternalBatchGetReq { + keys: missing_keys.clone(), + req_node_id: self.view.cluster_manager().get_self_info().id.clone(), + started_time, + transfer_concurrency, }, raw_bytes: Vec::new(), }; - let owner = self.shared_storage_node_id().await.ok_or_else(|| { KvError::SharedMem(SharedMemError::NotConfigured { node_id: None, @@ -1550,66 +5625,147 @@ impl ExternalInner { }) })?; let resp = match self - .rpc_caller_external_is_exist + .rpc_caller_external_batch_get .call(self.view.p2p_module(), owner.into(), req, None, 0) .await { Ok(resp) => resp, Err(e) => { - let err = KvError::from(e); - if matches!(&err, KvError::P2p(_)) - && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS - { - recover_attempts += 1; - tracing::warn!( - "is_exist: transient P2P error; retrying after owner-state recovery check: key={}, attempt={}/{}, err={}", - key, - recover_attempts, - EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, - err - ); - let _ = self - .recover_after_p2p_error(&mut prev_owner_start_time) - .await?; - continue; - } - return Err(err); - } - }; - - match resp.serialize_part.to_result() { - Ok(exists) => break Ok(exists), - Err(e) => { - if matches!(&e, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { - tracing::warn!("is_exist: OwnerStartTimeMismatch; remapping and retrying"); - let _ = self - .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) - .await?; - continue; - } - if matches!(&e, KvError::P2p(_)) + let err = KvError::from(e); + if matches!(&err, KvError::P2p(_)) && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS { recover_attempts += 1; tracing::warn!( - "is_exist: transient P2P error; retrying after owner-state recovery check: key={}, attempt={}/{}, err={}", - key, + "batch_get: transient P2P error; retrying after owner-state recovery check: batch_len={}, attempt={}/{}, err={}", + keys.len(), recover_attempts, EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, - e + err ); let _ = self .recover_after_p2p_error(&mut prev_owner_start_time) .await?; continue; } - tracing::warn!("External is_exist failed for key: {}, error: {}", key, e); - break Err(e); + return Err(err); + } + }; + if resp.serialize_part.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK + { + let err = KvError::from_json( + resp.serialize_part.error_code, + &resp.serialize_part.error_json, + ); + if matches!(&err, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { + tracing::warn!("batch_get: OwnerStartTimeMismatch; remapping and retrying"); + let _ = self + .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) + .await?; + continue; + } + if matches!(&err, KvError::P2p(_)) + && recover_attempts < EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS + { + recover_attempts += 1; + tracing::warn!( + "batch_get: transient P2P error; retrying after owner-state recovery check: batch_len={}, attempt={}/{}, err={}", + keys.len(), + recover_attempts, + EXTERNAL_RPC_P2P_RECOVER_MAX_ATTEMPTS, + err + ); + let _ = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + return Err(err); + } + if resp.serialize_part.items.len() != missing_indices.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_get response length mismatch: expected={} got={}", + missing_indices.len(), + resp.serialize_part.items.len() + ), + })); + } + + let handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("ExternalClientApi.batch_get expects IpcBandwidthAttributor handle to be attached"); + for (idx, item) in missing_indices + .iter() + .copied() + .zip(resp.serialize_part.items.into_iter()) + { + if item.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK { + match item.external_memholder_info { + Some(info) => { + if info.len > 0 { + handle.record_tx_bytes(info.len as u64); + } + let holder = Arc::new(ExternalMemHolder::new( + info.offset, + base_ptr + info.offset, + info.len, + info.holder_id, + keys[idx].clone(), + self.view.cluster_manager().get_self_info().id.clone(), + self.view.clone(), + started_time, + )); + self.key_weak_memholder_index + .insert(keys[idx].clone(), Arc::downgrade(&holder)); + results[idx] = Some(Ok(Some(holder))); + } + None => { + results[idx] = Some(Ok(None)); + } + } + continue; + } + if item.error_code + == crate::rpcresp_kvresult_convert::msg_and_error::codes_api::API_KEY_NOT_FOUND + { + results[idx] = Some(Ok(None)); + continue; } + results[idx] = Some(Err(KvError::from_json(item.error_code, &item.error_json))); } + + return Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "external batch_get result slot was not populated".to_string(), + })) + }) + }) + .collect()); } } + pub async fn is_exist_with_local_snapshot( + &self, + key: &str, + allow_local_snapshot: bool, + ) -> KvResult { + let mut results = self + .batch_is_exist(vec![key.to_string()], allow_local_snapshot) + .await?; + Ok(results.pop().unwrap_or(false)) + } + + /// Check if a key exists in the external storage (loop+wait) + pub async fn is_exist(&self, key: &str) -> KvResult { + self.is_exist_with_local_snapshot(key, false).await + } + /// External Get operation (outer): retry + wait wrapper around get_inner pub async fn get( &self, @@ -1620,7 +5776,7 @@ impl ExternalInner { // Ensure external mode configured; if not, block until owner is ready once let mut prev_owner_start_time = self.current_owner_start_time().await; if self.base_ptr().await.is_err() { - let path = self.share_mem_path(); + let path = self.shared_memory_path(); tracing::info!( "ExternalClientApi.get detected unmapped shared memory; waiting at: {}", path @@ -1803,6 +5959,8 @@ key={}, attempt={}/{}, err={}", ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let observe_sink = opts.test_observe_put_phases(); let observe_enabled = true; @@ -1816,7 +5974,7 @@ key={}, attempt={}/{}, err={}", let mut base_addr: usize = match self.base_ptr().await { Ok(addr) => addr, Err(_) => { - let path = self.share_mem_path(); + let path = self.shared_memory_path(); tracing::info!( "ExternalClientApi.put detected unmapped shared memory; waiting for owner to be ready at path: {}", path @@ -1838,6 +5996,8 @@ key={}, attempt={}/{}, err={}", base_addr, lease_id, reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, preferred_sub_cluster.as_deref(), observe_enabled, ) @@ -1890,6 +6050,8 @@ key={}, attempt={}/{}, err={}", ) -> KvResult<()> { let lease_id = opts.lease_id(); let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); let observe_sink = opts.test_observe_put_phases(); let observe_enabled = true; @@ -1905,7 +6067,7 @@ key={}, attempt={}/{}, err={}", let mut base_addr: usize = match self.base_ptr().await { Ok(addr) => addr, Err(_) => { - let path = self.share_mem_path(); + let path = self.shared_memory_path(); tracing::info!( "ExternalClientApi.put_flat_dict_ptrs detected unmapped shared memory; waiting for owner to be ready at path: {}", path @@ -1924,6 +6086,8 @@ key={}, attempt={}/{}, err={}", base_addr, lease_id, reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, preferred_sub_cluster.as_deref(), observe_enabled, ) @@ -1963,6 +6127,369 @@ key={}, attempt={}/{}, err={}", } } + pub async unsafe fn batch_put_flat_dict_ptrs( + &self, + keys: Vec, + ptrs_groups: Vec)>>, + opts: crate::client_kv_api::PutOptionalArgs, + transfer_concurrency: usize, + ) -> KvResult>> { + if keys.len() != ptrs_groups.len() { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "batch_put_flat_dict_ptrs requires keys and ptrs_groups to have the same length: keys={} ptrs_groups={}", + keys.len(), + ptrs_groups.len() + ), + })); + } + if keys.is_empty() { + return Ok(Vec::new()); + } + + let lease_id = opts.lease_id(); + let reject_if_inflight_same_key = opts.reject_if_inflight_same_key(); + let reject_if_exist_same_key = opts.reject_if_exist_same_key(); + let make_replica_task = opts.make_replica_task(); + let preferred_sub_cluster = opts.preferred_sub_cluster().map(|s| s.to_string()); + let mut payload_lens = Vec::with_capacity(ptrs_groups.len()); + for ptrs in ptrs_groups.iter() { + payload_lens.push(crate::memholder::kvclient_encode::calc_flat_dict_encoded_len(ptrs)?); + } + + let mut prev_owner_start_time = self.current_owner_start_time().await; + let mut base_addr: usize = match self.base_ptr().await { + Ok(addr) => addr, + Err(_) => { + let path = self.shared_memory_path(); + tracing::info!( + "ExternalClientApi.batch_put_flat_dict_ptrs waiting for owner at path: {}", + path + ); + self.ensure_owner_ready(&mut prev_owner_start_time).await? + } + }; + + loop { + match unsafe { + self.batch_put_inner_flat_dict_ptrs( + &keys, + &ptrs_groups, + &payload_lens, + prev_owner_start_time, + base_addr, + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster.as_deref(), + transfer_concurrency.max(1), + ) + .await + } { + Ok(results) => break Ok(results), + Err(e) => { + if matches!(&e, KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) { + tracing::warn!( + "batch_put_flat_dict_ptrs: OwnerStartTimeMismatch; remapping and retrying" + ); + base_addr = self + .recover_after_owner_start_time_mismatch(&mut prev_owner_start_time) + .await?; + continue; + } + if matches!(&e, KvError::P2p(_)) { + tracing::warn!( + "batch_put_flat_dict_ptrs: P2P error (owner/link likely offline); retrying after owner-state recovery check: {}", + e + ); + base_addr = self + .recover_after_p2p_error(&mut prev_owner_start_time) + .await?; + continue; + } + break Err(e); + } + } + } + } + + async unsafe fn batch_put_inner_flat_dict_ptrs( + &self, + keys: &[String], + ptrs_groups: &[Vec<(u8, usize, u32, u64, u32, Option)>], + payload_lens: &[u64], + started_time: i64, + base_addr: usize, + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, + preferred_sub_cluster: Option<&str>, + transfer_concurrency: usize, + ) -> KvResult>> { + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let start_req = MsgPack { + serialize_part: ExternalBatchPutStartReq { + items: keys + .iter() + .zip(payload_lens.iter()) + .map(|(key, payload_len)| ExternalBatchPutStartItemReq { + key: key.clone(), + len: *payload_len, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, + preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), + }) + .collect(), + atomic_group_lens: None, + started_time, + }, + raw_bytes: Vec::new(), + }; + let start_resp = self + .rpc_caller_external_batch_put_start + .call_with_transport_policy( + self.view.p2p_module(), + owner.clone().into(), + start_req, + Some(Duration::from_secs(EXTERNAL_PUT_START_RPC_TIMEOUT_SECS)), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if start_resp.serialize_part.error_code + != crate::rpcresp_kvresult_convert::msg_and_error::OK + { + return Err(KvError::from_json( + start_resp.serialize_part.error_code, + &start_resp.serialize_part.error_json, + )); + } + if start_resp.serialize_part.items.len() != keys.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_put_start response length mismatch: expected={} got={}", + keys.len(), + start_resp.serialize_part.items.len() + ), + })); + } + + let short_circuit_payload = self.short_circuit_put_payload_path_enabled(); + let mut results: Vec>> = (0..keys.len()).map(|_| None).collect(); + let mut commit_pending = Vec::new(); + let mut transfer_pending = Vec::new(); + let mut total_written_payload = 0u64; + + for (idx, (((key, ptrs), payload_len), start_item)) in keys + .iter() + .zip(ptrs_groups.iter()) + .zip(payload_lens.iter()) + .zip(start_resp.serialize_part.items.into_iter()) + .enumerate() + { + if start_item.error_code != crate::rpcresp_kvresult_convert::msg_and_error::OK { + results[idx] = Some(Err(KvError::from_json( + start_item.error_code, + &start_item.error_json, + ))); + continue; + } + let Some(put_id) = start_item.put_id else { + results[idx] = Some(Err(KvError::Unreachable( + crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { + rpc_input_json: format!( + "missing put_id in external batch_put_start success response; key={}", + key + ), + }, + ))); + continue; + }; + + if short_circuit_payload { + commit_pending.push(( + idx, + ExternalBatchPutCommitItemReq { + key: key.clone(), + len: *payload_len, + src_offset: start_item.src_offset, + remote_target: start_item.peer_id.is_some(), + put_id: Some(put_id), + lease_id, + }, + )); + continue; + } + + let write_ptr = (base_addr + start_item.target_offset as usize) as *mut u8; + unsafe { + crate::memholder::kvclient_encode::write_flat_dict_ptrs_to_ptr(write_ptr, ptrs); + } + total_written_payload = total_written_payload.saturating_add(*payload_len); + transfer_pending.push(( + idx, + ExternalBatchPutTransferEndItemReq { + key: key.clone(), + len: *payload_len, + src_offset: start_item.src_offset, + target_offset: start_item + .transfer_target_offset + .unwrap_or(start_item.target_offset), + peer_id: start_item.peer_id.clone(), + target_base_addr: if start_item.peer_id.is_some() { + Some(start_item.target_base_addr) + } else { + None + }, + put_id: Some(put_id), + lease_id, + }, + )); + } + + if total_written_payload > 0 { + let handle = self + .view + .cluster_manager() + .ipc_bandwidth_attributor_handle() + .expect("ExternalClientApi.batch_put_flat_dict_ptrs expects IpcBandwidthAttributor handle to be attached"); + handle.record_rx_bytes(total_written_payload); + } + + if short_circuit_payload { + if !commit_pending.is_empty() { + let commit_resp = self + .rpc_caller_external_batch_put_commit + .call_with_transport_policy( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: ExternalBatchPutCommitReq { + items: commit_pending + .iter() + .map(|(_, item)| item.clone()) + .collect(), + started_time, + }, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if commit_resp.serialize_part.error_code + != crate::rpcresp_kvresult_convert::msg_and_error::OK + { + return Err(KvError::from_json( + commit_resp.serialize_part.error_code, + &commit_resp.serialize_part.error_json, + )); + } + if commit_resp.serialize_part.items.len() != commit_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_put_commit response length mismatch: expected={} got={}", + commit_pending.len(), + commit_resp.serialize_part.items.len() + ), + })); + } + for ((idx, _), item_resp) in commit_pending + .into_iter() + .zip(commit_resp.serialize_part.items.into_iter()) + { + if item_resp.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK { + results[idx] = Some(Ok(())); + } else { + results[idx] = Some(Err(KvError::from_json( + item_resp.error_code, + &item_resp.error_json, + ))); + } + } + } + } else if !transfer_pending.is_empty() { + let transfer_resp = self + .rpc_caller_external_batch_put_transfer_end + .call_with_transport_policy( + self.view.p2p_module(), + owner.into(), + MsgPack { + serialize_part: ExternalBatchPutTransferEndReq { + items: transfer_pending + .iter() + .map(|(_, item)| item.clone()) + .collect(), + started_time, + transfer_concurrency, + }, + raw_bytes: Vec::new(), + }, + Some(Duration::from_secs( + EXTERNAL_PUT_TRANSFER_END_RPC_TIMEOUT_SECS, + )), + RpcTransportPolicy::ForceTransport, + 0, + ) + .await + .map_err(KvError::from)?; + if transfer_resp.serialize_part.error_code + != crate::rpcresp_kvresult_convert::msg_and_error::OK + { + return Err(KvError::from_json( + transfer_resp.serialize_part.error_code, + &transfer_resp.serialize_part.error_json, + )); + } + if transfer_resp.serialize_part.items.len() != transfer_pending.len() { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external batch_put_transfer_end response length mismatch: expected={} got={}", + transfer_pending.len(), + transfer_resp.serialize_part.items.len() + ), + })); + } + for ((idx, _), item_resp) in transfer_pending + .into_iter() + .zip(transfer_resp.serialize_part.items.into_iter()) + { + if item_resp.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK { + results[idx] = Some(Ok(())); + } else { + results[idx] = Some(Err(KvError::from_json( + item_resp.error_code, + &item_resp.error_json, + ))); + } + } + } + + Ok(results + .into_iter() + .map(|item| { + item.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "external batch_put result slot was not populated".to_string(), + })) + }) + }) + .collect()) + } + async unsafe fn put_inner_flat_dict_ptrs( &self, key: &str, @@ -1972,6 +6499,8 @@ key={}, attempt={}/{}, err={}", base_addr: usize, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, preferred_sub_cluster: Option<&str>, observe_enabled: bool, ) -> KvResult { @@ -1981,6 +6510,8 @@ key={}, attempt={}/{}, err={}", key: key.to_string(), len: payload_len, reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), started_time, test_observe_put_phases: true, @@ -1995,7 +6526,7 @@ key={}, attempt={}/{}, err={}", })?; let put_start_rpc_started_at = observe_enabled.then(Instant::now); let (put_resp, put_start_side) = self - .call_put_start_with_side_fallback(owner, put_start_req) + .call_put_start_with_side_fallback(owner.clone(), put_start_req) .await?; if let Some(started_at) = put_start_rpc_started_at { trace.external_put_start_rpc_us = duration_to_i64_us(started_at.elapsed()); @@ -2012,9 +6543,16 @@ key={}, attempt={}/{}, err={}", self.remember_side_transfer_binding(put_start_ok.put_id, put_start_side); if self.short_circuit_put_payload_path_enabled() { + let remote_target = put_start_ok + .peer_id + .as_deref() + .is_some_and(|peer| peer != owner.as_str()); let commit_req = MsgPack { serialize_part: ExternalPutCommitReq { key: key.to_string(), + len: payload_len, + src_offset: put_start_ok.src_offset, + remote_target, put_id: put_start_ok.put_id, lease_id, started_time, @@ -2129,6 +6667,8 @@ key={}, attempt={}/{}, err={}", base_addr: usize, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + make_replica_task: bool, preferred_sub_cluster: Option<&str>, observe_enabled: bool, ) -> KvResult { @@ -2139,6 +6679,8 @@ key={}, attempt={}/{}, err={}", key: key.to_string(), len: value.len() as u64, reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task, preferred_sub_cluster: preferred_sub_cluster.map(|s| s.to_string()), started_time, test_observe_put_phases: true, @@ -2153,7 +6695,7 @@ key={}, attempt={}/{}, err={}", })?; let put_start_rpc_started_at = observe_enabled.then(Instant::now); let (put_resp, put_start_side) = self - .call_put_start_with_side_fallback(owner, put_start_req) + .call_put_start_with_side_fallback(owner.clone(), put_start_req) .await?; if let Some(started_at) = put_start_rpc_started_at { trace.external_put_start_rpc_us = duration_to_i64_us(started_at.elapsed()); @@ -2169,9 +6711,16 @@ key={}, attempt={}/{}, err={}", self.remember_side_transfer_binding(put_start_ok.put_id, put_start_side); if self.short_circuit_put_payload_path_enabled() { + let remote_target = put_start_ok + .peer_id + .as_deref() + .is_some_and(|peer| peer != owner.as_str()); let commit_req = MsgPack { serialize_part: ExternalPutCommitReq { key: key.to_string(), + len: value.len() as u64, + src_offset: put_start_ok.src_offset, + remote_target, put_id: put_start_ok.put_id, lease_id, started_time, @@ -2295,7 +6844,7 @@ key={}, attempt={}/{}, err={}", let mut prev_owner_start_time = self.current_owner_start_time().await; let mut recover_attempts = 0usize; if self.base_ptr().await.is_err() { - let path = self.share_mem_path(); + let path = self.shared_memory_path(); tracing::info!("ExternalClientApi.delete waiting for owner at: {}", path); let _ = self.ensure_owner_ready(&mut prev_owner_start_time).await?; } @@ -2436,6 +6985,72 @@ key={}, attempt={}/{}, err={}", Ok(()) } + pub(crate) fn enqueue_external_delete_ack( + &self, + external_client_id: String, + holder_id: u64, + owner_start_time: i64, + ) -> Result<(), String> { + self.external_delete_ack_batch + .enqueue(ExternalDeleteAckItem { + external_client_id, + holder_id, + owner_start_time, + }) + } + + pub(crate) fn external_delete_ack_batch_snapshot(&self) -> ExternalDeleteAckBatchSnapshot { + self.external_delete_ack_batch.snapshot() + } + + pub(crate) async fn send_external_delete_ack_batch( + &self, + external_client_id: &str, + owner_start_time: i64, + holder_ids: Vec, + ) -> KvResult { + let item_count = u64::try_from(holder_ids.len()).unwrap_or(u64::MAX); + let req = MsgPack { + serialize_part: ExternalBatchDeleteAckReq { + external_client_id: external_client_id.to_string(), + holder_ids, + started_time: owner_start_time, + }, + raw_bytes: Vec::new(), + }; + let owner = self.shared_storage_node_id().await.ok_or_else(|| { + KvError::SharedMem(SharedMemError::NotConfigured { + node_id: None, + detail: Some("Shared storage node id unavailable".to_string()), + }) + })?; + let resp = self + .rpc_caller_external_batch_delete_ack + .call(self.view.p2p_module(), owner.into(), req, None, 0) + .await + .map_err(KvError::from)?; + match resp.serialize_part.to_result() { + Ok(resp) => { + let accounted = u64::from(resp.released_count) + u64::from(resp.missing_count); + if accounted != item_count { + return Err(KvError::Api(ApiError::Unknown { + detail: format!( + "external holder ACK batch accounted for {accounted} of {item_count} items" + ), + })); + } + Ok(ExternalDeleteAckBatchSendResult::Applied { + released: resp.released_count, + missing: resp.missing_count, + }) + } + Err(KvError::Api(ApiError::OwnerStartTimeMismatch { .. })) => { + Ok(ExternalDeleteAckBatchSendResult::OwnerGenerationChanged { items: item_count }) + } + Err(err) => Err(err), + } + } + /// Allocate a client lease (external role): send request to master via P2P. /// /// Semantics: @@ -2473,14 +7088,25 @@ async fn handle_external_invalidate_weak_index( let api = view.external_client_api(); let inner = api.inner(); let mut removed_total = 0usize; - for k in req.keys.iter() { - if let Some(_v) = inner.key_weak_memholder_index.remove(k) { + let items = if req.items.is_empty() { + req.keys + .iter() + .cloned() + .map(|key| ExternalInvalidateWeakIndexItem { key }) + .collect::>() + } else { + req.items.clone() + }; + for item in items.iter() { + let key = &item.key; + let weak_removed = inner.key_weak_memholder_index.remove(key).is_some(); + if weak_removed { removed_total += 1; } } tracing::debug!( - "External invalidated weak_index for keys: {:?} (removed {} entries)", - req.keys, + "External invalidated weak_index for items: {:?} (removed {} entries)", + items, removed_total ); @@ -2659,7 +7285,8 @@ async fn handle_sync_kv_to_file_external( async fn task_wait_owner_restart( view: ExternalClientApiView, - share_mem_path: String, + shared_memory_path: String, + shared_file_path: String, current_sig_snapshot: Option, wait_start_ts: i64, old_owner_id: Option, @@ -2668,7 +7295,7 @@ async fn task_wait_owner_restart( ) -> KvResult { let shutdown_poller = view.register_shutdown_poller(); let mut cluster_rx = view.cluster_manager().listen(); - let shared_meta_path = format!("{}/shared.json", &share_mem_path); + let shared_meta_path = format!("{}/shared.json", &shared_file_path); let mut waited = 0u64; loop { if !shutdown_poller.is_running() { @@ -2679,7 +7306,8 @@ async fn task_wait_owner_restart( match probe_owner_restart_payload( &view, - &share_mem_path, + &shared_memory_path, + &shared_file_path, &shared_meta_path, current_sig_snapshot.as_ref(), wait_start_ts, @@ -2728,7 +7356,8 @@ fn read_shared_json_snapshot( async fn probe_owner_restart_payload( view: &ExternalClientApiView, - share_mem_path: &str, + shared_memory_path: &str, + shared_file_path: &str, shared_meta_path: &str, current_sig_snapshot: Option<&SharedMetaSignature>, wait_start_ts: i64, @@ -2736,16 +7365,16 @@ async fn probe_owner_restart_payload( expected_cluster_name: &str, expected_protocol_version: &str, ) -> KvResult { - if !fluxon_util::fs_watch::are_files_ready(share_mem_path, &["mmap.file"]) { + if !fluxon_util::fs_watch::are_files_ready(shared_memory_path, &["mmap.file"]) { return Ok(OwnerRestartProbe::Pending(format!( "shared memory mmap.file not ready yet: path={}", - share_mem_path + shared_memory_path ))); } - if !fluxon_util::fs_watch::are_files_ready(share_mem_path, &["shared.json"]) { + if !fluxon_util::fs_watch::are_files_ready(shared_file_path, &["shared.json"]) { return Ok(OwnerRestartProbe::Pending(format!( "shared metadata shared.json not ready yet: path={}", - share_mem_path + shared_file_path ))); } @@ -2768,13 +7397,13 @@ async fn probe_owner_restart_payload( if meta.protocol_version != expected_protocol_version { return Ok(OwnerRestartProbe::Pending(format!( "shared.json protocol_version mismatch; waiting: shm_dir='{}' shared='{}' local='{}'", - share_mem_path, meta.protocol_version, expected_protocol_version + shared_memory_path, meta.protocol_version, expected_protocol_version ))); } if meta.cluster_name != expected_cluster_name { return Ok(OwnerRestartProbe::Pending(format!( "shared.json cluster_name mismatch; waiting: shm_dir='{}' shared='{}' local='{}'", - share_mem_path, meta.cluster_name, expected_cluster_name + shared_memory_path, meta.cluster_name, expected_cluster_name ))); } if let Some(old_owner_id) = old_owner_id { @@ -2844,7 +7473,7 @@ impl LogicalModule for ExternalClientApi { async fn shutdown(&self) -> Result<(), Self::Error> { // 只在ExternalClient模式下清理共享内存映射 let ext = &self.0; - if ext.share_mem_path().is_empty() { + if ext.shared_memory_path().is_empty() { tracing::info!("ExternalClientApi shutdown (no shared memory path configured)"); return Ok(()); } diff --git a/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs b/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs new file mode 100644 index 0000000..d5875a2 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/kv_ssd_storage.rs @@ -0,0 +1,1360 @@ +use crate::master_kv_router::put::PutIDForAKey; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}; +use ::tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard}; +use foyer::{ + BlockEngineConfig, DeviceBuilder, FsDeviceBuilder, HybridCache, HybridCacheBuilder, + HybridCacheEntry, HybridCachePolicy, PsyncIoEngineConfig, Source, +}; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, +}; +use std::time::Instant; + +const MEMORY_CAPACITY_BYTES: usize = 1; +const MEMORY_SHARDS: usize = 1; +const BLOCK_SIZE_BYTES: usize = 64 * 1024 * 1024; +const FLUSH_BUFFER_BYTES: usize = BLOCK_SIZE_BYTES; +const SUBMIT_QUEUE_BYTES: usize = 2 * FLUSH_BUFFER_BYTES; +// Thirteen workload-sized values (13 * 4.5 MiB) fit below one 64 MiB block. +// This is an I/O aggregation bound only; every entry keeps an independent +// single-KV generation and result. +pub(crate) const MAX_PERSIST_BATCH_ITEMS: usize = 13; +pub const MIN_CAPACITY_BYTES: u64 = BLOCK_SIZE_BYTES as u64; + +pub fn safe_path_component(raw: &str) -> String { + format!("v1-{}", hex::encode(Sha256::digest(raw.as_bytes()))) +} + +#[derive(Clone, Debug)] +pub struct KvSsdStorageRootLimit { + pub root_dir: PathBuf, + pub limit_bytes: u64, +} + +#[derive(Clone, Debug)] +pub struct KvSsdStorageInit { + pub roots: Vec, + pub write_rate_limit_bytes_per_sec: Option, + pub write_burst_bytes: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct KvSsdStorageDeviceUsage { + pub device: String, + pub capacity_bytes: u64, + pub used_bytes: u64, + pub persist_requests: u64, + pub persist_successes: u64, + pub persist_failures: u64, + pub persist_bytes: u64, + pub persist_duration_us: u64, + pub persist_batch_requests: u64, + pub persist_batch_items: u64, + pub persist_flush_batches: u64, + pub persist_busy_batches: u64, + pub persist_admission_skips: u64, + pub persist_batch_duration_us: u64, + pub write_candidate_items: u64, + pub write_candidate_bytes: u64, + pub write_admitted_items: u64, + pub write_admitted_bytes: u64, + pub write_dropped_items: u64, + pub write_dropped_bytes: u64, + pub write_refunded_items: u64, + pub write_refunded_bytes: u64, + pub load_requests: u64, + pub load_successes: u64, + pub load_misses: u64, + pub load_failures: u64, + pub load_bytes: u64, + pub load_duration_us: u64, + pub memory_hits: u64, + pub disk_hits: u64, + pub outer_hits: u64, + pub removals: u64, +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +struct KvSsdKey { + key: String, + put_id: PutIDForAKey, +} + +#[derive(Debug)] +struct WriteRateLimiter { + rate_bytes_per_sec: u64, + burst_bytes: u64, + available_bytes: u64, + refill_remainder: u128, + last_refill: Instant, +} + +impl WriteRateLimiter { + fn new(rate_bytes_per_sec: u64, burst_bytes: u64) -> Self { + Self { + rate_bytes_per_sec, + burst_bytes, + available_bytes: burst_bytes, + refill_remainder: 0, + last_refill: Instant::now(), + } + } + + fn refill(&mut self, now: Instant) { + let elapsed_ns = now.duration_since(self.last_refill).as_nanos(); + self.last_refill = now; + if self.available_bytes == self.burst_bytes { + self.refill_remainder = 0; + return; + } + let scaled = elapsed_ns + .saturating_mul(u128::from(self.rate_bytes_per_sec)) + .saturating_add(self.refill_remainder); + let refill = scaled / 1_000_000_000u128; + self.refill_remainder = scaled % 1_000_000_000u128; + let refill = u64::try_from(refill).unwrap_or(u64::MAX); + self.available_bytes = self + .available_bytes + .saturating_add(refill) + .min(self.burst_bytes); + if self.available_bytes == self.burst_bytes { + self.refill_remainder = 0; + } + } + + fn try_consume(&mut self, bytes: u64, now: Instant) -> bool { + self.refill(now); + if bytes > self.available_bytes { + return false; + } + self.available_bytes -= bytes; + true + } +} + +type SsdEntry = HybridCacheEntry>; + +/// Pins a just-persisted entry until the master has published its SSD backing. +pub(crate) struct KvSsdPersistGuard { + _entry: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct KvSsdPersistSource { + pub key: String, + pub put_id: PutIDForAKey, + pub addr: u64, + pub len: u64, +} + +pub(crate) struct KvSsdPersistCopy { + key: String, + put_id: PutIDForAKey, + data: Vec, +} + +impl KvSsdPersistCopy { + fn len(&self) -> u64 { + u64::try_from(self.data.len()).unwrap_or(u64::MAX) + } +} + +pub(crate) struct KvSsdPersistBatchPermit { + _guard: OwnedMutexGuard<()>, +} + +/// The only owner-local SSD store used by Fluxon KV. +/// +/// DRAM ownership, distributed routing, and Put/Get terminal state remain in +/// their existing modules. This type only persists and materializes bytes for +/// one `(key, put_id)` generation. +#[derive(Debug)] +pub struct KvSsdStorage { + cache: HybridCache>, + root_dir: PathBuf, + capacity_bytes: u64, + write_rate_limiter: Option>, + // Exactly one admitted write batch may own Foyer's global durability + // barrier. New unrelated pressure batches use try_lock and fail open to + // ordinary DRAM reclaim instead of becoming SSD backlog. + persist_batch_gate: Arc>, + entry_lengths: Mutex>, + logical_used_bytes: AtomicU64, + persist_requests: AtomicU64, + persist_successes: AtomicU64, + persist_failures: AtomicU64, + persist_bytes: AtomicU64, + persist_duration_us: AtomicU64, + persist_batch_requests: AtomicU64, + persist_batch_items: AtomicU64, + persist_flush_batches: AtomicU64, + persist_busy_batches: AtomicU64, + persist_admission_skips: AtomicU64, + persist_batch_duration_us: AtomicU64, + write_candidate_items: AtomicU64, + write_candidate_bytes: AtomicU64, + write_admitted_items: AtomicU64, + write_admitted_bytes: AtomicU64, + write_dropped_items: AtomicU64, + write_dropped_bytes: AtomicU64, + write_refunded_items: AtomicU64, + write_refunded_bytes: AtomicU64, + load_requests: AtomicU64, + load_successes: AtomicU64, + load_misses: AtomicU64, + load_failures: AtomicU64, + load_bytes: AtomicU64, + load_duration_us: AtomicU64, + memory_hits: AtomicU64, + disk_hits: AtomicU64, + outer_hits: AtomicU64, + removals: AtomicU64, +} + +impl KvSsdStorage { + pub async fn new(init: KvSsdStorageInit) -> KvResult { + let write_rate_limiter = match (init.write_rate_limit_bytes_per_sec, init.write_burst_bytes) + { + (None, None) => None, + (Some(rate), Some(burst)) if rate > 0 && burst > 0 => { + Some(Mutex::new(WriteRateLimiter::new(rate, burst))) + } + _ => { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "SSD write rate and burst must be configured together and be positive" + .to_string(), + })); + } + }; + let [root] = init.roots.as_slice() else { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd storage currently requires exactly one local root, got {}", + init.roots.len() + ), + })); + }; + if root.limit_bytes < MIN_CAPACITY_BYTES { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd capacity must be at least {} bytes, got {}", + BLOCK_SIZE_BYTES, root.limit_bytes + ), + })); + } + let capacity_bytes = usize::try_from(root.limit_bytes).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd capacity does not fit usize: {}", root.limit_bytes), + }) + })?; + + fs::create_dir_all(&root.root_dir) + .map_err(|err| file_error(&root.root_dir, "create root", err))?; + let storage_root = root.root_dir.join("foyer"); + match fs::remove_dir_all(&storage_root) { + Ok(()) => {} + Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} + Err(err) => return Err(file_error(&storage_root, "clear old store", err)), + } + fs::create_dir_all(&storage_root) + .map_err(|err| file_error(&storage_root, "create store", err))?; + + let device = FsDeviceBuilder::new(&storage_root) + .with_capacity(capacity_bytes) + .with_direct(true) + .build() + .map_err(|err| storage_error("build filesystem device", err))?; + let engine = BlockEngineConfig::new(device) + .with_block_size(BLOCK_SIZE_BYTES) + .with_buffer_pool_size(FLUSH_BUFFER_BYTES) + .with_submit_queue_size_threshold(SUBMIT_QUEUE_BYTES); + let cache = HybridCacheBuilder::new() + .with_name("fluxon_kv_ssd") + .with_policy(HybridCachePolicy::WriteOnInsertion) + .with_flush_on_close(false) + .memory(MEMORY_CAPACITY_BYTES) + .with_shards(MEMORY_SHARDS) + .with_weighter(|_key: &KvSsdKey, value: &Vec| value.len()) + .with_filter(|_key: &KvSsdKey, _value: &Vec| false) + .storage() + .with_io_engine_config(PsyncIoEngineConfig::new()) + .with_engine_config(engine) + .build() + .await + .map_err(|err| storage_error("build cache", err))?; + + tracing::info!( + root = %root.root_dir.display(), + capacity_bytes = root.limit_bytes, + write_rate_limit_bytes_per_sec = init.write_rate_limit_bytes_per_sec, + write_burst_bytes = init.write_burst_bytes, + direct_io = true, + "Initialized owner-local KV SSD backing" + ); + + Ok(Self { + cache, + root_dir: root.root_dir.clone(), + capacity_bytes: root.limit_bytes, + write_rate_limiter, + persist_batch_gate: Arc::new(AsyncMutex::new(())), + entry_lengths: Mutex::new(HashMap::new()), + logical_used_bytes: AtomicU64::new(0), + persist_requests: AtomicU64::new(0), + persist_successes: AtomicU64::new(0), + persist_failures: AtomicU64::new(0), + persist_bytes: AtomicU64::new(0), + persist_duration_us: AtomicU64::new(0), + persist_batch_requests: AtomicU64::new(0), + persist_batch_items: AtomicU64::new(0), + persist_flush_batches: AtomicU64::new(0), + persist_busy_batches: AtomicU64::new(0), + persist_admission_skips: AtomicU64::new(0), + persist_batch_duration_us: AtomicU64::new(0), + write_candidate_items: AtomicU64::new(0), + write_candidate_bytes: AtomicU64::new(0), + write_admitted_items: AtomicU64::new(0), + write_admitted_bytes: AtomicU64::new(0), + write_dropped_items: AtomicU64::new(0), + write_dropped_bytes: AtomicU64::new(0), + write_refunded_items: AtomicU64::new(0), + write_refunded_bytes: AtomicU64::new(0), + load_requests: AtomicU64::new(0), + load_successes: AtomicU64::new(0), + load_misses: AtomicU64::new(0), + load_failures: AtomicU64::new(0), + load_bytes: AtomicU64::new(0), + load_duration_us: AtomicU64::new(0), + memory_hits: AtomicU64::new(0), + disk_hits: AtomicU64::new(0), + outer_hits: AtomicU64::new(0), + removals: AtomicU64::new(0), + }) + } + + pub fn root_dir(&self) -> &Path { + &self.root_dir + } + + pub fn usage_snapshot(&self) -> KvSsdStorageDeviceUsage { + KvSsdStorageDeviceUsage { + device: self.root_dir.display().to_string(), + capacity_bytes: self.capacity_bytes, + used_bytes: self + .logical_used_bytes + .load(Ordering::Relaxed) + .min(self.capacity_bytes), + persist_requests: self.persist_requests.load(Ordering::Relaxed), + persist_successes: self.persist_successes.load(Ordering::Relaxed), + persist_failures: self.persist_failures.load(Ordering::Relaxed), + persist_bytes: self.persist_bytes.load(Ordering::Relaxed), + persist_duration_us: self.persist_duration_us.load(Ordering::Relaxed), + persist_batch_requests: self.persist_batch_requests.load(Ordering::Relaxed), + persist_batch_items: self.persist_batch_items.load(Ordering::Relaxed), + persist_flush_batches: self.persist_flush_batches.load(Ordering::Relaxed), + persist_busy_batches: self.persist_busy_batches.load(Ordering::Relaxed), + persist_admission_skips: self.persist_admission_skips.load(Ordering::Relaxed), + persist_batch_duration_us: self.persist_batch_duration_us.load(Ordering::Relaxed), + write_candidate_items: self.write_candidate_items.load(Ordering::Relaxed), + write_candidate_bytes: self.write_candidate_bytes.load(Ordering::Relaxed), + write_admitted_items: self.write_admitted_items.load(Ordering::Relaxed), + write_admitted_bytes: self.write_admitted_bytes.load(Ordering::Relaxed), + write_dropped_items: self.write_dropped_items.load(Ordering::Relaxed), + write_dropped_bytes: self.write_dropped_bytes.load(Ordering::Relaxed), + write_refunded_items: self.write_refunded_items.load(Ordering::Relaxed), + write_refunded_bytes: self.write_refunded_bytes.load(Ordering::Relaxed), + load_requests: self.load_requests.load(Ordering::Relaxed), + load_successes: self.load_successes.load(Ordering::Relaxed), + load_misses: self.load_misses.load(Ordering::Relaxed), + load_failures: self.load_failures.load(Ordering::Relaxed), + load_bytes: self.load_bytes.load(Ordering::Relaxed), + load_duration_us: self.load_duration_us.load(Ordering::Relaxed), + memory_hits: self.memory_hits.load(Ordering::Relaxed), + disk_hits: self.disk_hits.load(Ordering::Relaxed), + outer_hits: self.outer_hits.load(Ordering::Relaxed), + removals: self.removals.load(Ordering::Relaxed), + } + } + + pub async fn close(&self) -> KvResult<()> { + self.cache + .close() + .await + .map_err(|err| storage_error("close", err)) + } + + /// Select at most one bounded durability batch from master-confirmed + /// last-backing candidates. Smaller values win ties because they preserve + /// more independently reusable keys per admitted byte. The decision is + /// immediate: candidates outside the burst/rate budget are returned as + /// false and must be reclaimed without waiting for SSD bandwidth. + pub(crate) fn admit_owner_write_candidates(&self, lengths: &[u64]) -> Vec { + let mut order = (0..lengths.len()).collect::>(); + order.sort_by_key(|index| (lengths[*index], *index)); + let mut admitted = vec![false; lengths.len()]; + let mut admitted_items = 0usize; + let mut admitted_bytes = 0u64; + let now = Instant::now(); + let mut limiter = self + .write_rate_limiter + .as_ref() + .map(|limiter| limiter.lock()); + for index in order { + if admitted_items == MAX_PERSIST_BATCH_ITEMS { + break; + } + let len = lengths[index]; + let allowed = limiter + .as_mut() + .is_none_or(|limiter| limiter.try_consume(len, now)); + if allowed { + admitted[index] = true; + admitted_items += 1; + admitted_bytes = admitted_bytes.saturating_add(len); + } + } + let candidate_bytes = lengths.iter().copied().fold(0u64, u64::saturating_add); + let candidate_items = u64::try_from(lengths.len()).unwrap_or(u64::MAX); + let admitted_items = u64::try_from(admitted_items).unwrap_or(u64::MAX); + let dropped_items = candidate_items.saturating_sub(admitted_items); + let dropped_bytes = candidate_bytes.saturating_sub(admitted_bytes); + self.write_candidate_items + .fetch_add(candidate_items, Ordering::Relaxed); + self.write_candidate_bytes + .fetch_add(candidate_bytes, Ordering::Relaxed); + self.write_admitted_items + .fetch_add(admitted_items, Ordering::Relaxed); + self.write_admitted_bytes + .fetch_add(admitted_bytes, Ordering::Relaxed); + self.write_dropped_items + .fetch_add(dropped_items, Ordering::Relaxed); + self.write_dropped_bytes + .fetch_add(dropped_bytes, Ordering::Relaxed); + admitted + } + + /// Return rate-budget bytes when a provisionally admitted no-queue batch + /// loses the persist gate before any copy or I/O starts. This keeps gate + /// contention from silently burning future bandwidth while preserving the + /// immediate-Drop contract. + pub(crate) fn refund_owner_write_admission(&self, lengths: &[u64]) { + if lengths.is_empty() { + return; + } + let Some(limiter) = self.write_rate_limiter.as_ref() else { + return; + }; + let refunded_bytes = lengths.iter().copied().fold(0u64, u64::saturating_add); + let mut limiter = limiter.lock(); + limiter.refill(Instant::now()); + limiter.available_bytes = limiter + .available_bytes + .saturating_add(refunded_bytes) + .min(limiter.burst_bytes); + drop(limiter); + self.write_refunded_items.fetch_add( + u64::try_from(lengths.len()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + self.write_refunded_bytes + .fetch_add(refunded_bytes, Ordering::Relaxed); + } + + /// Copy exact source bytes into owner-owned buffers. Callers may release + /// source holders after this returns; durability uses only these copies. + pub(crate) fn copy_batch_from_addrs( + sources: &[KvSsdPersistSource], + ) -> Vec> { + sources + .iter() + .map(|source| { + let len = checked_len(source.len, "persist copy")?; + let data = + unsafe { std::slice::from_raw_parts(source.addr as *const u8, len).to_vec() }; + Ok(KvSsdPersistCopy { + key: source.key.clone(), + put_id: source.put_id, + data, + }) + }) + .collect() + } + + pub(crate) async fn persist_batch_from_addrs( + &self, + sources: &[KvSsdPersistSource], + ) -> Vec>> { + if sources.is_empty() { + return Vec::new(); + } + match self.try_acquire_persist_batch(sources.len()) { + Ok(Some(permit)) => { + self.persist_batch_from_copies_with_permit( + permit, + Self::copy_batch_from_addrs(sources), + ) + .await + } + Ok(None) => sources.iter().map(|_| Ok(None)).collect(), + Err(_) => sources + .iter() + .map(|_| Err(persist_batch_size_error(sources.len()))) + .collect(), + } + } + + pub(crate) fn try_acquire_persist_batch( + &self, + item_count: usize, + ) -> KvResult> { + self.persist_requests + .fetch_add(item_count as u64, Ordering::Relaxed); + self.persist_batch_requests.fetch_add(1, Ordering::Relaxed); + self.persist_batch_items + .fetch_add(item_count as u64, Ordering::Relaxed); + if item_count == 0 || item_count > MAX_PERSIST_BATCH_ITEMS { + self.persist_failures + .fetch_add(item_count as u64, Ordering::Relaxed); + return Err(persist_batch_size_error(item_count)); + } + let Ok(guard) = self.persist_batch_gate.clone().try_lock_owned() else { + self.persist_busy_batches.fetch_add(1, Ordering::Relaxed); + self.persist_admission_skips + .fetch_add(item_count as u64, Ordering::Relaxed); + return Ok(None); + }; + Ok(Some(KvSsdPersistBatchPermit { _guard: guard })) + } + + /// Persist owner-owned copies through the canonical durability path. + #[cfg(test)] + pub(crate) async fn persist_batch_from_copies( + &self, + copies: Vec>, + ) -> Vec>> { + if copies.is_empty() { + return Vec::new(); + } + let item_count = copies.len(); + match self.try_acquire_persist_batch(item_count) { + Ok(Some(permit)) => { + self.persist_batch_from_copies_with_permit(permit, copies) + .await + } + Ok(None) => copies.into_iter().map(|copy| copy.map(|_| None)).collect(), + Err(_) => copies + .into_iter() + .map(|copy| copy.and_then(|_| Err(persist_batch_size_error(item_count)))) + .collect(), + } + } + + pub(crate) async fn persist_batch_from_copies_with_permit( + &self, + _permit: KvSsdPersistBatchPermit, + copies: Vec>, + ) -> Vec>> { + let item_count = copies.len(); + let lengths = copies + .iter() + .map(|copy| copy.as_ref().map_or(0, KvSsdPersistCopy::len)) + .collect::>(); + let started_at = Instant::now(); + let mut results = std::iter::repeat_with(|| None) + .take(item_count) + .collect::>>>>(); + let mut batch_keys = std::collections::HashSet::with_capacity(item_count); + let mut pending = Vec::<(usize, KvSsdKey, u64, SsdEntry)>::new(); + + for (index, copy) in copies.into_iter().enumerate() { + let copy = match copy { + Ok(copy) => copy, + Err(err) => { + results[index] = Some(Err(err)); + continue; + } + }; + let len = lengths[index]; + let cache_key = KvSsdKey { + key: copy.key.clone(), + put_id: copy.put_id, + }; + if !batch_keys.insert(cache_key.clone()) { + results[index] = Some(Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd persist batch contains duplicate generation: key={} put_id=({},{})", + copy.key, copy.put_id.0, copy.put_id.1 + ), + }))); + continue; + } + let existing_len = self.entry_lengths.lock().get(&cache_key).copied(); + if let Some(existing_len) = existing_len { + if existing_len != len { + results[index] = Some(Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd duplicate persist length mismatch: key={} put_id=({},{}) existing={} requested={}", + copy.key, copy.put_id.0, copy.put_id.1, existing_len, len + ), + }))); + continue; + } + // Lengths are published only after the batch durability + // barrier, so a tracked generation can be replayed directly. + if self.cache.storage().may_contains(&cache_key) { + results[index] = Some(Ok(Some(KvSsdPersistGuard { _entry: None }))); + continue; + } + self.forget_entry(&cache_key); + } + + match self + .cache + .storage_writer(cache_key.clone()) + .force() + .insert(copy.data) + { + Some(entry) => pending.push((index, cache_key, len, entry)), + None => { + results[index] = Some(Err(KvError::Api(ApiError::FileWriteError { + path: self.root_dir.display().to_string(), + offset: 0, + detail: format!( + "SSD admission rejected: key={} put_id=({},{})", + copy.key, copy.put_id.0, copy.put_id.1 + ), + }))); + } + } + } + + if !pending.is_empty() { + self.persist_flush_batches.fetch_add(1, Ordering::Relaxed); + // One global barrier covers every independently keyed insertion in + // this admitted batch. No per-item Wait is interleaved into the + // Foyer submit queue. + self.cache.storage().wait().await; + } + for (index, cache_key, len, entry) in pending { + if !self.cache.storage().may_contains(&cache_key) { + results[index] = Some(Err(KvError::Api(ApiError::FileWriteError { + path: self.root_dir.display().to_string(), + offset: 0, + detail: format!( + "SSD commit missing: key={} put_id=({},{})", + cache_key.key, cache_key.put_id.0, cache_key.put_id.1 + ), + }))); + continue; + } + let mut lengths = self.entry_lengths.lock(); + if lengths.insert(cache_key, len).is_none() { + self.logical_used_bytes.fetch_add(len, Ordering::Relaxed); + } + drop(lengths); + results[index] = Some(Ok(Some(KvSsdPersistGuard { + _entry: Some(entry), + }))); + } + + let elapsed = elapsed_us(started_at); + self.persist_batch_duration_us + .fetch_add(elapsed, Ordering::Relaxed); + self.persist_duration_us + .fetch_add(elapsed.saturating_mul(item_count as u64), Ordering::Relaxed); + let mut successes = 0u64; + let mut failures = 0u64; + let mut bytes = 0u64; + let results = results + .into_iter() + .enumerate() + .map(|(index, result)| { + let result = result.unwrap_or_else(|| { + Err(KvError::Api(ApiError::Unknown { + detail: "kv ssd batch result was not populated".to_string(), + })) + }); + match &result { + Ok(Some(_)) => { + successes = successes.saturating_add(1); + bytes = bytes.saturating_add(lengths[index]); + } + Ok(None) => {} + Err(_) => failures = failures.saturating_add(1), + } + result + }) + .collect::>(); + self.persist_successes + .fetch_add(successes, Ordering::Relaxed); + self.persist_failures.fetch_add(failures, Ordering::Relaxed); + self.persist_bytes.fetch_add(bytes, Ordering::Relaxed); + results + } + + #[cfg(test)] + async fn persist( + &self, + key: &str, + put_id: PutIDForAKey, + data: &[u8], + ) -> KvResult { + let source = KvSsdPersistSource { + key: key.to_string(), + put_id, + addr: data.as_ptr() as u64, + len: data.len() as u64, + }; + match self + .persist_batch_from_addrs(std::slice::from_ref(&source)) + .await + .pop() + .expect("one persist source must produce one result")? + { + Some(guard) => Ok(guard), + None => Err(KvError::Api(ApiError::Unknown { + detail: "test persist was skipped by an unexpected busy batch".to_string(), + })), + } + } + + pub(crate) async fn load_into_addr( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_capacity: u64, + ) -> KvResult<()> { + self.load_requests.fetch_add(1, Ordering::Relaxed); + let started_at = Instant::now(); + let result = self + .load_into_addr_inner(key, put_id, target_addr, len, target_capacity) + .await; + self.load_duration_us + .fetch_add(elapsed_us(started_at), Ordering::Relaxed); + match &result { + Ok(()) => { + self.load_successes.fetch_add(1, Ordering::Relaxed); + self.load_bytes.fetch_add(len, Ordering::Relaxed); + } + Err(KvError::Api(ApiError::KeyNotFound { .. })) => { + self.load_misses.fetch_add(1, Ordering::Relaxed); + } + Err(_) => { + self.load_failures.fetch_add(1, Ordering::Relaxed); + } + } + result + } + + async fn load_into_addr_inner( + &self, + key: &str, + put_id: PutIDForAKey, + target_addr: u64, + len: u64, + target_capacity: u64, + ) -> KvResult<()> { + if target_capacity < len { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd target too small: key={} put_id=({},{}) len={} capacity={}", + key, put_id.0, put_id.1, len, target_capacity + ), + })); + } + let len_usize = checked_len(len, "load")?; + let cache_key = KvSsdKey { + key: key.to_string(), + put_id, + }; + let entry = self + .cache + .get(&cache_key) + .await + .map_err(|err| storage_error("load", err))? + .ok_or_else(|| { + self.forget_entry(&cache_key); + KvError::Api(ApiError::KeyNotFound { + key: key.to_string(), + }) + })?; + match entry.source() { + Source::Memory => self.memory_hits.fetch_add(1, Ordering::Relaxed), + Source::Disk => self.disk_hits.fetch_add(1, Ordering::Relaxed), + Source::Outer => self.outer_hits.fetch_add(1, Ordering::Relaxed), + }; + if entry.value().len() != len_usize { + tracing::warn!( + key, + put_time_ms = put_id.0, + put_version = put_id.1, + expected_len = len, + actual_len = entry.value().len(), + "Dropping corrupt/stale KV SSD entry with a length mismatch" + ); + self.forget_entry(&cache_key); + self.cache.remove(&cache_key); + return Err(KvError::Api(ApiError::KeyNotFound { + key: key.to_string(), + })); + } + unsafe { + std::ptr::copy_nonoverlapping( + entry.value().as_ptr(), + target_addr as *mut u8, + len_usize, + ); + } + Ok(()) + } + + pub(crate) async fn remove_exact(&self, key: &str, put_id: PutIDForAKey) -> bool { + let _batch_guard = self.persist_batch_gate.lock().await; + let cache_key = KvSsdKey { + key: key.to_string(), + put_id, + }; + let removed = self.entry_lengths.lock().remove(&cache_key); + self.cache.remove(&cache_key); + if let Some(len) = removed { + self.logical_used_bytes.fetch_sub(len, Ordering::Relaxed); + self.removals.fetch_add(1, Ordering::Relaxed); + true + } else { + false + } + } + + fn forget_entry(&self, key: &KvSsdKey) { + if let Some(len) = self.entry_lengths.lock().remove(key) { + self.logical_used_bytes.fetch_sub(len, Ordering::Relaxed); + } + } +} + +fn elapsed_us(started_at: Instant) -> u64 { + u64::try_from(started_at.elapsed().as_micros()).unwrap_or(u64::MAX) +} + +fn checked_len(len: u64, operation: &str) -> KvResult { + let len = usize::try_from(len).map_err(|_| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd {operation} len does not fit usize: {len}"), + }) + })?; + if len == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!("kv ssd {operation} len must be positive"), + })); + } + Ok(len) +} + +fn persist_batch_size_error(item_count: usize) -> KvError { + KvError::Api(ApiError::InvalidArgument { + detail: format!( + "kv ssd persist batch item count is outside 1..={}: items={}", + MAX_PERSIST_BATCH_ITEMS, item_count + ), + }) +} + +fn storage_error(operation: &str, err: impl std::fmt::Display) -> KvError { + KvError::Api(ApiError::Unknown { + detail: format!("kv ssd {operation} failed: {err}"), + }) +} + +fn file_error(path: &Path, operation: &str, err: std::io::Error) -> KvError { + KvError::Api(ApiError::FileWriteError { + path: path.display().to_string(), + offset: 0, + detail: format!("{operation}: {err}"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cluster_manager::NodeID; + use crate::master_kv_router::{ + CommittedSlotReplica, KvMemoryReplica, KvNodeReplicas, KvReplicaBacking, OneKvNodesRoutes, + SsdReplicaCommitStatus, + }; + use crate::master_seg_manager::NodeTombTag; + use parking_lot::RwLock; + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::AtomicU32; + use std::time::Duration; + + fn test_root(name: &str) -> PathBuf { + let target = std::env::var_os("CARGO_TARGET_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/mnt/nvme0/mjq_build/push_sglang_fluxon_target")); + target.join("kv_ssd_tests").join(format!( + "{}-{}-{}", + name, + std::process::id(), + uuid::Uuid::new_v4() + )) + } + + #[test] + fn write_rate_limiter_refills_without_queueing() { + let start = Instant::now(); + let mut limiter = WriteRateLimiter { + rate_bytes_per_sec: 100, + burst_bytes: 100, + available_bytes: 100, + refill_remainder: 0, + last_refill: start, + }; + assert!(limiter.try_consume(60, start)); + assert!(!limiter.try_consume(50, start)); + assert!(limiter.try_consume(50, start + Duration::from_millis(500))); + assert_eq!(limiter.available_bytes, 40); + } + + #[tokio::test] + async fn owner_candidate_admission_prefers_small_values_and_one_batch() { + let root = test_root("candidate-admission"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: Some(1), + write_burst_bytes: Some(10), + }) + .await + .unwrap(); + let admitted = store.admit_owner_write_candidates(&[8, 3, 3]); + assert_eq!(admitted, vec![false, true, true]); + let usage = store.usage_snapshot(); + assert_eq!(usage.write_candidate_items, 3); + assert_eq!(usage.write_candidate_bytes, 14); + assert_eq!(usage.write_admitted_items, 2); + assert_eq!(usage.write_admitted_bytes, 6); + assert_eq!(usage.write_dropped_items, 1); + assert_eq!(usage.write_dropped_bytes, 8); + + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn busy_gate_refund_restores_owner_write_budget() { + let root = test_root("candidate-refund"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: Some(1), + write_burst_bytes: Some(10), + }) + .await + .unwrap(); + + assert_eq!(store.admit_owner_write_candidates(&[8]), vec![true]); + store.refund_owner_write_admission(&[8]); + assert_eq!(store.admit_owner_write_candidates(&[8]), vec![true]); + let usage = store.usage_snapshot(); + assert_eq!(usage.write_refunded_items, 1); + assert_eq!(usage.write_refunded_bytes, 8); + + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn persist_and_load_one_generation() { + let root = test_root("roundtrip"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: BLOCK_SIZE_BYTES as u64, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let input = vec![0x5au8; 4096]; + let guard = store.persist("key", (7, 3), &input).await.unwrap(); + let mut output = vec![0u8; input.len()]; + store + .load_into_addr( + "key", + (7, 3), + output.as_mut_ptr() as u64, + input.len() as u64, + output.len() as u64, + ) + .await + .unwrap(); + assert_eq!(output, input); + let usage = store.usage_snapshot(); + assert_eq!(usage.used_bytes, input.len() as u64); + assert_eq!(usage.persist_requests, 1); + assert_eq!(usage.persist_successes, 1); + assert_eq!(usage.persist_failures, 0); + assert_eq!(usage.persist_bytes, input.len() as u64); + assert_eq!(usage.load_requests, 1); + assert_eq!(usage.load_successes, 1); + assert_eq!(usage.load_misses, 0); + assert_eq!(usage.load_bytes, input.len() as u64); + assert_eq!(usage.memory_hits + usage.disk_hits + usage.outer_hits, 1); + + assert!(store.remove_exact("key", (7, 3)).await); + assert!(!store.remove_exact("key", (7, 3)).await); + assert_eq!(store.usage_snapshot().used_bytes, 0); + assert!( + store + .load_into_addr( + "key", + (7, 3), + output.as_mut_ptr() as u64, + input.len() as u64, + output.len() as u64, + ) + .await + .is_err() + ); + let usage = store.usage_snapshot(); + assert_eq!(usage.load_requests, 2); + assert_eq!(usage.load_misses, 1); + assert_eq!(usage.removals, 1); + drop(guard); + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn copied_persist_no_longer_reads_the_source_buffer() { + let root = test_root("copied-persist"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: BLOCK_SIZE_BYTES as u64, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let mut input = vec![0x3cu8; 4096]; + let source = KvSsdPersistSource { + key: "copied-key".to_string(), + put_id: (8, 4), + addr: input.as_ptr() as u64, + len: input.len() as u64, + }; + let copies = KvSsdStorage::copy_batch_from_addrs(std::slice::from_ref(&source)); + input.fill(0xe7); + let guard = store + .persist_batch_from_copies(copies) + .await + .pop() + .unwrap() + .unwrap() + .expect("copied persist must acquire the durability gate"); + let mut output = vec![0u8; input.len()]; + store + .load_into_addr( + "copied-key", + (8, 4), + output.as_mut_ptr() as u64, + output.len() as u64, + output.len() as u64, + ) + .await + .unwrap(); + assert!(output.iter().all(|byte| *byte == 0x3c)); + + drop(guard); + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn one_batch_uses_one_flush_barrier_and_keeps_per_key_results() { + let root = test_root("one-batch-one-barrier"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let payloads = (0..MAX_PERSIST_BATCH_ITEMS) + .map(|index| vec![u8::try_from(index + 1).unwrap(); 1024 * 1024]) + .collect::>(); + let sources = payloads + .iter() + .enumerate() + .map(|(index, payload)| KvSsdPersistSource { + key: format!("batch-victim-{index}"), + put_id: (31, u32::try_from(index).unwrap()), + addr: payload.as_ptr() as u64, + len: payload.len() as u64, + }) + .collect::>(); + let guards = store + .persist_batch_from_addrs(&sources) + .await + .into_iter() + .map(|result| result.unwrap().expect("the only batch must be admitted")) + .collect::>(); + assert_eq!(guards.len(), MAX_PERSIST_BATCH_ITEMS); + + for (index, payload) in payloads.iter().enumerate() { + let mut output = vec![0; payload.len()]; + store + .load_into_addr( + &format!("batch-victim-{index}"), + (31, u32::try_from(index).unwrap()), + output.as_mut_ptr() as u64, + output.len() as u64, + output.len() as u64, + ) + .await + .unwrap(); + assert_eq!(&output, payload); + } + let usage = store.usage_snapshot(); + assert_eq!(usage.persist_successes, MAX_PERSIST_BATCH_ITEMS as u64); + assert_eq!(usage.persist_failures, 0); + assert_eq!(usage.persist_batch_requests, 1); + assert_eq!(usage.persist_flush_batches, 1); + assert_eq!(usage.persist_busy_batches, 0); + assert_eq!(usage.persist_admission_skips, 0); + assert_eq!(usage.load_successes, MAX_PERSIST_BATCH_ITEMS as u64); + + drop(guards); + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn retrying_a_multi_key_batch_reuses_durable_generations() { + let root = test_root("multi-key-batch-retry"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let original = (0..3) + .map(|index| vec![u8::try_from(0x41 + index).unwrap(); 4096]) + .collect::>(); + let original_sources = original + .iter() + .enumerate() + .map(|(index, payload)| KvSsdPersistSource { + key: format!("retry-victim-{index}"), + put_id: (35, u32::try_from(index).unwrap()), + addr: payload.as_ptr() as u64, + len: payload.len() as u64, + }) + .collect::>(); + let first_guards = store + .persist_batch_from_addrs(&original_sources) + .await + .into_iter() + .map(|result| result.unwrap().expect("the first batch must be admitted")) + .collect::>(); + + // A retry may arrive after the caller has reused its source buffer. The + // generation identity, not the retried bytes, owns the durable result. + let retry_payloads = original + .iter() + .map(|payload| vec![0xee; payload.len()]) + .collect::>(); + let retry_sources = retry_payloads + .iter() + .enumerate() + .map(|(index, payload)| KvSsdPersistSource { + key: format!("retry-victim-{index}"), + put_id: (35, u32::try_from(index).unwrap()), + addr: payload.as_ptr() as u64, + len: payload.len() as u64, + }) + .collect::>(); + let retry_guards = store + .persist_batch_from_addrs(&retry_sources) + .await + .into_iter() + .map(|result| result.unwrap().expect("a durable retry must be replayable")) + .collect::>(); + + for (index, expected) in original.iter().enumerate() { + let mut output = vec![0; expected.len()]; + store + .load_into_addr( + &format!("retry-victim-{index}"), + (35, u32::try_from(index).unwrap()), + output.as_mut_ptr() as u64, + output.len() as u64, + output.len() as u64, + ) + .await + .unwrap(); + assert_eq!(&output, expected); + } + let expected_bytes = original.iter().map(Vec::len).sum::() as u64; + let usage = store.usage_snapshot(); + assert_eq!(usage.used_bytes, expected_bytes); + assert_eq!(usage.persist_requests, 6); + assert_eq!(usage.persist_successes, 6); + assert_eq!(usage.persist_failures, 0); + assert_eq!(usage.persist_batch_requests, 2); + assert_eq!(usage.persist_flush_batches, 1); + assert_eq!(usage.persist_busy_batches, 0); + + drop(retry_guards); + drop(first_guards); + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn busy_write_batch_skips_without_copying_or_queueing() { + let root = test_root("busy-batch-skip"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let input = vec![0xabu8; 4096]; + let source = KvSsdPersistSource { + key: "busy-skip".to_string(), + put_id: (37, 1), + addr: input.as_ptr() as u64, + len: input.len() as u64, + }; + let held = store.persist_batch_gate.lock().await; + let result = store + .persist_batch_from_addrs(std::slice::from_ref(&source)) + .await + .pop() + .unwrap() + .unwrap(); + assert!(result.is_none()); + drop(held); + + let usage = store.usage_snapshot(); + assert_eq!(usage.persist_requests, 1); + assert_eq!(usage.persist_successes, 0); + assert_eq!(usage.persist_failures, 0); + assert_eq!(usage.persist_busy_batches, 1); + assert_eq!(usage.persist_admission_skips, 1); + assert_eq!(usage.persist_flush_batches, 0); + + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } + + #[tokio::test] + async fn persist_commit_memory_evict_and_ssd_load_roundtrip() { + let root = test_root("route-roundtrip"); + let store = KvSsdStorage::new(KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir: root.clone(), + limit_bytes: MIN_CAPACITY_BYTES, + }], + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }) + .await + .unwrap(); + let owner: NodeID = "owner".to_string().into(); + let put_id = (19, 4); + let input = vec![0xa5u8; 4096]; + let persist_guard = store.persist("route-key", put_id, &input).await.unwrap(); + + let route = Arc::new(OneKvNodesRoutes { + put_id, + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + owner.clone(), + KvNodeReplicas::memory( + NodeTombTag::new(), + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: owner.clone(), + grant_id: 7, + slot_index: 3, + slot_size: input.len() as u64, + addr: input.as_ptr() as u64, + len: input.len() as u64, + base_addr: input.as_ptr() as u64, + }), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation: None, + }, + ), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + assert_eq!( + route.commit_ssd_replica(&owner, input.len() as u64), + SsdReplicaCommitStatus::Committed + ); + + // This is the terminal state produced by the already-covered exact + // owner-memory reclaim transaction: the node and SSD route remain, + // while DRAM is no longer readable. + route.node_replicas.write().get_mut(&owner).unwrap().memory = None; + { + let replicas = route.node_replicas.read(); + let owner_backings = replicas.get(&owner).unwrap(); + assert!(owner_backings.memory.is_none()); + assert_eq!( + owner_backings.ssd.as_ref().map(|ssd| ssd.len), + Some(input.len() as u64) + ); + } + + let mut output = vec![0u8; input.len()]; + store + .load_into_addr( + "route-key", + put_id, + output.as_mut_ptr() as u64, + output.len() as u64, + output.len() as u64, + ) + .await + .unwrap(); + assert_eq!(output, input); + + drop(persist_guard); + store.close().await.unwrap(); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/fluxon_rs/fluxon_kv/src/kv_test.rs b/fluxon_rs/fluxon_kv/src/kv_test.rs index 5f0a9e2..a8c60c1 100644 --- a/fluxon_rs/fluxon_kv/src/kv_test.rs +++ b/fluxon_rs/fluxon_kv/src/kv_test.rs @@ -11,8 +11,9 @@ use crate::cluster_manager::ClusterManagerRdmaControlInit; use crate::config::{ - ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, LargeFilePaths, MasterConfig, MonitoringConfig, - ProtocolConfig, ProtocolType, TestSpecConfig, TestSpecTransportMode, TransferEngineType, + ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, LargeFilePaths, MasterConfig, + MonitoringConfig, ProtocolConfig, ProtocolType, ReplicaTaskPlacementConfig, TestSpecConfig, + TestSpecTransportMode, TransferEngineType, }; use crate::run_master_with_test_overrides; use crate::{ClientRunTestOverrides, MasterRunTestOverrides, run_client_with_test_overrides}; @@ -802,7 +803,6 @@ impl KvTestRoundOptions { kv_test_run_scope() ) } - } #[derive(Clone, Debug)] @@ -842,8 +842,7 @@ fn default_client_large_file_paths( instance_key: &str, contribute_to_cluster_pool_size: &ContributeToClusterPoolSize, ) -> LargeFilePaths { - if contribute_to_cluster_pool_size.dram == 0 - && contribute_to_cluster_pool_size.vram.is_empty() + if contribute_to_cluster_pool_size.dram == 0 && contribute_to_cluster_pool_size.vram.is_empty() { return LargeFilePaths { paths: Vec::new() }; } @@ -993,6 +992,9 @@ fn new_master_launch( network: None, log_dir, master_ui: None, + replica_task_placement: ReplicaTaskPlacementConfig::default(), + replica_cache_capacity_ratio: crate::config::DEFAULT_REPLICA_CACHE_CAPACITY_RATIO, + replica_writeback_tier1_capacity_ratio: None, // Keep kv_test self-describing: each round carries the intended transfer mode. test_spec_config: kv_test_round_test_spec_config(round.round_profile), }, @@ -1032,6 +1034,7 @@ fn build_client_launch( contribute_to_cluster_pool_size: contribute_to_cluster_pool_size.clone(), protocol: options.protocol_config.unwrap_or_else(tcp_protocol_config), pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: fluxonkv_etcd_addresses, @@ -1054,6 +1057,7 @@ fn build_client_launch( &instance_key, &contribute_to_cluster_pool_size, ), + ssd_storage: None, // Mirror round intent into the generated config so logs and runtime behavior // agree on whether this launch is transfer_only vs transfer_with_rpc. test_spec_config: kv_test_round_test_spec_config(round.round_profile), @@ -1381,7 +1385,10 @@ async fn key_meta_cache_check( } } - tracing::info!("🔍 Starting PUT and GET in parallel: {}", parallel_unique_key); + tracing::info!( + "🔍 Starting PUT and GET in parallel: {}", + parallel_unique_key + ); for i in 0..10 { let (put_client, other_client) = if i % 2 == 0 { (client, client2) @@ -1420,7 +1427,9 @@ async fn key_meta_cache_check( } assert!( - put_client.client_kv_api().has_cached_key(parallel_unique_key), + put_client + .client_kv_api() + .has_cached_key(parallel_unique_key), "put client should have immediate local cache metadata for key {} after put time {}", parallel_unique_key, i @@ -2462,7 +2471,7 @@ async fn run_kv_round(round: &KvTestRoundOptions) { .kv_routes .get(CLIENT_COMMUNICATION_KEY) { - let replicas = one_kv_nodes_routes.nodes_replicas.read(); + let replicas = one_kv_nodes_routes.node_replicas.read(); let active_replica_count = replicas .iter() .filter(|(_, kv_info)| !kv_info.tomb_tag.is_tomb()) diff --git a/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs b/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs index c74b64a..aba139c 100644 --- a/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs +++ b/fluxon_rs/fluxon_kv/src/kvcore_test_lib.rs @@ -10,7 +10,7 @@ use crate::client_kv_api::ClientKvApiView; use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::config::{ ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, MasterConfig, MonitoringConfig, - ProtocolConfig, ProtocolType, TestSpecConfig, TransferEngineType, + ProtocolConfig, ProtocolType, ReplicaTaskPlacementConfig, TestSpecConfig, TransferEngineType, }; use crate::{ConfigArg, Framework, run_client, run_master}; @@ -93,6 +93,9 @@ fn new_master_config_with_cluster( log_dir, pprof_duration_seconds: None, master_ui: None, + replica_task_placement: ReplicaTaskPlacementConfig::default(), + replica_cache_capacity_ratio: crate::config::DEFAULT_REPLICA_CACHE_CAPACITY_RATIO, + replica_writeback_tier1_capacity_ratio: None, test_spec_config: TestSpecConfig::default(), }; println!("fluxonkv core created master config for test: {:?}", conf); @@ -135,6 +138,7 @@ fn new_client_config_with_cluster_and_dram( rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: vec![etcd], @@ -148,6 +152,7 @@ fn new_client_config_with_cluster_and_dram( large_file_paths: crate::config::LargeFilePaths { paths: vec![format!("{}/large/{}", base, instance_key)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), }; println!("fluxonkv core created client config for test: {:?}", conf); @@ -221,6 +226,35 @@ pub async fn start_master_and_client_with_client_dram( (master_fw, client_fw) } +/// Start a test owner after applying one focused configuration mutation. This +/// keeps integration tests on the same canonical startup path while allowing +/// storage-tier tests to enable their local SSD root explicitly. +pub async fn start_master_and_client_with_client_config( + master_key: &str, + client_key: &str, + configure: impl FnOnce(&mut ClientConfig), +) -> (Arc, Arc) { + let cluster_name = test_cluster_name(master_key); + clean_etcd_members(&cluster_name).await; + + let (master_fw, _) = run_master(ConfigArg::Config(new_master_config_with_cluster( + master_key, + None, + &cluster_name, + ))) + .await + .expect("start master"); + let mut client_config = + new_client_config_with_cluster_and_dram(client_key, &cluster_name, 1024 * 1024 * 160); + configure(&mut client_config); + let (client_fw, _) = run_client(ConfigArg::Config(client_config)) + .await + .expect("start client"); + + sleep(Duration::from_secs(3)).await; + (master_fw, client_fw) +} + pub async fn stop_master_and_client(master_fw: Arc, client_fw: Arc) { let _ = client_fw.shutdown().await; let _ = master_fw.shutdown().await; diff --git a/fluxon_rs/fluxon_kv/src/lib.rs b/fluxon_rs/fluxon_kv/src/lib.rs index edaa386..e09061f 100644 --- a/fluxon_rs/fluxon_kv/src/lib.rs +++ b/fluxon_rs/fluxon_kv/src/lib.rs @@ -7,6 +7,7 @@ pub mod external_client_api; pub mod panel_proxy; // #[cfg(test)] pub mod key_prefix; +pub mod kv_ssd_storage; #[cfg(feature = "test_bins")] pub mod kv_test; pub mod kvlease; @@ -30,6 +31,39 @@ pub use crate::client_seg_pool::SharedJsonMeta; pub use crate::cluster_manager::{ClusterEvent, ClusterMember}; pub use fluxon_observability::types::FsMountKind; +pub const OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES: u64 = 512 * 1024 * 1024; +pub(crate) const OWNER_LOCAL_RESERVE_MIN_SLOT_SIZE_BYTES: u64 = 4 * 1024; + +pub(crate) fn owner_local_reserve_slot_size_bytes(value_len: u64) -> Option { + let normalized = value_len.max(OWNER_LOCAL_RESERVE_MIN_SLOT_SIZE_BYTES); + let alignment_mask = OWNER_LOCAL_RESERVE_MIN_SLOT_SIZE_BYTES - 1; + let slot_size = normalized + .checked_add(alignment_mask) + .map(|rounded| rounded & !alignment_mask)?; + (slot_size <= OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES).then_some(slot_size) +} + +pub(crate) fn owner_local_reserve_slots_per_grant(slot_size: u64) -> Option { + if slot_size == 0 || slot_size > OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES { + return None; + } + u32::try_from(OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES / slot_size).ok() +} + +pub(crate) fn owner_local_reserve_expected_grant_count( + value_len: u64, + payload_capacity_bytes: u64, +) -> Option { + if value_len == 0 || payload_capacity_bytes == 0 { + return None; + } + let slot_size = owner_local_reserve_slot_size_bytes(value_len)?; + let slots_per_grant = u64::from(owner_local_reserve_slots_per_grant(slot_size)?); + let value_count = + payload_capacity_bytes / value_len + u64::from(payload_capacity_bytes % value_len != 0); + Some(value_count / slots_per_grant + u64::from(value_count % slots_per_grant != 0)) +} + pub type MembershipEventReceiver = limit_thirdparty::tokio::sync::abroadcast::Receiver; @@ -53,6 +87,7 @@ use crate::cluster_manager::ClusterManagerViewTrait; use crate::external_client_api::ExternalClientApiAccessTrait; use crate::external_client_api::ExternalClientApiView; use crate::external_client_api::ExternalClientApiViewTrait; +use crate::kv_ssd_storage::{KvSsdStorageInit, KvSsdStorageRootLimit}; use crate::master_kv_router::MasterKvRouterAccessTrait; use crate::master_kv_router::MasterKvRouterView; use crate::master_kv_router::MasterKvRouterViewTrait; @@ -254,6 +289,16 @@ pub trait KvClientTrait { /// Existence check by role async fn kv_is_exist(&self, key: &str) -> KvResult; + /// Batch existence check by role. + /// + /// `allow_local_snapshot` only enables positive local-cache hits; misses + /// still fall through to the authority path. + async fn kv_batch_is_exist( + &self, + keys: Vec, + allow_local_snapshot: bool, + ) -> KvResult>; + /// Count keys by prefix via master-side radix index. async fn kv_count_prefix(&self, prefix: &str) -> KvResult; @@ -403,6 +448,26 @@ impl KvClientTrait for Framework { } } + async fn kv_batch_is_exist( + &self, + keys: Vec, + allow_local_snapshot: bool, + ) -> KvResult> { + if self.is_external_mode() { + self.external_client_api_view() + .external_client_api() + .inner() + .batch_is_exist(keys, allow_local_snapshot) + .await + } else { + self.client_kv_api_view() + .client_kv_api() + .inner() + .batch_is_exist(keys, allow_local_snapshot) + .await + } + } + async fn kv_count_prefix(&self, prefix: &str) -> KvResult { // Delegate to key_prefix helper that talks to master. crate::key_prefix::count_prefix_for_framework(self, prefix).await @@ -772,6 +837,7 @@ fn build_side_transfer_worker_config( test_spec_config.side_transfer_role = Some(SideTransferRole::Worker); test_spec_config.transport_mode = None; test_spec_config.rdma_device_names = None; + test_spec_config.owner_local_reserve_expected_capacity = None; Ok(ClientConfig { cluster_name: owner_config.cluster_name.clone(), @@ -786,6 +852,7 @@ fn build_side_transfer_worker_config( rdma_device_names: None, }, pprof_duration_seconds: owner_config.pprof_duration_seconds, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: Vec::new(), @@ -797,6 +864,7 @@ fn build_side_transfer_worker_config( }, share_mem_path: owner_config.share_mem_path.clone(), large_file_paths: owner_config.large_file_paths.clone(), + ssd_storage: None, test_spec_config, }) } @@ -836,11 +904,15 @@ fn build_side_transfer_worker_config_yaml( protocol: Some(side_config.protocol), contribute_to_cluster_pool_size: None, pprof_duration_seconds: side_config.pprof_duration_seconds, + replica_writeback_hot_capacity_ratio: None, fluxonkv_spec: crate::config::FluxonKvSpecYaml { etcd_addresses: None, cluster_name: side_config.cluster_name, share_mem_path: side_config.share_mem_path, large_file_paths: None, + large_limit_size: None, + ssd_write_rate_limit_bytes_per_sec: None, + ssd_write_burst_bytes: None, p2p_listen_port: side_config.fluxonkv_spec.p2p_listen_port, redis_compat: None, sub_cluster: None, @@ -1489,6 +1561,9 @@ async fn run_master_impl( master_seg_manager_arg: MasterSegManagerNewArg, master_kv_router_arg: MasterKvRouterNewArg { test_spec_config: config.test_spec_config.clone(), + replica_task_placement: config.replica_task_placement.clone(), + replica_cache_capacity_ratio: config.replica_cache_capacity_ratio, + replica_writeback_tier1_capacity_ratio: config.replica_writeback_tier1_capacity_ratio, }, metric_reporter_arg: MetricReporterNewArg { test_spec_config: config.test_spec_config.clone(), @@ -1923,6 +1998,12 @@ async fn run_client_impl( // Owner nodes are the global relay set. Keeping a single metadata marker converges // route selection, observability proxy eligibility, and ops topology rendering. metadata.insert("p2p_relay".to_string(), "true".to_string()); + if config.ssd_storage.is_some() { + metadata.insert( + crate::cluster_manager::META_KEY_KV_SSD_STORAGE.to_string(), + "true".to_string(), + ); + } } merge_startup_member_metadata(&mut metadata, startup_member_metadata)?; @@ -1957,6 +2038,24 @@ async fn run_client_impl( )); info!("Initializing client framework..."); + let ssd_storage_init = config + .ssd_storage + .as_ref() + .map(|ssd| { + config + .large_file_paths + .kv_ssd_storage_root(&config.cluster_name, &config.instance_key) + .map(|root_dir| KvSsdStorageInit { + roots: vec![KvSsdStorageRootLimit { + root_dir, + limit_bytes: ssd.limit_bytes, + }], + write_rate_limit_bytes_per_sec: ssd.write_rate_limit_bytes_per_sec, + write_burst_bytes: ssd.write_burst_bytes, + }) + }) + .transpose()?; + if is_external && !is_side_transfer_worker { let init_args = InitArgsExternal { cluster_manager_arg: ClusterManagerNewArg { @@ -1988,9 +2087,32 @@ async fn run_client_impl( metric_reporter_arg: MetricReporterNewArg { test_spec_config: config.test_spec_config.clone(), }, - external_client_api_arg: ExternalClientApiNewArg { + client_seg_pool_arg: ClientSegPoolNewArg { + contribute_size: config.contribute_to_cluster_pool_size.clone(), share_mem_path: config.share_mem_path.clone(), large_file_paths: config.large_file_paths.clone(), + cluster_name: config.cluster_name.clone(), + etcd_addresses: config.etcd_addresses_raw.clone(), + attach_existing_meta: None, + side_transfer_worker: false, + require_transfer_rpc_fast_path_ready_timeout: None, + }, + client_transfer_engine_arg: ClientTransferEngineNewArg { + metadata_uri: config.fluxonkv_spec.etcd_addresses[0].clone(), + instance_name: config.instance_key.clone(), + enable_transfer_rpc_fast_path: config.fluxonkv_spec.enable_transfer_rpc_fast_path, + rpc_port: 12345, + protocol_type: config.protocol.protocol_type.clone(), + rdma_device_names: transfer_engine_rdma_device_names_from_config( + &config.protocol, + Some(&config.test_spec_config), + ), + backend_activation_mode: transfer_backend_activation_mode, + transfer_engine: config.fluxonkv_spec.transfer_engine.clone(), + }, + external_client_api_arg: ExternalClientApiNewArg { + shared_memory_path: config.share_mem_path.clone(), + shared_file_path: config.share_mem_path.clone(), expected_cluster_name: config.cluster_name.clone(), expected_protocol_version: build_version.clone(), enable_side_transfer: config.test_spec_config.enable_side_transfer, @@ -2036,6 +2158,12 @@ async fn run_client_impl( }, client_kv_api_arg: ClientKvApiNewArg { test_spec_config: config.test_spec_config.clone(), + owner_hot_cache_capacity_bytes: config.replica_writeback_hot_capacity_ratio.map( + |ratio| { + (config.contribute_to_cluster_pool_size.dram as f64 * ratio).floor() as u64 + }, + ), + ssd_storage: ssd_storage_init, }, client_seg_pool_arg: ClientSegPoolNewArg { contribute_size: config.contribute_to_cluster_pool_size.clone(), @@ -2455,6 +2583,7 @@ mod tests { rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: vec!["http://127.0.0.1:2379".to_string()], @@ -2468,6 +2597,7 @@ mod tests { large_file_paths: crate::config::LargeFilePaths { paths: vec!["/tmp/fluxon_side_transfer_test_large".to_string()], }, + ssd_storage: None, test_spec_config: TestSpecConfig { enable_side_transfer: true, side_transfer_worker_count: 4, @@ -2736,8 +2866,8 @@ mod tests { large_file_paths: crate::config::LargeFilePaths { paths: vec![owner_large_root.to_string_lossy().into_owned()], }, - protocol_version: - fluxon_util::git_version_build_record::get_current_git_commitid().unwrap(), + protocol_version: fluxon_util::git_version_build_record::get_current_git_commitid() + .unwrap(), write_ts: Some(chrono::Utc::now().timestamp_micros()), }; let shared_meta_json = serde_json::to_string(&shared_meta).unwrap(); @@ -2762,6 +2892,7 @@ mod tests { rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: Vec::new(), @@ -2773,6 +2904,7 @@ mod tests { }, share_mem_path: share_mem_root.to_string_lossy().into_owned(), large_file_paths: crate::config::LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), }; diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/count_prefix_index.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/count_prefix_index.rs index 4c6e34a..ab7cdad 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/count_prefix_index.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/count_prefix_index.rs @@ -1,5 +1,7 @@ use std::collections::HashMap; +type KeyVersion = (u64, u32); + /// Radix tree node used for prefix counting. /// /// Each node tracks the total number of keys in its subtree via `count`. @@ -25,15 +27,33 @@ impl Node { #[derive(Default)] pub struct PrefixRadixTree { root: Node, + /// Current version for every logical key represented in the radix counts. + /// + /// Route publication and final-route reclaim are intentionally processed by separate async + /// actors. Their events can therefore be duplicated or observed out of order. Keeping the + /// version here makes both operations idempotent and prevents an old remove from deleting a + /// newer incarnation of the same key. + versions: HashMap, } impl PrefixRadixTree { pub fn new() -> Self { - Self { root: Node::new() } + Self { + root: Node::new(), + versions: HashMap::new(), + } } - /// Insert a new key. Must only be called once per logical key. - pub fn insert(&mut self, key: &str) { + /// Insert or update a key version. Returns true only when a new logical key was counted. + pub fn insert(&mut self, key: &str, version: KeyVersion) -> bool { + if let Some(current) = self.versions.get_mut(key) { + if *current != version { + *current = version; + } + return false; + } + + self.versions.insert(key.to_string(), version); let bytes = key.as_bytes(); let mut node = &mut self.root; node.count = node @@ -48,25 +68,42 @@ impl PrefixRadixTree { .checked_add(1) .expect("PrefixRadixTree count overflow on insert (child)"); } + true } - /// Remove an existing key. Must only be called for keys that were inserted. - pub fn remove(&mut self, key: &str) { + /// Remove exactly one key version. Missing, duplicate, or stale removes are harmless. + pub fn remove(&mut self, key: &str, version: KeyVersion) -> bool { + if !self + .versions + .get(key) + .is_some_and(|current| *current == version) + { + return false; + } + + // Validate the complete path before mutating any count. This keeps the operation + // non-panicking even if an index built by older code was inconsistent. + let mut current = &self.root; + for &byte in key.as_bytes() { + let Some(child) = current.children.get(&byte) else { + self.versions.remove(key); + return false; + }; + current = child; + } + self.versions.remove(key); + fn remove_inner(node: &mut Node, bytes: &[u8], idx: usize) -> bool { - node.count = node - .count - .checked_sub(1) - .expect("PrefixRadixTree underflow on remove"); + node.count = node.count.saturating_sub(1); if idx == bytes.len() { return node.count == 0; } let b = bytes[idx]; - let child = node - .children - .get_mut(&b) - .expect("PrefixRadixTree remove: missing child for existing key"); + let Some(child) = node.children.get_mut(&b) else { + return node.count == 0; + }; let should_prune = remove_inner(child, bytes, idx + 1); if should_prune { node.children.remove(&b); @@ -75,6 +112,7 @@ impl PrefixRadixTree { } remove_inner(&mut self.root, key.as_bytes(), 0); + true } /// Count keys whose name starts with the given prefix. @@ -89,3 +127,40 @@ impl PrefixRadixTree { node.count } } + +#[cfg(test)] +mod tests { + use super::PrefixRadixTree; + + #[test] + fn duplicate_insert_and_remove_are_idempotent() { + let mut tree = PrefixRadixTree::new(); + assert!(tree.insert("prefix/key", (1, 0))); + assert!(!tree.insert("prefix/key", (1, 0))); + assert_eq!(tree.count_prefix("prefix/"), 1); + + assert!(tree.remove("prefix/key", (1, 0))); + assert!(!tree.remove("prefix/key", (1, 0))); + assert_eq!(tree.count_prefix("prefix/"), 0); + } + + #[test] + fn stale_remove_cannot_delete_newer_key_version() { + let mut tree = PrefixRadixTree::new(); + assert!(tree.insert("prefix/key", (1, 0))); + assert!(!tree.insert("prefix/key", (2, 0))); + assert!(!tree.remove("prefix/key", (1, 0))); + assert_eq!(tree.count_prefix("prefix/"), 1); + + assert!(tree.remove("prefix/key", (2, 0))); + assert_eq!(tree.count_prefix("prefix/"), 0); + } + + #[test] + fn remove_before_delayed_insert_is_a_noop() { + let mut tree = PrefixRadixTree::new(); + assert!(!tree.remove("prefix/key", (3, 0))); + assert!(tree.insert("prefix/key", (3, 0))); + assert_eq!(tree.count_prefix("prefix/"), 1); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs index 12a55ee..2a2f758 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/delete.rs @@ -1,8 +1,7 @@ use super::{ MasterKvRouterView, msg_pack::{ - BatchDeleteAckReq, BatchDeleteAckResp, BatchDeleteClientKvMetaCacheReq, DeleteAckReq, - DeleteAckResp, DeleteClientKvMetaCacheItem, DeleteReq, DeleteResp, + BatchDeleteAckReq, BatchDeleteAckResp, DeleteAckReq, DeleteAckResp, DeleteReq, DeleteResp, }, }; use crate::master_kv_router::OneKvNodesRoutes; @@ -11,12 +10,18 @@ use crate::memholder::{ EnsureMemholderMgmtDeleteActorOwned, MasterOwnerMemMgr, MemholderManagerTrait, }; use crate::{ - cluster_manager::NodeID, - p2p::msg_pack::{MsgPack, RPCCaller}, + p2p::msg_pack::MsgPack, rpcresp_kvresult_convert::msg_and_error::{self, kv}, }; use limit_thirdparty::tokio; -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; + +#[cfg(any(test, feature = "test_bins"))] +use super::msg_pack::{BatchDeleteClientKvMetaCacheReq, DeleteClientKvMetaCacheItem}; +#[cfg(any(test, feature = "test_bins"))] +use crate::{cluster_manager::NodeID, p2p::msg_pack::RPCCaller}; +#[cfg(any(test, feature = "test_bins"))] +use std::time::Duration; /// Remove a key from master indices and trigger client cache invalidation broadcast. /// @@ -26,7 +31,7 @@ use std::{sync::Arc, time::Duration}; /// /// It removes the key from `kv_routes`, then asynchronously: /// - emits a `DeleteKeyInfo` to the shared delete broadcast actor for clients -/// - removes the key from every node's local `node_kv_cache_controller` +/// - point-removes the exact version from ring-B and tier1 metadata caches pub fn do_delete_one_kv_all_replicas( view: &MasterKvRouterView, key: String, @@ -45,7 +50,7 @@ pub fn do_delete_one_kv_all_replicas( if view.master_kv_router().prefix_index_enabled() { let inner = view.master_kv_router().inner(); let mut tree = inner.prefix_index.write().await; - tree.remove(&key_clone); + tree.remove(&key_clone, deleted_put_id); } if let Err(err) = view @@ -62,19 +67,9 @@ pub fn do_delete_one_kv_all_replicas( tracing::warn!("Failed to send delete broadcast: {}", err); } - // Remove from all node caches that hold replicas of this key - let nodes_replicas = kv_route_info.nodes_replicas.read(); - for (node_id, _kv_info) in nodes_replicas.iter() { - if let Some(cache) = view.master_kv_router().get_node_cache_controller(node_id) - { - let _ = cache.remove(&key_clone); - tracing::debug!( - "Removed key {} from node cache controller: {}", - key_clone, - node_id - ); - } - } + view.master_kv_router() + .remove_route_cache_entries_exact(&key_clone, &kv_route_info) + .await; } }); @@ -84,6 +79,8 @@ pub fn do_delete_one_kv_all_replicas( } } +/// Test-only direct route removal. Production eviction uses the fenced owner reclaim protocol. +#[cfg(any(test, feature = "test_bins"))] pub fn evict_one_kv_replica_for_node( view: &MasterKvRouterView, key: String, @@ -116,8 +113,8 @@ pub fn evict_one_kv_replica_for_node( } let removed_replica = { - let mut nodes_replicas = route.nodes_replicas.write(); - nodes_replicas.remove(&node_id).is_some() + let mut node_replicas = route.node_replicas.write(); + node_replicas.remove(&node_id).is_some() }; if !removed_replica { tracing::debug!( @@ -130,13 +127,17 @@ pub fn evict_one_kv_replica_for_node( return Ok(()); } - let last_replica_gone = route.nodes_replicas.read().is_empty(); + let last_replica_gone = route.node_replicas.read().is_empty(); if last_replica_gone { let removed = view .master_kv_router() .inner() .kv_routes - .remove_if(&key, |_, current| current.put_id == put_id) + .remove_if(&key, |_, current| { + Arc::ptr_eq(current, &route) + && current.put_id == put_id + && current.node_replicas.read().is_empty() + }) .is_some(); if removed && view.master_kv_router().prefix_index_enabled() { let view_task = view.clone(); @@ -144,7 +145,7 @@ pub fn evict_one_kv_replica_for_node( let _ = view.spawn("local_evict_remove_prefix_index", async move { let inner = view_task.master_kv_router().inner(); let mut tree = inner.prefix_index.write().await; - tree.remove(&key_for_prefix); + tree.remove(&key_for_prefix, put_id); }); } } diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs index 8c17155..5e08ec7 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/get.rs @@ -1,73 +1,1131 @@ use super::{ - InflightGetInfo, KvRouteInfo, MasterKvRouterView, NodeValueReplicaDesc, OwnerHoldingGetInfo, + CommittedSlotReplica, CompletedGetInfo, InflightGetInfo, InflightGetTarget, KvMemoryReplica, + KvNodeReplicas, MasterKeyActivityCompletionGuard, MasterKvRouterView, OwnerHoldingGetInfo, + ReservedCapacityReason, msg_pack::{ - GetAllocationMode, GetDoneReq, GetDoneResp, GetMetaReq, GetMetaResp, GetRevokeReq, - GetRevokeResp, GetStartReq, GetStartResp, + BatchGetBindItemReq, BatchGetBindReq, BatchGetBindResp, BatchGetDoneItemResp, + BatchGetDoneReq, BatchGetDoneResp, BatchGetPlanItemResp, BatchGetPlanReq, BatchGetPlanResp, + BatchGetRevokeItemResp, BatchGetRevokeReq, BatchGetRevokeResp, BatchGetStartItemResp, + BatchGetStartReq, BatchGetStartResp, BatchIsExistReq, BatchIsExistResp, GetAllocationMode, + GetBindTarget, GetDoneReq, GetDoneResp, GetExternalSinkTarget, GetMetaReq, GetMetaResp, + GetPreparedLocalReserveTarget, GetRevokeReq, GetRevokeResp, GetSourceKind, GetStartReq, + GetStartResp, MemHolderKeepAliveReq, MemHolderKeepAliveResp, MemHolderReleaseReq, + MemHolderReleaseResp, SsdStageBeginReq, SsdStageBeginResp, SsdStageDoneReq, + SsdStageDoneResp, }, + node_generation_is_current_live, publish_route_replica_tomb_fenced, + route_maintenance::{RoutePublishEvent, apply_post_route_maintenance_batch}, }; +use crate::config::SsdReadSourcePolicy; use crate::master_kv_router::OneKvNodesRoutes; use crate::master_kv_router::put::PutIDForAKey; use crate::memholder::MemholderManagerTrait; use crate::{ - cluster_manager::NodeID, master_seg_manager::one_seg_allocator::Allocation, - p2p::msg_pack::MsgPack, rpcresp_kvresult_convert::msg_and_error, + cluster_manager::{ClusterManagerAccessTrait, NodeID, NodeRole}, + master_seg_manager::{MasterSegManagerAccessTrait, NodeTombTag, one_seg_allocator::Allocation}, + p2p::msg_pack::MsgPack, + rpcresp_kvresult_convert::msg_and_error::{self, kv}, }; +use dashmap::DashMap; +use fluxon_commu::share_group_owner_ref_from_metadata; +use limit_thirdparty::tokio; use rand::Rng; use rand::seq::SliceRandom; use std::collections::HashSet; use std::{ collections::HashMap, sync::{Arc, atomic::Ordering}, + time::Instant, }; -fn update_moka_for_node( - view: MasterKvRouterView, - node_id: String, - key: String, - weight: u32, - put_id: PutIDForAKey, - new_inserted: bool, -) { +fn touch_moka_for_node(view: MasterKvRouterView, node_id: String, key: String) { if !view.master_kv_router().replica_cache_enabled() { return; } let view_task = view.clone(); - let _ = view.spawn("update_moka_for_node", async move { + view.spawn("touch_moka_for_node", async move { + let owner_cache_lock = view_task + .master_kv_router() + .inner() + .owner_cache_operation_locks + .get_lock(node_id.clone()); + let _owner_cache_guard = owner_cache_lock.lock().await; if let Some(cache) = view_task .master_kv_router() .get_node_cache_controller(&node_id) { - if new_inserted { - cache.insert( - key.clone(), - NodeValueReplicaDesc { - weight_bytes: weight, - put_id, - }, - ); - tracing::debug!( - "Inserted key: {:?} into node cache: {}, weight={}", - key, - node_id, - weight - ); - } else { - let _ = cache.get(&key); - tracing::debug!( - "Touched key: {:?} on node cache: {} (TTL refresh)", - key, - node_id - ); + // A get is a hit signal for ring B when the source is an + // unindexed Allocation. Owner-indexed routes are intentionally + // absent from this cache. + let _ = cache.get(&key); + if let Some(tier1_cache) = view_task + .master_kv_router() + .get_node_writeback_tier1_controller(&node_id) + { + // Tier1 has independent admission and replacement state; a + // hit only touches an already-admitted entry. + let _ = tier1_cache.get(&key); } + tracing::debug!( + "Touched key: {:?} on node cache: {} (TTL refresh)", + key, + node_id + ); } else { tracing::warn!( - "No cache controller found for node: {} when updating moka", + "No cache controller found for node: {} when touching moka", node_id ); } }); } +fn one_kv_routes_has_live_replica(one_kv_nodes_routes: &OneKvNodesRoutes) -> bool { + one_kv_nodes_routes + .node_replicas + .read() + .values() + .any(KvNodeReplicas::has_live_backing) +} + +fn validate_prepared_local_reserve_target( + view: &MasterKvRouterView, + req_node_id: &NodeID, + target: &GetPreparedLocalReserveTarget, + value_len: u64, +) -> Result<(CommittedSlotReplica, NodeTombTag), msg_and_error::KvError> { + let invalid = |detail: String| { + msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { detail }) + }; + if target.slot_size == 0 { + return Err(invalid( + "prepared local-reserve Get target has zero slot_size".to_string(), + )); + } + if value_len > target.slot_size { + return Err(invalid(format!( + "prepared local-reserve Get target is too small: value_len={} slot_size={}", + value_len, target.slot_size + ))); + } + let Some(grant) = view + .master_kv_router() + .inner() + .local_reserve_grants + .get(&target.grant_id) + else { + return Err(invalid(format!( + "prepared local-reserve Get target references unknown grant_id={}", + target.grant_id + ))); + }; + if grant.owner_node_id != *req_node_id { + return Err(invalid(format!( + "prepared local-reserve Get target owner mismatch: grant_id={} owner={} requester={}", + target.grant_id, grant.owner_node_id, req_node_id + ))); + } + let tomb_tag = grant.tomb_tag.clone(); + if !node_generation_is_current_live(view, req_node_id, &tomb_tag) { + return Err(invalid(format!( + "prepared local-reserve Get target belongs to a departed owner generation: grant_id={} requester={}", + target.grant_id, req_node_id + ))); + } + let grant_base_addr = grant.allocation.base_addr(); + let grant_addr = grant_base_addr + .checked_add(grant.allocation.addr()) + .ok_or_else(|| invalid("local-reserve grant address overflow".to_string()))?; + let slot_offset = target + .slot_size + .checked_mul(u64::from(target.slot_index)) + .ok_or_else(|| invalid("prepared local-reserve Get slot offset overflow".to_string()))?; + let slot_end = slot_offset + .checked_add(target.slot_size) + .ok_or_else(|| invalid("prepared local-reserve Get slot end overflow".to_string()))?; + if slot_end > grant.allocation.capcity() { + return Err(invalid(format!( + "prepared local-reserve Get target is outside grant: grant_id={} slot_index={} slot_size={} grant_len={}", + target.grant_id, + target.slot_index, + target.slot_size, + grant.allocation.capcity() + ))); + } + let expected_addr = grant_addr + .checked_add(slot_offset) + .ok_or_else(|| invalid("prepared local-reserve Get target address overflow".to_string()))?; + if target.base_addr != grant_base_addr || target.addr != expected_addr { + return Err(invalid(format!( + "prepared local-reserve Get target geometry mismatch: grant_id={} expected_base={:#x} got_base={:#x} expected_addr={:#x} got_addr={:#x}", + target.grant_id, grant_base_addr, target.base_addr, expected_addr, target.addr + ))); + } + Ok(( + CommittedSlotReplica { + owner_node_id: req_node_id.clone(), + grant_id: target.grant_id, + slot_index: target.slot_index, + slot_size: target.slot_size, + addr: target.addr, + len: value_len, + base_addr: target.base_addr, + }, + tomb_tag, + )) +} + +fn external_sink_local_owner_id( + view: &MasterKvRouterView, + req_node_id: &NodeID, + requester_node_start_time: i64, +) -> Option { + let requester = view + .cluster_manager() + .get_member_info_cached(req_node_id.as_ref())?; + if requester.node_start_time != requester_node_start_time + || requester.node_role() != NodeRole::External + { + return None; + } + let owner_ref = share_group_owner_ref_from_metadata(&requester.metadata)?; + let owner = view + .cluster_manager() + .get_member_info_cached(&owner_ref.owner_id)?; + (owner.node_start_time == owner_ref.owner_start_time && owner.node_role() == NodeRole::Client) + .then_some(owner_ref.owner_id) +} + +fn external_sink_requester_generation_is_current( + view: &MasterKvRouterView, + req_node_id: &NodeID, + requester_node_start_time: i64, +) -> bool { + external_sink_local_owner_id(view, req_node_id, requester_node_start_time).is_some() +} + +fn validate_external_sink_target( + view: &MasterKvRouterView, + req_node_id: &NodeID, + target: &GetExternalSinkTarget, + value_len: u64, +) -> Result<(), msg_and_error::KvError> { + let invalid = |detail: String| { + msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { detail }) + }; + if target.addr == 0 || target.capacity == 0 || target.registration_id == 0 { + return Err(invalid(format!( + "external Get sink requires non-zero addr/capacity/registration_id: addr={:#x} capacity={} registration_id={}", + target.addr, target.capacity, target.registration_id + ))); + } + if target.addr.checked_add(target.capacity).is_none() { + return Err(invalid(format!( + "external Get sink range overflows: addr={:#x} capacity={}", + target.addr, target.capacity + ))); + } + if value_len > target.capacity { + return Err(invalid(format!( + "external Get sink is too small: value_len={} capacity={}", + value_len, target.capacity + ))); + } + if !external_sink_requester_generation_is_current( + view, + req_node_id, + target.requester_node_start_time, + ) { + return Err(invalid(format!( + "external Get sink requester generation is not current: requester={} start_time={}", + req_node_id, target.requester_node_start_time + ))); + } + Ok(()) +} + +fn get_plan_item_error(err: &msg_and_error::KvError) -> BatchGetPlanItemResp { + let response: GetStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(err); + BatchGetPlanItemResp { + error_code: response.error_code, + error_json: response.error_json, + ..Default::default() + } +} + +fn get_bind_item_error(get_id: u64, err: &msg_and_error::KvError) -> BatchGetStartItemResp { + let response: GetStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(err); + BatchGetStartItemResp { + get_id, + error_code: response.error_code, + error_json: response.error_json, + ..Default::default() + } +} + +#[derive(Clone)] +struct PlannedGetSourceSnapshot { + node_id: NodeID, + tomb_tag: crate::master_seg_manager::NodeTombTag, + len: u64, + addr: u64, + base_addr: u64, + source_kind: GetSourceKind, +} + +fn planned_get_source_rank( + policy: SsdReadSourcePolicy, + source_kind: GetSourceKind, + requester_local: bool, +) -> Option { + match policy { + SsdReadSourcePolicy::LegacyRemoteFirst => match (source_kind, requester_local) { + (GetSourceKind::Memory, false) => Some(0), + (GetSourceKind::Memory, true) => Some(1), + (GetSourceKind::Ssd, false) => Some(2), + (GetSourceKind::Ssd, true) => Some(3), + }, + SsdReadSourcePolicy::LocalSsdOnlyFirst => match (source_kind, requester_local) { + (GetSourceKind::Memory, true) => Some(0), + (GetSourceKind::Ssd, true) => Some(1), + (GetSourceKind::Memory, false) => Some(2), + (GetSourceKind::Ssd, false) => None, + }, + } +} + +fn snapshot_live_get_sources(route: &OneKvNodesRoutes) -> Vec { + route + .node_replicas + .read() + .iter() + .filter_map(|(node_id, replicas)| { + if replicas.tomb_tag.is_tomb() { + return None; + } + if let Some(memory) = replicas.memory.as_ref() { + return Some(PlannedGetSourceSnapshot { + node_id: node_id.clone(), + tomb_tag: replicas.tomb_tag.clone(), + len: memory.backing.len(), + addr: memory.backing.abs_addr(), + base_addr: memory.backing.base_addr(), + source_kind: GetSourceKind::Memory, + }); + } + replicas.ssd.as_ref().map(|ssd| PlannedGetSourceSnapshot { + node_id: node_id.clone(), + tomb_tag: replicas.tomb_tag.clone(), + len: ssd.len, + addr: 0, + base_addr: 0, + source_kind: GetSourceKind::Ssd, + }) + }) + .collect() +} + +fn planned_get_source_is_current( + planned: &super::PlannedGetInfo, + route: &Arc, +) -> bool { + if planned.src_tomb_tag.is_tomb() || route.put_id != planned.put_id { + return false; + } + route + .node_replicas + .read() + .get(&planned.src_node_id) + .is_some_and(|replicas| { + !replicas.tomb_tag.is_tomb() + && replicas.tomb_tag.same_generation(&planned.src_tomb_tag) + && match planned.source_kind { + GetSourceKind::Memory => replicas.memory.as_ref().is_some_and(|memory| { + memory.backing.abs_addr() == planned.src_addr + && memory.backing.base_addr() == planned.src_base_addr + && memory.backing.len() == planned.len + }), + GetSourceKind::Ssd => replicas + .ssd + .as_ref() + .is_some_and(|ssd| ssd.len == planned.len), + } + }) +} + +fn allocate_ssd_source_stage( + view: &MasterKvRouterView, + source_node_id: &NodeID, + len: u64, +) -> Result, msg_and_error::KvError> { + let allocators = view + .master_seg_manager() + .get_node_allocators(source_node_id); + if allocators.is_empty() { + return Err(msg_and_error::KvError::Unreachable( + msg_and_error::UnreachableError::OwnerNoSeg { + detail: format!("SSD source owner has no registered segment: {source_node_id}"), + }, + )); + } + let allocator = allocators + .choose(&mut rand::thread_rng()) + .expect("non-empty SSD source allocators"); + for _ in 0..3 { + if let Ok(allocation) = allocator.allocate(len) { + return Ok(Arc::new(allocation)); + } + } + let capacity = allocator.node_pool_capacity_snapshot(); + Err(msg_and_error::KvError::Api( + msg_and_error::ApiError::NoSpace { + node: source_node_id.to_string(), + segment: allocator.seg_device_id.clone(), + total_capacity: capacity.active_capacity_bytes, + free_capacity: capacity.available_capacity_bytes, + }, + )) +} + +#[cfg(test)] +mod planned_get_tests { + use super::{ + planned_get_source_is_current, planned_get_source_rank, snapshot_live_get_sources, + }; + use crate::cluster_manager::NodeID; + use crate::config::SsdReadSourcePolicy; + use crate::master_kv_router::msg_pack::GetSourceKind; + use crate::master_kv_router::{ + CommittedSlotReplica, KvMemoryReplica, KvNodeReplicas, KvReplicaBacking, KvSsdReplica, + OneKvNodesRoutes, PlannedGetInfo, + }; + use crate::master_seg_manager::NodeTombTag; + use parking_lot::RwLock; + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::AtomicU32; + + fn planned_route() -> (PlannedGetInfo, Arc, NodeTombTag) { + let source: NodeID = "source".to_string().into(); + let source_tag = NodeTombTag::new(); + let route = Arc::new(OneKvNodesRoutes { + put_id: (7, 3), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + source.clone(), + KvNodeReplicas::memory( + source_tag.clone(), + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: source.clone(), + grant_id: 11, + slot_index: 2, + slot_size: 8192, + addr: 0x3000, + len: 4096, + base_addr: 0x1000, + }), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation: None, + }, + ), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let planned = PlannedGetInfo { + put_id: route.put_id, + src_node_id: source, + src_tomb_tag: source_tag.clone(), + key: "key".to_string(), + controller_node_id: "external".to_string().into(), + controller_node_start_time: 17, + len: 4096, + src_addr: 0x3000, + src_base_addr: 0x1000, + source_kind: GetSourceKind::Memory, + atomic_group: None, + }; + (planned, route, source_tag) + } + + #[test] + fn bind_revalidation_accepts_only_the_exact_source_generation() { + let (planned, route, source_tag) = planned_route(); + assert!(planned_get_source_is_current(&planned, &route)); + + route + .node_replicas + .write() + .get_mut(&planned.src_node_id) + .unwrap() + .memory + .as_mut() + .unwrap() + .backing = KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: planned.src_node_id.clone(), + grant_id: 11, + slot_index: 2, + slot_size: 8192, + addr: 0x4000, + len: 4096, + base_addr: 0x1000, + }); + assert!(!planned_get_source_is_current(&planned, &route)); + + route + .node_replicas + .write() + .get_mut(&planned.src_node_id) + .unwrap() + .memory + .as_mut() + .unwrap() + .backing = KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: planned.src_node_id.clone(), + grant_id: 11, + slot_index: 2, + slot_size: 8192, + addr: 0x3000, + len: 4096, + base_addr: 0x1000, + }); + source_tag.set_tomb(); + assert!(!planned_get_source_is_current(&planned, &route)); + } + + #[test] + fn bind_revalidation_rejects_a_replacement_route_generation() { + let (planned, route, _) = planned_route(); + let replacement_tag = NodeTombTag::new(); + route + .node_replicas + .write() + .get_mut(&planned.src_node_id) + .unwrap() + .tomb_tag = replacement_tag; + assert!(!planned_get_source_is_current(&planned, &route)); + } + + #[test] + fn ssd_plan_and_bind_revalidate_the_same_owner_backing() { + let (mut planned, route, _) = planned_route(); + planned.source_kind = GetSourceKind::Ssd; + planned.src_addr = 0; + planned.src_base_addr = 0; + let mut replicas = route.node_replicas.write(); + let source = replicas.get_mut(&planned.src_node_id).unwrap(); + source.memory = None; + source.ssd = Some(KvSsdReplica { len: planned.len }); + drop(replicas); + + let sources = snapshot_live_get_sources(&route); + assert_eq!(sources.len(), 1); + assert_eq!(sources[0].source_kind, GetSourceKind::Ssd); + assert_eq!(sources[0].addr, 0); + assert_eq!(sources[0].base_addr, 0); + assert!(planned_get_source_is_current(&planned, &route)); + + route + .node_replicas + .write() + .get_mut(&planned.src_node_id) + .unwrap() + .ssd + .as_mut() + .unwrap() + .len += 1; + assert!(!planned_get_source_is_current(&planned, &route)); + } + + #[test] + fn metadata_plan_does_not_retain_the_route() { + let (planned, route, _) = planned_route(); + let weak_route = Arc::downgrade(&route); + let sources = snapshot_live_get_sources(&route); + drop(route); + assert!(weak_route.upgrade().is_none()); + assert_eq!(sources.len(), 1); + assert_eq!(planned.key, "key"); + } + + #[test] + fn ssd_read_source_policies_have_distinct_explicit_orders() { + let order = |policy| { + let classes = vec![ + (GetSourceKind::Memory, false, "remote_memory"), + (GetSourceKind::Memory, true, "local_memory"), + (GetSourceKind::Ssd, false, "remote_ssd"), + (GetSourceKind::Ssd, true, "local_ssd"), + ]; + let mut classes = classes + .into_iter() + .filter_map(|(kind, local, label)| { + planned_get_source_rank(policy, kind, local).map(|rank| (rank, label)) + }) + .collect::>(); + classes.sort_by_key(|(rank, _)| *rank); + classes + .into_iter() + .map(|(_, label)| label) + .collect::>() + }; + + assert_eq!( + order(SsdReadSourcePolicy::LegacyRemoteFirst), + ["remote_memory", "local_memory", "remote_ssd", "local_ssd"] + ); + assert_eq!( + order(SsdReadSourcePolicy::LocalSsdOnlyFirst), + ["local_memory", "local_ssd", "remote_memory"] + ); + assert_eq!( + planned_get_source_rank( + SsdReadSourcePolicy::LocalSsdOnlyFirst, + GetSourceKind::Ssd, + false, + ), + None, + "remote SSD must not be a fallback for the local-only policy" + ); + } +} + +async fn handle_get_plan_item( + view: MasterKvRouterView, + key: String, + controller_node_id: NodeID, +) -> BatchGetPlanItemResp { + view.master_kv_router() + .inner() + .planned_get_counters + .plan_items + .fetch_add(1, Ordering::Relaxed); + let Some(controller) = view + .cluster_manager() + .get_member_info_cached(controller_node_id.as_ref()) + else { + return get_plan_item_error(&msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetPlan controller is not a current member: {}", + controller_node_id + ), + }, + )); + }; + if controller.node_role() != NodeRole::External { + return get_plan_item_error(&msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetPlan is supported only for external controllers: {}", + controller_node_id + ), + }, + )); + } + + let Some(route) = view + .master_kv_router() + .inner() + .kv_routes + .get(&key) + .map(|route| route.clone()) + else { + view.master_kv_router() + .inner() + .planned_get_counters + .plan_misses + .fetch_add(1, Ordering::Relaxed); + return get_plan_item_error(&msg_and_error::KvError::Api( + msg_and_error::ApiError::KeyNotFound { key }, + )); + }; + + let local_owner = + external_sink_local_owner_id(&view, &controller_node_id, controller.node_start_time); + let source_policy = view + .master_kv_router() + .inner() + .test_spec_config + .ssd_read_source_policy; + let mut filtered_remote_ssd_items = 0u64; + let mut filtered_remote_ssd_bytes = 0u64; + let mut candidates = snapshot_live_get_sources(&route) + .into_iter() + .filter_map(|source| { + let requester_local = local_owner + .as_deref() + .is_some_and(|owner| source.node_id.as_ref() == owner); + match planned_get_source_rank(source_policy, source.source_kind, requester_local) { + Some(rank) => Some((source, rank)), + None => { + filtered_remote_ssd_items = filtered_remote_ssd_items.saturating_add(1); + filtered_remote_ssd_bytes = + filtered_remote_ssd_bytes.saturating_add(source.len); + None + } + } + }) + .collect::>(); + if filtered_remote_ssd_items != 0 { + view.master_kv_router() + .inner() + .planned_get_counters + .remote_ssd_filtered_items + .fetch_add(filtered_remote_ssd_items, Ordering::Relaxed); + view.master_kv_router() + .inner() + .planned_get_counters + .remote_ssd_filtered_bytes + .fetch_add(filtered_remote_ssd_bytes, Ordering::Relaxed); + } + candidates.shuffle(&mut rand::thread_rng()); + candidates.sort_by_key(|(_, rank)| *rank); + let has_remote_memory_alternative = candidates.iter().any(|(candidate, _)| { + candidate.source_kind == GetSourceKind::Memory + && local_owner + .as_deref() + .is_none_or(|owner| candidate.node_id.as_ref() != owner) + }); + let Some((source, _)) = candidates.into_iter().next() else { + view.master_kv_router() + .inner() + .planned_get_counters + .plan_misses + .fetch_add(1, Ordering::Relaxed); + return get_plan_item_error(&msg_and_error::KvError::Api( + msg_and_error::ApiError::KeyNotFound { key }, + )); + }; + let selected_requester_local_ssd = source.source_kind == GetSourceKind::Ssd + && local_owner + .as_deref() + .is_some_and(|owner| source.node_id.as_ref() == owner); + if selected_requester_local_ssd { + let counters = &view.master_kv_router().inner().ssd_tier_counters; + let (items, bytes) = if has_remote_memory_alternative { + ( + &counters.local_ssd_selected_with_remote_memory_items, + &counters.local_ssd_selected_with_remote_memory_bytes, + ) + } else { + ( + &counters.local_ssd_selected_without_remote_memory_items, + &counters.local_ssd_selected_without_remote_memory_bytes, + ) + }; + items.fetch_add(1, Ordering::Relaxed); + bytes.fetch_add(source.len, Ordering::Relaxed); + } + + let get_id = view + .master_kv_router() + .inner() + .next_get_id + .fetch_add(1, Ordering::Relaxed); + let gpu_direct_eligible = source.source_kind == GetSourceKind::Memory + && local_owner + .as_deref() + .is_none_or(|owner| source.node_id.as_ref() != owner); + let planned = super::PlannedGetInfo { + put_id: route.put_id, + src_node_id: source.node_id.clone(), + src_tomb_tag: source.tomb_tag.clone(), + key: key.clone(), + controller_node_id: controller_node_id.clone(), + controller_node_start_time: controller.node_start_time, + len: source.len, + src_addr: source.addr, + src_base_addr: source.base_addr, + source_kind: source.source_kind, + atomic_group: route.atomic_group.as_deref().cloned(), + }; + drop(route); + view.master_kv_router() + .inner() + .planned_gets + .insert(get_id, planned.clone()) + .await; + view.master_kv_router() + .inner() + .planned_get_counters + .plan_hits + .fetch_add(1, Ordering::Relaxed); + BatchGetPlanItemResp { + get_id, + node_id: source.node_id.into(), + put_id: planned.put_id, + src_addr: planned.src_addr, + src_base_addr: planned.src_base_addr, + len: planned.len, + source_kind: planned.source_kind, + atomic_group: planned.atomic_group, + gpu_direct_eligible, + error_code: msg_and_error::OK, + error_json: String::new(), + } +} + +fn bound_get_matches_target(info: &InflightGetInfo, target: &GetBindTarget) -> bool { + match (target, &info.target) { + (GetBindTarget::ExternalSink(expected), InflightGetTarget::ExternalSink(actual)) => { + expected == actual + } + ( + GetBindTarget::PreparedLocalReserve(expected), + InflightGetTarget::PreparedLocalReserveSlot(actual), + ) => { + expected.grant_id == actual.grant_id + && expected.slot_index == actual.slot_index + && expected.slot_size == actual.slot_size + && expected.addr == actual.addr + && expected.base_addr == actual.base_addr + } + _ => false, + } +} + +fn bound_get_start_item(get_id: u64, info: &InflightGetInfo) -> BatchGetStartItemResp { + BatchGetStartItemResp { + get_id, + node_id: info.src_node_id.to_string().into(), + put_id: info.put_id, + target_addr: info.target.abs_addr(), + src_addr: info.src_addr, + target_base_addr: info.target.base_addr(), + src_base_addr: info.src_base_addr, + len: info.len, + source_kind: info.source_kind, + prepared_target: match &info.target { + InflightGetTarget::PreparedLocalReserveSlot(slot) => { + Some(GetPreparedLocalReserveTarget { + grant_id: slot.grant_id, + slot_index: slot.slot_index, + slot_size: slot.slot_size, + addr: slot.addr, + base_addr: slot.base_addr, + }) + } + _ => None, + }, + atomic_group: info.atomic_group.clone(), + error_code: msg_and_error::OK, + error_json: String::new(), + } +} + +async fn handle_get_bind_item( + view: MasterKvRouterView, + request: BatchGetBindItemReq, + req_node_id: NodeID, +) -> BatchGetStartItemResp { + let get_id = request.get_id; + let operation_lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(get_id); + let _operation_guard = operation_lock.lock().await; + + if let Some(bound) = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + { + if bound.req_node_id != req_node_id || !bound_get_matches_target(&bound, &request.target) { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetBind replay identity/target mismatch: get_id={} requester={}", + get_id, req_node_id + ), + }), + ); + } + return bound_get_start_item(get_id, &bound); + } + + let Some(planned) = view + .master_kv_router() + .inner() + .planned_gets + .get(&get_id) + .await + else { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::KeyNotFound { + key: format!("planned_get_id:{get_id}"), + }), + ); + }; + let (target, target_tomb_tag, allocation_mode, prepared_requester_lease) = match &request.target + { + GetBindTarget::ExternalSink(target) => { + if planned.source_kind == GetSourceKind::Ssd { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "SSD Get source requires owner-local CPU staging before an external GPU sink: get_id={get_id}" + ), + }), + ); + } + if req_node_id != planned.controller_node_id { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "external GetBind controller mismatch: get_id={} expected={} got={}", + get_id, planned.controller_node_id, req_node_id + ), + }), + ); + } + if let Err(err) = + validate_external_sink_target(&view, &req_node_id, target, planned.len) + { + return get_bind_item_error(get_id, &err); + } + ( + InflightGetTarget::ExternalSink(target.clone()), + None, + GetAllocationMode::ExternalSink, + None, + ) + } + GetBindTarget::PreparedLocalReserve(target) => { + let expected_owner = external_sink_local_owner_id( + &view, + &planned.controller_node_id, + planned.controller_node_start_time, + ); + if expected_owner.as_deref() != Some(req_node_id.as_ref()) { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "prepared GetBind executor is not the controller's owner: get_id={} controller={} expected_owner={:?} got={}", + get_id, planned.controller_node_id, expected_owner, req_node_id + ), + }), + ); + } + let requester_lease = match view.master_kv_router().reserve_prepared_get_requester( + &planned.key, + &req_node_id, + get_id, + ) { + Ok(lease) => lease, + Err(err) => return get_bind_item_error(get_id, &err), + }; + let (slot, tomb_tag) = match validate_prepared_local_reserve_target( + &view, + &req_node_id, + target, + planned.len, + ) { + Ok(value) => value, + Err(err) => return get_bind_item_error(get_id, &err), + }; + ( + InflightGetTarget::PreparedLocalReserveSlot(slot), + Some(tomb_tag), + GetAllocationMode::LocalCommittedSlot, + Some(requester_lease), + ) + } + GetBindTarget::Invalid => { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!("GetBind requires a concrete target: get_id={get_id}"), + }), + ); + } + }; + + let activity_lease = match view + .master_kv_router() + .reserve_inflight_get_key(&planned.key) + { + Ok(lease) => lease, + Err(err) => { + view.master_kv_router() + .inner() + .planned_get_counters + .bind_activity_busy + .fetch_add(1, Ordering::Relaxed); + return get_bind_item_error(get_id, &err); + } + }; + let current_route = view + .master_kv_router() + .inner() + .kv_routes + .get(&planned.key) + .map(|route| route.clone()); + let Some(current_route) = + current_route.filter(|route| planned_get_source_is_current(&planned, route)) + else { + view.master_kv_router() + .inner() + .planned_get_counters + .bind_stale + .fetch_add(1, Ordering::Relaxed); + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::StaleGetPlan { + get_id, + key: planned.key.clone(), + detail: format!( + "source route changed before Bind: source={}", + planned.src_node_id + ), + }), + ); + }; + if matches!(&request.target, GetBindTarget::PreparedLocalReserve(_)) + && current_route + .node_replicas + .read() + .get(&req_node_id) + .is_some_and(|replica| !replica.tomb_tag.is_tomb() && replica.memory.is_some()) + { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "prepared GetBind cannot replace a live owner replica: get_id={} key={} owner={}", + get_id, planned.key, req_node_id + ), + }), + ); + } + + let (src_addr, src_base_addr, ssd_source_allocation) = match planned.source_kind { + GetSourceKind::Memory => (planned.src_addr, planned.src_base_addr, None), + GetSourceKind::Ssd => { + let allocation = + match allocate_ssd_source_stage(&view, &planned.src_node_id, planned.len) { + Ok(allocation) => allocation, + Err(err) => return get_bind_item_error(get_id, &err), + }; + ( + allocation.base_addr() + allocation.addr(), + allocation.base_addr(), + Some(allocation), + ) + } + }; + + let Some(planned) = view + .master_kv_router() + .inner() + .planned_gets + .remove(&get_id) + .await + else { + return get_bind_item_error( + get_id, + &msg_and_error::KvError::Api(msg_and_error::ApiError::KeyNotFound { + key: format!("planned_get_id:{get_id}"), + }), + ); + }; + let inflight = InflightGetInfo { + put_id: planned.put_id, + src_node_id: planned.src_node_id.clone(), + key: planned.key.clone(), + req_node_id: req_node_id.clone(), + controller_node_id: Some(planned.controller_node_id), + len: planned.len, + src_addr, + src_base_addr, + source_kind: planned.source_kind, + ssd_source_allocation, + ssd_stage_lifecycle: (planned.source_kind == GetSourceKind::Ssd) + .then(|| Arc::new(super::SsdStageLifecycle::new())), + atomic_group: planned.atomic_group, + target, + target_tomb_tag, + route: current_route.clone(), + allocation_mode, + durable_reservation: None, + _activity_lease: activity_lease, + _prepared_requester_lease: prepared_requester_lease, + }; + let response = bound_get_start_item(get_id, &inflight); + view.master_kv_router().record_get_source_selection( + req_node_id.as_ref(), + inflight.src_node_id.as_ref(), + inflight.len, + allocation_mode, + inflight.source_kind, + inflight.src_node_id == inflight.req_node_id, + ); + view.master_kv_router() + .inner() + .inflight_gets + .insert(get_id, inflight) + .await; + view.master_kv_router() + .inner() + .planned_get_counters + .bind_succeeded + .fetch_add(1, Ordering::Relaxed); + if current_route.lease_id.is_none() && planned.source_kind == GetSourceKind::Memory { + touch_moka_for_node(view, response.node_id.to_string(), planned.key); + } + response +} + +pub async fn handle_batch_get_plan( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.keys.len()); + for key in req.serialize_part.keys { + items.push(handle_get_plan_item(view.clone(), key, req_node_id.clone()).await); + } + MsgPack { + serialize_part: BatchGetPlanResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_get_bind( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + items.push(handle_get_bind_item(view.clone(), item, req_node_id.clone()).await); + } + MsgPack { + serialize_part: BatchGetBindResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + pub async fn handle_get_start( view: MasterKvRouterView, req: MsgPack, @@ -82,7 +1140,7 @@ pub async fn handle_get_start( let mut remove_in_kv_routes = false; if let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(key) { one_kv_nodes_routes.clean_up_tomb_nodes_replicas(put_id, tombs, view); - if one_kv_nodes_routes.nodes_replicas.read().is_empty() { + if one_kv_nodes_routes.node_replicas.read().is_empty() { remove_in_kv_routes = true; } } @@ -116,11 +1174,41 @@ pub async fn handle_get_start( tracing::debug!("Handling GetStartReq: {:?}", req.serialize_part); + if req.serialize_part.prepared_target.is_some() + && req.serialize_part.external_sink_target.is_some() + { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: "GetStart prepared_target and external_sink_target are mutually exclusive" + .to_string(), + }); + return failed_resp_err(err, None, &view, &req.serialize_part.key); + } + + let activity_lease = match view + .master_kv_router() + .reserve_inflight_get_key(&req.serialize_part.key) + { + Ok(activity_lease) => activity_lease, + Err(err) => return failed_resp_err(err, None, &view, &req.serialize_part.key), + }; + let get_id = view .master_kv_router() .inner() .next_get_id .fetch_add(1, Ordering::Relaxed); + let prepared_requester_lease = if req.serialize_part.prepared_target.is_some() { + match view.master_kv_router().reserve_prepared_get_requester( + &req.serialize_part.key, + &req_node_id, + get_id, + ) { + Ok(lease) => Some(lease), + Err(err) => return failed_resp_err(err, None, &view, &req.serialize_part.key), + } + } else { + None + }; let one_kv_nodes_routes: Arc = if let Some(one_kv_nodes_routes) = view .master_kv_router() @@ -131,37 +1219,106 @@ pub async fn handle_get_start( one_kv_nodes_routes.clone() } else { // Key not found - tracing::info!("Key not found: {}", req.serialize_part.key); + tracing::debug!("Key not found: {}", req.serialize_part.key); let err = msg_and_error::KvError::Api(msg_and_error::ApiError::KeyNotFound { key: req.serialize_part.key.clone(), }); return failed_resp_err(err, None, &view, &req.serialize_part.key); }; - let replicas: HashMap = one_kv_nodes_routes.nodes_replicas.read().clone(); - // Currently we are holding the lock with `replicas` - // 选择一个replica (这里可以实现更复杂的选择逻辑) - let mut replica_keys = replicas.keys().collect::>(); - let mut tombs = HashSet::new(); - let mut target_allocations = None; + let replicas: HashMap = + one_kv_nodes_routes.node_replicas.read().clone(); + let prepared_target = req.serialize_part.prepared_target.clone(); + let external_sink_target = req.serialize_part.external_sink_target.clone(); + let external_sink_local_owner = external_sink_target.as_ref().and_then(|target| { + external_sink_local_owner_id(&view, &req_node_id, target.requester_node_start_time) + }); + let mut tombs = replicas + .iter() + .filter_map(|(node_id, replicas)| replicas.tomb_tag.is_tomb().then_some(node_id.clone())) + .collect::>(); + let mut memory_sources = replicas + .iter() + .filter_map(|(node_id, replicas)| { + (!replicas.tomb_tag.is_tomb() && replicas.memory.is_some()) + .then_some((node_id.clone(), GetSourceKind::Memory)) + }) + .collect::>(); + let mut ssd_sources = replicas + .iter() + .filter_map(|(node_id, replicas)| { + (!replicas.tomb_tag.is_tomb() && replicas.memory.is_none() && replicas.ssd.is_some()) + .then_some((node_id.clone(), GetSourceKind::Ssd)) + }) + .collect::>(); + memory_sources.shuffle(&mut rand::thread_rng()); + ssd_sources.shuffle(&mut rand::thread_rng()); + memory_sources.extend(ssd_sources); + let mut target = None; let mut allocation_mode = GetAllocationMode::Temporary; - for _ in 0..replicas.len() { - let to_remove_idx = rand::thread_rng().gen_range(0..replica_keys.len()); - let selected_replica_key = replica_keys.remove(to_remove_idx); - let selected_replica = replicas.get(&*selected_replica_key).unwrap(); - if selected_replica.tomb_tag.is_tomb() { - tombs.insert(selected_replica_key.to_owned()); + let mut durable_reservation = None; + for (selected_replica_key, source_kind) in memory_sources { + let selected_replicas = replicas + .get(&selected_replica_key) + .expect("selected Get source must exist"); + if external_sink_local_owner + .as_deref() + .is_some_and(|owner_id| selected_replica_key.as_ref() == owner_id) + { + // The explicit GPU path is RDMA-only. The requester's share-group + // owner is local IPC/P2P topology, whose fallback would require + // CPU access to a CUDA virtual address. Leave that replica for the + // ordinary CPU-buffered Get path and search for a remote route. + continue; + } + if source_kind == GetSourceKind::Ssd && external_sink_target.is_some() { continue; } - let src_allocation = selected_replica.allocation.clone(); - let src_node_id = selected_replica.node_id.clone(); + let src_node_id = selected_replica_key; + let (src_len, src_abs_addr, src_base, ssd_source_allocation) = match source_kind { + GetSourceKind::Memory => { + let selected_replica = selected_replicas + .memory + .as_ref() + .expect("memory source candidate must retain memory"); + ( + selected_replica.backing.len(), + selected_replica.backing.abs_addr(), + selected_replica.backing.base_addr(), + None, + ) + } + GetSourceKind::Ssd => { + let ssd = selected_replicas + .ssd + .as_ref() + .expect("SSD source candidate must retain SSD"); + let allocation = match allocate_ssd_source_stage(&view, &src_node_id, ssd.len) { + Ok(allocation) => allocation, + Err(err) => { + return failed_resp_err( + err, + Some((tombs.clone(), one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + ( + ssd.len, + allocation.base_addr() + allocation.addr(), + allocation.base_addr(), + Some(allocation), + ) + } + }; - // 为get调用方分配接收内存作为传输target - if target_allocations.is_none() { - target_allocations = if let Some(replica_on_recv_node) = replicas.get(&req_node_id) { - allocation_mode = GetAllocationMode::ReuseReplica; - Some(replica_on_recv_node.allocation.clone()) - } else { + // For committed-slot replicas on the requester node, we still allocate a + // normal target buffer here instead of reusing the committed slot as a + // MemHolder backing. That keeps the existing get-path carrier stable + // and avoids the old None->unwrap panic. + let mut allocate_request_target = + || -> Result)> { let target_allocation = { let req_node_allocators = view.master_seg_manager().get_node_allocators(&req_node_id); @@ -173,12 +1330,12 @@ pub async fn handle_get_start( let err = msg_and_error::KvError::Unreachable( msg_and_error::UnreachableError::OwnerNoSeg { detail: "config=0 initializes as external; non-zero initializes as owner; the owner must have memory space (segment)".to_string() } ); - return failed_resp_err( + return Err(failed_resp_err( err, - Some((tombs, one_kv_nodes_routes.put_id)), + Some((tombs.clone(), one_kv_nodes_routes.put_id)), &view, &req.serialize_part.key, - ); + )); } let target_allocator = @@ -186,7 +1343,7 @@ pub async fn handle_get_start( let mut allocated_addr: Option = None; for attempt in 1..=3 { - if let Ok(allocation) = target_allocator.allocate(src_allocation.size()) { + if let Ok(allocation) = target_allocator.allocate(src_len) { allocated_addr = Some(allocation); break; } else { @@ -199,51 +1356,167 @@ pub async fn handle_get_start( } if allocated_addr.is_none() { tracing::info!("No space left for target(Requesting node) allocation"); - let total = target_allocator.total_size_bytes(); - let used = target_allocator.used_size_bytes(); - let free = total.saturating_sub(used); + let capacity = target_allocator.node_pool_capacity_snapshot(); let err = msg_and_error::KvError::Api(msg_and_error::ApiError::NoSpace { node: req_node_id.as_ref().to_string(), segment: target_allocator.seg_device_id.clone(), - total_capacity: total, - free_capacity: free, + total_capacity: capacity.active_capacity_bytes, + free_capacity: capacity.available_capacity_bytes, }); - return failed_resp_err( + return Err(failed_resp_err( err, - Some((tombs, one_kv_nodes_routes.put_id)), + Some((tombs.clone(), one_kv_nodes_routes.put_id)), &view, &req.serialize_part.key, - ); + )); } allocated_addr.unwrap() }; - if one_kv_nodes_routes.try_reserve_get_durable_slot() { + if let Some(reservation) = one_kv_nodes_routes.try_reserve_get_durable_slot() { allocation_mode = GetAllocationMode::DurableReplica; + durable_reservation = Some(reservation); } else { allocation_mode = GetAllocationMode::Temporary; } - Some(Arc::new(target_allocation)) + Ok(InflightGetTarget::Allocation(Arc::new(target_allocation))) }; + + // 为get调用方分配接收内存作为传输target + if target.is_none() { + target = Some( + if let Some(external_sink_target) = external_sink_target.as_ref() { + if let Err(err) = validate_external_sink_target( + &view, + &req_node_id, + external_sink_target, + src_len, + ) { + return failed_resp_err( + err, + Some((tombs.clone(), one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + allocation_mode = GetAllocationMode::ExternalSink; + InflightGetTarget::ExternalSink(external_sink_target.clone()) + } else if let Some(prepared_target) = prepared_target.as_ref() { + if replicas.get(&req_node_id).is_some_and(|replicas| { + !replicas.tomb_tag.is_tomb() && replicas.memory.is_some() + }) { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidArgument { + detail: format!( + "prepared local-reserve Get target cannot replace a live replica: key={} requester={}", + req.serialize_part.key, req_node_id + ), + }, + ); + return failed_resp_err( + err, + Some((tombs.clone(), one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + let (slot, _prepared_tomb_tag) = match validate_prepared_local_reserve_target( + &view, + &req_node_id, + prepared_target, + src_len, + ) { + Ok(slot) => slot, + Err(err) => { + return failed_resp_err( + err, + Some((tombs.clone(), one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } + }; + allocation_mode = GetAllocationMode::LocalCommittedSlot; + InflightGetTarget::PreparedLocalReserveSlot(slot) + } else if let Some(replica_on_recv_node) = replicas + .get(&req_node_id) + .filter(|replicas| !replicas.tomb_tag.is_tomb()) + .and_then(|replicas| replicas.memory.as_ref()) + { + match &replica_on_recv_node.backing { + super::KvReplicaBacking::Allocation(allocation) => { + allocation_mode = GetAllocationMode::ReuseReplica; + InflightGetTarget::Allocation(allocation.clone()) + } + super::KvReplicaBacking::CommittedSlot(_) => { + match allocate_request_target() { + Ok(allocation) => allocation, + Err(resp) => return resp, + } + } + } + } else { + match allocate_request_target() { + Ok(allocation) => allocation, + Err(resp) => return resp, + } + }, + ); } - let target_allocation = target_allocations.unwrap(); + let target = target + .as_ref() + .expect("Get target must be selected before building response") + .clone(); + + // Bind the target to the exact registration generation that owns its + // allocator/grant. Looking up only by node id at GetDone would allow + // an old completion to publish addresses into a reconnected node. + let target_tomb_tag = match &target { + InflightGetTarget::Allocation(allocation) => view + .master_seg_manager() + .get_allocation_tomb_tag(&req_node_id, allocation), + InflightGetTarget::PreparedLocalReserveSlot(slot) => view + .master_kv_router() + .inner() + .local_reserve_grants + .get(&slot.grant_id) + .and_then(|grant| { + (grant.owner_node_id == req_node_id + && node_generation_is_current_live(&view, &req_node_id, &grant.tomb_tag)) + .then(|| grant.tomb_tag.clone()) + }), + InflightGetTarget::ExternalSink(_) => None, + }; + if !matches!(&target, InflightGetTarget::ExternalSink(_)) && target_tomb_tag.is_none() { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Get target generation changed before start publication: get_id={} key={} requester={}", + get_id, req.serialize_part.key, req_node_id + ), + }); + return failed_resp_err( + err, + Some((tombs.clone(), one_kv_nodes_routes.put_id)), + &view, + &req.serialize_part.key, + ); + } // Convert to absolute addresses for Mooncake (requires absolute) // Use allocation's allocator base directly - let src_base = src_allocation.base_addr(); - let target_base = target_allocation.base_addr(); + let target_base = target.base_addr(); // If we reuse existing target on requesting node, declare src=target on req node let (resp_node_id, resp_src_addr, resp_target_addr, resp_src_base, resp_target_base) = if allocation_mode == GetAllocationMode::ReuseReplica { - let addr = target_base + target_allocation.addr(); + let addr = target.abs_addr(); // both src/target are on requesting node's allocation in this reuse case (req_node_id.clone(), addr, addr, target_base, target_base) } else { ( src_node_id.clone(), - src_base + src_allocation.addr(), - target_base + target_allocation.addr(), + src_abs_addr, + target.abs_addr(), src_base, target_base, ) @@ -257,21 +1530,46 @@ pub async fn handle_get_start( target_addr: resp_target_addr, src_base_addr: resp_src_base, target_base_addr: resp_target_base, - len: src_allocation.size(), + len: src_len, + source_kind, + prepared_target: (allocation_mode == GetAllocationMode::LocalCommittedSlot) + .then(|| prepared_target.clone()) + .flatten(), + atomic_group: one_kv_nodes_routes.atomic_group.as_deref().cloned(), error_code: msg_and_error::OK, error_json: String::new(), server_process_us: 0, }; + view.master_kv_router().record_get_source_selection( + req_node_id.as_ref(), + resp_node_id.as_ref(), + src_len, + allocation_mode, + source_kind, + src_node_id == req_node_id, + ); // 创建在途的Get操作信息 let info = InflightGetInfo { put_id: one_kv_nodes_routes.put_id, src_node_id: src_node_id.clone(), key: req.serialize_part.key.clone(), req_node_id, - len: src_allocation.size(), - allocation: target_allocation, // 存储target allocation + controller_node_id: None, + len: src_len, + src_addr: resp_src_addr, + src_base_addr: resp_src_base, + source_kind, + ssd_source_allocation, + ssd_stage_lifecycle: (source_kind == GetSourceKind::Ssd) + .then(|| Arc::new(super::SsdStageLifecycle::new())), + atomic_group: one_kv_nodes_routes.atomic_group.as_deref().cloned(), + target, + target_tomb_tag, route: one_kv_nodes_routes.clone(), allocation_mode, + durable_reservation, + _activity_lease: activity_lease, + _prepared_requester_lease: prepared_requester_lease, }; view.master_kv_router() @@ -284,14 +1582,11 @@ pub async fn handle_get_start( // source node's moka to keep the kv alive during transfer (weight=0 => touch). // For leased keys, there should be no moka entry; skip touching to avoid // unnecessary cache work. - if one_kv_nodes_routes.lease_id.is_none() { - update_moka_for_node( + if one_kv_nodes_routes.lease_id.is_none() && source_kind == GetSourceKind::Memory { + touch_moka_for_node( view.clone(), src_node_id.to_string(), req.serialize_part.key.clone(), - 0, - one_kv_nodes_routes.put_id, - false, ); } @@ -308,7 +1603,7 @@ pub async fn handle_get_start( }, ); } - tracing::info!("Key not found: {}", req.serialize_part.key); + tracing::debug!("Key not found: {}", req.serialize_part.key); { let err = msg_and_error::KvError::Api(msg_and_error::ApiError::KeyNotFound { key: req.serialize_part.key.clone(), @@ -322,13 +1617,309 @@ pub async fn handle_get_start( } } +pub async fn handle_ssd_stage_begin( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let get_id = req.serialize_part.get_id; + let operation_lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(get_id); + let _operation_guard = operation_lock.lock().await; + let result = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + .and_then(|inflight| { + (inflight.source_kind == GetSourceKind::Ssd + && inflight.src_node_id == req_node_id + && inflight.ssd_source_allocation.is_some()) + .then(|| inflight.ssd_stage_lifecycle.as_ref().cloned()) + .flatten() + }) + .is_some_and(|lifecycle| lifecycle.begin()); + MsgPack { + serialize_part: SsdStageBeginResp { + started: result, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_ssd_stage_done( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let get_id = req.serialize_part.get_id; + let operation_lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(get_id); + let _operation_guard = operation_lock.lock().await; + let Some(inflight) = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + else { + return MsgPack { + serialize_part: SsdStageDoneResp { + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + }; + }; + if inflight.source_kind != GetSourceKind::Ssd + || !ssd_stage_done_request_authorized( + &inflight.src_node_id, + &inflight.req_node_id, + &req_node_id, + req.serialize_part.drop_ssd_source, + ) + { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "SSD stage terminal requester mismatch: get_id={} source={} target={} got={} drop_ssd_source={}", + get_id, + inflight.src_node_id, + inflight.req_node_id, + req_node_id, + req.serialize_part.drop_ssd_source, + ), + }); + return MsgPack { + serialize_part: SsdStageDoneResp { + error_code: err.code(), + error_json: err.to_json(), + }, + raw_bytes: Vec::new(), + }; + } + if let Some(lifecycle) = inflight.ssd_stage_lifecycle.as_ref() { + lifecycle.finish(); + } + + let mut removed_empty_route = None; + if req.serialize_part.drop_ssd_source { + let route = inflight.route.clone(); + if route.put_id == inflight.put_id && route.remove_ssd_replica(&inflight.src_node_id) { + let key = inflight.key.clone(); + if route.node_replicas.read().is_empty() + && view + .master_kv_router() + .inner() + .kv_routes + .remove_if(&key, |_, current| { + Arc::ptr_eq(current, &route) + && current.put_id == inflight.put_id + && current.node_replicas.read().is_empty() + }) + .is_some() + { + removed_empty_route = Some((key, inflight.put_id)); + } + } + } + drop(inflight); + if let Some((key, put_id)) = removed_empty_route + && view.master_kv_router().prefix_index_enabled() + { + view.master_kv_router() + .inner() + .prefix_index + .write() + .await + .remove(&key, put_id); + } + + MsgPack { + serialize_part: SsdStageDoneResp { + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +fn ssd_stage_done_request_authorized( + source_node_id: &NodeID, + target_node_id: &NodeID, + requester_node_id: &NodeID, + drop_ssd_source: bool, +) -> bool { + requester_node_id == source_node_id || (!drop_ssd_source && requester_node_id == target_node_id) +} + +#[cfg(test)] +mod ssd_stage_done_authorization_tests { + use super::ssd_stage_done_request_authorized; + use crate::cluster_manager::NodeID; + + #[test] + fn source_may_close_or_drop_but_target_may_only_close() { + let source: NodeID = "cpu-source".to_string().into(); + let target: NodeID = "gpu-target".to_string().into(); + let stranger: NodeID = "other-owner".to_string().into(); + + assert!(ssd_stage_done_request_authorized( + &source, &target, &source, false + )); + assert!(ssd_stage_done_request_authorized( + &source, &target, &source, true + )); + assert!(ssd_stage_done_request_authorized( + &source, &target, &target, false + )); + assert!(!ssd_stage_done_request_authorized( + &source, &target, &target, true + )); + assert!(!ssd_stage_done_request_authorized( + &source, &target, &stranger, false + )); + } +} + pub async fn handle_get_revoke( view: MasterKvRouterView, req: MsgPack, + req_node_id: NodeID, ) -> MsgPack { tracing::debug!("Handling GetRevokeReq: {:?}", req.serialize_part); let get_id = req.serialize_part.get_id; + let done_lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(get_id); + let _done_guard = loop { + let guard = done_lock.lock().await; + let active_lifecycle = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + .and_then(|inflight| { + inflight + .ssd_stage_lifecycle + .as_ref() + .filter(|lifecycle| lifecycle.is_active()) + .cloned() + }); + let Some(lifecycle) = active_lifecycle else { + break guard; + }; + drop(guard); + lifecycle.wait_until_not_active().await; + }; + + if let Some(planned) = view + .master_kv_router() + .inner() + .planned_gets + .get(&get_id) + .await + { + let controller_owner = external_sink_local_owner_id( + &view, + &planned.controller_node_id, + planned.controller_node_start_time, + ); + let authorized = planned.controller_node_id == req_node_id + || controller_owner.as_deref() == Some(req_node_id.as_ref()); + if !authorized { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetRevoke planned-operation requester mismatch: get_id={} controller={} got={}", + get_id, planned.controller_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + drop(planned); + if view + .master_kv_router() + .inner() + .planned_gets + .remove(&get_id) + .await + .is_some() + { + view.master_kv_router() + .inner() + .planned_get_counters + .plan_revoked + .fetch_add(1, Ordering::Relaxed); + } + return MsgPack { + serialize_part: GetRevokeResp { + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + }; + } + + if let Some(inflight_info) = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + { + if inflight_info.req_node_id != req_node_id + && inflight_info.controller_node_id.as_ref() != Some(&req_node_id) + { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetRevoke requester mismatch: get_id={} expected={} got={}", + get_id, inflight_info.req_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + } else if let Some(completed) = view + .master_kv_router() + .inner() + .completed_gets + .get(&get_id) + .await + { + let detail = if completed.req_node_id != req_node_id { + format!( + "GetRevoke requester mismatch after completion: get_id={} expected={} got={}", + get_id, completed.req_node_id, req_node_id + ) + } else { + format!( + "GetRevoke lost the Done race; committed target must not be released: get_id={}", + get_id + ) + }; + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { detail }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } // Remove from inflight_gets if let Some(inflight_info) = view @@ -338,8 +1929,9 @@ pub async fn handle_get_revoke( .remove(&get_id) .await { - inflight_info.release_durable_slot_if_needed(); - tracing::info!("Revoked get operation with get_id: {}", get_id); + let _activity_completion = + MasterKeyActivityCompletionGuard::new(inflight_info._activity_lease.clone()); + tracing::debug!("Revoked get operation with get_id: {}", get_id); } else { tracing::warn!("Get operation with get_id {} not found for revoke", get_id); } @@ -351,15 +1943,61 @@ pub async fn handle_get_revoke( }, raw_bytes: Vec::new(), } -} - -pub async fn handle_get_done( - view: MasterKvRouterView, - req: MsgPack, -) -> MsgPack { - tracing::debug!("Handling GetDoneReq: {:?}", req.serialize_part); - - let get_id = req.serialize_part.get_id; +} + +async fn handle_get_done_locked( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, + mut deferred_route_events: Option<&mut Vec>, + mut deferred_terminals: Option<&mut Vec<(u64, CompletedGetInfo)>>, +) -> MsgPack { + tracing::debug!("Handling GetDoneReq: {:?}", req.serialize_part); + + let get_id = req.serialize_part.get_id; + if let Some(inflight_info) = view + .master_kv_router() + .inner() + .inflight_gets + .get(&get_id) + .await + { + if inflight_info.req_node_id != req_node_id { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetDone requester mismatch: get_id={} expected={} got={}", + get_id, inflight_info.req_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + } else if let Some(completed) = view + .master_kv_router() + .inner() + .completed_gets + .get(&get_id) + .await + { + if completed.req_node_id != req_node_id { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetDone requester mismatch after completion: get_id={} expected={} got={}", + get_id, completed.req_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + return MsgPack { + serialize_part: completed.response, + raw_bytes: Vec::new(), + }; + } // Remove from inflight_gets and transfer to get_holding if let Some(inflight_info) = view .master_kv_router() @@ -368,101 +2006,252 @@ pub async fn handle_get_done( .remove(&get_id) .await { + let _activity_completion = + MasterKeyActivityCompletionGuard::new(inflight_info._activity_lease.clone()); let mut allocation_mode = inflight_info.allocation_mode; - let route = inflight_info.route.clone(); // clone req_node_id to avoid borrow/move conflict when inserting into kv_routes let req_node_id = inflight_info.req_node_id.clone(); - // capture allocation capacity before moving it - let alloc_cap = inflight_info.allocation.capcity(); - // Generate holder_id - let holder_id = view - .master_kv_router() - .inner() - .next_holder_id - .fetch_add(1, Ordering::Relaxed); - - let src_node_id = inflight_info.src_node_id; let key = inflight_info.key; + let target_cap = inflight_info.target.capacity(); + if allocation_mode == GetAllocationMode::ExternalSink { + let generation_is_current = match &inflight_info.target { + InflightGetTarget::ExternalSink(target) => { + external_sink_requester_generation_is_current( + &view, + &req_node_id, + target.requester_node_start_time, + ) + } + _ => false, + }; + let terminal = if generation_is_current { + view.master_kv_router() + .view() + .metric_reporter() + .metrics() + .inc_kv_get_done_allocation("external_sink"); + GetDoneResp { + holder_id: 0, + allocation_mode: GetAllocationMode::ExternalSink, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + } + } else { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "external Get sink requester generation departed before Done: get_id={} key={} requester={}", + get_id, key, req_node_id + ), + }, + ); + crate::rpcresp_kvresult_convert::FromError::from_error(&err) + }; + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + return MsgPack { + serialize_part: terminal, + raw_bytes: Vec::new(), + }; + } - // Create holding info - let holding_info = OwnerHoldingGetInfo { - key: key.clone(), - holding_node_id: inflight_info.req_node_id.clone(), - len: inflight_info.len, - allocation: inflight_info.allocation.clone(), + let Some(target_tomb_tag) = inflight_info.target_tomb_tag.as_ref() else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "allocator-backed Get lost its target generation: get_id={} key={} requester={}", + get_id, key, req_node_id + ), + }); + let terminal: GetDoneResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + return MsgPack { + serialize_part: terminal, + raw_bytes: Vec::new(), + }; + }; + if !node_generation_is_current_live(&view, &req_node_id, target_tomb_tag) { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "GetDone target generation departed: get_id={} key={} requester={}", + get_id, key, req_node_id + ), + }); + let terminal: GetDoneResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + return MsgPack { + serialize_part: terminal, + raw_bytes: Vec::new(), + }; + } + // Allocation-backed Gets need a master holder to keep their allocator guard alive. + // Local-reserve slots instead carry independent route and holder references in the + // owner's slot state, so they deliberately do not create a master Allocation holder. + let mut inserted_holder_key = None; + let holder_id = match &inflight_info.target { + InflightGetTarget::Allocation(allocation) => { + let holder_id = view + .master_kv_router() + .inner() + .next_holder_id + .fetch_add(1, Ordering::Relaxed); + let holder_key = + crate::memholder::NodeHolderKey::new(req_node_id.to_string(), holder_id); + view.master_kv_router().inner().get_holding.insert( + holder_key.clone(), + OwnerHoldingGetInfo { + key: key.clone(), + holding_node_id: inflight_info.req_node_id.clone(), + len: inflight_info.len, + allocation: allocation.clone(), + }, + ); + inserted_holder_key = Some(holder_key); + holder_id + } + InflightGetTarget::PreparedLocalReserveSlot(_) => 0, + InflightGetTarget::ExternalSink(_) => { + unreachable!("external Get sink must complete before holder publication") + } }; - // Store in get_holding cache (owned manager, flattened key) - view.master_kv_router().inner().get_holding.insert( - crate::memholder::NodeHolderKey::new(req_node_id.to_string(), holder_id), - holding_info, - ); + // Close the insertion-vs-MemberLeft cleanup race for the holder. If + // MemberLeft marked the shared tag before this check, remove the exact + // holder we just inserted. Otherwise its later cleanup must observe it. + if target_tomb_tag.is_tomb() { + if let Some(holder_key) = inserted_holder_key.as_ref() { + view.master_kv_router() + .inner() + .get_holding + .remove(holder_key); + } + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "GetDone target generation departed during holder publication: get_id={} key={} requester={}", + get_id, key, req_node_id + ), + }); + let terminal: GetDoneResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + return MsgPack { + serialize_part: terminal, + raw_bytes: Vec::new(), + }; + } if allocation_mode == GetAllocationMode::DurableReplica { let mut promote_committed = false; - if let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(&key) { + let mut route_publish_event = None; + if let Some(one_kv_nodes_routes) = view + .master_kv_router() + .inner() + .kv_routes + .get(&key) + .map(|route| route.clone()) + { if one_kv_nodes_routes.put_id == inflight_info.put_id { - let mut nodes_replicas = one_kv_nodes_routes.nodes_replicas.write(); - if let Some(tomb_tag) = - view.master_seg_manager().get_node_tomb_tag(&src_node_id) - { - if !tomb_tag.is_tomb() { - nodes_replicas.insert( - inflight_info.req_node_id.clone(), - KvRouteInfo { - node_id: inflight_info.req_node_id, - allocation: inflight_info.allocation, - tomb_tag, - }, - ); - promote_committed = true; - // Read lease binding from route snapshot: for this put_id, - // if the key is leased, we must NOT insert into moka. - if one_kv_nodes_routes.lease_id.is_none() { - // notify moka cache controller for requesting node after route insert - // See put.rs for rationale: saturate weight to avoid u32 truncation - let req_weight = if alloc_cap > u32::MAX as u64 { - tracing::warn!( - "moka weight saturation on get_done: key={} put_id=({},{}) cap={}B exceeds u32::MAX; weight set to u32::MAX", - key, - inflight_info.put_id.0, - inflight_info.put_id.1, - alloc_cap - ); - u32::MAX - } else { - alloc_cap as u32 - }; - update_moka_for_node( - view.clone(), - req_node_id.to_string(), + match view.master_kv_router().reserve_node_cache_capacity( + &req_node_id, + target_tomb_tag, + ReservedCapacityReason::OwnerIndexedAllocation, + target_cap, + ) { + Ok(capacity_reservation) => { + let replica = KvMemoryReplica { + backing: super::KvReplicaBacking::Allocation(match &inflight_info + .target + { + InflightGetTarget::Allocation(allocation) => allocation.clone(), + InflightGetTarget::PreparedLocalReserveSlot(_) => { + unreachable!("durable Get mode must use Allocation target") + } + InflightGetTarget::ExternalSink(_) => { + unreachable!("durable Get mode cannot use external sink") + } + }), + owner_local_indexed: true, + get_durable_reservation: inflight_info.durable_reservation.clone(), + capacity_reservation, + }; + if publish_route_replica_tomb_fenced( + &one_kv_nodes_routes, + req_node_id.clone(), + replica, + target_tomb_tag.clone(), + ) { + promote_committed = true; + route_publish_event = Some(RoutePublishEvent::replica_append( key.clone(), - req_weight, inflight_info.put_id, - true, - ); + one_kv_nodes_routes.lease_id, + req_node_id.clone(), + target_cap, + )); } else { - tracing::debug!( - "Skip moka insert for leased key={} put_id=({},{}) on node {}", - key, - inflight_info.put_id.0, - inflight_info.put_id.1, - req_node_id + tracing::warn!( + "durable Get replica publication rejected by generation/live-replica fence: get_id={} put_id={:?}", + get_id, + one_kv_nodes_routes.put_id ); } - } else { - tracing::warn!( - "get node is tomb, get_id: {}, put_id: {:?}", - get_id, - one_kv_nodes_routes.put_id - ); } - } else { - tracing::warn!( - "get node is tomb, get_id: {}, put_id: {:?}", + Err(err) => tracing::warn!( + "durable Get could not reserve owner-indexed Allocation capacity; keeping temporary: get_id={} key={} owner={} err={}", get_id, - one_kv_nodes_routes.put_id - ); + key, + req_node_id, + err, + ), } } else { tracing::warn!( @@ -479,29 +2268,247 @@ pub async fn handle_get_done( key ); } + if let Some(event) = route_publish_event { + if let Some(events) = deferred_route_events.as_deref_mut() { + events.push(event); + } else { + apply_post_route_maintenance_batch(&view, vec![event]).await; + } + } if !promote_committed { allocation_mode = GetAllocationMode::Temporary; - route.release_get_durable_slot(); + } + } else if allocation_mode == GetAllocationMode::ReuseReplica { + let mut local_index_published = false; + let mut route_lease_id = None; + let capacity_reservation = view.master_kv_router().reserve_node_cache_capacity( + &req_node_id, + target_tomb_tag, + ReservedCapacityReason::OwnerIndexedAllocation, + target_cap, + ); + match capacity_reservation { + Ok(capacity_reservation) => { + if let Some(current_route) = view + .master_kv_router() + .inner() + .kv_routes + .get(&key) + .map(|route| route.clone()) + { + if current_route.put_id == inflight_info.put_id { + route_lease_id = current_route.lease_id; + let mut replicas = current_route.node_replicas.write(); + if let Some(node_replicas) = replicas.get_mut(&req_node_id) { + local_index_published = !node_replicas.tomb_tag.is_tomb() + && node_replicas.tomb_tag.same_generation(target_tomb_tag) + && node_replicas.memory.as_ref().is_some_and(|replica| { + matches!( + &replica.backing, + super::KvReplicaBacking::Allocation(allocation) + if matches!( + &inflight_info.target, + InflightGetTarget::Allocation(target) + if Arc::ptr_eq(allocation, target) + ) + ) + }); + if local_index_published { + let replica = node_replicas + .memory + .as_mut() + .expect("matched live memory replica"); + replica.owner_local_indexed = true; + replica.capacity_reservation = capacity_reservation; + } + } + } + } + if local_index_published { + let old_ring_b_desc = super::NodeValueReplicaDesc { + weight_bytes: u32::try_from(target_cap).unwrap_or(u32::MAX), + put_id: inflight_info.put_id, + }; + let _ = view + .master_kv_router() + .remove_node_cache_entry_exact( + req_node_id.as_ref(), + &key, + &old_ring_b_desc, + ) + .await; + let event = RoutePublishEvent::replica_append( + key.clone(), + inflight_info.put_id, + route_lease_id, + req_node_id.clone(), + target_cap, + ); + if let Some(events) = deferred_route_events.as_deref_mut() { + events.push(event); + } else { + apply_post_route_maintenance_batch(&view, vec![event]).await; + } + } + } + Err(err) => tracing::warn!( + "reused Get allocation could not reserve owner-indexed capacity; keeping temporary: get_id={} key={} owner={} err={}", + get_id, + key, + req_node_id, + err, + ), + } + if !local_index_published { + tracing::warn!( + "Reused get allocation is no longer the current owner route; returning a temporary holder: get_id={} key={} put_id=({},{}) owner={}", + get_id, + key, + inflight_info.put_id.0, + inflight_info.put_id.1, + req_node_id + ); + allocation_mode = GetAllocationMode::Temporary; + } + } else if allocation_mode == GetAllocationMode::LocalCommittedSlot { + let slot = match &inflight_info.target { + InflightGetTarget::PreparedLocalReserveSlot(slot) => slot.clone(), + InflightGetTarget::Allocation(_) => { + unreachable!("local committed-slot Get mode must use a prepared slot") + } + InflightGetTarget::ExternalSink(_) => { + unreachable!("local committed-slot Get mode cannot use external sink") + } + }; + let mut published = false; + let mut route_publish_event = None; + if let Some(current_route) = view.master_kv_router().inner().kv_routes.get(&key) { + if current_route.put_id == inflight_info.put_id { + let replica = KvMemoryReplica { + backing: super::KvReplicaBacking::CommittedSlot(slot), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation: None, + }; + if publish_route_replica_tomb_fenced( + ¤t_route, + req_node_id.clone(), + replica, + target_tomb_tag.clone(), + ) { + published = true; + route_publish_event = Some(RoutePublishEvent::replica_append( + key.clone(), + inflight_info.put_id, + current_route.lease_id, + req_node_id.clone(), + target_cap, + )); + } + } + } + if !published { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::Unknown { + detail: format!( + "prepared local-reserve Get target could not publish current route: get_id={} key={} put_id=({},{}) owner={}", + get_id, key, inflight_info.put_id.0, inflight_info.put_id.1, req_node_id + ), + }); + let terminal: GetDoneResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + return MsgPack { + serialize_part: terminal, + raw_bytes: Vec::new(), + }; + } + if let Some(event) = route_publish_event { + if let Some(events) = deferred_route_events.as_deref_mut() { + events.push(event); + } else { + apply_post_route_maintenance_batch(&view, vec![event]).await; + } } } - tracing::info!( + tracing::debug!( "Completed get operation with get_id: {}, assigned holder_id: {}", get_id, holder_id ); + view.master_kv_router() + .view() + .metric_reporter() + .metrics() + .inc_kv_get_done_allocation(match allocation_mode { + GetAllocationMode::Temporary => "temporary", + GetAllocationMode::ReuseReplica => "reuse_replica", + GetAllocationMode::DurableReplica => "durable_replica", + GetAllocationMode::LocalCommittedSlot => "local_committed_slot", + GetAllocationMode::ExternalSink => "external_sink", + }); + let terminal = GetDoneResp { + holder_id, + allocation_mode, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }; + let completed = CompletedGetInfo { + req_node_id: req_node_id.clone(), + response: terminal.clone(), + }; + if let Some(terminals) = deferred_terminals.as_deref_mut() { + terminals.push((get_id, completed)); + } else { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } MsgPack { - serialize_part: GetDoneResp { - holder_id, - allocation_mode, - error_code: msg_and_error::OK, - error_json: String::new(), - server_process_us: 0, - }, + serialize_part: terminal, raw_bytes: Vec::new(), } } else { + if let Some(completed) = view + .master_kv_router() + .inner() + .completed_gets + .get(&get_id) + .await + { + if completed.req_node_id != req_node_id { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "GetDone requester mismatch after completion: get_id={} expected={} got={}", + get_id, completed.req_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + return MsgPack { + serialize_part: completed.response, + raw_bytes: Vec::new(), + }; + } tracing::warn!( "Get operation with get_id {} not found for completion", get_id @@ -523,6 +2530,20 @@ pub async fn handle_get_done( } } +pub async fn handle_get_done( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let done_lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(req.serialize_part.get_id); + let _done_guard = done_lock.lock().await; + handle_get_done_locked(view, req, req_node_id, None, None).await +} + // --- MemHolder Handler Functions --- // pub async fn handle_mem_holder_keep_alive( @@ -613,15 +2634,22 @@ pub async fn handle_get_meta( .get(&req.serialize_part.key) { // lock and clone, release the lock quickly - let nodes_replicas: HashMap = - (*one_kv_nodes_routes.nodes_replicas.read()).clone(); + let node_replicas: HashMap = + (*one_kv_nodes_routes.node_replicas.read()).clone(); // Key exists, get metadata from the first replica - for (_, kv_info) in nodes_replicas.iter() { - if kv_info.tomb_tag.is_tomb() { + for replicas in node_replicas.values() { + if replicas.tomb_tag.is_tomb() { continue; } - let len = kv_info.allocation.size(); + let Some(len) = replicas + .memory + .as_ref() + .map(|memory| memory.backing.len()) + .or_else(|| replicas.ssd.as_ref().map(|ssd| ssd.len)) + else { + continue; + }; return MsgPack { serialize_part: GetMetaResp { exists: true, @@ -684,3 +2712,265 @@ pub async fn handle_get_meta( } } } + +pub async fn handle_batch_is_exist( + view: MasterKvRouterView, + req: MsgPack, + _req_node_id: NodeID, +) -> MsgPack { + tracing::debug!( + "Handling BatchIsExistReq: batch_len={}", + req.serialize_part.keys.len() + ); + + let mut exists_list = Vec::with_capacity(req.serialize_part.keys.len()); + + for key in &req.serialize_part.keys { + if let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(key) { + let exists = one_kv_routes_has_live_replica(&one_kv_nodes_routes); + exists_list.push(exists); + } else { + exists_list.push(false); + } + } + + MsgPack { + serialize_part: BatchIsExistResp { + exists_list, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_get_start( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let BatchGetStartReq { + keys, + prepared_targets, + external_sink_targets, + } = req.serialize_part; + if !prepared_targets.is_empty() && prepared_targets.len() != keys.len() { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "batch_get_start prepared target length mismatch: keys={} targets={}", + keys.len(), + prepared_targets.len() + ), + }); + let error: GetStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return MsgPack { + serialize_part: BatchGetStartResp { + items: Vec::new(), + error_code: error.error_code, + error_json: error.error_json, + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }; + } + if !external_sink_targets.is_empty() && external_sink_targets.len() != keys.len() { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "batch_get_start external sink target length mismatch: keys={} targets={}", + keys.len(), + external_sink_targets.len() + ), + }); + let error: GetStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return MsgPack { + serialize_part: BatchGetStartResp { + items: Vec::new(), + error_code: error.error_code, + error_json: error.error_json, + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }; + } + let prepared_targets = if prepared_targets.is_empty() { + vec![None; keys.len()] + } else { + prepared_targets + }; + let external_sink_targets = if external_sink_targets.is_empty() { + vec![None; keys.len()] + } else { + external_sink_targets + }; + let mut items = Vec::with_capacity(keys.len()); + for ((key, prepared_target), external_sink_target) in keys + .into_iter() + .zip(prepared_targets) + .zip(external_sink_targets) + { + let (_get_id, resp) = handle_get_start( + view.clone(), + MsgPack { + serialize_part: GetStartReq { + key, + prepared_target, + external_sink_target, + }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + ) + .await; + let part = resp.serialize_part; + items.push(BatchGetStartItemResp { + get_id: part.get_id, + node_id: part.node_id, + put_id: part.put_id, + target_addr: part.target_addr, + src_addr: part.src_addr, + target_base_addr: part.target_base_addr, + src_base_addr: part.src_base_addr, + len: part.len, + source_kind: part.source_kind, + prepared_target: part.prepared_target, + atomic_group: part.atomic_group, + error_code: part.error_code, + error_json: part.error_json, + }); + } + MsgPack { + serialize_part: BatchGetStartResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_get_revoke( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.get_ids.len()); + for get_id in req.serialize_part.get_ids { + let resp = handle_get_revoke( + view.clone(), + MsgPack { + serialize_part: GetRevokeReq { get_id }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + ) + .await; + let part = resp.serialize_part; + items.push(BatchGetRevokeItemResp { + get_id, + error_code: part.error_code, + error_json: part.error_json, + }); + } + MsgPack { + serialize_part: BatchGetRevokeResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_get_done( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let started_at = Instant::now(); + let get_ids = req.serialize_part.get_ids; + + // Hold every per-get terminal lock until the combined route-maintenance + // batch and idempotency records are durable. Sorting gives overlapping + // retries one global acquisition order and avoids lock cycles. + let mut unique_get_ids = get_ids.clone(); + unique_get_ids.sort_unstable(); + unique_get_ids.dedup(); + let mut _done_guards = Vec::with_capacity(unique_get_ids.len()); + for get_id in unique_get_ids { + let lock = view + .master_kv_router() + .inner() + .get_done_locks + .get_lock(get_id); + _done_guards.push(lock.lock_owned().await); + } + + let mut items = Vec::with_capacity(get_ids.len()); + let mut route_events = Vec::new(); + let mut deferred_terminals = Vec::new(); + let mut response_by_get_id = HashMap::::new(); + for get_id in get_ids { + let part = if let Some(part) = response_by_get_id.get(&get_id) { + part.clone() + } else { + let part = handle_get_done_locked( + view.clone(), + MsgPack { + serialize_part: GetDoneReq { get_id }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + Some(&mut route_events), + Some(&mut deferred_terminals), + ) + .await + .serialize_part; + response_by_get_id.insert(get_id, part.clone()); + part + }; + items.push(BatchGetDoneItemResp { + get_id, + holder_id: part.holder_id, + allocation_mode: part.allocation_mode, + error_code: part.error_code, + error_json: part.error_json, + }); + } + + let route_event_count = route_events.len(); + let maintenance_started_at = Instant::now(); + if !route_events.is_empty() { + // One Moka capacity decision for the entire Done RPC replaces the old + // per-key full-LRU scan while preserving the rule that no success ACK + // is visible before every route is admitted to resident policy state. + apply_post_route_maintenance_batch(&view, route_events).await; + } + let maintenance_elapsed = maintenance_started_at.elapsed(); + for (get_id, completed) in deferred_terminals { + view.master_kv_router() + .inner() + .completed_gets + .insert(get_id, completed) + .await; + } + let elapsed = started_at.elapsed(); + if elapsed.as_millis() >= 100 { + tracing::warn!( + "slow BatchGetDone convergence: items={} route_events={} maintenance_ms={} total_ms={}", + items.len(), + route_event_count, + maintenance_elapsed.as_millis(), + elapsed.as_millis() + ); + } + MsgPack { + serialize_part: BatchGetDoneResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs index ee4ca2b..f09f84e 100644 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/mod.rs @@ -3,65 +3,150 @@ mod get; pub mod msg_pack; pub mod placement; pub mod put; +mod reclaim; mod count_prefix_index; +mod route_maintenance; +mod tiered_writeback; use self::{ count_prefix_index::PrefixRadixTree, delete::handle_batch_delete_ack, delete::handle_delete, delete::handle_delete_ack, - get::{handle_get_done, handle_get_meta, handle_get_revoke, handle_get_start}, + get::{ + handle_batch_get_bind, handle_batch_get_done, handle_batch_get_plan, + handle_batch_get_revoke, handle_batch_get_start, handle_batch_is_exist, handle_get_done, + handle_get_meta, handle_get_revoke, handle_get_start, handle_ssd_stage_begin, + handle_ssd_stage_done, + }, msg_pack::{ - BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, CountPrefixReq, CountPrefixResp, - DeleteAckReq, DeleteReq, GetAllocationMode, GetDoneReq, GetMetaReq, GetRevokeReq, - GetStartReq, PutDoneReq, PutRevokeReq, PutStartReq, + BatchDeleteAckReq, BatchDeleteClientKvMetaCacheReq, BatchEnqueueReplicaTaskReq, + BatchEvictOwnerSourceReq, BatchGetBindReq, BatchGetDoneReq, BatchGetPlanReq, + BatchGetRevokeReq, BatchGetStartReq, BatchIsExistReq, BatchOwnerReclaimReq, + BatchPreparePutKeysReq, BatchPublishOwnerSsdReq, BatchPutAppendDoneReq, + BatchPutAppendStartReq, BatchPutDoneReq, BatchPutRevokeReq, BatchPutStartReq, + BatchReleasePutKeyReservationsReq, CountPrefixReq, CountPrefixResp, DeleteAckReq, + DeleteReq, GetAllocationMode, GetDoneReq, GetDoneResp, GetMetaReq, GetRevokeReq, + GetSourceKind, GetStartReq, GroupedBatchPutDoneReq, OwnerReclaimItem, PutAppendDoneReq, + PutAppendRevokeReq, PutAppendStartReq, PutAtomicGroup, PutDoneReq, PutRevokeReq, + PutStartReq, ReleaseLocalGrantReq, ReserveLocalGrantReq, SsdStageBeginReq, SsdStageDoneReq, + }, + placement::{PlacementPolicy, build_placement_policy}, + put::{ + handle_batch_prepare_put_keys, handle_batch_publish_owner_ssd, + handle_batch_put_append_done, handle_batch_put_append_start, handle_batch_put_done, + handle_batch_put_revoke, handle_batch_put_start, handle_batch_release_put_key_reservations, + handle_grouped_batch_put_done, handle_put_append_done, handle_put_append_revoke, + handle_put_append_start, handle_put_done, handle_put_revoke, handle_put_start, + handle_release_local_grant, handle_reserve_local_grant, }, - placement::{PlacementDefault, PlacementPolicy}, - put::{handle_put_done, handle_put_revoke, handle_put_start}, + reclaim::handle_batch_evict_owner_source, }; use crate::ClientKvApiAccessTrait; use crate::client_kv_api::ClientKvApi; use crate::cluster_manager::{ ClusterEvent, ClusterManager, ClusterManagerAccessTrait, NodeID, NodeIDString, }; -use crate::config::TestSpecConfig; +use crate::config::{ReplicaTaskPlacementConfig, TestSpecConfig}; use crate::master_kv_router::delete::DeleteKeyInfo; use crate::master_kv_router::put::PutIDForAKey; use crate::master_lease_manager::{MasterLeaseManager, MasterLeaseManagerAccessTrait}; use crate::master_seg_manager::MasterSegManager; use crate::master_seg_manager::MasterSegManagerAccessTrait; use crate::master_seg_manager::NodeTombTag; -use crate::master_seg_manager::one_seg_allocator::Allocation; +use crate::master_seg_manager::one_seg_allocator::{Allocation, NodePoolCapacitySnapshot}; use crate::memholder::{EnsureMemholderMgmtDeleteHandle, MasterOwnerMemMgr, MemholderManagerTrait}; use crate::metric_reporter::{MetricReporter, MetricReporterAccessTrait}; use crate::p2p::msg_pack::{MsgPack, RPCCaller, RPCHandler}; -use crate::p2p::p2p_module::{P2pModule, P2pModuleAccessTrait}; -use crate::rpcresp_kvresult_convert::msg_and_error::{KvError, OK}; +use crate::p2p::p2p_module::{ + P2pModule, P2pModuleAccessTrait, UserRpcAsyncHandler, UserRpcFuture, + user_rpc_register_handler_async, +}; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult, OK}; use fluxon_framework::{LogicalModule, define_module}; +use fluxon_util::map_lock::AMapLock; +use fluxon_util::pin_aware_moka::{PinAwareMoka, PinGuard}; +use ::tokio::sync::watch; use async_trait::async_trait; use chrono::Utc; -use dashmap::DashMap; +use dashmap::{DashMap, DashSet}; use limit_thirdparty::tokio::sync::ARwLock; use limit_thirdparty::tokio::{self, sync::ampsc}; use moka::notification::RemovalCause; use parking_lot::Mutex; use parking_lot::RwLock; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::collections::HashSet; -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::{Arc, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tracing::{debug, error, info, warn}; -// Cache capacity policy: fraction of a node's space reserved for the KV cache. -// Keep as a single source of truth to avoid magic numbers scattered across methods. -const MOKA_CACHE_CAPACITY_RATIO: f32 = 0.8; const MAX_GET_DURABLE_REPLICA_SLOTS: u32 = 2; const PLACEMENT_REPORT_INTERVAL_SECS: u64 = 10; const INFLIGHT_PUT_TTL_SECONDS: u64 = 60; const INFLIGHT_PUT_TTL_SECONDS_SKIP_PUT_END_COMMIT: u64 = 5; +const POST_ROUTE_MAINTENANCE_QUEUE_CAPACITY: usize = 512; +const TIER1_WRITEBACK_QUEUE_CAPACITY: usize = 4096; +pub const NODE_POOL_CAPACITY_USER_RPC_PATH: &str = "fluxon_kv/node_pool_capacity"; + +fn subtract_pending_eviction_weight( + pending_weight: &AtomicU64, + owner_node_id: &str, + completed_weight: u64, +) { + if completed_weight == 0 { + return; + } + pending_weight + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current.checked_sub(completed_weight) + }) + .unwrap_or_else(|pending| { + panic!( + "eviction reclaim pending weight underflow for owner {}: pending={} completed={}", + owner_node_id, pending, completed_weight + ) + }); +} + +fn try_install_eviction_reclaim_identities( + inflight: &DashSet, + identities: Vec, +) -> bool { + let mut installed = Vec::with_capacity(identities.len()); + for identity in identities { + if inflight.insert(identity.clone()) { + installed.push(identity); + } else { + for identity in installed { + assert!(inflight.remove(&identity).is_some()); + } + return false; + } + } + true +} + +fn classify_existing_eviction_reclaim( + inflight: &DashSet, + identities: &[reclaim::EvictionReclaimIdentity], +) -> reclaim::EnqueueEvictionReclaimResult { + let inflight_count = identities + .iter() + .filter(|identity| inflight.contains(identity)) + .count(); + if inflight_count == identities.len() { + reclaim::EnqueueEvictionReclaimResult::AlreadyInProgress + } else if inflight_count == 0 { + reclaim::EnqueueEvictionReclaimResult::NotInProgress + } else { + reclaim::EnqueueEvictionReclaimResult::PartialOverlap + } +} #[derive(Clone, Copy, Debug)] pub enum PutPlacementMode { @@ -69,12 +154,298 @@ pub enum PutPlacementMode { Remote, } +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum ReservedCapacityReason { + LocalReserveGrant, + OwnerIndexedAllocation, + LeaseBoundKv, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct NodeCacheCapacityBoundaries { + ring_b_bytes: u64, + tier1_bytes: Option, +} + +fn node_cache_capacity_boundaries( + node_space_size: u64, + replica_cache_capacity_ratio: f64, + replica_writeback_tier1_capacity_ratio: Option, + reserved_capacity_bytes: u64, +) -> NodeCacheCapacityBoundaries { + let ring_b_base = (node_space_size as f64 * replica_cache_capacity_ratio).floor() as u64; + NodeCacheCapacityBoundaries { + ring_b_bytes: ring_b_base.saturating_sub(reserved_capacity_bytes), + // Tier1 contains only metadata and defines when pre-writeback starts. + // It is inclusive of owner residency, so physical ring-B reservations + // never reduce this logical policy window. + tier1_bytes: replica_writeback_tier1_capacity_ratio + .map(|ratio| (node_space_size as f64 * ratio).floor() as u64), + } +} + +#[derive(Debug)] +pub struct NodeCacheReservedCapacity { + /// Exact node registration generation owning these counters. + generation: NodeTombTag, + pub total_bytes: AtomicU64, + pub local_reserve_grant_bytes: AtomicU64, + pub owner_indexed_allocation_bytes: AtomicU64, + pub lease_bound_kv_bytes: AtomicU64, +} + +impl NodeCacheReservedCapacity { + fn new(generation: NodeTombTag) -> Self { + Self { + generation, + total_bytes: AtomicU64::new(0), + local_reserve_grant_bytes: AtomicU64::new(0), + owner_indexed_allocation_bytes: AtomicU64::new(0), + lease_bound_kv_bytes: AtomicU64::new(0), + } + } + + fn apply_delta(&self, reason: ReservedCapacityReason, delta_bytes: i64) { + if delta_bytes >= 0 { + let delta = delta_bytes as u64; + self.total_bytes.fetch_add(delta, Ordering::Relaxed); + match reason { + ReservedCapacityReason::LocalReserveGrant => { + self.local_reserve_grant_bytes + .fetch_add(delta, Ordering::Relaxed); + } + ReservedCapacityReason::OwnerIndexedAllocation => { + self.owner_indexed_allocation_bytes + .fetch_add(delta, Ordering::Relaxed); + } + ReservedCapacityReason::LeaseBoundKv => { + self.lease_bound_kv_bytes + .fetch_add(delta, Ordering::Relaxed); + } + } + } else { + let delta = (-delta_bytes) as u64; + self.total_bytes.fetch_sub(delta, Ordering::Relaxed); + match reason { + ReservedCapacityReason::LocalReserveGrant => { + self.local_reserve_grant_bytes + .fetch_sub(delta, Ordering::Relaxed); + } + ReservedCapacityReason::OwnerIndexedAllocation => { + self.owner_indexed_allocation_bytes + .fetch_sub(delta, Ordering::Relaxed); + } + ReservedCapacityReason::LeaseBoundKv => { + self.lease_bound_kv_bytes + .fetch_sub(delta, Ordering::Relaxed); + } + } + } + } + + fn total_reserved_bytes(&self) -> u64 { + self.total_bytes.load(Ordering::Relaxed) + } +} + +/// Version-scoped reservation that reduces the usable resident-cache capacity +/// while a lease-bound route is alive. The token owns the exact counter Arc; +/// dropping an old Allocation/route can therefore never decrement a newly +/// reconnected node's counters merely because it reused the same node id. +pub struct NodeCacheCapacityReservation { + view: MasterKvRouterView, + node_id: NodeIDString, + generation: NodeTombTag, + reserved_capacity: Arc, + reason: ReservedCapacityReason, + bytes: u64, + released: AtomicBool, +} + +impl std::fmt::Debug for NodeCacheCapacityReservation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NodeCacheCapacityReservation") + .field("node_id", &self.node_id) + .field("bytes", &self.bytes) + .field("released", &self.released.load(Ordering::Acquire)) + .finish() + } +} + +impl Drop for NodeCacheCapacityReservation { + fn drop(&mut self) { + if self.released.swap(true, Ordering::AcqRel) { + return; + } + let Some(_view_guard) = self.view.try_upgrade() else { + return; + }; + if let Err(err) = self + .view + .master_kv_router() + .adjust_node_cache_reserved_capacity_identity( + &self.node_id, + &self.generation, + &self.reserved_capacity, + self.reason, + -(self.bytes as i64), + ) + { + warn!( + "failed to release generation-scoped cache reservation: node={} bytes={} err={}", + self.node_id, self.bytes, err + ); + } + } +} + #[derive(Clone, Debug, Eq, Hash, PartialEq)] pub struct RequesterTargetPair { requester_node_id: NodeIDString, target_node_id: NodeIDString, } +#[derive(Clone, Debug, Default)] +pub struct ReplicaCacheNodeObserveSnapshot { + pub owner_node: String, + pub owner_node_start_time: i64, + pub pool_physical_capacity_bytes: u64, + pub pool_active_capacity_bytes: u64, + pub pool_used_capacity_bytes: u64, + pub pool_parked_capacity_bytes: u64, + pub pool_draining_capacity_bytes: u64, + pub pool_available_capacity_bytes: u64, + pub pool_capacity_epoch: u64, + pub entries: u64, + pub weighted_bytes: u64, + pub effective_capacity_bytes: u64, + pub reserved_capacity_bytes: u64, + pub base_capacity_bytes: u64, + pub pending_eviction_reclaim_bytes: u64, + pub writeback_tier1_entries: u64, + pub writeback_tier1_weighted_bytes: u64, + pub writeback_tier1_capacity_bytes: u64, + pub writeback_tier1_triggered: u64, + pub writeback_tier1_owner_accepted: u64, + pub writeback_tier1_failed: u64, + pub reclaim_master_activity_deferred: u64, + pub reclaim_owner_holder_deferred: u64, + pub reclaim_owner_other_deferred: u64, + pub reclaim_route_changed: u64, + pub reclaim_retry_queued: u64, + pub reclaim_retry_completed: u64, + pub reclaim_retry_restored: u64, + pub reclaim_completed: u64, + pub source_evict_rpc_requests: u64, + pub source_evict_victims: u64, + pub source_evict_requested_bytes: u64, + pub source_evict_accepted: u64, + pub source_evict_in_progress: u64, + pub source_evict_completed: u64, + pub source_evict_retryable_busy: u64, + pub source_evict_stale: u64, + pub source_evict_rejected: u64, + pub last_route_removed_members: u64, + pub last_route_removed_bytes: u64, + pub capacity_eviction_non_ring_b_entry_total: u64, + pub capacity_eviction_hit_committed_slot: u64, + pub eviction_reclaim_deduplicated: u64, +} + +#[derive(Clone, Debug, Default)] +pub struct MasterRuntimeObserveSnapshot { + pub get_holding_entries: u64, + pub get_holding_bytes: u64, + pub replica_cache_nodes: Vec, +} + +/// Generation- and epoch-fenced runtime control request for one owner's preallocated pool. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(tag = "operation", rename_all = "snake_case", deny_unknown_fields)] +pub enum NodePoolCapacityControlRequest { + Get { + owner_node_id: String, + }, + SetActive { + owner_node_id: String, + expected_owner_node_start_time: i64, + expected_capacity_epoch: u64, + active_capacity_bytes: u64, + }, +} + +/// One coherent snapshot returned by both query and mutation operations. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct NodePoolCapacityControlResponse { + pub owner_node_id: String, + pub owner_node_start_time: i64, + pub physical_capacity_bytes: u64, + pub active_capacity_bytes: u64, + pub used_capacity_bytes: u64, + pub parked_capacity_bytes: u64, + pub draining_capacity_bytes: u64, + pub available_capacity_bytes: u64, + pub capacity_epoch: u64, + pub ring_b_base_capacity_bytes: u64, + pub ring_b_effective_capacity_bytes: u64, + pub ring_b_weighted_bytes: u64, + pub ring_b_pending_reclaim_bytes: u64, + pub tier1_capacity_bytes: u64, + pub settled: bool, +} + +struct NodePoolCapacityUserRpcHandler { + view: MasterKvRouterView, +} + +impl UserRpcAsyncHandler for NodePoolCapacityUserRpcHandler { + fn handle(&self, from_node: NodeID, payload: Vec) -> UserRpcFuture { + let view = self.view.clone(); + Box::pin(async move { + let request: NodePoolCapacityControlRequest = serde_json::from_slice(&payload) + .map_err(|err| { + KvError::Api(ApiError::InvalidArgument { + detail: format!("invalid node pool capacity request: {err}"), + }) + })?; + let response = match request { + NodePoolCapacityControlRequest::Get { owner_node_id } => view + .master_kv_router() + .node_pool_capacity_control_snapshot(&owner_node_id)?, + NodePoolCapacityControlRequest::SetActive { + owner_node_id, + expected_owner_node_start_time, + expected_capacity_epoch, + active_capacity_bytes, + } => { + tracing::info!( + caller = %from_node, + owner = %owner_node_id, + expected_owner_node_start_time, + expected_capacity_epoch, + active_capacity_bytes, + "applying node pool active capacity" + ); + view.master_kv_router() + .set_node_pool_active_capacity( + owner_node_id.into(), + expected_owner_node_start_time, + expected_capacity_epoch, + active_capacity_bytes, + ) + .await? + } + }; + serde_json::to_vec(&response).map_err(|err| { + KvError::Api(ApiError::Unknown { + detail: format!("serialize node pool capacity response failed: {err}"), + }) + }) + }) + } +} + impl RequesterTargetPair { fn new(requester_node_id: &str, target_node_id: &str) -> Self { Self { @@ -88,23 +459,70 @@ impl RequesterTargetPair { } } +#[derive(Clone, Debug)] +pub struct CommittedSlotReplica { + pub owner_node_id: NodeID, + pub grant_id: u64, + pub slot_index: u32, + pub slot_size: u64, + pub addr: u64, + pub len: u64, + pub base_addr: u64, +} + /// Information about a `put` operation that is currently in progress. pub enum InflightPutAllocation { /// Local fast path: the same allocation is used as both src (staging) and target (final). Local(Allocation), /// Remote path: separate allocations for src (on requester) and target (on selected node). Remote { src: Allocation, target: Allocation }, + /// Local-first fast path: data is already committed in owner-local reserve slot. + LocalCommittedSlot(CommittedSlotReplica), +} + +#[derive(Clone)] +pub struct InflightPutCommitInfo { + pub node_id: NodeID, + /// Exact target registration generation captured while the allocation was current. + pub target_tomb_tag: NodeTombTag, + pub src_target_allocation: Arc>>, + pub replica_target: Option, } /// Information about a `put` operation that is currently in progress. #[derive(Clone)] pub struct InflightPutInfo { - pub node_id: NodeID, - // seg_name: String, pub key: String, pub req_node_id: NodeID, pub len: u64, - pub src_target_allocation: Arc>>, + pub commit_info: InflightPutCommitInfo, + pub(crate) _activity_lease: Arc, +} + +#[derive(Clone)] +pub struct InflightReplicaTaskInfo { + /// Distinguishes repeated remote-copy attempts for the same `(key, + /// put_id)` generation. A generation may lose its remote route and need + /// another append while the previous attempt's replayable terminal result + /// is still cached. + pub operation_id: u64, + pub node_id: NodeID, + /// Exact target registration generation that owns `target_allocation`. + pub target_tomb_tag: NodeTombTag, + pub source_node_id: NodeID, + pub key: String, + pub put_id: PutIDForAKey, + pub target_allocation: Arc>>, + /// Protect the source-owner copy only after this inclusive replica has been + /// published successfully. Backend-admitted replicas set this bit; tiered + /// write-back and owner-hot exclusive demotion deliberately do not. + pub protect_source_on_remote_complete: bool, + pub(crate) _activity_lease: Arc, +} + +#[derive(Clone, Debug)] +pub struct CompletedReplicaTaskInfo { + pub appended: bool, } /// Information about a `get` operation that is currently in progress. @@ -114,17 +532,158 @@ pub struct InflightGetInfo { pub src_node_id: NodeID, pub key: String, pub req_node_id: NodeID, + /// Optional external controller for a late-bound operation whose executor + /// is the share-group owner in `req_node_id`. + pub controller_node_id: Option, pub len: u64, - pub allocation: Arc, + pub src_addr: u64, + pub src_base_addr: u64, + pub source_kind: crate::master_kv_router::msg_pack::GetSourceKind, + /// Registered source staging owned by the master only for SSD reads. + pub ssd_source_allocation: Option>, + pub(crate) ssd_stage_lifecycle: Option>, + pub atomic_group: Option, + pub target: InflightGetTarget, + /// Exact requester segment-registration generation that owns an + /// allocator-backed target. External sinks are caller-owned and instead + /// carry their membership generation inside `InflightGetTarget`. + pub target_tomb_tag: Option, pub route: Arc, pub allocation_mode: GetAllocationMode, + pub durable_reservation: Option>, + pub(crate) _activity_lease: Arc, + /// Requester-scoped guard for prepared local-reserve Gets. Different GPU + /// owners may fetch the same key concurrently, but one owner must never + /// materialize two candidate committed slots for the same key. + pub(crate) _prepared_requester_lease: Option>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SsdStagePhase { + NotStarted, + Active, + Quiescent, +} + +pub(crate) struct SsdStageLifecycle { + phase: Mutex, + phase_changed: watch::Sender, +} + +impl SsdStageLifecycle { + fn new() -> Self { + let (phase_changed, _receiver) = watch::channel(SsdStagePhase::NotStarted); + Self { + phase: Mutex::new(SsdStagePhase::NotStarted), + phase_changed, + } + } + + pub(crate) fn begin(&self) -> bool { + let mut phase = self.phase.lock(); + match *phase { + SsdStagePhase::NotStarted => { + *phase = SsdStagePhase::Active; + self.phase_changed.send_replace(SsdStagePhase::Active); + true + } + // An RPC response can be lost after the master accepted Begin. + // Replaying the exact source/get operation must not make the owner + // abandon a stage which the master already considers active. + SsdStagePhase::Active => true, + SsdStagePhase::Quiescent => false, + } + } + + pub(crate) fn finish(&self) -> bool { + let mut phase = self.phase.lock(); + if *phase != SsdStagePhase::Active { + return false; + } + *phase = SsdStagePhase::Quiescent; + self.phase_changed.send_replace(SsdStagePhase::Quiescent); + true + } + + pub(crate) fn is_active(&self) -> bool { + *self.phase.lock() == SsdStagePhase::Active + } + + pub(crate) async fn wait_until_not_active(&self) { + let mut phase_changed = self.phase_changed.subscribe(); + loop { + if *phase_changed.borrow_and_update() != SsdStagePhase::Active { + return; + } + if phase_changed.changed().await.is_err() { + return; + } + } + } +} + +/// A metadata-only Get source snapshot before caller-owned destination binding. +/// +/// This type must not own a route, allocation, activity lease, cache pin, or +/// destination resource. Bind re-reads the current route and installs the real +/// Get activity before accepting the snapshot. +#[derive(Clone)] +pub struct PlannedGetInfo { + pub put_id: PutIDForAKey, + pub src_node_id: NodeID, + pub src_tomb_tag: NodeTombTag, + pub key: String, + pub controller_node_id: NodeID, + pub controller_node_start_time: i64, + pub len: u64, + pub src_addr: u64, + pub src_base_addr: u64, + pub source_kind: crate::master_kv_router::msg_pack::GetSourceKind, + pub atomic_group: Option, +} + +#[derive(Clone)] +pub struct CompletedGetInfo { + pub req_node_id: NodeID, + pub response: GetDoneResp, +} + +#[derive(Clone, Debug)] +pub enum InflightGetTarget { + Allocation(Arc), + PreparedLocalReserveSlot(CommittedSlotReplica), + ExternalSink(crate::master_kv_router::msg_pack::GetExternalSinkTarget), +} + +impl InflightGetTarget { + pub fn abs_addr(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.base_addr() + allocation.addr(), + Self::PreparedLocalReserveSlot(slot) => slot.addr, + Self::ExternalSink(target) => target.addr, + } + } + + pub fn base_addr(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.base_addr(), + Self::PreparedLocalReserveSlot(slot) => slot.base_addr, + Self::ExternalSink(target) => target.addr, + } + } + + pub fn capacity(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.capcity(), + Self::PreparedLocalReserveSlot(slot) => slot.slot_size, + Self::ExternalSink(target) => target.capacity, + } + } } impl InflightGetInfo { pub fn release_durable_slot_if_needed(&self) { - if self.allocation_mode == GetAllocationMode::DurableReplica { - self.route.release_get_durable_slot(); - } + // Durable-slot capacity is returned by the reservation token's Drop. } } @@ -137,726 +696,4032 @@ pub struct OwnerHoldingGetInfo { pub allocation: Arc, // The target allocation where data was transferred } -async fn handle_count_prefix( - view: &MasterKvRouterView, - msg: MsgPack, -) -> MsgPack { - let prefix = msg.serialize_part.prefix.clone(); - let inner = view.master_kv_router().inner(); +pub struct LocalReserveGrantInfo { + pub owner_node_id: NodeID, + /// Exact owner registration generation that owns the grant allocation. + pub tomb_tag: NodeTombTag, + pub allocation: Allocation, + /// Excludes this whole grant from the master unindexed-Allocation domain. + pub capacity_reservation: Option>, +} - let count = { - if view.master_kv_router().prefix_index_enabled() { - let tree = inner.prefix_index.read().await; - tree.count_prefix(&prefix) - } else { - inner - .kv_routes - .iter() - .filter(|entry| entry.key().starts_with(&prefix)) - .count() as u64 - } - }; +pub struct PreparedPutKeyReservationInfo { + pub owner_node_id: NodeID, + pub key: String, + pub(crate) _activity_lease: Arc, +} - MsgPack { - serialize_part: CountPrefixResp { - count, - error_code: OK, - error_json: String::new(), - }, - raw_bytes: Vec::new(), - } +#[derive(Default)] +pub(crate) struct EvictionReclaimCounters { + pub master_activity_deferred: AtomicU64, + pub owner_holder_deferred: AtomicU64, + pub owner_other_deferred: AtomicU64, + pub route_changed: AtomicU64, + pub retry_queued: AtomicU64, + pub retry_completed: AtomicU64, + pub retry_restored: AtomicU64, + pub completed: AtomicU64, + pub source_evict_rpc_requests: AtomicU64, + pub source_evict_victims: AtomicU64, + pub source_evict_requested_bytes: AtomicU64, + pub source_evict_accepted: AtomicU64, + pub source_evict_in_progress: AtomicU64, + pub source_evict_completed: AtomicU64, + pub source_evict_retryable_busy: AtomicU64, + pub source_evict_stale: AtomicU64, + pub source_evict_rejected: AtomicU64, + /// Exact reclaim commits that removed the final readable route for a key. + pub last_route_removed_members: AtomicU64, + /// Physical backing bytes represented by `last_route_removed_members`. + pub last_route_removed_bytes: AtomicU64, + /// A Size event from the ring-B controller resolved to any current route + /// outside `Allocation && !owner_local_indexed`. + pub capacity_eviction_non_ring_b_entry_total: AtomicU64, + /// More specific subset retained for diagnostics/backward-compatible + /// metrics: the non-ring-B route used a CommittedSlot backing. + pub capacity_eviction_hit_committed_slot: AtomicU64, + /// Duplicate listener/victim events suppressed while the same physical + /// cache version already has an outstanding reclaim lifecycle. + pub eviction_reclaim_deduplicated: AtomicU64, } -// --- MasterKvRouter Module --- +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct MasterKeyActivitySnapshot { + pub puts: u32, + pub gets: u32, + pub replicas: u32, + pub reclaim_installed: bool, +} -define_module!( - MasterKvRouter, - (master_kv_router, MasterKvRouter), - (p2p, P2pModule), - (master_seg_manager, MasterSegManager), - (cluster_manager, ClusterManager), - (metric_reporter, MetricReporter), - (client_kv_api, ClientKvApi), - (master_lease_manager, MasterLeaseManager) -); +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct MasterKeyActivityObserveSnapshot { + active_keys: u64, + put_keys: u64, + get_keys: u64, + replica_keys: u64, + reclaim_keys: u64, + inflight_puts: u64, + inflight_gets: u64, + inflight_replicas: u64, +} -/// MasterKvRouter module creation parameters -#[derive(Clone, Debug, Default)] -pub struct MasterKvRouterNewArg { - pub test_spec_config: TestSpecConfig, +#[derive(Default)] +pub(crate) struct MasterPlannedGetCounters { + plan_items: AtomicU64, + plan_hits: AtomicU64, + plan_misses: AtomicU64, + bind_succeeded: AtomicU64, + bind_stale: AtomicU64, + bind_activity_busy: AtomicU64, + plan_revoked: AtomicU64, + remote_ssd_filtered_items: AtomicU64, + remote_ssd_filtered_bytes: AtomicU64, } -#[derive(Clone)] -pub struct NodeValueReplicaDesc { - pub weight_bytes: u32, - pub put_id: PutIDForAKey, +#[derive(Default)] +pub(crate) struct MasterSsdTierCounters { + local_ssd_selected_with_remote_memory_items: AtomicU64, + local_ssd_selected_with_remote_memory_bytes: AtomicU64, + local_ssd_selected_without_remote_memory_items: AtomicU64, + local_ssd_selected_without_remote_memory_bytes: AtomicU64, + local_ssd_published_with_remote_memory_items: AtomicU64, + local_ssd_published_with_remote_memory_bytes: AtomicU64, + local_ssd_published_without_remote_memory_items: AtomicU64, + local_ssd_published_without_remote_memory_bytes: AtomicU64, + memory_removed_ssd_survived_items: AtomicU64, + memory_removed_ssd_survived_bytes: AtomicU64, + memory_removed_ssd_became_only_items: AtomicU64, + memory_removed_ssd_became_only_bytes: AtomicU64, } -/// Information about a completed `put` operation that can be retrieved via `get`. -/// Now supports multiple replicas per key. -#[derive(Clone, Debug)] -pub struct KvRouteInfo { - pub node_id: NodeID, - pub allocation: Arc, - pub tomb_tag: NodeTombTag, +#[derive(Clone, Copy, Debug)] +enum MasterKeyActivityKind { + Put, + Get, + Replica, } -#[derive(Debug)] -pub struct OneKvNodesRoutes { - /// the version id for a kv put operation - pub put_id: PutIDForAKey, +#[derive(Default)] +struct MasterKeyActivityState { + puts: u32, + gets: u32, + replicas: u32, + reclaim: Option, +} - /// Lease binding of this key-version on the master. This is an explicit - /// contract set at PutDone time only when the caller provides a lease_id. - /// - /// Semantics and rationale (read before modifying): - /// - This field records whether the current key route (identified by - /// `put_id`) is associated with a lease and which lease it is. - /// - The association is written once during `put_done` together with the - /// route update. Subsequent `get_done` replica additions read this field - /// to decide cache behavior deterministically without consulting the - /// lease manager again on the hot path. - /// - We do not use any fallback or implicit default. Absence (`None`) - /// means "not leased" for this exact `put_id`. Presence (`Some(lease)`) - /// means "leased" and must be respected by cache-controller to prevent - /// eviction-driven global deletes for leased keys. - /// - When a newer `put` arrives, we rebuild a fresh `OneKvNodesRoutes` with - /// the new `put_id` and its (possibly different) lease binding. This keeps - /// the binding strictly version-scoped and avoids state leakage. - /// - Lease expiry/cleanup still owns deletion of leased keys. If a lease - /// expires, the cleanup task deletes keys via master delete. Until then, - /// nodes must not insert leased keys into moka caches. - pub lease_id: Option, +#[derive(Default)] +pub(crate) struct MasterKeyActivityTable { + states: Mutex>, +} - /// node_id -> KvRouteInfo - pub nodes_replicas: RwLock>, - pub get_durable_slots_used: AtomicU32, +pub(crate) struct MasterKeyActivityLease { + table: Arc, + key: String, + kind: MasterKeyActivityKind, + cache_pins: Mutex>, + released: AtomicBool, } -impl OneKvNodesRoutes { - fn clean_up_tomb_nodes_replicas( - &self, - verify_put_id: PutIDForAKey, - tombs: HashSet, - _view: &MasterKvRouterView, - ) -> bool { - if self.put_id != verify_put_id { - return false; +impl MasterKeyActivityLease { + pub(crate) fn attach_cache_pin(&self, pin: PinGuard) { + let mut pins = self.cache_pins.lock(); + if !self.released.load(Ordering::Acquire) { + pins.push(pin); } + } - let mut nodes_replicas = self.nodes_replicas.write(); - nodes_replicas.retain(|_, kv_info| !tombs.contains(&kv_info.node_id)); + pub(crate) fn release_now(&self) { + if self + .released + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.cache_pins.lock().clear(); + self.table.release(&self.key, self.kind); + } + } +} - return true; +impl Drop for MasterKeyActivityLease { + fn drop(&mut self) { + self.release_now(); } +} - fn try_reserve_get_durable_slot(&self) -> bool { - self.get_durable_slots_used - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { - if current < MAX_GET_DURABLE_REPLICA_SLOTS { - Some(current + 1) - } else { - None - } - }) - .is_ok() +pub(crate) struct MasterKeyActivityCompletionGuard { + lease: Option>, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct PreparedGetRequesterKey { + key: String, + requester: NodeID, +} + +#[derive(Default)] +pub(crate) struct PreparedGetRequesterTable { + active: Mutex>, +} + +pub(crate) struct PreparedGetRequesterLease { + table: Arc, + identity: PreparedGetRequesterKey, + get_id: u64, + released: AtomicBool, +} + +impl PreparedGetRequesterTable { + fn reserve( + self: &Arc, + key: &str, + requester: &NodeID, + get_id: u64, + ) -> Option> { + let identity = PreparedGetRequesterKey { + key: key.to_string(), + requester: requester.clone(), + }; + let mut active = self.active.lock(); + if active.contains_key(&identity) { + return None; + } + active.insert(identity.clone(), get_id); + Some(Arc::new(PreparedGetRequesterLease { + table: self.clone(), + identity, + get_id, + released: AtomicBool::new(false), + })) } - fn release_get_durable_slot(&self) { - self.get_durable_slots_used - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { - current.checked_sub(1) + #[cfg(test)] + fn active_get_id(&self, key: &str, requester: &NodeID) -> Option { + self.active + .lock() + .get(&PreparedGetRequesterKey { + key: key.to_string(), + requester: requester.clone(), }) - .unwrap_or_else(|_| panic!("get durable slot underflow indicates a logic bug")); + .copied() } } -#[cfg(test)] -mod tests { - use super::*; - use crate::cluster_manager::ClusterMember; - use std::collections::HashMap; +impl PreparedGetRequesterLease { + pub(crate) fn release_now(&self) { + if self + .released + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + let mut active = self.table.active.lock(); + if active.get(&self.identity) == Some(&self.get_id) { + active.remove(&self.identity); + } + } +} - #[test] - fn one_kv_nodes_routes_only_reserves_two_get_durable_slots() { - let routes = OneKvNodesRoutes { - put_id: (1, 0), - lease_id: None, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), +impl Drop for PreparedGetRequesterLease { + fn drop(&mut self) { + self.release_now(); + } +} + +impl MasterKeyActivityCompletionGuard { + pub(crate) fn new(lease: Arc) -> Self { + Self { lease: Some(lease) } + } + + pub(crate) fn disarm(&mut self) { + self.lease = None; + } +} + +impl Drop for MasterKeyActivityCompletionGuard { + fn drop(&mut self) { + if let Some(lease) = self.lease.as_ref() { + lease.release_now(); + } + } +} + +impl MasterKeyActivityTable { + fn reserve( + self: &Arc, + key: &str, + kind: MasterKeyActivityKind, + reject_same_kind_if_active: bool, + ) -> Option> { + let mut states = self.states.lock(); + let state = states.entry(key.to_string()).or_default(); + if state.reclaim.is_some() { + return None; + } + let counter = match kind { + MasterKeyActivityKind::Put => &mut state.puts, + MasterKeyActivityKind::Get => &mut state.gets, + MasterKeyActivityKind::Replica => &mut state.replicas, }; + if reject_same_kind_if_active && *counter > 0 { + return None; + } + *counter = counter + .checked_add(1) + .expect("master per-key activity counter overflow"); + Some(Arc::new(MasterKeyActivityLease { + table: self.clone(), + key: key.to_string(), + kind, + cache_pins: Mutex::new(Vec::new()), + released: AtomicBool::new(false), + })) + } - assert!(routes.try_reserve_get_durable_slot()); - assert!(routes.try_reserve_get_durable_slot()); - assert!(!routes.try_reserve_get_durable_slot()); + fn release(&self, key: &str, kind: MasterKeyActivityKind) { + let mut states = self.states.lock(); + let remove = { + let state = states + .get_mut(key) + .expect("master per-key activity state missing on release"); + let counter = match kind { + MasterKeyActivityKind::Put => &mut state.puts, + MasterKeyActivityKind::Get => &mut state.gets, + MasterKeyActivityKind::Replica => &mut state.replicas, + }; + *counter = counter + .checked_sub(1) + .expect("master per-key activity counter underflow"); + state.puts == 0 && state.gets == 0 && state.replicas == 0 && state.reclaim.is_none() + }; + if remove { + states.remove(key); + } + } - routes.release_get_durable_slot(); - assert!(routes.try_reserve_get_durable_slot()); + pub(crate) fn try_install_reclaim( + &self, + item: &OwnerReclaimItem, + ) -> Result<(), MasterKeyActivitySnapshot> { + let mut states = self.states.lock(); + let state = states.entry(item.key.clone()).or_default(); + if state.puts != 0 || state.gets != 0 || state.replicas != 0 || state.reclaim.is_some() { + return Err(MasterKeyActivitySnapshot { + puts: state.puts, + gets: state.gets, + replicas: state.replicas, + reclaim_installed: state.reclaim.is_some(), + }); + } + state.reclaim = Some(item.clone()); + Ok(()) } - fn new_test_member(metadata: HashMap) -> ClusterMember { - ClusterMember { - id: "node-a".to_string(), - addresses: Vec::new(), - port: None, - node_start_time: 1, - metadata, - sub_cluster: Some("owner".to_string()), - network: None, + pub(crate) fn is_quiescent(&self, key: &str) -> bool { + let states = self.states.lock(); + match states.get(key) { + None => true, + Some(state) => { + state.puts == 0 && state.gets == 0 && state.replicas == 0 && state.reclaim.is_none() + } } } - #[test] - fn segment_registration_readiness_accepts_owner_without_local_ipc_root() { - let member = new_test_member(HashMap::from([ - ("client".to_string(), "true".to_string()), - ("p2p_relay".to_string(), "true".to_string()), - ])); + fn observe_snapshot(&self) -> MasterKeyActivityObserveSnapshot { + let states = self.states.lock(); + let mut snapshot = MasterKeyActivityObserveSnapshot { + active_keys: u64::try_from(states.len()).unwrap_or(u64::MAX), + ..Default::default() + }; + for state in states.values() { + snapshot.put_keys += u64::from(state.puts != 0); + snapshot.get_keys += u64::from(state.gets != 0); + snapshot.replica_keys += u64::from(state.replicas != 0); + snapshot.reclaim_keys += u64::from(state.reclaim.is_some()); + snapshot.inflight_puts = snapshot.inflight_puts.saturating_add(u64::from(state.puts)); + snapshot.inflight_gets = snapshot.inflight_gets.saturating_add(u64::from(state.gets)); + snapshot.inflight_replicas = snapshot + .inflight_replicas + .saturating_add(u64::from(state.replicas)); + } + snapshot + } - assert!(MasterKvRouter::member_ready_for_segment_registration( - &member - )); + pub(crate) fn reclaim_matches(&self, item: &OwnerReclaimItem) -> bool { + self.states + .lock() + .get(&item.key) + .and_then(|state| state.reclaim.as_ref()) + == Some(item) } - #[test] - fn segment_registration_readiness_rejects_external_and_side_worker() { - let external = new_test_member(HashMap::from([( - "external_client".to_string(), - "true".to_string(), - )])); - assert!(!MasterKvRouter::member_ready_for_segment_registration( - &external - )); + pub(crate) fn clear_reclaim(&self, item: &OwnerReclaimItem) -> bool { + let mut states = self.states.lock(); + let Some(state) = states.get_mut(&item.key) else { + return false; + }; + if state.reclaim.as_ref() != Some(item) { + return false; + } + state.reclaim = None; + if state.puts == 0 && state.gets == 0 && state.replicas == 0 { + states.remove(&item.key); + } + true + } - let side_worker = new_test_member(HashMap::from([ - ("client".to_string(), "true".to_string()), - ("side_transfer_worker".to_string(), "true".to_string()), - ("p2p_relay".to_string(), "true".to_string()), - ])); - assert!(!MasterKvRouter::member_ready_for_segment_registration( - &side_worker - )); + #[cfg(test)] + fn has_reclaim(&self, key: &str) -> bool { + self.states + .lock() + .get(key) + .is_some_and(|state| state.reclaim.is_some()) } } -pub struct MasterKvRouterInner { - view: std::sync::OnceLock, - pub policy: Box, - test_spec_config: TestSpecConfig, +async fn handle_count_prefix( + view: &MasterKvRouterView, + msg: MsgPack, +) -> MsgPack { + let prefix = msg.serialize_part.prefix.clone(); + let inner = view.master_kv_router().inner(); - /// (key, put_time_ms, put_version) -> inflight_put_info - pub inflight_puts: moka::future::Cache<(String, u64, u32), InflightPutInfo>, - /// key -> inflight put count - pub inflight_put_key_counts: Arc>, - pub inflight_gets: moka::future::Cache, + let count = { + if view.master_kv_router().prefix_index_enabled() { + let tree = inner.prefix_index.read().await; + tree.count_prefix(&prefix) + } else { + inner + .kv_routes + .iter() + .filter(|entry| entry.key().starts_with(&prefix)) + .count() as u64 + } + }; - /// Cache for holding get operations (owned, flattened by (node_id, holder_id)) - pub get_holding: MasterOwnerMemMgr, + MsgPack { + serialize_part: CountPrefixResp { + count, + error_code: OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} - /// Counter for get_id - pub next_get_id: AtomicU64, +// --- MasterKvRouter Module --- - /// Counter for holder_id - pub next_holder_id: AtomicU64, +define_module!( + MasterKvRouter, + (master_kv_router, MasterKvRouter), + (p2p, P2pModule), + (master_seg_manager, MasterSegManager), + (cluster_manager, ClusterManager), + (metric_reporter, MetricReporter), + (client_kv_api, ClientKvApi), + (master_lease_manager, MasterLeaseManager) +); - /// Latest version of key-value replicas - pub kv_routes: DashMap>, +/// MasterKvRouter module creation parameters +#[derive(Clone, Debug, Default)] +pub struct MasterKvRouterNewArg { + pub test_spec_config: TestSpecConfig, + pub replica_task_placement: ReplicaTaskPlacementConfig, + pub replica_cache_capacity_ratio: f64, + pub replica_writeback_tier1_capacity_ratio: Option, +} - /// Prefix-counting index derived from `kv_routes`. - /// - /// It is updated through async follow-up maintenance, so it does not guarantee - /// immediate strong-consistency visibility for a freshly committed put. - /// The generic surface is `CountPrefix` RPC; the current primary use case is - /// MQ capacity backpressure / prefix counting. - pub prefix_index: ARwLock, +#[derive(Clone, Debug)] +pub struct NodeValueReplicaDesc { + pub weight_bytes: u32, + pub put_id: PutIDForAKey, +} - /// Support replicas: node_id -> key -> route_info - pub node_kv_cache_controller: - DashMap>>, +#[derive(Clone, Debug, Hash, Eq, PartialEq)] +pub struct MasterPinAlias { + key: String, + put_time_ms: u64, + put_version: u32, +} - /// Per-node total bytes reserved for leased replicas. We subtract this from - /// the base max capacity of each node's moka cache. Acts like a fetch_sub/add counter. - pub lease_reserved_bytes: DashMap>, +impl MasterPinAlias { + fn new(key: &str, put_id: PutIDForAKey) -> Self { + Self { + key: key.to_string(), + put_time_ms: put_id.0, + put_version: put_id.1, + } + } +} - /// Historical final put placement decisions by target node. - pub put_target_decision_counts: DashMap>, +pub type MasterNodeCache = PinAwareMoka; - /// Historical final put placement decisions by requester->target pair. - pub put_requester_target_decision_counts: DashMap>, +/// Information about a completed `put` operation that can be retrieved via `get`. +/// Now supports multiple replicas per key. +#[derive(Clone, Debug)] +pub enum KvReplicaBacking { + Allocation(Arc), + CommittedSlot(CommittedSlotReplica), +} - /// Historical final put placement decisions grouped by placement mode. - pub put_placement_mode_counts: DashMap<&'static str, Arc>, +impl KvReplicaBacking { + pub fn abs_addr(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.base_addr() + allocation.addr(), + Self::CommittedSlot(slot) => slot.addr, + } + } - /// Support replicas: key -> version_id - recent_key_versionid_allocator: moka::sync::SegmentedCache>, + pub fn base_addr(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.base_addr(), + Self::CommittedSlot(slot) => slot.base_addr, + } + } - pub delete_broadcast: EnsureMemholderMgmtDeleteHandle, -} + pub fn len(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.size(), + Self::CommittedSlot(slot) => slot.len, + } + } -impl MasterKvRouterInner { - fn view(&self) -> &MasterKvRouterView { - self.view.get().unwrap() + /// Bytes owned by the physical backing. Capacity accounting must use + /// this value rather than the logical payload length: an Allocation owns + /// its allocator capacity and a committed local-reserve slot owns the + /// whole slot. + pub fn capacity_bytes(&self) -> u64 { + match self { + Self::Allocation(allocation) => allocation.capcity(), + Self::CommittedSlot(slot) => slot.slot_size, + } } } -pub struct MasterKvRouter(MasterKvRouterInner); +#[derive(Clone, Debug)] +pub struct KvMemoryReplica { + pub backing: KvReplicaBacking, + /// Whether this owner also published the route backing into its local key index. + /// Replica-task and remote-put targets are raw master-owned allocations and have no + /// owner-side key entry to fence during capacity eviction. + pub owner_local_indexed: bool, + /// Present only for an allocation replica created by GetDone. The shared + /// token returns the per-key durable-replica budget when this route entry's + /// final clone is dropped. + pub get_durable_reservation: Option>, + /// Excludes a non-ring-B backing from the unindexed-Allocation budget and + /// is released with the exact route lifetime. + pub capacity_reservation: Option>, +} -#[async_trait] -impl LogicalModule for MasterKvRouter { - type View = MasterKvRouterView; - type NewArg = MasterKvRouterNewArg; - type Error = KvError; +#[derive(Clone, Debug)] +pub struct KvSsdReplica { + pub len: u64, +} - fn name(&self) -> &str { - "MasterKvRouter" +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum SsdReplicaCommitStatus { + Committed, + MissingMemory, + TombedNode, + LengthMismatch, +} + +/// All physical backings contributed by one owner generation for one key-version. +/// +/// Route, version, and owner generation stay common. Storage tiers only own their +/// medium-specific state, so adding SSD does not create a second route lifecycle. +#[derive(Clone, Debug)] +pub struct KvNodeReplicas { + pub tomb_tag: NodeTombTag, + pub memory: Option, + pub ssd: Option, +} + +impl KvNodeReplicas { + pub fn memory(tomb_tag: NodeTombTag, memory: KvMemoryReplica) -> Self { + Self { + tomb_tag, + memory: Some(memory), + ssd: None, + } } - fn attach_view(&self, view: Self::View) { - MasterKvRouter::attach_view(self, view); + pub fn has_live_backing(&self) -> bool { + !self.tomb_tag.is_tomb() && (self.memory.is_some() || self.ssd.is_some()) } +} - async fn shutdown(&self) -> Result<(), Self::Error> { - info!("Shutting down MasterKvRouter"); - // Send shutdown signal to delete broadcast task to flush and exit. - if let Err(e) = self - .0 - .delete_broadcast - .sender() - .send(crate::master_kv_router::delete::DeleteKeyInfo::Shutdown) - .await - { - warn!("Failed to send delete broadcast shutdown signal: {}", e); +pub struct GetDurableSlotReservation { + route: Weak, + released: AtomicBool, +} + +impl std::fmt::Debug for GetDurableSlotReservation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GetDurableSlotReservation") + .field("released", &self.released.load(Ordering::Acquire)) + .finish() + } +} + +impl Drop for GetDurableSlotReservation { + fn drop(&mut self) { + if self.released.swap(true, Ordering::AcqRel) { + return; + } + if let Some(route) = self.route.upgrade() { + route.release_get_durable_slot(); } - Ok(()) } } -impl MasterKvRouter { - fn member_ready_for_segment_registration( - member: &crate::cluster_manager::ClusterMember, - ) -> bool { - // Segment registration must follow owner role semantics, not local IPC capability. - // - // Causal chain: - // - owner/external topology is already encoded in cluster member metadata - // (`client`, `external_client`, `side_transfer_worker`, `p2p_relay`); - // - `disable_local_ipc=true` intentionally suppresses `local_ipc_root` so the planner - // does not create same-machine IPC lanes; - // - owner shared bundle publication still depends on segment registration; - // - therefore the registration gate must stay tied to "real owner client" identity and - // must not reuse `local_ipc_root` as a readiness proxy. - member.metadata.get("client").is_some_and(|v| v == "true") - && member - .metadata - .get("p2p_relay") - .is_some_and(|v| v == "true") - && !member - .metadata - .get("external_client") - .is_some_and(|v| v == "true") - && !member - .metadata - .get("side_transfer_worker") - .is_some_and(|v| v == "true") - } +#[derive(Debug)] +pub struct OneKvNodesRoutes { + /// the version id for a kv put operation + pub put_id: PutIDForAKey, - pub fn attach_view(&self, view: MasterKvRouterView) { - // The framework attaches a module's PostView exactly once at the init barrier. - // A second attach indicates a programming error. - self.0 - .view - .set(view) - .unwrap_or_else(|_| panic!("MasterKvRouter view attached twice")); + /// Lease binding of this key-version on the master. This is an explicit + /// contract set at PutDone time only when the caller provides a lease_id. + /// + /// Semantics and rationale (read before modifying): + /// - This field records whether the current key route (identified by + /// `put_id`) is associated with a lease and which lease it is. + /// - The association is written once during `put_done` together with the + /// route update. Subsequent `get_done` replica additions read this field + /// to decide cache behavior deterministically without consulting the + /// lease manager again on the hot path. + /// - We do not use any fallback or implicit default. Absence (`None`) + /// means "not leased" for this exact `put_id`. Presence (`Some(lease)`) + /// means "leased" and must be respected by cache-controller to prevent + /// eviction-driven global deletes for leased keys. + /// - When a newer `put` arrives, we rebuild a fresh `OneKvNodesRoutes` with + /// the new `put_id` and its (possibly different) lease binding. This keeps + /// the binding strictly version-scoped and avoids state leakage. + /// - Lease expiry/cleanup still owns deletion of leased keys. If a lease + /// expires, the cleanup task deletes keys via master delete. Until then, + /// nodes must not insert leased keys into moka caches. + pub lease_id: Option, + + /// Version-scoped multi-key group supplied by the put caller. + pub atomic_group: Option>, + + /// node_id -> all backings owned by that exact node generation. + pub node_replicas: RwLock>, + pub get_durable_slots_used: AtomicU32, +} + +impl OneKvNodesRoutes { + pub(crate) fn commit_ssd_replica(&self, node_id: &NodeID, len: u64) -> SsdReplicaCommitStatus { + let mut replicas = self.node_replicas.write(); + let Some(node_replicas) = replicas.get_mut(node_id) else { + return SsdReplicaCommitStatus::MissingMemory; + }; + if node_replicas.tomb_tag.is_tomb() { + return SsdReplicaCommitStatus::TombedNode; + } + let Some(memory) = node_replicas.memory.as_ref() else { + return SsdReplicaCommitStatus::MissingMemory; + }; + if memory.backing.len() != len { + return SsdReplicaCommitStatus::LengthMismatch; + } + node_replicas.ssd = Some(KvSsdReplica { len }); + SsdReplicaCommitStatus::Committed } - pub async fn construct(arg: MasterKvRouterNewArg) -> Result { - let policy_impl: Box = Box::new(PlacementDefault::new()); - let inflight_put_ttl_seconds = if arg.test_spec_config.skip_put_end_commit { - INFLIGHT_PUT_TTL_SECONDS_SKIP_PUT_END_COMMIT - } else { - INFLIGHT_PUT_TTL_SECONDS + pub(crate) fn remove_ssd_replica(&self, node_id: &NodeID) -> bool { + let mut replicas = self.node_replicas.write(); + let Some(node_replicas) = replicas.get_mut(node_id) else { + return false; }; - let inflight_put_key_counts: Arc> = Arc::new(DashMap::new()); - let inflight_put_key_counts_for_listener = inflight_put_key_counts.clone(); - let inflight_puts = moka::future::Cache::builder() - .time_to_live(Duration::from_secs(inflight_put_ttl_seconds)) - .eviction_listener(move |_put_id, inflight_info: InflightPutInfo, cause| { - if cause == RemovalCause::Expired { - MasterKvRouter::release_inflight_put_key_count_map( - &inflight_put_key_counts_for_listener, - &inflight_info.key, - ); + let removed = node_replicas.ssd.take().is_some(); + if node_replicas.memory.is_none() { + replicas.remove(node_id); + } + removed + } + + fn clean_up_tomb_nodes_replicas( + &self, + verify_put_id: PutIDForAKey, + tombs: HashSet, + _view: &MasterKvRouterView, + ) -> bool { + if self.put_id != verify_put_id { + return false; + } + + let mut node_replicas = self.node_replicas.write(); + node_replicas + .retain(|node_id, replicas| !tombs.contains(node_id) || !replicas.tomb_tag.is_tomb()); + + return true; + } + + fn try_reserve_get_durable_slot(self: &Arc) -> Option> { + self.get_durable_slots_used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + if current < MAX_GET_DURABLE_REPLICA_SLOTS { + Some(current + 1) + } else { + None } }) - .build(); - let inflight_gets = moka::future::Cache::builder() - .time_to_live(Duration::from_secs(60)) - .eviction_listener(|_get_id, inflight_info: InflightGetInfo, cause| { - if cause == RemovalCause::Expired { - inflight_info.release_durable_slot_if_needed(); - } + .ok() + .map(|_| { + Arc::new(GetDurableSlotReservation { + route: Arc::downgrade(self), + released: AtomicBool::new(false), + }) }) - .build(); - let inner = MasterKvRouterInner { - view: std::sync::OnceLock::new(), - policy: policy_impl, - test_spec_config: arg.test_spec_config, - inflight_puts, - inflight_put_key_counts, - inflight_gets, - get_holding: MasterOwnerMemMgr::default(), - next_get_id: AtomicU64::new(0), - next_holder_id: AtomicU64::new(0), - kv_routes: DashMap::new(), - prefix_index: ARwLock::new(PrefixRadixTree::new()), - node_kv_cache_controller: DashMap::new(), - lease_reserved_bytes: DashMap::new(), - put_target_decision_counts: DashMap::new(), - put_requester_target_decision_counts: DashMap::new(), - put_placement_mode_counts: DashMap::new(), - recent_key_versionid_allocator: moka::sync::SegmentedCache::builder(8) - .time_to_idle(Duration::from_secs(5)) - .build(), - delete_broadcast: EnsureMemholderMgmtDeleteHandle::new( - MasterOwnerMemMgr::DELETE_SUBMIT_QUEUE_CAPACITY, - ), - }; - Ok(Self(inner)) } - pub async fn init2_for_init_dag(&self) -> Result<(), KvError> { - info!("MasterKvRouter init2_for_init_dag"); - self.register_rpc_handlers(); - self.register_rpc_callers(); - let view = self.0.view().clone(); + fn release_get_durable_slot(&self) { + self.get_durable_slots_used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current.checked_sub(1) + }) + .unwrap_or_else(|_| panic!("get durable slot underflow indicates a logic bug")); + } +} - self.spawn_cluster_listener(); - self.spawn_put_placement_reporter(); +/// Ring B contains exactly master-only, unindexed Allocation routes. Node +/// role is deliberately absent from this predicate: placement and capacity +/// ownership are orthogonal. +fn ring_b_route_replica_desc( + route: &OneKvNodesRoutes, + node_id: &str, +) -> Option { + if route.lease_id.is_some() { + return None; + } + let replicas = route.node_replicas.read(); + let node_replicas = replicas.get(node_id)?; + let replica = node_replicas.memory.as_ref()?; + if node_replicas.tomb_tag.is_tomb() + || replica.owner_local_indexed + || !matches!(&replica.backing, KvReplicaBacking::Allocation(_)) + { + return None; + } + Some(NodeValueReplicaDesc { + weight_bytes: u32::try_from(replica.backing.capacity_bytes()).unwrap_or(u32::MAX), + put_id: route.put_id, + }) +} - let delete_broadcast_rx = self - .0 - .delete_broadcast - .take_rx() - .expect("delete_broadcast rx already taken, that's impossible"); - delete::spawn_delete_broadcast(view, delete_broadcast_rx); - Ok(()) +/// True only while `tag` is the live registration generation currently bound to `node_id`. +/// Completion paths must validate their captured tag instead of borrowing the tag of a later +/// registration that happens to reuse the same node id. +pub(crate) fn node_generation_is_current_live( + view: &MasterKvRouterView, + node_id: &NodeID, + tag: &NodeTombTag, +) -> bool { + !tag.is_tomb() + && view + .master_seg_manager() + .get_node_tomb_tag(node_id) + .is_some_and(|current| !current.is_tomb() && current.same_generation(tag)) +} + +/// Publish one replica under the route-local lock and close the MemberLeft snapshot race. +/// +/// There are only two possible linearizations: +/// - publication wins the final tag check, so a later MemberLeft snapshot observes the route; +/// - MemberLeft marks the shared tag first, so this function rolls back the exact generation +/// before releasing the route lock. +/// +/// A previous live generation is never overwritten on rollback. In normal operation such a +/// generation cannot coexist with `replica` (registration replacement tombs the old tag first), +/// but preserving it makes the identity rule explicit and keeps tests resistant to ABA setup. +pub(crate) fn publish_route_replica_tomb_fenced( + route: &OneKvNodesRoutes, + node_id: NodeID, + replica: KvMemoryReplica, + publish_tag: NodeTombTag, +) -> bool { + if publish_tag.is_tomb() { + return false; } - pub(crate) fn view(&self) -> &MasterKvRouterView { - self.inner().view() + let mut replicas = route.node_replicas.write(); + if publish_tag.is_tomb() { + return false; } + let previous = match replicas.get_mut(&node_id) { + Some(current) + if !current.tomb_tag.is_tomb() && current.tomb_tag.same_generation(&publish_tag) => + { + // A completion pre-check is necessarily racy with another completion. + // Preserve an already-published memory backing, while allowing the same + // generation to regain DRAM from an SSD-only route. + if current.memory.is_some() { + return false; + } + current.memory = Some(replica); + None + } + Some(current) if !current.tomb_tag.is_tomb() => return false, + _ => replicas.insert( + node_id.clone(), + KvNodeReplicas::memory(publish_tag.clone(), replica), + ), + }; - pub fn inner(&self) -> &MasterKvRouterInner { - &self.0 + if publish_tag.is_tomb() { + let published_is_current = replicas + .get(&node_id) + .is_some_and(|current| current.tomb_tag.same_generation(&publish_tag)); + if published_is_current { + replicas.remove(&node_id); + if let Some(previous) = previous.filter(|old| !old.tomb_tag.is_tomb()) { + replicas.insert(node_id, previous); + } + } + return false; } + true +} - pub fn replica_cache_enabled(&self) -> bool { - !self.0.test_spec_config.disable_master_replica_cache +/// Replace one key's primary route while closing the MemberLeft snapshot race. +/// +/// The DashMap entry guard is held through the final tomb check. Therefore a +/// MemberLeft cleanup that starts after a successful check must observe the new +/// route, while a MemberLeft that marks the generation first causes the exact +/// replacement to be rolled back. Restoring `previous` also protects a newer +/// key version from being lost when a stale completion races node departure. +pub(crate) fn publish_primary_route_tomb_fenced( + routes: &DashMap>, + key: &str, + new_route: Arc, + publish_tag: &NodeTombTag, +) -> Result>, ()> { + if publish_tag.is_tomb() { + return Err(()); } - pub fn prefix_index_enabled(&self) -> bool { - !self.0.test_spec_config.disable_prefix_index + let mut inserted = false; + let mut current = routes.entry(key.to_string()).or_insert_with(|| { + inserted = true; + new_route.clone() + }); + let previous = if inserted { + None + } else { + Some(std::mem::replace(&mut *current, new_route.clone())) + }; + + if publish_tag.is_tomb() { + if let Some(previous) = previous { + *current = previous; + } else { + // DashMap's vacant-entry guard cannot remove itself. Drop the + // shard guard first, then remove only our exact Arc. The route is + // tombed throughout this tiny window, so readers cannot use it. + drop(current); + routes.remove_if(key, |_, route| Arc::ptr_eq(route, &new_route)); + } + return Err(()); } - /// return (put_time_ms, put_version) - pub fn get_recent_key_versionid(&self, key: String) -> (u64, u32) { - let put_time_ms = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards") - .as_millis() as u64; - let put_version = self + Ok(previous) +} + +fn member_left_can_forward_to_registration_actor(current_node_start_time: Option) -> bool { + // MemberLeft carries only a node id. If membership already exposes a live generation, + // forwarding the ambiguous old leave would cancel that generation's registration actor. + current_node_start_time.is_none() +} + +/// Remove one departed node generation from a single route. +/// +/// The route map and replica map deliberately use their normal fine-grained locking only: +/// MemberLeft is a cold path and must not introduce a process-wide lock around a full-table scan. +/// The tomb-tag identity fences a delayed cleanup from deleting a live replica published by a +/// reconnected generation that reused the same node id. +fn remove_departed_generation_from_route( + routes: &DashMap>, + key: &str, + route: &Arc, + node_id: &str, + departed_tag: &NodeTombTag, +) -> Option { + let became_empty = { + let mut replicas = route.node_replicas.write(); + let remove = replicas.get(node_id).is_some_and(|replica| { + replica.tomb_tag.is_tomb() && replica.tomb_tag.same_generation(departed_tag) + }); + if !remove { + return None; + } + replicas.remove(node_id); + replicas.is_empty() + }; + + if !became_empty { + return None; + } + + routes + .remove_if(key, |_, current| { + // `Arc::ptr_eq` prevents an old cleanup from removing a replacement route (ABA). + // Recheck emptiness under the per-route read lock because another replica may have + // joined after the write lock above was released. + Arc::ptr_eq(current, route) && current.node_replicas.read().is_empty() + }) + .map(|(_, removed)| removed.put_id) +} + +const MEMBER_LEFT_ROUTE_CLEANUP_BATCH: usize = 512; + +async fn cleanup_departed_generation_routes( + view: MasterKvRouterView, + node_id: NodeIDString, + departed_tag: NodeTombTag, +) { + // Grants own master-side Allocation guards and therefore must be released + // for the exact departed generation as well. Collect ids without keeping + // DashMap guards across a removal or yield. + let departed_grant_ids: Vec = view + .master_kv_router() + .inner() + .local_reserve_grants + .iter() + .filter_map(|entry| { + (entry.value().owner_node_id.as_ref() == node_id.as_str() + && entry.value().tomb_tag.is_tomb() + && entry.value().tomb_tag.same_generation(&departed_tag)) + .then_some(*entry.key()) + }) + .collect(); + let mut removed_grants = 0usize; + for grant_id in departed_grant_ids { + if view + .master_kv_router() .inner() - .recent_key_versionid_allocator - .get_with(key, || Arc::new(AtomicU32::new(0))) - .fetch_add(1, Ordering::Relaxed); - (put_time_ms, put_version) + .local_reserve_grants + .remove_if(&grant_id, |_, grant| { + grant.owner_node_id.as_ref() == node_id.as_str() + && grant.tomb_tag.is_tomb() + && grant.tomb_tag.same_generation(&departed_tag) + }) + .is_some() + { + removed_grants = removed_grants.saturating_add(1); + } } - fn release_inflight_put_key_count_map(counts: &DashMap, key: &str) { - if let Some(mut entry) = counts.get_mut(key) { - if *entry <= 1 { - drop(entry); - counts.remove(key); - } else { - *entry -= 1; + // Weak snapshots avoid pinning every route Allocation for the duration of a large scan. + // No DashMap guard or replica lock is held across an await. + let route_snapshot: Vec<(String, Weak)> = view + .master_kv_router() + .inner() + .kv_routes + .iter() + .map(|entry| (entry.key().clone(), Arc::downgrade(entry.value()))) + .collect(); + let mut removed_empty_routes = Vec::new(); + let mut removed_replicas = 0usize; + + for batch in route_snapshot.chunks(MEMBER_LEFT_ROUTE_CLEANUP_BATCH) { + for (key, weak_route) in batch { + let Some(route) = weak_route.upgrade() else { + continue; + }; + let had_departed_replica = route + .node_replicas + .read() + .get(node_id.as_str()) + .is_some_and(|replica| { + replica.tomb_tag.is_tomb() && replica.tomb_tag.same_generation(&departed_tag) + }); + if !had_departed_replica { + continue; + } + if let Some(put_id) = remove_departed_generation_from_route( + &view.master_kv_router().inner().kv_routes, + key, + &route, + node_id.as_str(), + &departed_tag, + ) { + removed_empty_routes.push((key.clone(), put_id)); } + removed_replicas = removed_replicas.saturating_add(1); } + tokio::task::yield_now().await; } - pub fn reserve_inflight_put_key( - &self, - key: &str, - reject_if_inflight_same_key: bool, - ) -> Result<(), KvError> { - let counts = &self.inner().inflight_put_key_counts; - let mut entry = counts.entry(key.to_string()).or_insert(0); - if reject_if_inflight_same_key && *entry > 0 { - return Err(KvError::Api( - crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { - key: key.to_string(), - }, - )); + if view.master_kv_router().prefix_index_enabled() { + // Prefix cleanup is batched so one MemberLeft does not spawn one task or acquire one + // async write lock per key. + for batch in removed_empty_routes.chunks(MEMBER_LEFT_ROUTE_CLEANUP_BATCH) { + let mut tree = view.master_kv_router().inner().prefix_index.write().await; + for (key, put_id) in batch { + tree.remove(key, *put_id); + } + drop(tree); + tokio::task::yield_now().await; } - *entry += 1; - Ok(()) } - pub fn release_inflight_put_key(&self, key: &str) { - Self::release_inflight_put_key_count_map(&self.inner().inflight_put_key_counts, key); + info!( + "MemberLeft route cleanup completed: node={} removed_grants={} removed_replicas={} removed_empty_routes={}", + node_id, + removed_grants, + removed_replicas, + removed_empty_routes.len() + ); +} + +fn remove_exact_cache_entry( + cache: &MasterNodeCache, + key: &str, + expected_desc: &NodeValueReplicaDesc, +) -> bool { + cache + .take_if(&key.to_string(), |entry| { + entry.put_id == expected_desc.put_id && entry.weight_bytes == expected_desc.weight_bytes + }) + .is_some() +} + +fn insert_master_cache_entry(cache: &MasterNodeCache, key: String, desc: NodeValueReplicaDesc) { + let alias = MasterPinAlias::new(&key, desc.put_id); + cache.insert(key, [alias], desc); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cluster_manager::ClusterMember; + use crate::master_kv_router::msg_pack::{OwnerReclaimBacking, OwnerReclaimReason}; + use std::collections::HashMap; + + #[test] + fn one_kv_nodes_routes_only_reserves_two_get_durable_slots() { + let routes = Arc::new(OneKvNodesRoutes { + put_id: (1, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::new()), + get_durable_slots_used: AtomicU32::new(0), + }); + + let first = routes.try_reserve_get_durable_slot().unwrap(); + let second = routes.try_reserve_get_durable_slot().unwrap(); + assert!(routes.try_reserve_get_durable_slot().is_none()); + + drop(first); + assert!(routes.try_reserve_get_durable_slot().is_some()); + drop(second); } - fn register_rpc_callers(&self) { - RPCCaller::::new().regist(self.0.view().p2p_module()); + #[test] + fn ssd_stage_lifecycle_replays_begin_and_never_loses_quiescence() { + futures::executor::block_on(async { + let lifecycle = Arc::new(SsdStageLifecycle::new()); + assert!(lifecycle.begin()); + assert!( + lifecycle.begin(), + "a lost Begin response must replay Active as accepted" + ); + + let waiter = lifecycle.wait_until_not_active(); + let finisher = async { + assert!(lifecycle.is_active()); + assert!(lifecycle.finish()); + }; + futures::future::join(waiter, finisher).await; + + assert!(!lifecycle.finish()); + assert!(!lifecycle.begin()); + lifecycle.wait_until_not_active().await; + }); } - fn register_rpc_handlers(&self) { - let p2p = self.0.view().p2p_module(); + #[test] + fn durable_slot_token_returns_capacity_across_ten_fill_demote_cycles() { + let routes = Arc::new(OneKvNodesRoutes { + put_id: (1, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::new()), + get_durable_slots_used: AtomicU32::new(0), + }); - // --- Get Handlers --- - let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + for cycle in 0..10 { + let reservation = routes + .try_reserve_get_durable_slot() + .expect("each fill must reacquire durable capacity after demotion"); + assert_eq!(routes.get_durable_slots_used.load(Ordering::Acquire), 1); + let route_clone = reservation.clone(); + drop(reservation); + assert_eq!( + routes.get_durable_slots_used.load(Ordering::Acquire), + 1, + "route clone must retain the token during cycle {cycle}" + ); + drop(route_clone); + assert_eq!( + routes.get_durable_slots_used.load(Ordering::Acquire), + 0, + "demotion must return durable capacity during cycle {cycle}" + ); + } + } + + #[test] + fn prepared_get_singleflight_is_same_owner_only_and_aba_safe() { + let table = Arc::new(PreparedGetRequesterTable::default()); + let gpu0: NodeID = "gpu-owner-0".to_string().into(); + let gpu1: NodeID = "gpu-owner-1".to_string().into(); + + let gpu0_first = table + .reserve("shared-key", &gpu0, 11) + .expect("first prepared Get on gpu0 must lead"); + assert!( + table.reserve("shared-key", &gpu0, 12).is_none(), + "same owner cannot materialize a second candidate slot" + ); + let gpu1_first = table + .reserve("shared-key", &gpu1, 13) + .expect("different requester owners remain parallel"); + assert_eq!(table.active_get_id("shared-key", &gpu0), Some(11)); + assert_eq!(table.active_get_id("shared-key", &gpu1), Some(13)); + + gpu0_first.release_now(); + let gpu0_second = table + .reserve("shared-key", &gpu0, 14) + .expect("released owner identity can be reused"); + // A delayed Drop/release of the old generation must not delete get 14. + gpu0_first.release_now(); + assert_eq!(table.active_get_id("shared-key", &gpu0), Some(14)); + + drop(gpu0_second); + drop(gpu1_first); + assert_eq!(table.active_get_id("shared-key", &gpu0), None); + assert_eq!(table.active_get_id("shared-key", &gpu1), None); + } + + #[test] + fn reclaim_fence_is_atomic_with_all_master_key_activity() { + let table = Arc::new(MasterKeyActivityTable::default()); + let item = OwnerReclaimItem { + key: "k".to_string(), + put_id: (7, 3), + epoch: 11, + backing: OwnerReclaimBacking::CommittedSlot { + grant_id: 13, + slot_index: 17, + slot_size: 8 * 1024 * 1024, + }, + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + + let get_lease = table + .reserve("k", MasterKeyActivityKind::Get, false) + .expect("get activity should be admitted before reclaim"); + assert!(!table.is_quiescent("k")); + assert_eq!( + table.try_install_reclaim(&item), + Err(MasterKeyActivitySnapshot { + puts: 0, + gets: 1, + replicas: 0, + reclaim_installed: false, + }) + ); + drop(get_lease); + + assert!(table.is_quiescent("k")); + assert!(table.try_install_reclaim(&item).is_ok()); + assert!(!table.is_quiescent("k")); + assert!(table.has_reclaim("k")); + assert!( + table + .reserve("k", MasterKeyActivityKind::Put, false) + .is_none() + ); + assert!( + table + .reserve("k", MasterKeyActivityKind::Get, false) + .is_none() + ); + assert!( + table + .reserve("k", MasterKeyActivityKind::Replica, false) + .is_none() + ); + + assert!(table.clear_reclaim(&item)); + assert!(!table.has_reclaim("k")); + assert!( + table + .reserve("k", MasterKeyActivityKind::Put, true) + .is_some() + ); + } + + #[test] + fn master_key_activity_observe_snapshot_counts_keys_and_leases() { + let table = Arc::new(MasterKeyActivityTable::default()); + let put = table + .reserve("shared", MasterKeyActivityKind::Put, false) + .expect("put activity must be admitted"); + let get_a = table + .reserve("shared", MasterKeyActivityKind::Get, false) + .expect("first get activity must be admitted"); + let get_b = table + .reserve("shared", MasterKeyActivityKind::Get, false) + .expect("second get activity must be admitted"); + let replica = table + .reserve("replica", MasterKeyActivityKind::Replica, false) + .expect("replica activity must be admitted"); + let reclaim = OwnerReclaimItem { + key: "reclaim".to_string(), + put_id: (9, 1), + epoch: 17, + backing: OwnerReclaimBacking::CommittedSlot { + grant_id: 19, + slot_index: 23, + slot_size: 8 * 1024 * 1024, + }, + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + table + .try_install_reclaim(&reclaim) + .expect("idle key must accept reclaim fence"); + + assert_eq!( + table.observe_snapshot(), + MasterKeyActivityObserveSnapshot { + active_keys: 3, + put_keys: 1, + get_keys: 1, + replica_keys: 1, + reclaim_keys: 1, + inflight_puts: 1, + inflight_gets: 2, + inflight_replicas: 1, + } + ); + + drop(put); + drop(get_a); + drop(get_b); + drop(replica); + assert!(table.clear_reclaim(&reclaim)); + assert_eq!( + table.observe_snapshot(), + MasterKeyActivityObserveSnapshot::default() + ); + } + + #[test] + fn explicit_activity_completion_is_idempotent_with_retired_cache_clones() { + let table = Arc::new(MasterKeyActivityTable::default()); + let lease = table + .reserve("retired-clone", MasterKeyActivityKind::Get, false) + .expect("get activity should be admitted"); + let retired_cache_clone = lease.clone(); + + lease.release_now(); + lease.release_now(); + assert!(table.is_quiescent("retired-clone")); + + drop(lease); + drop(retired_cache_clone); + assert!(table.is_quiescent("retired-clone")); + } + + #[test] + fn replica_terminal_result_is_scoped_to_one_append_attempt() { + let terminals = moka::sync::Cache::builder() + .time_to_live(Duration::from_secs(120)) + .build(); + let key = "same-kv-generation".to_string(); + let put_id = (17, 3); + let first_operation_id = 41; + let second_operation_id = 42; + + terminals.insert( + (key.clone(), put_id.0, put_id.1, first_operation_id), + CompletedReplicaTaskInfo { appended: true }, + ); + + assert!( + terminals + .get(&(key.clone(), put_id.0, put_id.1, first_operation_id)) + .is_some(), + "a retry of the same append attempt must replay its terminal result" + ); + assert!( + terminals + .get(&(key, put_id.0, put_id.1, second_operation_id)) + .is_none(), + "an old terminal result must not complete a later remote-copy attempt" + ); + } + + fn test_memory_replica(node_id: &str) -> KvMemoryReplica { + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: node_id.to_string().into(), + grant_id: 1, + slot_index: 0, + slot_size: 1024, + addr: 0, + len: 1024, + base_addr: 0, + }), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation: None, + } + } + + fn test_route_info_with_tag(node_id: &str, tomb_tag: NodeTombTag) -> KvNodeReplicas { + KvNodeReplicas::memory(tomb_tag, test_memory_replica(node_id)) + } + + fn test_route_info(node_id: &str, tomb: bool) -> KvNodeReplicas { + let tomb_tag = NodeTombTag::new(); + if tomb { + tomb_tag.set_tomb(); + } + test_route_info_with_tag(node_id, tomb_tag) + } + + fn test_allocation_route_info(node_id: &str, owner_local_indexed: bool) -> KvNodeReplicas { + let allocator = Arc::new( + crate::master_seg_manager::one_seg_allocator::OneSegAllocator::new( + format!("{node_id}-segment"), + crate::master_seg_manager::msg_pack::SegmentDeviceDescription::Cpu, + 0, + 4096, + ) + .unwrap(), + ); + let allocation = allocator.allocate(1024).unwrap(); + KvNodeReplicas::memory( + NodeTombTag::new(), + KvMemoryReplica { + backing: KvReplicaBacking::Allocation(Arc::new(allocation)), + owner_local_indexed, + get_durable_reservation: None, + capacity_reservation: None, + }, + ) + } + + fn test_route( + put_id: PutIDForAKey, + lease_id: Option, + replicas: Vec<(&str, bool)>, + ) -> OneKvNodesRoutes { + OneKvNodesRoutes { + put_id, + lease_id, + atomic_group: None, + node_replicas: RwLock::new( + replicas + .into_iter() + .map(|(node_id, tomb)| { + (node_id.to_string().into(), test_route_info(node_id, tomb)) + }) + .collect(), + ), + get_durable_slots_used: AtomicU32::new(0), + } + } + + #[test] + fn ring_b_admission_depends_on_backing_and_local_index_not_node_role() { + let node: NodeID = "same-node".to_string().into(); + let make_route = |lease_id, replica| OneKvNodesRoutes { + put_id: (91, 3), + lease_id, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([(node.clone(), replica)])), + get_durable_slots_used: AtomicU32::new(0), + }; + + let unindexed = make_route(None, test_allocation_route_info("same-node", false)); + let desc = ring_b_route_replica_desc(&unindexed, "same-node") + .expect("unindexed Allocation must enter ring B on any node role"); + assert_eq!(desc.put_id, (91, 3)); + assert_eq!(desc.weight_bytes, 4096); + + let indexed = make_route(None, test_allocation_route_info("same-node", true)); + assert!(ring_b_route_replica_desc(&indexed, "same-node").is_none()); + + let leased = make_route(Some(7), test_allocation_route_info("same-node", false)); + assert!(ring_b_route_replica_desc(&leased, "same-node").is_none()); + + let mut committed = test_route_info("same-node", false); + committed + .memory + .as_mut() + .expect("test route has memory") + .owner_local_indexed = false; + let committed = make_route(None, committed); + assert!(ring_b_route_replica_desc(&committed, "same-node").is_none()); + } + + #[test] + fn ssd_commit_reuses_the_live_owner_route_and_validates_memory_length() { + let owner: NodeID = "owner".to_string().into(); + let route = test_route((92, 1), None, vec![("owner", false)]); + + assert_eq!( + route.commit_ssd_replica(&owner, 2048), + SsdReplicaCommitStatus::LengthMismatch + ); + assert_eq!( + route.commit_ssd_replica(&owner, 1024), + SsdReplicaCommitStatus::Committed + ); + { + let replicas = route.node_replicas.read(); + let owner_replicas = replicas.get(&owner).unwrap(); + assert!(owner_replicas.memory.is_some()); + assert_eq!(owner_replicas.ssd.as_ref().map(|ssd| ssd.len), Some(1024)); + } + + route.node_replicas.write().get_mut(&owner).unwrap().memory = None; + assert_eq!( + route.commit_ssd_replica(&owner, 1024), + SsdReplicaCommitStatus::MissingMemory + ); + assert!(route.remove_ssd_replica(&owner)); + assert!( + route.node_replicas.read().is_empty(), + "a stale SSD-only owner entry must converge to an empty route" + ); + assert!(!route.remove_ssd_replica(&owner)); + } + + #[test] + fn member_left_cleanup_removes_only_the_exact_tomb_generation() { + let routes = DashMap::new(); + let departed_tag = NodeTombTag::new(); + departed_tag.set_tomb(); + let live_reconnect_tag = NodeTombTag::new(); + let newer_tomb_tag = NodeTombTag::new(); + newer_tomb_tag.set_tomb(); + + let departed_route = Arc::new(OneKvNodesRoutes { + put_id: (10, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", departed_tag.clone()), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let live_route = Arc::new(OneKvNodesRoutes { + put_id: (11, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", live_reconnect_tag), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let newer_tomb_route = Arc::new(OneKvNodesRoutes { + put_id: (12, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", newer_tomb_tag), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + routes.insert("departed".to_string(), departed_route.clone()); + routes.insert("live".to_string(), live_route.clone()); + routes.insert("newer-tomb".to_string(), newer_tomb_route.clone()); + + assert_eq!( + remove_departed_generation_from_route( + &routes, + "departed", + &departed_route, + "owner", + &departed_tag, + ), + Some((10, 0)) + ); + assert!(routes.get("departed").is_none()); + assert_eq!( + remove_departed_generation_from_route( + &routes, + "live", + &live_route, + "owner", + &departed_tag, + ), + None + ); + assert!( + live_route + .node_replicas + .read() + .get("owner") + .is_some_and(|replica| !replica.tomb_tag.is_tomb()) + ); + assert_eq!( + remove_departed_generation_from_route( + &routes, + "newer-tomb", + &newer_tomb_route, + "owner", + &departed_tag, + ), + None + ); + assert!(newer_tomb_route.node_replicas.read().contains_key("owner")); + } + + #[test] + fn member_left_empty_route_removal_is_arc_identity_aba_safe() { + let routes = DashMap::new(); + let departed_tag = NodeTombTag::new(); + departed_tag.set_tomb(); + let old_route = Arc::new(OneKvNodesRoutes { + put_id: (20, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", departed_tag.clone()), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let replacement = Arc::new(test_route((21, 0), None, vec![("peer", false)])); + routes.insert("aba".to_string(), old_route.clone()); + routes.insert("aba".to_string(), replacement.clone()); + + assert_eq!( + remove_departed_generation_from_route( + &routes, + "aba", + &old_route, + "owner", + &departed_tag, + ), + None + ); + let current = routes.get("aba").expect("replacement route must survive"); + assert!(Arc::ptr_eq(current.value(), &replacement)); + assert!( + current + .node_replicas + .read() + .get("peer") + .is_some_and(|replica| !replica.tomb_tag.is_tomb()) + ); + } + + #[test] + fn replica_publish_fence_never_overwrites_live_or_publishes_tomb_generation() { + let old_tag = NodeTombTag::new(); + let route = test_route((30, 0), None, vec![]); + route.node_replicas.write().insert( + "owner".to_string().into(), + test_route_info_with_tag("owner", old_tag.clone()), + ); + + let contender_tag = NodeTombTag::new(); + assert!(!publish_route_replica_tomb_fenced( + &route, + "owner".to_string().into(), + test_memory_replica("owner"), + contender_tag, + )); + assert!( + route + .node_replicas + .read() + .get("owner") + .is_some_and(|replica| replica.tomb_tag.same_generation(&old_tag)) + ); + + old_tag.set_tomb(); + let replacement_tag = NodeTombTag::new(); + assert!(publish_route_replica_tomb_fenced( + &route, + "owner".to_string().into(), + test_memory_replica("owner"), + replacement_tag.clone(), + )); + assert!( + route + .node_replicas + .read() + .get("owner") + .is_some_and(|replica| replica.tomb_tag.same_generation(&replacement_tag)) + ); + + let departed_tag = NodeTombTag::new(); + departed_tag.set_tomb(); + assert!(!publish_route_replica_tomb_fenced( + &route, + "departed".to_string().into(), + test_memory_replica("departed"), + departed_tag, + )); + assert!(!route.node_replicas.read().contains_key("departed")); + } + + #[test] + fn primary_publish_fence_restores_previous_route_on_tomb_generation() { + let routes = DashMap::new(); + let previous = Arc::new(test_route((40, 0), None, vec![("peer", false)])); + routes.insert("key".to_string(), previous.clone()); + + let departed_tag = NodeTombTag::new(); + departed_tag.set_tomb(); + let rejected = Arc::new(OneKvNodesRoutes { + put_id: (41, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", departed_tag.clone()), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + assert!( + publish_primary_route_tomb_fenced(&routes, "key", rejected, &departed_tag).is_err() + ); + assert!( + routes + .get("key") + .is_some_and(|route| Arc::ptr_eq(route.value(), &previous)) + ); + + let live_tag = NodeTombTag::new(); + let accepted = Arc::new(OneKvNodesRoutes { + put_id: (42, 0), + lease_id: None, + atomic_group: None, + node_replicas: RwLock::new(HashMap::from([( + "owner".to_string().into(), + test_route_info_with_tag("owner", live_tag.clone()), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let replaced = + publish_primary_route_tomb_fenced(&routes, "key", accepted.clone(), &live_tag) + .expect("live generation must publish") + .expect("previous route must be returned"); + assert!(Arc::ptr_eq(&replaced, &previous)); + assert!( + routes + .get("key") + .is_some_and(|route| Arc::ptr_eq(route.value(), &accepted)) + ); + } + + #[test] + fn reserved_capacity_counters_are_generation_and_arc_scoped() { + let old_tag = NodeTombTag::new(); + let new_tag = NodeTombTag::new(); + let old = Arc::new(NodeCacheReservedCapacity::new(old_tag.clone())); + let new = Arc::new(NodeCacheReservedCapacity::new(new_tag.clone())); + + old.apply_delta(ReservedCapacityReason::LeaseBoundKv, 4096); + new.apply_delta(ReservedCapacityReason::LeaseBoundKv, 8192); + old_tag.set_tomb(); + old.apply_delta(ReservedCapacityReason::LeaseBoundKv, -4096); + + assert_eq!(old.total_reserved_bytes(), 0); + assert_eq!(new.total_reserved_bytes(), 8192); + assert!(!old.generation.same_generation(&new.generation)); + assert!(!new_tag.is_tomb()); + } + + #[test] + fn tier1_capacity_is_independent_of_ring_b_reservations() { + let node_space_size = 128 * 1024 * 1024 * 1024; + let boundaries = + node_cache_capacity_boundaries(node_space_size, 0.95, Some(0.75), 124_554_051_584); + + assert_eq!(boundaries.ring_b_bytes, 6_012_954_214); + assert_eq!(boundaries.tier1_bytes, Some(96 * 1024 * 1024 * 1024)); + assert!(boundaries.tier1_bytes.unwrap() > boundaries.ring_b_bytes); + } + + #[test] + fn delayed_member_left_is_not_forwarded_to_a_live_generation_actor() { + assert!(member_left_can_forward_to_registration_actor(None)); + assert!(!member_left_can_forward_to_registration_actor(Some(17))); + } + + #[test] + fn ring_b_cache_is_bounded_for_every_node_role() { + let cache = MasterNodeCache::builder(64) + .weigher(|_key: &String, desc: &NodeValueReplicaDesc| desc.weight_bytes) + .build(); + assert_eq!(cache.max_capacity(), Some(64)); + + for key in 0..128 { + insert_master_cache_entry( + &cache, + key.to_string(), + NodeValueReplicaDesc { + weight_bytes: 8, + put_id: (10, key), + }, + ); + } + cache.run_pending_tasks(); + assert!(cache.weighted_size() <= 64); + } + + #[test] + fn ring_b_live_capacity_update_never_clears_boundary() { + let cache = MasterNodeCache::builder(1024) + .weigher(|_key: &String, desc: &NodeValueReplicaDesc| desc.weight_bytes) + .build(); + assert_eq!(cache.max_capacity(), Some(1024)); + cache.set_max_capacity(512).unwrap(); + assert_eq!(cache.max_capacity(), Some(512)); + cache.set_max_capacity(2048).unwrap(); + assert_eq!(cache.max_capacity(), Some(2048)); + } + + #[test] + fn ring_b_shrink_selects_single_entries_to_the_new_active_boundary() { + let selected = Arc::new(AtomicU64::new(0)); + let selected_for_listener = selected.clone(); + let cache = MasterNodeCache::builder(1024) + .weigher(|_key: &String, desc: &NodeValueReplicaDesc| desc.weight_bytes) + .eviction_listener(move |_key, _desc, cause| { + assert_eq!(cause, RemovalCause::Size); + selected_for_listener.fetch_add(1, Ordering::Relaxed); + }) + .build(); + for key in 0..8 { + insert_master_cache_entry( + &cache, + key.to_string(), + NodeValueReplicaDesc { + weight_bytes: 128, + put_id: (10, key), + }, + ); + } + cache.run_pending_tasks(); + assert_eq!(cache.weighted_size(), 1024); + + cache.set_max_capacity(512).unwrap(); + cache.run_pending_tasks(); + assert!(cache.weighted_size() <= 512); + assert_eq!(selected.load(Ordering::Relaxed), 4); + } + + #[test] + fn node_pool_capacity_control_json_is_a_finite_generation_fenced_contract() { + let request = NodePoolCapacityControlRequest::SetActive { + owner_node_id: "cpu-owner".to_string(), + expected_owner_node_start_time: 17, + expected_capacity_epoch: 3, + active_capacity_bytes: 240 * 1024 * 1024 * 1024, + }; + let encoded = serde_json::to_vec(&request).unwrap(); + let decoded: NodePoolCapacityControlRequest = serde_json::from_slice(&encoded).unwrap(); + assert_eq!(decoded, request); + let json: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + assert_eq!(json["operation"], "set_active"); + assert_eq!(json["expected_capacity_epoch"], 3); + } + + #[test] + fn synchronous_restore_re_eviction_keeps_new_pending_weight() { + let weight = 8 * 1024 * 1024; + let pending = AtomicU64::new(weight); + + // A synchronous Moka Size listener enqueues the restored entry again before the old + // request completes. Completing the old request must leave exactly the new request. + pending.fetch_add(weight, Ordering::AcqRel); + subtract_pending_eviction_weight(&pending, "owner", weight); + assert_eq!(pending.load(Ordering::Acquire), weight); + + subtract_pending_eviction_weight(&pending, "owner", weight); + assert_eq!(pending.load(Ordering::Acquire), 0); + } + + #[test] + fn multi_identity_registration_rolls_back_partial_insert() { + let request = |keys: &[&str]| reclaim::EvictionReclaimRequest { + owner_node_id: "cpu0".to_string(), + owner_node_start_time: None, + members: keys + .iter() + .enumerate() + .map(|(index, key)| reclaim::EvictionReclaimMember { + key: (*key).to_string(), + desc: NodeValueReplicaDesc { + weight_bytes: 1024, + put_id: (7, index as u32), + }, + expected_backing: None, + }) + .collect(), + origin: reclaim::EvictionReclaimOrigin::MasterAllocationCapacity, + retry_count: 0, + }; + let inflight = DashSet::new(); + let first = request(&["a", "b"]); + assert!(try_install_eviction_reclaim_identities( + &inflight, + first.identities(), + )); + assert_eq!(inflight.len(), 2); + assert_eq!( + classify_existing_eviction_reclaim(&inflight, &first.identities()), + reclaim::EnqueueEvictionReclaimResult::AlreadyInProgress + ); + + // `c` is tentatively installed before the collision on `a`; failure + // must roll it back and preserve only the original request. + let overlapping = reclaim::EvictionReclaimRequest { + members: vec![ + reclaim::EvictionReclaimMember { + key: "c".to_string(), + desc: NodeValueReplicaDesc { + weight_bytes: 1024, + put_id: (7, 2), + }, + expected_backing: None, + }, + first.members[0].clone(), + ], + ..request(&[]) + }; + let tentative_identity = overlapping.identities()[0].clone(); + assert!(!try_install_eviction_reclaim_identities( + &inflight, + overlapping.identities(), + )); + assert_eq!(inflight.len(), 2); + assert!(!inflight.contains(&tentative_identity)); + assert_eq!( + classify_existing_eviction_reclaim(&inflight, &overlapping.identities()), + reclaim::EnqueueEvictionReclaimResult::PartialOverlap + ); + + for identity in first.identities() { + assert!(inflight.remove(&identity).is_some()); + } + assert_eq!( + classify_existing_eviction_reclaim(&inflight, &first.identities()), + reclaim::EnqueueEvictionReclaimResult::NotInProgress + ); + } + + #[test] + fn eviction_reclaim_metadata_channel_does_not_drop_at_old_queue_limit() { + let (tx, mut rx) = ampsc::unbounded_channel(); + let count = 4096 * 2 + 1; + for index in 0..count { + tx.send(reclaim::EvictionReclaimRequest { + owner_node_id: "cpu0".to_string(), + owner_node_start_time: None, + members: vec![reclaim::EvictionReclaimMember { + key: format!("key-{index}"), + desc: NodeValueReplicaDesc { + weight_bytes: 4096, + put_id: (7, index as u32), + }, + expected_backing: None, + }], + origin: reclaim::EvictionReclaimOrigin::MasterAllocationCapacity, + retry_count: 0, + }) + .unwrap(); + } + assert_eq!((0..count).filter(|_| rx.try_recv().is_ok()).count(), count); + } + + #[test] + fn closed_lossless_channel_rolls_back_identity_accounting_and_metadata() { + let cache = MasterNodeCache::builder(1024 * 1024).build(); + let request = reclaim::EvictionReclaimRequest { + owner_node_id: "cpu0".to_string(), + owner_node_start_time: None, + members: vec![reclaim::EvictionReclaimMember { + key: "closed-channel".to_string(), + desc: NodeValueReplicaDesc { + weight_bytes: 4096, + put_id: (7, 1), + }, + expected_backing: None, + }], + origin: reclaim::EvictionReclaimOrigin::MasterAllocationCapacity, + retry_count: 0, + }; + let inflight = DashSet::new(); + assert!(try_install_eviction_reclaim_identities( + &inflight, + request.identities(), + )); + let pending = AtomicU64::new(request.weight_bytes()); + let (tx, rx) = ampsc::unbounded_channel(); + drop(rx); + + let returned = tx.send(request).unwrap_err().0; + for identity in returned.identities() { + assert!(inflight.remove(&identity).is_some()); + } + subtract_pending_eviction_weight(&pending, "cpu0", returned.weight_bytes()); + for member in returned.members { + insert_master_cache_entry(&cache, member.key, member.desc); + } + + assert!(inflight.is_empty()); + assert_eq!(pending.load(Ordering::Acquire), 0); + assert!(cache.get(&"closed-channel".to_string()).is_some()); + } + + fn new_test_member(metadata: HashMap) -> ClusterMember { + ClusterMember { + id: "node-a".to_string(), + addresses: Vec::new(), + port: None, + node_start_time: 1, + metadata, + sub_cluster: Some("owner".to_string()), + network: None, + } + } + + #[test] + fn segment_registration_readiness_accepts_owner_without_local_ipc_root() { + let member = new_test_member(HashMap::from([ + ("client".to_string(), "true".to_string()), + ("p2p_relay".to_string(), "true".to_string()), + ])); + + assert!(MasterKvRouter::member_ready_for_segment_registration( + &member + )); + } + + #[test] + fn segment_registration_readiness_rejects_external_and_side_worker() { + let external = new_test_member(HashMap::from([( + "external_client".to_string(), + "true".to_string(), + )])); + assert!(!MasterKvRouter::member_ready_for_segment_registration( + &external + )); + + let side_worker = new_test_member(HashMap::from([ + ("client".to_string(), "true".to_string()), + ("side_transfer_worker".to_string(), "true".to_string()), + ("p2p_relay".to_string(), "true".to_string()), + ])); + assert!(!MasterKvRouter::member_ready_for_segment_registration( + &side_worker + )); + } + + #[test] + fn tier1_source_accepts_gpu_owner_role_and_rejects_remote_cache() { + let placement_config = ReplicaTaskPlacementConfig::default(); + + // This is the production GPU-owner shape: it deliberately does not match the + // prefill/decode roles carried by the associated zero-contribution client. + let gpu_owner = new_test_member(HashMap::from([ + ("client".to_string(), "true".to_string()), + ("p2p_relay".to_string(), "true".to_string()), + ])); + assert!(MasterKvRouter::tier1_source_member_eligible( + Some(&gpu_owner), + &placement_config, + )); + + let mut remote_cache = gpu_owner.clone(); + remote_cache.sub_cluster = Some("remote_cache".to_string()); + assert!(!MasterKvRouter::tier1_source_member_eligible( + Some(&remote_cache), + &placement_config, + )); + assert!(!MasterKvRouter::tier1_source_member_eligible( + None, + &placement_config, + )); + } +} + +pub struct MasterKvRouterInner { + view: std::sync::OnceLock, + pub policy: Box, + test_spec_config: TestSpecConfig, + pub replica_task_placement: ReplicaTaskPlacementConfig, + replica_cache_capacity_ratio: f64, + replica_writeback_tier1_capacity_ratio: Option, + + /// (key, put_time_ms, put_version) -> inflight_put_info + pub inflight_puts: moka::future::Cache<(String, u64, u32), InflightPutInfo>, + pub inflight_replica_tasks: moka::future::Cache<(String, u64, u32), InflightReplicaTaskInfo>, + /// Idempotent terminal PutAppendDone results retained across response + /// loss and batch-to-individual fallback. + pub completed_replica_tasks: + moka::future::Cache<(String, u64, u32, u64), CompletedReplicaTaskInfo>, + /// Serializes Start/Done/Revoke for one replica operation identity. The + /// lock is per `(key, put_id)`, so unrelated remote writes never share a + /// queue or actor. + pub replica_operation_locks: AMapLock<(String, u64, u32)>, + pub inflight_gets: moka::future::Cache, + /// Target-free source plans. Bind atomically moves an entry into + /// `inflight_gets` after re-reading and validating the current route. + pub planned_gets: moka::future::Cache, + pub(crate) planned_get_counters: MasterPlannedGetCounters, + pub(crate) ssd_tier_counters: MasterSsdTierCounters, + /// Idempotent terminal GetDone results retained across response loss/retry. + pub completed_gets: moka::future::Cache, + pub get_done_locks: AMapLock, + pub(crate) key_activity: Arc, + pub(crate) prepared_get_requesters: Arc, + + /// Cache for holding get operations (owned, flattened by (node_id, holder_id)) + pub get_holding: MasterOwnerMemMgr, + + /// Counter for get_id + pub next_get_id: AtomicU64, + + /// Counter for holder_id + pub next_holder_id: AtomicU64, + + /// Counter for local reserve grant identifiers. + pub next_local_reserve_grant_id: AtomicU64, + + /// Counter for prepared put-key reservation identifiers. + pub next_prepared_put_key_reservation_id: AtomicU64, + + /// Counter for concrete replica append attempts. This is separate from + /// `put_id`: one KV generation may legitimately be copied remotely again + /// after its previous remote route is reclaimed. + pub next_replica_operation_id: AtomicU64, + + /// Counter for two-sided owner reclaim epochs. + pub next_owner_reclaim_epoch: AtomicU64, + + /// Latest version of key-value replicas + pub kv_routes: DashMap>, + + /// Interns recent multi-key put groups so all member routes share one descriptor. + put_atomic_groups: moka::sync::SegmentedCache<(String, u64, u32), Arc>, + + /// Grants reserved for owner-local hot-path staging. + pub local_reserve_grants: DashMap, + + /// Prepared key reservations for owner-local hot-path staging. + pub prepared_put_key_reservations: DashMap, + + /// Prefix-counting index for keys, used by CountPrefix RPC. + pub prefix_index: ARwLock, + + /// Support replicas: node_id -> key -> route_info + pub node_kv_cache_controller: DashMap>, + + /// Independent pre-writeback tier. Its Size eviction starts a replica task + /// while owner residency remains governed by the owner's local hot cache; + /// `node_kv_cache_controller` tracks ring B only. + pub node_writeback_tier1_controller: DashMap>, + + /// Per-node bytes reserved out of moka usable capacity. + /// The reservation is reason-grouped, but `total_bytes` is the authority + /// used to derive the effective moka max_capacity for the node. + pub node_cache_reserved_capacity: DashMap>, + + /// Moka weight already removed and queued for owner-side safe reclaim. + pub eviction_reclaim_pending_weight: DashMap>, + + /// Exact, versioned metadata identities currently owned by the lossless + /// eviction-reclaim pipeline. This bounds duplicate Size/victim events; + /// payload memory is not retained here. + pub(crate) eviction_reclaim_inflight: DashSet, + + /// Per-owner reclaim lifecycle counters. These distinguish transient holder/activity + /// deferrals from terminal route changes and bounded retry restoration. + pub(crate) eviction_reclaim_counters: DashMap>, + + /// Async admission gate in front of Moka's synchronous housekeeper lock. + /// Bounded maintenance can still process a batch of Size evictions; without + /// this gate, concurrent async tasks can park multiple Tokio workers on the + /// same blocking Moka lock. It is a scheduling gate, not a correctness or + /// request lock, and no cache/global scan occurs under it. + pub(crate) owner_cache_operation_locks: AMapLock, + + /// Serializes the active pool target, reservation counters, and their derived Moka limits. + /// The critical section contains only local metadata/housekeeping and is shared with + /// synchronous reservation Drop paths, so it intentionally uses a short parking_lot lock. + node_capacity_boundary_locks: DashMap>>, + + /// Historical final put placement decisions by target node. + pub put_target_decision_counts: DashMap>, + + /// Historical final put placement decisions by requester->target pair. + pub put_requester_target_decision_counts: DashMap>, + + /// Historical final put placement decisions grouped by placement mode. + pub put_placement_mode_counts: DashMap<&'static str, Arc>, + + /// Historical accepted replica task reservations by target node. + pub replica_task_target_counts: DashMap>, + /// PutAppendDone responses served from the terminal cache after an RPC + /// replay or batch fallback. + pub replica_done_terminal_replay_count: AtomicU64, + + /// Historical get source choices by requester->source pair. + pub get_requester_source_counts: DashMap>, + pub get_requester_source_bytes: DashMap>, + /// Source-tier/locality matrix for planned and ordinary Gets. This keeps + /// local SSD, remote SSD, local memory, and remote memory separable in + /// evidence instead of inferring them from per-owner aggregate loads. + pub get_source_class_counts: DashMap<&'static str, Arc>, + pub get_source_class_bytes: DashMap<&'static str, Arc>, + pub get_allocation_mode_counts: DashMap<&'static str, Arc>, + + /// Support replicas: key -> version_id + recent_key_versionid_allocator: moka::sync::SegmentedCache>, + + pub delete_broadcast: EnsureMemholderMgmtDeleteHandle, + post_route_maintenance_tx: ampsc::Sender, + post_route_maintenance_rx: Mutex>>, + eviction_reclaim_tx: ampsc::UnboundedSender, + eviction_reclaim_rx: Mutex>>, + tier1_writeback_tx: ampsc::Sender, + tier1_writeback_rx: Mutex>>, + tier1_writeback_dedupe: moka::sync::SegmentedCache<(String, u64, u32), ()>, + tier1_writeback_trigger_counts: DashMap>, + tier1_writeback_owner_accepted_counts: DashMap>, + tier1_writeback_failed_counts: DashMap>, +} + +impl MasterKvRouterInner { + fn view(&self) -> &MasterKvRouterView { + self.view.get().unwrap() + } +} + +pub struct MasterKvRouter(MasterKvRouterInner); + +#[async_trait] +impl LogicalModule for MasterKvRouter { + type View = MasterKvRouterView; + type NewArg = MasterKvRouterNewArg; + type Error = KvError; + + fn name(&self) -> &str { + "MasterKvRouter" + } + + fn attach_view(&self, view: Self::View) { + MasterKvRouter::attach_view(self, view); + } + + async fn shutdown(&self) -> Result<(), Self::Error> { + info!("Shutting down MasterKvRouter"); + // Send shutdown signal to delete broadcast task to flush and exit. + if let Err(e) = self + .0 + .delete_broadcast + .sender() + .send(crate::master_kv_router::delete::DeleteKeyInfo::Shutdown) + .await + { + warn!("Failed to send delete broadcast shutdown signal: {}", e); + } + Ok(()) + } +} + +impl MasterKvRouter { + fn member_ready_for_segment_registration( + member: &crate::cluster_manager::ClusterMember, + ) -> bool { + // Segment registration must follow owner role semantics, not local IPC capability. + // + // Causal chain: + // - owner/external topology is already encoded in cluster member metadata + // (`client`, `external_client`, `side_transfer_worker`, `p2p_relay`); + // - `disable_local_ipc=true` intentionally suppresses `local_ipc_root` so the planner + // does not create same-machine IPC lanes; + // - owner shared bundle publication still depends on segment registration; + // - therefore the registration gate must stay tied to "real owner client" identity and + // must not reuse `local_ipc_root` as a readiness proxy. + member.metadata.get("client").is_some_and(|v| v == "true") + && member + .metadata + .get("p2p_relay") + .is_some_and(|v| v == "true") + && !member + .metadata + .get("external_client") + .is_some_and(|v| v == "true") + && !member + .metadata + .get("side_transfer_worker") + .is_some_and(|v| v == "true") + } + + pub fn attach_view(&self, view: MasterKvRouterView) { + // The framework attaches a module's PostView exactly once at the init barrier. + // A second attach indicates a programming error. + self.0 + .view + .set(view) + .unwrap_or_else(|_| panic!("MasterKvRouter view attached twice")); + } + + pub async fn construct(arg: MasterKvRouterNewArg) -> Result { + let policy_impl = build_placement_policy(arg.replica_task_placement.clone()); + let inflight_put_ttl_seconds = if arg.test_spec_config.skip_put_end_commit { + INFLIGHT_PUT_TTL_SECONDS_SKIP_PUT_END_COMMIT + } else { + INFLIGHT_PUT_TTL_SECONDS + }; + let key_activity = Arc::new(MasterKeyActivityTable::default()); + let prepared_get_requesters = Arc::new(PreparedGetRequesterTable::default()); + let inflight_puts = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(inflight_put_ttl_seconds)) + .eviction_listener(|_put_id, inflight_info: InflightPutInfo, cause| { + if cause == RemovalCause::Expired { + inflight_info._activity_lease.release_now(); + if let Some(replica_target) = inflight_info.commit_info.replica_target.as_ref() + { + replica_target._activity_lease.release_now(); + } + } + }) + .build(); + let inflight_gets = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(60)) + .eviction_listener(|_get_id, inflight_info: InflightGetInfo, cause| { + if cause == RemovalCause::Expired { + inflight_info.release_durable_slot_if_needed(); + inflight_info._activity_lease.release_now(); + } + }) + .build(); + let planned_gets: moka::future::Cache = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(60)) + .build(); + let completed_gets = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(120)) + .build(); + let inflight_replica_tasks = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(60)) + .eviction_listener(|_put_id, inflight_info: InflightReplicaTaskInfo, cause| { + if cause == RemovalCause::Expired { + inflight_info._activity_lease.release_now(); + } + }) + .build(); + let completed_replica_tasks = moka::future::Cache::builder() + .time_to_live(Duration::from_secs(120)) + .build(); + // A synchronous Moka listener must perform only constant-size metadata + // work and must never block or drop an Allocation reclaim event. The + // versioned inflight set deduplicates this unbounded channel. + let (eviction_reclaim_tx, eviction_reclaim_rx) = ampsc::unbounded_channel(); + let (post_route_maintenance_tx, post_route_maintenance_rx) = + ampsc::channel(POST_ROUTE_MAINTENANCE_QUEUE_CAPACITY); + let (tier1_writeback_tx, tier1_writeback_rx) = + ampsc::channel(TIER1_WRITEBACK_QUEUE_CAPACITY); + let inner = MasterKvRouterInner { + view: std::sync::OnceLock::new(), + policy: policy_impl, + test_spec_config: arg.test_spec_config, + replica_task_placement: arg.replica_task_placement, + replica_cache_capacity_ratio: arg.replica_cache_capacity_ratio, + replica_writeback_tier1_capacity_ratio: arg.replica_writeback_tier1_capacity_ratio, + inflight_puts, + inflight_replica_tasks, + completed_replica_tasks, + replica_operation_locks: AMapLock::new(Duration::from_secs(10 * 60)), + inflight_gets, + planned_gets, + planned_get_counters: MasterPlannedGetCounters::default(), + ssd_tier_counters: MasterSsdTierCounters::default(), + completed_gets, + get_done_locks: AMapLock::new(Duration::from_secs(10 * 60)), + key_activity, + prepared_get_requesters, + get_holding: MasterOwnerMemMgr::default(), + next_get_id: AtomicU64::new(0), + next_holder_id: AtomicU64::new(0), + next_local_reserve_grant_id: AtomicU64::new(1), + next_prepared_put_key_reservation_id: AtomicU64::new(1), + next_replica_operation_id: AtomicU64::new(1), + next_owner_reclaim_epoch: AtomicU64::new(1), + kv_routes: DashMap::new(), + put_atomic_groups: moka::sync::SegmentedCache::builder(8) + .max_capacity(262_144) + .time_to_idle(Duration::from_secs(30 * 60)) + .build(), + local_reserve_grants: DashMap::new(), + prepared_put_key_reservations: DashMap::new(), + prefix_index: ARwLock::new(PrefixRadixTree::new()), + node_kv_cache_controller: DashMap::new(), + node_writeback_tier1_controller: DashMap::new(), + node_cache_reserved_capacity: DashMap::new(), + eviction_reclaim_pending_weight: DashMap::new(), + eviction_reclaim_inflight: DashSet::new(), + eviction_reclaim_counters: DashMap::new(), + owner_cache_operation_locks: AMapLock::new(Duration::from_secs(10 * 60)), + node_capacity_boundary_locks: DashMap::new(), + put_target_decision_counts: DashMap::new(), + put_requester_target_decision_counts: DashMap::new(), + put_placement_mode_counts: DashMap::new(), + replica_task_target_counts: DashMap::new(), + replica_done_terminal_replay_count: AtomicU64::new(0), + get_requester_source_counts: DashMap::new(), + get_requester_source_bytes: DashMap::new(), + get_source_class_counts: DashMap::new(), + get_source_class_bytes: DashMap::new(), + get_allocation_mode_counts: DashMap::new(), + recent_key_versionid_allocator: moka::sync::SegmentedCache::builder(8) + .time_to_idle(Duration::from_secs(5)) + .build(), + delete_broadcast: EnsureMemholderMgmtDeleteHandle::new( + MasterOwnerMemMgr::DELETE_SUBMIT_QUEUE_CAPACITY, + ), + post_route_maintenance_tx, + post_route_maintenance_rx: Mutex::new(Some(post_route_maintenance_rx)), + eviction_reclaim_tx, + eviction_reclaim_rx: Mutex::new(Some(eviction_reclaim_rx)), + tier1_writeback_tx, + tier1_writeback_rx: Mutex::new(Some(tier1_writeback_rx)), + tier1_writeback_dedupe: moka::sync::SegmentedCache::builder(8) + .time_to_live(Duration::from_secs(60)) + .build(), + tier1_writeback_trigger_counts: DashMap::new(), + tier1_writeback_owner_accepted_counts: DashMap::new(), + tier1_writeback_failed_counts: DashMap::new(), + }; + Ok(Self(inner)) + } + + pub async fn init2_for_init_dag(&self) -> Result<(), KvError> { + info!("MasterKvRouter init2_for_init_dag"); + self.register_rpc_handlers(); + self.register_rpc_callers(); + let view = self.0.view().clone(); + + self.spawn_cluster_listener(); + self.spawn_put_placement_reporter(); + self.spawn_runtime_observe_reporter(); + + let delete_broadcast_rx = self + .0 + .delete_broadcast + .take_rx() + .expect("delete_broadcast rx already taken, that's impossible"); + delete::spawn_delete_broadcast(view, delete_broadcast_rx); + if let Some(post_route_maintenance_rx) = self.0.post_route_maintenance_rx.lock().take() { + route_maintenance::spawn_post_route_maintenance_actor( + self.0.view().clone(), + post_route_maintenance_rx, + ); + } else { + warn!("post_route_maintenance_rx already taken for MasterKvRouter"); + } + if let Some(eviction_reclaim_rx) = self.0.eviction_reclaim_rx.lock().take() { + reclaim::spawn_eviction_reclaim_actor(self.0.view().clone(), eviction_reclaim_rx); + } else { + warn!("eviction_reclaim_rx already taken for MasterKvRouter"); + } + if let Some(tier1_writeback_rx) = self.0.tier1_writeback_rx.lock().take() { + tiered_writeback::spawn_tier1_writeback_actor( + self.0.view().clone(), + tier1_writeback_rx, + ); + } else { + warn!("tier1_writeback_rx already taken for MasterKvRouter"); + } + Ok(()) + } + + pub(crate) fn view(&self) -> &MasterKvRouterView { + self.inner().view() + } + + pub fn inner(&self) -> &MasterKvRouterInner { + &self.0 + } + + fn replica_cache_base_capacity(&self, node_space_size: u64) -> u64 { + node_cache_capacity_boundaries( + node_space_size, + self.inner().replica_cache_capacity_ratio, + self.inner().replica_writeback_tier1_capacity_ratio, + 0, + ) + .ring_b_bytes + } + + fn replica_cache_effective_capacity(&self, node_id: &str, node_space_size: u64) -> u64 { + let reserved_capacity = self + .inner() + .node_cache_reserved_capacity + .get(node_id) + .filter(|reserved| !reserved.generation.is_tomb()) + .map(|reserved| reserved.total_reserved_bytes()) + .unwrap_or(0); + node_cache_capacity_boundaries( + node_space_size, + self.inner().replica_cache_capacity_ratio, + self.inner().replica_writeback_tier1_capacity_ratio, + reserved_capacity, + ) + .ring_b_bytes + } + + fn node_capacity_boundary_lock(&self, node_id: &str) -> Arc> { + self.inner() + .node_capacity_boundary_locks + .entry(node_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .value() + .clone() + } + + /// Apply the current active pool size to every derived Moka boundary. + fn apply_node_cache_capacity_locked(&self, node_id: &str) -> KvResult<()> { + if !self.replica_cache_enabled() { + return Ok(()); + } + let node_space_size = self + .inner() + .view() + .master_seg_manager() + .get_node_active_space_size(node_id); + if node_space_size == 0 { + return Err(KvError::Api(ApiError::NodeNotFound { + desc: format!("{node_id} (no live active pool)"), + })); + } + let capacity = self.replica_cache_effective_capacity(node_id, node_space_size); + let cache = self.get_node_cache_controller(node_id).ok_or_else(|| { + KvError::Api(ApiError::NodeNotFound { + desc: format!("{node_id} (ring-B controller unavailable)"), + }) + })?; + cache.set_max_capacity(capacity).map_err(|err| { + KvError::Api(ApiError::Allocator { + detail: format!( + "failed to set ring-B capacity: node={} capacity={} err={}", + node_id, capacity, err + ), + }) + })?; + // Drive Size selection now. Physical release remains asynchronous and continues through + // the existing generation-fenced single-KV reclaim actor. + cache.run_pending_tasks(); + if let Some(cache) = self + .inner() + .node_writeback_tier1_controller + .get(node_id) + .map(|entry| entry.value().clone()) + { + let tier1_capacity = self + .writeback_tier1_base_capacity(node_space_size) + .unwrap_or(0); + cache.set_max_capacity(tier1_capacity).map_err(|err| { + KvError::Api(ApiError::Allocator { + detail: format!( + "failed to set tier1 capacity: node={} capacity={} err={}", + node_id, tier1_capacity, err + ), + }) + })?; + cache.run_pending_tasks(); + } + Ok(()) + } + + fn apply_node_cache_capacity(&self, node_id: &str) -> KvResult<()> { + let capacity_lock = self.node_capacity_boundary_lock(node_id); + let _capacity_guard = capacity_lock.lock(); + self.apply_node_cache_capacity_locked(node_id) + } + + /// Refresh controllers after registration or reservation changes. Runtime resize uses the + /// fallible helper directly so a control response never reports an unapplied target. + fn reconcile_node_cache_capacity(&self, node_id: &str) { + // Member discovery can precede segment registration. There is no pool boundary to + // reconcile during that normal startup window; registration success will call us again. + // Keep the fallible runtime-control path strict, but do not emit a false operational + // error for this transient internal reconcile. + if self + .inner() + .view() + .master_seg_manager() + .get_node_active_space_size(node_id) + == 0 + { + debug!( + "deferring node cache capacity refresh until pool registration: node={}", + node_id + ); + return; + } + if let Err(err) = self.apply_node_cache_capacity(node_id) { + error!( + "failed to refresh node cache capacity: node={} err={}", + node_id, err + ); + } + } + + fn writeback_tier1_base_capacity(&self, node_space_size: u64) -> Option { + node_cache_capacity_boundaries( + node_space_size, + self.inner().replica_cache_capacity_ratio, + self.inner().replica_writeback_tier1_capacity_ratio, + 0, + ) + .tier1_bytes + } + + fn node_pool_capacity_response( + &self, + owner_node_id: &str, + owner_node_start_time: i64, + pool: NodePoolCapacitySnapshot, + ) -> NodePoolCapacityControlResponse { + let ring_b_base_capacity_bytes = + self.replica_cache_base_capacity(pool.active_capacity_bytes); + let reserved_capacity_bytes = self + .inner() + .node_cache_reserved_capacity + .get(owner_node_id) + .filter(|reserved| !reserved.generation.is_tomb()) + .map(|reserved| reserved.total_reserved_bytes()) + .unwrap_or(0); + let derived_effective = ring_b_base_capacity_bytes.saturating_sub(reserved_capacity_bytes); + let cache = self + .inner() + .node_kv_cache_controller + .get(owner_node_id) + .map(|entry| entry.value().clone()); + let ring_b_effective_capacity_bytes = cache + .as_ref() + .and_then(|cache| cache.max_capacity()) + .unwrap_or(derived_effective); + let ring_b_weighted_bytes = cache + .as_ref() + .map(|cache| cache.weighted_size()) + .unwrap_or(0); + let ring_b_pending_reclaim_bytes = self.eviction_reclaim_pending_weight(owner_node_id); + let tier1_capacity_bytes = self + .inner() + .node_writeback_tier1_controller + .get(owner_node_id) + .and_then(|entry| entry.value().max_capacity()) + .unwrap_or_else(|| { + self.writeback_tier1_base_capacity(pool.active_capacity_bytes) + .unwrap_or(0) + }); + let settled = pool.draining_capacity_bytes == 0 + && ring_b_weighted_bytes <= ring_b_effective_capacity_bytes + && ring_b_pending_reclaim_bytes == 0; + NodePoolCapacityControlResponse { + owner_node_id: owner_node_id.to_string(), + owner_node_start_time, + physical_capacity_bytes: pool.physical_capacity_bytes, + active_capacity_bytes: pool.active_capacity_bytes, + used_capacity_bytes: pool.used_capacity_bytes, + parked_capacity_bytes: pool.parked_capacity_bytes, + draining_capacity_bytes: pool.draining_capacity_bytes, + available_capacity_bytes: pool.available_capacity_bytes, + capacity_epoch: pool.capacity_epoch, + ring_b_base_capacity_bytes, + ring_b_effective_capacity_bytes, + ring_b_weighted_bytes, + ring_b_pending_reclaim_bytes, + tier1_capacity_bytes, + settled, + } + } + + pub fn node_pool_capacity_control_snapshot( + &self, + owner_node_id: &str, + ) -> KvResult { + let capacity_lock = self.node_capacity_boundary_lock(owner_node_id); + let _capacity_guard = capacity_lock.lock(); + let (owner_node_start_time, pool) = self + .inner() + .view() + .master_seg_manager() + .get_node_pool_capacity(owner_node_id) + .ok_or_else(|| { + KvError::Api(ApiError::NodeNotFound { + desc: owner_node_id.to_string(), + }) + })?; + Ok(self.node_pool_capacity_response(owner_node_id, owner_node_start_time, pool)) + } + + /// Atomically close the allocator admission gate before lowering Moka's derived boundaries. + /// The returned response is an initiation snapshot; callers poll `Get` until `settled=true`. + pub async fn set_node_pool_active_capacity( + &self, + owner_node_id: NodeID, + expected_owner_node_start_time: i64, + expected_capacity_epoch: u64, + active_capacity_bytes: u64, + ) -> KvResult { + let owner_cache_lock = self + .inner() + .owner_cache_operation_locks + .get_lock(owner_node_id.to_string()); + let _owner_cache_guard = owner_cache_lock.lock().await; + let capacity_lock = self.node_capacity_boundary_lock(owner_node_id.as_ref()); + let _capacity_guard = capacity_lock.lock(); + + let (_, before) = self + .inner() + .view() + .master_seg_manager() + .get_node_pool_capacity(owner_node_id.as_ref()) + .ok_or_else(|| { + KvError::Api(ApiError::NodeNotFound { + desc: owner_node_id.to_string(), + }) + })?; + let reserved_capacity_bytes = self + .inner() + .node_cache_reserved_capacity + .get(owner_node_id.as_ref()) + .filter(|reserved| !reserved.generation.is_tomb()) + .map(|reserved| reserved.total_reserved_bytes()) + .unwrap_or(0); + let prospective_ring_b = node_cache_capacity_boundaries( + active_capacity_bytes, + self.inner().replica_cache_capacity_ratio, + self.inner().replica_writeback_tier1_capacity_ratio, + reserved_capacity_bytes, + ) + .ring_b_bytes; + if self.replica_cache_enabled() && prospective_ring_b == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "active capacity leaves no ring-B allocation budget: owner={} active={} reserved={}", + owner_node_id, active_capacity_bytes, reserved_capacity_bytes + ), + })); + } + + let updated = self + .inner() + .view() + .master_seg_manager() + .set_node_active_capacity( + &owner_node_id, + expected_owner_node_start_time, + expected_capacity_epoch, + active_capacity_bytes, + )?; + if let Err(apply_err) = self.apply_node_cache_capacity_locked(owner_node_id.as_ref()) { + let rollback = self + .inner() + .view() + .master_seg_manager() + .set_node_active_capacity( + &owner_node_id, + expected_owner_node_start_time, + updated.capacity_epoch, + before.active_capacity_bytes, + ) + .and_then(|_| self.apply_node_cache_capacity_locked(owner_node_id.as_ref())); + error!( + owner = %owner_node_id, + requested_active_capacity_bytes = active_capacity_bytes, + error = %apply_err, + rollback = ?rollback, + "node pool capacity apply failed" + ); + return Err(apply_err); + } + + let response = self.node_pool_capacity_response( + owner_node_id.as_ref(), + expected_owner_node_start_time, + updated, + ); + info!( + owner = %owner_node_id, + owner_node_start_time = response.owner_node_start_time, + capacity_epoch = response.capacity_epoch, + physical_capacity_bytes = response.physical_capacity_bytes, + active_capacity_bytes = response.active_capacity_bytes, + parked_capacity_bytes = response.parked_capacity_bytes, + used_capacity_bytes = response.used_capacity_bytes, + draining_capacity_bytes = response.draining_capacity_bytes, + ring_b_effective_capacity_bytes = response.ring_b_effective_capacity_bytes, + ring_b_weighted_bytes = response.ring_b_weighted_bytes, + ring_b_pending_reclaim_bytes = response.ring_b_pending_reclaim_bytes, + settled = response.settled, + "node pool active capacity applied" + ); + Ok(response) + } + + pub fn tiered_writeback_enabled(&self) -> bool { + self.inner() + .replica_writeback_tier1_capacity_ratio + .is_some() + && self.replica_cache_enabled() + } + + pub fn replica_cache_enabled(&self) -> bool { + !self.0.test_spec_config.disable_master_replica_cache + } + + pub fn prefix_index_enabled(&self) -> bool { + !self.0.test_spec_config.disable_prefix_index + } + + /// return (put_time_ms, put_version) + pub fn get_recent_key_versionid(&self, key: String) -> (u64, u32) { + let put_time_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time went backwards") + .as_millis() as u64; + let put_version = self + .inner() + .recent_key_versionid_allocator + .get_with(key, || Arc::new(AtomicU32::new(0))) + .fetch_add(1, Ordering::Relaxed); + (put_time_ms, put_version) + } + + pub(crate) fn reserve_inflight_put_key( + &self, + key: &str, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + ) -> Result, KvError> { + if reject_if_exist_same_key && self.key_has_live_replica(key) { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyAlreadyExists { + key: key.to_string(), + }, + )); + } + let lease = self + .inner() + .key_activity + .reserve(key, MasterKeyActivityKind::Put, reject_if_inflight_same_key) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { + key: key.to_string(), + }, + ) + })?; + self.pin_current_master_cache_entries_for_activity(&lease, key); + Ok(lease) + } + + pub(crate) fn reserve_inflight_get_key( + &self, + key: &str, + ) -> Result, KvError> { + let lease = self + .inner() + .key_activity + .reserve(key, MasterKeyActivityKind::Get, false) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyNotFound { + key: key.to_string(), + }, + ) + })?; + self.pin_current_master_cache_entries_for_activity(&lease, key); + Ok(lease) + } + + pub(crate) fn reserve_prepared_get_requester( + &self, + key: &str, + requester: &NodeID, + get_id: u64, + ) -> Result, KvError> { + self.inner() + .prepared_get_requesters + .reserve(key, requester, get_id) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { + key: key.to_string(), + }, + ) + }) + } + + pub(crate) fn reserve_inflight_replica_key( + &self, + key: &str, + ) -> Result, KvError> { + let lease = self + .inner() + .key_activity + .reserve(key, MasterKeyActivityKind::Replica, false) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::KeyBeingWritten { + key: key.to_string(), + }, + ) + })?; + self.pin_current_master_cache_entries_for_activity(&lease, key); + Ok(lease) + } + + fn attach_cache_pin_if_present( + &self, + lease: &Arc, + cache: &MasterNodeCache, + key: &str, + desc: &NodeValueReplicaDesc, + ) { + let alias = MasterPinAlias::new(key, desc.put_id); + if let Some(pin) = cache.try_pin_alias_if(alias, |entry| { + entry.put_id == desc.put_id && entry.weight_bytes == desc.weight_bytes + }) { + lease.attach_cache_pin(pin); + } + } + + pub(crate) fn pin_master_cache_identity_for_activity( + &self, + lease: &Arc, + owner_node_id: &str, + key: &str, + desc: &NodeValueReplicaDesc, + ) { + if let Some(cache) = self + .inner() + .node_kv_cache_controller + .get(owner_node_id) + .map(|entry| entry.value().clone()) + { + self.attach_cache_pin_if_present(lease, cache.as_ref(), key, desc); + } + if let Some(cache) = self + .inner() + .node_writeback_tier1_controller + .get(owner_node_id) + .map(|entry| entry.value().clone()) + { + self.attach_cache_pin_if_present(lease, cache.as_ref(), key, desc); + } + } + + pub(crate) fn pin_current_master_cache_identity_for_activity( + &self, + lease: &Arc, + owner_node_id: &str, + key: &str, + put_id: PutIDForAKey, + ) { + let desc = self.inner().kv_routes.get(key).and_then(|route| { + (route.put_id == put_id).then(|| ring_b_route_replica_desc(&route, owner_node_id))? + }); + if let Some(desc) = desc { + self.pin_master_cache_identity_for_activity(lease, owner_node_id, key, &desc); + } + } + + fn pin_current_master_cache_entries_for_activity( + &self, + lease: &Arc, + key: &str, + ) { + let entries = self + .inner() + .kv_routes + .get(key) + .map(|route| { + let put_id = route.put_id; + route + .node_replicas + .read() + .iter() + .filter_map(|(node_id, replicas)| { + let memory = (!replicas.tomb_tag.is_tomb()) + .then(|| replicas.memory.as_ref()) + .flatten()?; + Some(( + node_id.as_ref().to_string(), + NodeValueReplicaDesc { + weight_bytes: u32::try_from(memory.backing.capacity_bytes()) + .unwrap_or(u32::MAX), + put_id, + }, + )) + }) + .collect::>() + }) + .unwrap_or_default(); + for (node_id, desc) in entries { + self.pin_master_cache_identity_for_activity(lease, node_id.as_str(), key, &desc); + } + } + + pub fn key_has_live_replica(&self, key: &str) -> bool { + self.inner() + .kv_routes + .get(key) + .map(|one_kv_nodes_routes| { + one_kv_nodes_routes + .node_replicas + .read() + .values() + .any(|kv_info| !kv_info.tomb_tag.is_tomb()) + }) + .unwrap_or(false) + } + + pub(crate) fn resolve_put_atomic_group( + &self, + key: &str, + put_id: PutIDForAKey, + group: Option, + ) -> Result>, KvError> { + let Some(group) = group else { + return Ok(None); + }; + if !(2..=4096).contains(&group.members.len()) { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: format!( + "put atomic group must contain 2..=4096 members; got={}", + group.members.len() + ), + }, + )); + } + let mut keys = HashSet::with_capacity(group.members.len()); + let mut current_member_count = 0usize; + for member in &group.members { + if member.key.is_empty() || !keys.insert(member.key.as_str()) { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: "put atomic group member keys must be non-empty and unique" + .to_string(), + }, + )); + } + if member.key == key && member.put_id == put_id { + current_member_count += 1; + } + } + if current_member_count != 1 { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: format!( + "put atomic group must contain current member exactly once: key={} put_id=({},{}) count={}", + key, put_id.0, put_id.1, current_member_count + ), + }, + )); + } + + let anchor = group + .members + .first() + .expect("validated non-empty put atomic group"); + let cache_key = (anchor.key.clone(), anchor.put_id.0, anchor.put_id.1); + if let Some(existing) = self.inner().put_atomic_groups.get(&cache_key) { + if existing.as_ref() != &group { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: format!( + "put atomic group anchor was reused with different membership: anchor_key={} anchor_put_id=({},{})", + anchor.key, anchor.put_id.0, anchor.put_id.1 + ), + }, + )); + } + return Ok(Some(existing)); + } + let group = Arc::new(group); + self.inner() + .put_atomic_groups + .insert(cache_key, group.clone()); + Ok(Some(group)) + } + + pub(crate) fn next_owner_reclaim_epoch(&self) -> u64 { + self.inner() + .next_owner_reclaim_epoch + .fetch_add(1, Ordering::Relaxed) + } + + fn enqueue_eviction_reclaim( + &self, + owner_node_id: NodeIDString, + key: String, + desc: NodeValueReplicaDesc, + origin: reclaim::EvictionReclaimOrigin, + ) -> bool { + self.enqueue_eviction_reclaim_request( + owner_node_id, + vec![reclaim::EvictionReclaimMember { + key, + desc, + expected_backing: None, + }], + origin, + ) + } + + fn enqueue_eviction_reclaim_request( + &self, + owner_node_id: NodeIDString, + members: Vec, + origin: reclaim::EvictionReclaimOrigin, + ) -> bool { + matches!( + self.enqueue_eviction_reclaim_request_exact( + owner_node_id, + None, + members, + origin, + true, + ), + reclaim::EnqueueEvictionReclaimResult::Accepted + | reclaim::EnqueueEvictionReclaimResult::AlreadyInProgress + ) + } + + pub(crate) fn enqueue_owner_capacity_eviction_victim( + &self, + owner_node_id: NodeIDString, + owner_node_start_time: i64, + victim: reclaim::EvictionReclaimMember, + ) -> reclaim::EnqueueEvictionReclaimResult { + self.enqueue_eviction_reclaim_request_exact( + owner_node_id, + Some(owner_node_start_time), + vec![victim], + reclaim::EvictionReclaimOrigin::OwnerCapacityEviction, + true, + ) + } + + fn enqueue_eviction_reclaim_request_exact( + &self, + owner_node_id: NodeIDString, + owner_node_start_time: Option, + members: Vec, + origin: reclaim::EvictionReclaimOrigin, + allow_new_request: bool, + ) -> reclaim::EnqueueEvictionReclaimResult { + if members.len() != 1 { + return reclaim::EnqueueEvictionReclaimResult::PartialOverlap; + } + let request = reclaim::EvictionReclaimRequest { + owner_node_id, + owner_node_start_time, + members, + origin, + retry_count: 0, + }; + if !allow_new_request { + let identities = request.identities(); + return classify_existing_eviction_reclaim( + &self.inner().eviction_reclaim_inflight, + &identities, + ); + } + if !self.register_eviction_reclaim(&request) { + self.eviction_reclaim_counters(&request.owner_node_id) + .eviction_reclaim_deduplicated + .fetch_add(1, Ordering::Relaxed); + // An idempotent retry is accepted only when the exact victim is + // already owned by the pipeline. + return if request + .identities() + .iter() + .all(|identity| self.inner().eviction_reclaim_inflight.contains(identity)) + { + reclaim::EnqueueEvictionReclaimResult::AlreadyInProgress + } else { + reclaim::EnqueueEvictionReclaimResult::PartialOverlap + }; + } + if let Err(err) = self.inner().eviction_reclaim_tx.send(request) { + let request = err.0; + self.complete_eviction_reclaim(&request); + if request.origin == reclaim::EvictionReclaimOrigin::MasterAllocationCapacity { + // This branch can run inside Moka's synchronous Size listener. + // Never re-enter that cache while its housekeeper lock is + // held; restore by key after yielding out of the callback. + let restore_view = self.inner().view().clone(); + let restore_request = request.clone(); + let spawn_view = restore_view.clone(); + let _ = spawn_view.spawn("closed_eviction_reclaim_restore", async move { + tokio::task::yield_now().await; + let mut restored = 0usize; + for member in &restore_request.members { + if restore_view.master_kv_router().restore_eviction_cache_entry_if_current( + &restore_request.owner_node_id, + member.key.clone(), + member.desc.clone(), + ) { + restored += 1; + } + } + warn!( + "restored master Allocation metadata after reclaim actor closed: owner={} members={} restored={}", + restore_request.owner_node_id, + restore_request.members.len(), + restored, + ); + }); + } + warn!( + "lossless eviction reclaim actor is closed: owner={} members={} origin={:?} restore_deferred={}", + request.owner_node_id, + request.members.len(), + request.origin, + request.origin == reclaim::EvictionReclaimOrigin::MasterAllocationCapacity, + ); + return reclaim::EnqueueEvictionReclaimResult::Closed; + } + reclaim::EnqueueEvictionReclaimResult::Accepted + } + + pub(crate) fn register_eviction_reclaim( + &self, + request: &reclaim::EvictionReclaimRequest, + ) -> bool { + if !try_install_eviction_reclaim_identities( + &self.inner().eviction_reclaim_inflight, + request.identities(), + ) { + return false; + } + let pending_weight = self + .inner() + .eviction_reclaim_pending_weight + .entry(request.owner_node_id.clone()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .clone(); + pending_weight.fetch_add(request.weight_bytes(), Ordering::AcqRel); + true + } + + pub(crate) fn eviction_reclaim_pending_weight(&self, owner_node_id: &str) -> u64 { + self.inner() + .eviction_reclaim_pending_weight + .get(owner_node_id) + .map(|weight| weight.load(Ordering::Acquire)) + .unwrap_or(0) + } + + pub(crate) fn eviction_reclaim_counters( + &self, + owner_node_id: &str, + ) -> Arc { + self.inner() + .eviction_reclaim_counters + .entry(owner_node_id.to_string()) + .or_insert_with(|| Arc::new(EvictionReclaimCounters::default())) + .value() + .clone() + } + + pub(crate) fn complete_eviction_reclaim(&self, request: &reclaim::EvictionReclaimRequest) { + for identity in request.identities() { + assert!( + self.inner() + .eviction_reclaim_inflight + .remove(&identity) + .is_some(), + "eviction reclaim identity completed without registration: {:?}", + identity, + ); + } + let completed_weight = request.weight_bytes(); + let pending_weight = self + .inner() + .eviction_reclaim_pending_weight + .get(&request.owner_node_id) + .unwrap_or_else(|| { + panic!( + "eviction reclaim pending weight missing for owner {}", + request.owner_node_id + ) + }); + subtract_pending_eviction_weight( + pending_weight.value().as_ref(), + &request.owner_node_id, + completed_weight, + ); + } + + pub(crate) fn restore_eviction_cache_entry_if_current( + &self, + owner_node_id: &str, + key: String, + desc: NodeValueReplicaDesc, + ) -> bool { + if !self.eviction_cache_entry_is_current(owner_node_id, &key, &desc) { + return false; + } + self.insert_node_cache_entry(owner_node_id, key, desc) + } + + pub(crate) fn eviction_cache_entry_is_current( + &self, + owner_node_id: &str, + key: &str, + desc: &NodeValueReplicaDesc, + ) -> bool { + self.inner().kv_routes.get(key).is_some_and(|route| { + ring_b_route_replica_desc(&route, owner_node_id).is_some_and(|current| { + current.put_id == desc.put_id && current.weight_bytes == desc.weight_bytes + }) + }) + } + + /// Point-remove one exact ring-B metadata identity. Callers use the + /// async per-node gate before entering Moka's synchronous housekeeper, so + /// route conversion never falls back to a cache scan. + pub(crate) async fn remove_node_cache_entry_exact( + &self, + owner_node_id: &str, + key: &str, + desc: &NodeValueReplicaDesc, + ) -> bool { + let owner_cache_lock = self + .inner() + .owner_cache_operation_locks + .get_lock(owner_node_id.to_string()); + let _owner_cache_guard = owner_cache_lock.lock().await; + self.inner() + .node_kv_cache_controller + .get(owner_node_id) + .is_some_and(|cache| remove_exact_cache_entry(cache.value(), key, desc)) + } + + /// Remove one superseded route version from both metadata policies using + /// only the route's own replica list. This is O(replica count), never a + /// Moka or global-route scan, and cannot delete a newer same-key version. + pub(crate) async fn remove_route_cache_entries_exact( + &self, + key: &str, + route: &OneKvNodesRoutes, + ) { + let replicas = route + .node_replicas + .read() + .iter() + .filter_map(|(node_id, replicas)| { + let memory = (!replicas.tomb_tag.is_tomb()) + .then(|| replicas.memory.as_ref()) + .flatten()?; + Some(( + node_id.as_ref().to_string(), + NodeValueReplicaDesc { + weight_bytes: u32::try_from(memory.backing.capacity_bytes()) + .unwrap_or(u32::MAX), + put_id: route.put_id, + }, + )) + }) + .collect::>(); + for (node_id, desc) in replicas { + let owner_cache_lock = self + .inner() + .owner_cache_operation_locks + .get_lock(node_id.clone()); + let _owner_cache_guard = owner_cache_lock.lock().await; + if let Some(cache) = self.inner().node_kv_cache_controller.get(&node_id) { + let _ = remove_exact_cache_entry(cache.value(), key, &desc); + } + if let Some(cache) = self.inner().node_writeback_tier1_controller.get(&node_id) { + let _ = remove_exact_cache_entry(cache.value(), key, &desc); + } + } + } + + pub(crate) fn insert_node_cache_entry( + &self, + owner_node_id: &str, + key: String, + desc: NodeValueReplicaDesc, + ) -> bool { + let Some(cache) = self.get_node_cache_controller(owner_node_id) else { + return false; + }; + // Restoration is an exact metadata repair after a failed reclaim + // dispatch. It must never search for or evict an unrelated victim. + insert_master_cache_entry(cache.as_ref(), key, desc); + true + } + + fn register_rpc_callers(&self) { + RPCCaller::::new().regist(self.0.view().p2p_module()); + RPCCaller::::new().regist(self.0.view().p2p_module()); + RPCCaller::::new().regist(self.0.view().p2p_module()); + } + + fn register_rpc_handlers(&self) { + let p2p = self.0.view().p2p_module(); + + user_rpc_register_handler_async( + p2p, + NODE_POOL_CAPACITY_USER_RPC_PATH.to_string(), + Arc::new(NodePoolCapacityUserRpcHandler { + view: self.0.view().clone(), + }), + ); + + // --- Get Handlers --- + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let cleanup_view = view.clone(); + let _ = view.spawn("rpc_get_start", async move { + let t0 = Utc::now().timestamp_micros(); + let (get_id, mut ack) = + handle_get_start(view_task, msg, resp.node_id().clone()).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + let get_started = ack.serialize_part.error_code + == crate::rpcresp_kvresult_convert::msg_and_error::OK; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send GetStartResp: {:?}", e); + if get_started { + if let Some(inflight_info) = cleanup_view + .master_kv_router() + .inner() + .inflight_gets + .remove(&get_id) + .await + { + inflight_info.release_durable_slot_if_needed(); + inflight_info._activity_lease.release_now(); + } + } + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let req_node_id = resp.node_id().clone(); + let _ = view.spawn("rpc_get_revoke", async move { + let ack = handle_get_revoke(view_task, msg, req_node_id).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send GetRevokeResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let req_node_id = resp.node_id().clone(); + let _ = view.spawn("rpc_get_done", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_get_done(view_task, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send GetDoneResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let task_view = view.clone(); + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_ssd_stage_begin", async move { + let ack = handle_ssd_stage_begin(task_view, msg, req_node_id).await; + let _ = resp.send_resp(ack).await; + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let task_view = view.clone(); + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_ssd_stage_done", async move { + let ack = handle_ssd_stage_done(task_view, msg, req_node_id).await; + let _ = resp.send_resp(ack).await; + }); + Ok(()) + }); + + // --- CountPrefix Handler --- + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_count_prefix", async move { + let ack = handle_count_prefix(&view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send CountPrefixResp: {:?}", e); + } + }); + Ok(()) + }); + + // --- GetMasterOnlyMetricPart Handler (metrics module registers) --- + crate::metrics::datasource::register_master_only_metric_handler(self.0.view()); + + // --- Put Handlers --- + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_reserve_local_grant", async move { + let t0 = Utc::now().timestamp_micros(); + let (grant_id, mut ack) = + handle_reserve_local_grant(view_task.clone(), msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send ReserveLocalGrantResp: {:?}", e); + if grant_id != 0 { + let _ = view_task + .master_kv_router() + .take_local_reserve_grant(grant_id); + } + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_release_local_grant", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_release_local_grant(view_task, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send ReleaseLocalGrantResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_batch_prepare_put_keys", async move { + let t0 = Utc::now().timestamp_micros(); + let (reservation_ids, mut ack) = + handle_batch_prepare_put_keys(view_task.clone(), msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchPreparePutKeysResp: {:?}", e); + for reservation_id in reservation_ids { + let _ = view_task + .master_kv_router() + .take_prepared_put_key_reservation(reservation_id); + } + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_batch_release_put_key_reservations", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = + handle_batch_release_put_key_reservations(view_task, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchReleasePutKeyReservationsResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_put_start", async move { + let key = msg.serialize_part.key.clone(); + #[cfg(feature = "test_bins")] + tracing::info!( + "rpc_put_start handler begin: self={} peer={} task_id={} key={} len={}", + view_task.cluster_manager().get_self_info().id, + req_node_id, + resp.task_id(), + key, + msg.serialize_part.len + ); + let t0 = Utc::now().timestamp_micros(); + let (put_id, mut ack) = handle_put_start(view_task.clone(), msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send PutStartResp: {:?}", e); + if put_id != (0, 0) { + if let Some(inflight_info) = view_task + .master_kv_router() + .inner() + .inflight_puts + .remove(&(key, put_id.0, put_id.1)) + .await + { + inflight_info._activity_lease.release_now(); + if let Some(replica_target) = + inflight_info.commit_info.replica_target.as_ref() + { + replica_target._activity_lease.release_now(); + } + } + } + } else { + #[cfg(feature = "test_bins")] + tracing::info!( + "rpc_put_start response sent: self={} peer={} task_id={} key={} put_id=({},{})", + view_task.cluster_manager().get_self_info().id, + resp.node_id(), + resp.task_id(), + key, + put_id.0, + put_id.1 + ); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_put_revoke", async move { + let ack = handle_put_revoke(view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send PutRevokeResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let req_node_id = resp.node_id().clone(); + let _ = view.spawn("rpc_put_done", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_put_done(view_task, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send PutDoneResp: {:?}", e); + } + }); + Ok(()) + }); + + // --- MemHolder Handlers --- + // let view = inner.view.clone(); + // RPCHandler::::new().regist(p2p, move |resp, msg| { + // let view = view.clone(); + // tokio::spawn(async move { + // let ack = handle_mem_holder_keep_alive(view, msg).await; + // if let Err(e) = resp.send_resp(ack).await { + // error!("Failed to send MemHolderKeepAliveResp: {}", e); + // } + // }); + // Ok(()) + // }); + + // let view = inner.view.clone(); + // RPCHandler::::new().regist(p2p, move |resp, msg| { + // let view = view.clone(); + // tokio::spawn(async move { + // let ack = handle_mem_holder_release(view, msg).await; + // if let Err(e) = resp.send_resp(ack).await { + // error!("Failed to send MemHolderReleaseResp: {}", e); + // } + // }); + // Ok(()) + // }); + + // --- Delete Handler --- + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_delete", async move { + let ack = handle_delete(view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send DeleteResp: {:?}", e); + } + }); + Ok(()) + }); + + // --- DeleteAck Handler --- + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + view.spawn("rpc_delete_ack", async move { + let ack = handle_delete_ack(view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send DeleteAckResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let view_task = view2.clone(); + view.spawn("rpc_batch_delete_ack", async move { + let ack = handle_batch_delete_ack(view_task, msg).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchDeleteAckResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let owner = resp.node_id().clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_batch_evict_owner_source", async move { + let ack = handle_batch_evict_owner_source(&view_task, msg, owner).await; + if let Err(err) = resp.send_resp(ack).await { + error!("Failed to send BatchEvictOwnerSourceResp: {:?}", err); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let owner = resp.node_id().clone(); + let view_task = view.clone(); + let _ = view.spawn("rpc_batch_publish_owner_ssd", async move { + let ack = handle_batch_publish_owner_ssd(view_task, msg, owner).await; + if let Err(err) = resp.send_resp(ack).await { + error!("Failed to send BatchPublishOwnerSsdResp: {:?}", err); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_put_append_start", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_put_append_start(view_task.clone(), msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send PutAppendStartResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + let view_task = view2.clone(); + let _ = view.spawn("rpc_batch_put_append_start", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_batch_put_append_start(view_task, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchPutAppendStartResp: {:?}", e); + } + }); + Ok(()) + }); + + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); let view_task = view2.clone(); - let cleanup_view = view.clone(); - let _ = view.spawn("rpc_get_start", async move { - let t0 = Utc::now().timestamp_micros(); - let (get_id, mut ack) = - handle_get_start(view_task, msg, resp.node_id().clone()).await; - ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + let _ = view.spawn("rpc_put_append_revoke", async move { + let ack = handle_put_append_revoke(view_task, msg).await; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send GetStartResp: {:?}", e); - if get_id != 0 { - if let Some(inflight_info) = cleanup_view - .master_kv_router() - .inner() - .inflight_gets - .remove(&get_id) - .await - { - inflight_info.release_durable_slot_if_needed(); - } - } + error!("Failed to send PutAppendRevokeResp: {:?}", e); } }); Ok(()) }); let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); let view_task = view2.clone(); - let _ = view.spawn("rpc_get_revoke", async move { - let ack = handle_get_revoke(view_task, msg).await; + let _ = view.spawn("rpc_put_append_done", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_put_append_done(view_task, msg).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send GetRevokeResp: {:?}", e); + error!("Failed to send PutAppendDoneResp: {:?}", e); } }); Ok(()) }); let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); let view_task = view2.clone(); - let _ = view.spawn("rpc_get_done", async move { + let _ = view.spawn("rpc_batch_put_append_done", async move { let t0 = Utc::now().timestamp_micros(); - let mut ack = handle_get_done(view_task, msg).await; + let mut ack = handle_batch_put_append_done(view_task, msg).await; ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send GetDoneResp: {:?}", e); + error!("Failed to send BatchPutAppendDoneResp: {:?}", e); } }); Ok(()) }); - // --- CountPrefix Handler --- + // --- GetMeta Handler --- let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); - let view_task = view.clone(); - let _ = view.spawn("rpc_count_prefix", async move { - let ack = handle_count_prefix(&view_task, msg).await; + let view2 = view.clone(); + view.spawn("rpc_get_meta", async move { + let ack = handle_get_meta(view2, msg, resp.node_id().clone()).await; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send CountPrefixResp: {:?}", e); + error!("Failed to send GetMetaResp: {:?}", e); } }); Ok(()) }); - // --- GetMasterOnlyMetricPart Handler (metrics module registers) --- - crate::metrics::datasource::register_master_only_metric_handler(self.0.view()); + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + view.spawn("rpc_batch_is_exist", async move { + let ack = handle_batch_is_exist(view2, msg, resp.node_id().clone()).await; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchIsExistResp: {:?}", e); + } + }); + Ok(()) + }); - // --- Put Handlers --- let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); let req_node_id = resp.node_id().clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_put_start", async move { - let key = msg.serialize_part.key.clone(); - #[cfg(feature = "test_bins")] - tracing::info!( - "rpc_put_start handler begin: self={} peer={} task_id={} key={} len={}", - view_task.cluster_manager().get_self_info().id, - req_node_id, - resp.task_id(), - key, - msg.serialize_part.len - ); + view.spawn("rpc_batch_put_start", async move { let t0 = Utc::now().timestamp_micros(); - let (put_id, mut ack) = handle_put_start(view_task.clone(), msg, req_node_id).await; + let mut ack = handle_batch_put_start(view2, msg, req_node_id).await; ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send PutStartResp: {:?}", e); - if put_id != (0, 0) { - view_task - .master_kv_router() - .inner() - .inflight_puts - .remove(&(key, put_id.0, put_id.1)) - .await; - } - } else { - #[cfg(feature = "test_bins")] - tracing::info!( - "rpc_put_start response sent: self={} peer={} task_id={} key={} put_id=({},{})", - view_task.cluster_manager().get_self_info().id, - resp.node_id(), - resp.task_id(), - key, - put_id.0, - put_id.1 - ); + error!("Failed to send BatchPutStartResp: {:?}", e); } }); Ok(()) }); let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_put_revoke", async move { - let ack = handle_put_revoke(view_task, msg).await; + view.spawn("rpc_batch_put_revoke", async move { + let ack = handle_batch_put_revoke(view2, msg).await; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send PutRevokeResp: {:?}", e); + error!("Failed to send BatchPutRevokeResp: {:?}", e); } }); Ok(()) }); let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_put_done", async move { + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_put_done", async move { let t0 = Utc::now().timestamp_micros(); - let mut ack = handle_put_done(view_task, msg).await; + let mut ack = handle_batch_put_done(view2, msg, req_node_id).await; ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send PutDoneResp: {:?}", e); + error!("Failed to send BatchPutDoneResp: {:?}", e); } }); Ok(()) }); - // --- MemHolder Handlers --- - // let view = inner.view.clone(); - // RPCHandler::::new().regist(p2p, move |resp, msg| { - // let view = view.clone(); - // tokio::spawn(async move { - // let ack = handle_mem_holder_keep_alive(view, msg).await; - // if let Err(e) = resp.send_resp(ack).await { - // error!("Failed to send MemHolderKeepAliveResp: {}", e); - // } - // }); - // Ok(()) - // }); + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_grouped_batch_put_done", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_grouped_batch_put_done(view2, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send GroupedBatchPutDoneResp: {:?}", e); + } + }); + Ok(()) + }); - // let view = inner.view.clone(); - // RPCHandler::::new().regist(p2p, move |resp, msg| { - // let view = view.clone(); - // tokio::spawn(async move { - // let ack = handle_mem_holder_release(view, msg).await; - // if let Err(e) = resp.send_resp(ack).await { - // error!("Failed to send MemHolderReleaseResp: {}", e); - // } - // }); - // Ok(()) - // }); + let view = self.0.view().clone(); + RPCHandler::::new().regist(p2p, move |resp, msg| { + let view = view.clone(); + let view2 = view.clone(); + let cleanup_view = view.clone(); + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_get_plan", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_batch_get_plan(view2, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + let planned_get_ids = ack + .serialize_part + .items + .iter() + .filter(|item| { + item.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK + }) + .map(|item| item.get_id) + .collect::>(); + if let Err(e) = resp.send_resp(ack).await { + error!("Failed to send BatchGetPlanResp: {:?}", e); + for get_id in planned_get_ids { + cleanup_view + .master_kv_router() + .inner() + .planned_gets + .remove(&get_id) + .await; + } + } + }); + Ok(()) + }); - // --- Delete Handler --- let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_delete", async move { - let ack = handle_delete(view_task, msg).await; + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_get_bind", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_batch_get_bind(view2, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send DeleteResp: {:?}", e); + error!("Failed to send BatchGetBindResp: {:?}", e); } }); Ok(()) }); - // --- DeleteAck Handler --- let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_delete_ack", async move { - let ack = handle_delete_ack(view_task, msg).await; + let cleanup_view = view.clone(); + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_get_start", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_batch_get_start(view2, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; + let started_get_ids = ack + .serialize_part + .items + .iter() + .filter(|item| { + item.error_code == crate::rpcresp_kvresult_convert::msg_and_error::OK + }) + .map(|item| item.get_id) + .collect::>(); if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send DeleteAckResp: {:?}", e); + error!("Failed to send BatchGetStartResp: {:?}", e); + for get_id in started_get_ids { + if let Some(inflight_info) = cleanup_view + .master_kv_router() + .inner() + .inflight_gets + .remove(&get_id) + .await + { + inflight_info.release_durable_slot_if_needed(); + inflight_info._activity_lease.release_now(); + } + } } }); Ok(()) }); let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let view_task = view2.clone(); - let _ = view.spawn("rpc_batch_delete_ack", async move { - let ack = handle_batch_delete_ack(view_task, msg).await; + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_get_revoke", async move { + let ack = handle_batch_get_revoke(view2, msg, req_node_id).await; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send BatchDeleteAckResp: {:?}", e); + error!("Failed to send BatchGetRevokeResp: {:?}", e); } }); Ok(()) }); - // --- GetMeta Handler --- let view = self.0.view().clone(); - RPCHandler::::new().regist(p2p, move |resp, msg| { + RPCHandler::::new().regist(p2p, move |resp, msg| { let view = view.clone(); let view2 = view.clone(); - let _ = view.spawn("rpc_get_meta", async move { - let ack = handle_get_meta(view2, msg, resp.node_id().clone()).await; + let req_node_id = resp.node_id().clone(); + view.spawn("rpc_batch_get_done", async move { + let t0 = Utc::now().timestamp_micros(); + let mut ack = handle_batch_get_done(view2, msg, req_node_id).await; + ack.serialize_part.server_process_us = Utc::now().timestamp_micros() - t0; if let Err(e) = resp.send_resp(ack).await { - error!("Failed to send GetMetaResp: {:?}", e); + error!("Failed to send BatchGetDoneResp: {:?}", e); } }); Ok(()) }); } + /// Start cleanup for one exact departed membership generation. + /// + /// Tomb publication and O(1) controller detachment happen synchronously in the cluster + /// listener. The potentially large route scan is then performed by a yielding async task. + fn begin_departed_generation_cleanup( + &self, + node_id: &str, + expected_node_start_time: Option, + ) -> bool { + let node: NodeID = node_id.to_string().into(); + let Some(departed_tag) = self + .inner() + .view() + .master_seg_manager() + .mark_node_tomb_generation(&node, expected_node_start_time) + else { + debug!( + "MemberLeft generation cleanup skipped because the registered generation changed or no segment was registered: node={} expected_node_start_time={:?}", + node_id, expected_node_start_time + ); + return false; + }; + + // Detach generation-scoped controllers before starting the full route scan. A tombed + // segment reports zero usable space, so old-generation work cannot recreate the caches. + let resident_cache_detached = self + .inner() + .node_kv_cache_controller + .remove(node_id) + .is_some(); + let tier1_cache_detached = self + .inner() + .node_writeback_tier1_controller + .remove(node_id) + .is_some(); + self.inner() + .node_cache_reserved_capacity + .remove_if(node_id, |_, reserved| { + reserved.generation.same_generation(&departed_tag) + }); + + let removed_holdings = self.inner().get_holding.cleanup_node(node_id); + let view = self.inner().view().clone(); + let node_id_owned = node_id.to_string(); + let _ = view.clone().spawn("member_left_route_cleanup", async move { + cleanup_departed_generation_routes(view, node_id_owned, departed_tag).await; + }); + + info!( + "MemberLeft generation marked and controllers detached: node={} expected_node_start_time={:?} resident_cache_detached={} tier1_cache_detached={} removed_holdings={}", + node_id, + expected_node_start_time, + resident_cache_detached, + tier1_cache_detached, + removed_holdings + ); + true + } + fn spawn_node_segement_registration_caller(&self) -> ampsc::Sender { const KEEP_ALIVE_TIME: Duration = Duration::from_secs(30); const NODE_EVENT_QUEUE_CAPACITY: usize = 64; let (tx, mut rx) = ampsc::channel::(NODE_EVENT_QUEUE_CAPACITY); let view = self.inner().view().clone(); let view_task = view.clone(); - let _ = view.spawn("node_segment_registration_caller", async move { + view.spawn("node_segment_registration_caller", async move { use std::future::Future; use std::pin::Pin; @@ -916,8 +4781,8 @@ impl MasterKvRouter { let epoch = member.node_start_time; if let Some(prev) = *last_seen_epoch { if prev != epoch { - view.master_seg_manager() - .mark_node_tomb(&node_id.clone().into()); + view.master_kv_router() + .begin_departed_generation_cleanup(&node_id, Some(prev)); } } *last_seen_epoch = Some(epoch); @@ -975,10 +4840,9 @@ impl MasterKvRouter { inflight = None; debug!( - "MasterKvRouter received node leave event: {:?}, mark it as tomb", + "MasterKvRouter registration actor canceled departed node: {:?}", node_id ); - view_task.master_seg_manager().mark_node_tomb(&node_id.into()); } } } @@ -1054,10 +4918,9 @@ impl MasterKvRouter { inflight = None; debug!( - "MasterKvRouter received node leave event: {:?}, mark it as tomb", + "MasterKvRouter registration actor canceled departed node: {:?}", node_id ); - view_task.master_seg_manager().mark_node_tomb(&node_id.into()); } } } @@ -1082,6 +4945,11 @@ impl MasterKvRouter { registered_epoch = Some(epoch); desired_epoch = None; backoff = INITIAL_BACKOFF; + if let Some(node_id) = actor_node_id.as_deref() { + view_task + .master_kv_router() + .reconcile_node_cache_capacity(node_id); + } info!( "Successfully requested segment registration from client {}", actor_node_id.clone().unwrap_or_default() @@ -1139,10 +5007,9 @@ impl MasterKvRouter { inflight = None; debug!( - "MasterKvRouter received node leave event: {:?}, mark it as tomb", + "MasterKvRouter registration actor canceled departed node: {:?}", node_id ); - view_task.master_seg_manager().mark_node_tomb(&node_id.into()); } } } @@ -1171,7 +5038,7 @@ impl MasterKvRouter { > = HashMap::new(); async fn send_event_with_warn( - _view: &MasterKvRouterView, + view: &MasterKvRouterView, node_id: &str, tx: ampsc::Sender, event: ClusterEvent, @@ -1208,17 +5075,66 @@ impl MasterKvRouter { event: ClusterEvent, ) { match &event { + ClusterEvent::MemberJoined(member) | ClusterEvent::MemberUpdated(member) => { + view.master_kv_router() + .reconcile_node_cache_capacity(&member.id); + } ClusterEvent::MemberLeft(node_id) => { - let removed = view + let departed_epoch = view + .cluster_manager() + .get_prev_member_info(node_id) + .map(|member| member.node_start_time); + let current_member = view + .cluster_manager() + .get_member_info_cached(node_id); + let current_epoch = current_member + .as_ref() + .map(|member| member.node_start_time); + + // MemberLeft has no epoch. Once a live generation is visible, this leave + // is ambiguous and must neither clean state nor reach the per-node actor: + // forwarding it would clear desired/registered_epoch for the reconnect. + if !member_left_can_forward_to_registration_actor(current_epoch) { + debug!( + "ignoring delayed MemberLeft after reconnect: node={} current_node_start_time={}", + node_id, + current_epoch.unwrap_or_default() + ); + return; + } + + let registered_tag_before = view + .master_seg_manager() + .get_node_tomb_tag(&node_id.clone().into()); + if !view .master_kv_router() - .inner() - .get_holding - .cleanup_node(&node_id); - if removed > 0 { - info!("Cleaned up {} holdings for left member {}", removed, node_id); + .begin_departed_generation_cleanup(node_id, departed_epoch) + { + // A registered segment exists but did not match `departed_epoch`: + // preserve it and do not forward the old leave to the actor. + if registered_tag_before.is_some() { + debug!( + "ignoring generation-mismatched MemberLeft: node={} departed_epoch={:?}", + node_id, departed_epoch + ); + return; + } + + // External/zero-contribution members do not register a segment or + // own route backing, but they can still own get holdings. + let removed = view + .master_kv_router() + .inner() + .get_holding + .cleanup_node(node_id); + if removed > 0 { + info!( + "Cleaned up {} holdings for segmentless left member {}", + removed, node_id + ); + } } } - _ => {} } let node_id = event.node_id(); @@ -1298,30 +5214,105 @@ impl MasterKvRouter { .or_insert_with(|| Arc::new(AtomicU64::new(0))) .value() .clone(); - target_counter.fetch_add(1, Ordering::Relaxed); + target_counter.fetch_add(1, Ordering::Relaxed); + + let requester_target_key = RequesterTargetPair::new(requester_node_id, target_node_id); + let requester_target_counter = self + .inner() + .put_requester_target_decision_counts + .entry(requester_target_key) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .clone(); + requester_target_counter.fetch_add(1, Ordering::Relaxed); + + let mode_key = match placement_mode { + PutPlacementMode::Local => "local", + PutPlacementMode::Remote => "remote", + }; + let mode_counter = self + .inner() + .put_placement_mode_counts + .entry(mode_key) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .clone(); + mode_counter.fetch_add(1, Ordering::Relaxed); + } + + pub fn record_replica_task_target(&self, target_node_id: &str) { + let target_counter = self + .inner() + .replica_task_target_counts + .entry(target_node_id.to_string()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .clone(); + target_counter.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn record_get_source_selection( + &self, + requester_node_id: &str, + source_node_id: &str, + bytes: u64, + allocation_mode: GetAllocationMode, + source_kind: GetSourceKind, + source_is_requester_local_owner: bool, + ) { + let requester_source_key = RequesterTargetPair::new(requester_node_id, source_node_id); + let source_count = self + .inner() + .get_requester_source_counts + .entry(requester_source_key.clone()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .clone(); + source_count.fetch_add(1, Ordering::Relaxed); - let requester_target_key = RequesterTargetPair::new(requester_node_id, target_node_id); - let requester_target_counter = self + let source_bytes = self .inner() - .put_requester_target_decision_counts - .entry(requester_target_key) + .get_requester_source_bytes + .entry(requester_source_key) .or_insert_with(|| Arc::new(AtomicU64::new(0))) .value() .clone(); - requester_target_counter.fetch_add(1, Ordering::Relaxed); + source_bytes.fetch_add(bytes, Ordering::Relaxed); - let mode_key = match placement_mode { - PutPlacementMode::Local => "local", - PutPlacementMode::Remote => "remote", + let source_class = match (source_kind, source_is_requester_local_owner) { + (GetSourceKind::Memory, true) => "local_memory", + (GetSourceKind::Memory, false) => "remote_memory", + (GetSourceKind::Ssd, true) => "local_ssd", + (GetSourceKind::Ssd, false) => "remote_ssd", }; - let mode_counter = self + self.inner() + .get_source_class_counts + .entry(source_class) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .fetch_add(1, Ordering::Relaxed); + self.inner() + .get_source_class_bytes + .entry(source_class) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .value() + .fetch_add(bytes, Ordering::Relaxed); + + let mode_key = match allocation_mode { + GetAllocationMode::Temporary => "temporary", + GetAllocationMode::ReuseReplica => "reuse_replica", + GetAllocationMode::DurableReplica => "durable_replica", + GetAllocationMode::LocalCommittedSlot => "local_committed_slot", + GetAllocationMode::ExternalSink => "external_sink", + }; + let mode_count = self .inner() - .put_placement_mode_counts + .get_allocation_mode_counts .entry(mode_key) .or_insert_with(|| Arc::new(AtomicU64::new(0))) .value() .clone(); - mode_counter.fetch_add(1, Ordering::Relaxed); + mode_count.fetch_add(1, Ordering::Relaxed); } fn spawn_put_placement_reporter(&self) { @@ -1361,11 +5352,70 @@ impl MasterKvRouter { .collect(); mode_counts.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + let mut replica_task_target_counts: Vec<(String, u64)> = router + .inner() + .replica_task_target_counts + .iter() + .map(|entry| (entry.key().clone(), entry.value().load(Ordering::Relaxed))) + .collect(); + replica_task_target_counts.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + let replica_done_terminal_replays = router + .inner() + .replica_done_terminal_replay_count + .load(Ordering::Relaxed); + + let mut get_source_counts: Vec<(String, u64)> = router + .inner() + .get_requester_source_counts + .iter() + .map(|entry| (entry.key().as_log_key(), entry.value().load(Ordering::Relaxed))) + .collect(); + get_source_counts.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + + let mut get_source_bytes: Vec<(String, u64)> = router + .inner() + .get_requester_source_bytes + .iter() + .map(|entry| (entry.key().as_log_key(), entry.value().load(Ordering::Relaxed))) + .collect(); + get_source_bytes.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + + let mut get_source_class_counts: Vec<(String, u64)> = router + .inner() + .get_source_class_counts + .iter() + .map(|entry| (entry.key().to_string(), entry.value().load(Ordering::Relaxed))) + .collect(); + get_source_class_counts.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + + let mut get_source_class_bytes: Vec<(String, u64)> = router + .inner() + .get_source_class_bytes + .iter() + .map(|entry| (entry.key().to_string(), entry.value().load(Ordering::Relaxed))) + .collect(); + get_source_class_bytes.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + + let mut get_allocation_mode_counts: Vec<(String, u64)> = router + .inner() + .get_allocation_mode_counts + .iter() + .map(|entry| (entry.key().to_string(), entry.value().load(Ordering::Relaxed))) + .collect(); + get_allocation_mode_counts.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + info!( - "put placement historical distribution | target_counts={:?} | mode_counts={:?} | requester_target_counts={:?}", + "placement historical distribution | put_target_counts={:?} | put_mode_counts={:?} | put_requester_target_counts={:?} | replica_task_target_counts={:?} | replica_done_terminal_replays={} | get_requester_source_counts={:?} | get_requester_source_bytes={:?} | get_source_class_counts={:?} | get_source_class_bytes={:?} | get_allocation_mode_counts={:?}", target_counts, mode_counts, requester_target_counts, + replica_task_target_counts, + replica_done_terminal_replays, + get_source_counts, + get_source_bytes, + get_source_class_counts, + get_source_class_bytes, + get_allocation_mode_counts, ); } _ = shutdown_waiter.wait() => { @@ -1376,10 +5426,7 @@ impl MasterKvRouter { }); } - pub fn get_node_cache_controller( - &self, - node_id: &str, - ) -> Option>> { + pub fn get_node_cache_controller(&self, node_id: &str) -> Option> { if !self.replica_cache_enabled() { return None; } @@ -1387,10 +5434,15 @@ impl MasterKvRouter { .inner() .view() .master_seg_manager() - .get_node_space_size(node_id); + .get_node_active_space_size(node_id); if node_space_size == 0 { return None; } + // Ring B is a backing/index domain, not a placement-role domain. A + // GPU owner can also hold master-only Allocation replicas, so every + // live segment gets the same bounded controller. Non-ring-B bytes + // are subtracted through generation-scoped reservation tokens. + let allocation_capacity = self.replica_cache_effective_capacity(node_id, node_space_size); let view = self.inner().view().clone(); let node_id_owned = node_id.to_string(); Some( @@ -1399,50 +5451,106 @@ impl MasterKvRouter { .entry(node_id_owned.clone()) .or_insert_with(move || { let view = view.clone(); + let cache_node_id = node_id_owned.clone(); + let builder = MasterNodeCache::builder(allocation_capacity) + // Admission is restricted to unindexed Allocation + // routes. CommittedSlot and owner-indexed Allocation + // entries belong to ring A and never enter this cache. + .weigher(|_key: &String, value: &NodeValueReplicaDesc| value.weight_bytes) + .eviction_listener( + move |key: Arc, + value: NodeValueReplicaDesc, + cause: RemovalCause| { + if cause == RemovalCause::Size { + // Listener work is deliberately O(1): + // clone fixed metadata, dedupe, and send + // to a lossless channel. Route/victim + // lookup belongs to the async actor. + let _ = view.master_kv_router().enqueue_eviction_reclaim( + cache_node_id.clone(), + (*key).clone(), + value, + reclaim::EvictionReclaimOrigin::MasterAllocationCapacity, + ); + } + }, + ); + Arc::new(builder.build()) + }) + .value() + .clone(), + ) + } + + pub fn tier1_source_node_eligible(&self, node_id: &str) -> bool { + if !self.tiered_writeback_enabled() { + return false; + } + let member = self + .inner() + .view() + .cluster_manager() + .get_member_info_cached(node_id); + Self::tier1_source_member_eligible(member.as_ref(), &self.inner().replica_task_placement) + } + + fn tier1_source_member_eligible( + member: Option<&crate::cluster_manager::ClusterMember>, + placement_config: &ReplicaTaskPlacementConfig, + ) -> bool { + // Route entries are keyed by the storage owner that holds the payload. In the + // external-SGLang topology that owner is tagged `sglang_owner`; the separate + // zero-contribution SGLang client carries the `prefill`/`decode` role. Requiring + // the owner itself to match active_node_roles therefore disables T1 entirely. + // + // A T1 source only needs to be a known, non-remote-only owner. The caller also + // requires a non-zero registered segment before constructing the tier cache. + member.is_some() + && !placement::member_matches_roles(member, &placement_config.remote_only_node_roles) + } + + pub fn get_node_writeback_tier1_controller( + &self, + node_id: &str, + ) -> Option> { + if !self.tier1_source_node_eligible(node_id) { + return None; + } + let node_space_size = self + .inner() + .view() + .master_seg_manager() + .get_node_active_space_size(node_id); + let base_capacity = self.writeback_tier1_base_capacity(node_space_size)?; + if base_capacity == 0 { + return None; + } + let view = self.inner().view().clone(); + let node_id_owned = node_id.to_string(); + Some( + self.inner() + .node_writeback_tier1_controller + .entry(node_id_owned.clone()) + .or_insert_with(move || { let cache_node_id = node_id_owned.clone(); Arc::new( - moka::sync::SegmentedCache::builder(8) - .max_capacity((node_space_size as f32 * MOKA_CACHE_CAPACITY_RATIO) as u64) - // Use the actual allocated/rounded size as weight to - // make eviction reflect real memory usage. - .weigher(Box::new(|_key: &String, value: &NodeValueReplicaDesc| { + MasterNodeCache::builder(base_capacity) + .weigher(|_key: &String, value: &NodeValueReplicaDesc| { value.weight_bytes - })) - .eviction_listener(Box::new( + }) + .eviction_listener( move |key: Arc, - _value: NodeValueReplicaDesc, + value: NodeValueReplicaDesc, cause: RemovalCause| { - debug!("Evicted key: {:?}, caused by: {:?}", key, cause); - match cause { - // timeout or size exceed - RemovalCause::Size | RemovalCause::Expired => { - let k = (*key).clone(); - let evicted_put_id = _value.put_id; - tracing::debug!( - "Eviction-triggered local replica cleanup for key {} on node {} put_id=({},{})", - k, - cache_node_id, - evicted_put_id.0, - evicted_put_id.1 - ); - if let Err(code) = crate::master_kv_router::delete::evict_one_kv_replica_for_node( - &view, - k.clone(), - cache_node_id.clone().into(), - evicted_put_id, - ) { - warn!( - "Eviction-triggered local replica cleanup failed for key {} on node {}: {:?}", - k, - cache_node_id, - code - ); - } - } - _ => {} + if cause == RemovalCause::Size { + view.master_kv_router().enqueue_tier1_writeback( + cache_node_id.clone(), + (*key).clone(), + value, + ); } }, - )) + ) .build(), ) }) @@ -1451,116 +5559,744 @@ impl MasterKvRouter { ) } - /// Atomically adjust a node's cache capacity reservation by `delta_bytes`. - /// Positive delta reserves capacity (fetch_sub from usable capacity), - /// negative delta releases reservation (fetch_add back to usable capacity). - pub fn adjust_node_cache_capacity_for_lease( + pub(crate) fn tier1_writeback_entry_is_current( + &self, + source_node_id: &str, + key: &str, + desc: &NodeValueReplicaDesc, + ) -> bool { + if self.inner().inflight_replica_tasks.contains_key(&( + key.to_string(), + desc.put_id.0, + desc.put_id.1, + )) { + return false; + } + self.inner().kv_routes.get(key).is_some_and(|route| { + if route.put_id != desc.put_id || route.lease_id.is_some() { + return false; + } + let replicas = route.node_replicas.read(); + replicas.iter().any(|(node_id, replica)| { + node_id.as_ref() == source_node_id + && !replica.tomb_tag.is_tomb() + && replica.memory.is_some() + }) && replicas.iter().all(|(node_id, replica)| { + !replica.has_live_backing() || node_id.as_ref() == source_node_id + }) + }) + } + + fn enqueue_tier1_writeback( + &self, + source_node_id: NodeIDString, + key: String, + desc: NodeValueReplicaDesc, + ) { + if !self.tier1_writeback_entry_is_current(&source_node_id, &key, &desc) { + return; + } + let dedupe_key = (key.clone(), desc.put_id.0, desc.put_id.1); + if self + .inner() + .tier1_writeback_dedupe + .get(&dedupe_key) + .is_some() + { + return; + } + self.inner() + .tier1_writeback_dedupe + .insert(dedupe_key.clone(), ()); + Self::increment_tier1_writeback_counter( + &self.inner().tier1_writeback_trigger_counts, + &source_node_id, + 1, + ); + let request = tiered_writeback::Tier1WritebackRequest { + source_node_id: source_node_id.clone(), + key, + desc, + }; + if let Err(err) = self.inner().tier1_writeback_tx.try_send(request) { + self.inner().tier1_writeback_dedupe.remove(&dedupe_key); + self.record_tier1_writeback_failed(&source_node_id, 1); + tracing::warn!("tier1 write-back queue is full or closed: {}", err); + } + } + + fn increment_tier1_writeback_counter( + counters: &DashMap>, + source_node_id: &str, + count: u64, + ) { + if count == 0 { + return; + } + counters + .entry(source_node_id.to_string()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .fetch_add(count, Ordering::Relaxed); + } + + pub(crate) fn record_tier1_writeback_owner_accepted(&self, source_node_id: &str, count: u64) { + Self::increment_tier1_writeback_counter( + &self.inner().tier1_writeback_owner_accepted_counts, + source_node_id, + count, + ); + } + + pub(crate) fn record_tier1_writeback_failed(&self, source_node_id: &str, count: u64) { + Self::increment_tier1_writeback_counter( + &self.inner().tier1_writeback_failed_counts, + source_node_id, + count, + ); + } + + fn tier1_writeback_counter( + counters: &DashMap>, + source_node_id: &str, + ) -> u64 { + counters + .get(source_node_id) + .map(|counter| counter.load(Ordering::Relaxed)) + .unwrap_or(0) + } + + pub(crate) fn finish_tier1_writeback_request( + &self, + request: tiered_writeback::Tier1WritebackRequest, + ) { + self.inner().tier1_writeback_dedupe.remove(&( + request.key, + request.desc.put_id.0, + request.desc.put_id.1, + )); + } + + /// Adjust one exact generation/counter identity and refresh its live cache + /// boundary. Negative releases always stay applied to the captured Arc, + /// even when the node has already left and its controller was detached. + fn adjust_node_cache_reserved_capacity_identity( &self, node_id: &str, + generation: &NodeTombTag, + reserved_capacity: &Arc, + reason: ReservedCapacityReason, delta_bytes: i64, ) -> crate::rpcresp_kvresult_convert::msg_and_error::KvResult<()> { if !self.replica_cache_enabled() { return Ok(()); } - // Track per-node reserved bytes with an atomic counter - let reserved_counter = self - .inner() - .lease_reserved_bytes - .entry(node_id.to_string()) - .or_insert_with(|| Arc::new(AtomicU64::new(0))) - .value() - .clone(); - - // Apply delta to the counter with simple fetch ops per user's preference - if delta_bytes >= 0 { - reserved_counter.fetch_add(delta_bytes as u64, Ordering::Relaxed); - } else { - let sub = (-delta_bytes) as u64; - reserved_counter.fetch_sub(sub, Ordering::Relaxed); + let capacity_lock = self.node_capacity_boundary_lock(node_id); + let _capacity_guard = capacity_lock.lock(); + if !reserved_capacity.generation.same_generation(generation) { + return Err( + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "cache reservation generation/counter identity mismatch: node_id={}", + node_id + ), + }, + ), + ); } - // Recompute target capacity: base(=MOKA_CACHE_CAPACITY_RATIO * node_space) - reserved_total - let reserved_total = reserved_counter.load(Ordering::Relaxed); - let node_space_size = self + reserved_capacity.apply_delta(reason, delta_bytes); + + let current_identity = self .inner() - .view() - .master_seg_manager() - .get_node_space_size(node_id); - if node_space_size == 0 { - // Node not ready: this should not happen in a successful put_done path. - // Revert the counter delta before returning error. + .node_cache_reserved_capacity + .get(node_id) + .is_some_and(|current| { + Arc::ptr_eq(current.value(), reserved_capacity) + && current.generation.same_generation(generation) + }) + && node_generation_is_current_live( + self.inner().view(), + &node_id.to_string().into(), + generation, + ); + if !current_identity { if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); - } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); + reserved_capacity.apply_delta(reason, -delta_bytes); + return Err( + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "cache reservation target generation changed: node_id={}", + node_id + ), + }, + ), + ); } + // The old generation has been detached. Its exact counter Arc is + // now balanced, and no live controller must be modified. + return Ok(()); + } + + if let Err(err) = self.apply_node_cache_capacity_locked(node_id) { + if delta_bytes >= 0 { + reserved_capacity.apply_delta(reason, -delta_bytes); + if let Err(rollback_err) = self.apply_node_cache_capacity_locked(node_id) { + error!( + node = node_id, + error = %rollback_err, + "failed to restore node cache boundary after reservation rollback" + ); + } + } + return Err(err); + } + Ok(()) + } + + /// Reserve cache capacity for one exact live node generation. The returned + /// route-lifetime token releases the same counter identity on Drop. + pub fn reserve_node_cache_capacity( + &self, + node_id: &NodeID, + generation: &NodeTombTag, + reason: ReservedCapacityReason, + bytes: u64, + ) -> crate::rpcresp_kvresult_convert::msg_and_error::KvResult< + Option>, + > { + if !self.replica_cache_enabled() { + return Ok(None); + } + if !node_generation_is_current_live(self.inner().view(), node_id, generation) { return Err( - crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidPutMasterState { detail: format!( - "node_id={} has no segment (node_space_size=0) while adjusting cache capacity", + "cannot reserve cache capacity for departed generation: node_id={}", node_id ), }, ), ); } - let base_capacity = (node_space_size as f32 * MOKA_CACHE_CAPACITY_RATIO) as u64; - let new_capacity = base_capacity.saturating_sub(reserved_total); - if let Some(cache) = self.get_node_cache_controller(node_id) { - if let Err(e) = cache.set_max_capacity(new_capacity) { - // Revert counter and return error. - if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); + let reserved_capacity = match self + .inner() + .node_cache_reserved_capacity + .entry(node_id.to_string()) + { + dashmap::mapref::entry::Entry::Occupied(mut entry) => { + if entry.get().generation.same_generation(generation) { + entry.get().clone() + } else if entry.get().generation.is_tomb() { + let replacement = Arc::new(NodeCacheReservedCapacity::new(generation.clone())); + entry.insert(replacement.clone()); + replacement } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); - } - return Err(crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::RpcDecodeError { - rpc_input_json: format!( - "moka.set_max_capacity failed: node_id={}, new_capacity={}, err={}", - node_id, new_capacity, e + return Err( + crate::rpcresp_kvresult_convert::msg_and_error::KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "live cache reservation counter belongs to another generation: node_id={}", + node_id + ), + }, ), - } - )); + ); + } } - Ok(()) - } else { - // Revert counter and return error. - if delta_bytes >= 0 { - reserved_counter.fetch_sub(delta_bytes as u64, Ordering::Relaxed); - } else { - let add = (-delta_bytes) as u64; - reserved_counter.fetch_add(add, Ordering::Relaxed); + dashmap::mapref::entry::Entry::Vacant(entry) => { + let counter = Arc::new(NodeCacheReservedCapacity::new(generation.clone())); + entry.insert(counter.clone()); + counter } - Err( - crate::rpcresp_kvresult_convert::msg_and_error::KvError::Unreachable( - crate::rpcresp_kvresult_convert::msg_and_error::UnreachableError::OwnerNoSeg { - detail: format!("node_id={} cache_controller not found", node_id), - }, - ), - ) + }; + + self.adjust_node_cache_reserved_capacity_identity( + node_id.as_ref(), + generation, + &reserved_capacity, + reason, + bytes as i64, + )?; + Ok(Some(Arc::new(NodeCacheCapacityReservation { + view: self.inner().view().clone(), + node_id: node_id.to_string(), + generation: generation.clone(), + reserved_capacity, + reason, + bytes, + released: AtomicBool::new(false), + }))) + } + + pub fn next_local_reserve_grant_id(&self) -> u64 { + self.inner() + .next_local_reserve_grant_id + .fetch_add(1, Ordering::Relaxed) + } + + pub fn next_prepared_put_key_reservation_id(&self) -> u64 { + self.inner() + .next_prepared_put_key_reservation_id + .fetch_add(1, Ordering::Relaxed) + } + + pub fn install_local_reserve_grant(&self, grant_id: u64, grant: LocalReserveGrantInfo) { + if self + .inner() + .local_reserve_grants + .insert(grant_id, grant) + .is_some() + { + panic!( + "duplicate local reserve grant id indicates a logic bug: grant_id={}", + grant_id + ); } } - // Note: no additional getters for reserved bytes; policy currently relies only on adjust calls. -} + pub fn take_local_reserve_grant(&self, grant_id: u64) -> Option { + self.inner() + .local_reserve_grants + .remove(&grant_id) + .map(|(_, grant)| grant) + } -impl MasterKvRouterView { - pub fn try_adjust_node_cache_capacity_for_lease( + pub fn install_prepared_put_key_reservation( &self, - node_id: &str, - delta_bytes: i64, - ) -> Option> { - let _view_guard = self.try_upgrade()?; - Some( - self.master_kv_router() - .adjust_node_cache_capacity_for_lease(node_id, delta_bytes), - ) + reservation_id: u64, + info: PreparedPutKeyReservationInfo, + ) { + if self + .inner() + .prepared_put_key_reservations + .insert(reservation_id, info) + .is_some() + { + panic!( + "duplicate prepared put key reservation id indicates a logic bug: reservation_id={}", + reservation_id + ); + } + } + + pub fn take_prepared_put_key_reservation( + &self, + reservation_id: u64, + ) -> Option { + self.inner() + .prepared_put_key_reservations + .remove(&reservation_id) + .map(|(_, info)| info) + } + + pub fn runtime_observe_snapshot(&self) -> MasterRuntimeObserveSnapshot { + let mut get_holding_bytes = 0u64; + for entry in self.inner().get_holding.inner().iter() { + get_holding_bytes = get_holding_bytes.saturating_add(entry.value().len); + } + + let mut replica_cache_nodes = Vec::new(); + for entry in self.inner().node_kv_cache_controller.iter() { + let owner_node = entry.key().clone(); + let cache = entry.value().clone(); + let Some((owner_node_start_time, pool_capacity)) = self + .inner() + .view() + .master_seg_manager() + .get_node_pool_capacity(owner_node.as_str()) + else { + continue; + }; + let node_space_size = pool_capacity.active_capacity_bytes; + let base_capacity_bytes = self.replica_cache_base_capacity(node_space_size); + let reserved_capacity_bytes = self + .inner() + .node_cache_reserved_capacity + .get(owner_node.as_str()) + .filter(|reserved| !reserved.generation.is_tomb()) + .map(|reserved| reserved.total_reserved_bytes()) + .unwrap_or(0); + // Every live node's ring-B controller is bounded, independent of + // its placement role. + let effective_capacity_bytes = cache + .max_capacity() + .expect("ring-B controller must always be bounded"); + let pending_eviction_reclaim_bytes = + self.eviction_reclaim_pending_weight(owner_node.as_str()); + let reclaim_counters = self.eviction_reclaim_counters(owner_node.as_str()); + let tier1_cache = self + .inner() + .node_writeback_tier1_controller + .get(owner_node.as_str()) + .map(|entry| entry.value().clone()); + replica_cache_nodes.push(ReplicaCacheNodeObserveSnapshot { + owner_node: owner_node.clone(), + owner_node_start_time, + pool_physical_capacity_bytes: pool_capacity.physical_capacity_bytes, + pool_active_capacity_bytes: pool_capacity.active_capacity_bytes, + pool_used_capacity_bytes: pool_capacity.used_capacity_bytes, + pool_parked_capacity_bytes: pool_capacity.parked_capacity_bytes, + pool_draining_capacity_bytes: pool_capacity.draining_capacity_bytes, + pool_available_capacity_bytes: pool_capacity.available_capacity_bytes, + pool_capacity_epoch: pool_capacity.capacity_epoch, + entries: cache.entry_count(), + weighted_bytes: cache.weighted_size(), + effective_capacity_bytes, + reserved_capacity_bytes, + base_capacity_bytes, + pending_eviction_reclaim_bytes, + writeback_tier1_entries: tier1_cache + .as_ref() + .map(|cache| cache.entry_count()) + .unwrap_or(0), + writeback_tier1_weighted_bytes: tier1_cache + .as_ref() + .map(|cache| cache.weighted_size()) + .unwrap_or(0), + writeback_tier1_capacity_bytes: tier1_cache + .as_ref() + .and_then(|cache| cache.max_capacity()) + .unwrap_or(0), + writeback_tier1_triggered: Self::tier1_writeback_counter( + &self.inner().tier1_writeback_trigger_counts, + owner_node.as_str(), + ), + writeback_tier1_owner_accepted: Self::tier1_writeback_counter( + &self.inner().tier1_writeback_owner_accepted_counts, + owner_node.as_str(), + ), + writeback_tier1_failed: Self::tier1_writeback_counter( + &self.inner().tier1_writeback_failed_counts, + owner_node.as_str(), + ), + reclaim_master_activity_deferred: reclaim_counters + .master_activity_deferred + .load(Ordering::Relaxed), + reclaim_owner_holder_deferred: reclaim_counters + .owner_holder_deferred + .load(Ordering::Relaxed), + reclaim_owner_other_deferred: reclaim_counters + .owner_other_deferred + .load(Ordering::Relaxed), + reclaim_route_changed: reclaim_counters.route_changed.load(Ordering::Relaxed), + reclaim_retry_queued: reclaim_counters.retry_queued.load(Ordering::Relaxed), + reclaim_retry_completed: reclaim_counters.retry_completed.load(Ordering::Relaxed), + reclaim_retry_restored: reclaim_counters.retry_restored.load(Ordering::Relaxed), + reclaim_completed: reclaim_counters.completed.load(Ordering::Relaxed), + source_evict_rpc_requests: reclaim_counters + .source_evict_rpc_requests + .load(Ordering::Relaxed), + source_evict_victims: reclaim_counters + .source_evict_victims + .load(Ordering::Relaxed), + source_evict_requested_bytes: reclaim_counters + .source_evict_requested_bytes + .load(Ordering::Relaxed), + source_evict_accepted: reclaim_counters + .source_evict_accepted + .load(Ordering::Relaxed), + source_evict_in_progress: reclaim_counters + .source_evict_in_progress + .load(Ordering::Relaxed), + source_evict_completed: reclaim_counters + .source_evict_completed + .load(Ordering::Relaxed), + source_evict_retryable_busy: reclaim_counters + .source_evict_retryable_busy + .load(Ordering::Relaxed), + source_evict_stale: reclaim_counters.source_evict_stale.load(Ordering::Relaxed), + source_evict_rejected: reclaim_counters + .source_evict_rejected + .load(Ordering::Relaxed), + last_route_removed_members: reclaim_counters + .last_route_removed_members + .load(Ordering::Relaxed), + last_route_removed_bytes: reclaim_counters + .last_route_removed_bytes + .load(Ordering::Relaxed), + capacity_eviction_non_ring_b_entry_total: reclaim_counters + .capacity_eviction_non_ring_b_entry_total + .load(Ordering::Relaxed), + capacity_eviction_hit_committed_slot: reclaim_counters + .capacity_eviction_hit_committed_slot + .load(Ordering::Relaxed), + eviction_reclaim_deduplicated: reclaim_counters + .eviction_reclaim_deduplicated + .load(Ordering::Relaxed), + }); + } + + MasterRuntimeObserveSnapshot { + get_holding_entries: self.inner().get_holding.total() as u64, + get_holding_bytes, + replica_cache_nodes, + } + } + + fn spawn_runtime_observe_reporter(&self) { + let view = self.0.view().clone(); + let view_task = view.clone(); + view.spawn("master_runtime_observe_reporter", async move { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + tokio::select! { + _ = shutdown_waiter.wait() => break, + _ = interval.tick() => { + let snapshot = view_task.master_kv_router().runtime_observe_snapshot(); + let activity = view_task + .master_kv_router() + .inner() + .key_activity + .observe_snapshot(); + let metrics = view_task.metric_reporter().metrics(); + metrics.set_kv_holding_entries( + "master_get_holding", + snapshot.get_holding_entries, + ); + metrics.set_kv_holding_bytes( + "master_get_holding", + snapshot.get_holding_bytes, + ); + tracing::info!( + "master get holding runtime: entries={} bytes={}", + snapshot.get_holding_entries, + snapshot.get_holding_bytes + ); + tracing::info!( + active_keys = activity.active_keys, + put_keys = activity.put_keys, + get_keys = activity.get_keys, + replica_keys = activity.replica_keys, + reclaim_keys = activity.reclaim_keys, + inflight_puts = activity.inflight_puts, + inflight_gets = activity.inflight_gets, + inflight_replicas = activity.inflight_replicas, + "master key activity runtime" + ); + tracing::info!( + active_plans = view_task + .master_kv_router() + .inner() + .planned_gets + .entry_count(), + plan_items = view_task + .master_kv_router() + .inner() + .planned_get_counters + .plan_items + .load(Ordering::Relaxed), + plan_hits = view_task + .master_kv_router() + .inner() + .planned_get_counters + .plan_hits + .load(Ordering::Relaxed), + plan_misses = view_task + .master_kv_router() + .inner() + .planned_get_counters + .plan_misses + .load(Ordering::Relaxed), + bind_succeeded = view_task + .master_kv_router() + .inner() + .planned_get_counters + .bind_succeeded + .load(Ordering::Relaxed), + bind_stale = view_task + .master_kv_router() + .inner() + .planned_get_counters + .bind_stale + .load(Ordering::Relaxed), + bind_activity_busy = view_task + .master_kv_router() + .inner() + .planned_get_counters + .bind_activity_busy + .load(Ordering::Relaxed), + plan_revoked = view_task + .master_kv_router() + .inner() + .planned_get_counters + .plan_revoked + .load(Ordering::Relaxed), + remote_ssd_filtered_items = view_task + .master_kv_router() + .inner() + .planned_get_counters + .remote_ssd_filtered_items + .load(Ordering::Relaxed), + remote_ssd_filtered_bytes = view_task + .master_kv_router() + .inner() + .planned_get_counters + .remote_ssd_filtered_bytes + .load(Ordering::Relaxed), + "master metadata-only Get plan runtime" + ); + let ssd_tier = &view_task.master_kv_router().inner().ssd_tier_counters; + tracing::info!( + local_ssd_selected_with_remote_memory_items = ssd_tier + .local_ssd_selected_with_remote_memory_items + .load(Ordering::Relaxed), + local_ssd_selected_with_remote_memory_bytes = ssd_tier + .local_ssd_selected_with_remote_memory_bytes + .load(Ordering::Relaxed), + local_ssd_selected_without_remote_memory_items = ssd_tier + .local_ssd_selected_without_remote_memory_items + .load(Ordering::Relaxed), + local_ssd_selected_without_remote_memory_bytes = ssd_tier + .local_ssd_selected_without_remote_memory_bytes + .load(Ordering::Relaxed), + local_ssd_published_with_remote_memory_items = ssd_tier + .local_ssd_published_with_remote_memory_items + .load(Ordering::Relaxed), + local_ssd_published_with_remote_memory_bytes = ssd_tier + .local_ssd_published_with_remote_memory_bytes + .load(Ordering::Relaxed), + local_ssd_published_without_remote_memory_items = ssd_tier + .local_ssd_published_without_remote_memory_items + .load(Ordering::Relaxed), + local_ssd_published_without_remote_memory_bytes = ssd_tier + .local_ssd_published_without_remote_memory_bytes + .load(Ordering::Relaxed), + memory_removed_ssd_survived_items = ssd_tier + .memory_removed_ssd_survived_items + .load(Ordering::Relaxed), + memory_removed_ssd_survived_bytes = ssd_tier + .memory_removed_ssd_survived_bytes + .load(Ordering::Relaxed), + memory_removed_ssd_became_only_items = ssd_tier + .memory_removed_ssd_became_only_items + .load(Ordering::Relaxed), + memory_removed_ssd_became_only_bytes = ssd_tier + .memory_removed_ssd_became_only_bytes + .load(Ordering::Relaxed), + "master SSD substitution runtime" + ); + for node in snapshot.replica_cache_nodes { + tracing::info!( + "node pool capacity runtime: owner={} owner_node_start_time={} capacity_epoch={} physical_capacity_bytes={} active_capacity_bytes={} used_capacity_bytes={} parked_capacity_bytes={} draining_capacity_bytes={} available_capacity_bytes={}", + node.owner_node, + node.owner_node_start_time, + node.pool_capacity_epoch, + node.pool_physical_capacity_bytes, + node.pool_active_capacity_bytes, + node.pool_used_capacity_bytes, + node.pool_parked_capacity_bytes, + node.pool_draining_capacity_bytes, + node.pool_available_capacity_bytes, + ); + tracing::info!( + "replica cache runtime: owner={} entries={} weighted_bytes={} effective_capacity_bytes={} base_capacity_bytes={} reserved_capacity_bytes={} pending_eviction_reclaim_bytes={} writeback_tier1_entries={} writeback_tier1_weighted_bytes={} writeback_tier1_capacity_bytes={} writeback_tier1_triggered={} writeback_tier1_owner_accepted={} writeback_tier1_failed={} reclaim_master_activity_deferred={} reclaim_owner_holder_deferred={} reclaim_owner_other_deferred={} reclaim_route_changed={} reclaim_retry_queued={} reclaim_retry_completed={} reclaim_retry_restored={} reclaim_completed={} source_evict_rpc_requests={} source_evict_victims={} source_evict_requested_bytes={} source_evict_accepted={} source_evict_in_progress={} source_evict_completed={} source_evict_retryable_busy={} source_evict_stale={} source_evict_rejected={} last_route_removed_members={} last_route_removed_bytes={} capacity_eviction_non_ring_b_entry_total={} capacity_eviction_hit_committed_slot={} eviction_reclaim_deduplicated={}", + node.owner_node, + node.entries, + node.weighted_bytes, + node.effective_capacity_bytes, + node.base_capacity_bytes, + node.reserved_capacity_bytes, + node.pending_eviction_reclaim_bytes, + node.writeback_tier1_entries, + node.writeback_tier1_weighted_bytes, + node.writeback_tier1_capacity_bytes, + node.writeback_tier1_triggered, + node.writeback_tier1_owner_accepted, + node.writeback_tier1_failed, + node.reclaim_master_activity_deferred, + node.reclaim_owner_holder_deferred, + node.reclaim_owner_other_deferred, + node.reclaim_route_changed, + node.reclaim_retry_queued, + node.reclaim_retry_completed, + node.reclaim_retry_restored, + node.reclaim_completed, + node.source_evict_rpc_requests, + node.source_evict_victims, + node.source_evict_requested_bytes, + node.source_evict_accepted, + node.source_evict_in_progress, + node.source_evict_completed, + node.source_evict_retryable_busy, + node.source_evict_stale, + node.source_evict_rejected, + node.last_route_removed_members, + node.last_route_removed_bytes, + node.capacity_eviction_non_ring_b_entry_total, + node.capacity_eviction_hit_committed_slot, + node.eviction_reclaim_deduplicated, + ); + metrics.set_kv_replica_cache_entries( + node.owner_node.as_str(), + node.entries, + ); + metrics.set_kv_replica_cache_weighted_bytes( + node.owner_node.as_str(), + node.weighted_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "effective", + node.effective_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "reserved", + node.reserved_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "base", + node.base_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pending_eviction_reclaim", + node.pending_eviction_reclaim_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "writeback_tier1", + node.writeback_tier1_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pool_physical", + node.pool_physical_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pool_active", + node.pool_active_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pool_used", + node.pool_used_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pool_parked", + node.pool_parked_capacity_bytes, + ); + metrics.set_kv_replica_cache_capacity_bytes( + node.owner_node.as_str(), + "pool_draining", + node.pool_draining_capacity_bytes, + ); + } + } + } + } + }); } + + // Note: no additional getters for reserved bytes; policy currently relies only on adjust calls. } // moved to crate::metrics::client diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs index 9d5eb1d..6cbfca9 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/msg_pack.rs @@ -1,10 +1,11 @@ use crate::{ cluster_manager::NodeIDString, p2p::msg_pack::{MsgPackSerializePart, RPCReq}, - rpcresp_kvresult_convert::msg_and_error::{ErrorCode, MsgId}, + rpcresp_kvresult_convert::msg_and_error::{ErrorCode, MsgId, OK}, }; use bitcode::{Decode, Encode}; use std::collections::HashMap; +use std::sync::Arc; use super::put::PutIDForAKey; @@ -16,238 +17,1490 @@ pub enum GetAllocationMode { Temporary = 0, ReuseReplica = 1, DurableReplica = 2, + LocalCommittedSlot = 3, + /// Caller-owned memory used only as the terminal data sink. The master + /// neither allocates it nor publishes it as a cache route on GetDone. + ExternalSink = 4, +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum GetSourceKind { + #[default] + Memory = 0, + Ssd = 1, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct GetPreparedLocalReserveTarget { + pub grant_id: u64, + pub slot_index: u32, + pub slot_size: u64, + pub addr: u64, + pub base_addr: u64, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct GetExternalSinkTarget { + /// Exact destination address in the requester's registered memory. + pub addr: u64, + /// Caller-validated writable bytes starting at `addr`. + pub capacity: u64, + /// Opaque requester-side registration generation, retained for identity + /// and observability. The requester remains authoritative for MR lifetime. + pub registration_id: u64, + /// Requester membership generation captured with the GPU registration. + pub requester_node_start_time: i64, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetStartReq { + pub key: String, + pub prepared_target: Option, + pub external_sink_target: Option, +} +impl MsgPackSerializePart for GetStartReq { + fn msg_id(&self) -> u32 { + MsgId::GetStartReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetStartResp { + pub get_id: u64, + pub node_id: NodeIDString, + pub put_id: PutIDForAKey, + // absolute addresses because Mooncake transfer engine requires absolute addresses (not offsets) + pub target_addr: u64, + pub src_addr: u64, + // base addresses to allow callers to convert abs->offset when needed + pub target_base_addr: u64, + pub src_base_addr: u64, + pub len: u64, + pub source_kind: GetSourceKind, + /// Echoes the owner-local slot accepted as this Get's target. + pub prepared_target: Option, + pub atomic_group: Option, + pub error_code: ErrorCode, + pub error_json: String, + /// Server-side processing time in microseconds for this RPC handler + pub server_process_us: i64, +} +impl MsgPackSerializePart for GetStartResp { + fn msg_id(&self) -> u32 { + MsgId::GetStartResp as u32 + } +} +impl RPCReq for GetStartReq { + type Resp = GetStartResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetRevokeReq { + pub get_id: u64, +} +impl MsgPackSerializePart for GetRevokeReq { + fn msg_id(&self) -> u32 { + MsgId::GetRevokeReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetRevokeResp { + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for GetRevokeResp { + fn msg_id(&self) -> u32 { + MsgId::GetRevokeResp as u32 + } +} +impl RPCReq for GetRevokeReq { + type Resp = GetRevokeResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetDoneReq { + pub get_id: u64, +} +impl MsgPackSerializePart for GetDoneReq { + fn msg_id(&self) -> u32 { + MsgId::GetDoneReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetDoneResp { + pub holder_id: u64, + pub allocation_mode: GetAllocationMode, + pub error_code: ErrorCode, + pub error_json: String, + /// Server-side processing time in microseconds for this RPC handler + pub server_process_us: i64, +} +impl MsgPackSerializePart for GetDoneResp { + fn msg_id(&self) -> u32 { + MsgId::GetDoneResp as u32 + } +} +impl RPCReq for GetDoneReq { + type Resp = GetDoneResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageBeginReq { + pub get_id: u64, +} + +impl MsgPackSerializePart for SsdStageBeginReq { + fn msg_id(&self) -> u32 { + MsgId::SsdStageBeginReq as u32 + } +} + +impl RPCReq for SsdStageBeginReq { + type Resp = SsdStageBeginResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageBeginResp { + pub started: bool, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for SsdStageBeginResp { + fn msg_id(&self) -> u32 { + MsgId::SsdStageBeginResp as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageDoneReq { + pub get_id: u64, + pub drop_ssd_source: bool, +} + +impl MsgPackSerializePart for SsdStageDoneReq { + fn msg_id(&self) -> u32 { + MsgId::SsdStageDoneReq as u32 + } +} + +impl RPCReq for SsdStageDoneReq { + type Resp = SsdStageDoneResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct SsdStageDoneResp { + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for SsdStageDoneResp { + fn msg_id(&self) -> u32 { + MsgId::SsdStageDoneResp as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetStartReq { + pub keys: Vec, + /// Empty selects ordinary master allocations. Otherwise this must contain + /// exactly one entry per key. + pub prepared_targets: Vec>, + /// Empty selects no external sinks. Otherwise this must contain exactly + /// one entry per key and is mutually exclusive with prepared targets. + pub external_sink_targets: Vec>, +} +impl MsgPackSerializePart for BatchGetStartReq { + fn msg_id(&self) -> u32 { + MsgId::BatchGetStartReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetStartItemResp { + pub get_id: u64, + pub node_id: NodeIDString, + pub put_id: PutIDForAKey, + pub target_addr: u64, + pub src_addr: u64, + pub target_base_addr: u64, + pub src_base_addr: u64, + pub len: u64, + pub source_kind: GetSourceKind, + pub prepared_target: Option, + pub atomic_group: Option, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetStartResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchGetStartResp { + fn msg_id(&self) -> u32 { + MsgId::BatchGetStartResp as u32 + } +} +impl RPCReq for BatchGetStartReq { + type Resp = BatchGetStartResp; +} + +/// Target-free Get planning. Successful items retain one exact source +/// generation until Bind, Revoke, or expiry. +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetPlanReq { + pub keys: Vec, +} +impl MsgPackSerializePart for BatchGetPlanReq { + fn msg_id(&self) -> u32 { + MsgId::BatchGetPlanReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetPlanItemResp { + pub get_id: u64, + pub node_id: NodeIDString, + pub put_id: PutIDForAKey, + pub src_addr: u64, + pub src_base_addr: u64, + pub len: u64, + pub source_kind: GetSourceKind, + pub atomic_group: Option, + /// True only for remote memory. Requester-local memory and every SSD + /// source must materialize through an owner CPU holder instead of binding + /// the RDMA-only GPU sink. + pub gpu_direct_eligible: bool, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetPlanResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchGetPlanResp { + fn msg_id(&self) -> u32 { + MsgId::BatchGetPlanResp as u32 + } +} +impl RPCReq for BatchGetPlanReq { + type Resp = BatchGetPlanResp; +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub enum GetBindTarget { + #[default] + Invalid, + PreparedLocalReserve(GetPreparedLocalReserveTarget), + ExternalSink(GetExternalSinkTarget), +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetBindItemReq { + pub get_id: u64, + pub target: GetBindTarget, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetBindReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchGetBindReq { + fn msg_id(&self) -> u32 { + MsgId::BatchGetBindReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetBindResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchGetBindResp { + fn msg_id(&self) -> u32 { + MsgId::BatchGetBindResp as u32 + } +} +impl RPCReq for BatchGetBindReq { + type Resp = BatchGetBindResp; +} + +#[cfg(test)] +mod planned_get_wire_tests { + use super::{ + BatchGetBindItemReq, BatchGetBindReq, BatchGetPlanItemResp, BatchGetPlanResp, + GetBindTarget, GetExternalSinkTarget, + }; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + + #[test] + fn first_get_id_and_late_gpu_binding_round_trip() { + let plan = BatchGetPlanResp { + items: vec![BatchGetPlanItemResp { + get_id: 0, + node_id: "source-a".to_string(), + src_addr: 0x1000, + src_base_addr: 0x800, + len: 4096, + gpu_direct_eligible: true, + error_code: OK, + ..Default::default() + }], + error_code: OK, + ..Default::default() + }; + let decoded: BatchGetPlanResp = + bitcode::decode(&bitcode::encode(&plan)).expect("decode GetPlan response"); + assert_eq!(decoded.items[0].get_id, 0); + assert!(decoded.items[0].gpu_direct_eligible); + assert_eq!(decoded.items[0].src_addr, 0x1000); + + let bind = BatchGetBindReq { + items: vec![BatchGetBindItemReq { + get_id: 0, + target: GetBindTarget::ExternalSink(GetExternalSinkTarget { + addr: 0x2000, + capacity: 4096, + registration_id: 7, + requester_node_start_time: 11, + }), + }], + }; + let decoded: BatchGetBindReq = + bitcode::decode(&bitcode::encode(&bind)).expect("decode GetBind request"); + assert!(matches!( + &decoded.items[0].target, + GetBindTarget::ExternalSink(target) + if decoded.items[0].get_id == 0 + && target.registration_id == 7 + && target.requester_node_start_time == 11 + )); + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetRevokeReq { + pub get_ids: Vec, +} +impl MsgPackSerializePart for BatchGetRevokeReq { + fn msg_id(&self) -> u32 { + MsgId::BatchGetRevokeReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetRevokeItemResp { + pub get_id: u64, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetRevokeResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for BatchGetRevokeResp { + fn msg_id(&self) -> u32 { + MsgId::BatchGetRevokeResp as u32 + } +} +impl RPCReq for BatchGetRevokeReq { + type Resp = BatchGetRevokeResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetDoneReq { + pub get_ids: Vec, +} +impl MsgPackSerializePart for BatchGetDoneReq { + fn msg_id(&self) -> u32 { + MsgId::BatchGetDoneReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetDoneItemResp { + pub get_id: u64, + pub holder_id: u64, + pub allocation_mode: GetAllocationMode, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchGetDoneResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchGetDoneResp { + fn msg_id(&self) -> u32 { + MsgId::BatchGetDoneResp as u32 + } +} +impl RPCReq for BatchGetDoneReq { + type Resp = BatchGetDoneResp; +} + +// --- RPC for CountPrefix --- + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct CountPrefixReq { + pub prefix: String, +} +impl MsgPackSerializePart for CountPrefixReq { + fn msg_id(&self) -> u32 { + MsgId::CountPrefixReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct CountPrefixResp { + pub count: u64, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for CountPrefixResp { + fn msg_id(&self) -> u32 { + MsgId::CountPrefixResp as u32 + } +} +impl RPCReq for CountPrefixReq { + type Resp = CountPrefixResp; +} + +// --- RPC for Master-only metric parts (authoritative snapshots) --- + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetMasterOnlyMetricPartReq { + pub part: String, // e.g. "segment_bytes" +} +impl MsgPackSerializePart for GetMasterOnlyMetricPartReq { + fn msg_id(&self) -> u32 { + MsgId::GetMasterOnlyMetricPartReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GetMasterOnlyMetricPartResp { + pub seg_bytes_map: HashMap, // used when part=="segment_bytes" + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for GetMasterOnlyMetricPartResp { + fn msg_id(&self) -> u32 { + MsgId::GetMasterOnlyMetricPartResp as u32 + } +} +impl RPCReq for GetMasterOnlyMetricPartReq { + type Resp = GetMasterOnlyMetricPartResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ReserveLocalGrantReq {} +impl MsgPackSerializePart for ReserveLocalGrantReq { + fn msg_id(&self) -> u32 { + MsgId::ReserveLocalGrantReq as u32 + } +} + +#[allow(unused_assignments)] +#[derive(Default, Debug, Clone, Encode, Decode)] +pub enum ReserveLocalGrantOutcome { + #[default] + None, + Granted { + grant_id: u64, + node_id: NodeIDString, + addr: u64, + base_addr: u64, + len: u64, + }, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ReserveLocalGrantResp { + pub outcome: ReserveLocalGrantOutcome, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for ReserveLocalGrantResp { + fn msg_id(&self) -> u32 { + MsgId::ReserveLocalGrantResp as u32 + } +} +impl RPCReq for ReserveLocalGrantReq { + type Resp = ReserveLocalGrantResp; +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum OwnerReclaimPhase { + #[default] + Prepare, + Commit, + Abort, + Finalize, +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum OwnerReclaimItemState { + #[default] + Busy, + Prepared, + Committed, + Aborted, + Finalized, + Stale, +} + +#[derive(Default, Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)] +pub enum OwnerReclaimBacking { + #[default] + Allocation, + CommittedSlot { + grant_id: u64, + slot_index: u32, + slot_size: u64, + }, + /// A master-owned allocation with no owner-side key index. + /// + /// SSD-capable owners use this exact source identity to persist the bytes while the + /// master's route and key-activity fence keep the allocation alive. Owners without SSD + /// still skip owner coordination and let the master reclaim the allocation directly. + UnindexedAllocation { + /// Absolute address in the owner's registered CPU segment. + addr: u64, + /// Base address of the exact registered segment generation. + base_addr: u64, + /// Logical KV payload length. + len: u64, + /// Physical allocator capacity released when the master drops the allocation. + capacity_bytes: u64, + }, +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum OwnerReclaimReason { + #[default] + OwnerCapacityEviction, + MasterAllocationCapacity, +} + +/// Per-victim SSD action for owner-local capacity reclaim. Selection and +/// deletion remain single-KV decisions even when one RPC carries a vector. +#[derive(Default, Debug, Clone, Copy, Hash, PartialEq, Eq, Encode, Decode)] +pub enum OwnerSourceSsdPolicy { + /// Delete the exact memory source without preserving it on SSD. + #[default] + Drop, + /// Ask the master to return `SsdCandidate` only if this is the last live + /// backing. Sources with another backing are deleted immediately. + SelectLastLive, + /// The owner has durably persisted this generation and supplied its exact + /// length in `ssd_backing_len`. + Persisted, +} + +/// One exact owner-local source selected for capacity eviction. +#[derive(Default, Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)] +pub struct OwnerSourceEvictionVictim { + pub key: String, + pub put_id: PutIDForAKey, + pub backing: OwnerReclaimBacking, + /// Durable owner-local SSD bytes prepared under the exact source fence. + /// The master installs this backing immediately before deleting `memory`. + pub ssd_backing_len: Option, + pub ssd_policy: OwnerSourceSsdPolicy, +} + +pub(crate) fn owner_source_eviction_epoch(operation_id: u64, victim_index: usize) -> u64 { + operation_id + .rotate_left(32) + .wrapping_add(u64::try_from(victim_index).unwrap_or(u64::MAX)) +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchEvictOwnerSourceReq { + pub operation_id: u64, + /// Membership generation of the authenticated source owner. + pub owner_node_start_time: i64, + pub victims: Vec, +} + +impl MsgPackSerializePart for BatchEvictOwnerSourceReq { + fn msg_id(&self) -> u32 { + MsgId::BatchEvictOwnerSourceReq as u32 + } +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum OwnerSourceEvictionOutcome { + #[default] + Unspecified, + Accepted, + AlreadyInProgress, + Completed, + SsdCandidate, + RetryableBusy, + Stale, + RejectedNotEvictable, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct OwnerSourceEvictionVictimResp { + pub victim_index: u32, + pub outcome: OwnerSourceEvictionOutcome, + pub ssd_backing_committed: bool, + pub detail: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchEvictOwnerSourceResp { + pub operation_id: u64, + pub victims: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for BatchEvictOwnerSourceResp { + fn msg_id(&self) -> u32 { + MsgId::BatchEvictOwnerSourceResp as u32 + } +} + +impl RPCReq for BatchEvictOwnerSourceReq { + type Resp = BatchEvictOwnerSourceResp; +} + +/// One durable same-owner SSD generation to publish without removing memory. +#[derive(Default, Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)] +pub struct OwnerSsdPublishItem { + pub key: String, + pub put_id: PutIDForAKey, + pub len: u64, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPublishOwnerSsdReq { + /// Membership generation of the authenticated owner. + pub owner_node_start_time: i64, + pub items: Vec, +} + +impl MsgPackSerializePart for BatchPublishOwnerSsdReq { + fn msg_id(&self) -> u32 { + MsgId::BatchPublishOwnerSsdReq as u32 + } +} + +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum OwnerSsdPublishOutcome { + #[default] + Unspecified, + Published, + AlreadyPresent, + RetryableBusy, + Obsolete, + Rejected, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct OwnerSsdPublishItemResp { + pub key: String, + pub put_id: PutIDForAKey, + pub outcome: OwnerSsdPublishOutcome, + pub detail: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPublishOwnerSsdResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for BatchPublishOwnerSsdResp { + fn msg_id(&self) -> u32 { + MsgId::BatchPublishOwnerSsdResp as u32 + } +} + +impl RPCReq for BatchPublishOwnerSsdReq { + type Resp = BatchPublishOwnerSsdResp; +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct OwnerReclaimItem { + pub key: String, + pub put_id: (u64, u32), + pub epoch: u64, + pub backing: OwnerReclaimBacking, + pub reason: OwnerReclaimReason, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchOwnerReclaimReq { + pub phase: OwnerReclaimPhase, + pub items: Vec, +} + +impl MsgPackSerializePart for BatchOwnerReclaimReq { + fn msg_id(&self) -> u32 { + MsgId::BatchOwnerReclaimReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct OwnerReclaimItemResp { + pub key: String, + pub epoch: u64, + pub state: OwnerReclaimItemState, + /// Durable bytes persisted by the owner while the exact memory source is + /// hidden behind the Prepare fence. The master publishes this backing on + /// the existing route before it asks the owner to Commit/free DRAM. + pub ssd_backing_len: Option, + pub detail: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchOwnerReclaimResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for BatchOwnerReclaimResp { + fn msg_id(&self) -> u32 { + MsgId::BatchOwnerReclaimResp as u32 + } +} + +impl RPCReq for BatchOwnerReclaimReq { + type Resp = BatchOwnerReclaimResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct EnqueueReplicaTaskItem { + pub key: String, + pub put_id: PutIDForAKey, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchEnqueueReplicaTaskReq { + pub items: Vec, +} + +impl MsgPackSerializePart for BatchEnqueueReplicaTaskReq { + fn msg_id(&self) -> u32 { + MsgId::BatchEnqueueReplicaTaskReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct EnqueueReplicaTaskItemResp { + pub key: String, + pub put_id: PutIDForAKey, + pub accepted: bool, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchEnqueueReplicaTaskResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} + +impl MsgPackSerializePart for BatchEnqueueReplicaTaskResp { + fn msg_id(&self) -> u32 { + MsgId::BatchEnqueueReplicaTaskResp as u32 + } +} + +impl RPCReq for BatchEnqueueReplicaTaskReq { + type Resp = BatchEnqueueReplicaTaskResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ReleaseLocalGrantReq { + pub grant_id: u64, +} +impl MsgPackSerializePart for ReleaseLocalGrantReq { + fn msg_id(&self) -> u32 { + MsgId::ReleaseLocalGrantReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct ReleaseLocalGrantResp { + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for ReleaseLocalGrantResp { + fn msg_id(&self) -> u32 { + MsgId::ReleaseLocalGrantResp as u32 + } +} +impl RPCReq for ReleaseLocalGrantReq { + type Resp = ReleaseLocalGrantResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPreparePutKeyItemReq { + pub key: String, + pub reject_if_inflight_same_key: bool, + pub reject_if_exist_same_key: bool, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPreparePutKeysReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPreparePutKeysReq { + fn msg_id(&self) -> u32 { + MsgId::BatchPreparePutKeysReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPreparePutKeysResp { + pub reservation_ids: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchPreparePutKeysResp { + fn msg_id(&self) -> u32 { + MsgId::BatchPreparePutKeysResp as u32 + } +} +impl RPCReq for BatchPreparePutKeysReq { + type Resp = BatchPreparePutKeysResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchReleasePutKeyReservationsReq { + pub reservation_ids: Vec, +} +impl MsgPackSerializePart for BatchReleasePutKeyReservationsReq { + fn msg_id(&self) -> u32 { + MsgId::BatchReleasePutKeyReservationsReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchReleasePutKeyReservationsResp { + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchReleasePutKeyReservationsResp { + fn msg_id(&self) -> u32 { + MsgId::BatchReleasePutKeyReservationsResp as u32 + } +} +impl RPCReq for BatchReleasePutKeyReservationsReq { + type Resp = BatchReleasePutKeyReservationsResp; +} + +// --- RPC for Put --- + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutStartReq { + pub key: String, + pub len: u64, + pub reject_if_inflight_same_key: bool, + pub reject_if_exist_same_key: bool, + pub make_replica_task: bool, + /// Prefer placing the target allocation on any kvclient within this sub_cluster. + pub preferred_sub_cluster: Option, + /// Optional source-node override for side-transfer workers that share an owner's mmap. + pub source_node_id: Option, +} +impl MsgPackSerializePart for PutStartReq { + fn msg_id(&self) -> u32 { + MsgId::PutStartReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutReplicaTarget { + pub node_id: NodeIDString, + pub target_addr: u64, + pub target_base_addr: u64, + pub len: u64, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutStartResp { + pub put_id: PutIDForAKey, + pub node_id: NodeIDString, + // absolute addresses because Mooncake transfer engine requires absolute addresses (not offsets) + pub target_addr: u64, + pub src_addr: u64, + // base addresses to allow callers to convert abs->offset when needed + pub target_base_addr: u64, + pub src_base_addr: u64, + pub len: u64, + pub error_code: ErrorCode, + pub error_json: String, + /// Server-side processing time in microseconds for this RPC handler + pub server_process_us: i64, + pub replica_target: Option, +} +impl MsgPackSerializePart for PutStartResp { + fn msg_id(&self) -> u32 { + MsgId::PutStartResp as u32 + } +} +impl RPCReq for PutStartReq { + type Resp = PutStartResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutRevokeReq { + pub key: String, + pub put_id: PutIDForAKey, +} +impl MsgPackSerializePart for PutRevokeReq { + fn msg_id(&self) -> u32 { + MsgId::PutRevokeReq as u32 + } +} +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutRevokeResp { + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for PutRevokeResp { + fn msg_id(&self) -> u32 { + MsgId::PutRevokeResp as u32 + } +} +impl RPCReq for PutRevokeReq { + type Resp = PutRevokeResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct PutDoneCommittedSlot { + pub grant_id: u64, + pub slot_index: u32, + pub slot_size: u64, + pub addr: u64, + pub base_addr: u64, + pub len: u64, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct PutAtomicGroupMember { + pub key: String, + pub put_id: PutIDForAKey, +} + +/// Version-scoped members of one caller-declared atomic put group. +/// +/// Groups with one member are represented as `None` on the route. Multi-member +/// groups let eviction require one common remote-cache owner to hold every member. +#[derive(Default, Debug, Clone, PartialEq, Eq, Encode, Decode)] +pub struct PutAtomicGroup { + pub members: Vec, +} + +pub fn build_put_atomic_group_assignments( + keys_and_put_ids: &[(String, PutIDForAKey)], + atomic_group_lens: &[usize], +) -> Result>, String> { + build_shared_put_atomic_group_assignments(keys_and_put_ids, atomic_group_lens).map( + |assignments| { + assignments + .into_iter() + .map(|group| group.map(|group| group.as_ref().clone())) + .collect() + }, + ) +} + +/// Builds one shared descriptor per multi-key group and assigns cheap `Arc` +/// clones to its members. This is the grouped-put representation used by the +/// V2 route-publish protocol; unlike the legacy wire representation, it is +/// linear in the number of keys rather than the sum of squared group sizes. +pub fn build_shared_put_atomic_group_assignments( + keys_and_put_ids: &[(String, PutIDForAKey)], + atomic_group_lens: &[usize], +) -> Result>>, String> { + if atomic_group_lens.is_empty() && !keys_and_put_ids.is_empty() { + return Err("atomic_group_lens must be non-empty".to_string()); + } + let mut offset = 0usize; + let mut assignments = Vec::with_capacity(keys_and_put_ids.len()); + for (group_index, group_len) in atomic_group_lens.iter().copied().enumerate() { + if group_len == 0 { + return Err(format!( + "atomic_group_lens entries must be > 0; index={group_index}" + )); + } + if group_len > 4096 { + return Err(format!( + "atomic_group_lens entries must be <= 4096; index={group_index} len={group_len}" + )); + } + let end = offset + .checked_add(group_len) + .ok_or_else(|| "atomic_group_lens sum overflowed usize".to_string())?; + let members = keys_and_put_ids.get(offset..end).ok_or_else(|| { + format!( + "atomic_group_lens exceeds keys length; end={} keys={}", + end, + keys_and_put_ids.len() + ) + })?; + if group_len == 1 { + assignments.push(None); + } else { + let group = Arc::new(PutAtomicGroup { + members: members + .iter() + .map(|(key, put_id)| PutAtomicGroupMember { + key: key.clone(), + put_id: *put_id, + }) + .collect(), + }); + assignments.extend((0..group_len).map(|_| Some(group.clone()))); + } + offset = end; + } + if offset != keys_and_put_ids.len() { + return Err(format!( + "atomic_group_lens must sum to keys length; sum={} keys={}", + offset, + keys_and_put_ids.len() + )); + } + Ok(assignments) } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetStartReq { +pub struct PutDoneReq { pub key: String, + pub put_id: PutIDForAKey, + /// Optional lease to attach this key to on commit + pub lease_id: Option, + /// Optional local committed slot descriptor for local-first publish path. + pub committed_slot: Option, + /// Ask master to keep a local read holder for the committing node. + pub publish_local_cache: bool, + /// Multi-key atomic group for this exact key version. + pub atomic_group: Option, } -impl MsgPackSerializePart for GetStartReq { +impl MsgPackSerializePart for PutDoneReq { fn msg_id(&self) -> u32 { - MsgId::GetStartReq as u32 + MsgId::PutDoneReq as u32 } } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetStartResp { - pub get_id: u64, - pub node_id: NodeIDString, +pub struct PutDoneResp { + pub error_code: ErrorCode, + pub error_json: String, + /// Server-side processing time in microseconds for this RPC handler + pub server_process_us: i64, + /// Holder id for an owner-local cache view, present only when requested. + pub local_cache_holder_id: Option, +} +impl MsgPackSerializePart for PutDoneResp { + fn msg_id(&self) -> u32 { + MsgId::PutDoneResp as u32 + } +} +impl RPCReq for PutDoneReq { + type Resp = PutDoneResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutStartItemReq { + pub key: String, + pub len: u64, + pub reject_if_inflight_same_key: bool, + pub reject_if_exist_same_key: bool, + pub make_replica_task: bool, + pub preferred_sub_cluster: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutStartReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPutStartReq { + fn msg_id(&self) -> u32 { + MsgId::BatchPutStartReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutStartItemResp { pub put_id: PutIDForAKey, - // absolute addresses because Mooncake transfer engine requires absolute addresses (not offsets) + pub node_id: NodeIDString, pub target_addr: u64, pub src_addr: u64, - // base addresses to allow callers to convert abs->offset when needed pub target_base_addr: u64, pub src_base_addr: u64, pub len: u64, pub error_code: ErrorCode, pub error_json: String, - /// Server-side processing time in microseconds for this RPC handler + pub replica_target: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutStartResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, pub server_process_us: i64, } -impl MsgPackSerializePart for GetStartResp { +impl MsgPackSerializePart for BatchPutStartResp { fn msg_id(&self) -> u32 { - MsgId::GetStartResp as u32 + MsgId::BatchPutStartResp as u32 } } -impl RPCReq for GetStartReq { - type Resp = GetStartResp; +impl RPCReq for BatchPutStartReq { + type Resp = BatchPutStartResp; } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetRevokeReq { - pub get_id: u64, +pub struct BatchPutRevokeItemReq { + pub key: String, + pub put_id: PutIDForAKey, } -impl MsgPackSerializePart for GetRevokeReq { + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutRevokeReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPutRevokeReq { fn msg_id(&self) -> u32 { - MsgId::GetRevokeReq as u32 + MsgId::BatchPutRevokeReq as u32 } } + #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetRevokeResp { +pub struct BatchPutRevokeItemResp { + pub key: String, + pub put_id: PutIDForAKey, pub error_code: ErrorCode, pub error_json: String, } -impl MsgPackSerializePart for GetRevokeResp { + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutRevokeResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for BatchPutRevokeResp { fn msg_id(&self) -> u32 { - MsgId::GetRevokeResp as u32 + MsgId::BatchPutRevokeResp as u32 } } -impl RPCReq for GetRevokeReq { - type Resp = GetRevokeResp; +impl RPCReq for BatchPutRevokeReq { + type Resp = BatchPutRevokeResp; } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetDoneReq { - pub get_id: u64, +pub struct BatchPutDoneItemReq { + pub key: String, + pub put_id: PutIDForAKey, + pub lease_id: Option, + pub committed_slot: Option, + pub publish_local_cache: bool, + pub atomic_group: Option, } -impl MsgPackSerializePart for GetDoneReq { + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutDoneReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPutDoneReq { fn msg_id(&self) -> u32 { - MsgId::GetDoneReq as u32 + MsgId::BatchPutDoneReq as u32 } } + #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetDoneResp { - pub holder_id: u64, - pub allocation_mode: GetAllocationMode, +pub struct BatchPutDoneItemResp { + pub key: String, + pub put_id: PutIDForAKey, + pub error_code: ErrorCode, + pub error_json: String, + pub local_cache_holder_id: Option, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutDoneResp { + pub items: Vec, pub error_code: ErrorCode, pub error_json: String, - /// Server-side processing time in microseconds for this RPC handler pub server_process_us: i64, } -impl MsgPackSerializePart for GetDoneResp { +impl MsgPackSerializePart for BatchPutDoneResp { fn msg_id(&self) -> u32 { - MsgId::GetDoneResp as u32 + MsgId::BatchPutDoneResp as u32 } } -impl RPCReq for GetDoneReq { - type Resp = GetDoneResp; +impl RPCReq for BatchPutDoneReq { + type Resp = BatchPutDoneResp; } -// --- RPC for CountPrefix --- +/// Linear-size V2 batch route publication. `atomic_group_lens` partitions the +/// ordered items; the master reconstructs each group once from the item keys +/// and put ids, then shares one interned descriptor across member routes. +/// +/// The V1 `BatchPutDoneReq` remains registered unchanged for rolling and API +/// compatibility. +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct GroupedBatchPutDoneItemReq { + pub key: String, + pub put_id: PutIDForAKey, + pub lease_id: Option, + pub committed_slot: Option, + pub publish_local_cache: bool, +} #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct CountPrefixReq { - pub prefix: String, +pub struct GroupedBatchPutDoneReq { + pub items: Vec, + pub atomic_group_lens: Vec, } -impl MsgPackSerializePart for CountPrefixReq { +impl MsgPackSerializePart for GroupedBatchPutDoneReq { fn msg_id(&self) -> u32 { - MsgId::CountPrefixReq as u32 + MsgId::GroupedBatchPutDoneReq as u32 } } + #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct CountPrefixResp { - pub count: u64, +pub struct GroupedBatchPutDoneResp { + pub items: Vec, pub error_code: ErrorCode, pub error_json: String, + pub server_process_us: i64, } -impl MsgPackSerializePart for CountPrefixResp { +impl MsgPackSerializePart for GroupedBatchPutDoneResp { fn msg_id(&self) -> u32 { - MsgId::CountPrefixResp as u32 + MsgId::GroupedBatchPutDoneResp as u32 } } -impl RPCReq for CountPrefixReq { - type Resp = CountPrefixResp; +impl RPCReq for GroupedBatchPutDoneReq { + type Resp = GroupedBatchPutDoneResp; } -// --- RPC for Master-only metric parts (authoritative snapshots) --- - #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetMasterOnlyMetricPartReq { - pub part: String, // e.g. "segment_bytes" +pub struct PutAppendStartReq { + pub key: String, + pub put_id: PutIDForAKey, + pub len: u64, + pub preferred_sub_cluster: Option, + pub protect_source_on_remote_complete: bool, } -impl MsgPackSerializePart for GetMasterOnlyMetricPartReq { +impl MsgPackSerializePart for PutAppendStartReq { fn msg_id(&self) -> u32 { - MsgId::GetMasterOnlyMetricPartReq as u32 + MsgId::PutAppendStartReq as u32 } } +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)] +pub enum PutAppendStartOutcome { + /// Missing/old peers must not accidentally interpret a zero value as + /// successful completion. + #[default] + Unspecified, + Scheduled, + /// A complete non-source replica already exists for this exact put_id. + AlreadySatisfied, + /// The source route/version no longer exists; retry would target stale data. + Obsolete, + /// No remote allocation is available now. The owner keeps its local slot + /// and retries with backoff; this is never a demotion/drop instruction. + RetryableNoSpace, +} + #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct GetMasterOnlyMetricPartResp { - pub seg_bytes_map: HashMap, // used when part=="segment_bytes" +pub struct PutAppendStartResp { + pub outcome: PutAppendStartOutcome, + /// Master-issued identity for this concrete replica append attempt. + /// + /// `put_id` identifies the KV generation, but one generation may need to + /// be copied remotely more than once after an earlier remote route is + /// reclaimed. Done/Revoke must echo this value so an old replayable + /// terminal result cannot complete a later reservation. + pub operation_id: u64, + pub node_id: NodeIDString, + pub target_addr: u64, + pub target_base_addr: u64, + pub len: u64, pub error_code: ErrorCode, pub error_json: String, + pub server_process_us: i64, } -impl MsgPackSerializePart for GetMasterOnlyMetricPartResp { +impl MsgPackSerializePart for PutAppendStartResp { fn msg_id(&self) -> u32 { - MsgId::GetMasterOnlyMetricPartResp as u32 + MsgId::PutAppendStartResp as u32 } } -impl RPCReq for GetMasterOnlyMetricPartReq { - type Resp = GetMasterOnlyMetricPartResp; +impl RPCReq for PutAppendStartReq { + type Resp = PutAppendStartResp; } -// --- RPC for Put --- - #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutStartReq { +pub struct BatchPutAppendStartItemReq { pub key: String, + pub put_id: PutIDForAKey, pub len: u64, - pub reject_if_inflight_same_key: bool, - /// Prefer placing the target allocation on any kvclient within this sub_cluster. pub preferred_sub_cluster: Option, - /// Optional source-node override for side-transfer workers that share an owner's mmap. - pub source_node_id: Option, + pub protect_source_on_remote_complete: bool, } -impl MsgPackSerializePart for PutStartReq { + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendStartReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPutAppendStartReq { fn msg_id(&self) -> u32 { - MsgId::PutStartReq as u32 + MsgId::BatchPutAppendStartReq as u32 } } + #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutStartResp { +pub struct BatchPutAppendStartItemResp { + pub key: String, pub put_id: PutIDForAKey, + pub outcome: PutAppendStartOutcome, + pub operation_id: u64, pub node_id: NodeIDString, - // absolute addresses because Mooncake transfer engine requires absolute addresses (not offsets) pub target_addr: u64, - pub src_addr: u64, - // base addresses to allow callers to convert abs->offset when needed pub target_base_addr: u64, - pub src_base_addr: u64, pub len: u64, pub error_code: ErrorCode, pub error_json: String, - /// Server-side processing time in microseconds for this RPC handler +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendStartResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, pub server_process_us: i64, } -impl MsgPackSerializePart for PutStartResp { +impl MsgPackSerializePart for BatchPutAppendStartResp { fn msg_id(&self) -> u32 { - MsgId::PutStartResp as u32 + MsgId::BatchPutAppendStartResp as u32 } } -impl RPCReq for PutStartReq { - type Resp = PutStartResp; +impl RPCReq for BatchPutAppendStartReq { + type Resp = BatchPutAppendStartResp; } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutRevokeReq { +pub struct PutAppendRevokeReq { pub key: String, pub put_id: PutIDForAKey, + pub operation_id: u64, } -impl MsgPackSerializePart for PutRevokeReq { +impl MsgPackSerializePart for PutAppendRevokeReq { fn msg_id(&self) -> u32 { - MsgId::PutRevokeReq as u32 + MsgId::PutAppendRevokeReq as u32 } } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutRevokeResp { +pub struct PutAppendRevokeResp { pub error_code: ErrorCode, pub error_json: String, } -impl MsgPackSerializePart for PutRevokeResp { +impl MsgPackSerializePart for PutAppendRevokeResp { fn msg_id(&self) -> u32 { - MsgId::PutRevokeResp as u32 + MsgId::PutAppendRevokeResp as u32 } } -impl RPCReq for PutRevokeReq { - type Resp = PutRevokeResp; +impl RPCReq for PutAppendRevokeReq { + type Resp = PutAppendRevokeResp; } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutDoneReq { +pub struct PutAppendDoneReq { pub key: String, pub put_id: PutIDForAKey, - /// Optional lease to attach this key to on commit - pub lease_id: Option, + pub operation_id: u64, } -impl MsgPackSerializePart for PutDoneReq { +impl MsgPackSerializePart for PutAppendDoneReq { fn msg_id(&self) -> u32 { - MsgId::PutDoneReq as u32 + MsgId::PutAppendDoneReq as u32 } } #[derive(Default, Debug, Clone, Encode, Decode)] -pub struct PutDoneResp { +pub struct PutAppendDoneResp { + pub appended: bool, pub error_code: ErrorCode, pub error_json: String, - /// Server-side processing time in microseconds for this RPC handler pub server_process_us: i64, } -impl MsgPackSerializePart for PutDoneResp { +impl MsgPackSerializePart for PutAppendDoneResp { fn msg_id(&self) -> u32 { - MsgId::PutDoneResp as u32 + MsgId::PutAppendDoneResp as u32 } } -impl RPCReq for PutDoneReq { - type Resp = PutDoneResp; +impl RPCReq for PutAppendDoneReq { + type Resp = PutAppendDoneResp; +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendDoneItemReq { + pub key: String, + pub put_id: PutIDForAKey, + pub operation_id: u64, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendDoneReq { + pub items: Vec, +} +impl MsgPackSerializePart for BatchPutAppendDoneReq { + fn msg_id(&self) -> u32 { + MsgId::BatchPutAppendDoneReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendDoneItemResp { + pub key: String, + pub put_id: PutIDForAKey, + pub appended: bool, + pub error_code: ErrorCode, + pub error_json: String, +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchPutAppendDoneResp { + pub items: Vec, + pub error_code: ErrorCode, + pub error_json: String, + pub server_process_us: i64, +} +impl MsgPackSerializePart for BatchPutAppendDoneResp { + fn msg_id(&self) -> u32 { + MsgId::BatchPutAppendDoneResp as u32 + } +} +impl RPCReq for BatchPutAppendDoneReq { + type Resp = BatchPutAppendDoneResp; } // --- RPC for MemHolder KeepAlive --- @@ -414,6 +1667,31 @@ impl RPCReq for GetMetaReq { type Resp = GetMetaResp; } +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchIsExistReq { + pub keys: Vec, +} +impl MsgPackSerializePart for BatchIsExistReq { + fn msg_id(&self) -> u32 { + MsgId::BatchIsExistReq as u32 + } +} + +#[derive(Default, Debug, Clone, Encode, Decode)] +pub struct BatchIsExistResp { + pub exists_list: Vec, + pub error_code: ErrorCode, + pub error_json: String, +} +impl MsgPackSerializePart for BatchIsExistResp { + fn msg_id(&self) -> u32 { + MsgId::BatchIsExistResp as u32 + } +} +impl RPCReq for BatchIsExistReq { + type Resp = BatchIsExistResp; +} + // --- RPC for Batch Delete Client KV Meta Cache --- #[derive(Debug, Clone, Encode, Decode, Default)] @@ -451,3 +1729,128 @@ impl MsgPackSerializePart for BatchDeleteClientKvMetaCacheResp { impl RPCReq for BatchDeleteClientKvMetaCacheReq { type Resp = BatchDeleteClientKvMetaCacheResp; } + +#[cfg(test)] +mod put_atomic_group_tests { + use super::*; + + #[test] + fn assignments_expand_only_multi_member_groups() { + let keys_and_put_ids = vec![ + ("a".to_string(), (1, 0)), + ("b".to_string(), (1, 1)), + ("c".to_string(), (1, 2)), + ]; + let assignments = build_put_atomic_group_assignments(&keys_and_put_ids, &[2, 1]).unwrap(); + assert_eq!(assignments.len(), 3); + assert_eq!(assignments[0], assignments[1]); + assert_eq!(assignments[0].as_ref().unwrap().members.len(), 2); + assert!(assignments[2].is_none()); + } + + #[test] + fn assignments_reject_invalid_partitions() { + let keys_and_put_ids = vec![("a".to_string(), (1, 0)), ("b".to_string(), (1, 1))]; + assert!(build_put_atomic_group_assignments(&keys_and_put_ids, &[1]).is_err()); + assert!(build_put_atomic_group_assignments(&keys_and_put_ids, &[0, 2]).is_err()); + } + + #[test] + fn batch_put_done_group_round_trips_on_wire() { + let group = PutAtomicGroup { + members: vec![ + PutAtomicGroupMember { + key: "a".to_string(), + put_id: (1, 0), + }, + PutAtomicGroupMember { + key: "b".to_string(), + put_id: (1, 1), + }, + ], + }; + let req = BatchPutDoneReq { + items: vec![BatchPutDoneItemReq { + key: "a".to_string(), + put_id: (1, 0), + lease_id: None, + committed_slot: None, + publish_local_cache: false, + atomic_group: Some(group.clone()), + }], + }; + let decoded: BatchPutDoneReq = + bitcode::decode(&bitcode::encode(&req)).expect("decode atomic put group"); + assert_eq!(decoded.items[0].atomic_group.as_ref(), Some(&group)); + } + + #[test] + fn owner_ssd_publish_only_batch_round_trips_exact_generation() { + let req = BatchPublishOwnerSsdReq { + owner_node_start_time: 41, + items: vec![OwnerSsdPublishItem { + key: "ssd-key".to_string(), + put_id: (17, 3), + len: 4_718_592, + }], + }; + let decoded: BatchPublishOwnerSsdReq = + bitcode::decode(&bitcode::encode(&req)).expect("decode owner SSD publication"); + assert_eq!(decoded.owner_node_start_time, 41); + assert_eq!(decoded.items, req.items); + } + + #[test] + fn grouped_batch_put_done_wire_is_linear_and_round_trips() { + let keys_and_put_ids = (0..128) + .map(|index| (format!("page-{index:03}"), (7, index))) + .collect::>(); + let group = PutAtomicGroup { + members: keys_and_put_ids + .iter() + .map(|(key, put_id)| PutAtomicGroupMember { + key: key.clone(), + put_id: *put_id, + }) + .collect(), + }; + let legacy = BatchPutDoneReq { + items: keys_and_put_ids + .iter() + .map(|(key, put_id)| BatchPutDoneItemReq { + key: key.clone(), + put_id: *put_id, + lease_id: None, + committed_slot: None, + publish_local_cache: false, + atomic_group: Some(group.clone()), + }) + .collect(), + }; + let grouped = GroupedBatchPutDoneReq { + items: keys_and_put_ids + .iter() + .map(|(key, put_id)| GroupedBatchPutDoneItemReq { + key: key.clone(), + put_id: *put_id, + lease_id: None, + committed_slot: None, + publish_local_cache: false, + }) + .collect(), + atomic_group_lens: vec![128], + }; + let legacy_bytes = bitcode::encode(&legacy); + let grouped_bytes = bitcode::encode(&grouped); + assert!( + grouped_bytes.len() * 16 < legacy_bytes.len(), + "grouped={} legacy={}", + grouped_bytes.len(), + legacy_bytes.len() + ); + let decoded: GroupedBatchPutDoneReq = + bitcode::decode(&grouped_bytes).expect("decode grouped put done"); + assert_eq!(decoded.items.len(), 128); + assert_eq!(decoded.atomic_group_lens, vec![128]); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs index d9f509f..161d950 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/placement.rs @@ -1,7 +1,10 @@ +use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::sync::atomic::Ordering as AtomicOrdering; use crate::{ - cluster_manager::NodeID, + cluster_manager::{ClusterMember, NodeID}, + config::{ReplicaTaskPlacementConfig, ReplicaTaskPlacementPolicyKind}, master_seg_manager::one_seg_allocator::{Allocation, OneSegAllocator}, rpcresp_kvresult_convert::msg_and_error::KvError, }; @@ -17,7 +20,6 @@ pub enum PutPlacementTarget { /// Place remotely with a pre-allocated target allocation. Remote { node_id: NodeID, - allocator: Arc, allocation: Allocation, }, } @@ -33,15 +35,29 @@ pub trait PlacementPolicy: Send + Sync { preferred_sub_cluster: Option<&str>, len: u64, ) -> Result; + + /// Selects a remote-only target for replica task placement. + fn select_remote_target( + &self, + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + len: u64, + ) -> Result<(NodeID, Allocation), KvError> { + choose_random_remote_target( + view, + source_node_id, + excluded_nodes, + preferred_sub_cluster, + len, + ) + } } -/// Compile-time switch for the master placement default. -/// -/// Change the type alias to switch behavior (and rebuild): -/// - `LocalFirstPlacementPolicy` prefers local placement when possible. -/// - `RandomPlacementPolicy` selects a random eligible target. -// pub type PlacementDefault = LocalFirstPlacementPolicy; -pub type PlacementDefault = RandomPlacementPolicy; +pub fn build_placement_policy(config: ReplicaTaskPlacementConfig) -> Box { + Box::new(ReplicaTaskPlacementPolicy::new(config)) +} /// A policy that prefers placing on the requesting node when possible. pub struct LocalFirstPlacementPolicy; @@ -103,9 +119,9 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { continue; }; - let total = allocator.total_size_bytes(); - let used = allocator.used_size_bytes(); - let free = total.saturating_sub(used); + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; last_no_space_ctx = Some(( node_id.as_ref().to_string(), allocator.seg_device_id.clone(), @@ -116,7 +132,6 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { node_id, - allocator, allocation, }); } @@ -137,9 +152,9 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { let all_segs = seg_manager.get_all_segments_allocator(); if let Some((nodeid, allocator)) = all_segs.choose(&mut rand::thread_rng()).cloned() { let node_id: NodeID = nodeid.into(); - let total = allocator.total_size_bytes(); - let used = allocator.used_size_bytes(); - let free = total.saturating_sub(used); + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; last_no_space_ctx = Some(( node_id.as_ref().to_string(), allocator.seg_device_id.clone(), @@ -149,7 +164,6 @@ impl PlacementPolicy for LocalFirstPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { node_id, - allocator, allocation, }); } @@ -246,9 +260,9 @@ impl PlacementPolicy for RandomPlacementPolicy { continue; }; - let total = allocator.total_size_bytes(); - let used = allocator.used_size_bytes(); - let free = total.saturating_sub(used); + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; last_no_space_ctx = Some(( node_id.as_ref().to_string(), allocator.seg_device_id.clone(), @@ -259,7 +273,6 @@ impl PlacementPolicy for RandomPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { node_id, - allocator, allocation, }); } @@ -281,9 +294,9 @@ impl PlacementPolicy for RandomPlacementPolicy { continue; } - let total = allocator.total_size_bytes(); - let used = allocator.used_size_bytes(); - let free = total.saturating_sub(used); + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; last_no_space_ctx = Some(( node_id.as_ref().to_string(), allocator.seg_device_id.clone(), @@ -293,7 +306,6 @@ impl PlacementPolicy for RandomPlacementPolicy { if let Ok(allocation) = allocator.allocate(len) { return Ok(PutPlacementTarget::Remote { node_id, - allocator, allocation, }); } @@ -322,3 +334,681 @@ impl PlacementPolicy for RandomPlacementPolicy { Err(err) } } + +#[derive(Clone)] +struct PlacementCandidate { + node_id: NodeID, + allocator: Arc, + total_bytes: u64, + free_bytes: u64, + used_bytes: u64, + node_write_count: u64, + requester_target_count: u64, + is_remote_only_role: bool, + is_active_role: bool, + preferred_sub_cluster_match: bool, +} + +impl PlacementCandidate { + fn queue_wait_ms(&self) -> f64 { + self.node_write_count as f64 + (self.requester_target_count as f64 * 0.01) + } + + fn mem_pressure(&self) -> f64 { + if self.total_bytes == 0 { + 1.0 + } else { + 1.0 - (self.free_bytes as f64 / self.total_bytes as f64) + } + } + + fn queue_score(&self) -> f64 { + self.queue_wait_ms() + self.mem_pressure() * 0.001 + self.used_bytes as f64 * 1e-9 + } +} + +type NoSpaceCtx = (String, String, u64, u64); + +fn no_space_error(last_no_space_ctx: Option) -> KvError { + if let Some((node, segment, total_capacity, free_capacity)) = last_no_space_ctx { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::NoSpace { + node, + segment, + total_capacity, + free_capacity, + }, + ) + } else { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::NoSpace { + node: "unknown".to_string(), + segment: "unknown".to_string(), + total_capacity: 0, + free_capacity: 0, + }, + ) + } +} + +pub(super) fn member_matches_roles(member: Option<&ClusterMember>, roles: &[String]) -> bool { + let Some(member) = member else { + return false; + }; + let metadata_role = member.metadata.get("role").map(|v| v.as_str()); + let metadata_node_role = member.metadata.get("node_role").map(|v| v.as_str()); + let sub_cluster = member.sub_cluster.as_deref(); + roles.iter().any(|role| { + let role = role.as_str(); + metadata_role == Some(role) || metadata_node_role == Some(role) || sub_cluster == Some(role) + }) +} + +fn is_role_aware_policy(policy: ReplicaTaskPlacementPolicyKind) -> bool { + matches!( + policy, + ReplicaTaskPlacementPolicyKind::WeightedRoleAware + | ReplicaTaskPlacementPolicyKind::BoundedRoleQueueAware + | ReplicaTaskPlacementPolicyKind::PressureRoleQueueAware + ) +} + +fn filter_remote_only_candidates(candidates: &[PlacementCandidate]) -> Vec { + candidates + .iter() + .filter(|candidate| candidate.is_remote_only_role) + .cloned() + .collect() +} + +fn collect_remote_candidates( + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let members_by_id: HashMap = view + .cluster_manager() + .get_client_members() + .into_iter() + .map(|member| (member.id.clone(), member)) + .collect(); + + let mut candidates = Vec::new(); + for (node_id, allocator) in view.master_seg_manager().get_all_segments_allocator() { + if node_id.as_ref() == source_node_id.as_ref() || excluded_nodes.contains(&node_id) { + continue; + } + + let member = members_by_id.get(node_id.as_ref()); + let preferred_sub_cluster_match = preferred_sub_cluster + .map(|sc| member.and_then(|m| m.sub_cluster.as_deref()) == Some(sc)) + .unwrap_or(false); + if preferred_sub_cluster.is_some() && !preferred_sub_cluster_match { + continue; + } + + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let used = capacity.used_capacity_bytes; + let free = capacity.available_capacity_bytes; + let node_key = node_id.as_ref().to_string(); + let node_write_count = view + .master_kv_router() + .inner() + .put_target_decision_counts + .get(&node_key) + .map(|entry| entry.value().load(AtomicOrdering::Relaxed)) + .unwrap_or(0); + let requester_target_count = view + .master_kv_router() + .inner() + .put_requester_target_decision_counts + .get(&super::RequesterTargetPair::new( + source_node_id.as_ref(), + node_id.as_ref(), + )) + .map(|entry| entry.value().load(AtomicOrdering::Relaxed)) + .unwrap_or(0); + + candidates.push(PlacementCandidate { + node_id, + allocator, + total_bytes: total, + free_bytes: free, + used_bytes: used, + node_write_count, + requester_target_count, + is_remote_only_role: member_matches_roles(member, &config.remote_only_node_roles), + is_active_role: member_matches_roles(member, &config.active_node_roles), + preferred_sub_cluster_match, + }); + } + candidates +} + +fn choose_candidate_pool( + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let global = collect_remote_candidates(view, source_node_id, excluded_nodes, None, config); + let global_remote_only = filter_remote_only_candidates(&global); + let Some(sc) = preferred_sub_cluster else { + if config.restrict_to_remote_only_node_roles { + return global_remote_only; + } + return global; + }; + + let preferred = + collect_remote_candidates(view, source_node_id, excluded_nodes, Some(sc), config); + let preferred_remote_only = filter_remote_only_candidates(&preferred); + if config.restrict_to_remote_only_node_roles { + if !preferred_remote_only.is_empty() { + return preferred_remote_only; + } + if !global_remote_only.is_empty() { + tracing::warn!( + "preferred_sub_cluster has no eligible remote-only kvclients; using global remote-only candidates: source_node_id={} sub_cluster={:?} remote_only_node_roles={:?}", + source_node_id, + sc, + config.remote_only_node_roles + ); + return global_remote_only; + } + tracing::warn!( + "strict remote-only placement has no eligible candidates: source_node_id={} sub_cluster={:?} remote_only_node_roles={:?}", + source_node_id, + sc, + config.remote_only_node_roles + ); + return Vec::new(); + } + + if preferred.is_empty() { + tracing::warn!( + "preferred_sub_cluster has no eligible remote kvclients: source_node_id={} sub_cluster={:?}", + source_node_id, + sc + ); + return global; + } + + if is_role_aware_policy(config.policy) + && !preferred + .iter() + .any(|candidate| candidate.is_remote_only_role) + && global.iter().any(|candidate| candidate.is_remote_only_role) + { + return global; + } + + preferred +} + +fn sort_by_queue_score(candidates: &mut [PlacementCandidate]) { + candidates.sort_by(|a, b| { + a.queue_score() + .total_cmp(&b.queue_score()) + .then_with(|| { + b.preferred_sub_cluster_match + .cmp(&a.preferred_sub_cluster_match) + }) + .then_with(|| a.node_id.as_ref().cmp(b.node_id.as_ref())) + .then_with(|| { + a.allocator + .seg_device_id + .as_str() + .cmp(b.allocator.seg_device_id.as_str()) + }) + }); +} + +fn remote_only_first(candidates: Vec) -> Vec { + let has_remote_only = candidates + .iter() + .any(|candidate| candidate.is_remote_only_role); + if !has_remote_only { + return candidates; + } + candidates + .into_iter() + .filter(|candidate| candidate.is_remote_only_role) + .collect() +} + +fn order_queue_aware(mut candidates: Vec) -> Vec { + sort_by_queue_score(&mut candidates); + candidates +} + +fn order_weighted_role_aware( + candidates: Vec, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let mut candidates = remote_only_first(candidates); + candidates.sort_by(|a, b| { + let a_role_weight = if a.is_remote_only_role { + config.remote_only_shard_weight + } else { + 1.0 + }; + let b_role_weight = if b.is_remote_only_role { + config.remote_only_shard_weight + } else { + 1.0 + }; + (a.queue_score() / a_role_weight) + .total_cmp(&(b.queue_score() / b_role_weight)) + .then_with(|| b.is_remote_only_role.cmp(&a.is_remote_only_role)) + .then_with(|| { + b.preferred_sub_cluster_match + .cmp(&a.preferred_sub_cluster_match) + }) + .then_with(|| a.node_id.as_ref().cmp(b.node_id.as_ref())) + .then_with(|| { + a.allocator + .seg_device_id + .as_str() + .cmp(b.allocator.seg_device_id.as_str()) + }) + }); + candidates +} + +fn order_bounded_role_queue_aware( + candidates: Vec, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let Some(best_wait) = candidates + .iter() + .map(|candidate| candidate.queue_wait_ms()) + .min_by(|a, b| a.total_cmp(b)) + else { + return candidates; + }; + let max_wait = best_wait + config.role_queue_window_ms; + let mut eligible: Vec = candidates + .into_iter() + .filter(|candidate| candidate.queue_wait_ms() <= max_wait) + .collect(); + eligible = remote_only_first(eligible); + eligible.sort_by(|a, b| { + let a_weight = if a.is_remote_only_role { + config.remote_only_shard_weight + } else { + 1.0 + } - ((a.used_bytes as f64 / 4096.0) * 1e-6); + let b_weight = if b.is_remote_only_role { + config.remote_only_shard_weight + } else { + 1.0 + } - ((b.used_bytes as f64 / 4096.0) * 1e-6); + b_weight + .total_cmp(&a_weight) + .then_with(|| a.queue_score().total_cmp(&b.queue_score())) + .then_with(|| { + b.preferred_sub_cluster_match + .cmp(&a.preferred_sub_cluster_match) + }) + .then_with(|| a.node_id.as_ref().cmp(b.node_id.as_ref())) + .then_with(|| { + a.allocator + .seg_device_id + .as_str() + .cmp(b.allocator.seg_device_id.as_str()) + }) + }); + eligible +} + +fn average_queue_wait<'a>(candidates: impl Iterator) -> Option { + let mut sum = 0.0; + let mut count = 0usize; + for candidate in candidates { + sum += candidate.queue_wait_ms(); + count += 1; + } + (count > 0).then_some(sum / count as f64) +} + +fn filter_remote_imbalance( + candidates: Vec, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let Some(min_count) = candidates + .iter() + .map(|candidate| candidate.node_write_count) + .min() + else { + return candidates; + }; + let max_next = (min_count as f64 + 1.0) * config.role_max_shard_imbalance; + let filtered: Vec = candidates + .iter() + .filter(|candidate| candidate.node_write_count as f64 + 1.0 <= max_next) + .cloned() + .collect(); + if filtered.is_empty() { + candidates + } else { + filtered + } +} + +fn order_pressure_role_queue_aware( + candidates: Vec, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + let remote_candidates: Vec = candidates + .iter() + .filter(|candidate| candidate.is_remote_only_role) + .cloned() + .collect(); + if remote_candidates.is_empty() { + return order_queue_aware(candidates); + } + + let active_wait = average_queue_wait(candidates.iter().filter(|candidate| { + candidate.is_active_role || (!candidate.is_remote_only_role && !candidate.is_active_role) + })); + let remote_wait = average_queue_wait(remote_candidates.iter()); + if let (Some(active_wait), Some(remote_wait)) = (active_wait, remote_wait) { + let gap = active_wait - remote_wait; + let fabric_guard_ok = + config.role_fabric_guard_ms == 0.0 || gap <= config.role_fabric_guard_ms; + if gap >= config.role_pressure_gap_ms && fabric_guard_ok { + let mut remote_candidates = filter_remote_imbalance(remote_candidates, config); + sort_by_queue_score(&mut remote_candidates); + return remote_candidates; + } + } + + order_queue_aware(candidates) +} + +fn order_remote_candidates( + candidates: Vec, + config: &ReplicaTaskPlacementConfig, +) -> Vec { + match config.policy { + ReplicaTaskPlacementPolicyKind::LocalFirst | ReplicaTaskPlacementPolicyKind::Random => { + let mut candidates = candidates; + candidates.shuffle(&mut rand::thread_rng()); + candidates + } + ReplicaTaskPlacementPolicyKind::QueueAware => order_queue_aware(candidates), + ReplicaTaskPlacementPolicyKind::WeightedRoleAware => { + order_weighted_role_aware(candidates, config) + } + ReplicaTaskPlacementPolicyKind::BoundedRoleQueueAware => { + order_bounded_role_queue_aware(candidates, config) + } + ReplicaTaskPlacementPolicyKind::PressureRoleQueueAware => { + order_pressure_role_queue_aware(candidates, config) + } + } +} + +fn try_allocate_candidates( + candidates: Vec, + len: u64, + last_no_space_ctx: &mut Option, +) -> Option<(NodeID, Arc, Allocation)> { + for candidate in candidates { + *last_no_space_ctx = Some(( + candidate.node_id.as_ref().to_string(), + candidate.allocator.seg_device_id.clone(), + candidate.total_bytes, + candidate.free_bytes, + )); + if let Ok(allocation) = candidate.allocator.allocate(len) { + return Some((candidate.node_id, candidate.allocator, allocation)); + } + } + None +} + +fn choose_random_remote_target_with_allocator( + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + len: u64, +) -> Result<(NodeID, Arc, Allocation), KvError> { + let seg_manager = view.master_seg_manager(); + let mut last_no_space_ctx: Option = None; + + if let Some(sc) = preferred_sub_cluster { + let mut preferred_nodes: Vec = view + .cluster_manager() + .get_client_members() + .into_iter() + .filter_map(|m| (m.sub_cluster.as_deref() == Some(sc)).then_some(m.id.into())) + .collect(); + preferred_nodes.retain(|node_id| { + node_id.as_ref() != source_node_id.as_ref() && !excluded_nodes.contains(node_id) + }); + preferred_nodes.shuffle(&mut rand::thread_rng()); + for node_id in preferred_nodes { + let node_allocators = seg_manager.get_node_allocators(&node_id); + let Some(allocator) = node_allocators.choose(&mut rand::thread_rng()).cloned() else { + continue; + }; + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; + last_no_space_ctx = Some(( + node_id.as_ref().to_string(), + allocator.seg_device_id.clone(), + total, + free, + )); + if let Ok(allocation) = allocator.allocate(len) { + return Ok((node_id, allocator, allocation)); + } + } + } + + let all_segs = seg_manager.get_all_segments_allocator(); + let mut candidates: Vec<(NodeID, Arc)> = all_segs + .into_iter() + .filter_map(|(node_id, allocator)| { + if node_id.as_ref() == source_node_id.as_ref() || excluded_nodes.contains(&node_id) { + return None; + } + Some((node_id, allocator)) + }) + .collect(); + candidates.shuffle(&mut rand::thread_rng()); + for (node_id, allocator) in candidates { + let capacity = allocator.node_pool_capacity_snapshot(); + let total = capacity.active_capacity_bytes; + let free = capacity.available_capacity_bytes; + last_no_space_ctx = Some(( + node_id.as_ref().to_string(), + allocator.seg_device_id.clone(), + total, + free, + )); + if let Ok(allocation) = allocator.allocate(len) { + return Ok((node_id, allocator, allocation)); + } + } + + Err(no_space_error(last_no_space_ctx)) +} + +fn choose_random_remote_target( + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + len: u64, +) -> Result<(NodeID, Allocation), KvError> { + choose_random_remote_target_with_allocator( + view, + source_node_id, + excluded_nodes, + preferred_sub_cluster, + len, + ) + .map(|(node_id, _allocator, allocation)| (node_id, allocation)) +} + +pub struct ReplicaTaskPlacementPolicy { + config: ReplicaTaskPlacementConfig, +} + +impl ReplicaTaskPlacementPolicy { + pub fn new(config: ReplicaTaskPlacementConfig) -> Self { + Self { config } + } + + fn select_remote_target_with_allocator( + &self, + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + len: u64, + ) -> Result<(NodeID, Arc, Allocation), KvError> { + let mut last_no_space_ctx: Option = None; + let candidates = choose_candidate_pool( + view, + source_node_id, + excluded_nodes, + preferred_sub_cluster, + &self.config, + ); + let ordered = order_remote_candidates(candidates, &self.config); + if let Some(selected) = try_allocate_candidates(ordered, len, &mut last_no_space_ctx) { + return Ok(selected); + } + + if preferred_sub_cluster.is_some() { + let global_candidates = + choose_candidate_pool(view, source_node_id, excluded_nodes, None, &self.config); + let ordered = order_remote_candidates(global_candidates, &self.config); + if let Some(selected) = try_allocate_candidates(ordered, len, &mut last_no_space_ctx) { + return Ok(selected); + } + } + + Err(no_space_error(last_no_space_ctx)) + } +} + +#[async_trait] +impl PlacementPolicy for ReplicaTaskPlacementPolicy { + async fn select_put_target( + &self, + view: &MasterKvRouterView, + req_node_id: &NodeID, + preferred_sub_cluster: Option<&str>, + len: u64, + ) -> Result { + self.select_remote_target_with_allocator( + view, + req_node_id, + &HashSet::new(), + preferred_sub_cluster, + len, + ) + .map( + |(node_id, _allocator, allocation)| PutPlacementTarget::Remote { + node_id, + allocation, + }, + ) + } + + fn select_remote_target( + &self, + view: &MasterKvRouterView, + source_node_id: &NodeID, + excluded_nodes: &HashSet, + preferred_sub_cluster: Option<&str>, + len: u64, + ) -> Result<(NodeID, Allocation), KvError> { + self.select_remote_target_with_allocator( + view, + source_node_id, + excluded_nodes, + preferred_sub_cluster, + len, + ) + .map(|(node_id, _allocator, allocation)| (node_id, allocation)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{ReplicaTaskPlacementConfig, ReplicaTaskPlacementPolicyKind}; + use crate::master_seg_manager::msg_pack::SegmentDeviceDescription; + + fn test_allocator(id: &str) -> Arc { + Arc::new( + OneSegAllocator::new( + id.to_string(), + SegmentDeviceDescription::Cpu, + 0, + 1024 * 1024, + ) + .unwrap(), + ) + } + + fn candidate( + node_id: &str, + is_remote_only_role: bool, + node_write_count: u64, + ) -> PlacementCandidate { + PlacementCandidate { + node_id: node_id.to_string().into(), + allocator: test_allocator(node_id), + total_bytes: 1024 * 1024, + free_bytes: 1024 * 1024, + used_bytes: 0, + node_write_count, + requester_target_count: 0, + is_remote_only_role, + is_active_role: !is_remote_only_role, + preferred_sub_cluster_match: false, + } + } + + #[test] + fn bounded_role_queue_aware_prefers_remote_only_within_window() { + let mut config = ReplicaTaskPlacementConfig::default(); + config.policy = ReplicaTaskPlacementPolicyKind::BoundedRoleQueueAware; + config.role_queue_window_ms = 2.0; + config.remote_only_shard_weight = 1.02; + + let ordered = order_bounded_role_queue_aware( + vec![ + candidate("active-a", false, 0), + candidate("remote-a", true, 1), + ], + &config, + ); + assert_eq!(ordered[0].node_id.as_ref(), "remote-a"); + } + + #[test] + fn strict_remote_only_candidates_filter_out_active_nodes() { + let filtered = filter_remote_only_candidates(&[ + candidate("prefill-a", false, 0), + candidate("decode-a", false, 0), + candidate("remote-cache-a", true, 10), + ]); + + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].node_id.as_ref(), "remote-cache-a"); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs index 70d8858..4ce99a6 100755 --- a/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/put.rs @@ -1,55 +1,210 @@ -use super::NodeValueReplicaDesc; use super::{ - InflightPutAllocation, InflightPutInfo, KvRouteInfo, MasterKvRouterView, PutPlacementMode, - msg_pack::{PutDoneReq, PutDoneResp, PutRevokeReq, PutRevokeResp, PutStartReq, PutStartResp}, + CommittedSlotReplica, CompletedReplicaTaskInfo, InflightPutAllocation, InflightPutCommitInfo, + InflightPutInfo, InflightReplicaTaskInfo, KvMemoryReplica, KvNodeReplicas, KvReplicaBacking, + LocalReserveGrantInfo, MasterKeyActivityCompletionGuard, MasterKvRouterView, + NodeCacheCapacityReservation, OwnerHoldingGetInfo, PreparedPutKeyReservationInfo, + PutPlacementMode, ReservedCapacityReason, SsdReplicaCommitStatus, + msg_pack::{ + BatchPreparePutKeysReq, BatchPreparePutKeysResp, BatchPublishOwnerSsdReq, + BatchPublishOwnerSsdResp, BatchPutAppendDoneItemResp, BatchPutAppendDoneReq, + BatchPutAppendDoneResp, BatchPutAppendStartItemResp, BatchPutAppendStartReq, + BatchPutAppendStartResp, BatchPutDoneItemResp, BatchPutDoneReq, BatchPutDoneResp, + BatchPutRevokeItemResp, BatchPutRevokeReq, BatchPutRevokeResp, BatchPutStartItemResp, + BatchPutStartReq, BatchPutStartResp, BatchReleasePutKeyReservationsReq, + BatchReleasePutKeyReservationsResp, GroupedBatchPutDoneReq, GroupedBatchPutDoneResp, + OwnerSsdPublishItem, OwnerSsdPublishItemResp, OwnerSsdPublishOutcome, PutAppendDoneReq, + PutAppendDoneResp, PutAppendRevokeReq, PutAppendRevokeResp, PutAppendStartOutcome, + PutAppendStartReq, PutAppendStartResp, PutAtomicGroup, PutDoneCommittedSlot, PutDoneReq, + PutDoneResp, PutRevokeReq, PutRevokeResp, PutStartReq, PutStartResp, ReleaseLocalGrantReq, + ReleaseLocalGrantResp, ReserveLocalGrantOutcome, ReserveLocalGrantReq, + ReserveLocalGrantResp, build_shared_put_atomic_group_assignments, + }, + node_generation_is_current_live, placement::PutPlacementTarget, + publish_primary_route_tomb_fenced, publish_route_replica_tomb_fenced, + route_maintenance::{ + RoutePublishEvent, apply_post_route_maintenance_batch, enqueue_post_route_maintenance, + }, }; use crate::master_kv_router::OneKvNodesRoutes; use crate::master_kv_router::delete::DeleteKeyInfo; +use crate::memholder::MemholderManagerTrait; use crate::{ - cluster_manager::{META_KEY_LOCAL_IPC_ROOT, NodeID}, - master_seg_manager::one_seg_allocator::Allocation, + OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES, + cluster_manager::{ + META_KEY_LOCAL_IPC_ROOT, META_KEY_SHARED_STORAGE_NODE_ID, + META_KEY_SHARED_STORAGE_NODE_START_TIME, NodeID, + }, + master_seg_manager::{MasterSegManagerAccessTrait, one_seg_allocator::Allocation}, p2p::msg_pack::MsgPack, - rpcresp_kvresult_convert::msg_and_error, + rpcresp_kvresult_convert::msg_and_error::{self, kv}, }; -use fluxon_commu::{META_KEY_SHARED_STORAGE_NODE_ID, META_KEY_SHARED_STORAGE_NODE_START_TIME}; +use chrono::Utc; +use limit_thirdparty::tokio; use parking_lot::Mutex; use parking_lot::RwLock; use rand::seq::SliceRandom; use std::{ - collections::HashMap, - sync::{Arc, atomic::AtomicU32}, + collections::{HashMap, HashSet}, + sync::{ + Arc, + atomic::{AtomicU32, Ordering}, + }, }; pub type PutIDForAKey = (u64, u32); -struct InflightPutKeyReservation { - view: MasterKvRouterView, - key: String, - active: bool, -} - -impl InflightPutKeyReservation { - fn new(view: MasterKvRouterView, key: String) -> Self { - Self { - view, - key, - active: true, +fn publish_owner_ssd_item( + view: &MasterKvRouterView, + owner: &NodeID, + item: OwnerSsdPublishItem, +) -> OwnerSsdPublishItemResp { + let response = |outcome, detail: String| OwnerSsdPublishItemResp { + key: item.key.clone(), + put_id: item.put_id, + outcome, + detail, + }; + let _activity = match view + .master_kv_router() + .reserve_inflight_replica_key(&item.key) + { + Ok(activity) => activity, + Err(err) => { + return response( + OwnerSsdPublishOutcome::RetryableBusy, + format!("master key activity is busy: {err}"), + ); } + }; + + let Some(route) = view.master_kv_router().inner().kv_routes.get(&item.key) else { + return response( + OwnerSsdPublishOutcome::Obsolete, + "route is absent".to_string(), + ); + }; + if route.put_id != item.put_id { + return response( + OwnerSsdPublishOutcome::Obsolete, + format!( + "route generation changed: current=({},{})", + route.put_id.0, route.put_id.1 + ), + ); } - fn disarm(&mut self) { - self.active = false; + let has_remote_memory = { + let replicas = route.node_replicas.read(); + let Some(owner_replicas) = replicas.get(owner) else { + return response( + OwnerSsdPublishOutcome::Obsolete, + "same-owner route entry is absent".to_string(), + ); + }; + if owner_replicas.tomb_tag.is_tomb() { + return response( + OwnerSsdPublishOutcome::Obsolete, + "same-owner generation is tombed".to_string(), + ); + } + if let Some(existing) = owner_replicas.ssd.as_ref() { + return if existing.len == item.len { + response( + OwnerSsdPublishOutcome::AlreadyPresent, + "same-owner SSD backing is already published".to_string(), + ) + } else { + response( + OwnerSsdPublishOutcome::Rejected, + format!( + "existing SSD length mismatch: existing={} requested={}", + existing.len, item.len + ), + ) + }; + } + replicas.iter().any(|(node_id, replicas)| { + node_id != owner && !replicas.tomb_tag.is_tomb() && replicas.memory.is_some() + }) + }; + + match route.commit_ssd_replica(owner, item.len) { + SsdReplicaCommitStatus::Committed => { + let counters = &view.master_kv_router().inner().ssd_tier_counters; + let (items, bytes) = if has_remote_memory { + ( + &counters.local_ssd_published_with_remote_memory_items, + &counters.local_ssd_published_with_remote_memory_bytes, + ) + } else { + ( + &counters.local_ssd_published_without_remote_memory_items, + &counters.local_ssd_published_without_remote_memory_bytes, + ) + }; + items.fetch_add(1, Ordering::Relaxed); + bytes.fetch_add(item.len, Ordering::Relaxed); + response( + OwnerSsdPublishOutcome::Published, + "same-owner SSD backing published while memory remains live".to_string(), + ) + } + SsdReplicaCommitStatus::MissingMemory | SsdReplicaCommitStatus::TombedNode => response( + OwnerSsdPublishOutcome::Obsolete, + "exact live same-owner memory generation is absent".to_string(), + ), + SsdReplicaCommitStatus::LengthMismatch => response( + OwnerSsdPublishOutcome::Rejected, + format!( + "SSD length does not match same-owner memory: requested={}", + item.len + ), + ), } } -impl Drop for InflightPutKeyReservation { - fn drop(&mut self) { - if self.active { - self.view - .master_kv_router() - .release_inflight_put_key(&self.key); - } +/// Publish durable same-owner SSD bytes onto the current route without +/// deleting or replacing the memory backing. +pub async fn handle_batch_publish_owner_ssd( + view: MasterKvRouterView, + req: MsgPack, + owner: NodeID, +) -> MsgPack { + let current_generation = view + .cluster_manager() + .get_member_info_cached(owner.as_ref()) + .map(|member| member.node_start_time); + if current_generation != Some(req.serialize_part.owner_node_start_time) { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "owner SSD publication generation mismatch: owner={} requested={} current={:?}", + owner, req.serialize_part.owner_node_start_time, current_generation + ), + }); + return MsgPack { + serialize_part: BatchPublishOwnerSsdResp { + items: Vec::new(), + error_code: err.code(), + error_json: err.to_json(), + }, + raw_bytes: Vec::new(), + }; + } + + let items = req + .serialize_part + .items + .into_iter() + .map(|item| publish_owner_ssd_item(&view, &owner, item)) + .collect(); + MsgPack { + serialize_part: BatchPublishOwnerSsdResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), } } @@ -153,26 +308,523 @@ fn validate_put_start_source_node_override( Ok(()) } +fn current_route_append_outcome( + route: &OneKvNodesRoutes, + source_node_id: &NodeID, + verify_put_id: PutIDForAKey, +) -> PutAppendStartOutcome { + let has_complete_remote = route + .node_replicas + .read() + .iter() + .any(|(node_id, replicas)| node_id != source_node_id && replicas.has_live_backing()); + classify_put_append_start_outcome(route.put_id == verify_put_id, has_complete_remote) +} + +fn classify_put_append_start_outcome( + current_identity: bool, + has_complete_remote: bool, +) -> PutAppendStartOutcome { + if !current_identity { + PutAppendStartOutcome::Obsolete + } else if has_complete_remote { + PutAppendStartOutcome::AlreadySatisfied + } else { + PutAppendStartOutcome::Scheduled + } +} + +#[cfg(test)] +mod append_start_outcome_tests { + use super::{PutAppendStartOutcome, classify_put_append_start_outcome}; + + #[test] + fn append_start_never_conflates_no_remote_with_already_satisfied() { + assert_eq!( + classify_put_append_start_outcome(true, false), + PutAppendStartOutcome::Scheduled + ); + assert_eq!( + classify_put_append_start_outcome(true, true), + PutAppendStartOutcome::AlreadySatisfied + ); + assert_eq!( + classify_put_append_start_outcome(false, true), + PutAppendStartOutcome::Obsolete + ); + assert_ne!( + PutAppendStartOutcome::RetryableNoSpace, + PutAppendStartOutcome::AlreadySatisfied + ); + } +} + +fn append_current_route_replica_if_matching( + view: &MasterKvRouterView, + key: &str, + put_id: PutIDForAKey, + node_id: NodeID, + target_tomb_tag: crate::master_seg_manager::NodeTombTag, + allocation: Allocation, +) -> Option { + let Some(one_kv_nodes_routes) = view.master_kv_router().inner().kv_routes.get(key) else { + tracing::debug!( + "append_current_route_replica_if_matching skipped because route disappeared: key={} put_id=({},{})", + key, + put_id.0, + put_id.1 + ); + return None; + }; + if one_kv_nodes_routes.put_id != put_id { + tracing::debug!( + "append_current_route_replica_if_matching skipped because version changed: key={} current_put_id=({},{}) append_put_id=({},{})", + key, + one_kv_nodes_routes.put_id.0, + one_kv_nodes_routes.put_id.1, + put_id.0, + put_id.1 + ); + return None; + } + if !node_generation_is_current_live(view, &node_id, &target_tomb_tag) { + tracing::warn!( + "append_current_route_replica_if_matching skipped because target generation departed: key={} put_id=({},{}) node_id={}", + key, + put_id.0, + put_id.1, + node_id + ); + return None; + } + let capacity_bytes = allocation.capcity(); + let lease_id = one_kv_nodes_routes.lease_id; + let capacity_reservation = match lease_id { + Some(_) => match view.master_kv_router().reserve_node_cache_capacity( + &node_id, + &target_tomb_tag, + ReservedCapacityReason::LeaseBoundKv, + capacity_bytes, + ) { + Ok(reservation) => reservation, + Err(err) => { + tracing::warn!( + "append_current_route_replica_if_matching could not reserve lease-bound capacity: key={} put_id=({},{}) node_id={} err={}", + key, + put_id.0, + put_id.1, + node_id, + err, + ); + return None; + } + }, + None => None, + }; + let published = publish_route_replica_tomb_fenced( + &one_kv_nodes_routes, + node_id.clone(), + KvMemoryReplica { + backing: KvReplicaBacking::Allocation(Arc::new(allocation)), + owner_local_indexed: false, + get_durable_reservation: None, + capacity_reservation, + }, + target_tomb_tag, + ); + if !published { + tracing::warn!( + "append_current_route_replica_if_matching rejected by generation/live-replica fence: key={} put_id=({},{}) node_id={}", + key, + put_id.0, + put_id.1, + node_id + ); + return None; + } + Some(RoutePublishEvent::replica_append( + key.to_string(), + put_id, + lease_id, + node_id, + capacity_bytes, + )) +} + +fn allocate_from_node_local_segment( + view: &MasterKvRouterView, + node_id: &NodeID, + len: u64, + op_name: &str, +) -> msg_and_error::KvResult { + let node_allocators = view.master_seg_manager().get_node_allocators(node_id); + if node_allocators.is_empty() { + tracing::warn!("No allocators found for {} node={}", op_name, node_id); + return Err(msg_and_error::KvError::Api( + msg_and_error::ApiError::RegisterSegmentFailed { + detail: format!( + "{} node has no registered segments: node={}", + op_name, node_id + ), + }, + )); + } + + let allocator = node_allocators.choose(&mut rand::thread_rng()).unwrap(); + for attempt in 1..=3 { + if let Ok(allocation) = allocator.allocate(len) { + return Ok(allocation); + } + tracing::warn!( + "Allocation attempt {}/3 failed for {} node={} len={}", + attempt, + op_name, + node_id, + len + ); + } + + let capacity = allocator.node_pool_capacity_snapshot(); + Err(msg_and_error::KvError::Api( + msg_and_error::ApiError::NoSpace { + node: node_id.as_ref().to_string(), + segment: allocator.seg_device_id.clone(), + total_capacity: capacity.active_capacity_bytes, + free_capacity: capacity.available_capacity_bytes, + }, + )) +} + +fn validate_put_done_committed_slot( + view: &MasterKvRouterView, + node_id: &NodeID, + slot: &PutDoneCommittedSlot, + expected_tomb_tag: Option<&crate::master_seg_manager::NodeTombTag>, +) -> msg_and_error::KvResult<(CommittedSlotReplica, crate::master_seg_manager::NodeTombTag)> { + let invalid = |detail: String| { + msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { detail }) + }; + if slot.slot_size == 0 || slot.len > slot.slot_size { + return Err(invalid(format!( + "invalid committed local slot geometry: grant_id={} slot_size={} len={}", + slot.grant_id, slot.slot_size, slot.len + ))); + } + let Some(grant) = view + .master_kv_router() + .inner() + .local_reserve_grants + .get(&slot.grant_id) + else { + return Err(invalid(format!( + "committed local slot references unknown grant_id={}", + slot.grant_id + ))); + }; + if grant.owner_node_id != *node_id { + return Err(invalid(format!( + "committed local slot owner mismatch: grant_id={} owner={} requester={}", + slot.grant_id, grant.owner_node_id, node_id + ))); + } + let tomb_tag = grant.tomb_tag.clone(); + if !node_generation_is_current_live(view, node_id, &tomb_tag) + || expected_tomb_tag.is_some_and(|expected| !expected.same_generation(&tomb_tag)) + { + return Err(invalid(format!( + "committed local slot belongs to a different/departed owner generation: grant_id={} requester={}", + slot.grant_id, node_id + ))); + } + + let grant_base = grant.allocation.base_addr(); + let grant_addr = grant_base + .checked_add(grant.allocation.addr()) + .ok_or_else(|| invalid("committed local slot grant address overflow".to_string()))?; + let slot_offset = slot + .slot_size + .checked_mul(u64::from(slot.slot_index)) + .ok_or_else(|| invalid("committed local slot offset overflow".to_string()))?; + let slot_end = slot_offset + .checked_add(slot.slot_size) + .ok_or_else(|| invalid("committed local slot end overflow".to_string()))?; + if slot_end > grant.allocation.capcity() { + return Err(invalid(format!( + "committed local slot is outside grant: grant_id={} slot_index={} slot_size={} grant_len={}", + slot.grant_id, + slot.slot_index, + slot.slot_size, + grant.allocation.capcity() + ))); + } + let expected_addr = grant_addr + .checked_add(slot_offset) + .ok_or_else(|| invalid("committed local slot address overflow".to_string()))?; + if slot.base_addr != grant_base || slot.addr != expected_addr { + return Err(invalid(format!( + "committed local slot address mismatch: grant_id={} expected_base={:#x} got_base={:#x} expected_addr={:#x} got_addr={:#x}", + slot.grant_id, grant_base, slot.base_addr, expected_addr, slot.addr + ))); + } + + Ok(( + CommittedSlotReplica { + owner_node_id: node_id.clone(), + grant_id: slot.grant_id, + slot_index: slot.slot_index, + slot_size: slot.slot_size, + addr: slot.addr, + len: slot.len, + base_addr: slot.base_addr, + }, + tomb_tag, + )) +} + +async fn prepare_route_state( + view: &MasterKvRouterView, + lease_id: Option, + key: &str, + put_id: PutIDForAKey, + node_id: &NodeID, + tomb_tag: &crate::master_seg_manager::NodeTombTag, + reservation_reason: Option, + target_cap_bytes: u64, +) -> msg_and_error::KvResult>> { + // Reserve first. If lease attachment fails, dropping this local token + // restores the exact generation-scoped counter automatically. Committed + // slots do not reserve again because their whole grant is already excluded. + let reservation = match reservation_reason { + Some(reason) => view.master_kv_router().reserve_node_cache_capacity( + node_id, + tomb_tag, + reason, + target_cap_bytes, + )?, + None => None, + }; + if let Some(lease_id) = lease_id { + view.master_lease_manager() + .attach_key(lease_id, key.to_string(), put_id) + .await + .map_err(|err| -> msg_and_error::KvError { err.into() })?; + } + Ok(reservation) +} + +fn reserve_replica_task( + view: &MasterKvRouterView, + key: &str, + put_id: PutIDForAKey, + source_node_id: &NodeID, + preferred_sub_cluster: Option<&str>, + len: u64, +) -> msg_and_error::KvResult { + reserve_replica_task_excluding( + view, + key, + put_id, + source_node_id, + preferred_sub_cluster, + len, + &HashSet::new(), + true, + ) +} + +fn reserve_replica_task_excluding( + view: &MasterKvRouterView, + key: &str, + put_id: PutIDForAKey, + source_node_id: &NodeID, + preferred_sub_cluster: Option<&str>, + len: u64, + excluded_nodes: &HashSet, + protect_source_on_remote_complete: bool, +) -> msg_and_error::KvResult { + let activity_lease = view.master_kv_router().reserve_inflight_replica_key(key)?; + let operation_id = view + .master_kv_router() + .inner() + .next_replica_operation_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + assert_ne!( + operation_id, 0, + "master replica operation identifier overflow" + ); + view.master_kv_router() + .pin_current_master_cache_identity_for_activity( + &activity_lease, + source_node_id.as_ref(), + key, + put_id, + ); + let (target_node_id, target_allocation) = view + .master_kv_router() + .inner() + .policy + .select_remote_target( + view, + source_node_id, + excluded_nodes, + preferred_sub_cluster, + len, + )?; + let Some(target_tomb_tag) = view + .master_seg_manager() + .get_allocation_tomb_tag(&target_node_id, &target_allocation) + else { + return Err(msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "replica target generation changed during reservation: key={} put_id=({},{}) target_node_id={}", + key, put_id.0, put_id.1, target_node_id + ), + }, + )); + }; + tracing::debug!( + "replica task reserved: key={} put_id=({},{}) operation_id={} source_node_id={} target_node_id={} preferred_sub_cluster={:?} len={}", + key, + put_id.0, + put_id.1, + operation_id, + source_node_id, + target_node_id, + preferred_sub_cluster, + len + ); + Ok(InflightReplicaTaskInfo { + operation_id, + node_id: target_node_id, + target_tomb_tag, + source_node_id: source_node_id.clone(), + key: key.to_string(), + put_id, + target_allocation: Arc::new(Mutex::new(Some(target_allocation))), + protect_source_on_remote_complete, + _activity_lease: activity_lease, + }) +} + +async fn publish_completed_put_route( + view: MasterKvRouterView, + key: String, + put_id: PutIDForAKey, + lease_id_opt: Option, + atomic_group: Option>, + node_id: NodeID, + publish_tag: crate::master_seg_manager::NodeTombTag, + completed_info: KvMemoryReplica, + target_cap_bytes: u64, + local_cache_holder_id: Option, +) -> MsgPack { + let new_route = Arc::new(OneKvNodesRoutes { + put_id, + lease_id: lease_id_opt, + atomic_group: atomic_group.clone(), + node_replicas: RwLock::new(HashMap::from([( + node_id.clone(), + KvNodeReplicas::memory(publish_tag.clone(), completed_info), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let old_one_kv_routes = match publish_primary_route_tomb_fenced( + &view.master_kv_router().inner().kv_routes, + &key, + new_route.clone(), + &publish_tag, + ) { + Ok(previous) => previous, + Err(()) => { + if let Some(lease_id) = lease_id_opt { + view.master_lease_manager() + .detach_key_if_version(lease_id, &key, put_id); + } + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "primary route publication rejected because target generation departed: key={} put_id=({},{}) node_id={}", + key, put_id.0, put_id.1, node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + }; + + if let Some(old) = old_one_kv_routes { + view.master_kv_router() + .remove_route_cache_entries_exact(&key, &old) + .await; + if let Err(err) = view + .master_kv_router() + .inner() + .delete_broadcast + .sender() + .send(DeleteKeyInfo::Key { + key: key.clone(), + nodes_kv_route_info: old, + }) + .await + { + tracing::warn!("Failed to send delete broadcast: {}", err); + } + } + + apply_post_route_maintenance_batch( + &view, + vec![RoutePublishEvent::primary_put( + key.clone(), + put_id, + lease_id_opt, + node_id.clone(), + target_cap_bytes, + )], + ) + .await; + + tracing::debug!( + "Completed put operation with put_id: {:?}, key: {:?}", + put_id, + key + ); + + MsgPack { + serialize_part: PutDoneResp { + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + local_cache_holder_id, + }, + raw_bytes: Vec::new(), + } +} + pub async fn handle_put_start( view: MasterKvRouterView, req: MsgPack, req_node_id: NodeID, ) -> (PutIDForAKey, MsgPack) { let key = req.serialize_part.key.clone(); - if let Err(err) = view - .master_kv_router() - .reserve_inflight_put_key(&key, req.serialize_part.reject_if_inflight_same_key) - { - let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); - return ( - (0, 0), - MsgPack { - serialize_part: resp, - raw_bytes: Vec::new(), - }, - ); - } - let mut key_reservation = InflightPutKeyReservation::new(view.clone(), key.clone()); + let activity_lease = match view.master_kv_router().reserve_inflight_put_key( + &key, + req.serialize_part.reject_if_inflight_same_key, + req.serialize_part.reject_if_exist_same_key, + ) { + Ok(activity_lease) => activity_lease, + Err(err) => { + let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + (0, 0), + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + }; let source_node_id = match req.serialize_part.source_node_id.as_ref() { Some(source_node_id) => { let source_node_id: NodeID = source_node_id.clone().into(); @@ -199,58 +851,14 @@ pub async fn handle_put_start( let inflight_put_key: (String, u64, u32) = (key.clone(), put_id.0, put_id.1); - // randomly select one src_allocator - let src_allocation = { - let src_node_allocators = view - .master_seg_manager() - .get_node_allocators(&source_node_id); - if src_node_allocators.is_empty() { - tracing::warn!( - "No allocators found for put_start source node: requester={} source={}", - req_node_id, - source_node_id - ); - let err = msg_and_error::KvError::Api(msg_and_error::ApiError::RegisterSegmentFailed { - detail: format!( - "put_start source node has no registered segments: requester={} source={}", - req_node_id, source_node_id - ), - }); - let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); - return ( - (0, 0), - MsgPack { - serialize_part: resp, - raw_bytes: Vec::new(), - }, - ); - } - - let src_allocator = src_node_allocators.choose(&mut rand::thread_rng()).unwrap(); - - let mut allocated_addr: Option = None; - for attempt in 1..=3 { - if let Ok(allocation) = src_allocator.allocate(req.serialize_part.len) { - allocated_addr = Some(allocation); - break; - } else { - tracing::warn!( - "Allocation attempt {}/3 failed for put_id {:?}", - attempt, - put_id - ); - } - } - if allocated_addr.is_none() { - let total = src_allocator.total_size_bytes(); - let used = src_allocator.used_size_bytes(); - let free = total.saturating_sub(used); - let err = msg_and_error::KvError::Api(msg_and_error::ApiError::NoSpace { - node: source_node_id.as_ref().to_string(), - segment: src_allocator.seg_device_id.clone(), - total_capacity: total, - free_capacity: free, - }); + let src_allocation = match allocate_from_node_local_segment( + &view, + &source_node_id, + req.serialize_part.len, + "put_start source", + ) { + Ok(allocation) => allocation, + Err(err) => { let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); return ( (0, 0), @@ -260,25 +868,32 @@ pub async fn handle_put_start( }, ); } - allocated_addr.unwrap() }; // Keep src allocation alive across retry attempts until we have a successful target. let mut src_allocation = Some(src_allocation); - let finalize = |node_id: NodeID, + let finalize = |commit_node_id: NodeID, + target_tomb_tag: crate::master_seg_manager::NodeTombTag, + response_node_id: NodeID, inflight_alloc: InflightPutAllocation, src_addr: u64, target_addr: u64, src_base_addr: u64, target_base_addr: u64, - len: u64| { + len: u64, + replica_target: Option| { let info = InflightPutInfo { - node_id: node_id.clone(), key: key.clone(), len, req_node_id: req_node_id.clone(), - src_target_allocation: Arc::new(Mutex::new(Some(inflight_alloc))), + commit_info: InflightPutCommitInfo { + node_id: commit_node_id, + target_tomb_tag, + src_target_allocation: Arc::new(Mutex::new(Some(inflight_alloc))), + replica_target: replica_target.clone(), + }, + _activity_lease: activity_lease.clone(), }; let view_task = view.clone(); @@ -291,9 +906,22 @@ pub async fn handle_put_start( .insert(inflight_put_key, info) .await; + let response_replica_target = replica_target.as_ref().map(|target| { + let target_allocation_guard = target.target_allocation.lock(); + let target_allocation = target_allocation_guard.as_ref().expect( + "replica target allocation must exist while building put_start response", + ); + super::msg_pack::PutReplicaTarget { + node_id: target.node_id.clone().into(), + target_addr: target_allocation.base_addr() + target_allocation.addr(), + target_base_addr: target_allocation.base_addr(), + len: target_allocation.size(), + } + }); + let resp = PutStartResp { put_id, - node_id: node_id.into(), + node_id: response_node_id.into(), src_addr, target_addr, src_base_addr, @@ -302,6 +930,7 @@ pub async fn handle_put_start( error_code: msg_and_error::OK, error_json: String::new(), server_process_us: 0, + replica_target: response_replica_target, }; ( @@ -314,17 +943,22 @@ pub async fn handle_put_start( } }; - let put_target = view - .master_kv_router() - .inner() - .policy - .select_put_target( - &view, - &source_node_id, - req.serialize_part.preferred_sub_cluster.as_deref(), - req.serialize_part.len, - ) - .await; + let put_target = if req.serialize_part.make_replica_task { + Ok(PutPlacementTarget::Local { + node_id: source_node_id.clone(), + }) + } else { + view.master_kv_router() + .inner() + .policy + .select_put_target( + &view, + &source_node_id, + req.serialize_part.preferred_sub_cluster.as_deref(), + req.serialize_part.len, + ) + .await + }; match put_target { Ok(PutPlacementTarget::Local { node_id }) => { @@ -362,18 +996,91 @@ pub async fn handle_put_start( let src = src_allocation .take() .expect("src_allocation must exist when finalizing local put"); - let fut = finalize( - node_id, - InflightPutAllocation::Local(src), - abs, - abs, - src_base, - src_base, - allocation_size, + let Some(target_tomb_tag) = view + .master_seg_manager() + .get_allocation_tomb_tag(&node_id, &src) + else { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "local put target generation changed before start publication: key={} put_id=({},{}) node_id={}", + key, put_id.0, put_id.1, node_id + ), + }, + ); + return ( + (0, 0), + MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }, + ); + }; + let replica_target = if req.serialize_part.make_replica_task { + match reserve_replica_task( + &view, + &key, + put_id, + &source_node_id, + req.serialize_part.preferred_sub_cluster.as_deref(), + req.serialize_part.len, + ) { + Ok(reservation) => { + view.master_kv_router() + .record_replica_task_target(reservation.node_id.as_ref()); + Some(reservation) + } + Err(msg_and_error::KvError::Api(msg_and_error::ApiError::NoSpace { + node, + segment, + total_capacity, + free_capacity, + })) => { + tracing::info!( + "replica task not pre-reserved; local-only commit remains valid: key={} put_id=({},{}) source_node_id={} preferred_sub_cluster={:?} node={} segment={} total_capacity={} free_capacity={}", + key, + put_id.0, + put_id.1, + source_node_id, + req.serialize_part.preferred_sub_cluster, + node, + segment, + total_capacity, + free_capacity + ); + None + } + Err(err) => { + tracing::warn!( + "replica task pre-reserve failed; local-only commit remains valid: key={} put_id=({},{}) source_node_id={} preferred_sub_cluster={:?} err={}", + key, + put_id.0, + put_id.1, + source_node_id, + req.serialize_part.preferred_sub_cluster, + err + ); + None + } + } + } else { + None + }; + let fut = finalize( + node_id.clone(), + target_tomb_tag, + node_id, + InflightPutAllocation::Local(src), + abs, + abs, + src_base, + src_base, + allocation_size, + replica_target, ); - let result = fut.await; - key_reservation.disarm(); - return result; + return fut.await; } Ok(PutPlacementTarget::Remote { node_id, @@ -412,7 +1119,31 @@ pub async fn handle_put_start( let src = src_allocation .take() .expect("src_allocation must exist when finalizing remote put"); + let Some(target_tomb_tag) = view + .master_seg_manager() + .get_allocation_tomb_tag(&node_id, &target_allocation) + else { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "remote put target generation changed before start publication: key={} put_id=({},{}) node_id={}", + key, put_id.0, put_id.1, node_id + ), + }, + ); + return ( + (0, 0), + MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }, + ); + }; let fut = finalize( + node_id.clone(), + target_tomb_tag, node_id, InflightPutAllocation::Remote { src, @@ -423,10 +1154,9 @@ pub async fn handle_put_start( src_base, target_base, allocation_size, + None, ); - let result = fut.await; - key_reservation.disarm(); - return result; + return fut.await; } Err(err) => { let resp: PutStartResp = crate::rpcresp_kvresult_convert::FromError::from_error(&err); @@ -441,6 +1171,290 @@ pub async fn handle_put_start( } } +pub async fn handle_reserve_local_grant( + view: MasterKvRouterView, + _req: MsgPack, + req_node_id: NodeID, +) -> (u64, MsgPack) { + let allocation = match allocate_from_node_local_segment( + &view, + &req_node_id, + OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES, + "reserve_local_grant", + ) { + Ok(allocation) => allocation, + Err(err) => { + let resp: ReserveLocalGrantResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + 0, + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + }; + + let grant_id = view.master_kv_router().next_local_reserve_grant_id(); + let grant_base_addr = allocation.base_addr(); + let grant_abs_addr = grant_base_addr + allocation.addr(); + let grant_len = allocation.capcity(); + let Some(tomb_tag) = view + .master_seg_manager() + .get_allocation_tomb_tag(&req_node_id, &allocation) + else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "local reserve grant generation changed before publication: owner={}", + req_node_id + ), + }); + let resp: ReserveLocalGrantResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + 0, + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + }; + let capacity_reservation = match view.master_kv_router().reserve_node_cache_capacity( + &req_node_id, + &tomb_tag, + ReservedCapacityReason::LocalReserveGrant, + grant_len, + ) { + Ok(reservation) => reservation, + Err(err) => { + let resp: ReserveLocalGrantResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + 0, + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + }; + view.master_kv_router().install_local_reserve_grant( + grant_id, + LocalReserveGrantInfo { + owner_node_id: req_node_id.clone(), + tomb_tag: tomb_tag.clone(), + allocation, + capacity_reservation, + }, + ); + if tomb_tag.is_tomb() { + // MemberLeft may mark the shared tag between allocation and map + // insertion. Remove the exact grant before returning so the old + // allocator cannot be pinned after its cleanup snapshot. + drop(view.master_kv_router().take_local_reserve_grant(grant_id)); + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "local reserve grant owner departed during publication: owner={} grant_id={}", + req_node_id, grant_id + ), + }); + let resp: ReserveLocalGrantResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + 0, + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + + ( + grant_id, + MsgPack { + serialize_part: ReserveLocalGrantResp { + outcome: ReserveLocalGrantOutcome::Granted { + grant_id, + node_id: req_node_id.into_owned(), + addr: grant_abs_addr, + base_addr: grant_base_addr, + len: grant_len, + }, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }, + ) +} + +pub async fn handle_release_local_grant( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let grant_id = req.serialize_part.grant_id; + let Some(grant) = view.master_kv_router().take_local_reserve_grant(grant_id) else { + tracing::info!( + "release_local_grant ignored missing grant_id={} requester_node_id={}", + grant_id, + req_node_id + ); + return MsgPack { + serialize_part: ReleaseLocalGrantResp::default(), + raw_bytes: Vec::new(), + }; + }; + + if grant.owner_node_id.as_ref() != req_node_id.as_ref() { + let owner_node_id = grant.owner_node_id.to_string(); + view.master_kv_router() + .install_local_reserve_grant(grant_id, grant); + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "release_local_grant owner mismatch: grant_id={} owner_node_id={} requester_node_id={}", + grant_id, owner_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + + if !node_generation_is_current_live(&view, &req_node_id, &grant.tomb_tag) { + // Do not reinsert an old-generation allocation: MemberLeft cleanup may + // already have taken its snapshot. Dropping it here is the exact + // generation-safe release. + drop(grant); + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "release_local_grant references a departed owner generation: grant_id={} requester_node_id={}", + grant_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + + drop(grant); + MsgPack { + serialize_part: ReleaseLocalGrantResp::default(), + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_prepare_put_keys( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> (Vec, MsgPack) { + let mut reservation_ids = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let activity_lease = match view.master_kv_router().reserve_inflight_put_key( + &item.key, + item.reject_if_inflight_same_key, + item.reject_if_exist_same_key, + ) { + Ok(activity_lease) => activity_lease, + Err(err) => { + for reservation_id in reservation_ids.drain(..) { + let _ = view + .master_kv_router() + .take_prepared_put_key_reservation(reservation_id); + } + let resp: BatchPreparePutKeysResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return ( + Vec::new(), + MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }, + ); + } + }; + let reservation_id = view + .master_kv_router() + .next_prepared_put_key_reservation_id(); + view.master_kv_router() + .install_prepared_put_key_reservation( + reservation_id, + PreparedPutKeyReservationInfo { + owner_node_id: req_node_id.clone(), + key: item.key, + _activity_lease: activity_lease, + }, + ); + reservation_ids.push(reservation_id); + } + + ( + reservation_ids.clone(), + MsgPack { + serialize_part: BatchPreparePutKeysResp { + reservation_ids, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }, + ) +} + +pub async fn handle_batch_release_put_key_reservations( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut taken = Vec::with_capacity(req.serialize_part.reservation_ids.len()); + for reservation_id in req.serialize_part.reservation_ids { + let Some(info) = view + .master_kv_router() + .take_prepared_put_key_reservation(reservation_id) + else { + tracing::info!( + "batch_release_put_key_reservations ignored missing reservation_id={} requester_node_id={}", + reservation_id, + req_node_id + ); + continue; + }; + if info.owner_node_id.as_ref() != req_node_id.as_ref() { + let owner_node_id = info.owner_node_id.to_string(); + view.master_kv_router() + .install_prepared_put_key_reservation(reservation_id, info); + for (restore_id, restore_info) in taken { + view.master_kv_router() + .install_prepared_put_key_reservation(restore_id, restore_info); + } + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidArgument { + detail: format!( + "batch_release_put_key_reservations owner mismatch: reservation_id={} owner_node_id={} requester_node_id={}", + reservation_id, owner_node_id, req_node_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + taken.push((reservation_id, info)); + } + + drop(taken); + + MsgPack { + serialize_part: BatchReleasePutKeyReservationsResp::default(), + raw_bytes: Vec::new(), + } +} + pub async fn handle_put_revoke( view: MasterKvRouterView, req: MsgPack, @@ -458,8 +1472,13 @@ pub async fn handle_put_revoke( .remove(&kvrouter_key) .await { - view.master_kv_router() - .release_inflight_put_key(&inflight_info.key); + let _activity_completion = + MasterKeyActivityCompletionGuard::new(inflight_info._activity_lease.clone()); + let _replica_activity_completion = inflight_info + .commit_info + .replica_target + .as_ref() + .map(|target| MasterKeyActivityCompletionGuard::new(target._activity_lease.clone())); tracing::info!("Revoked put operation with put_id: {:?}", kvrouter_key); } else { tracing::warn!( @@ -477,18 +1496,46 @@ pub async fn handle_put_revoke( pub async fn handle_put_done( view: MasterKvRouterView, req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + handle_put_done_with_resolved_group(view, req, req_node_id, None).await +} + +async fn handle_put_done_with_resolved_group( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, + resolved_atomic_group: Option>, ) -> MsgPack { tracing::debug!("Handling PutDoneReq: {:?}", req.serialize_part); let put_id = req.serialize_part.put_id; let lease_id_opt = req.serialize_part.lease_id; let full_put_id: (String, u64, u32) = (req.serialize_part.key.clone(), put_id.0, put_id.1); + let mut local_cache_holder_id: Option; + let atomic_group = if let Some(group) = resolved_atomic_group { + Some(group) + } else { + match view.master_kv_router().resolve_put_atomic_group( + &req.serialize_part.key, + put_id, + req.serialize_part.atomic_group.clone(), + ) { + Ok(group) => group, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + } + }; // Remove from inflight_puts and store in completed_puts if let Some(InflightPutInfo { - node_id, key, - src_target_allocation, + commit_info, + _activity_lease, .. }) = view .master_kv_router() @@ -497,8 +1544,14 @@ pub async fn handle_put_done( .remove(&full_put_id) .await { - view.master_kv_router().release_inflight_put_key(&key); - let Some(allocs) = src_target_allocation.lock().take() else { + let _activity_completion = MasterKeyActivityCompletionGuard::new(_activity_lease); + let mut replica_activity_completion = commit_info + .replica_target + .as_ref() + .map(|target| MasterKeyActivityCompletionGuard::new(target._activity_lease.clone())); + let node_id = commit_info.node_id; + let tomb_tag = commit_info.target_tomb_tag.clone(); + let Some(allocs) = commit_info.src_target_allocation.lock().take() else { tracing::warn!( "Put operation with put_id {:?} not found for completion", full_put_id @@ -515,19 +1568,14 @@ pub async fn handle_put_done( }; }; - let mut target_allocation = match allocs { - InflightPutAllocation::Local(target) => target, - InflightPutAllocation::Remote { src: _src, target } => target, - }; - - let Some(tomb_tag) = view.master_seg_manager().get_node_tomb_tag(&node_id) else { - tracing::warn!( - "Put operation with put_id {:?} not found for completion", + if !node_generation_is_current_live(&view, &node_id, &tomb_tag) { + tracing::info!( + "Put operation with put_id {:?} belongs to a departed target generation, skip", put_id ); let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { detail: format!( - "Put operation with put_id {:?} not found for completion", + "Put operation with put_id {:?} belongs to a departed target generation", put_id ), }); @@ -535,140 +1583,306 @@ pub async fn handle_put_done( serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), raw_bytes: Vec::new(), }; - }; + } - if tomb_tag.is_tomb() { - tracing::info!("Put operation with put_id {:?} is tomb, skip", put_id); + let route_committed_slot = req.serialize_part.committed_slot.clone(); + if req.serialize_part.publish_local_cache + && (!matches!(&allocs, InflightPutAllocation::Local(_)) + || route_committed_slot.is_some()) + { let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { - detail: format!("Put operation with put_id {:?} is tomb, skip", put_id), + detail: format!( + "publish_local_cache requires owner-local allocation backing; key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), }); return MsgPack { serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), raw_bytes: Vec::new(), }; } - - let target_cap_bytes = target_allocation.capcity(); - // NOTE on weight sizing for moka cache: - // - moka's `weigher` returns a u32 per-entry weight while the cache's - // `max_capacity` and `weighted_size()` use u64. If an allocation's - // capacity exceeds u32::MAX (e.g., >= 4 GiB), a naive `as u32` cast - // would truncate and could become 0 for ~exact 4 GiB multiples. - // That would effectively disable size-based eviction because such - // entries would contribute 0 to the cache weight and the cache would - // never reach its configured capacity. This directly causes the - // observed "non‑lease mode eviction not working; puts fill to full". - // - To make eviction robust, we saturate the per-entry weight at - // u32::MAX when `capcity()` is larger than u32::MAX. This keeps the - // cache accounting conservative (evicts earlier rather than later) - // and prevents weight=0 due to truncation. - let saturated_weight_u32 = if target_cap_bytes > u32::MAX as u64 { - tracing::warn!( - "moka weight saturation: key={} put_id=({},{}) cap={}B exceeds u32::MAX; weight set to u32::MAX", - key, - put_id.0, - put_id.1, - target_cap_bytes - ); - u32::MAX - } else { - target_cap_bytes as u32 - }; - // Note: moka cache insertion happens after commit in a spawned task - // using the same saturated weight; avoid unused local here. - // If lease is provided, attach first and fail fast on error - if let Some(lease_id) = lease_id_opt { - if let Err(e) = view - .master_lease_manager() - .attach_key(lease_id, key.clone(), put_id) + let (target_cap_bytes, completed_info, local_cache_publish_supported) = match allocs { + InflightPutAllocation::Local(target_allocation) => { + if let Some(slot) = route_committed_slot { + let (committed_slot, slot_tomb_tag) = match validate_put_done_committed_slot( + &view, + &node_id, + &slot, + Some(&tomb_tag), + ) { + Ok(validated) => validated, + Err(err) => { + return MsgPack { + serialize_part: + crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + }; + let target_cap_bytes = committed_slot.slot_size; + let capacity_reservation = match prepare_route_state( + &view, + lease_id_opt, + &key, + put_id, + &node_id, + &slot_tomb_tag, + None, + target_cap_bytes, + ) + .await + { + Ok(reservation) => reservation, + Err(err) => { + return MsgPack { + serialize_part: + crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + }; + drop(target_allocation); + ( + target_cap_bytes, + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(committed_slot), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation, + }, + false, + ) + } else { + let target_cap_bytes = target_allocation.capcity(); + let reservation_reason = if req.serialize_part.publish_local_cache { + Some(ReservedCapacityReason::OwnerIndexedAllocation) + } else if lease_id_opt.is_some() { + Some(ReservedCapacityReason::LeaseBoundKv) + } else { + None + }; + let capacity_reservation = match prepare_route_state( + &view, + lease_id_opt, + &key, + put_id, + &node_id, + &tomb_tag, + reservation_reason, + target_cap_bytes, + ) + .await + { + Ok(reservation) => reservation, + Err(err) => { + return MsgPack { + serialize_part: + crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + }; + ( + target_cap_bytes, + KvMemoryReplica { + backing: KvReplicaBacking::Allocation(Arc::new(target_allocation)), + owner_local_indexed: req.serialize_part.publish_local_cache, + get_durable_reservation: None, + capacity_reservation, + }, + true, + ) + } + } + InflightPutAllocation::Remote { src: _src, target } => { + let target_cap_bytes = target.capcity(); + let capacity_reservation = match prepare_route_state( + &view, + lease_id_opt, + &key, + put_id, + &node_id, + &tomb_tag, + lease_id_opt.map(|_| ReservedCapacityReason::LeaseBoundKv), + target_cap_bytes, + ) .await - { - let kv_err: crate::rpcresp_kvresult_convert::msg_and_error::KvError = e.into(); + { + Ok(reservation) => reservation, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }; + } + }; + ( + target_cap_bytes, + KvMemoryReplica { + backing: KvReplicaBacking::Allocation(Arc::new(target)), + owner_local_indexed: false, + get_durable_reservation: None, + capacity_reservation, + }, + false, + ) + } + InflightPutAllocation::LocalCommittedSlot(slot) => { + let target_cap_bytes = slot.slot_size; + let capacity_reservation = match prepare_route_state( + &view, + lease_id_opt, + &key, + put_id, + &node_id, + &tomb_tag, + None, + target_cap_bytes, + ) + .await + { + Ok(reservation) => reservation, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }; + } + }; + ( + target_cap_bytes, + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(slot), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation, + }, + false, + ) + } + }; + + local_cache_holder_id = if req.serialize_part.publish_local_cache { + if !local_cache_publish_supported { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "publish_local_cache requires owner-local allocation backing; key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }, + ); return MsgPack { - serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&kv_err), + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), raw_bytes: Vec::new(), }; - } - // Reserve cache capacity on this node for the leased allocation now (fetch_sub semantics) - if let Err(e) = view + }; + let KvReplicaBacking::Allocation(allocation) = &completed_info.backing else { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "publish_local_cache requires allocation backing; key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }, + ); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + }; + let holder_id = view .master_kv_router() - .adjust_node_cache_capacity_for_lease(node_id.as_ref(), target_cap_bytes as i64) - { - let kv_err: crate::rpcresp_kvresult_convert::msg_and_error::KvError = e.into(); + .inner() + .next_holder_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + view.master_kv_router().inner().get_holding.insert( + crate::memholder::NodeHolderKey::new(node_id.to_string(), holder_id), + OwnerHoldingGetInfo { + key: key.clone(), + holding_node_id: node_id.clone(), + len: allocation.size(), + allocation: allocation.clone(), + }, + ); + Some(holder_id) + } else { + None + }; + + // Publish the primary route under the node-generation fence. This + // closes the gap where MemberLeft marked/snapshotted routes between a + // pre-check and a later DashMap insert. + let publish_tag = tomb_tag.clone(); + let new_route = Arc::new(OneKvNodesRoutes { + put_id, + lease_id: lease_id_opt, + atomic_group: atomic_group.clone(), + node_replicas: RwLock::new(HashMap::from([( + node_id.clone(), + KvNodeReplicas::memory(publish_tag.clone(), completed_info), + )])), + get_durable_slots_used: AtomicU32::new(0), + }); + let old_one_kv_routes = match publish_primary_route_tomb_fenced( + &view.master_kv_router().inner().kv_routes, + &key, + new_route.clone(), + &publish_tag, + ) { + Ok(previous) => previous, + Err(()) => { + if let Some(lease_id) = lease_id_opt { + view.master_lease_manager() + .detach_key_if_version(lease_id, &key, put_id); + } + if let Some(holder_id) = local_cache_holder_id.take() { + view.master_kv_router().inner().get_holding.remove( + &crate::memholder::NodeHolderKey::new(node_id.to_string(), holder_id), + ); + } + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "primary route publication rejected because target generation departed: key={} put_id=({},{}) node_id={}", + key, put_id.0, put_id.1, node_id + ), + }, + ); return MsgPack { - serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&kv_err), + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), raw_bytes: Vec::new(), }; } - // And attach an on-drop hook to restore it (fetch_add on Allocation drop) - let view_clone = view.clone(); - let node_id_string = node_id.as_ref().to_string(); - target_allocation.set_on_drop(move || { - match view_clone.try_adjust_node_cache_capacity_for_lease( - &node_id_string, - -(target_cap_bytes as i64), - ) { - Some(Ok(())) => {} - Some(Err(e)) => { - tracing::warn!( - "Failed to restore moka capacity on drop: node_id={}, bytes={}, err={}", - node_id_string, - target_cap_bytes, - e - ); - } - None => { - tracing::debug!( - "Skipped restoring moka capacity on drop because MasterKvRouterView is gone: node_id={}, bytes={}", - node_id_string, - target_cap_bytes - ); - } - } - }); - } - - let completed_info = KvRouteInfo { - node_id: node_id.clone(), - allocation: Arc::new(target_allocation), - tomb_tag, }; - // Insert into kv_routes with replica support - let mut old_one_kv_routes: Option> = None; - let mut inserted = false; - { - let mut one_kv_routes = view - .master_kv_router() + if let Some(replica_target) = commit_info.replica_target { + view.master_kv_router() .inner() - .kv_routes - .entry(key.clone()) - .or_insert_with(|| { - inserted = true; - Arc::new(OneKvNodesRoutes { - put_id, - lease_id: lease_id_opt, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), - }) - }); - // we need to take out old one_kv_routes if it is not inserted - if !inserted { - old_one_kv_routes = Some(one_kv_routes.clone()); - *one_kv_routes = Arc::new(OneKvNodesRoutes { - put_id, - lease_id: lease_id_opt, - nodes_replicas: RwLock::new(HashMap::new()), - get_durable_slots_used: AtomicU32::new(0), - }); - } - one_kv_routes - .nodes_replicas - .write() - .insert(node_id.clone(), completed_info); + .inflight_replica_tasks + .insert( + ( + replica_target.key.clone(), + replica_target.put_id.0, + replica_target.put_id.1, + ), + replica_target, + ) + .await; + replica_activity_completion + .as_mut() + .expect("replica target activity guard must exist") + .disarm(); } if let Some(old) = old_one_kv_routes { + view.master_kv_router() + .remove_route_cache_entries_exact(&key, &old) + .await; if let Err(err) = view .master_kv_router() .inner() @@ -684,65 +1898,112 @@ pub async fn handle_put_done( } } - // Post-commit maintenance: update prefix-count index (for CountPrefix RPC) - // and, if applicable, update per-node cache controller. Run both in a - // spawned task to keep the PutDone RPC path lean and consistent with - // other async cache control operations. Deletion path already removes - // the index entry in delete.rs (do_delete_one_kv_all_replicas). - { - let view_task = view.clone(); - let key_for_spawn = key.clone(); - let node_for_spawn = node_id.clone(); - let do_prefix_index_update = view.master_kv_router().prefix_index_enabled(); - let do_cache_insert = - lease_id_opt.is_none() && view.master_kv_router().replica_cache_enabled(); - // Reuse the saturated weight computed above for moka insertion - let cap_bytes_u32 = saturated_weight_u32; - let _ = view.spawn("post_put_done_maintenance", async move { - // 1) Update the derived prefix-counting index asynchronously. - // This keeps PutDone lean, but also means CountPrefix visibility - // is not an immediate strong-consistency signal for this put. - if do_prefix_index_update { - let inner = view_task.master_kv_router().inner(); - let mut tree = inner.prefix_index.write().await; - tree.insert(&key_for_spawn); - } - - // 2) Optionally update node cache controller (non-leased keys) - if do_cache_insert { - let cache = view_task - .master_kv_router() - .get_node_cache_controller(&node_for_spawn); - if let Some(cache) = cache { - let desc = NodeValueReplicaDesc { - weight_bytes: cap_bytes_u32, - put_id, - }; - tracing::debug!("Inserting key: {:?} into cache", key_for_spawn); - cache.insert(key_for_spawn.clone(), desc); - tracing::debug!( - "Inserted key: {:?} into cache, current cache size: {}", - key_for_spawn, - cache.weighted_size() - ); - } else { - tracing::warn!( - "No cache controller found for node: {}, node is not ready", - node_for_spawn - ); - } - } - }); - } + enqueue_post_route_maintenance( + &view, + RoutePublishEvent::primary_put( + key.clone(), + put_id, + lease_id_opt, + node_id.clone(), + target_cap_bytes, + ), + ) + .await; // Lease attach is handled before kv_routes insertion - tracing::info!( + tracing::debug!( "Completed put operation with put_id: {:?}, key: {:?}", put_id, key ); } else { + if let Some(slot) = req.serialize_part.committed_slot.clone() { + let key = req.serialize_part.key.clone(); + let _activity_lease = match view + .master_kv_router() + .reserve_inflight_put_key(&key, false, false) + { + Ok(activity_lease) => activity_lease, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }; + } + }; + let node_id = req_node_id; + let (committed_slot, tomb_tag) = + match validate_put_done_committed_slot(&view, &node_id, &slot, None) { + Ok(validated) => validated, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }; + } + }; + if req.serialize_part.publish_local_cache { + let err = msg_and_error::KvError::Api( + msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "local-first put_done does not support publish_local_cache: key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }, + ); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + let target_cap_bytes = committed_slot.slot_size; + let capacity_reservation = match prepare_route_state( + &view, + lease_id_opt, + &key, + put_id, + &node_id, + &tomb_tag, + None, + target_cap_bytes, + ) + .await + { + Ok(reservation) => reservation, + Err(err) => { + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error( + &err, + ), + raw_bytes: Vec::new(), + }; + } + }; + let completed_info = KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(committed_slot), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation, + }; + return publish_completed_put_route( + view, + key, + put_id, + lease_id_opt, + atomic_group, + node_id, + tomb_tag, + completed_info, + target_cap_bytes, + None, + ) + .await; + } tracing::warn!( "Put operation with put_id {:?} not found for completion", put_id @@ -761,6 +2022,698 @@ pub async fn handle_put_done( error_code: msg_and_error::OK, error_json: String::new(), server_process_us: 0, + local_cache_holder_id, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_put_start( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let (_put_id, resp) = handle_put_start( + view.clone(), + MsgPack { + serialize_part: PutStartReq { + key: item.key, + len: item.len, + reject_if_inflight_same_key: item.reject_if_inflight_same_key, + reject_if_exist_same_key: item.reject_if_exist_same_key, + make_replica_task: item.make_replica_task, + preferred_sub_cluster: item.preferred_sub_cluster, + source_node_id: None, + }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutStartItemResp { + put_id: part.put_id, + node_id: part.node_id, + target_addr: part.target_addr, + src_addr: part.src_addr, + target_base_addr: part.target_base_addr, + src_base_addr: part.src_base_addr, + len: part.len, + error_code: part.error_code, + error_json: part.error_json, + replica_target: part.replica_target, + }); + } + MsgPack { + serialize_part: BatchPutStartResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_put_revoke( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let key = item.key.clone(); + let put_id = item.put_id; + let resp = handle_put_revoke( + view.clone(), + MsgPack { + serialize_part: PutRevokeReq { key, put_id }, + raw_bytes: Vec::new(), + }, + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutRevokeItemResp { + key: item.key, + put_id: item.put_id, + error_code: part.error_code, + error_json: part.error_json, + }); + } + MsgPack { + serialize_part: BatchPutRevokeResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_batch_put_done( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let key = item.key.clone(); + let put_id = item.put_id; + let lease_id = item.lease_id; + let resp = handle_put_done( + view.clone(), + MsgPack { + serialize_part: PutDoneReq { + key, + put_id, + lease_id, + committed_slot: item.committed_slot, + publish_local_cache: item.publish_local_cache, + atomic_group: item.atomic_group, + }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutDoneItemResp { + key: item.key, + put_id: item.put_id, + error_code: part.error_code, + error_json: part.error_json, + local_cache_holder_id: part.local_cache_holder_id, + }); + } + MsgPack { + serialize_part: BatchPutDoneResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +/// V2 route publication for local-first puts. The wire carries each key once +/// plus a compact ordered partition. The master materializes one shared group +/// descriptor per partition and passes cheap `Arc` clones to member routes, +/// avoiding both repeated wire descriptors and repeated group validation. +pub async fn handle_grouped_batch_put_done( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let GroupedBatchPutDoneReq { + items: request_items, + atomic_group_lens, + } = req.serialize_part; + let keys_and_put_ids = request_items + .iter() + .map(|item| (item.key.clone(), item.put_id)) + .collect::>(); + let assignments = + match build_shared_put_atomic_group_assignments(&keys_and_put_ids, &atomic_group_lens) { + Ok(assignments) => assignments, + Err(detail) => { + let err = msg_and_error::ApiError::InvalidArgument { detail }; + let (error_code, error_json) = err.to_code_and_json(); + return MsgPack { + serialize_part: GroupedBatchPutDoneResp { + items: Vec::new(), + error_code, + error_json, + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }; + } + }; + + // The partition builder derives membership from these exact ordered items. + // Reject duplicate/empty keys once per group so every member is represented + // exactly once before any route becomes visible. + let mut offset = 0usize; + for group_len in atomic_group_lens.iter().copied() { + if group_len > 1 { + let mut unique = HashSet::with_capacity(group_len); + let end = offset + group_len; + if keys_and_put_ids[offset..end] + .iter() + .any(|(key, _)| key.is_empty() || !unique.insert(key.as_str())) + { + let err = msg_and_error::ApiError::InvalidArgument { + detail: format!( + "grouped put member keys must be non-empty and unique: offset={} len={}", + offset, group_len + ), + }; + let (error_code, error_json) = err.to_code_and_json(); + return MsgPack { + serialize_part: GroupedBatchPutDoneResp { + items: Vec::new(), + error_code, + error_json, + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }; + } + } + offset += group_len; + } + + let mut items = Vec::with_capacity(request_items.len()); + for (item, atomic_group) in request_items.into_iter().zip(assignments) { + let key = item.key.clone(); + let put_id = item.put_id; + let resp = handle_put_done_with_resolved_group( + view.clone(), + MsgPack { + serialize_part: PutDoneReq { + key, + put_id, + lease_id: item.lease_id, + committed_slot: item.committed_slot, + publish_local_cache: item.publish_local_cache, + atomic_group: None, + }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + atomic_group, + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutDoneItemResp { + key: item.key, + put_id: item.put_id, + error_code: part.error_code, + error_json: part.error_json, + local_cache_holder_id: part.local_cache_holder_id, + }); + } + MsgPack { + serialize_part: GroupedBatchPutDoneResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +async fn handle_put_append_start_inner( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let key = req.serialize_part.key.clone(); + let put_id = req.serialize_part.put_id; + let append_key = (key.clone(), put_id.0, put_id.1); + let operation_lock = view + .master_kv_router() + .inner() + .replica_operation_locks + .get_lock(append_key.clone()); + let _operation_guard = operation_lock.lock().await; + let route_snapshot = view + .master_kv_router() + .inner() + .kv_routes + .get(&key) + .map(|route| route.clone()); + let current_outcome = route_snapshot + .as_ref() + .map(|route| current_route_append_outcome(route, &req_node_id, put_id)) + .unwrap_or(PutAppendStartOutcome::Obsolete); + if current_outcome != PutAppendStartOutcome::Scheduled { + if let Some(inflight) = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .remove(&(key.clone(), put_id.0, put_id.1)) + .await + { + inflight._activity_lease.release_now(); + } + return MsgPack { + serialize_part: PutAppendStartResp { + outcome: current_outcome, + error_code: msg_and_error::OK, + error_json: String::new(), + ..Default::default() + }, + raw_bytes: Vec::new(), + }; + } + + let inflight = if let Some(existing) = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .get(&append_key) + .await + { + existing + } else { + let excluded_nodes = route_snapshot + .as_ref() + .map(|route| { + route + .node_replicas + .read() + .iter() + .filter_map(|(node_id, replicas)| { + replicas.has_live_backing().then_some(node_id.clone()) + }) + .collect::>() + }) + .unwrap_or_default(); + let reservation = match reserve_replica_task_excluding( + &view, + &key, + put_id, + &req_node_id, + req.serialize_part.preferred_sub_cluster.as_deref(), + req.serialize_part.len, + &excluded_nodes, + req.serialize_part.protect_source_on_remote_complete, + ) { + Ok(reservation) => reservation, + Err(msg_and_error::KvError::Api(msg_and_error::ApiError::NoSpace { + node, + segment, + total_capacity, + free_capacity, + })) => { + tracing::info!( + "replica task not scheduled; local-only commit remains valid: key={} put_id=({},{}) source_node_id={} preferred_sub_cluster={:?} node={} segment={} total_capacity={} free_capacity={}", + key, + put_id.0, + put_id.1, + req_node_id, + req.serialize_part.preferred_sub_cluster, + node, + segment, + total_capacity, + free_capacity + ); + return MsgPack { + serialize_part: PutAppendStartResp { + outcome: PutAppendStartOutcome::RetryableNoSpace, + error_code: msg_and_error::OK, + error_json: String::new(), + ..Default::default() + }, + raw_bytes: Vec::new(), + }; + } + Err(err) => { + let resp: PutAppendStartResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }; + } + }; + view.master_kv_router() + .record_replica_task_target(reservation.node_id.as_ref()); + view.master_kv_router() + .inner() + .inflight_replica_tasks + .insert(append_key.clone(), reservation.clone()) + .await; + reservation + }; + let (target_base_addr, target_addr, allocation_size) = { + let target_allocation_guard = inflight.target_allocation.lock(); + let Some(target_allocation): Option<&Allocation> = target_allocation_guard.as_ref() else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Replica task reservation allocation missing: key={} put_id=({},{}) requester_node_id={}", + key, put_id.0, put_id.1, req_node_id + ), + }); + let resp: PutAppendStartResp = + crate::rpcresp_kvresult_convert::FromError::from_error(&err); + return MsgPack { + serialize_part: resp, + raw_bytes: Vec::new(), + }; + }; + let target_base_addr = target_allocation.base_addr(); + let target_addr = target_base_addr + target_allocation.addr(); + let allocation_size = target_allocation.size(); + (target_base_addr, target_addr, allocation_size) + }; + + MsgPack { + serialize_part: PutAppendStartResp { + outcome: PutAppendStartOutcome::Scheduled, + operation_id: inflight.operation_id, + node_id: inflight.node_id.clone().into(), + target_addr, + target_base_addr, + len: allocation_size, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_put_append_start( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + handle_put_append_start_inner(view, req, req_node_id).await +} + +pub async fn handle_batch_put_append_start( + view: MasterKvRouterView, + req: MsgPack, + req_node_id: NodeID, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let key = item.key.clone(); + let put_id = item.put_id; + let resp = handle_put_append_start_inner( + view.clone(), + MsgPack { + serialize_part: PutAppendStartReq { + key, + put_id, + len: item.len, + preferred_sub_cluster: item.preferred_sub_cluster, + protect_source_on_remote_complete: item.protect_source_on_remote_complete, + }, + raw_bytes: Vec::new(), + }, + req_node_id.clone(), + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutAppendStartItemResp { + key: item.key, + put_id: item.put_id, + outcome: part.outcome, + operation_id: part.operation_id, + node_id: part.node_id, + target_addr: part.target_addr, + target_base_addr: part.target_base_addr, + len: part.len, + error_code: part.error_code, + error_json: part.error_json, + }); + } + MsgPack { + serialize_part: BatchPutAppendStartResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_put_append_revoke( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + let put_id = req.serialize_part.put_id; + let key = req.serialize_part.key; + let operation_id = req.serialize_part.operation_id; + let generation_identity = (key.clone(), put_id.0, put_id.1); + let operation_identity = (key.clone(), put_id.0, put_id.1, operation_id); + let operation_lock = view + .master_kv_router() + .inner() + .replica_operation_locks + .get_lock(generation_identity.clone()); + let _operation_guard = operation_lock.lock().await; + if view + .master_kv_router() + .inner() + .completed_replica_tasks + .get(&operation_identity) + .await + .is_some() + { + return MsgPack { + serialize_part: PutAppendRevokeResp { + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + }; + } + let inflight = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .get(&generation_identity) + .await; + if inflight + .as_ref() + .is_some_and(|inflight| inflight.operation_id == operation_id) + { + if let Some(inflight) = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .remove(&generation_identity) + .await + { + inflight._activity_lease.release_now(); + } + } + MsgPack { + serialize_part: PutAppendRevokeResp { + error_code: msg_and_error::OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +async fn handle_put_append_done_inner( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + let put_id = req.serialize_part.put_id; + let key = req.serialize_part.key.clone(); + let operation_id = req.serialize_part.operation_id; + let generation_identity = (key.clone(), put_id.0, put_id.1); + let operation_identity = (key.clone(), put_id.0, put_id.1, operation_id); + let operation_lock = view + .master_kv_router() + .inner() + .replica_operation_locks + .get_lock(generation_identity.clone()); + let _operation_guard = operation_lock.lock().await; + if let Some(completed) = view + .master_kv_router() + .inner() + .completed_replica_tasks + .get(&operation_identity) + .await + { + view.master_kv_router() + .inner() + .replica_done_terminal_replay_count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return MsgPack { + serialize_part: PutAppendDoneResp { + appended: completed.appended, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + }; + } + let Some(current) = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .get(&generation_identity) + .await + else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Put append operation not found for completion: key={} put_id=({},{}) operation_id={}", + key, put_id.0, put_id.1, operation_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + }; + if current.operation_id != operation_id { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Put append operation generation mismatch: key={} put_id=({},{}) requested_operation_id={} current_operation_id={}", + key, put_id.0, put_id.1, operation_id, current.operation_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + } + let Some(inflight) = view + .master_kv_router() + .inner() + .inflight_replica_tasks + .remove(&generation_identity) + .await + else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Put append operation disappeared during completion: key={} put_id=({},{}) operation_id={}", + key, put_id.0, put_id.1, operation_id + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + }; + let _activity_completion = + MasterKeyActivityCompletionGuard::new(inflight._activity_lease.clone()); + let Some(allocation) = inflight.target_allocation.lock().take() else { + let err = msg_and_error::KvError::Api(msg_and_error::ApiError::InvalidPutMasterState { + detail: format!( + "Replica task append target allocation already taken: key={} put_id=({},{})", + key, put_id.0, put_id.1 + ), + }); + return MsgPack { + serialize_part: crate::rpcresp_kvresult_convert::FromError::from_error(&err), + raw_bytes: Vec::new(), + }; + }; + let published = append_current_route_replica_if_matching( + &view, + &key, + inflight.put_id, + inflight.node_id, + inflight.target_tomb_tag, + allocation, + ); + let appended = published.is_some(); + view.master_kv_router() + .inner() + .completed_replica_tasks + .insert(operation_identity, CompletedReplicaTaskInfo { appended }) + .await; + if let Some(event) = published { + enqueue_post_route_maintenance(&view, event).await; + } + MsgPack { + serialize_part: PutAppendDoneResp { + appended, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, + }, + raw_bytes: Vec::new(), + } +} + +pub async fn handle_put_append_done( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + handle_put_append_done_inner(view, req).await +} + +pub async fn handle_batch_put_append_done( + view: MasterKvRouterView, + req: MsgPack, +) -> MsgPack { + let mut items = Vec::with_capacity(req.serialize_part.items.len()); + for item in req.serialize_part.items { + let key = item.key.clone(); + let put_id = item.put_id; + let operation_id = item.operation_id; + let resp = handle_put_append_done_inner( + view.clone(), + MsgPack { + serialize_part: PutAppendDoneReq { + key, + put_id, + operation_id, + }, + raw_bytes: Vec::new(), + }, + ) + .await; + let part = resp.serialize_part; + items.push(BatchPutAppendDoneItemResp { + key: item.key, + put_id: item.put_id, + appended: part.appended, + error_code: part.error_code, + error_json: part.error_json, + }); + } + MsgPack { + serialize_part: BatchPutAppendDoneResp { + items, + error_code: msg_and_error::OK, + error_json: String::new(), + server_process_us: 0, }, raw_bytes: Vec::new(), } diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/reclaim.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/reclaim.rs new file mode 100644 index 0000000..ac5736b --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/reclaim.rs @@ -0,0 +1,3174 @@ +use super::{KvReplicaBacking, MasterKvRouterView, NodeValueReplicaDesc}; +use crate::cluster_manager::{NodeID, NodeIDString}; +use crate::master_kv_router::msg_pack::{ + BatchEvictOwnerSourceReq, BatchEvictOwnerSourceResp, BatchOwnerReclaimReq, OwnerReclaimBacking, + OwnerReclaimItem, OwnerReclaimItemResp, OwnerReclaimItemState, OwnerReclaimPhase, + OwnerReclaimReason, OwnerSourceEvictionOutcome, OwnerSourceEvictionVictim, + OwnerSourceEvictionVictimResp, OwnerSourceSsdPolicy, owner_source_eviction_epoch, +}; +use crate::p2p::msg_pack::{MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, RPCCaller}; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, OK}; +use limit_thirdparty::tokio; +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +const OWNER_RECLAIM_RPC_TIMEOUT: Duration = Duration::from_secs(MIN_EXPLICIT_RPC_TIMEOUT_SECS); +const OWNER_RECLAIM_MAX_BATCH: usize = 256; +// Keep each transport/SSD transaction small enough to release physical slots +// continuously. Every entry remains an independently fenced single-KV +// victim; this is only an RPC and bounded-I/O aggregation limit. +const OWNER_RECLAIM_RPC_BATCH: usize = 32; +const OWNER_RECLAIM_MERGE_WINDOW: Duration = Duration::from_millis(5); +const EVICTION_RECLAIM_RETRY_INITIAL: Duration = Duration::from_millis(100); +const EVICTION_RECLAIM_RETRY_MAX: Duration = Duration::from_secs(1); +const EVICTION_RECLAIM_MAX_RETRY_COUNT: u32 = 5; + +fn should_restore_after_retry(retry_count: u32) -> bool { + retry_count >= EVICTION_RECLAIM_MAX_RETRY_COUNT +} + +fn eviction_reclaim_retry_delay(retry_count: u32) -> Duration { + let multiplier = 1u32 << retry_count.saturating_sub(1).min(16); + EVICTION_RECLAIM_RETRY_INITIAL + .saturating_mul(multiplier) + .min(EVICTION_RECLAIM_RETRY_MAX) +} + +#[cfg(test)] +mod timeout_contract_tests { + use super::*; + use crate::p2p::msg_pack::validate_explicit_rpc_timeout; + + #[test] + fn owner_reclaim_timeout_satisfies_rpc_contract() { + validate_explicit_rpc_timeout(Some(OWNER_RECLAIM_RPC_TIMEOUT)).unwrap(); + } + + #[test] + fn eviction_reclaim_retry_restoration_is_bounded() { + assert!(!should_restore_after_retry( + EVICTION_RECLAIM_MAX_RETRY_COUNT - 1 + )); + assert!(should_restore_after_retry(EVICTION_RECLAIM_MAX_RETRY_COUNT)); + assert!(should_restore_after_retry(u32::MAX)); + } + + #[test] + fn eviction_reclaim_retry_paces_restore_view_holders() { + assert_eq!(eviction_reclaim_retry_delay(1), Duration::from_millis(100)); + assert_eq!(eviction_reclaim_retry_delay(2), Duration::from_millis(200)); + assert_eq!(eviction_reclaim_retry_delay(3), Duration::from_millis(400)); + assert_eq!(eviction_reclaim_retry_delay(4), Duration::from_millis(800)); + assert_eq!( + eviction_reclaim_retry_delay(u32::MAX), + Duration::from_secs(1) + ); + } +} + +#[cfg(test)] +mod single_victim_transaction_tests { + use super::*; + use crate::master_kv_router::msg_pack::{ + OwnerSourceEvictionVictim, PutAtomicGroup, PutAtomicGroupMember, + }; + use crate::master_kv_router::{ + CommittedSlotReplica, KvMemoryReplica, KvNodeReplicas, KvSsdReplica, MasterKeyActivityKind, + MasterKeyActivityTable, OneKvNodesRoutes, + }; + use crate::master_seg_manager::NodeTombTag; + use parking_lot::RwLock; + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicU32}; + + fn item(key: &str, epoch: u64) -> OwnerReclaimItem { + OwnerReclaimItem { + key: key.to_string(), + put_id: (7, epoch as u32), + epoch, + backing: OwnerReclaimBacking::CommittedSlot { + grant_id: 1, + slot_index: epoch as u32, + slot_size: 4096, + }, + reason: OwnerReclaimReason::OwnerCapacityEviction, + } + } + + #[test] + fn single_victim_master_fence_rejects_a_busy_key() { + let activity = Arc::new(MasterKeyActivityTable::default()); + let items = vec![item("victim", 1)]; + let _busy = activity + .reserve("victim", MasterKeyActivityKind::Get, false) + .unwrap(); + + assert!(try_install_master_fences(&activity, &items).is_err()); + assert!(!activity.has_reclaim("victim")); + } + + #[test] + fn single_victim_master_fence_installs_and_clears() { + let activity = Arc::new(MasterKeyActivityTable::default()); + let items = vec![item("victim", 1)]; + + try_install_master_fences(&activity, &items).unwrap(); + assert!(activity.has_reclaim("victim")); + for item in &items { + assert!(activity.clear_reclaim(item)); + } + } + + #[test] + fn master_capacity_origin_rejects_committed_slot_backing() { + let backing = KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: "gpu0".to_string().into(), + grant_id: 1, + slot_index: 2, + slot_size: 4096, + addr: 0, + len: 4096, + base_addr: 0, + }); + assert_eq!( + master_allocation_capacity_weight(&backing), + Err(MasterCapacityPlanError::CommittedSlot), + ); + } + + fn source_victim( + key: &str, + put_id: (u64, u32), + grant_id: u64, + slot_index: u32, + ) -> OwnerSourceEvictionVictim { + OwnerSourceEvictionVictim { + key: key.to_string(), + put_id, + backing: OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size: 4096, + }, + ssd_backing_len: None, + ssd_policy: OwnerSourceSsdPolicy::Drop, + } + } + + fn source_route( + owner: &NodeID, + member: &OwnerSourceEvictionVictim, + atomic_group: Option>, + include_cpu_replica: bool, + ) -> Arc { + let OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size, + } = &member.backing + else { + unreachable!() + }; + let owner_replica = KvNodeReplicas::memory( + NodeTombTag::new(), + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: owner.clone(), + grant_id: *grant_id, + slot_index: *slot_index, + slot_size: *slot_size, + addr: 0, + len: *slot_size, + base_addr: 0, + }), + owner_local_indexed: true, + get_durable_reservation: None, + capacity_reservation: None, + }, + ); + let mut replicas = HashMap::from([(owner.clone(), owner_replica)]); + if include_cpu_replica { + let cpu: NodeID = "cpu0".to_string().into(); + replicas.insert( + cpu.clone(), + KvNodeReplicas::memory( + NodeTombTag::new(), + KvMemoryReplica { + backing: KvReplicaBacking::CommittedSlot(CommittedSlotReplica { + owner_node_id: cpu, + grant_id: 900 + *grant_id, + slot_index: *slot_index, + slot_size: *slot_size, + addr: 0, + len: *slot_size, + base_addr: 0, + }), + owner_local_indexed: false, + get_durable_reservation: None, + capacity_reservation: None, + }, + ), + ); + } + Arc::new(OneKvNodesRoutes { + put_id: member.put_id, + lease_id: None, + atomic_group, + node_replicas: RwLock::new(replicas), + get_durable_slots_used: AtomicU32::new(0), + }) + } + + fn source_reclaim_item(member: &OwnerSourceEvictionVictim) -> OwnerReclaimItem { + OwnerReclaimItem { + key: member.key.clone(), + put_id: member.put_id, + epoch: 1, + backing: member.backing.clone(), + reason: OwnerReclaimReason::OwnerCapacityEviction, + } + } + + #[test] + fn exact_source_plan_accepts_with_or_without_a_cpu_replica() { + let owner: NodeID = "gpu0".to_string().into(); + let member = source_victim("single", (10, 1), 7, 3); + + for include_cpu_replica in [false, true] { + let route = source_route(&owner, &member, None, include_cpu_replica); + let plan = plan_exact_owner_source_victim_with(&owner, &member, &|key| { + (key == "single").then(|| route.clone()) + }); + match plan { + OwnerSourceVictimPlan::Ready(planned) => { + assert_eq!(planned.key, "single"); + assert_eq!(planned.expected_backing, Some(member.backing.clone())); + } + _ => panic!("exact current owner source must be accepted"), + } + } + } + + #[test] + fn exact_source_removal_deletes_only_gpu_when_cpu_exists_and_last_route_otherwise() { + let owner: NodeID = "gpu0".to_string().into(); + let member = source_victim("with-cpu", (10, 2), 7, 4); + let routes = dashmap::DashMap::new(); + routes.insert( + member.key.clone(), + source_route(&owner, &member, None, true), + ); + let removed = + remove_exact_owner_source_route(&routes, &owner, &source_reclaim_item(&member)) + .expect("exact GPU source must be removed"); + let counters = crate::master_kv_router::EvictionReclaimCounters::default(); + record_last_route_removal(&counters, &removed); + assert!(!removed.removed_last_route); + assert!(!removed.ssd_survived); + assert!(!removed.ssd_became_only_backing); + assert_eq!( + counters.last_route_removed_members.load(Ordering::Relaxed), + 0 + ); + assert_eq!(counters.last_route_removed_bytes.load(Ordering::Relaxed), 0); + let remaining = routes.get(&member.key).expect("CPU route must remain"); + assert!(!remaining.node_replicas.read().contains_key(&owner)); + assert!(remaining.node_replicas.read().contains_key("cpu0")); + + let last = source_victim("last", (10, 3), 7, 5); + routes.insert(last.key.clone(), source_route(&owner, &last, None, false)); + let removed = remove_exact_owner_source_route(&routes, &owner, &source_reclaim_item(&last)) + .expect("last exact GPU source must be removed"); + record_last_route_removal(&counters, &removed); + assert!(removed.removed_last_route); + assert!(!routes.contains_key(&last.key)); + assert_eq!( + counters.last_route_removed_members.load(Ordering::Relaxed), + 1 + ); + assert_eq!( + counters.last_route_removed_bytes.load(Ordering::Relaxed), + removed.capacity_bytes + ); + + let stale = source_victim("stale", (10, 4), 8, 6); + routes.insert(stale.key.clone(), source_route(&owner, &stale, None, false)); + let wrong_identity = source_victim("stale", stale.put_id, 999, 6); + assert!( + remove_exact_owner_source_route( + &routes, + &owner, + &source_reclaim_item(&wrong_identity), + ) + .is_none() + ); + assert!(routes.contains_key(&stale.key)); + } + + #[test] + fn exact_source_memory_removal_preserves_same_owner_ssd_backing() { + let owner: NodeID = "gpu0".to_string().into(); + let member = source_victim("ssd-backed", (10, 5), 7, 7); + let route = source_route(&owner, &member, None, false); + route + .node_replicas + .write() + .get_mut(&owner) + .expect("owner route must exist") + .ssd = Some(KvSsdReplica { len: 4096 }); + let routes = dashmap::DashMap::new(); + routes.insert(member.key.clone(), route.clone()); + + let removed = + remove_exact_owner_source_route(&routes, &owner, &source_reclaim_item(&member)) + .expect("exact memory backing must be removed"); + + assert!(!removed.removed_last_route); + assert!(removed.ssd_survived); + assert!(removed.ssd_became_only_backing); + let current = routes + .get(&member.key) + .expect("SSD backing must keep the key route alive"); + let replicas = current.node_replicas.read(); + let owner_backings = replicas.get(&owner).expect("owner route must remain"); + assert!(owner_backings.memory.is_none()); + assert_eq!(owner_backings.ssd.as_ref().map(|ssd| ssd.len), Some(4096)); + } + + #[test] + fn ssd_preservation_selects_only_an_exact_last_live_backing() { + let owner: NodeID = "gpu0".to_string().into(); + let last = source_victim("last-for-ssd", (20, 1), 11, 1); + let last_route = source_route(&owner, &last, None, false); + let last_item = source_reclaim_item(&last); + assert!(exact_memory_reclaim_needs_ssd_with( + &owner, + &last_item, + &|key| (key == last.key).then(|| last_route.clone()), + )); + + let with_other = source_victim("other-live", (20, 2), 12, 2); + let with_other_route = source_route(&owner, &with_other, None, true); + let with_other_item = source_reclaim_item(&with_other); + assert!(!exact_memory_reclaim_needs_ssd_with( + &owner, + &with_other_item, + &|key| (key == with_other.key).then(|| with_other_route.clone()), + )); + + let already_on_ssd = source_victim("already-on-ssd", (20, 3), 13, 3); + let already_on_ssd_route = source_route(&owner, &already_on_ssd, None, false); + already_on_ssd_route + .node_replicas + .write() + .get_mut(&owner) + .expect("owner route must exist") + .ssd = Some(KvSsdReplica { len: 4096 }); + let already_on_ssd_item = source_reclaim_item(&already_on_ssd); + assert!(!exact_memory_reclaim_needs_ssd_with( + &owner, + &already_on_ssd_item, + &|key| (key == already_on_ssd.key).then(|| already_on_ssd_route.clone()), + )); + + let stale = source_victim("stale-for-ssd", (20, 4), 14, 4); + let stale_route = source_route(&owner, &stale, None, false); + let wrong_identity = source_victim("stale-for-ssd", stale.put_id, 999, 4); + assert!(!exact_memory_reclaim_needs_ssd_with( + &owner, + &source_reclaim_item(&wrong_identity), + &|key| (key == stale.key).then(|| stale_route.clone()), + )); + } + + #[test] + fn direct_delete_filters_before_ssd_and_drop_finishes_without_waiting() { + let owner: NodeID = "gpu0".to_string().into(); + let mut last = source_victim("last-candidate", (21, 1), 15, 1); + last.ssd_policy = OwnerSourceSsdPolicy::SelectLastLive; + let mut backed = source_victim("already-backed", (21, 2), 16, 2); + backed.ssd_policy = OwnerSourceSsdPolicy::SelectLastLive; + let routes = dashmap::DashMap::new(); + routes.insert(last.key.clone(), source_route(&owner, &last, None, false)); + routes.insert( + backed.key.clone(), + source_route(&owner, &backed, None, true), + ); + let activity = MasterKeyActivityTable::default(); + let (responses, busy) = direct_delete_exact_owner_source_batch_with( + &activity, + &owner, + 81, + &[last.clone(), backed.clone()], + &|key| routes.get(key).map(|route| route.clone()), + |item| remove_exact_owner_source_route(&routes, &owner, item).is_some(), + ); + assert_eq!(busy, DirectDeleteBatchBusySummary::default()); + assert_eq!( + responses + .iter() + .map(|response| response.outcome) + .collect::>(), + vec![ + OwnerSourceEvictionOutcome::SsdCandidate, + OwnerSourceEvictionOutcome::Completed, + ] + ); + assert!(routes.contains_key(&last.key)); + assert!( + routes + .get(&backed.key) + .unwrap() + .node_replicas + .read() + .contains_key("cpu0") + ); + assert!(!activity.has_reclaim(&last.key)); + + last.ssd_policy = OwnerSourceSsdPolicy::Drop; + let dropped = direct_delete_exact_owner_source_with( + &activity, + &owner, + &last, + owner_source_eviction_epoch(82, 0), + &|key| routes.get(key).map(|route| route.clone()), + |item| remove_exact_owner_source_route(&routes, &owner, item).is_some(), + ); + assert_eq!(dropped.outcome, OwnerSourceEvictionOutcome::Completed); + assert!(!dropped.ssd_backing_committed); + assert!(!routes.contains_key(&last.key)); + assert!(!activity.has_reclaim(&last.key)); + } + + #[test] + fn singleton_source_is_independent_of_atomic_group_siblings() { + let owner: NodeID = "gpu0".to_string().into(); + let a = source_victim("a", (12, 0), 9, 0); + let b = source_victim("b", (12, 1), 9, 1); + let group = Arc::new(PutAtomicGroup { + members: vec![ + PutAtomicGroupMember { + key: a.key.clone(), + put_id: a.put_id, + }, + PutAtomicGroupMember { + key: b.key.clone(), + put_id: b.put_id, + }, + ], + }); + let route_a = source_route(&owner, &a, Some(group.clone()), false); + match plan_exact_owner_source_victim_with(&owner, &a, &|key| { + (key == "a").then(|| route_a.clone()) + }) { + OwnerSourceVictimPlan::Ready(planned) => assert_eq!(planned.key, "a"), + _ => panic!("one current key must be reclaimable without its siblings"), + } + + let changed_a = source_victim("a", a.put_id, 99, 0); + let changed_route = source_route(&owner, &changed_a, Some(group), false); + assert!(matches!( + plan_exact_owner_source_victim_with(&owner, &a, &|key| { + (key == "a").then(|| changed_route.clone()) + }), + OwnerSourceVictimPlan::Stale(_) + )); + } + + #[test] + fn absent_single_source_is_already_completed() { + let owner: NodeID = "gpu0".to_string().into(); + let victim = source_victim("gone", (13, 0), 10, 0); + assert!(matches!( + plan_exact_owner_source_victim_with(&owner, &victim, &|_| None), + OwnerSourceVictimPlan::Completed(_) + )); + } + + #[test] + fn direct_delete_batch_keeps_results_independent_and_replay_idempotent() { + let owner: NodeID = "gpu0".to_string().into(); + let mut ready = source_victim("ready", (14, 0), 11, 0); + ready.ssd_backing_len = Some(4096); + ready.ssd_policy = OwnerSourceSsdPolicy::Persisted; + let mut busy = source_victim("busy", (14, 1), 11, 1); + busy.ssd_backing_len = Some(4096); + busy.ssd_policy = OwnerSourceSsdPolicy::Persisted; + let stale = source_victim("stale", (14, 2), 11, 2); + let routes = dashmap::DashMap::new(); + for victim in [&ready, &busy, &stale] { + routes.insert( + victim.key.clone(), + source_route(&owner, victim, None, false), + ); + } + let activity = Arc::new(MasterKeyActivityTable::default()); + let _busy_get = activity + .reserve(&busy.key, MasterKeyActivityKind::Get, false) + .expect("busy victim must hold a master Get lease"); + let mut stale_request = source_victim("stale", stale.put_id, 999, 2); + stale_request.ssd_backing_len = Some(4096); + stale_request.ssd_policy = OwnerSourceSsdPolicy::Persisted; + let victims = vec![ready.clone(), busy.clone(), stale_request]; + let (responses, busy_summary) = direct_delete_exact_owner_source_batch_with( + activity.as_ref(), + &owner, + 77, + &victims, + &|key| routes.get(key).map(|route| route.clone()), + |item| remove_exact_owner_source_route(&routes, &owner, item).is_some(), + ); + + assert_eq!( + responses + .iter() + .map(|response| response.outcome) + .collect::>(), + vec![ + OwnerSourceEvictionOutcome::Completed, + OwnerSourceEvictionOutcome::RetryableBusy, + OwnerSourceEvictionOutcome::Stale, + ] + ); + assert_eq!( + responses + .iter() + .map(|response| response.victim_index) + .collect::>(), + vec![0, 1, 2], + "one batch response vector must stay aligned with every input victim" + ); + assert_eq!( + responses + .iter() + .map(|response| response.ssd_backing_committed) + .collect::>(), + vec![true, false, false], + "only the exact fenced write-back victim may publish SSD backing" + ); + assert_eq!( + busy_summary, + DirectDeleteBatchBusySummary { + activity_busy_items: 1, + get_busy_items: 1, + inflight_gets: 1, + ..Default::default() + } + ); + let ready_route = routes + .get(&ready.key) + .expect("SSD write-back must keep the ready route alive"); + let ready_replicas = ready_route.node_replicas.read(); + let ready_backings = ready_replicas + .get(&owner) + .expect("same-owner SSD backing must remain"); + assert!(ready_backings.memory.is_none()); + assert_eq!(ready_backings.ssd.as_ref().map(|ssd| ssd.len), Some(4096)); + drop(ready_replicas); + drop(ready_route); + assert!(routes.contains_key(&busy.key)); + assert!(routes.contains_key(&stale.key)); + assert!( + routes + .get(&busy.key) + .unwrap() + .node_replicas + .read() + .get(&owner) + .unwrap() + .ssd + .is_none(), + "a busy key must not publish SSD backing before its master fence" + ); + assert!( + routes + .get(&stale.key) + .unwrap() + .node_replicas + .read() + .get(&owner) + .unwrap() + .ssd + .is_none(), + "a stale exact backing must not publish SSD backing" + ); + assert!(!activity.has_reclaim(&ready.key)); + + let replay_delete_called = AtomicBool::new(false); + let replay = direct_delete_exact_owner_source_with( + activity.as_ref(), + &owner, + &ready, + owner_source_eviction_epoch(77, 0), + &|key| routes.get(key).map(|route| route.clone()), + |_| { + replay_delete_called.store(true, Ordering::Relaxed); + false + }, + ); + assert_eq!(replay.outcome, OwnerSourceEvictionOutcome::Completed); + assert!( + replay.ssd_backing_committed, + "a lost direct-delete response must replay the already-published SSD terminal" + ); + assert_eq!(replay.busy_cause, None); + assert!(!replay_delete_called.load(Ordering::Relaxed)); + } +} + +/// Why an entry entered the shared safe-reclaim pipeline. +/// +/// Only `MasterAllocationCapacity` may originate from the master resident +/// cache's Size listener. `OwnerCapacityEviction` is an exact source-deletion +/// request selected by the owner and may resolve to a CommittedSlot. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) enum EvictionReclaimOrigin { + MasterAllocationCapacity, + OwnerCapacityEviction, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum EnqueueEvictionReclaimResult { + Accepted, + AlreadyInProgress, + PartialOverlap, + NotInProgress, + Closed, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct EvictionReclaimIdentity { + owner_node_id: NodeIDString, + owner_node_start_time: Option, + key: String, + put_id: (u64, u32), + weight_bytes: u32, + expected_backing: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct EvictionReclaimMember { + pub key: String, + pub desc: NodeValueReplicaDesc, + pub expected_backing: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct EvictionReclaimRequest { + pub owner_node_id: NodeIDString, + pub owner_node_start_time: Option, + pub members: Vec, + pub origin: EvictionReclaimOrigin, + pub retry_count: u32, +} + +impl EvictionReclaimRequest { + pub(crate) fn identities(&self) -> Vec { + self.members + .iter() + .map(|member| EvictionReclaimIdentity { + owner_node_id: self.owner_node_id.clone(), + owner_node_start_time: self.owner_node_start_time, + key: member.key.clone(), + put_id: member.desc.put_id, + weight_bytes: member.desc.weight_bytes, + expected_backing: member.expected_backing.clone(), + }) + .collect() + } + + pub(crate) fn weight_bytes(&self) -> u64 { + self.members + .iter() + .map(|member| u64::from(member.desc.weight_bytes)) + .fold(0u64, u64::saturating_add) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum MasterCapacityPlanError { + RouteChanged, + WrongRole, + CommittedSlot, +} + +fn master_allocation_capacity_weight( + backing: &KvReplicaBacking, +) -> Result { + match backing { + KvReplicaBacking::Allocation(allocation) => { + Ok(u32::try_from(allocation.capcity()).unwrap_or(u32::MAX)) + } + KvReplicaBacking::CommittedSlot(_) => Err(MasterCapacityPlanError::CommittedSlot), + } +} + +fn allocation_member_from_route( + view: &MasterKvRouterView, + owner: &NodeID, + key: &str, +) -> Result< + ( + EvictionReclaimMember, + std::sync::Arc, + ), + MasterCapacityPlanError, +> { + let route = view + .master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|entry| entry.clone()) + .ok_or(MasterCapacityPlanError::RouteChanged)?; + if route.lease_id.is_some() { + return Err(MasterCapacityPlanError::RouteChanged); + } + let replicas = route.node_replicas.read(); + let node_replicas = replicas + .get(owner) + .filter(|replicas| !replicas.tomb_tag.is_tomb()) + .ok_or(MasterCapacityPlanError::RouteChanged)?; + let replica = node_replicas + .memory + .as_ref() + .ok_or(MasterCapacityPlanError::RouteChanged)?; + if replica.owner_local_indexed { + return Err(MasterCapacityPlanError::WrongRole); + } + let weight_bytes = master_allocation_capacity_weight(&replica.backing)?; + drop(replicas); + Ok(( + EvictionReclaimMember { + key: key.to_string(), + desc: NodeValueReplicaDesc { + weight_bytes, + put_id: route.put_id, + }, + expected_backing: None, + }, + route, + )) +} + +/// Validate one exact key popped by the master Allocation Moka. +fn plan_master_allocation_capacity_victim( + view: &MasterKvRouterView, + request: &EvictionReclaimRequest, +) -> Result { + if request.origin != EvictionReclaimOrigin::MasterAllocationCapacity + || request.members.len() != 1 + { + return Err(MasterCapacityPlanError::WrongRole); + } + let owner: NodeID = request.owner_node_id.clone().into(); + let anchor = &request.members[0]; + let (current_anchor, _route) = allocation_member_from_route(view, &owner, &anchor.key)?; + if current_anchor.desc.put_id != anchor.desc.put_id + || current_anchor.desc.weight_bytes != anchor.desc.weight_bytes + { + return Err(MasterCapacityPlanError::RouteChanged); + } + Ok(current_anchor) +} + +fn route_item( + view: &MasterKvRouterView, + owner_node_id: &NodeID, + key: &str, + expected_put_id: Option<(u64, u32)>, + required_slot_size: Option, + reason: OwnerReclaimReason, + epoch: u64, +) -> Option { + let route = view.master_kv_router().inner().kv_routes.get(key)?.clone(); + if expected_put_id.is_some_and(|put_id| put_id != route.put_id) || route.lease_id.is_some() { + return None; + } + let replicas = route.node_replicas.read(); + let node_replicas = replicas.get(owner_node_id)?; + if node_replicas.tomb_tag.is_tomb() { + return None; + } + let target = node_replicas.memory.as_ref()?; + let backing = match &target.backing { + KvReplicaBacking::Allocation(allocation) if required_slot_size.is_none() => { + if target.owner_local_indexed { + OwnerReclaimBacking::Allocation + } else { + OwnerReclaimBacking::UnindexedAllocation { + addr: allocation.base_addr().checked_add(allocation.addr())?, + base_addr: allocation.base_addr(), + len: allocation.size(), + capacity_bytes: allocation.capcity(), + } + } + } + KvReplicaBacking::Allocation(_) => return None, + KvReplicaBacking::CommittedSlot(slot) + if slot.owner_node_id == *owner_node_id + && required_slot_size.map_or(true, |slot_size| slot.slot_size == slot_size) => + { + OwnerReclaimBacking::CommittedSlot { + grant_id: slot.grant_id, + slot_index: slot.slot_index, + slot_size: slot.slot_size, + } + } + KvReplicaBacking::CommittedSlot(_) => return None, + }; + drop(replicas); + Some(OwnerReclaimItem { + key: key.to_string(), + put_id: route.put_id, + epoch, + backing, + reason, + }) +} + +fn item_still_valid(view: &MasterKvRouterView, owner: &NodeID, item: &OwnerReclaimItem) -> bool { + if !view + .master_kv_router() + .inner() + .key_activity + .reclaim_matches(item) + { + return false; + } + route_item( + view, + owner, + &item.key, + Some(item.put_id), + match &item.backing { + OwnerReclaimBacking::Allocation | OwnerReclaimBacking::UnindexedAllocation { .. } => { + None + } + OwnerReclaimBacking::CommittedSlot { slot_size, .. } => Some(*slot_size), + }, + item.reason, + item.epoch, + ) + .is_some_and(|current| current.backing == item.backing && current.reason == item.reason) +} + +/// Publish bytes already persisted by the owner onto the existing route while +/// both master and owner Prepare fences still protect the exact memory +/// generation. Memory remains present until the subsequent Commit succeeds, +/// so Get never observes a route gap between DRAM and SSD. +fn publish_prepared_ssd_backing( + view: &MasterKvRouterView, + owner: &NodeID, + item: &OwnerReclaimItem, + ssd_backing_len: Option, +) -> Result { + let Some(len) = ssd_backing_len else { + return Ok(false); + }; + if !item_still_valid(view, owner, item) { + return Err(format!( + "exact memory route changed before SSD publication: key={} epoch={}", + item.key, item.epoch + )); + } + let route = view + .master_kv_router() + .inner() + .kv_routes + .get(&item.key) + .map(|route| route.clone()) + .ok_or_else(|| format!("route disappeared before SSD publication: key={}", item.key))?; + let existing_ssd_len = route + .node_replicas + .read() + .get(owner) + .and_then(|replicas| replicas.ssd.as_ref()) + .map(|ssd| ssd.len); + if let Some(existing_len) = existing_ssd_len { + return if existing_len == len { + Ok(false) + } else { + Err(format!( + "existing SSD backing length mismatch: key={} existing={} prepared={}", + item.key, existing_len, len + )) + }; + } + match route.commit_ssd_replica(owner, len) { + super::SsdReplicaCommitStatus::Committed => Ok(true), + super::SsdReplicaCommitStatus::MissingMemory => Err(format!( + "memory backing disappeared before SSD publication: key={}", + item.key + )), + super::SsdReplicaCommitStatus::TombedNode => Err(format!( + "owner generation was tombed before SSD publication: owner={} key={}", + owner, item.key + )), + super::SsdReplicaCommitStatus::LengthMismatch => Err(format!( + "persisted SSD length mismatches memory route: key={} len={}", + item.key, len + )), + } +} + +fn rollback_new_ssd_backing( + view: &MasterKvRouterView, + owner: &NodeID, + item: &OwnerReclaimItem, +) -> bool { + if !item_still_valid(view, owner, item) { + return false; + } + view.master_kv_router() + .inner() + .kv_routes + .get(&item.key) + .filter(|route| route.put_id == item.put_id) + .is_some_and(|route| route.remove_ssd_replica(owner)) +} + +async fn call_owner_phase( + view: &MasterKvRouterView, + owner: &NodeID, + phase: OwnerReclaimPhase, + items: Vec, +) -> Result, String> { + if items.is_empty() { + return Ok(Vec::new()); + } + debug_assert!(items.iter().all(|item| { + !matches!( + &item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + ) || item.reason == OwnerReclaimReason::MasterAllocationCapacity + })); + let caller = RPCCaller::::new(); + caller.regist(view.p2p_module()); + let resp = caller + .call( + view.p2p_module(), + owner.clone(), + MsgPack { + serialize_part: BatchOwnerReclaimReq { + phase, + items: items.clone(), + }, + raw_bytes: Vec::new(), + }, + Some(OWNER_RECLAIM_RPC_TIMEOUT), + 1, + ) + .await + .map_err(|err| format!("{err:?}"))?; + if resp.serialize_part.error_code != OK { + return Err(format!( + "code={} error={}", + resp.serialize_part.error_code, resp.serialize_part.error_json + )); + } + if resp.serialize_part.items.len() != items.len() { + return Err(format!( + "owner reclaim response length mismatch: phase={phase:?} expected={} got={}", + items.len(), + resp.serialize_part.items.len() + )); + } + for (request, response) in items.iter().zip(resp.serialize_part.items.iter()) { + if request.key != response.key || request.epoch != response.epoch { + return Err(format!( + "owner reclaim response identity mismatch: phase={phase:?} request=({}, {}) response=({}, {})", + request.key, request.epoch, response.key, response.epoch + )); + } + } + Ok(resp.serialize_part.items) +} + +fn clear_master_fence(view: &MasterKvRouterView, item: &OwnerReclaimItem) { + let cleared = view + .master_kv_router() + .inner() + .key_activity + .clear_reclaim(item); + if !cleared { + tracing::warn!( + "owner reclaim master fence did not match during cleanup: key={} epoch={}", + item.key, + item.epoch + ); + } +} + +async fn abort_prepared( + view: &MasterKvRouterView, + owner: &NodeID, + items: Vec, + newly_published_ssd: HashSet<(String, u64)>, +) -> Vec { + if items.is_empty() { + return Vec::new(); + } + match call_owner_phase(view, owner, OwnerReclaimPhase::Abort, items.clone()).await { + Ok(responses) => { + let mut already_committed = Vec::new(); + for (item, response) in items.into_iter().zip(responses.into_iter()) { + match response.state { + OwnerReclaimItemState::Committed => already_committed.push(item), + OwnerReclaimItemState::Aborted => { + if newly_published_ssd.contains(&(item.key.clone(), item.epoch)) + && !rollback_new_ssd_backing(view, owner, &item) + { + tracing::warn!( + owner = %owner, + key = %item.key, + epoch = item.epoch, + "new SSD backing was not removable after owner reclaim abort" + ); + } + clear_master_fence(view, &item); + } + OwnerReclaimItemState::Stale | OwnerReclaimItemState::Finalized => { + clear_master_fence(view, &item) + } + state => tracing::warn!( + "owner reclaim abort returned unresolved state: key={} epoch={} state={:?} detail={}", + item.key, + item.epoch, + state, + response.detail + ), + } + } + already_committed + } + Err(err) => { + tracing::warn!( + "owner reclaim abort RPC failed; retaining master fences: owner={} keys={} err={}", + owner, + items.len(), + err + ); + spawn_abort_retry(view.clone(), owner.clone(), items, newly_published_ssd); + Vec::new() + } + } +} + +fn spawn_abort_retry( + view: MasterKvRouterView, + owner: NodeID, + items: Vec, + newly_published_ssd: HashSet<(String, u64)>, +) { + if items.is_empty() { + return; + } + let spawn_view = view.clone(); + let _ = spawn_view.spawn("owner_reclaim_abort_retry", async move { + let mut pending = items; + let mut committed = Vec::new(); + let mut delay = Duration::from_millis(25); + for _attempt in 1..=8 { + tokio::time::sleep(delay).await; + match call_owner_phase(&view, &owner, OwnerReclaimPhase::Abort, pending.clone()).await { + Ok(responses) => { + let mut next = Vec::new(); + for (item, response) in pending.into_iter().zip(responses.into_iter()) { + match response.state { + OwnerReclaimItemState::Committed => committed.push(item), + OwnerReclaimItemState::Aborted => { + if newly_published_ssd.contains(&(item.key.clone(), item.epoch)) + && !rollback_new_ssd_backing(&view, &owner, &item) + { + tracing::warn!( + owner = %owner, + key = %item.key, + epoch = item.epoch, + "new SSD backing was not removable after retried owner reclaim abort" + ); + } + clear_master_fence(&view, &item); + } + OwnerReclaimItemState::Stale | OwnerReclaimItemState::Finalized => { + clear_master_fence(&view, &item) + } + _ => next.push(item), + } + } + pending = next; + if pending.is_empty() { + break; + } + } + Err(err) => tracing::warn!( + "owner reclaim abort retry failed: owner={} keys={} err={}", + owner, + pending.len(), + err + ), + } + delay = (delay * 2).min(Duration::from_secs(1)); + } + if !committed.is_empty() { + let _ = finish_committed(&view, &owner, committed).await; + } + if !pending.is_empty() { + tracing::error!( + "owner reclaim abort retry exhausted; fences retained: owner={} keys={}", + owner, + pending.len() + ); + } + }); +} + +fn reclaim_backing_matches( + replica: &super::KvMemoryReplica, + expected: &OwnerReclaimBacking, +) -> bool { + match (&replica.backing, expected) { + (KvReplicaBacking::Allocation(_), OwnerReclaimBacking::Allocation) => { + replica.owner_local_indexed + } + ( + KvReplicaBacking::Allocation(allocation), + OwnerReclaimBacking::UnindexedAllocation { + addr, + base_addr, + len, + capacity_bytes, + }, + ) => { + !replica.owner_local_indexed + && allocation.base_addr().checked_add(allocation.addr()) == Some(*addr) + && allocation.base_addr() == *base_addr + && allocation.size() == *len + && allocation.capcity() == *capacity_bytes + } + ( + KvReplicaBacking::CommittedSlot(slot), + OwnerReclaimBacking::CommittedSlot { + grant_id, + slot_index, + slot_size, + }, + ) => { + slot.grant_id == *grant_id + && slot.slot_index == *slot_index + && slot.slot_size == *slot_size + } + _ => false, + } +} + +fn owner_source_member_weight(backing: &OwnerReclaimBacking) -> Option { + match backing { + OwnerReclaimBacking::CommittedSlot { slot_size, .. } => u32::try_from(*slot_size).ok(), + // Allocation does not carry an address/generation identity in the + // current wire contract, so accepting it would not be an exact delete. + OwnerReclaimBacking::Allocation | OwnerReclaimBacking::UnindexedAllocation { .. } => None, + } +} + +enum OwnerSourceVictimPlan { + Ready(EvictionReclaimMember), + Completed(&'static str), + Stale(String), + Rejected(String), +} + +fn plan_exact_owner_source_victim_with( + owner: &NodeID, + victim: &OwnerSourceEvictionVictim, + route_lookup: &dyn Fn(&str) -> Option>, +) -> OwnerSourceVictimPlan { + let Some(weight_bytes) = owner_source_member_weight(&victim.backing) else { + return OwnerSourceVictimPlan::Rejected(format!( + "source backing is not an exact committed slot: key={}", + victim.key + )); + }; + let desc = NodeValueReplicaDesc { + weight_bytes, + put_id: victim.put_id, + }; + let planned = EvictionReclaimMember { + key: victim.key.clone(), + desc: desc.clone(), + expected_backing: Some(victim.backing.clone()), + }; + + let Some(route) = route_lookup(&victim.key) else { + return OwnerSourceVictimPlan::Completed("exact source replica is already absent"); + }; + if route.put_id != victim.put_id { + return OwnerSourceVictimPlan::Stale(format!( + "route version changed: key={} expected=({},{}) current=({},{})", + victim.key, victim.put_id.0, victim.put_id.1, route.put_id.0, route.put_id.1, + )); + } + if route.lease_id.is_some() { + return OwnerSourceVictimPlan::Rejected(format!( + "leased route is not cache-evictable: key={}", + victim.key + )); + } + let replica_matches = { + let replicas = route.node_replicas.read(); + match replicas.get(owner) { + Some(node_replicas) if !node_replicas.tomb_tag.is_tomb() => { + let Some(replica) = node_replicas.memory.as_ref() else { + return OwnerSourceVictimPlan::Completed( + "exact source memory replica is already absent", + ); + }; + if !replica.owner_local_indexed { + return OwnerSourceVictimPlan::Rejected(format!( + "source route is not owner-local indexed: key={}", + victim.key + )); + } + reclaim_backing_matches(replica, &victim.backing) + } + _ => { + return OwnerSourceVictimPlan::Completed("exact source replica is already absent"); + } + } + }; + if !replica_matches { + return OwnerSourceVictimPlan::Stale(format!( + "source backing changed: key={} put_id=({},{})", + victim.key, victim.put_id.0, victim.put_id.1 + )); + } + OwnerSourceVictimPlan::Ready(planned) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum DirectDeleteBusyCause { + MasterActivity(super::MasterKeyActivitySnapshot), + DeleteUnderFence, +} + +struct DirectDeleteResult { + outcome: OwnerSourceEvictionOutcome, + ssd_backing_committed: bool, + detail: String, + busy_cause: Option, +} + +impl DirectDeleteResult { + fn terminal(outcome: OwnerSourceEvictionOutcome, detail: impl Into) -> Self { + Self { + outcome, + ssd_backing_committed: false, + detail: detail.into(), + busy_cause: None, + } + } + + fn activity_busy(snapshot: super::MasterKeyActivitySnapshot) -> Self { + Self { + outcome: OwnerSourceEvictionOutcome::RetryableBusy, + ssd_backing_committed: false, + detail: format!( + "master key activity is busy: puts={} gets={} replicas={} reclaim_installed={}", + snapshot.puts, snapshot.gets, snapshot.replicas, snapshot.reclaim_installed + ), + busy_cause: Some(DirectDeleteBusyCause::MasterActivity(snapshot)), + } + } + + fn delete_under_fence_busy() -> Self { + Self { + outcome: OwnerSourceEvictionOutcome::RetryableBusy, + ssd_backing_committed: false, + detail: "exact source route could not be deleted under its master fence".to_string(), + busy_cause: Some(DirectDeleteBusyCause::DeleteUnderFence), + } + } +} + +fn exact_ssd_writeback_is_published( + owner: &NodeID, + victim: &OwnerSourceEvictionVictim, + route_lookup: &dyn Fn(&str) -> Option>, +) -> bool { + if victim.ssd_policy != OwnerSourceSsdPolicy::Persisted { + return false; + } + let Some(expected_len) = victim.ssd_backing_len else { + return false; + }; + let Some(route) = route_lookup(&victim.key) else { + return false; + }; + if route.put_id != victim.put_id { + return false; + } + route + .node_replicas + .read() + .get(owner) + .is_some_and(|replicas| { + !replicas.tomb_tag.is_tomb() + && replicas + .ssd + .as_ref() + .is_some_and(|ssd| ssd.len == expected_len) + }) +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct DirectDeleteBatchBusySummary { + activity_busy_items: u64, + put_busy_items: u64, + get_busy_items: u64, + replica_busy_items: u64, + reclaim_busy_items: u64, + inflight_puts: u64, + inflight_gets: u64, + inflight_replicas: u64, + delete_under_fence_busy_items: u64, +} + +impl DirectDeleteBatchBusySummary { + fn record(&mut self, cause: Option) { + match cause { + Some(DirectDeleteBusyCause::MasterActivity(snapshot)) => { + self.activity_busy_items = self.activity_busy_items.saturating_add(1); + self.put_busy_items = self + .put_busy_items + .saturating_add(u64::from(snapshot.puts != 0)); + self.get_busy_items = self + .get_busy_items + .saturating_add(u64::from(snapshot.gets != 0)); + self.replica_busy_items = self + .replica_busy_items + .saturating_add(u64::from(snapshot.replicas != 0)); + self.reclaim_busy_items = self + .reclaim_busy_items + .saturating_add(u64::from(snapshot.reclaim_installed)); + self.inflight_puts = self.inflight_puts.saturating_add(u64::from(snapshot.puts)); + self.inflight_gets = self.inflight_gets.saturating_add(u64::from(snapshot.gets)); + self.inflight_replicas = self + .inflight_replicas + .saturating_add(u64::from(snapshot.replicas)); + } + Some(DirectDeleteBusyCause::DeleteUnderFence) => { + self.delete_under_fence_busy_items = + self.delete_under_fence_busy_items.saturating_add(1); + } + None => {} + } + } +} + +fn direct_delete_exact_owner_source_with( + activity: &super::MasterKeyActivityTable, + owner: &NodeID, + victim: &OwnerSourceEvictionVictim, + epoch: u64, + route_lookup: &dyn Fn(&str) -> Option>, + delete: impl FnOnce(&OwnerReclaimItem) -> bool, +) -> DirectDeleteResult { + let member = match plan_exact_owner_source_victim_with(owner, victim, route_lookup) { + OwnerSourceVictimPlan::Ready(member) => member, + OwnerSourceVictimPlan::Completed(detail) => { + let mut result = + DirectDeleteResult::terminal(OwnerSourceEvictionOutcome::Completed, detail); + result.ssd_backing_committed = + exact_ssd_writeback_is_published(owner, victim, route_lookup); + return result; + } + OwnerSourceVictimPlan::Stale(detail) => { + return DirectDeleteResult::terminal(OwnerSourceEvictionOutcome::Stale, detail); + } + OwnerSourceVictimPlan::Rejected(detail) => { + return DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::RejectedNotEvictable, + detail, + ); + } + }; + let item = OwnerReclaimItem { + key: member.key, + put_id: member.desc.put_id, + epoch, + backing: member + .expected_backing + .expect("exact owner source plan must retain its backing"), + reason: OwnerReclaimReason::OwnerCapacityEviction, + }; + if let Err(snapshot) = activity.try_install_reclaim(&item) { + return DirectDeleteResult::activity_busy(snapshot); + } + + let result = match plan_exact_owner_source_victim_with(owner, victim, route_lookup) { + OwnerSourceVictimPlan::Ready(_) => { + let ssd_ready = match victim.ssd_policy { + OwnerSourceSsdPolicy::Drop => { + if victim.ssd_backing_len.is_none() { + Ok(()) + } else { + Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::RejectedNotEvictable, + format!("drop policy must not carry SSD bytes: key={}", victim.key), + )) + } + } + OwnerSourceSsdPolicy::SelectLastLive => { + if victim.ssd_backing_len.is_some() { + Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::RejectedNotEvictable, + format!( + "SSD selection policy must not carry prepared bytes: key={}", + victim.key + ), + )) + } else if exact_memory_reclaim_needs_ssd_with(owner, &item, route_lookup) { + Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::SsdCandidate, + "exact source is the last live backing; owner SSD admission required", + )) + } else { + Ok(()) + } + } + OwnerSourceSsdPolicy::Persisted => match victim.ssd_backing_len { + None => Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::RejectedNotEvictable, + format!( + "persisted SSD policy requires exact bytes: key={}", + victim.key + ), + )), + Some(len) => match route_lookup(&victim.key) + .map(|route| route.commit_ssd_replica(owner, len)) + { + Some(super::SsdReplicaCommitStatus::Committed) => Ok(()), + Some(super::SsdReplicaCommitStatus::LengthMismatch) => { + Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::RejectedNotEvictable, + format!( + "SSD write-back length does not match the exact source: key={} len={}", + victim.key, len + ), + )) + } + Some( + super::SsdReplicaCommitStatus::MissingMemory + | super::SsdReplicaCommitStatus::TombedNode, + ) + | None => Err(DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::Stale, + format!( + "exact source changed before SSD write-back publication: key={}", + victim.key + ), + )), + }, + }, + }; + match ssd_ready { + Err(result) => result, + Ok(()) => { + let mut result = if delete(&item) { + DirectDeleteResult::terminal( + OwnerSourceEvictionOutcome::Completed, + if victim.ssd_policy == OwnerSourceSsdPolicy::Persisted { + "SSD backing published and exact memory source deleted by batch handler" + } else { + "exact source route deleted by batch handler" + }, + ) + } else { + DirectDeleteResult::delete_under_fence_busy() + }; + result.ssd_backing_committed = + victim.ssd_policy == OwnerSourceSsdPolicy::Persisted; + result + } + } + } + OwnerSourceVictimPlan::Completed(detail) => { + let mut result = + DirectDeleteResult::terminal(OwnerSourceEvictionOutcome::Completed, detail); + result.ssd_backing_committed = + exact_ssd_writeback_is_published(owner, victim, route_lookup); + result + } + OwnerSourceVictimPlan::Stale(detail) => { + DirectDeleteResult::terminal(OwnerSourceEvictionOutcome::Stale, detail) + } + OwnerSourceVictimPlan::Rejected(detail) => { + DirectDeleteResult::terminal(OwnerSourceEvictionOutcome::RejectedNotEvictable, detail) + } + }; + assert!( + activity.clear_reclaim(&item), + "direct-delete master fence must remain installed until route deletion completes" + ); + result +} + +fn direct_delete_exact_owner_source_batch_with( + activity: &super::MasterKeyActivityTable, + owner: &NodeID, + operation_id: u64, + victims: &[OwnerSourceEvictionVictim], + route_lookup: &dyn Fn(&str) -> Option>, + delete: impl Fn(&OwnerReclaimItem) -> bool, +) -> ( + Vec, + DirectDeleteBatchBusySummary, +) { + let mut responses = Vec::with_capacity(victims.len()); + let mut busy = DirectDeleteBatchBusySummary::default(); + for (index, victim) in victims.iter().enumerate() { + let result = direct_delete_exact_owner_source_with( + activity, + owner, + victim, + owner_source_eviction_epoch(operation_id, index), + route_lookup, + |item| delete(item), + ); + busy.record(result.busy_cause); + responses.push(OwnerSourceEvictionVictimResp { + victim_index: u32::try_from(index).unwrap_or(u32::MAX), + outcome: result.outcome, + ssd_backing_committed: result.ssd_backing_committed, + detail: result.detail, + }); + } + (responses, busy) +} + +pub(crate) async fn handle_batch_evict_owner_source( + view: &MasterKvRouterView, + req: MsgPack, + owner: NodeID, +) -> MsgPack { + let operation_id = req.serialize_part.operation_id; + let counters = view + .master_kv_router() + .eviction_reclaim_counters(owner.as_ref()); + counters + .source_evict_rpc_requests + .fetch_add(1, Ordering::Relaxed); + counters.source_evict_victims.fetch_add( + u64::try_from(req.serialize_part.victims.len()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + let requested_bytes = req + .serialize_part + .victims + .iter() + .filter_map(|victim| owner_source_member_weight(&victim.backing)) + .map(u64::from) + .fold(0u64, u64::saturating_add); + counters + .source_evict_requested_bytes + .fetch_add(requested_bytes, Ordering::Relaxed); + let current_generation = view + .cluster_manager() + .get_member_info_cached(owner.as_ref()) + .map(|member| member.node_start_time); + if current_generation != Some(req.serialize_part.owner_node_start_time) { + counters.source_evict_rejected.fetch_add( + u64::try_from(req.serialize_part.victims.len()).unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + let err = KvError::Api(ApiError::InvalidArgument { + detail: format!( + "owner source-eviction generation mismatch: owner={} requested={} current={:?}", + owner, req.serialize_part.owner_node_start_time, current_generation + ), + }); + return MsgPack { + serialize_part: BatchEvictOwnerSourceResp { + operation_id, + victims: Vec::new(), + error_code: err.code(), + error_json: err.to_json(), + }, + raw_bytes: Vec::new(), + }; + } + + let (responses, busy) = direct_delete_exact_owner_source_batch_with( + &view.master_kv_router().inner().key_activity, + &owner, + operation_id, + &req.serialize_part.victims, + &|key| { + view.master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|route| route.clone()) + }, + |item| remove_reclaimed_replica(view, &owner, item), + ); + for response in &responses { + let outcome_counter = match response.outcome { + OwnerSourceEvictionOutcome::Accepted => &counters.source_evict_accepted, + OwnerSourceEvictionOutcome::AlreadyInProgress => &counters.source_evict_in_progress, + OwnerSourceEvictionOutcome::Completed => &counters.source_evict_completed, + OwnerSourceEvictionOutcome::SsdCandidate => &counters.source_evict_accepted, + OwnerSourceEvictionOutcome::RetryableBusy | OwnerSourceEvictionOutcome::Unspecified => { + &counters.source_evict_retryable_busy + } + OwnerSourceEvictionOutcome::Stale => &counters.source_evict_stale, + OwnerSourceEvictionOutcome::RejectedNotEvictable => &counters.source_evict_rejected, + }; + outcome_counter.fetch_add(1, Ordering::Relaxed); + } + + let completed = responses + .iter() + .filter(|response| response.outcome == OwnerSourceEvictionOutcome::Completed) + .count(); + let retryable = responses + .iter() + .filter(|response| response.outcome == OwnerSourceEvictionOutcome::RetryableBusy) + .count(); + let ssd_candidates = responses + .iter() + .filter(|response| response.outcome == OwnerSourceEvictionOutcome::SsdCandidate) + .count(); + tracing::info!( + owner = %owner, + operation_id, + victims = responses.len(), + completed, + ssd_candidates, + retryable, + activity_busy_items = busy.activity_busy_items, + put_busy_items = busy.put_busy_items, + get_busy_items = busy.get_busy_items, + replica_busy_items = busy.replica_busy_items, + reclaim_busy_items = busy.reclaim_busy_items, + inflight_puts = busy.inflight_puts, + inflight_gets = busy.inflight_gets, + inflight_replicas = busy.inflight_replicas, + delete_under_fence_busy_items = busy.delete_under_fence_busy_items, + "owner source direct-delete batch completed" + ); + + MsgPack { + serialize_part: BatchEvictOwnerSourceResp { + operation_id, + victims: responses, + error_code: OK, + error_json: String::new(), + }, + raw_bytes: Vec::new(), + } +} + +struct RemovedOwnerSource { + desc: NodeValueReplicaDesc, + capacity_bytes: u64, + logical_bytes: u64, + removed_last_route: bool, + ssd_survived: bool, + ssd_became_only_backing: bool, +} + +fn record_last_route_removal( + counters: &super::EvictionReclaimCounters, + removed: &RemovedOwnerSource, +) { + if !removed.removed_last_route { + return; + } + counters + .last_route_removed_members + .fetch_add(1, Ordering::Relaxed); + counters + .last_route_removed_bytes + .fetch_add(removed.capacity_bytes, Ordering::Relaxed); +} + +fn remove_exact_owner_source_route( + routes: &dashmap::DashMap>, + owner: &NodeID, + item: &OwnerReclaimItem, +) -> Option { + let route = routes.get(&item.key).map(|route| route.clone())?; + if route.put_id != item.put_id { + return None; + } + let removed_desc = { + let mut replicas = route.node_replicas.write(); + let Some(node_replicas) = replicas.get_mut(owner) else { + return None; + }; + if node_replicas.tomb_tag.is_tomb() { + return None; + } + let Some(replica) = node_replicas.memory.as_ref() else { + return None; + }; + if !reclaim_backing_matches(replica, &item.backing) { + return None; + } + let capacity_bytes = replica.backing.capacity_bytes(); + let logical_bytes = replica.backing.len(); + let desc = NodeValueReplicaDesc { + weight_bytes: u32::try_from(capacity_bytes).unwrap_or(u32::MAX), + put_id: route.put_id, + }; + node_replicas.memory.take(); + let ssd_survived = node_replicas.ssd.is_some(); + let remove_node_entry = !ssd_survived; + if remove_node_entry { + replicas.remove(owner); + } + let live_backings = replicas + .values() + .filter(|replicas| !replicas.tomb_tag.is_tomb()) + .map(|replicas| { + usize::from(replicas.memory.is_some()) + usize::from(replicas.ssd.is_some()) + }) + .sum::(); + let ssd_became_only_backing = ssd_survived && live_backings == 1; + Some(( + desc, + capacity_bytes, + logical_bytes, + ssd_survived, + ssd_became_only_backing, + )) + }; + let (removed_desc, capacity_bytes, logical_bytes, ssd_survived, ssd_became_only_backing) = + removed_desc?; + + let removed_last_route = if route.node_replicas.read().is_empty() { + routes + .remove_if(&item.key, |_, current| { + Arc::ptr_eq(current, &route) + && current.put_id == item.put_id + && current.node_replicas.read().is_empty() + }) + .is_some() + } else { + false + }; + Some(RemovedOwnerSource { + desc: removed_desc, + capacity_bytes, + logical_bytes, + removed_last_route, + ssd_survived, + ssd_became_only_backing, + }) +} + +fn remove_reclaimed_replica( + view: &MasterKvRouterView, + owner: &NodeID, + item: &OwnerReclaimItem, +) -> bool { + if !view + .master_kv_router() + .inner() + .key_activity + .reclaim_matches(item) + { + return false; + } + let removed = + remove_exact_owner_source_route(&view.master_kv_router().inner().kv_routes, owner, item); + if let Some(removed) = removed { + let ssd_tier = &view.master_kv_router().inner().ssd_tier_counters; + if removed.ssd_survived { + ssd_tier + .memory_removed_ssd_survived_items + .fetch_add(1, Ordering::Relaxed); + ssd_tier + .memory_removed_ssd_survived_bytes + .fetch_add(removed.logical_bytes, Ordering::Relaxed); + } + if removed.ssd_became_only_backing { + ssd_tier + .memory_removed_ssd_became_only_items + .fetch_add(1, Ordering::Relaxed); + ssd_tier + .memory_removed_ssd_became_only_bytes + .fetch_add(removed.logical_bytes, Ordering::Relaxed); + } + let counters = view + .master_kv_router() + .eviction_reclaim_counters(owner.as_ref()); + record_last_route_removal(counters.as_ref(), &removed); + if removed.removed_last_route && view.master_kv_router().prefix_index_enabled() { + let view_task = view.clone(); + let key = item.key.clone(); + let put_id = item.put_id; + let spawn_view = view.clone(); + let _ = spawn_view.spawn("owner_reclaim_remove_prefix", async move { + let mut tree = view_task + .master_kv_router() + .inner() + .prefix_index + .write() + .await; + tree.remove(&key, put_id); + }); + } + if let Some(cache) = view + .master_kv_router() + .inner() + .node_kv_cache_controller + .get(owner.as_ref()) + { + let _ = super::remove_exact_cache_entry(cache.value(), &item.key, &removed.desc); + } + if let Some(cache) = view + .master_kv_router() + .inner() + .node_writeback_tier1_controller + .get(owner.as_ref()) + { + let _ = super::remove_exact_cache_entry(cache.value(), &item.key, &removed.desc); + } + } + true +} + +fn finish_unindexed_allocations( + view: &MasterKvRouterView, + owner: &NodeID, + items: Vec, +) -> u32 { + let mut reclaimed = 0u32; + for item in items { + debug_assert!(matches!( + &item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + )); + if remove_reclaimed_replica(view, owner, &item) { + clear_master_fence(view, &item); + reclaimed = reclaimed.saturating_add(1); + } else { + tracing::error!( + "unindexed allocation reclaim backing could not be removed from master route: owner={} key={} epoch={}", + owner, + item.key, + item.epoch + ); + } + } + reclaimed +} + +/// Return true only when removing this exact memory source would leave the key +/// with no live backing. The caller holds the per-key master reclaim fence, so +/// this route snapshot cannot race another Put/Get/replica/reclaim mutation. +fn exact_memory_reclaim_needs_ssd_with( + owner: &NodeID, + item: &OwnerReclaimItem, + route_lookup: &dyn Fn(&str) -> Option>, +) -> bool { + let Some(route) = route_lookup(&item.key) else { + return false; + }; + if route.put_id != item.put_id || route.lease_id.is_some() { + return false; + } + let replicas = route.node_replicas.read(); + let Some(target) = replicas.get(owner) else { + return false; + }; + if target.tomb_tag.is_tomb() + || !target + .memory + .as_ref() + .is_some_and(|memory| reclaim_backing_matches(memory, &item.backing)) + { + return false; + } + + replicas.iter().all(|(node_id, node_replicas)| { + if node_replicas.tomb_tag.is_tomb() { + return true; + } + if node_id == owner { + // The exact memory above is the source being removed. An existing + // same-owner SSD backing already preserves this route. + node_replicas.ssd.is_none() + } else { + !node_replicas.has_live_backing() + } + }) +} + +fn exact_memory_reclaim_needs_ssd( + view: &MasterKvRouterView, + owner: &NodeID, + item: &OwnerReclaimItem, +) -> bool { + exact_memory_reclaim_needs_ssd_with(owner, item, &|key| { + view.master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|route| route.clone()) + }) +} + +fn partition_reclaim_coordination( + items: Vec, + should_coordinate_unindexed: F, +) -> (Vec, Vec) +where + F: Fn(&OwnerReclaimItem) -> bool, +{ + items.into_iter().partition(|item| { + matches!( + &item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + ) && !should_coordinate_unindexed(item) + }) +} + +fn owner_has_ssd_storage(view: &MasterKvRouterView, owner: &NodeID) -> bool { + view.cluster_manager() + .get_member_info_cached(owner.as_ref()) + .and_then(|member| { + member + .metadata + .get(crate::cluster_manager::META_KEY_KV_SSD_STORAGE) + .cloned() + }) + .is_some_and(|value| value == "true") +} + +async fn finish_committed( + view: &MasterKvRouterView, + owner: &NodeID, + items: Vec, +) -> u32 { + let mut removed = Vec::new(); + for item in items { + if remove_reclaimed_replica(view, owner, &item) { + removed.push(item); + } else { + tracing::error!( + "owner reclaim backing could not be removed from master route: owner={} key={} epoch={}", + owner, + item.key, + item.epoch + ); + } + } + if removed.is_empty() { + return 0; + } + match call_owner_phase(view, owner, OwnerReclaimPhase::Finalize, removed.clone()).await { + Ok(responses) => { + let mut finalized = 0u32; + let mut retry = Vec::new(); + for (item, response) in removed.into_iter().zip(responses.into_iter()) { + if response.state == OwnerReclaimItemState::Finalized { + clear_master_fence(view, &item); + finalized = finalized.saturating_add(1); + } else { + tracing::warn!( + "owner reclaim finalize returned unresolved state: owner={} key={} epoch={} state={:?} detail={}", + owner, + item.key, + item.epoch, + response.state, + response.detail + ); + retry.push(item); + } + } + spawn_finalize_retry(view.clone(), owner.clone(), retry); + finalized + } + Err(err) => { + tracing::warn!( + "owner reclaim finalize RPC failed; retaining both fences: owner={} keys={} err={}", + owner, + removed.len(), + err + ); + spawn_finalize_retry(view.clone(), owner.clone(), removed); + 0 + } + } +} + +fn spawn_finalize_retry(view: MasterKvRouterView, owner: NodeID, items: Vec) { + if items.is_empty() { + return; + } + let spawn_view = view.clone(); + let _ = spawn_view.spawn("owner_reclaim_finalize_retry", async move { + let mut pending = items; + let mut delay = Duration::from_millis(25); + for _attempt in 1..=8 { + tokio::time::sleep(delay).await; + match call_owner_phase(&view, &owner, OwnerReclaimPhase::Finalize, pending.clone()) + .await + { + Ok(responses) => { + let mut next = Vec::new(); + for (item, response) in pending.into_iter().zip(responses.into_iter()) { + if response.state == OwnerReclaimItemState::Finalized { + clear_master_fence(&view, &item); + } else { + next.push(item); + } + } + pending = next; + if pending.is_empty() { + return; + } + } + Err(err) => tracing::warn!( + "owner reclaim finalize retry failed: owner={} keys={} err={}", + owner, + pending.len(), + err + ), + } + delay = (delay * 2).min(Duration::from_secs(1)); + } + tracing::error!( + "owner reclaim finalize retry exhausted; fences retained: owner={} keys={}", + owner, + pending.len() + ); + }); +} + +async fn reclaim_items( + view: &MasterKvRouterView, + owner: &NodeID, + candidates: Vec, +) -> u32 { + let counters = view + .master_kv_router() + .eviction_reclaim_counters(owner.as_ref()); + let mut fenced = Vec::new(); + for item in candidates { + match view + .master_kv_router() + .inner() + .key_activity + .try_install_reclaim(&item) + { + Ok(()) => { + if item_still_valid(view, owner, &item) { + fenced.push(item); + } else { + counters.route_changed.fetch_add(1, Ordering::Relaxed); + clear_master_fence(view, &item); + } + } + Err(activity) => { + counters + .master_activity_deferred + .fetch_add(1, Ordering::Relaxed); + tracing::trace!( + "owner reclaim deferred by master activity: owner={} key={} puts={} gets={} replicas={} reclaim_installed={}", + owner, + item.key, + activity.puts, + activity.gets, + activity.replicas, + activity.reclaim_installed + ); + } + } + } + if fenced.is_empty() { + return 0; + } + + let owner_has_ssd = owner_has_ssd_storage(view, owner); + let (master_only, owner_coordinated) = partition_reclaim_coordination(fenced, |item| { + owner_has_ssd && exact_memory_reclaim_needs_ssd(view, owner, item) + }); + let master_reclaimed = finish_unindexed_allocations(view, owner, master_only); + if owner_coordinated.is_empty() { + counters + .completed + .fetch_add(u64::from(master_reclaimed), Ordering::Relaxed); + return master_reclaimed; + } + + let prepare_responses = match call_owner_phase( + view, + owner, + OwnerReclaimPhase::Prepare, + owner_coordinated.clone(), + ) + .await + { + Ok(responses) => responses, + Err(err) => { + tracing::warn!( + "owner reclaim prepare RPC failed; aborting batch: owner={} keys={} err={}", + owner, + owner_coordinated.len(), + err + ); + let _ = abort_prepared(view, owner, owner_coordinated, HashSet::new()).await; + counters + .completed + .fetch_add(u64::from(master_reclaimed), Ordering::Relaxed); + return master_reclaimed; + } + }; + let mut prepared = Vec::new(); + let mut committed = Vec::new(); + let mut ssd_publish_failed = Vec::new(); + let mut newly_published_ssd = HashSet::new(); + for (item, response) in owner_coordinated + .into_iter() + .zip(prepare_responses.into_iter()) + { + match response.state { + OwnerReclaimItemState::Prepared => { + match publish_prepared_ssd_backing(view, owner, &item, response.ssd_backing_len) { + Ok(newly_published) => { + if newly_published { + newly_published_ssd.insert((item.key.clone(), item.epoch)); + } + prepared.push(item); + } + Err(detail) => { + counters.route_changed.fetch_add(1, Ordering::Relaxed); + tracing::warn!( + owner = %owner, + key = %item.key, + epoch = item.epoch, + detail = %detail, + "aborting owner reclaim after SSD backing publication validation failed" + ); + ssd_publish_failed.push(item); + } + } + } + OwnerReclaimItemState::Committed => committed.push(item), + OwnerReclaimItemState::Busy => { + if response.detail == "owner local memory still has active holders" { + counters + .owner_holder_deferred + .fetch_add(1, Ordering::Relaxed); + } else { + counters + .owner_other_deferred + .fetch_add(1, Ordering::Relaxed); + } + clear_master_fence(view, &item); + } + _ => { + counters + .owner_other_deferred + .fetch_add(1, Ordering::Relaxed); + clear_master_fence(view, &item); + } + } + } + committed.extend(abort_prepared(view, owner, ssd_publish_failed, HashSet::new()).await); + + let mut invalid_prepared = Vec::new(); + prepared.retain(|item| { + let valid = item_still_valid(view, owner, item); + if !valid { + counters.route_changed.fetch_add(1, Ordering::Relaxed); + invalid_prepared.push(item.clone()); + } + valid + }); + let invalid_published = invalid_prepared + .iter() + .filter_map(|item| { + let identity = (item.key.clone(), item.epoch); + newly_published_ssd.contains(&identity).then_some(identity) + }) + .collect(); + committed.extend(abort_prepared(view, owner, invalid_prepared, invalid_published).await); + + if !prepared.is_empty() { + match call_owner_phase(view, owner, OwnerReclaimPhase::Commit, prepared.clone()).await { + Ok(responses) => { + let mut unresolved = Vec::new(); + for (item, response) in prepared.into_iter().zip(responses.into_iter()) { + if response.state == OwnerReclaimItemState::Committed { + committed.push(item); + } else { + unresolved.push(item); + } + } + let unresolved_published = unresolved + .iter() + .filter_map(|item| { + let identity = (item.key.clone(), item.epoch); + newly_published_ssd.contains(&identity).then_some(identity) + }) + .collect(); + committed + .extend(abort_prepared(view, owner, unresolved, unresolved_published).await); + } + Err(err) => { + tracing::warn!( + "owner reclaim commit RPC failed; resolving with abort: owner={} keys={} err={}", + owner, + prepared.len(), + err + ); + let prepared_published = prepared + .iter() + .filter_map(|item| { + let identity = (item.key.clone(), item.epoch); + newly_published_ssd.contains(&identity).then_some(identity) + }) + .collect(); + committed.extend(abort_prepared(view, owner, prepared, prepared_published).await); + } + } + } + let reclaimed = master_reclaimed.saturating_add(finish_committed(view, owner, committed).await); + counters + .completed + .fetch_add(u64::from(reclaimed), Ordering::Relaxed); + reclaimed +} + +fn clear_master_fences(view: &MasterKvRouterView, items: &[OwnerReclaimItem]) { + for item in items { + clear_master_fence(view, item); + } +} + +fn try_install_master_fences( + activity: &super::MasterKeyActivityTable, + items: &[OwnerReclaimItem], +) -> Result<(), (usize, super::MasterKeyActivitySnapshot)> { + let mut installed = Vec::with_capacity(items.len()); + for (index, item) in items.iter().enumerate() { + match activity.try_install_reclaim(item) { + Ok(()) => installed.push(item), + Err(snapshot) => { + for installed_item in installed { + assert!(activity.clear_reclaim(installed_item)); + } + return Err((index, snapshot)); + } + } + } + Ok(()) +} + +/// Reclaim one independently selected key. +async fn reclaim_single_victim( + view: &MasterKvRouterView, + owner: &NodeID, + items: Vec, +) -> u32 { + if items.len() != 1 { + tracing::error!( + owner = %owner, + victims = items.len(), + "single-key reclaim received a non-singleton request" + ); + return 0; + } + let counters = view + .master_kv_router() + .eviction_reclaim_counters(owner.as_ref()); + if let Err((failed_index, activity)) = + try_install_master_fences(&view.master_kv_router().inner().key_activity, &items) + { + counters + .master_activity_deferred + .fetch_add(1, Ordering::Relaxed); + tracing::trace!( + "single-key reclaim deferred by master activity: owner={} key={} puts={} gets={} replicas={} reclaim_installed={}", + owner, + items[failed_index].key, + activity.puts, + activity.gets, + activity.replicas, + activity.reclaim_installed, + ); + return 0; + } + let fenced = items; + if fenced + .iter() + .any(|item| !item_still_valid(view, owner, item)) + { + clear_master_fences(view, &fenced); + return 0; + } + + let owner_has_ssd = owner_has_ssd_storage(view, owner); + let (master_only, owner_coordinated) = partition_reclaim_coordination(fenced.clone(), |item| { + owner_has_ssd && exact_memory_reclaim_needs_ssd(view, owner, item) + }); + if owner_coordinated.is_empty() { + // All master-owned allocations are fenced and revalidated before the + // first route mutation, so no member can be admitted independently. + if master_only.len() != fenced.len() + || master_only + .iter() + .any(|item| !item_still_valid(view, owner, item)) + { + clear_master_fences(view, &fenced); + return 0; + } + let reclaimed = finish_unindexed_allocations(view, owner, master_only); + counters + .completed + .fetch_add(u64::from(reclaimed), Ordering::Relaxed); + return reclaimed; + } + if !master_only.is_empty() || owner_coordinated.len() != fenced.len() { + tracing::error!( + "BUG: one single-key reclaim mixed master-only and owner-coordinated backings: owner={} victims={} master_only={} owner_coordinated={}", + owner, + fenced.len(), + master_only.len(), + owner_coordinated.len(), + ); + clear_master_fences(view, &fenced); + return 0; + } + + let all_by_key = owner_coordinated + .iter() + .cloned() + .map(|item| (item.key.clone(), item)) + .collect::>(); + let mut committed_keys = HashSet::new(); + let mut delay = Duration::from_millis(25); + let mut rounds = 0u32; + loop { + let pending = all_by_key + .iter() + .filter(|(key, _)| !committed_keys.contains(*key)) + .map(|(_, item)| item.clone()) + .collect::>(); + if pending.is_empty() { + let reclaimed = finish_committed(view, owner, owner_coordinated).await; + counters + .completed + .fetch_add(u64::from(reclaimed), Ordering::Relaxed); + return reclaimed; + } + + rounds = rounds.saturating_add(1); + let prepare = + call_owner_phase(view, owner, OwnerReclaimPhase::Prepare, pending.clone()).await; + let Ok(prepare_responses) = prepare else { + // No member is known committed yet. Abort is both rollback and + // response-loss resolution: a Committed response moves us onto + // the mandatory roll-forward branch. + if committed_keys.is_empty() { + if let Ok(abort_responses) = + call_owner_phase(view, owner, OwnerReclaimPhase::Abort, pending.clone()).await + { + for (item, response) in pending.iter().zip(abort_responses) { + if response.state == OwnerReclaimItemState::Committed { + committed_keys.insert(item.key.clone()); + } + } + if committed_keys.is_empty() { + clear_master_fences(view, &fenced); + return 0; + } + } + } + tokio::time::sleep(delay).await; + delay = (delay * 2).min(Duration::from_secs(1)); + continue; + }; + + let mut prepared = Vec::new(); + let mut blocked = false; + for (item, response) in pending.iter().cloned().zip(prepare_responses) { + match response.state { + OwnerReclaimItemState::Prepared => { + match publish_prepared_ssd_backing(view, owner, &item, response.ssd_backing_len) + { + Ok(_) => prepared.push(item), + Err(detail) => { + counters.route_changed.fetch_add(1, Ordering::Relaxed); + tracing::warn!( + owner = %owner, + key = %item.key, + epoch = item.epoch, + detail = %detail, + "single-key reclaim SSD publication validation failed" + ); + blocked = true; + } + } + } + OwnerReclaimItemState::Committed => { + committed_keys.insert(item.key); + } + _ => blocked = true, + } + } + + if blocked && committed_keys.is_empty() { + // Nothing irreversible happened. Abort every possibly-prepared + // member, and only roll back after the response proves that none + // had crossed Commit during a lost response. + match call_owner_phase(view, owner, OwnerReclaimPhase::Abort, pending.clone()).await { + Ok(responses) => { + for (item, response) in pending.iter().zip(responses) { + if response.state == OwnerReclaimItemState::Committed { + committed_keys.insert(item.key.clone()); + } + } + if committed_keys.is_empty() { + clear_master_fences(view, &fenced); + return 0; + } + } + Err(_) => { + tokio::time::sleep(delay).await; + delay = (delay * 2).min(Duration::from_secs(1)); + continue; + } + } + } + + // With no blocked member this is the first atomic Commit attempt. If + // another member was already observed Committed, this is mandatory + // roll-forward for the rest of the transaction. + if !prepared.is_empty() && (!blocked || !committed_keys.is_empty()) { + if let Ok(commit_responses) = + call_owner_phase(view, owner, OwnerReclaimPhase::Commit, prepared.clone()).await + { + for (item, response) in prepared.iter().zip(commit_responses) { + if response.state == OwnerReclaimItemState::Committed { + committed_keys.insert(item.key.clone()); + } + } + } + } + + if rounds == 8 && !committed_keys.is_empty() { + tracing::warn!( + "owner single-key reclaim is rolling forward after uncertain commit: owner={} victims={} committed={}", + owner, + owner_coordinated.len(), + committed_keys.len(), + ); + } + tokio::time::sleep(delay).await; + delay = (delay * 2).min(Duration::from_secs(1)); + } +} + +#[cfg(test)] +mod reclaim_partition_tests { + use super::{ + OwnerReclaimBacking, OwnerReclaimItem, OwnerReclaimReason, partition_reclaim_coordination, + }; + + fn candidate(index: u32) -> OwnerReclaimItem { + OwnerReclaimItem { + key: format!("candidate-{index}"), + put_id: (u64::from(index), 0), + epoch: u64::from(index), + backing: OwnerReclaimBacking::CommittedSlot { + grant_id: u64::from(index), + slot_index: index, + slot_size: 8 * 1024 * 1024, + }, + reason: OwnerReclaimReason::OwnerCapacityEviction, + } + } + + #[test] + fn only_unindexed_allocations_skip_owner_coordination() { + let mut indexed_allocation = candidate(1); + indexed_allocation.backing = OwnerReclaimBacking::Allocation; + indexed_allocation.reason = OwnerReclaimReason::OwnerCapacityEviction; + let mut unindexed_allocation = candidate(2); + unindexed_allocation.backing = OwnerReclaimBacking::UnindexedAllocation { + addr: 0x2000, + base_addr: 0x1000, + len: 4096, + capacity_bytes: 4096, + }; + unindexed_allocation.reason = OwnerReclaimReason::MasterAllocationCapacity; + let committed_slot = candidate(3); + + let candidates = vec![ + indexed_allocation.clone(), + unindexed_allocation.clone(), + committed_slot.clone(), + ]; + let (master_only, owner_coordinated) = + partition_reclaim_coordination(candidates.clone(), |_| false); + + assert_eq!(master_only, vec![unindexed_allocation.clone()]); + assert_eq!(owner_coordinated.len(), 2); + assert!(owner_coordinated.iter().all(|item| !matches!( + &item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + ))); + + let (master_only, owner_coordinated) = partition_reclaim_coordination(candidates, |_| true); + assert!(master_only.is_empty()); + assert_eq!( + owner_coordinated, + vec![indexed_allocation, unindexed_allocation, committed_slot] + ); + } +} + +#[cfg(test)] +mod owner_get_holding_reclaim_tests { + use super::{OwnerReclaimBacking, OwnerReclaimReason, reclaim_items, route_item}; + use crate::client_kv_api::PutOptionalArgs; + use crate::config::KvSsdStorageConfig; + use crate::kv_ssd_storage::MIN_CAPACITY_BYTES; + use crate::kvcore_test_lib::{ + integration_test_lock, start_master_and_client, start_master_and_client_with_client_config, + stop_master_and_client, + }; + use crate::master_kv_router::msg_pack::{PutDoneReq, PutStartReq}; + use crate::master_kv_router::put::{handle_put_done, handle_put_start}; + use crate::memholder::{MemholderManagerTrait, NodeHolderKey}; + use crate::p2p::msg_pack::MsgPack; + use crate::rpcresp_kvresult_convert::msg_and_error::OK; + use std::time::{Duration, Instant}; + + #[limit_thirdparty::tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn completed_get_holding_does_not_block_two_sided_owner_reclaim() { + let _test_guard = integration_test_lock().await; + let (master, client) = + start_master_and_client("reclaim_get_holding_master", "reclaim_get_holding_owner") + .await; + let key = "completed_get_holding_reclaim_key"; + let owner_view = client.client_kv_api_view(); + let owner_api = owner_view.client_kv_api(); + owner_api + .inner() + .put(key, &[7u8; 4096], PutOptionalArgs::default()) + .await + .expect("owner put"); + let (holder, _get_info) = owner_api + .inner() + .get(key) + .await + .expect("owner get") + .expect("owner get should hit"); + + let owner_id = client + .cluster_manager_view() + .cluster_manager() + .get_self_info() + .id; + let holding_key = NodeHolderKey::new(owner_id.clone(), holder.holder_id()); + let master_view = master.master_kv_router_view().clone(); + assert!( + master_view + .master_kv_router() + .inner() + .get_holding + .inner() + .contains_key(&holding_key), + "get_done must install the Allocation lifetime holder" + ); + + assert!( + master_view + .master_kv_router() + .inner() + .key_activity + .is_quiescent(key), + "completed get must release its master key-activity lease" + ); + let owner_node: crate::cluster_manager::NodeID = owner_id.clone().into(); + let busy_item = route_item( + &master_view, + &owner_node, + key, + None, + None, + OwnerReclaimReason::OwnerCapacityEviction, + master_view.master_kv_router().next_owner_reclaim_epoch(), + ) + .expect("active-holder owner route should be reclaimable after the reader exits"); + assert_eq!( + reclaim_items(&master_view, &owner_node, vec![busy_item]).await, + 0, + "owner Prepare must reject reclaim while the user holder is live" + ); + + drop(holder); + limit_thirdparty::tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + master_view + .master_kv_router() + .inner() + .get_holding + .inner() + .contains_key(&holding_key), + "the committed local index intentionally keeps MemoryInfo and its ACK holder alive" + ); + + let item = route_item( + &master_view, + &owner_node, + key, + None, + None, + OwnerReclaimReason::OwnerCapacityEviction, + master_view.master_kv_router().next_owner_reclaim_epoch(), + ) + .expect("current owner route should be reclaimable"); + assert_eq!( + item.backing, + OwnerReclaimBacking::Allocation, + "reuse-replica get_done must publish the owner-local index on the route" + ); + assert_eq!( + reclaim_items(&master_view, &owner_node, vec![item]).await, + 1 + ); + + let wait_started = Instant::now(); + while master_view + .master_kv_router() + .inner() + .get_holding + .inner() + .contains_key(&holding_key) + { + assert!( + wait_started.elapsed() < Duration::from_secs(5), + "owner reclaim must drop MemoryInfo and deliver its delete ACK" + ); + limit_thirdparty::tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!( + !master_view + .master_kv_router() + .inner() + .kv_routes + .contains_key(key), + "the reclaimed last replica route must be removed" + ); + + stop_master_and_client(master, client).await; + } + + #[limit_thirdparty::tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn master_unindexed_allocation_reclaim_publishes_ssd_and_reloads_exact_bytes() { + let _test_guard = integration_test_lock().await; + let (master, client) = start_master_and_client_with_client_config( + "master_capacity_ssd_reclaim_master", + "master_capacity_ssd_reclaim_owner", + |config| { + config.ssd_storage = Some(KvSsdStorageConfig { + limit_bytes: MIN_CAPACITY_BYTES, + write_rate_limit_bytes_per_sec: None, + write_burst_bytes: None, + }); + let target = std::env::var("CARGO_TARGET_DIR") + .expect("SSD integration test requires the NVMe Cargo target"); + let root = format!( + "{target}/kv_ssd_integration/master_capacity_reclaim-{}", + std::process::id() + ); + config.share_mem_path = format!("{root}/sharemem"); + config.large_file_paths.paths = vec![format!("{root}/large")]; + }, + ) + .await; + let key = "master_capacity_ssd_reclaim_key"; + let payload = vec![0x5au8; 4096]; + let owner_view = client.client_kv_api_view(); + let owner_api = owner_view.client_kv_api(); + let owner_info = client + .cluster_manager_view() + .cluster_manager() + .get_self_info(); + let owner_id = owner_info.id; + let owner_node: crate::cluster_manager::NodeID = owner_id.clone().into(); + let master_view = master.master_kv_router_view().clone(); + assert!( + master_view + .cluster_manager() + .get_member_info_cached(&owner_id) + .and_then(|member| { + member + .metadata + .get(crate::cluster_manager::META_KEY_KV_SSD_STORAGE) + .cloned() + }) + .is_some_and(|value| value == "true"), + "master must observe the owner's SSD capability before selecting coordination" + ); + + let (_put_id, start) = handle_put_start( + master_view.clone(), + MsgPack { + serialize_part: PutStartReq { + key: key.to_string(), + len: payload.len() as u64, + reject_if_inflight_same_key: false, + reject_if_exist_same_key: false, + make_replica_task: true, + preferred_sub_cluster: None, + source_node_id: None, + }, + raw_bytes: Vec::new(), + }, + owner_node.clone(), + ) + .await; + assert_eq!( + start.serialize_part.error_code, OK, + "PutStart failed: {}", + start.serialize_part.error_json + ); + assert_eq!(start.serialize_part.node_id, owner_id); + owner_view + .client_seg_pool() + .copy_into_segment(start.serialize_part.target_addr, &payload) + .await + .expect("copy payload into the master-owned owner segment Allocation"); + let done = handle_put_done( + master_view.clone(), + MsgPack { + serialize_part: PutDoneReq { + key: key.to_string(), + put_id: start.serialize_part.put_id, + lease_id: None, + committed_slot: None, + publish_local_cache: false, + atomic_group: None, + }, + raw_bytes: Vec::new(), + }, + owner_node.clone(), + ) + .await; + assert_eq!(done.serialize_part.error_code, OK); + assert_eq!(done.serialize_part.local_cache_holder_id, None); + let route = master_view + .master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|route| route.clone()) + .expect("PutDone must publish the production unindexed Allocation route"); + assert!( + route + .node_replicas + .read() + .get(&owner_node) + .and_then(|replicas| replicas.memory.as_ref()) + .is_some_and(|memory| !memory.owner_local_indexed), + "production remote backing must have no owner-local key index" + ); + drop(route); + + let item = route_item( + &master_view, + &owner_node, + key, + None, + None, + OwnerReclaimReason::MasterAllocationCapacity, + master_view.master_kv_router().next_owner_reclaim_epoch(), + ) + .expect("master-capacity victim must resolve to the current owner Allocation"); + assert!(matches!( + item.backing, + OwnerReclaimBacking::UnindexedAllocation { .. } + )); + assert_eq!( + reclaim_items(&master_view, &owner_node, vec![item]).await, + 1 + ); + + let route = master_view + .master_kv_router() + .inner() + .kv_routes + .get(key) + .map(|route| route.clone()) + .expect("SSD backing must keep the route alive after DRAM reclaim"); + { + let replicas = route.node_replicas.read(); + let owner_backings = replicas + .get(&owner_node) + .expect("the SSD-only owner route must remain"); + assert!(owner_backings.memory.is_none()); + assert_eq!( + owner_backings.ssd.as_ref().map(|ssd| ssd.len), + Some(payload.len() as u64) + ); + } + let persisted = owner_api + .inner() + .kv_ssd_storage_usage_snapshot() + .expect("test owner SSD must be configured"); + assert_eq!(persisted.persist_successes, 1); + assert_eq!(persisted.persist_failures, 0); + assert_eq!(persisted.used_bytes, payload.len() as u64); + + let (holder, _get_info) = owner_api + .inner() + .get(key) + .await + .expect("SSD-backed get") + .expect("SSD-only route must be readable"); + assert_eq!(holder.bytes(), payload); + drop(holder); + let loaded = owner_api + .inner() + .kv_ssd_storage_usage_snapshot() + .expect("test owner SSD must remain configured"); + assert_eq!(loaded.load_successes, 1); + assert_eq!(loaded.load_failures, 0); + assert_eq!(loaded.load_bytes, payload.len() as u64); + assert_eq!(loaded.memory_hits + loaded.disk_hits + loaded.outer_hits, 1); + + stop_master_and_client(master, client).await; + } +} + +fn request_is_current(view: &MasterKvRouterView, request: &EvictionReclaimRequest) -> bool { + if let Some(expected_generation) = request.owner_node_start_time + && view + .cluster_manager() + .get_member_info_cached(&request.owner_node_id) + .map(|member| member.node_start_time) + != Some(expected_generation) + { + return false; + } + match request.origin { + EvictionReclaimOrigin::MasterAllocationCapacity => request.members.iter().all(|member| { + view.master_kv_router().eviction_cache_entry_is_current( + &request.owner_node_id, + &member.key, + &member.desc, + ) + }), + EvictionReclaimOrigin::OwnerCapacityEviction => { + let owner: NodeID = request.owner_node_id.clone().into(); + request.members.iter().all(|member| { + let Some(expected_backing) = member.expected_backing.as_ref() else { + return false; + }; + let Some(route) = view + .master_kv_router() + .inner() + .kv_routes + .get(&member.key) + .map(|entry| entry.clone()) + else { + return false; + }; + if route.put_id != member.desc.put_id || route.lease_id.is_some() { + return false; + } + route + .node_replicas + .read() + .get(&owner) + .is_some_and(|node_replicas| { + !node_replicas.tomb_tag.is_tomb() + && node_replicas.memory.as_ref().is_some_and(|replica| { + replica.owner_local_indexed + && reclaim_backing_matches(replica, expected_backing) + }) + }) + }) + } + } +} + +fn restore_request_entries(view: &MasterKvRouterView, request: &EvictionReclaimRequest) -> usize { + request + .members + .iter() + .filter(|member| { + view.master_kv_router() + .restore_eviction_cache_entry_if_current( + &request.owner_node_id, + member.key.clone(), + member.desc.clone(), + ) + }) + .count() +} + +fn spawn_eviction_reclaim_retry(view: MasterKvRouterView, requests: Vec) { + if requests.is_empty() { + return; + } + let mut delayed = Vec::with_capacity(requests.len()); + let mut restored_count = 0usize; + let mut restored_weight = 0u64; + for mut request in requests { + request.retry_count = request.retry_count.saturating_add(1); + let counters = view + .master_kv_router() + .eviction_reclaim_counters(&request.owner_node_id); + let weight = request.weight_bytes(); + if !request_is_current(&view, &request) { + view.master_kv_router().complete_eviction_reclaim(&request); + counters.route_changed.fetch_add(1, Ordering::Relaxed); + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + continue; + } + if request.origin == EvictionReclaimOrigin::MasterAllocationCapacity + && should_restore_after_retry(request.retry_count) + { + // Release the old identity before reinsertion. A bounded remote + // cache may synchronously produce a fresh Size event; that event + // must own a new, independently-accounted lifecycle. + view.master_kv_router().complete_eviction_reclaim(&request); + let restored = restore_request_entries(&view, &request); + if restored == request.members.len() { + counters.retry_restored.fetch_add(1, Ordering::Relaxed); + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + restored_count += restored; + restored_weight = restored_weight.saturating_add(weight); + continue; + } + counters.route_changed.fetch_add(1, Ordering::Relaxed); + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + continue; + } + counters.retry_queued.fetch_add(1, Ordering::Relaxed); + delayed.push(request); + } + if restored_count != 0 { + tracing::info!( + "safe eviction reclaim restored current cache entries after bounded retry: entries={} weight_bytes={} max_retry_count={}", + restored_count, + restored_weight, + EVICTION_RECLAIM_MAX_RETRY_COUNT + ); + } + if delayed.is_empty() { + return; + } + let max_retry_count = delayed + .iter() + .map(|request| request.retry_count) + .max() + .unwrap_or(1); + let retry_delay = eviction_reclaim_retry_delay(max_retry_count); + let spawn_view = view.clone(); + let _ = spawn_view.spawn("eviction_reclaim_retry", async move { + tokio::time::sleep(retry_delay).await; + let tx = view.master_kv_router().inner().eviction_reclaim_tx.clone(); + for request in delayed { + let counters = view + .master_kv_router() + .eviction_reclaim_counters(&request.owner_node_id); + if !request_is_current(&view, &request) { + view.master_kv_router().complete_eviction_reclaim(&request); + counters.route_changed.fetch_add(1, Ordering::Relaxed); + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + continue; + } + if let Err(err) = tx.send(request) { + let request = err.0; + view.master_kv_router().complete_eviction_reclaim(&request); + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + tracing::warn!( + "lossless eviction reclaim retry channel closed: owner={} members={}", + request.owner_node_id, + request.members.len(), + ); + } + } + }); +} + +pub(crate) fn spawn_eviction_reclaim_actor( + view: MasterKvRouterView, + mut rx: limit_thirdparty::tokio::sync::ampsc::UnboundedReceiver, +) { + let view_task = view.clone(); + let _ = view.spawn("eviction_reclaim_actor", async move { + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + let first = tokio::select! { + _ = shutdown_waiter.wait() => break, + request = rx.recv() => { + let Some(request) = request else { break; }; + request + } + }; + let mut batch = Vec::with_capacity(OWNER_RECLAIM_MAX_BATCH); + batch.push(first); + let mut merge_window = Box::pin(tokio::time::sleep(OWNER_RECLAIM_MERGE_WINDOW)); + while batch.len() < OWNER_RECLAIM_MAX_BATCH { + tokio::select! { + _ = &mut merge_window => break, + request = rx.recv() => { + let Some(request) = request else { break; }; + batch.push(request); + } + } + } + + let mut groups: HashMap> = HashMap::new(); + for request in batch { + groups + .entry(request.owner_node_id.clone()) + .or_default() + .push(request); + } + for (owner_node_id, requests) in groups { + let owner: NodeID = owner_node_id.clone().into(); + let counters = view_task + .master_kv_router() + .eviction_reclaim_counters(&owner_node_id); + let mut pending = std::collections::VecDeque::from(requests); + let mut retry_requests = Vec::new(); + while !pending.is_empty() { + let mut accounting_requests = Vec::with_capacity(OWNER_RECLAIM_RPC_BATCH); + let mut items = Vec::with_capacity(OWNER_RECLAIM_RPC_BATCH); + for _ in 0..OWNER_RECLAIM_RPC_BATCH { + let Some(request) = pending.pop_front() else { + break; + }; + let (members, reason) = match request.origin { + EvictionReclaimOrigin::OwnerCapacityEviction => ( + request.members.clone(), + OwnerReclaimReason::OwnerCapacityEviction, + ), + EvictionReclaimOrigin::MasterAllocationCapacity => { + let member = match plan_master_allocation_capacity_victim( + &view_task, + &request, + ) { + Ok(member) => member, + Err(MasterCapacityPlanError::CommittedSlot) => { + counters + .capacity_eviction_non_ring_b_entry_total + .fetch_add(1, Ordering::Relaxed); + counters + .capacity_eviction_hit_committed_slot + .fetch_add(1, Ordering::Relaxed); + view_task + .master_kv_router() + .complete_eviction_reclaim(&request); + let restored = + restore_request_entries(&view_task, &request); + tracing::error!( + "BUG: master Allocation capacity event resolved to CommittedSlot; restored metadata: owner={} members={} restored={}", + owner_node_id, + request.members.len(), + restored, + ); + continue; + } + Err(MasterCapacityPlanError::RouteChanged) => { + view_task + .master_kv_router() + .complete_eviction_reclaim(&request); + counters.route_changed.fetch_add(1, Ordering::Relaxed); + continue; + } + Err(MasterCapacityPlanError::WrongRole) => { + counters + .capacity_eviction_non_ring_b_entry_total + .fetch_add(1, Ordering::Relaxed); + view_task + .master_kv_router() + .complete_eviction_reclaim(&request); + let restored = + restore_request_entries(&view_task, &request); + tracing::error!( + "BUG: master Allocation Size event resolved to a non-ring-B route; restored metadata: owner={} members={} restored={}", + owner_node_id, + request.members.len(), + restored, + ); + continue; + } + }; + (vec![member], OwnerReclaimReason::MasterAllocationCapacity) + } + }; + + let planned = members + .iter() + .map(|member| { + let item = route_item( + &view_task, + &owner, + &member.key, + Some(member.desc.put_id), + None, + reason, + view_task.master_kv_router().next_owner_reclaim_epoch(), + )?; + if member + .expected_backing + .as_ref() + .is_some_and(|expected| expected != &item.backing) + { + return None; + } + Some(item) + }) + .collect::>>(); + if let Some(mut planned) = planned { + items.append(&mut planned); + } + accounting_requests.push(request); + } + + if !items.is_empty() { + let _ = reclaim_items(&view_task, &owner, items).await; + } + for accounting_request in accounting_requests { + if request_is_current(&view_task, &accounting_request) { + retry_requests.push(accounting_request); + } else { + view_task + .master_kv_router() + .complete_eviction_reclaim(&accounting_request); + counters.route_changed.fetch_add(1, Ordering::Relaxed); + if accounting_request.retry_count != 0 { + counters.retry_completed.fetch_add(1, Ordering::Relaxed); + } + } + } + } + let retry_count = retry_requests.len(); + spawn_eviction_reclaim_retry(view_task.clone(), retry_requests); + tracing::trace!( + "batched single-key eviction reclaim completed: owner={} retry_deferred={}", + owner_node_id, + retry_count, + ); + } + } + }); +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/route_maintenance.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/route_maintenance.rs new file mode 100644 index 0000000..04df8d5 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/route_maintenance.rs @@ -0,0 +1,331 @@ +use super::{MasterKvRouterView, NodeValueReplicaDesc, put::PutIDForAKey}; +use crate::cluster_manager::NodeID; +use limit_thirdparty::tokio; +use std::collections::HashMap; +use std::time::Duration; + +const POST_ROUTE_MAINTENANCE_MAX_BATCH: usize = 512; +const POST_ROUTE_MAINTENANCE_MERGE_WINDOW: Duration = Duration::from_millis(2); + +#[derive(Clone, Copy)] +enum RoutePublishKind { + PrimaryPut, + ReplicaAppend, +} + +pub(super) struct RoutePublishEvent { + kind: RoutePublishKind, + key: String, + put_id: PutIDForAKey, + lease_id: Option, + node_id: NodeID, + capacity_bytes: u64, +} + +impl RoutePublishEvent { + pub(super) fn primary_put( + key: String, + put_id: PutIDForAKey, + lease_id: Option, + node_id: NodeID, + capacity_bytes: u64, + ) -> Self { + Self { + kind: RoutePublishKind::PrimaryPut, + key, + put_id, + lease_id, + node_id, + capacity_bytes, + } + } + + pub(super) fn replica_append( + key: String, + put_id: PutIDForAKey, + lease_id: Option, + node_id: NodeID, + capacity_bytes: u64, + ) -> Self { + Self { + kind: RoutePublishKind::ReplicaAppend, + key, + put_id, + lease_id, + node_id, + capacity_bytes, + } + } +} + +fn saturating_moka_weight_bytes(key: &str, put_id: PutIDForAKey, capacity_bytes: u64) -> u32 { + if capacity_bytes > u32::MAX as u64 { + tracing::warn!( + "moka weight saturation after route publish: key={} put_id=({},{}) cap={}B exceeds u32::MAX; weight set to u32::MAX", + key, + put_id.0, + put_id.1, + capacity_bytes + ); + u32::MAX + } else { + capacity_bytes as u32 + } +} + +fn deduplicate_owner_events( + events: Vec<(usize, String, NodeValueReplicaDesc)>, +) -> Vec<(String, NodeValueReplicaDesc)> { + let mut latest_by_key = HashMap::with_capacity(events.len()); + for (sequence, key, desc) in events { + latest_by_key.insert(key, (sequence, desc)); + } + let mut latest = latest_by_key.into_iter().collect::>(); + latest.sort_unstable_by_key(|(_, (sequence, _))| *sequence); + latest + .into_iter() + .map(|(key, (_, desc))| (key, desc)) + .collect() +} + +/// Applies index and cache work after route guards have been released. +pub(super) async fn apply_post_route_maintenance_batch( + view: &MasterKvRouterView, + events: Vec, +) { + if view.master_kv_router().prefix_index_enabled() + && events + .iter() + .any(|event| matches!(event.kind, RoutePublishKind::PrimaryPut)) + { + let inner = view.master_kv_router().inner(); + let mut tree = inner.prefix_index.write().await; + for event in &events { + if matches!(event.kind, RoutePublishKind::PrimaryPut) { + let event_is_current = inner.kv_routes.get(&event.key).is_some_and(|route| { + route.put_id == event.put_id + && route + .node_replicas + .read() + .values() + .any(|replica| !replica.tomb_tag.is_tomb()) + }); + if event_is_current { + tree.insert(&event.key, event.put_id); + } + } + } + } + + if !view.master_kv_router().replica_cache_enabled() { + return; + } + let mut ring_b_events_by_owner = + HashMap::>::new(); + let mut tier1_events_by_owner = + HashMap::>::new(); + for (sequence, event) in events.into_iter().enumerate() { + if event.lease_id.is_some() { + continue; + } + let weight_bytes = + saturating_moka_weight_bytes(&event.key, event.put_id, event.capacity_bytes); + let desc = NodeValueReplicaDesc { + weight_bytes, + put_id: event.put_id, + }; + if view.master_kv_router().eviction_cache_entry_is_current( + event.node_id.as_ref(), + &event.key, + &desc, + ) { + ring_b_events_by_owner + .entry(event.node_id.as_ref().to_string()) + .or_default() + .push((sequence, event.key.clone(), desc.clone())); + } + if view.master_kv_router().tier1_writeback_entry_is_current( + event.node_id.as_ref(), + &event.key, + &desc, + ) { + tier1_events_by_owner + .entry(event.node_id.as_ref().to_string()) + .or_default() + .push((sequence, event.key, desc)); + } + } + + for (owner_node_id, owner_events) in ring_b_events_by_owner { + let entries = deduplicate_owner_events(owner_events); + // Moka's sync housekeeper uses a blocking mutex. Serialize before + // entering it with an async owner-level gate so waiting route-publish + // tasks yield instead of occupying every Tokio worker thread. + let owner_cache_lock = view + .master_kv_router() + .inner() + .owner_cache_operation_locks + .get_lock(owner_node_id.clone()); + let _owner_cache_guard = owner_cache_lock.lock().await; + let Some(cache) = view + .master_kv_router() + .get_node_cache_controller(&owner_node_id) + else { + tracing::warn!( + "No cache controller found for node: {}, node is not ready", + owner_node_id + ); + continue; + }; + + // Drain prior writes once, then admit this owner batch as one unit. + // Every node's controller is the bounded authority for that node's + // unindexed Allocation domain; placement role is irrelevant. + cache.run_pending_tasks(); + for (key, desc) in entries { + if !view + .master_kv_router() + .eviction_cache_entry_is_current(&owner_node_id, &key, &desc) + { + continue; + } + tracing::debug!("Inserting key: {:?} into cache", key); + super::insert_master_cache_entry(cache.as_ref(), key.clone(), desc.clone()); + tracing::debug!( + "Inserted key: {:?} into cache, current cache size: {}", + key, + cache.weighted_size() + ); + } + // Do not search Moka for a recoverable victim here. CPU append Done + // already carries the exact source key/atomic_batch and performs a validated + // point demotion. Local-reserve Free/Prepared/Pending/Committed state is + // the physical capacity authority while that writeback is in flight. + } + + // Tier1 is a separate pre-writeback policy. Its admission rules remain + // owner-route based and must not be coupled to ring-B backing admission. + for (owner_node_id, owner_events) in tier1_events_by_owner { + let entries = deduplicate_owner_events(owner_events); + let owner_cache_lock = view + .master_kv_router() + .inner() + .owner_cache_operation_locks + .get_lock(owner_node_id.clone()); + let _owner_cache_guard = owner_cache_lock.lock().await; + let Some(tier1_cache) = view + .master_kv_router() + .get_node_writeback_tier1_controller(&owner_node_id) + else { + continue; + }; + tier1_cache.run_pending_tasks(); + for (key, desc) in entries { + if !view.master_kv_router().tier1_writeback_entry_is_current( + &owner_node_id, + &key, + &desc, + ) { + continue; + } + super::insert_master_cache_entry(tier1_cache.as_ref(), key, desc); + } + } +} + +pub(super) fn spawn_post_route_maintenance_actor( + view: MasterKvRouterView, + mut rx: tokio::sync::ampsc::Receiver, +) { + let view_task = view.clone(); + view.spawn("post_route_maintenance_actor", async move { + tracing::info!( + "post-route maintenance actor started: max_batch={} merge_window_ms={}", + POST_ROUTE_MAINTENANCE_MAX_BATCH, + POST_ROUTE_MAINTENANCE_MERGE_WINDOW.as_millis(), + ); + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + let first = tokio::select! { + _ = shutdown_waiter.wait() => break, + event = rx.recv() => { + let Some(event) = event else { break; }; + event + } + }; + let mut events = Vec::with_capacity(POST_ROUTE_MAINTENANCE_MAX_BATCH); + events.push(first); + let mut merge_window = + Box::pin(tokio::time::sleep(POST_ROUTE_MAINTENANCE_MERGE_WINDOW)); + while events.len() < POST_ROUTE_MAINTENANCE_MAX_BATCH { + tokio::select! { + _ = &mut merge_window => break, + event = rx.recv() => { + let Some(event) = event else { break; }; + events.push(event); + } + } + } + apply_post_route_maintenance_batch(&view_task, events).await; + } + tracing::info!("post-route maintenance actor stopped"); + }); +} + +/// Queues index and cache work with bounded backpressure after route publication. +pub(super) async fn enqueue_post_route_maintenance( + view: &MasterKvRouterView, + event: RoutePublishEvent, +) { + let update_prefix_index = matches!(event.kind, RoutePublishKind::PrimaryPut) + && view.master_kv_router().prefix_index_enabled(); + let insert_replica_cache = + event.lease_id.is_none() && view.master_kv_router().replica_cache_enabled(); + if !update_prefix_index && !insert_replica_cache { + return; + } + view.master_kv_router() + .inner() + .post_route_maintenance_tx + .send(event) + .await + .expect("post-route maintenance actor stopped while master is serving requests"); +} + +#[cfg(test)] +mod tests { + use super::{deduplicate_owner_events, saturating_moka_weight_bytes}; + use crate::master_kv_router::NodeValueReplicaDesc; + + #[test] + fn moka_weight_saturates_without_truncating() { + assert_eq!( + saturating_moka_weight_bytes("key", (1, 2), u32::MAX as u64), + u32::MAX + ); + assert_eq!( + saturating_moka_weight_bytes("key", (1, 2), u32::MAX as u64 + 1), + u32::MAX + ); + } + + #[test] + fn owner_batch_keeps_only_the_last_event_for_each_key() { + let desc = |weight_bytes, put_id| NodeValueReplicaDesc { + weight_bytes, + put_id, + }; + let events = vec![ + (0, "a".to_string(), desc(10, (1, 0))), + (1, "b".to_string(), desc(20, (1, 0))), + (2, "a".to_string(), desc(30, (2, 0))), + ]; + + let deduplicated = deduplicate_owner_events(events); + assert_eq!(deduplicated.len(), 2); + assert_eq!(deduplicated[0].0, "b"); + assert_eq!(deduplicated[1].0, "a"); + assert_eq!(deduplicated[1].1.weight_bytes, 30); + assert_eq!(deduplicated[1].1.put_id, (2, 0)); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_kv_router/tiered_writeback.rs b/fluxon_rs/fluxon_kv/src/master_kv_router/tiered_writeback.rs new file mode 100644 index 0000000..97ebfd1 --- /dev/null +++ b/fluxon_rs/fluxon_kv/src/master_kv_router/tiered_writeback.rs @@ -0,0 +1,224 @@ +use super::{MasterKvRouterView, NodeValueReplicaDesc}; +use crate::cluster_manager::{NodeID, NodeIDString}; +use crate::master_kv_router::msg_pack::{BatchEnqueueReplicaTaskReq, EnqueueReplicaTaskItem}; +use crate::p2p::msg_pack::{MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, RPCCaller}; +use crate::rpcresp_kvresult_convert::msg_and_error::OK; +use limit_thirdparty::tokio; +use std::collections::HashMap; +use std::time::Duration; + +const TIER1_WRITEBACK_RPC_TIMEOUT: Duration = Duration::from_secs(MIN_EXPLICIT_RPC_TIMEOUT_SECS); +const TIER1_WRITEBACK_MAX_BATCH: usize = 256; +const TIER1_WRITEBACK_MERGE_WINDOW: Duration = Duration::from_millis(2); + +#[derive(Clone, Debug)] +pub(crate) struct Tier1WritebackRequest { + pub source_node_id: NodeIDString, + pub key: String, + pub desc: NodeValueReplicaDesc, +} + +async fn dispatch_owner_batch( + view: &MasterKvRouterView, + source_node_id: NodeIDString, + requests: Vec, +) { + let mut current = Vec::with_capacity(requests.len()); + for request in requests { + if view.master_kv_router().tier1_writeback_entry_is_current( + &request.source_node_id, + &request.key, + &request.desc, + ) { + current.push(request); + } else { + view.master_kv_router() + .finish_tier1_writeback_request(request); + } + } + if current.is_empty() { + return; + } + + let items = current + .iter() + .map(|request| EnqueueReplicaTaskItem { + key: request.key.clone(), + put_id: request.desc.put_id, + }) + .collect::>(); + let caller = RPCCaller::::new(); + caller.regist(view.p2p_module()); + let result = caller + .call( + view.p2p_module(), + NodeID::from(source_node_id.clone()), + MsgPack { + serialize_part: BatchEnqueueReplicaTaskReq { items }, + raw_bytes: Vec::new(), + }, + Some(TIER1_WRITEBACK_RPC_TIMEOUT), + 1, + ) + .await; + + let response = match result { + Ok(response) + if response.serialize_part.error_code == OK + && response.serialize_part.items.len() == current.len() => + { + response.serialize_part.items + } + Ok(response) => { + tracing::warn!( + "tier1 write-back owner response rejected: owner={} requested={} returned={} code={} error={}", + source_node_id, + current.len(), + response.serialize_part.items.len(), + response.serialize_part.error_code, + response.serialize_part.error_json + ); + view.master_kv_router().record_tier1_writeback_failed( + &source_node_id, + u64::try_from(current.len()).unwrap_or(u64::MAX), + ); + for request in current { + view.master_kv_router() + .finish_tier1_writeback_request(request); + } + return; + } + Err(err) => { + tracing::warn!( + "tier1 write-back owner RPC failed: owner={} requested={} err={:?}", + source_node_id, + current.len(), + err + ); + view.master_kv_router().record_tier1_writeback_failed( + &source_node_id, + u64::try_from(current.len()).unwrap_or(u64::MAX), + ); + for request in current { + view.master_kv_router() + .finish_tier1_writeback_request(request); + } + return; + } + }; + + let requested = current.len(); + let mut accepted = 0usize; + for (request, item) in current.into_iter().zip(response.into_iter()) { + let identity_matches = request.key == item.key && request.desc.put_id == item.put_id; + if identity_matches && item.accepted { + accepted += 1; + continue; + } + view.master_kv_router() + .record_tier1_writeback_failed(&source_node_id, 1); + if !identity_matches { + tracing::warn!( + "tier1 write-back owner response identity mismatch: owner={} request_key={} request_put_id=({},{}) response_key={} response_put_id=({},{})", + source_node_id, + request.key, + request.desc.put_id.0, + request.desc.put_id.1, + item.key, + item.put_id.0, + item.put_id.1 + ); + } + view.master_kv_router() + .finish_tier1_writeback_request(request); + } + view.master_kv_router() + .record_tier1_writeback_owner_accepted( + &source_node_id, + u64::try_from(accepted).unwrap_or(u64::MAX), + ); + tracing::debug!( + "tier1 write-back owner batch dispatched: owner={} requested={} accepted={}", + source_node_id, + requested, + accepted + ); +} + +pub(crate) fn spawn_tier1_writeback_actor( + view: MasterKvRouterView, + mut rx: tokio::sync::ampsc::Receiver, +) { + let view_task = view.clone(); + let _ = view.spawn("tier1_writeback_actor", async move { + let mut shutdown_waiter = view_task.register_shutdown_waiter(); + loop { + let first = tokio::select! { + _ = shutdown_waiter.wait() => break, + request = rx.recv() => { + let Some(request) = request else { break; }; + request + } + }; + let mut batch = Vec::with_capacity(TIER1_WRITEBACK_MAX_BATCH); + batch.push(first); + let mut merge_window = Box::pin(tokio::time::sleep(TIER1_WRITEBACK_MERGE_WINDOW)); + while batch.len() < TIER1_WRITEBACK_MAX_BATCH { + tokio::select! { + _ = &mut merge_window => break, + request = rx.recv() => { + let Some(request) = request else { break; }; + batch.push(request); + } + } + } + + let mut groups: HashMap> = HashMap::new(); + for request in batch { + groups + .entry(request.source_node_id.clone()) + .or_default() + .push(request); + } + for (source_node_id, requests) in groups { + dispatch_owner_batch(&view_task, source_node_id, requests).await; + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::p2p::msg_pack::validate_explicit_rpc_timeout; + + #[test] + fn tier1_writeback_timeout_satisfies_rpc_contract() { + validate_explicit_rpc_timeout(Some(TIER1_WRITEBACK_RPC_TIMEOUT)).unwrap(); + } + + #[test] + fn inclusive_hot_tier_does_not_reduce_resident_capacity() { + let build = |capacity| { + moka::sync::SegmentedCache::builder(1) + .max_capacity(capacity) + .weigher(Box::new(|_key: &String, weight: &u32| *weight)) + .build() + }; + let resident = build(100); + let tier1 = build(60); + + for key in ["a", "b", "c"] { + resident.insert(key.to_string(), 30); + tier1.insert(key.to_string(), 30); + } + resident.run_pending_tasks(); + tier1.run_pending_tasks(); + + assert_eq!(resident.policy().max_capacity(), Some(100)); + assert_eq!(resident.weighted_size(), 90); + assert_eq!(tier1.policy().max_capacity(), Some(60)); + assert!(tier1.weighted_size() <= 60); + assert!(resident.entry_count() > tier1.entry_count()); + } +} diff --git a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease.rs b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease.rs index 1df23e3..eeb24e6 100755 --- a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease.rs +++ b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease.rs @@ -146,6 +146,15 @@ impl Lease { Ok(()) } + /// Remove only the exact version installed by a failed route publication. + /// A newer Put may already have rebound the same key, so key-only removal + /// would be an ABA bug. + pub fn detach_key_if_version(&self, key: &str, put_id: PutIDForAKey) -> bool { + self.keys + .remove_if(key, |_, current_put_id| *current_put_id == put_id) + .is_some() + } + // Note: no detach_all_keys semantics; deletion paths iterate keys directly. /// Get all keys with their put_id information diff --git a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs index 5c20cc1..ed7f6be 100755 --- a/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs +++ b/fluxon_rs/fluxon_kv/src/master_lease_manager/lease_manager_test.rs @@ -22,7 +22,8 @@ async fn test1_lease_expire_removes_keys() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t1", "lease_client_t1").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t1", "lease_client_t1").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -82,7 +83,8 @@ async fn test2_rebind_to_new_lease_preserves_until_new_expire() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t2", "lease_client_t2").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t2", "lease_client_t2").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -161,7 +163,8 @@ async fn test3_keepalive() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t3", "lease_client_t3").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t3", "lease_client_t3").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -236,7 +239,8 @@ async fn test4_delete_under_lease_then_get_fails() { unsafe { std::env::set_var("FLUXON_LOG", "debug"); } - let (master_fw, client_fw) = start_master_and_client("lease_master_t4", "lease_client_t4").await; + let (master_fw, client_fw) = + start_master_and_client("lease_master_t4", "lease_client_t4").await; let client_view = client_fw.client_kv_api_view(); wait_master_ready(&client_view).await; @@ -346,6 +350,36 @@ async fn test5_eviction_when_lease_consumes_space() { leased_success_keys.len() ); + // Rounded allocations and the three reclaimable normal keys mean that a + // fixed 19 * 5MiB loop does not itself prove the physical segment is full. + // Establish the precondition explicitly before asserting later NoSpace + // behavior; otherwise the first post-check lease put can legitimately use + // the final free 5MiB slot. + let mut reached_lease_full = false; + for i in 0..32u32 { + let k = format!("t5_lease_topoff_{}", i); + match client_view + .client_kv_api() + .put(&k, &leased_value, Some(lease_id)) + .await + { + Ok(()) => leased_success_keys.push(k), + Err(e) => { + assert!( + e.to_string().contains("NoSpace"), + "expected top-off to end in NoSpace, got: {}", + e + ); + reached_lease_full = true; + break; + } + } + } + assert!( + reached_lease_full, + "lease top-off did not establish a full physical segment" + ); + // Canvas assertion: assert 放成功的所有 lease key 都存在（TTL 尚未到期） tracing::info!( "[test5-assert-all-leased-keys-exist] verifying all successfully stored leased keys still exist before TTL expiry" diff --git a/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs b/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs index 86d6736..5eb2120 100644 --- a/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs +++ b/fluxon_rs/fluxon_kv/src/master_lease_manager/master_lease_manager.rs @@ -503,6 +503,21 @@ impl MasterLeaseManager { } } + /// Roll back one exact lease binding when its route publication did not + /// linearize. This is synchronous because it only touches the lease's + /// DashMap and is used from completion rollback paths. + pub fn detach_key_if_version( + &self, + lease_id: LeaseID, + key: &str, + put_id: PutIDForAKey, + ) -> bool { + self.inner() + .leases + .get(&lease_id) + .is_some_and(|lease| lease.detach_key_if_version(key, put_id)) + } + /// Get lease information /// pub fn get_lease(&self, lease_id: LeaseID) -> Option> { /// let leases_guard = self.leases.read(); diff --git a/fluxon_rs/fluxon_kv/src/master_seg_manager/mod.rs b/fluxon_rs/fluxon_kv/src/master_seg_manager/mod.rs index 8a50032..45d7059 100644 --- a/fluxon_rs/fluxon_kv/src/master_seg_manager/mod.rs +++ b/fluxon_rs/fluxon_kv/src/master_seg_manager/mod.rs @@ -2,7 +2,9 @@ pub mod msg_pack; pub mod one_seg_allocator; use self::msg_pack::RequestSegmentRegistrationReq; use self::msg_pack::SegmentDeviceDescription; -use self::one_seg_allocator::OneSegAllocator; +use self::one_seg_allocator::{ + Allocation, NodePoolCapacityBudget, NodePoolCapacitySnapshot, OneSegAllocator, +}; use crate::cluster_manager::NodeID; use crate::p2p::p2p_module::P2pModuleAccessTrait; use crate::rpcresp_kvresult_convert::msg_and_error::OK; @@ -22,11 +24,52 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; +fn build_node_segments_manager( + node_start_time: i64, + seg_map: std::collections::HashMap< + SegmentDeviceID, + (SegmentDeviceDescription, msg_pack::SegmentDeviceMemInfo), + >, +) -> KvResult { + let total_size = seg_map.values().try_fold(0u64, |total, (_, info)| { + total.checked_add(info.len).ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: "registered node segment capacity overflows u64".to_string(), + }, + ) + }) + })?; + let capacity_budget = Arc::new(NodePoolCapacityBudget::new(total_size)?); + let mut device_id_2_allocator: HashMap> = HashMap::new(); + for (device_id, (seg_device_desc, seg_mem_info)) in seg_map { + let allocator = OneSegAllocator::new_with_capacity_budget( + device_id.clone(), + seg_device_desc, + seg_mem_info.addr, + seg_mem_info.len, + capacity_budget.clone(), + ) + .map_err(|e| { + tracing::error!("Failed to create OneSegAllocator: {}", e); + e + })?; + device_id_2_allocator.insert(device_id, Arc::new(allocator)); + } + Ok(NodeSegmentsManager::new( + node_start_time, + total_size, + device_id_2_allocator, + capacity_budget, + )) +} + // --- Handler Functions --- /// https://qcnoe3hd7k5c.feishu.cn/wiki/KkeXwBbP4iCRN8kWSDccP5GBnrd#share-AuMbdrSaXoadUbxRmUncooKnnQd fn register_node_segments( view: &MasterSegManagerView, node_id: NodeID, + node_start_time: i64, seg_map: std::collections::HashMap< SegmentDeviceID, (SegmentDeviceDescription, msg_pack::SegmentDeviceMemInfo), @@ -52,60 +95,43 @@ fn register_node_segments( match alloc_map.entry(node_id.clone()) { dashmap::mapref::entry::Entry::Vacant(v) => { - let mut total_size: u64 = 0; - let mut device_id_2_allocator: HashMap> = - HashMap::new(); - - for (device_id, (seg_device_desc, seg_mem_info)) in seg_map { - let allocator = OneSegAllocator::new( - device_id.clone(), - seg_device_desc, - seg_mem_info.addr, - seg_mem_info.len, - ) - .map_err(|e| { - tracing::error!("Failed to create OneSegAllocator: {}", e); - e - })?; - - total_size = total_size.saturating_add(seg_mem_info.len); - device_id_2_allocator.insert(device_id, Arc::new(allocator)); - } - - v.insert(NodeSegmentsManager::new(total_size, device_id_2_allocator)); + v.insert(build_node_segments_manager(node_start_time, seg_map)?); } dashmap::mapref::entry::Entry::Occupied(mut occ) => { let node_segments_manager = occ.get_mut(); // Tomb means the previous instance has left/restarted; replace the full segment set. if node_segments_manager.tomb_tag.is_tomb() { - let mut total_size: u64 = 0; - let mut device_id_2_allocator: HashMap> = - HashMap::new(); - - for (device_id, (seg_device_desc, seg_mem_info)) in seg_map { - let allocator = OneSegAllocator::new( - device_id.clone(), - seg_device_desc, - seg_mem_info.addr, - seg_mem_info.len, - ) - .map_err(|e| { - tracing::error!("Failed to create OneSegAllocator: {}", e); - e - })?; - - total_size = total_size.saturating_add(seg_mem_info.len); - device_id_2_allocator.insert(device_id, Arc::new(allocator)); + // An RPC response already in flight at MemberLeft must not resurrect the + // departed generation. Only a genuinely newer epoch may replace a tomb. + if node_segments_manager.node_start_time == node_start_time { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::RegisterSegmentFailed { + detail: format!( + "stale segment registration for tombed node generation: node={} node_start_time={}", + node_id, node_start_time + ), + }, + )); } - - *node_segments_manager = - NodeSegmentsManager::new(total_size, device_id_2_allocator); + *node_segments_manager = build_node_segments_manager(node_start_time, seg_map)?; tracing::info!("RegisterSegment replaced tombed node: {}", node_id); return Ok(()); } - // Non-tomb: allow re-entrant registration (idempotent) to tolerate transient retries. + if node_segments_manager.node_start_time != node_start_time { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::RegisterSegmentFailed { + detail: format!( + "new segment generation attempted to replace a live generation: node={} live_node_start_time={} requested_node_start_time={}", + node_id, node_segments_manager.node_start_time, node_start_time + ), + }, + )); + } + + // Non-tomb, same generation: allow re-entrant registration (idempotent) to + // tolerate transient retries. for (device_id, (seg_device_desc, seg_mem_info)) in seg_map { if let Some(existing) = node_segments_manager.device_id_2_allocator.get(&device_id) { @@ -125,23 +151,38 @@ fn register_node_segments( )); } - let allocator = OneSegAllocator::new( + let new_total_size = node_segments_manager + .total_size + .checked_add(seg_mem_info.len) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::InvalidArgument { + detail: format!( + "registered node segment capacity overflows u64: node={}", + node_id + ), + }, + ) + })?; + let allocator = OneSegAllocator::new_with_capacity_budget( device_id.clone(), seg_device_desc, seg_mem_info.addr, seg_mem_info.len, + node_segments_manager.capacity_budget.clone(), ) .map_err(|e| { tracing::error!("Failed to create OneSegAllocator: {}", e); e })?; + node_segments_manager + .capacity_budget + .extend_physical_capacity(seg_mem_info.len)?; node_segments_manager .device_id_2_allocator .insert(device_id, Arc::new(allocator)); - node_segments_manager.total_size = node_segments_manager - .total_size - .saturating_add(seg_mem_info.len); + node_segments_manager.total_size = new_total_size; } } } @@ -181,22 +222,33 @@ impl NodeTombTag { pub fn set_tomb(&self) { self.0.store(true, Ordering::Release); } + + /// True only when both tags belong to the same node registration generation. + pub fn same_generation(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } } pub struct NodeSegmentsManager { + node_start_time: i64, total_size: u64, device_id_2_allocator: HashMap>, + capacity_budget: Arc, tomb_tag: NodeTombTag, } impl NodeSegmentsManager { - pub fn new( + fn new( + node_start_time: i64, total_size: u64, device_id_2_allocator: HashMap>, + capacity_budget: Arc, ) -> Self { Self { + node_start_time, total_size, device_id_2_allocator, + capacity_budget, tomb_tag: NodeTombTag::new(), } } @@ -298,11 +350,30 @@ impl MasterSegManager { // allocator.clone().allocate(size) // } - pub fn mark_node_tomb(&self, node_id: &NodeID) { + pub fn mark_node_tomb(&self, node_id: &NodeID) -> Option { + self.mark_node_tomb_generation(node_id, None) + } + + /// Mark one exact membership generation as departed and return its shared tomb tag. + /// A delayed leave for an older epoch must not tomb a newly registered generation that + /// happens to reuse the same node id. + pub fn mark_node_tomb_generation( + &self, + node_id: &NodeID, + expected_node_start_time: Option, + ) -> Option { if let Some(allocators_and_tomb_tag) = self.inner().node_allocators_and_tomb_tag.get(node_id) { + if expected_node_start_time + .is_some_and(|expected| allocators_and_tomb_tag.node_start_time != expected) + { + return None; + } allocators_and_tomb_tag.tomb_tag.set_tomb(); + Some(allocators_and_tomb_tag.tomb_tag.clone()) + } else { + None } } @@ -316,6 +387,28 @@ impl MasterSegManager { } } + /// Resolve the registration generation that owns an already-created allocation. + /// + /// A plain `get_node_tomb_tag(node_id)` is insufficient: the node may have left and + /// re-registered between allocation and completion. In that case the current tag belongs + /// to a different allocator set and must never be attached to the old allocation. + pub fn get_allocation_tomb_tag( + &self, + node_id: &NodeID, + allocation: &Allocation, + ) -> Option { + let node = self.inner().node_allocators_and_tomb_tag.get(node_id)?; + if node.tomb_tag.is_tomb() + || !node + .device_id_2_allocator + .values() + .any(|allocator| allocation.belongs_to_allocator(allocator)) + { + return None; + } + Some(node.tomb_tag.clone()) + } + pub fn get_node_allocators(&self, node_id: &NodeID) -> Vec> { let mut ret = Vec::new(); if let Some(node_allocators) = self.inner().node_allocators_and_tomb_tag.get(node_id) { @@ -417,7 +510,12 @@ impl MasterSegManager { ); // Now, register these segments in the master. - match register_node_segments(inner.view(), node_id.clone(), resp.serialize_part.seg_map) { + match register_node_segments( + inner.view(), + node_id.clone(), + expected_node_start_time, + resp.serialize_part.seg_map, + ) { Ok(()) => { tracing::info!("Successfully registered segments for node {}", node_id); } @@ -430,11 +528,117 @@ impl MasterSegManager { Ok(()) } - pub fn get_node_space_size(&self, node_id: &str) -> u64 { + pub fn get_node_physical_space_size(&self, node_id: &str) -> u64 { self.inner() .node_allocators_and_tomb_tag .get(node_id) + .filter(|node_segments_manager| !node_segments_manager.tomb_tag.is_tomb()) .map(|node_segments_manager| node_segments_manager.total_size) .unwrap_or(0) } + + pub fn get_node_active_space_size(&self, node_id: &str) -> u64 { + self.get_node_pool_capacity(node_id) + .map(|(_, snapshot)| snapshot.active_capacity_bytes) + .unwrap_or(0) + } + + /// Return the exact live node generation and its shared active/parked capacity state. + pub fn get_node_pool_capacity(&self, node_id: &str) -> Option<(i64, NodePoolCapacitySnapshot)> { + let node = self.inner().node_allocators_and_tomb_tag.get(node_id)?; + if node.tomb_tag.is_tomb() { + return None; + } + Some((node.node_start_time, node.capacity_budget.snapshot())) + } + + /// Update one live node generation with optimistic epoch fencing. + pub fn set_node_active_capacity( + &self, + node_id: &NodeID, + expected_node_start_time: i64, + expected_capacity_epoch: u64, + active_capacity_bytes: u64, + ) -> KvResult { + let node = self + .inner() + .node_allocators_and_tomb_tag + .get(node_id) + .ok_or_else(|| { + KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::NodeNotFound { + desc: node_id.to_string(), + }, + ) + })?; + if node.tomb_tag.is_tomb() { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::NodeNotFound { + desc: format!("{} (departed generation)", node_id), + }, + )); + } + if node.node_start_time != expected_node_start_time { + return Err(KvError::Api( + crate::rpcresp_kvresult_convert::msg_and_error::ApiError::OwnerStartTimeMismatch { + expected: expected_node_start_time, + got: node.node_start_time, + }, + )); + } + node.capacity_budget + .set_active_capacity(expected_capacity_epoch, active_capacity_bytes) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::master_seg_manager::msg_pack::SegmentDeviceMemInfo; + + #[test] + fn one_node_generation_shares_one_budget_across_registered_segments() { + let manager = build_node_segments_manager( + 17, + HashMap::from([ + ( + "cpu0".to_string(), + ( + SegmentDeviceDescription::Cpu, + SegmentDeviceMemInfo { + addr: 0, + len: 8 * 1024, + }, + ), + ), + ( + "cpu1".to_string(), + ( + SegmentDeviceDescription::Cpu, + SegmentDeviceMemInfo { + addr: 8 * 1024, + len: 8 * 1024, + }, + ), + ), + ]), + ) + .unwrap(); + let initial = manager.capacity_budget.snapshot(); + assert_eq!(initial.physical_capacity_bytes, 16 * 1024); + assert_eq!(initial.active_capacity_bytes, 16 * 1024); + + let allocators = manager + .device_id_2_allocator + .values() + .cloned() + .collect::>(); + let _first = allocators[0].allocate(8 * 1024).unwrap(); + let _second = allocators[1].allocate(8 * 1024).unwrap(); + assert_eq!( + manager.capacity_budget.snapshot().used_capacity_bytes, + 16 * 1024 + ); + assert!(allocators[0].allocate(1).is_err()); + } } diff --git a/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs b/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs index 6dcd136..817e278 100755 --- a/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs +++ b/fluxon_rs/fluxon_kv/src/master_seg_manager/one_seg_allocator.rs @@ -1,10 +1,209 @@ use fluxon_util::vallocator::VirtualAllocator; +use parking_lot::Mutex; -use crate::rpcresp_kvresult_convert::msg_and_error::KvResult; +use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}; use std::sync::Arc; use super::msg_pack::{SegmentDeviceDescription, SegmentDeviceID}; +/// Runtime capacity state shared by every segment registered by one node generation. +/// +/// Physical capacity describes the already allocated and registered memory. Active capacity is +/// the allocation budget currently exposed to Fluxon; the remainder is parked without changing +/// the underlying mapping or registration. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct NodePoolCapacitySnapshot { + pub physical_capacity_bytes: u64, + pub active_capacity_bytes: u64, + pub used_capacity_bytes: u64, + pub parked_capacity_bytes: u64, + pub draining_capacity_bytes: u64, + pub available_capacity_bytes: u64, + pub capacity_epoch: u64, +} + +#[derive(Debug)] +struct NodePoolCapacityState { + physical_capacity_bytes: u64, + active_capacity_bytes: u64, + used_capacity_bytes: u64, + capacity_epoch: u64, +} + +/// One generation-scoped allocation budget shared by all of a node's segment allocators. +#[derive(Debug)] +pub(crate) struct NodePoolCapacityBudget { + state: Mutex, +} + +impl NodePoolCapacityBudget { + pub(crate) fn new(physical_capacity_bytes: u64) -> KvResult { + if physical_capacity_bytes == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "node pool physical capacity must be greater than zero".to_string(), + })); + } + Ok(Self { + state: Mutex::new(NodePoolCapacityState { + physical_capacity_bytes, + active_capacity_bytes: physical_capacity_bytes, + used_capacity_bytes: 0, + capacity_epoch: 1, + }), + }) + } + + fn snapshot_locked(state: &NodePoolCapacityState) -> NodePoolCapacitySnapshot { + NodePoolCapacitySnapshot { + physical_capacity_bytes: state.physical_capacity_bytes, + active_capacity_bytes: state.active_capacity_bytes, + used_capacity_bytes: state.used_capacity_bytes, + parked_capacity_bytes: state + .physical_capacity_bytes + .saturating_sub(state.active_capacity_bytes), + draining_capacity_bytes: state + .used_capacity_bytes + .saturating_sub(state.active_capacity_bytes), + available_capacity_bytes: state + .active_capacity_bytes + .saturating_sub(state.used_capacity_bytes), + capacity_epoch: state.capacity_epoch, + } + } + + pub(crate) fn snapshot(&self) -> NodePoolCapacitySnapshot { + Self::snapshot_locked(&self.state.lock()) + } + + /// Change only the active allocation budget. Existing allocations above a lower target are + /// accounted as draining bytes and remain valid until the normal single-KV reclaim path drops + /// them. + pub(crate) fn set_active_capacity( + &self, + expected_capacity_epoch: u64, + active_capacity_bytes: u64, + ) -> KvResult { + let mut state = self.state.lock(); + if state.capacity_epoch != expected_capacity_epoch { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "stale node pool capacity epoch: expected={} current={}", + expected_capacity_epoch, state.capacity_epoch + ), + })); + } + if active_capacity_bytes == 0 || active_capacity_bytes > state.physical_capacity_bytes { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: format!( + "active capacity must be in 1..={}: got {}", + state.physical_capacity_bytes, active_capacity_bytes + ), + })); + } + if state.active_capacity_bytes != active_capacity_bytes { + state.active_capacity_bytes = active_capacity_bytes; + state.capacity_epoch = state + .capacity_epoch + .checked_add(1) + .expect("node pool capacity epoch overflow"); + } + Ok(Self::snapshot_locked(&state)) + } + + /// Extend a live generation when it registers an additional segment. A fully active pool + /// grows with its physical mapping; an already parked pool preserves its active byte target. + pub(crate) fn extend_physical_capacity(&self, additional_bytes: u64) -> KvResult<()> { + if additional_bytes == 0 { + return Err(KvError::Api(ApiError::InvalidArgument { + detail: "additional physical capacity must be greater than zero".to_string(), + })); + } + let mut state = self.state.lock(); + let was_fully_active = state.active_capacity_bytes == state.physical_capacity_bytes; + state.physical_capacity_bytes = state + .physical_capacity_bytes + .checked_add(additional_bytes) + .ok_or_else(|| { + KvError::Api(ApiError::InvalidArgument { + detail: "node pool physical capacity overflow".to_string(), + }) + })?; + if was_fully_active { + state.active_capacity_bytes = state.physical_capacity_bytes; + } + state.capacity_epoch = state + .capacity_epoch + .checked_add(1) + .expect("node pool capacity epoch overflow"); + Ok(()) + } + + fn allocate(&self, allocator: &VirtualAllocator, size: u64) -> KvResult<(u64, u64)> { + // The budget lock spans the physical allocator mutation. Capacity changes and allocations + // therefore have one total order; a resize cannot be bypassed by a stale pre-check. + let mut state = self.state.lock(); + if size + > state + .active_capacity_bytes + .saturating_sub(state.used_capacity_bytes) + { + return Err(KvError::Api(ApiError::Allocator { + detail: format!( + "node pool active capacity exhausted: requested={} active={} used={} epoch={}", + size, + state.active_capacity_bytes, + state.used_capacity_bytes, + state.capacity_epoch + ), + })); + } + let region = allocator.alloc(size)?; + let Some(next_used) = state.used_capacity_bytes.checked_add(region.size) else { + let _ = allocator.free(region.start_addr, region.size); + return Err(KvError::Api(ApiError::Allocator { + detail: "node pool used capacity overflow".to_string(), + })); + }; + if next_used > state.active_capacity_bytes { + // VirtualAllocator rounds allocations to its alignment. Roll the physical mutation + // back when the aligned size crosses the active byte boundary. + let _ = allocator.free(region.start_addr, region.size); + return Err(KvError::Api(ApiError::Allocator { + detail: format!( + "node pool active capacity exhausted after alignment: requested={} aligned={} active={} used={} epoch={}", + size, + region.size, + state.active_capacity_bytes, + state.used_capacity_bytes, + state.capacity_epoch + ), + })); + } + state.used_capacity_bytes = next_used; + Ok((region.start_addr, region.size)) + } + + fn free(&self, allocator: &VirtualAllocator, addr: u64, capacity: u64) { + let mut state = self.state.lock(); + match allocator.free(addr, capacity) { + Ok(freed) => { + state.used_capacity_bytes = state + .used_capacity_bytes + .checked_sub(freed) + .expect("node pool used capacity underflow"); + } + Err(err) => { + tracing::error!( + addr, + capacity, + error = %err, + "failed to free segment allocation; preserving node pool used accounting" + ); + } + } + } +} + /// An RAII guard for a memory allocation from a `OneSegAllocator`. /// /// When this guard is dropped, it attempts to free the memory block @@ -72,6 +271,15 @@ impl Allocation { self.allocator.base_addr } + /// Returns whether this allocation was created by the exact allocator instance. + /// + /// Segment names and node ids are reusable after a node reconnects. Pointer identity is + /// therefore the only safe way for master-side completion paths to bind an allocation back + /// to the registration generation that created it. + pub(crate) fn belongs_to_allocator(&self, allocator: &Arc) -> bool { + Arc::ptr_eq(&self.allocator, allocator) + } + /// Attach an on-drop callback. It will be executed exactly once /// when this allocation is dropped. pub fn set_on_drop(&mut self, f: F) @@ -99,6 +307,7 @@ pub struct OneSegAllocator { pub seg_device_desc: SegmentDeviceDescription, pub base_addr: u64, inner: VirtualAllocator, + capacity_budget: Arc, } impl OneSegAllocator { @@ -108,6 +317,23 @@ impl OneSegAllocator { seg_device_desc: SegmentDeviceDescription, base_addr: u64, size: u64, + ) -> KvResult { + let capacity_budget = Arc::new(NodePoolCapacityBudget::new(size)?); + Self::new_with_capacity_budget( + seg_device_id, + seg_device_desc, + base_addr, + size, + capacity_budget, + ) + } + + pub(crate) fn new_with_capacity_budget( + seg_device_id: SegmentDeviceID, + seg_device_desc: SegmentDeviceDescription, + base_addr: u64, + size: u64, + capacity_budget: Arc, ) -> KvResult { let inner = VirtualAllocator::new(size as u64)?; Ok(Self { @@ -115,18 +341,19 @@ impl OneSegAllocator { seg_device_desc, base_addr, inner, + capacity_budget, }) } /// Allocates a block of memory of `size` bytes. /// Returns an RAII guard for the allocation. pub fn allocate(self: &Arc, size: u64) -> KvResult { - let region = self.inner.alloc(size as u64)?; + let (start_addr, capacity) = self.capacity_budget.allocate(&self.inner, size)?; // return base0 offset in addr (pure offset); base address is carried separately Ok(Allocation::new( - region.start_addr, + start_addr, size, - region.size, + capacity, Arc::clone(self), )) } @@ -134,7 +361,7 @@ impl OneSegAllocator { /// Frees a block of memory. fn free(&self, addr: u64, capcity: u64) { // addr is offset (base0); free directly - let _ = self.inner.free(addr, capcity as u64); + self.capacity_budget.free(&self.inner, addr, capcity); } /// Returns total capacity (bytes) of this segment. @@ -146,4 +373,82 @@ impl OneSegAllocator { pub fn used_size_bytes(&self) -> u64 { self.inner.get_allocated_size() as u64 } + + /// Returns the active/parked capacity state shared by this node generation. + pub fn node_pool_capacity_snapshot(&self) -> NodePoolCapacitySnapshot { + self.capacity_budget.snapshot() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn allocator_with_budget( + id: &str, + size: u64, + budget: Arc, + ) -> Arc { + Arc::new( + OneSegAllocator::new_with_capacity_budget( + id.to_string(), + SegmentDeviceDescription::Cpu, + 0, + size, + budget, + ) + .unwrap(), + ) + } + + #[test] + fn shrinking_below_used_blocks_new_allocations_until_normal_drop_drains() { + let budget = Arc::new(NodePoolCapacityBudget::new(16 * 1024).unwrap()); + let allocator = allocator_with_budget("cpu", 16 * 1024, budget.clone()); + let first = allocator.allocate(8 * 1024).unwrap(); + + let shrinking = budget.set_active_capacity(1, 4 * 1024).unwrap(); + assert_eq!(shrinking.capacity_epoch, 2); + assert_eq!(shrinking.draining_capacity_bytes, 4 * 1024); + assert!(allocator.allocate(1).is_err()); + + drop(first); + let drained = budget.snapshot(); + assert_eq!(drained.used_capacity_bytes, 0); + assert_eq!(drained.draining_capacity_bytes, 0); + assert!(allocator.allocate(4 * 1024).is_ok()); + } + + #[test] + fn shared_node_budget_limits_allocations_across_segments_and_expands_immediately() { + let budget = Arc::new(NodePoolCapacityBudget::new(16 * 1024).unwrap()); + let first_allocator = allocator_with_budget("cpu0", 8 * 1024, budget.clone()); + let second_allocator = allocator_with_budget("cpu1", 8 * 1024, budget.clone()); + let first = first_allocator.allocate(8 * 1024).unwrap(); + let _second = second_allocator.allocate(8 * 1024).unwrap(); + assert!(first_allocator.allocate(1).is_err()); + + drop(first); + budget.set_active_capacity(1, 12 * 1024).unwrap(); + assert!(first_allocator.allocate(4 * 1024).is_ok()); + assert!(second_allocator.allocate(1).is_err()); + + let expanded = budget.set_active_capacity(2, 16 * 1024).unwrap(); + assert_eq!(expanded.capacity_epoch, 3); + assert!(first_allocator.allocate(4 * 1024).is_ok()); + } + + #[test] + fn stale_epoch_and_out_of_physical_range_updates_are_rejected() { + let budget = NodePoolCapacityBudget::new(16 * 1024).unwrap(); + assert!(budget.set_active_capacity(0, 8 * 1024).is_err()); + assert!(budget.set_active_capacity(1, 0).is_err()); + assert!(budget.set_active_capacity(1, 32 * 1024).is_err()); + assert_eq!(budget.snapshot().capacity_epoch, 1); + + let updated = budget.set_active_capacity(1, 8 * 1024).unwrap(); + assert_eq!(updated.capacity_epoch, 2); + assert!(budget.set_active_capacity(1, 4 * 1024).is_err()); + assert_eq!(budget.snapshot().active_capacity_bytes, 8 * 1024); + } } diff --git a/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs b/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs index ad23b4d..1ea89ec 100755 --- a/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/lifetime.rs @@ -1,5 +1,7 @@ use super::DeleteTargetMember; -use crate::client_kv_api::msg_pack::ExternalInvalidateWeakIndexReq; +use crate::client_kv_api::msg_pack::{ + ExternalInvalidateWeakIndexItem, ExternalInvalidateWeakIndexReq, +}; use crate::client_kv_api::{ClientKvApiView, ClientKvApiViewTrait, ExternalHoldingGetInfo}; use crate::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use crate::external_client_api::{ExternalClientApiView, ExternalClientApiViewTrait}; @@ -18,8 +20,6 @@ use std::collections::HashSet; use std::hash::Hash; use std::time::Duration; -const EXTERNAL_DELETE_ACK_TIMEOUT_SECS: u64 = 5; - pub type OwnerDeleteAckItem = DeleteAckItem; /// Spawn an async task only if the framework/view is still running. @@ -126,7 +126,6 @@ impl MemholderDropAck for OwnerDeleteAckCtx { pub struct ExternalDeleteAckCtx { pub view: ExternalClientApiView, - pub key: String, pub external_client_id: String, pub holder_id: u64, pub started_time: i64, @@ -146,75 +145,25 @@ impl MemholderDropAck for ExternalDeleteAckCtx { } fn ack_future( &self, - guard: Self::Guard, + _guard: Self::Guard, ) -> std::pin::Pin + Send>> { - let v = self.view.clone(); - let key = self.key.clone(); - let external_client_id = self.external_client_id.clone(); - let holder_id = self.holder_id; - let started_time = self.started_time; - Box::pin(async move { - let _keep = guard; - // Best-effort drop ACK must not outlive framework shutdown. - // During step8 teardown, the owner/external peers can already be concurrently exiting, - // and the ACK RPC can otherwise stall long enough to block task_registry shutdown. - match tokio::time::timeout( - Duration::from_secs(EXTERNAL_DELETE_ACK_TIMEOUT_SECS), - v.external_client_api().inner().send_external_delete_ack( - &key, - &external_client_id, - holder_id, - started_time, - ), - ) - .await - { - Err(_) => { - tracing::warn!( - "Timed out sending external_delete_ack for key={}, holder_id={}, external_client_id={} after {}s", - key, - holder_id, - external_client_id, - EXTERNAL_DELETE_ACK_TIMEOUT_SECS - ); - } - Ok(Err(e)) => { - tracing::warn!( - "Failed to send external_delete_ack for key={}, holder_id={}, external_client_id={}: {}", - key, - holder_id, - external_client_id, - e - ); - } - Ok(Ok(())) => { - tracing::debug!( - "Successfully sent external_delete_ack for key={}, holder_id={}, external_client_id={}", - key, - holder_id, - external_client_id - ); - } - } - }) + Box::pin(async { unreachable!("external holder ACKs are enqueued synchronously") }) } fn on_view_dropped(&self) { tracing::warn!( - "ExternalClientApiView has been dropped, cannot send external_delete_ack for key='{}', holder_id {}.", - self.key, + "ExternalClientApiView has been dropped, cannot enqueue external_delete_ack for holder_id {}.", self.holder_id ); } fn on_skip_shutdown(&self) { tracing::info!( - "Skipping external_delete_ack for key={}, holder_id={} due to shutdown", - self.key, + "Skipping external_delete_ack for holder_id={} due to shutdown", self.holder_id ); } fn run_drop_ack(&self) { let v = self.view().clone(); - let Some(g) = ::try_upgrade(&v) else { + let Some(_guard) = ::try_upgrade(&v) else { self.on_view_dropped(); return; }; @@ -222,8 +171,18 @@ impl MemholderDropAck for ExternalDeleteAckCtx { self.on_skip_shutdown(); return; } - let fut = self.ack_future(g); - let _ = v.spawn("memholder_drop_ack_external", async move { fut.await }); + if let Err(err) = v.external_client_api().inner().enqueue_external_delete_ack( + self.external_client_id.clone(), + self.holder_id, + self.started_time, + ) { + tracing::warn!( + holder_id = self.holder_id, + external_client_id = %self.external_client_id, + error = %err, + "failed to enqueue external holder ACK" + ); + } } } @@ -474,8 +433,8 @@ impl MemholderManagerTrait for MasterOwnerMemMgr { }; let mut targets = Vec::new(); - let nodes_replicas = nodes_kv_route_info.nodes_replicas.read(); - for (node_id, kv_route_info) in nodes_replicas.iter() { + let node_replicas = nodes_kv_route_info.node_replicas.read(); + for (node_id, kv_route_info) in node_replicas.iter() { if kv_route_info.tomb_tag.is_tomb() { continue; } @@ -517,8 +476,8 @@ impl MemholderManagerTrait for MasterOwnerMemMgr { continue; }; - let nodes_replicas = nodes_kv_route_info.nodes_replicas.read(); - let Some(kv_route_info) = nodes_replicas.get(target.node_id.as_str()) else { + let node_replicas = nodes_kv_route_info.node_replicas.read(); + let Some(kv_route_info) = node_replicas.get(target.node_id.as_str()) else { continue; }; if kv_route_info.tomb_tag.is_tomb() { @@ -719,6 +678,86 @@ pub struct OwnerExternalMemMgr { inner: MemholderManagerInner, } +impl OwnerExternalMemMgr { + /// Remove only holdings owned by the departed membership generation. + /// Node ids are reusable, so node-only cleanup can release memory still in + /// use by a reconnected external requester. + pub fn cleanup_node_generation(&self, node_id: &str, node_start_time: i64) -> usize { + let keys = self + .inner + .as_map() + .iter() + .filter_map(|entry| { + owner_external_holding_matches_generation( + &entry.value().req_node_id, + entry.value().requester_node_start_time, + node_id, + node_start_time, + ) + .then(|| entry.key().clone()) + }) + .collect::>(); + + keys.into_iter() + .filter(|key| { + self.inner + .as_map() + .remove_if(key, |_, holding| { + owner_external_holding_matches_generation( + &holding.req_node_id, + holding.requester_node_start_time, + node_id, + node_start_time, + ) + }) + .is_some() + }) + .count() + } +} + +fn owner_external_holding_matches_generation( + holding_node_id: &str, + holding_node_start_time: Option, + departed_node_id: &str, + departed_node_start_time: i64, +) -> bool { + holding_node_id == departed_node_id && holding_node_start_time == Some(departed_node_start_time) +} + +#[cfg(test)] +mod owner_external_generation_tests { + use super::owner_external_holding_matches_generation; + + #[test] + fn holding_cleanup_never_matches_reconnected_or_unknown_generation() { + assert!(owner_external_holding_matches_generation( + "external-a", + Some(11), + "external-a", + 11, + )); + assert!(!owner_external_holding_matches_generation( + "external-a", + Some(12), + "external-a", + 11, + )); + assert!(!owner_external_holding_matches_generation( + "external-a", + None, + "external-a", + 11, + )); + assert!(!owner_external_holding_matches_generation( + "external-b", + Some(11), + "external-a", + 11, + )); + } +} + impl Default for OwnerExternalMemMgr { fn default() -> Self { Self { @@ -760,15 +799,21 @@ impl MemholderManagerTrait for OwnerExternalMemMgr { if holding.key != task.key { continue; } + let Some(holding_epoch) = holding.requester_node_start_time else { + continue; + }; let Some(member) = ctx .cluster_manager() .get_member_info_cached(&holding.req_node_id) else { continue; }; + if member.node_start_time != holding_epoch { + continue; + } targets.insert(DeleteTargetMember::new( holding.req_node_id.clone(), - member.node_start_time, + holding_epoch, )); } targets.into_iter().collect() @@ -786,18 +831,24 @@ impl MemholderManagerTrait for OwnerExternalMemMgr { target: Self::DeleteTarget, tasks: Vec, ) -> Result<(), String> { - let mut keys: Vec = tasks.into_iter().map(|task| task.key).collect(); - keys.sort(); - keys.dedup(); - - if keys.is_empty() { + let mut items: Vec = tasks + .into_iter() + .map(|task| ExternalInvalidateWeakIndexItem { key: task.key }) + .collect(); + items.sort_by(|a, b| a.key.cmp(&b.key)); + items.dedup_by(|a, b| a.key == b.key); + + if items.is_empty() { return Ok(()); } let rpc_caller = RPCCaller::::new(); rpc_caller.regist(ctx.p2p_module()); let req = MsgPack { - serialize_part: ExternalInvalidateWeakIndexReq { keys }, + serialize_part: ExternalInvalidateWeakIndexReq { + keys: Vec::new(), + items, + }, raw_bytes: Vec::new(), }; diff --git a/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs b/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs index 692a9a0..07fdf7b 100644 --- a/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/memholder_test.rs @@ -15,7 +15,7 @@ use limit_thirdparty::tokio::time::sleep; use crate::config::{ ClientConfig, ContributeToClusterPoolSize, FluxonKvSpec, MasterConfig, MonitoringConfig, - ProtocolConfig, ProtocolType, TestSpecConfig, TransferEngineType, + ProtocolConfig, ProtocolType, ReplicaTaskPlacementConfig, TestSpecConfig, TransferEngineType, }; use crate::master_kv_router::MasterKvRouterView; use crate::master_seg_manager::one_seg_allocator::Allocation; @@ -58,6 +58,9 @@ fn new_master_config( log_dir: "/tmp/fluxon_master_logs".to_string(), pprof_duration_seconds: None, master_ui: None, + replica_task_placement: ReplicaTaskPlacementConfig::default(), + replica_cache_capacity_ratio: crate::config::DEFAULT_REPLICA_CACHE_CAPACITY_RATIO, + replica_writeback_tier1_capacity_ratio: None, test_spec_config: TestSpecConfig::default(), } } @@ -88,6 +91,7 @@ fn new_client_config_with_size( rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: vec![etcd.to_string()], @@ -101,6 +105,7 @@ fn new_client_config_with_size( large_file_paths: crate::config::LargeFilePaths { paths: vec![format!("/tmp/kvcache_large/{}", instance_key)], }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -123,6 +128,7 @@ fn new_zero_contribution_client_config( rdma_device_names: None, }, pprof_duration_seconds: None, + replica_writeback_hot_capacity_ratio: None, redis_compat_listen_addr: None, fluxonkv_spec: FluxonKvSpec { etcd_addresses: Vec::new(), @@ -134,6 +140,7 @@ fn new_zero_contribution_client_config( }, share_mem_path: format!("/tmp/kvcache_shared_memory/{}", owner_instance_key), large_file_paths: crate::config::LargeFilePaths { paths: Vec::new() }, + ssd_storage: None, test_spec_config: TestSpecConfig::default(), } } @@ -192,7 +199,7 @@ fn evict_all_replicas_for_key(master_view: &MasterKvRouterView, key: &str) -> us .get(key) .map(|route| { let node_ids = route - .nodes_replicas + .node_replicas .read() .keys() .cloned() @@ -241,6 +248,7 @@ pub mod test_memholder { new_master_config, new_zero_contribution_client_config, read_etcd, unique_cluster_name, wait_node_allocators, wait_weak_drop, }; + use crate::memholder::{MemholderManagerTrait, NodeHolderKey}; use crate::{ConfigArg, run_client, run_master}; use std::time::{Duration, Instant}; use tokio::time::sleep; @@ -587,10 +595,16 @@ pub mod test_memholder { .unwrap() .unwrap(); let ext_hold_holder_id = eh1.holder_id; - // Note: master's get_holding is keyed by the requesting node of the inner owner get - // for external GET, the owner performs an inner get, so the node is the owner + let ext_hold_owner_holder_id = owner_api + .inner() + .external_get_holding + .inner_map() + .get(&NodeHolderKey::new(external_id.clone(), ext_hold_holder_id)) + .expect("capture ext_hold owner holding") + .memory_info + .holder_id; let ext_hold_weak = - capture_master_holding_allocation(&master_view, &owner_id, ext_hold_holder_id) + capture_master_holding_allocation(&master_view, &owner_id, ext_hold_owner_holder_id) .expect("capture ext_hold weak"); drop(eh2); @@ -607,9 +621,19 @@ pub mod test_memholder { .unwrap() .unwrap(); let ext_release_holder_id = er1.holder_id; - // Similar rationale as above: use owner_id when capturing master's holding entry + let ext_release_owner_holder_id = owner_api + .inner() + .external_get_holding + .inner_map() + .get(&NodeHolderKey::new( + external_id.clone(), + ext_release_holder_id, + )) + .expect("capture ext_release owner holding") + .memory_info + .holder_id; let ext_release_weak = - capture_master_holding_allocation(&master_view, &owner_id, ext_release_holder_id) + capture_master_holding_allocation(&master_view, &owner_id, ext_release_owner_holder_id) .expect("capture ext_release weak"); drop(er1); drop(er2); // drop all for release key diff --git a/fluxon_rs/fluxon_kv/src/memholder/mod.rs b/fluxon_rs/fluxon_kv/src/memholder/mod.rs index db0ada2..55cb437 100755 --- a/fluxon_rs/fluxon_kv/src/memholder/mod.rs +++ b/fluxon_rs/fluxon_kv/src/memholder/mod.rs @@ -23,6 +23,18 @@ use crate::client_kv_api::ClientKvApiView; use crate::{cluster_manager::NodeID, external_client_api::ExternalClientApiView}; use bitcode::{Decode, Encode}; +use fluxon_util::pin_aware_moka::PinGuard; + +#[derive(Clone, Debug)] +pub enum MemoryInfoDropAction { + None, + OwnerDeleteAck, + ReleaseLocalReserveResidentSlot { + slot_size: u64, + grant_id: u64, + slot_index: u32, + }, +} /// Memory metadata for owner/client user holders #[derive(Clone)] @@ -35,6 +47,7 @@ pub struct MemoryInfo { pub key: String, pub master_node_id: NodeID, pub view: ClientKvApiView, + pub drop_action: MemoryInfoDropAction, } impl std::fmt::Debug for MemoryInfo { @@ -46,6 +59,7 @@ impl std::fmt::Debug for MemoryInfo { .field("holder_id", &self.holder_id) .field("key", &self.key) .field("master_node_id", &self.master_node_id) + .field("drop_action", &self.drop_action) .finish() } } @@ -78,6 +92,7 @@ pub struct UserMemHolder { pub memory_info: Arc, pub refcount: Arc, expose_kind: UserMemHolderExposeKind, + _owner_hot_pin: Option, } impl UserMemHolder { pub fn holder_id(&self) -> u64 { @@ -104,6 +119,14 @@ impl UserMemHolder { refcount: Arc, expose_kind: UserMemHolderExposeKind, ) -> Self { + let owner_hot_pin = match expose_kind { + UserMemHolderExposeKind::SegPtr => refcount + .view + .client_kv_api() + .inner() + .owner_hot_pin_memory_info(&memory_info), + UserMemHolderExposeKind::OwnedCopy => None, + }; tracing::debug!( "Creating UserMemHolder for key '{}', _holder_id_ {}, expose_kind={:?}.", memory_info.key, @@ -114,6 +137,7 @@ impl UserMemHolder { memory_info, refcount, expose_kind, + _owner_hot_pin: owner_hot_pin, } } pub fn memory_info(&self) -> Arc { @@ -159,8 +183,74 @@ impl MemoryInfo { key, master_node_id, view, + drop_action: MemoryInfoDropAction::OwnerDeleteAck, + } + } + + pub async fn new_local_reserve_resident( + addr: u64, + len: u32, + key: String, + master_node_id: NodeID, + view: ClientKvApiView, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> Self { + let base_addr = { + let base_guard = view + .client_seg_pool() + .cpu_mem_read_guard() + .await + .expect("segment cpu mem must be available when creating resident MemoryInfo"); + base_guard.allocated_addr + }; + Self::new_local_reserve_resident_with_base( + addr, + base_addr, + len, + key, + master_node_id, + view, + slot_size, + grant_id, + slot_index, + ) + } + + pub fn new_local_reserve_resident_with_base( + addr: u64, + base_addr: u64, + len: u32, + key: String, + master_node_id: NodeID, + view: ClientKvApiView, + slot_size: u64, + grant_id: u64, + slot_index: u32, + ) -> Self { + assert!( + addr >= base_addr, + "resident local reserve addr must be within client segment: addr={:#x} base_addr={:#x}", + addr, + base_addr + ); + Self { + offset: addr - base_addr, + addr, + len, + holder_id: 0, + key, + master_node_id, + view, + drop_action: MemoryInfoDropAction::ReleaseLocalReserveResidentSlot { + slot_size, + grant_id, + slot_index, + }, } } + pub fn bytes(&self) -> &[u8] { tracing::debug!( "MemHolder accessing memory: addr={:#x}, len={}", @@ -169,16 +259,87 @@ impl MemoryInfo { ); unsafe { std::slice::from_raw_parts(self.addr as *const u8, self.len as usize) } } + + pub fn local_reserve_resident_slot_ref(&self) -> Option<(u64, u64, u32)> { + match self.drop_action { + MemoryInfoDropAction::ReleaseLocalReserveResidentSlot { + slot_size, + grant_id, + slot_index, + } => Some((slot_size, grant_id, slot_index)), + MemoryInfoDropAction::None | MemoryInfoDropAction::OwnerDeleteAck => None, + } + } + + pub(crate) fn take_local_reserve_resident_slot_ref(&mut self) -> Option<(u64, u64, u32)> { + let action = std::mem::replace(&mut self.drop_action, MemoryInfoDropAction::None); + match action { + MemoryInfoDropAction::ReleaseLocalReserveResidentSlot { + slot_size, + grant_id, + slot_index, + } => Some((slot_size, grant_id, slot_index)), + other => { + self.drop_action = other; + None + } + } + } } + /// Represents a memory holder that keeps a reference to transferred data impl Drop for MemoryInfo { fn drop(&mut self) { - let ctx = OwnerDeleteAckCtx { - view: self.view.clone(), - key: self.key.clone(), - holder_id: self.holder_id, - }; - ctx.run_drop_ack(); + match &self.drop_action { + MemoryInfoDropAction::None => {} + MemoryInfoDropAction::OwnerDeleteAck => { + let ctx = OwnerDeleteAckCtx { + view: self.view.clone(), + key: self.key.clone(), + holder_id: self.holder_id, + }; + ctx.run_drop_ack(); + } + MemoryInfoDropAction::ReleaseLocalReserveResidentSlot { + slot_size, + grant_id, + slot_index, + } => { + let Some(_view_guard) = self.view.try_upgrade() else { + tracing::debug!( + key = self.key, + "skipping resident local reserve slot release after view drop" + ); + return; + }; + if !self.view.register_shutdown_poller().is_running() { + tracing::debug!( + key = self.key, + "skipping resident local reserve slot release during shutdown" + ); + return; + } + if let Err(err) = self + .view + .client_kv_api() + .inner() + .owner_release_local_reserve_resident_slot_holder( + *slot_size, + *grant_id, + *slot_index, + ) + { + tracing::warn!( + "failed to release resident local reserve slot holder on MemoryInfo drop: key={} slot_size={} grant_id={} slot_index={} err={}", + self.key, + slot_size, + grant_id, + slot_index, + err + ); + } + } + } } } @@ -253,7 +414,6 @@ impl Drop for ExternalMemHolder { let ctx = ExternalDeleteAckCtx { view: self.view.clone(), - key: self.key.clone(), external_client_id: self.external_client_id.clone(), holder_id: self.holder_id, started_time: self.owner_start_time, diff --git a/fluxon_rs/fluxon_kv/src/metrics.rs b/fluxon_rs/fluxon_kv/src/metrics.rs index 88f1b9c..4f5d88e 100644 --- a/fluxon_rs/fluxon_kv/src/metrics.rs +++ b/fluxon_rs/fluxon_kv/src/metrics.rs @@ -26,6 +26,8 @@ pub struct MetricsHandle { // Pending put stats captured at put_start to attribute external-owner path at put_end. pending_put_stats: DashMap, + + locality_counters: KvLocalityCounters, } #[derive(Clone, Debug)] @@ -43,6 +45,7 @@ pub struct PendingPutStat { pub transfer_poll_wait_us: i64, pub transfer_poll_iters: i64, pub transfer_used_fast_path: bool, + pub transfer_used_nixl: bool, pub transfer_local_noop: bool, pub transfer_remote_transfer: bool, pub top_emitted: bool, @@ -54,6 +57,85 @@ pub enum OperationKind { Get, } +#[derive(Clone, Debug, Default)] +pub struct KvIoLocalitySnapshot { + pub op_count: u64, + pub bytes: u64, + pub transfer_us: u64, +} + +impl KvIoLocalitySnapshot { + pub fn bandwidth_gbps(&self) -> f64 { + if self.transfer_us == 0 { + return 0.0; + } + (self.bytes as f64) / (self.transfer_us as f64 / 1_000_000.0) / 1_000_000_000.0 + } +} + +#[derive(Clone, Debug, Default)] +pub struct KvLocalitySnapshot { + pub l2_local_hit_pages: u64, + pub l2_local_hit_bytes: u64, + pub l2_remote_hit_pages: u64, + pub l2_remote_hit_bytes: u64, + pub put_local: KvIoLocalitySnapshot, + pub put_remote: KvIoLocalitySnapshot, + pub get_local: KvIoLocalitySnapshot, + pub get_remote: KvIoLocalitySnapshot, +} + +#[derive(Default)] +struct KvLocalityCounters { + l2_local_hit_pages: AtomicU64, + l2_local_hit_bytes: AtomicU64, + l2_remote_hit_pages: AtomicU64, + l2_remote_hit_bytes: AtomicU64, + put_local_ops: AtomicU64, + put_local_bytes: AtomicU64, + put_local_transfer_us: AtomicU64, + put_remote_ops: AtomicU64, + put_remote_bytes: AtomicU64, + put_remote_transfer_us: AtomicU64, + get_local_ops: AtomicU64, + get_local_bytes: AtomicU64, + get_local_transfer_us: AtomicU64, + get_remote_ops: AtomicU64, + get_remote_bytes: AtomicU64, + get_remote_transfer_us: AtomicU64, +} + +impl KvLocalityCounters { + fn snapshot(&self) -> KvLocalitySnapshot { + KvLocalitySnapshot { + l2_local_hit_pages: self.l2_local_hit_pages.load(Ordering::Relaxed), + l2_local_hit_bytes: self.l2_local_hit_bytes.load(Ordering::Relaxed), + l2_remote_hit_pages: self.l2_remote_hit_pages.load(Ordering::Relaxed), + l2_remote_hit_bytes: self.l2_remote_hit_bytes.load(Ordering::Relaxed), + put_local: KvIoLocalitySnapshot { + op_count: self.put_local_ops.load(Ordering::Relaxed), + bytes: self.put_local_bytes.load(Ordering::Relaxed), + transfer_us: self.put_local_transfer_us.load(Ordering::Relaxed), + }, + put_remote: KvIoLocalitySnapshot { + op_count: self.put_remote_ops.load(Ordering::Relaxed), + bytes: self.put_remote_bytes.load(Ordering::Relaxed), + transfer_us: self.put_remote_transfer_us.load(Ordering::Relaxed), + }, + get_local: KvIoLocalitySnapshot { + op_count: self.get_local_ops.load(Ordering::Relaxed), + bytes: self.get_local_bytes.load(Ordering::Relaxed), + transfer_us: self.get_local_transfer_us.load(Ordering::Relaxed), + }, + get_remote: KvIoLocalitySnapshot { + op_count: self.get_remote_ops.load(Ordering::Relaxed), + bytes: self.get_remote_bytes.load(Ordering::Relaxed), + transfer_us: self.get_remote_transfer_us.load(Ordering::Relaxed), + }, + } + } +} + impl OperationKind { pub const fn as_label(self) -> &'static str { match self { @@ -147,6 +229,7 @@ impl MetricsHandle { get_metrics_queue: SegQueue::new(), latest_metrics_snapshot: RwLock::new(HashMap::new()), pending_put_stats: DashMap::new(), + locality_counters: KvLocalityCounters::default(), } } @@ -184,6 +267,78 @@ impl MetricsHandle { } } + pub fn get_locality_snapshot(&self) -> KvLocalitySnapshot { + self.locality_counters.snapshot() + } + + pub fn record_l2_hit_locality(&self, remote: bool, bytes: u64) { + if remote { + self.locality_counters + .l2_remote_hit_pages + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .l2_remote_hit_bytes + .fetch_add(bytes, Ordering::Relaxed); + } else { + self.locality_counters + .l2_local_hit_pages + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .l2_local_hit_bytes + .fetch_add(bytes, Ordering::Relaxed); + } + } + + pub fn record_put_io_locality(&self, remote: bool, bytes: u64, transfer_us: i64) { + let transfer_us = transfer_us.max(0) as u64; + if remote { + self.locality_counters + .put_remote_ops + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .put_remote_bytes + .fetch_add(bytes, Ordering::Relaxed); + self.locality_counters + .put_remote_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } else { + self.locality_counters + .put_local_ops + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .put_local_bytes + .fetch_add(bytes, Ordering::Relaxed); + self.locality_counters + .put_local_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } + } + + pub fn record_get_io_locality(&self, remote: bool, bytes: u64, transfer_us: i64) { + let transfer_us = transfer_us.max(0) as u64; + if remote { + self.locality_counters + .get_remote_ops + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .get_remote_bytes + .fetch_add(bytes, Ordering::Relaxed); + self.locality_counters + .get_remote_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } else { + self.locality_counters + .get_local_ops + .fetch_add(1, Ordering::Relaxed); + self.locality_counters + .get_local_bytes + .fetch_add(bytes, Ordering::Relaxed); + self.locality_counters + .get_local_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } + } + pub fn emit_op_end_bytes_pulse( &self, op: OperationKind, @@ -252,6 +407,39 @@ impl MetricsHandle { }); } + pub fn set_kv_holding_entries(&self, scope: &str, entries: u64) { + let _ = (scope, entries); + } + + pub fn set_kv_holding_bytes(&self, scope: &str, bytes: u64) { + let _ = (scope, bytes); + } + + pub fn set_kv_external_pending_put_entries(&self, entries: u64) { + let _ = entries; + } + + pub fn set_kv_replica_cache_entries(&self, owner_node: &str, entries: u64) { + let _ = (owner_node, entries); + } + + pub fn set_kv_replica_cache_weighted_bytes(&self, owner_node: &str, bytes: u64) { + let _ = (owner_node, bytes); + } + + pub fn set_kv_replica_cache_capacity_bytes( + &self, + owner_node: &str, + capacity_kind: &str, + bytes: u64, + ) { + let _ = (owner_node, capacity_kind, bytes); + } + + pub fn inc_kv_get_done_allocation(&self, allocation_mode: &str) { + let _ = allocation_mode; + } + pub fn set_fs_mount_fs_bytes( &self, mount_kind: FsMountKind, @@ -333,6 +521,7 @@ impl MetricsHandle { transfer_poll_wait_us, transfer_poll_iters, transfer_used_fast_path, + transfer_used_nixl, transfer_local_noop, transfer_remote_transfer, } = &metric @@ -356,7 +545,7 @@ impl MetricsHandle { if let Some(seq) = sampled_debug_seq(&PUT_METRIC_SUBMIT_SEQ) { tracing::debug!( - "submit kv op metric seq={} kind=put put_id={} key={} whole_us={} start_us={} transfer_us={} end_us={} rpc_us={} start_handle_us={} end_handle_us={} t1_us={} t2_us={} t3_us={} t4_us={} transfer_submit_blocking_us={} transfer_create_xfer_req_us={} transfer_post_xfer_req_us={} transfer_poll_wait_us={} transfer_poll_iters={} transfer_used_fast_path={} transfer_local_noop={} transfer_remote_transfer={}", + "submit kv op metric seq={} kind=put put_id={} key={} whole_us={} start_us={} transfer_us={} end_us={} rpc_us={} start_handle_us={} end_handle_us={} t1_us={} t2_us={} t3_us={} t4_us={} transfer_submit_blocking_us={} transfer_create_xfer_req_us={} transfer_post_xfer_req_us={} transfer_poll_wait_us={} transfer_poll_iters={} transfer_used_fast_path={} transfer_used_nixl={} transfer_local_noop={} transfer_remote_transfer={}", seq, put_id, sanitize_key(key), @@ -377,6 +566,7 @@ impl MetricsHandle { transfer_poll_wait_us, transfer_poll_iters, transfer_used_fast_path, + transfer_used_nixl, transfer_local_noop, transfer_remote_transfer ); @@ -537,6 +727,7 @@ impl MetricsHandle { transfer_poll_wait_us: 0, transfer_poll_iters: 0, transfer_used_fast_path: false, + transfer_used_nixl: false, transfer_local_noop: false, transfer_remote_transfer: false, top_emitted: false, @@ -580,6 +771,7 @@ impl MetricsHandle { transfer_poll_wait_us: i64, transfer_poll_iters: i64, transfer_used_fast_path: bool, + transfer_used_nixl: bool, transfer_local_noop: bool, transfer_remote_transfer: bool, ) { @@ -593,6 +785,7 @@ impl MetricsHandle { guard.transfer_poll_wait_us = transfer_poll_wait_us; guard.transfer_poll_iters = transfer_poll_iters; guard.transfer_used_fast_path = transfer_used_fast_path; + guard.transfer_used_nixl = transfer_used_nixl; guard.transfer_local_noop = transfer_local_noop; guard.transfer_remote_transfer = transfer_remote_transfer; } diff --git a/fluxon_rs/fluxon_kv/src/observe_kvope.rs b/fluxon_rs/fluxon_kv/src/observe_kvope.rs index 69c724d..d67ec67 100755 --- a/fluxon_rs/fluxon_kv/src/observe_kvope.rs +++ b/fluxon_rs/fluxon_kv/src/observe_kvope.rs @@ -245,6 +245,7 @@ pub fn obe_put_done_success( transfer_poll_wait_us, transfer_poll_iters, transfer_used_fast_path, + transfer_used_nixl: false, transfer_local_noop, transfer_remote_transfer, }); diff --git a/fluxon_rs/fluxon_kv/src/p2p/msg_pack.rs b/fluxon_rs/fluxon_kv/src/p2p/msg_pack.rs index 6f774ca..a057609 100644 --- a/fluxon_rs/fluxon_kv/src/p2p/msg_pack.rs +++ b/fluxon_rs/fluxon_kv/src/p2p/msg_pack.rs @@ -1,6 +1,7 @@ pub use fluxon_commu::p2p::rpc::{ - MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, MsgPackSerializePart, RPCCaller, RPCHandler, RPCReq, - RPCResponsor, Responser, RpcCallObserveTrace, RpcCallObservedOutput, call_rpc, - call_rpc_observed, + MIN_EXPLICIT_RPC_TIMEOUT_MS, MIN_EXPLICIT_RPC_TIMEOUT_SECS, MsgPack, MsgPackSerializePart, + RPCCaller, RPCHandler, RPCReq, RPCResponsor, Responser, RpcCallObserveTrace, + RpcCallObservedOutput, call_rpc, call_rpc_observed, validate_explicit_rpc_timeout, + validate_explicit_rpc_timeout_ms, }; pub use fluxon_commu::p2p::{MsgId, MsgPackHeadMeta, MsgPackRelay, TaskId, WireMessageBody}; diff --git a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs index 42a9cbc..d9d24a8 100644 --- a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs +++ b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/msg_and_error.rs @@ -43,6 +43,56 @@ pub enum MsgId { CountPrefixResp = 3026, GetMasterOnlyMetricPartReq = 3027, GetMasterOnlyMetricPartResp = 3028, + BatchIsExistReq = 3031, + BatchIsExistResp = 3032, + PutAppendStartReq = 3033, + PutAppendStartResp = 3034, + PutAppendRevokeReq = 3035, + PutAppendRevokeResp = 3036, + PutAppendDoneReq = 3037, + PutAppendDoneResp = 3038, + BatchPutStartReq = 3039, + BatchPutStartResp = 3040, + BatchPutRevokeReq = 3041, + BatchPutRevokeResp = 3042, + BatchPutDoneReq = 3043, + BatchPutDoneResp = 3044, + BatchGetStartReq = 3045, + BatchGetStartResp = 3046, + BatchGetRevokeReq = 3047, + BatchGetRevokeResp = 3048, + BatchGetDoneReq = 3049, + BatchGetDoneResp = 3050, + ReserveLocalGrantReq = 3051, + ReserveLocalGrantResp = 3052, + ReleaseLocalGrantReq = 3053, + ReleaseLocalGrantResp = 3054, + BatchPreparePutKeysReq = 3055, + BatchPreparePutKeysResp = 3056, + BatchReleasePutKeyReservationsReq = 3057, + BatchReleasePutKeyReservationsResp = 3058, + BatchOwnerReclaimReq = 3059, + BatchOwnerReclaimResp = 3060, + BatchEnqueueReplicaTaskReq = 3061, + BatchEnqueueReplicaTaskResp = 3062, + GroupedBatchPutDoneReq = 3063, + GroupedBatchPutDoneResp = 3064, + BatchPutAppendStartReq = 3065, + BatchPutAppendStartResp = 3066, + BatchPutAppendDoneReq = 3067, + BatchPutAppendDoneResp = 3068, + BatchEvictOwnerSourceReq = 3069, + BatchEvictOwnerSourceResp = 3070, + BatchGetPlanReq = 3071, + BatchGetPlanResp = 3072, + BatchGetBindReq = 3073, + BatchGetBindResp = 3074, + BatchPublishOwnerSsdReq = 3075, + BatchPublishOwnerSsdResp = 3076, + SsdStageBeginReq = 3077, + SsdStageBeginResp = 3078, + SsdStageDoneReq = 3079, + SsdStageDoneResp = 3080, ExternalGetReq = 4001, ExternalGetResp = 4002, ExternalPutStartReq = 4003, @@ -51,6 +101,8 @@ pub enum MsgId { ExternalPutTransferEndResp = 4006, ExternalDeleteReq = 4009, ExternalDeleteResp = 4010, + ExternalBatchIsExistReq = 4015, + ExternalBatchIsExistResp = 4016, AllocateClientLeaseReq = 5001, AllocateClientLeaseResp = 5002, ClientLeaseKeepaliveReq = 5003, @@ -704,6 +756,12 @@ crate::define_err_group! { , (122, KeyBeingWritten { key: String }, msg: "Key is currently being written: key={key}") + , + (123, KeyAlreadyExists { key: String }, + msg: "Key already exists: key={key}") + , + (124, StaleGetPlan { get_id: u64, key: String, detail: String }, + msg: "Get plan is stale: get_id={get_id}, key={key}, detail={detail}") } } diff --git a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs index b6eb7d6..77efaac 100755 --- a/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs +++ b/fluxon_rs/fluxon_kv/src/rpcresp_kvresult_convert/rpcresp_kvresult_convert.rs @@ -1,13 +1,23 @@ use super::msg_and_error::OK; -use super::msg_and_error::{ErrorCode, KvError, KvResult}; +use super::msg_and_error::{ + ApiError, ConfigError, ErrorCode, KvError, KvResult, SharedMemError, UnreachableError, +}; use crate::client_kv_api::msg_pack::{ - ExternalDeleteAckResp, ExternalDeleteResp, ExternalGetResp, ExternalIsExistResp, - ExternalPutCommitResp, ExternalPutRevokeResp, ExternalPutStartResp, ExternalPutTransferEndResp, + ExternalBatchDeleteAckResp, ExternalBatchGetCancelResp, ExternalBatchGetItemResp, + ExternalBatchGetResp, ExternalBatchGetStartResp, ExternalBatchGetTransferResp, + ExternalBatchIsExistResp, ExternalBatchPutCommitItemResp, ExternalBatchPutCommitResp, + ExternalBatchPutStartItemResp, ExternalBatchPutStartResp, ExternalBatchPutTransferEndItemResp, + ExternalBatchPutTransferEndResp, ExternalDeleteAckResp, ExternalDeleteResp, ExternalGetResp, + ExternalIsExistResp, ExternalObservabilitySnapshotResp, ExternalPutCommitResp, + ExternalPutRevokeResp, ExternalPutStartResp, ExternalPutTransferEndResp, }; use crate::master_kv_router::msg_pack::{ - BatchDeleteAckResp, BatchDeleteClientKvMetaCacheResp, DeleteAckResp, DeleteResp, GetDoneResp, - GetMasterOnlyMetricPartResp, GetMetaResp, GetRevokeResp, GetStartResp, MemHolderKeepAliveResp, - MemHolderReleaseResp, PutDoneResp, PutRevokeResp, PutStartResp, + BatchDeleteAckResp, BatchDeleteClientKvMetaCacheResp, BatchIsExistResp, + BatchPreparePutKeysResp, BatchReleasePutKeyReservationsResp, DeleteAckResp, DeleteResp, + GetDoneResp, GetMasterOnlyMetricPartResp, GetMetaResp, GetRevokeResp, GetStartResp, + MemHolderKeepAliveResp, MemHolderReleaseResp, PutAppendDoneResp, PutAppendRevokeResp, + PutAppendStartResp, PutDoneResp, PutRevokeResp, PutStartResp, ReleaseLocalGrantResp, + ReserveLocalGrantResp, }; use crate::master_seg_manager::msg_pack::RequestSegmentRegistrationResp; use crate::memholder::ExternalMemHolderInfo; @@ -58,6 +68,16 @@ impl ToResult for ExternalIsExistResp { } } +impl ToResult for ExternalBatchIsExistResp { + type Ok = Vec; + fn to_result(self) -> KvResult { + if self.error_code == OK { + return Ok(self.exists_list); + } + Err(KvError::from_json(self.error_code, &self.error_json)) + } +} + /// Convert ExternalDeleteResp into KvResult<()> impl ToResult for ExternalDeleteResp { type Ok = (); @@ -131,6 +151,27 @@ impl ToResult for ExternalDeleteAckResp { } } +impl ToResult for ExternalBatchDeleteAckResp { + type Ok = ExternalBatchDeleteAckResp; + fn to_result(self) -> KvResult { + if self.error_code == OK { + Ok(self) + } else { + Err(KvError::from_json(self.error_code, &self.error_json)) + } + } +} + +impl ToResult for BatchIsExistResp { + type Ok = Vec; + fn to_result(self) -> KvResult { + if self.error_code == OK { + return Ok(self.exists_list); + } + Err(KvError::from_json(self.error_code, &self.error_json)) + } +} + /// Centralized: convert code+desc to KvError or Ok if code==Ok pub fn try_from_code(code: ErrorCode, json: String) -> KvResult<()> { if code == OK { @@ -222,6 +263,135 @@ impl FromError for ExternalIsExistResp { } } } +impl FromError for ExternalBatchIsExistResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchGetResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchGetItemResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchGetStartResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchGetTransferResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchGetCancelResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + } + } +} +impl FromError for ExternalBatchPutStartResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchPutStartItemResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchPutTransferEndResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchPutTransferEndItemResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchPutCommitResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalBatchPutCommitItemResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ExternalObservabilitySnapshotResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} impl FromError for ExternalDeleteAckResp { fn from_error(e: &KvError) -> Self { let code = e.code(); @@ -232,6 +402,16 @@ impl FromError for ExternalDeleteAckResp { } } } +impl FromError for ExternalBatchDeleteAckResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} // ---- FromError for Master KV Router Resps ---- impl FromError for GetStartResp { @@ -294,6 +474,76 @@ impl FromError for PutDoneResp { } } } +impl FromError for ReserveLocalGrantResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for ReleaseLocalGrantResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for BatchPreparePutKeysResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for BatchReleasePutKeyReservationsResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for PutAppendStartResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for PutAppendRevokeResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} +impl FromError for PutAppendDoneResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} impl FromError for MemHolderKeepAliveResp { fn from_error(e: &KvError) -> Self { let code = e.code(); @@ -354,6 +604,16 @@ impl FromError for GetMetaResp { } } } +impl FromError for BatchIsExistResp { + fn from_error(e: &KvError) -> Self { + let code = e.code(); + Self { + error_code: code, + error_json: e.to_json(), + ..Default::default() + } + } +} impl FromError for BatchDeleteClientKvMetaCacheResp { fn from_error(e: &KvError) -> Self { let code = e.code(); diff --git a/fluxon_rs/fluxon_kv/src/user_api/rpc.rs b/fluxon_rs/fluxon_kv/src/user_api/rpc.rs index eff064f..dfbda39 100644 --- a/fluxon_rs/fluxon_kv/src/user_api/rpc.rs +++ b/fluxon_rs/fluxon_kv/src/user_api/rpc.rs @@ -44,16 +44,7 @@ pub trait UserRpcServer: Send + Sync { pub type UserRpcFlatDictFuture = Pin> + Send + 'static>>; pub(crate) fn validate_timeout_ms(timeout_ms: u64) -> KvResult<()> { - if timeout_ms < user_rpc::USER_RPC_MIN_TIMEOUT_MS { - return Err(KvError::Api(ApiError::InvalidArgument { - detail: format!( - "timeout_ms must be >= {} (got {})", - user_rpc::USER_RPC_MIN_TIMEOUT_MS, - timeout_ms - ), - })); - } - Ok(()) + user_rpc::validate_timeout_ms(timeout_ms) } pub(crate) struct FluxonUserRpcImpl { diff --git a/fluxon_rs/fluxon_kv/src/user_rpc.rs b/fluxon_rs/fluxon_kv/src/user_rpc.rs index 7bff55b..e9587ac 100644 --- a/fluxon_rs/fluxon_kv/src/user_rpc.rs +++ b/fluxon_rs/fluxon_kv/src/user_rpc.rs @@ -2,10 +2,10 @@ use std::time::Duration; use std::time::Instant; use fluxon_commu::p2p::rpc::{ - USER_RPC_OBSERVE_TRACE_RAW_BYTES_INDEX, USER_RPC_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX, - UserRpcObserveTrace, UserRpcOwner1ObserveTrace, UserRpcTransportPathKind, - current_cross_process_monotonic_us, decode_user_rpc_observe_trace, - decode_user_rpc_owner1_observe_trace, + MIN_EXPLICIT_RPC_TIMEOUT_MS, USER_RPC_OBSERVE_TRACE_RAW_BYTES_INDEX, + USER_RPC_OWNER1_OBSERVE_TRACE_RAW_BYTES_INDEX, UserRpcObserveTrace, UserRpcOwner1ObserveTrace, + UserRpcTransportPathKind, current_cross_process_monotonic_us, decode_user_rpc_observe_trace, + decode_user_rpc_owner1_observe_trace, validate_explicit_rpc_timeout_ms, }; use prost::bytes::Bytes; @@ -15,7 +15,34 @@ use crate::p2p::p2p_module::UserRpcReq; use crate::rpcresp_kvresult_convert; use crate::rpcresp_kvresult_convert::msg_and_error::{ApiError, KvError, KvResult}; -pub const USER_RPC_MIN_TIMEOUT_MS: u64 = 10_000; +pub const USER_RPC_MIN_TIMEOUT_MS: u64 = MIN_EXPLICIT_RPC_TIMEOUT_MS; + +pub fn validate_timeout_ms(timeout_ms: u64) -> KvResult<()> { + validate_explicit_rpc_timeout_ms(timeout_ms).map_err(KvError::from) +} + +#[cfg(test)] +mod timeout_validation_tests { + use super::*; + use fluxon_commu::p2p::P2pError; + + #[test] + fn invalid_user_rpc_timeout_preserves_typed_p2p_error() { + let timeout_ms = USER_RPC_MIN_TIMEOUT_MS - 1; + let err = validate_timeout_ms(timeout_ms).unwrap_err(); + match err { + KvError::P2p(P2pError::InvalidRpcTimeout { + timeout_ms: actual_timeout_ms, + min_timeout_ms, + .. + }) => { + assert_eq!(actual_timeout_ms, timeout_ms); + assert_eq!(min_timeout_ms, USER_RPC_MIN_TIMEOUT_MS); + } + other => panic!("expected typed InvalidRpcTimeout, got {other:?}"), + } + } +} #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum UserRpcOwnerPathKind { @@ -140,14 +167,7 @@ pub async fn user_rpc_call_observed( payload: Vec, timeout_ms: u64, ) -> KvResult { - if timeout_ms < USER_RPC_MIN_TIMEOUT_MS { - return Err(KvError::Api(ApiError::InvalidArgument { - detail: format!( - "timeout_ms must be >= {} (got {})", - USER_RPC_MIN_TIMEOUT_MS, timeout_ms - ), - })); - } + validate_timeout_ms(timeout_ms)?; let path_for_error = path.clone(); let req = MsgPack { diff --git a/fluxon_rs/fluxon_ops/build.rs b/fluxon_rs/fluxon_ops/build.rs index 585fbfc..51e95c4 100644 --- a/fluxon_rs/fluxon_ops/build.rs +++ b/fluxon_rs/fluxon_ops/build.rs @@ -59,9 +59,17 @@ print( } fn render_log_shard_helper(repo_root: &Path) -> String { - let helper_path = repo_root.join("deployment").join("utils").join("log_shard.py"); - fs::read_to_string(&helper_path) - .unwrap_or_else(|e| panic!("read log shard helper failed: {} ({})", helper_path.display(), e)) + let helper_path = repo_root + .join("deployment") + .join("utils") + .join("log_shard.py"); + fs::read_to_string(&helper_path).unwrap_or_else(|e| { + panic!( + "read log shard helper failed: {} ({})", + helper_path.display(), + e + ) + }) } fn main() { @@ -87,6 +95,10 @@ fn main() { ); println!( "cargo:rerun-if-changed={}", - repo_root.join("deployment").join("utils").join("log_shard.py").display() + repo_root + .join("deployment") + .join("utils") + .join("log_shard.py") + .display() ); } diff --git a/fluxon_rs/fluxon_ops/src/lib.rs b/fluxon_rs/fluxon_ops/src/lib.rs index 29d9434..3adb053 100644 --- a/fluxon_rs/fluxon_ops/src/lib.rs +++ b/fluxon_rs/fluxon_ops/src/lib.rs @@ -80,7 +80,8 @@ const DELETE_APPLY_NO_WAIT_DELAY_SECONDS: u64 = 30; const EMBEDDED_SELECTION_SUPERVISOR_SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/selection_supervisor.py")); -const EMBEDDED_LOG_SHARD_HELPER_SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/log_shard.py")); +const EMBEDDED_LOG_SHARD_HELPER_SOURCE: &str = + include_str!(concat!(env!("OUT_DIR"), "/log_shard.py")); // Ops controller uses Fluxon user-RPC to talk to ops agents. // Keep the timeout as a fixed constant to avoid config surface area. @@ -351,7 +352,10 @@ fn workload_log_latest_shard_identity(logical_path: &Path) -> anyhow::Result anyhow::Result Ok(resolved) } -fn ensure_embedded_selection_supervisor_runtime(workdir: &Path) -> anyhow::Result<(PathBuf, PathBuf)> { +fn ensure_embedded_selection_supervisor_runtime( + workdir: &Path, +) -> anyhow::Result<(PathBuf, PathBuf)> { let runtime_dir = workdir.join(OPS_SELECTION_SUPERVISOR_DIR_NAME); std::fs::create_dir_all(&runtime_dir).with_context(|| { format!( @@ -1657,10 +1663,11 @@ fn selection_owner_supervisor( scope_key: Option<&str>, exclude_pid: Option, ) -> anyhow::Result> { - let owners: Vec = live_selection_supervisors(snapshot, Some(label), scope_key)? - .into_iter() - .filter(|supervisor| exclude_pid != Some(supervisor.pid())) - .collect(); + let owners: Vec = + live_selection_supervisors(snapshot, Some(label), scope_key)? + .into_iter() + .filter(|supervisor| exclude_pid != Some(supervisor.pid())) + .collect(); if owners.is_empty() { return Ok(None); } @@ -2068,7 +2075,16 @@ fn wait_for_selection_attached( argv: &[String], cwd: Option<&str>, ) -> anyhow::Result { - wait_for_selection_attached_for_scope(kind, name, authority, None, apply_id, owner_ts_ms, argv, cwd) + wait_for_selection_attached_for_scope( + kind, + name, + authority, + None, + apply_id, + owner_ts_ms, + argv, + cwd, + ) } fn wait_for_selection_attached_without_present_for_scope( @@ -2803,10 +2819,9 @@ impl SupervisorBackedWorkloads { fn list_workloads(&self) -> anyhow::Result> { let mut out: Vec = Vec::new(); let snapshot = selection_supervisor_proc_snapshot()?; - for status in observe_all_selection_statuses_for_snapshot( - &snapshot, - Some(self.scope_key.as_str()), - )? { + for status in + observe_all_selection_statuses_for_snapshot(&snapshot, Some(self.scope_key.as_str()))? + { let kind = status.kind.with_context(|| { format!( "selection supervisor list item missing kind: label={}", @@ -3159,7 +3174,10 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { None => { let Some(path) = resolve_readable_log_path(&logical_path) else { let resp = make_err_resp( - format!("log file is not available yet: logical_path={}", logical_path.display()), + format!( + "log file is not available yet: logical_path={}", + logical_path.display() + ), None, ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3186,138 +3204,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { }; let file_size = meta.len(); - let (start, end, start_cursor, end_cursor, effective_path, effective_file_size) = - match req.direction { - LogReadDirection::Forward => { - if let Some(cursor) = req.cursor.as_ref() { - if cursor.offset > file_size { - let resp = make_err_resp( - format!( - "cursor out of range: shard={} cursor={} file_size={}", - cursor.shard, cursor.offset, file_size - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - let mut effective_path = path.clone(); - let mut effective_shard = shard.clone(); - let mut effective_file_size = file_size; - let mut start = cursor.offset; - if cursor.offset == file_size { - if let Ok(Some(next_shard)) = - workload_log_next_shard(&logical_path, &cursor.shard) - { - let next_path = match workload_log_path_for_shard(&logical_path, &next_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&next_path) { - Ok(next_meta) => { - effective_file_size = next_meta.len(); - effective_path = next_path; - effective_shard = next_shard; - start = 0; - } - Err(e) => { - let resp = make_err_resp( - format!( - "stat next log shard failed: path={} err={}", - next_path.display(), - e - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - } - } else if let Ok(Some(latest_shard)) = - workload_log_latest_shard_identity(&logical_path) - { - if latest_shard != cursor.shard { - let latest_path = - match workload_log_path_for_shard(&logical_path, &latest_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&latest_path) { - Ok(latest_meta) => { - effective_file_size = latest_meta.len(); - effective_path = latest_path; - effective_shard = latest_shard; - start = 0; - } - Err(e) => { - let resp = make_err_resp( - format!( - "stat latest log shard failed: path={} err={}", - latest_path.display(), - e - ), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - } - } - } - } - let end = match max_bytes { - Some(max_bytes) => { - std::cmp::min(effective_file_size, start.saturating_add(max_bytes)) - } - None => effective_file_size, - }; - ( - start, - end, - Some(WorkloadLogCursor { - shard: effective_shard.clone(), - offset: start, - }), - Some(WorkloadLogCursor { - shard: effective_shard.clone(), - offset: end, - }), - effective_path, - effective_file_size, - ) - } else { - let end = file_size; - let start = match max_bytes { - Some(max_bytes) => end.saturating_sub(max_bytes), - None => 0, - }; - ( - start, - end, - Some(WorkloadLogCursor { - shard: shard.clone(), - offset: start, - }), - Some(WorkloadLogCursor { - shard: shard.clone(), - offset: end, - }), - path.clone(), - file_size, - ) - } - } - LogReadDirection::Backward => { - let Some(cursor) = req.cursor.as_ref() else { - let resp = make_err_resp( - "cursor is required for Backward reads".to_string(), - Some(file_size), - ); - return Ok(serde_json::to_vec(&resp).unwrap()); - }; + let (start, end, start_cursor, end_cursor, effective_path, effective_file_size) = match req + .direction + { + LogReadDirection::Forward => { + if let Some(cursor) = req.cursor.as_ref() { if cursor.offset > file_size { let resp = make_err_resp( format!( @@ -3331,30 +3222,31 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { let mut effective_path = path.clone(); let mut effective_shard = shard.clone(); let mut effective_file_size = file_size; - let mut end = cursor.offset; - if cursor.offset == 0 { - if let Ok(Some(prev_shard)) = - workload_log_previous_shard(&logical_path, &cursor.shard) + let mut start = cursor.offset; + if cursor.offset == file_size { + if let Ok(Some(next_shard)) = + workload_log_next_shard(&logical_path, &cursor.shard) { - let prev_path = match workload_log_path_for_shard(&logical_path, &prev_shard) { - Ok(v) => v, - Err(e) => { - let resp = make_err_resp(format!("{}", e), Some(file_size)); - return Ok(serde_json::to_vec(&resp).unwrap()); - } - }; - match std::fs::metadata(&prev_path) { - Ok(prev_meta) => { - effective_file_size = prev_meta.len(); - effective_path = prev_path; - effective_shard = prev_shard; - end = effective_file_size; + let next_path = + match workload_log_path_for_shard(&logical_path, &next_shard) { + Ok(v) => v, + Err(e) => { + let resp = make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&next_path) { + Ok(next_meta) => { + effective_file_size = next_meta.len(); + effective_path = next_path; + effective_shard = next_shard; + start = 0; } Err(e) => { let resp = make_err_resp( format!( - "stat previous log shard failed: path={} err={}", - prev_path.display(), + "stat next log shard failed: path={} err={}", + next_path.display(), e ), Some(file_size), @@ -3362,11 +3254,47 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { return Ok(serde_json::to_vec(&resp).unwrap()); } } + } else if let Ok(Some(latest_shard)) = + workload_log_latest_shard_identity(&logical_path) + { + if latest_shard != cursor.shard { + let latest_path = + match workload_log_path_for_shard(&logical_path, &latest_shard) + { + Ok(v) => v, + Err(e) => { + let resp = + make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&latest_path) { + Ok(latest_meta) => { + effective_file_size = latest_meta.len(); + effective_path = latest_path; + effective_shard = latest_shard; + start = 0; + } + Err(e) => { + let resp = make_err_resp( + format!( + "stat latest log shard failed: path={} err={}", + latest_path.display(), + e + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + } + } } } - let start = match max_bytes { - Some(max_bytes) => end.saturating_sub(max_bytes), - None => 0, + let end = match max_bytes { + Some(max_bytes) => { + std::cmp::min(effective_file_size, start.saturating_add(max_bytes)) + } + None => effective_file_size, }; ( start, @@ -3382,8 +3310,103 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { effective_path, effective_file_size, ) + } else { + let end = file_size; + let start = match max_bytes { + Some(max_bytes) => end.saturating_sub(max_bytes), + None => 0, + }; + ( + start, + end, + Some(WorkloadLogCursor { + shard: shard.clone(), + offset: start, + }), + Some(WorkloadLogCursor { + shard: shard.clone(), + offset: end, + }), + path.clone(), + file_size, + ) } - }; + } + LogReadDirection::Backward => { + let Some(cursor) = req.cursor.as_ref() else { + let resp = make_err_resp( + "cursor is required for Backward reads".to_string(), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + }; + if cursor.offset > file_size { + let resp = make_err_resp( + format!( + "cursor out of range: shard={} cursor={} file_size={}", + cursor.shard, cursor.offset, file_size + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + let mut effective_path = path.clone(); + let mut effective_shard = shard.clone(); + let mut effective_file_size = file_size; + let mut end = cursor.offset; + if cursor.offset == 0 { + if let Ok(Some(prev_shard)) = + workload_log_previous_shard(&logical_path, &cursor.shard) + { + let prev_path = + match workload_log_path_for_shard(&logical_path, &prev_shard) { + Ok(v) => v, + Err(e) => { + let resp = make_err_resp(format!("{}", e), Some(file_size)); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + }; + match std::fs::metadata(&prev_path) { + Ok(prev_meta) => { + effective_file_size = prev_meta.len(); + effective_path = prev_path; + effective_shard = prev_shard; + end = effective_file_size; + } + Err(e) => { + let resp = make_err_resp( + format!( + "stat previous log shard failed: path={} err={}", + prev_path.display(), + e + ), + Some(file_size), + ); + return Ok(serde_json::to_vec(&resp).unwrap()); + } + } + } + } + let start = match max_bytes { + Some(max_bytes) => end.saturating_sub(max_bytes), + None => 0, + }; + ( + start, + end, + Some(WorkloadLogCursor { + shard: effective_shard.clone(), + offset: start, + }), + Some(WorkloadLogCursor { + shard: effective_shard.clone(), + offset: end, + }), + effective_path, + effective_file_size, + ) + } + }; if end < start { let resp = make_err_resp( @@ -3416,7 +3439,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { Ok(v) => v, Err(e) => { let resp = make_err_resp( - format!("open log failed: path={} err={}", effective_path.display(), e), + format!( + "open log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3425,7 +3452,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { if let Err(e) = std::io::Seek::seek(&mut f, std::io::SeekFrom::Start(start)) { let resp = make_err_resp( - format!("seek log failed: path={} err={}", effective_path.display(), e), + format!( + "seek log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -3434,7 +3465,11 @@ impl UserRpcHandler for ReadWorkloadLogChunkHandler { let mut buf: Vec = vec![0; len]; if let Err(e) = std::io::Read::read_exact(&mut f, &mut buf) { let resp = make_err_resp( - format!("read log failed: path={} err={}", effective_path.display(), e), + format!( + "read log failed: path={} err={}", + effective_path.display(), + e + ), Some(effective_file_size), ); return Ok(serde_json::to_vec(&resp).unwrap()); @@ -4067,8 +4102,7 @@ fn desired_workload_matches_running( &desired.name, &desired.authority, Some(workloads.scope_key.as_str()), - ) - else { + ) else { return false; }; desired_workload_status_matches_goal(&status, desired) @@ -14337,14 +14371,14 @@ mod tests { assert_eq!(scoped_b.len(), 1); assert_eq!(scoped_b[0].pid(), 22); - let listed_a = observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-a")) - .unwrap(); + let listed_a = + observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-a")).unwrap(); assert_eq!(listed_a.len(), 1); assert_eq!(listed_a[0].label, "DaemonSet/target"); assert_eq!(listed_a[0].pid, Some(11)); - let listed_b = observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-b")) - .unwrap(); + let listed_b = + observe_all_selection_statuses_for_snapshot(&snapshot, Some("/tmp/scope-b")).unwrap(); assert_eq!(listed_b.len(), 1); assert_eq!(listed_b[0].label, "DaemonSet/target"); assert_eq!(listed_b[0].pid, Some(22)); @@ -14548,8 +14582,8 @@ mod tests { zombie_infos: Vec::new(), }; - let listed = observe_apply_runtime_statuses_for_snapshot("apply-1", &snapshot, None) - .unwrap(); + let listed = + observe_apply_runtime_statuses_for_snapshot("apply-1", &snapshot, None).unwrap(); assert_eq!(listed.len(), 1); assert_eq!(listed[0].name.as_deref(), Some("target-present")); assert!(listed[0].present); @@ -14774,12 +14808,8 @@ mod tests { None )); - let delete_old = workloads.delete_generation( - WorkloadKind::Deployment, - &name, - &name, - Some("apply-1"), - ); + let delete_old = + workloads.delete_generation(WorkloadKind::Deployment, &name, &name, Some("apply-1")); if !delete_old.ok { let err = delete_old.err.as_deref().unwrap_or_default(); assert!( @@ -14813,8 +14843,7 @@ mod tests { delete_current.ok, "unguarded delete should bind and retire the current visible generation: {delete_current:?}" ); - wait_for_selection_absent(WorkloadKind::Deployment, &name, &name, Some("apply-2")) - .unwrap(); + wait_for_selection_absent(WorkloadKind::Deployment, &name, &name, Some("apply-2")).unwrap(); } #[test] @@ -14826,9 +14855,12 @@ mod tests { python_exe.display() ); let workdir = tempfile::tempdir().unwrap(); - let runtime = - SelectionSupervisorRuntime::materialize(workdir.path(), workdir.path(), python_exe.as_path()) - .unwrap(); + let runtime = SelectionSupervisorRuntime::materialize( + workdir.path(), + workdir.path(), + python_exe.as_path(), + ) + .unwrap(); assert!(runtime.script_path.exists()); assert!( runtime @@ -14849,9 +14881,12 @@ mod tests { python_exe.display() ); let workdir = tempfile::tempdir().unwrap(); - let runtime = - SelectionSupervisorRuntime::materialize(workdir.path(), workdir.path(), python_exe.as_path()) - .unwrap(); + let runtime = SelectionSupervisorRuntime::materialize( + workdir.path(), + workdir.path(), + python_exe.as_path(), + ) + .unwrap(); let log_path = workdir.path().join("startup.log"); let command = vec![ python_exe.display().to_string(), @@ -14876,7 +14911,9 @@ mod tests { "--".to_string(), "/bin/true".to_string(), ]; - let pid = runtime.spawn_detached_command(&log_path, command.as_slice()).unwrap(); + let pid = runtime + .spawn_detached_command(&log_path, command.as_slice()) + .unwrap(); let deadline = Instant::now() + Duration::from_secs(10); let expected = "owner-ts-ms must be positive"; let mut saw_expected = false; @@ -15164,7 +15201,9 @@ mod tests { }), max_bytes: Some(65536), }; - let raw = handler.handle("n1".into(), &serde_json::to_vec(&req).unwrap()).unwrap(); + let raw = handler + .handle("n1".into(), &serde_json::to_vec(&req).unwrap()) + .unwrap(); let resp: ReadWorkloadLogResp = serde_json::from_slice(&raw).unwrap(); assert!(resp.ok, "{resp:?}"); assert_eq!(resp.text.as_deref(), Some("new\n")); @@ -15209,7 +15248,9 @@ mod tests { }), max_bytes: Some(65536), }; - let raw = handler.handle("n1".into(), &serde_json::to_vec(&req).unwrap()).unwrap(); + let raw = handler + .handle("n1".into(), &serde_json::to_vec(&req).unwrap()) + .unwrap(); let resp: ReadWorkloadLogResp = serde_json::from_slice(&raw).unwrap(); assert!(resp.ok, "{resp:?}"); assert_eq!(resp.text.as_deref(), Some("old\n")); diff --git a/fluxon_rs/fluxon_pyo3/src/error.rs b/fluxon_rs/fluxon_pyo3/src/error.rs index 97ab680..3848124 100644 --- a/fluxon_rs/fluxon_pyo3/src/error.rs +++ b/fluxon_rs/fluxon_pyo3/src/error.rs @@ -221,6 +221,23 @@ pub(crate) fn new_key_being_written_error( error_class.call((), Some(&kwargs)).unwrap().into() } +pub(crate) fn new_key_already_exists_error( + py: Python<'_>, + message: &str, + key: Option<&str>, +) -> PyObject { + let api_error_module = py.import_bound("fluxon_py.api_error").unwrap(); + let error_class = api_error_module.getattr("KeyAlreadyExistsError").unwrap(); + + let kwargs = PyDict::new_bound(py); + kwargs.set_item("message", message).unwrap(); + if let Some(k) = key { + kwargs.set_item("key", k).unwrap(); + } + + error_class.call((), Some(&kwargs)).unwrap().into() +} + pub(crate) fn new_storage_full_error( py: Python<'_>, message: &str, @@ -352,6 +369,7 @@ pub(crate) fn new_payload_lease_not_found_error( /// - TransferEngine::TransferFailedForBlock -> TransferBlockFailedError（可重试） /// - Api::InvalidPutMasterState -> PutDoneFailedError /// - Api::KeyNotFound -> KeyNotFoundError（携带 key） +/// - Api::KeyAlreadyExists -> KeyAlreadyExistsError（携带 key） /// - Api::KeyBeingWritten -> KeyBeingWrittenError（携带 key） /// - Api::NoSpace -> StorageFullError（available_space=free_capacity） /// - 其他 -> NetworkError（携带格式化消息） @@ -361,7 +379,8 @@ pub(crate) fn py_error_from_kv_error( prefix: &str, ) -> PyObject { use fluxon_kv::rpcresp_kvresult_convert::msg_and_error::{ - ApiError as CoreApiError, KvError, LeaseMgrError, TransferEngineError, + ApiError as CoreApiError, KvError, LeaseMgrError, P2pError as CoreP2pError, + TransferEngineError, }; let msg = format!("{}: {}", prefix, e); match e { @@ -381,6 +400,11 @@ pub(crate) fn py_error_from_kv_error( &format!("{}: Key not found: {}", prefix, key), Some(key), ), + KvError::Api(CoreApiError::KeyAlreadyExists { key }) => new_key_already_exists_error( + py, + &format!("{}: Key already exists: {}", prefix, key), + Some(key), + ), KvError::Api(CoreApiError::KeyBeingWritten { key }) => new_key_being_written_error( py, &format!("{}: Key is currently being written: {}", prefix, key), @@ -389,6 +413,9 @@ pub(crate) fn py_error_from_kv_error( KvError::Api(CoreApiError::InvalidArgument { detail }) => { new_invalid_argument_error(py, &format!("{}: Invalid argument: {}", prefix, detail)) } + KvError::P2p(CoreP2pError::InvalidRpcTimeout { .. }) => { + new_invalid_argument_error(py, &msg) + } KvError::Api(CoreApiError::FileWriteError { path, offset, diff --git a/fluxon_rs/fluxon_pyo3/src/fixed_slab_allocator.rs b/fluxon_rs/fluxon_pyo3/src/fixed_slab_allocator.rs new file mode 100644 index 0000000..bbcd7df --- /dev/null +++ b/fluxon_rs/fluxon_pyo3/src/fixed_slab_allocator.rs @@ -0,0 +1,52 @@ +use fluxon_util::fixed_slab_allocator::FixedSlabAllocator as RustFixedSlabAllocator; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +/// Thread-safe fixed-slot slab allocator. +/// +/// `try_reserve` is all-or-none. `release` rejects invalid input atomically. +#[pyclass(module = "fluxon_pyo3")] +pub struct FixedSlabAllocator { + inner: RustFixedSlabAllocator, +} + +#[pymethods] +impl FixedSlabAllocator { + #[new] + fn new(slot_count: u32) -> PyResult { + let inner = RustFixedSlabAllocator::new(slot_count) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(Self { inner }) + } + + /// Reserve exactly `count` slots, returning `None` if capacity is unavailable. + fn try_reserve(&self, count: u32) -> Option> { + self.inner.try_reserve(count) + } + + /// Release all slots atomically after bounds, duplicate, and live-state checks. + fn release(&self, slots: Vec) -> PyResult<()> { + self.inner + .release(&slots) + .map_err(|error| PyValueError::new_err(error.to_string())) + } + + #[getter] + fn capacity(&self) -> u32 { + self.inner.capacity() + } + + #[getter] + fn free_count(&self) -> u32 { + self.inner.free_count() + } + + #[getter] + fn live_count(&self) -> u32 { + self.inner.live_count() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} diff --git a/fluxon_rs/fluxon_pyo3/src/lib.rs b/fluxon_rs/fluxon_pyo3/src/lib.rs index 631f1ea..0165f9f 100644 --- a/fluxon_rs/fluxon_pyo3/src/lib.rs +++ b/fluxon_rs/fluxon_pyo3/src/lib.rs @@ -1,9 +1,9 @@ -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::panic::{self, AssertUnwindSafe}; use std::path::{Path, PathBuf}; use std::slice; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; use std::thread; use bytes::Bytes; @@ -14,11 +14,31 @@ use fluxon_fs::config::{ FS_MASTER_CONFIG_RPC_PATH, FluxonFsRequestIdentity, FluxonFsS3KvMissPolicy, FsAgentDeclaredExportWire, }; +use fluxon_kv::OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES; use fluxon_kv::client_kv_api::ClientKvApiViewTrait; +use fluxon_kv::client_kv_api::msg_pack::{ + ExternalBatchPutCommitItemReq, ExternalBatchPutCommitReq, ExternalBatchPutStartItemReq, + ExternalBatchPutStartReq, ExternalBatchPutTransferEndItemReq, ExternalBatchPutTransferEndReq, + ExternalPutRevokeReq, +}; +use fluxon_kv::client_kv_api::{ + ClientKvApiInner, OwnerLocalPublishItem, OwnerLocalPublishJob, OwnerLocalReserveSlotLease, +}; +use fluxon_kv::client_seg_pool::ClientSegPoolViewTrait; +use fluxon_kv::client_transfer_engine::ClientTransferEngineAccessTrait; use fluxon_kv::cluster_manager::ClusterManagerViewTrait; use fluxon_kv::cluster_manager::app_logic_ext::ClusterManagerAppLogicExt; use fluxon_kv::config::{ClientConfigYaml, MasterConfigYaml}; +use fluxon_kv::external_client_api::ExternalClientApiViewTrait; +use fluxon_kv::master_kv_router::msg_pack::{ + BatchPreparePutKeyItemReq, PutDoneCommittedSlot, ReserveLocalGrantOutcome, + build_put_atomic_group_assignments, +}; use fluxon_kv::master_lease_manager::msg_pack::{AllocateClientLeaseReq, ClientLeaseKeepaliveReq}; +use fluxon_kv::memholder::kvclient_encode::{BorrowedFlatKvValueRange, flat_kv_decode_borrowed}; +use fluxon_kv::memholder::{ + ExternalMemHolder as RustExternalMemHolder, UserMemHolder as RustUserMemHolder, +}; use fluxon_kv::p2p::msg_pack::{MsgPack, RPCCaller, call_rpc}; use fluxon_kv::p2p::p2p_module::P2pModuleViewTrait; use fluxon_kv::p2p::p2p_module::{UserRpcHandler, user_rpc_register_handler}; @@ -40,7 +60,9 @@ use fluxon_util::{ FluxonCliProxyDescriptorV2, FluxonCliProxyTransportV2, fluxon_cli_proxy_desc_etcd_key_v2, }; use futures::Future; -use pyo3::exceptions::{PyOSError, PyPermissionError, PyRuntimeError, PyValueError}; +use pyo3::exceptions::{ + PyKeyboardInterrupt, PyOSError, PyPermissionError, PyRuntimeError, PyValueError, +}; use pyo3::prelude::*; use pyo3::pybacked::PyBackedBytes; use pyo3::types::{PyAny, PyBytes, PyDict, PyList, PyModule, PyString, PyTuple}; @@ -57,7 +79,10 @@ include!(env!("FLUXON_PYO3_TEST_PYTHON_LINK_RS")); mod memholder; pub use memholder::{ExternalMemHolder, MemHolder}; +mod fixed_slab_allocator; +pub use fixed_slab_allocator::FixedSlabAllocator; mod flatdict_zerocopy; +use flatdict_zerocopy::{FlatDictDataOwner, decode_flat_dict_to_wrapped_py_object}; mod kvfuture; pub use kvfuture::KvFuture; mod error; @@ -1930,6 +1955,415 @@ fn new_store_closed_error(py: Python, message: &str) -> PyObject { crate::error::new_store_closed_error(py, message) } +#[derive(Clone, Copy, Debug)] +struct RegisteredBufferRange { + start: usize, + end: usize, +} + +impl RegisteredBufferRange { + fn contains(&self, ptr: usize, len: usize) -> bool { + if len == 0 { + return ptr >= self.start && ptr <= self.end; + } + let Some(req_end) = ptr.checked_add(len) else { + return false; + }; + ptr >= self.start && req_end <= self.end + } +} + +#[derive(Default, Debug)] +struct RegisteredBufferRegistry { + ranges: Vec, +} + +impl RegisteredBufferRegistry { + fn register(&mut self, ptr: usize, len: usize) { + let end = ptr.saturating_add(len); + self.ranges.push(RegisteredBufferRange { start: ptr, end }); + } + + fn contains(&self, ptr: usize, len: usize) -> bool { + self.ranges.iter().any(|range| range.contains(ptr, len)) + } +} + +const DEFAULT_PYO3_BATCH_CONCURRENCY: usize = 128; + +fn kv_error_to_ret_code(err: &CoreKvError) -> i32 { + match err { + CoreKvError::Config(inner) => -(inner.code() as i32), + CoreKvError::Api(inner) => -(inner.code() as i32), + CoreKvError::P2p(inner) => -(inner.code() as i32), + CoreKvError::SharedMem(inner) => -(inner.code() as i32), + CoreKvError::Unreachable(inner) => -(inner.code() as i32), + CoreKvError::ClusterManagerExt(inner) => -(inner.code() as i32), + CoreKvError::Metric(inner) => -(inner.code() as i32), + CoreKvError::LeaseMgr(inner) => -(inner.code() as i32), + CoreKvError::TransferEngine(inner) => -(inner.code() as i32), + CoreKvError::P2pTransfer(inner) => -(inner.code() as i32), + } +} + +fn rpc_item_error_to_ret_code(error_code: u32, error_json: &str) -> i32 { + if error_code == OK { + 0 + } else { + let err = CoreKvError::from_json(error_code, error_json); + kv_error_to_ret_code(&err) + } +} + +fn payload_slice_from_get_result<'a>(key: &str, get_result: &'a KvGetResult) -> KvResult<&'a [u8]> { + let data = match get_result { + KvGetResult::Owner(Some(holder)) => holder.bytes(), + KvGetResult::External(Some(holder)) => holder.bytes(), + KvGetResult::Owner(None) | KvGetResult::External(None) => { + return Err(CoreKvError::Api(CoreApiError::KeyNotFound { + key: key.to_string(), + })); + } + }; + + let entries = flat_kv_decode_borrowed(data).map_err(|detail| { + CoreKvError::Api(CoreApiError::Unknown { + detail: format!("decode flat dict for key {} failed: {}", key, detail), + }) + })?; + for entry in entries { + if entry.key != "payload" { + continue; + } + if let BorrowedFlatKvValueRange::BytesRange { start, len } = entry.value { + let end = start.checked_add(len).ok_or_else(|| { + CoreKvError::Api(CoreApiError::Unknown { + detail: format!("payload range overflow for key {}", key), + }) + })?; + if end > data.len() { + return Err(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "payload range out of bounds for key {}: end={} data_len={}", + key, + end, + data.len() + ), + })); + } + return Ok(&data[start..end]); + } + return Err(CoreKvError::Api(CoreApiError::InvalidArgument { + detail: format!("key {} payload field is not bytes", key), + })); + } + + Err(CoreKvError::Api(CoreApiError::InvalidArgument { + detail: format!("key {} does not contain payload field", key), + })) +} + +const FLUXON_PLAN_BLOB_MAGIC: u64 = 0x4658_504c_414e_5631; + +#[derive(Clone)] +enum FluxonPlanRegistryEntry { + Put(Arc), + Get(Arc), +} + +enum FluxonPutPlanState { + Prepared(StagedPutPlanData), + Committing, +} + +struct FluxonPutPlan { + blob: Box<[u64]>, + state: Mutex, +} + +struct FluxonGetViewsPlan { + blob: Box<[u64]>, + _holders: Vec, +} + +fn defer_drop_to_runtime(runtime: Option<&Runtime>, value: T) +where + T: Send + 'static, +{ + if let Some(runtime) = runtime { + // Dropping a read plan can release hundreds of ExternalMemHolder + // instances. Each holder schedules its delete ACK, so doing the full + // drop on the Python scheduler thread adds milliseconds to every + // restore completion. The CUDA consumer has already finished before + // release_views is called; removing the plan from the registry makes + // the pointer invalid immediately, while the actual holder release can + // safely run on Tokio's blocking pool. + let _join = runtime.spawn_blocking(move || drop(value)); + } else { + // A post-close caller has no runtime on which to defer cleanup. + drop(value); + } +} + +enum StagedGetViewHolder { + Owner { _holder: Arc }, + External { _holder: Arc }, +} + +fn external_get_start_result_to_py( + started: &fluxon_kv::external_client_api::ExternalClientGetStartResp, + py: Python, +) -> PyObject { + let out = PyDict::new_bound(py); + out.set_item("handle", started.handle).expect("set handle"); + out.set_item("raw_prefix_hit_len", started.raw_prefix_hit_len as u64) + .expect("set raw_prefix_hit_len"); + out.into_py(py) +} + +#[derive(Clone)] +enum StagedPutPlanData { + Owner(StagedOwnerPutPlanData), + External(StagedExternalPutPlanData), +} + +#[derive(Clone)] +struct StagedOwnerPutPlanData { + keys: Vec, + value_len: u64, + key_reservation_ids: Vec, + slot_lease: OwnerLocalReserveSlotLease, + make_replica_task_mask: Vec, + atomic_group_lens: Vec, +} + +#[derive(Clone)] +struct StagedExternalPutPlanData { + keys: Vec, + value_len: u64, + started_time: i64, + items: Vec, + short_circuit_payload: bool, +} + +#[derive(Clone)] +struct StagedExternalPutItem { + put_id: fluxon_kv::master_kv_router::put::PutIDForAKey, + src_offset: u64, + target_offset: u64, + transfer_target_offset: Option, + peer_id: Option, + target_base_addr: u64, +} + +#[derive(Clone, Debug, Default)] +struct FluxonIoLocalitySnapshot { + op_count: u64, + bytes: u64, + transfer_us: u64, +} + +impl FluxonIoLocalitySnapshot { + fn bandwidth_gbps(&self) -> f64 { + if self.transfer_us == 0 { + return 0.0; + } + (self.bytes as f64) / (self.transfer_us as f64 / 1_000_000.0) / 1_000_000_000.0 + } +} + +#[derive(Clone, Debug, Default)] +struct FluxonLocalitySnapshot { + l2_local_hit_pages: u64, + l2_local_hit_bytes: u64, + l2_remote_hit_pages: u64, + l2_remote_hit_bytes: u64, + put_local: FluxonIoLocalitySnapshot, + put_remote: FluxonIoLocalitySnapshot, + get_local: FluxonIoLocalitySnapshot, + get_remote: FluxonIoLocalitySnapshot, +} + +#[derive(Default)] +struct FluxonLocalityCounters { + l2_local_hit_pages: AtomicU64, + l2_local_hit_bytes: AtomicU64, + l2_remote_hit_pages: AtomicU64, + l2_remote_hit_bytes: AtomicU64, + put_local_ops: AtomicU64, + put_local_bytes: AtomicU64, + put_local_transfer_us: AtomicU64, + put_remote_ops: AtomicU64, + put_remote_bytes: AtomicU64, + put_remote_transfer_us: AtomicU64, + get_local_ops: AtomicU64, + get_local_bytes: AtomicU64, + get_local_transfer_us: AtomicU64, + get_remote_ops: AtomicU64, + get_remote_bytes: AtomicU64, + get_remote_transfer_us: AtomicU64, +} + +impl FluxonLocalityCounters { + fn snapshot(&self) -> FluxonLocalitySnapshot { + FluxonLocalitySnapshot { + l2_local_hit_pages: self.l2_local_hit_pages.load(Ordering::Relaxed), + l2_local_hit_bytes: self.l2_local_hit_bytes.load(Ordering::Relaxed), + l2_remote_hit_pages: self.l2_remote_hit_pages.load(Ordering::Relaxed), + l2_remote_hit_bytes: self.l2_remote_hit_bytes.load(Ordering::Relaxed), + put_local: FluxonIoLocalitySnapshot { + op_count: self.put_local_ops.load(Ordering::Relaxed), + bytes: self.put_local_bytes.load(Ordering::Relaxed), + transfer_us: self.put_local_transfer_us.load(Ordering::Relaxed), + }, + put_remote: FluxonIoLocalitySnapshot { + op_count: self.put_remote_ops.load(Ordering::Relaxed), + bytes: self.put_remote_bytes.load(Ordering::Relaxed), + transfer_us: self.put_remote_transfer_us.load(Ordering::Relaxed), + }, + get_local: FluxonIoLocalitySnapshot { + op_count: self.get_local_ops.load(Ordering::Relaxed), + bytes: self.get_local_bytes.load(Ordering::Relaxed), + transfer_us: self.get_local_transfer_us.load(Ordering::Relaxed), + }, + get_remote: FluxonIoLocalitySnapshot { + op_count: self.get_remote_ops.load(Ordering::Relaxed), + bytes: self.get_remote_bytes.load(Ordering::Relaxed), + transfer_us: self.get_remote_transfer_us.load(Ordering::Relaxed), + }, + } + } + + fn record_l2_hit(&self, remote: bool, bytes: u64) { + if remote { + self.l2_remote_hit_pages.fetch_add(1, Ordering::Relaxed); + self.l2_remote_hit_bytes.fetch_add(bytes, Ordering::Relaxed); + } else { + self.l2_local_hit_pages.fetch_add(1, Ordering::Relaxed); + self.l2_local_hit_bytes.fetch_add(bytes, Ordering::Relaxed); + } + } + + fn record_put(&self, remote: bool, bytes: u64, transfer_us: i64) { + let transfer_us = transfer_us.max(0) as u64; + if remote { + self.put_remote_ops.fetch_add(1, Ordering::Relaxed); + self.put_remote_bytes.fetch_add(bytes, Ordering::Relaxed); + self.put_remote_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } else { + self.put_local_ops.fetch_add(1, Ordering::Relaxed); + self.put_local_bytes.fetch_add(bytes, Ordering::Relaxed); + self.put_local_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } + } + + fn record_get(&self, remote: bool, bytes: u64, transfer_us: i64) { + let transfer_us = transfer_us.max(0) as u64; + if remote { + self.get_remote_ops.fetch_add(1, Ordering::Relaxed); + self.get_remote_bytes.fetch_add(bytes, Ordering::Relaxed); + self.get_remote_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } else { + self.get_local_ops.fetch_add(1, Ordering::Relaxed); + self.get_local_bytes.fetch_add(bytes, Ordering::Relaxed); + self.get_local_transfer_us + .fetch_add(transfer_us, Ordering::Relaxed); + } + } +} + +fn build_plan_blob(value_ptrs: &[u64]) -> Box<[u64]> { + let mut words = Vec::with_capacity(2 + value_ptrs.len()); + words.push(FLUXON_PLAN_BLOB_MAGIC); + words.push(value_ptrs.len() as u64); + words.extend_from_slice(value_ptrs); + words.into_boxed_slice() +} + +fn plan_ptr_from_blob(blob: &[u64]) -> usize { + blob.as_ptr() as usize +} + +fn plan_ptr_u64_to_usize(plan_ptr: u64, py: Python) -> Result { + let plan_ptr_usize = usize::try_from(plan_ptr) + .map_err(|_| new_invalid_argument_error(py, "plan_ptr out of range"))?; + if plan_ptr_usize == 0 { + return Err(new_invalid_argument_error(py, "plan_ptr must be non-zero")); + } + Ok(plan_ptr_usize) +} + +fn cleanup_plan_registry_entry( + registry: &Arc>>, + plan_ptr: usize, +) { + let mut guard = registry.write().expect("plan_registry poisoned"); + guard.remove(&plan_ptr); +} + +async fn release_staged_put_resources( + inner: &ClientKvApiInner, + key_reservation_ids: Vec, + slot_lease: OwnerLocalReserveSlotLease, +) -> KvResult<()> { + let release_keys_result = inner + .batch_release_put_key_reservations(key_reservation_ids) + .await; + let release_slot_result = inner + .owner_release_local_reserve_slot_lease(slot_lease) + .await; + match (release_keys_result, release_slot_result) { + (Ok(_), Ok(_)) => Ok(()), + (Err(err), Ok(_)) => Err(err), + (Ok(_), Err(err)) => Err(err), + (Err(primary_err), Err(cleanup_err)) => Err(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "release_staged_put_resources failed twice: batch_release_put_key_reservations_err={} release_local_slot_lease_err={}", + primary_err, cleanup_err + ), + })), + } +} + +async fn release_put_key_reservations_only( + inner: &ClientKvApiInner, + key_reservation_ids: Vec, +) -> KvResult<()> { + inner + .batch_release_put_key_reservations(key_reservation_ids) + .await + .map(|_| ()) +} + +fn combine_kv_errors( + context: &str, + primary_err: CoreKvError, + secondary_err: CoreKvError, +) -> CoreKvError { + CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "{}: primary_err={} secondary_err={}", + context, primary_err, secondary_err + ), + }) +} + +#[pyfunction] +fn decode_flat_dict_payload(payload: Vec, py: Python<'_>) -> PyObject { + match decode_flat_dict_to_wrapped_py_object(py, FlatDictDataOwner::from_owned_bytes(payload)) { + Ok(value) => ApiResult::new_success(value).into_py_object(py), + Err(err) => ApiResult::::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "flat dict decode failed", + )) + .into_py_object(py), + } +} + #[pyfunction] fn monitor_render_cli(config_path: String, workdir: String) -> PyResult { let cfg_yaml = MonitorCliConfigYaml::from_file(std::path::Path::new(&config_path)) @@ -2676,6 +3110,108 @@ pub struct KvClient { // - Futures must not hold Arc; they should spawn via Handle clones only. runtime: Option, config: ClientConfig, + registered_buffers: RwLock, + plan_registry: Arc>>, + locality_counters: Arc, +} + +fn build_put_optional_args( + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, +) -> fluxon_kv::client_kv_api::PutOptionalArgs { + let mut opts = fluxon_kv::client_kv_api::PutOptionalArgs::new(); + if let Some(id) = lease_id { + opts.0 + .push(fluxon_kv::client_kv_api::PutOptionalArg::LeaseId(id)); + } + if reject_if_inflight_same_key { + opts.0 + .push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfInflightSameKey); + } + if reject_if_exist_same_key { + opts.0 + .push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfExistSameKey); + } + if write_through || !make_replica_task { + opts.0 + .push(fluxon_kv::client_kv_api::PutOptionalArg::SkipMakeReplicaTask); + } + opts +} + +fn resolve_make_replica_task_mask( + key_count: usize, + write_through: bool, + make_replica_task: bool, + requested_mask: Option>, + atomic_group_lens: &[usize], +) -> Result, String> { + let requested_mask = match requested_mask { + Some(mask) => { + if mask.len() != key_count { + return Err(format!( + "make_replica_task_mask length must match keys length; keys={} mask={}", + key_count, + mask.len() + )); + } + mask + } + None => vec![true; key_count], + }; + let mut offset = 0usize; + for (group_index, group_len) in atomic_group_lens.iter().copied().enumerate() { + let end = offset + group_len; + let group_mask = &requested_mask[offset..end]; + if group_mask[1..].iter().any(|item| *item != group_mask[0]) { + return Err(format!( + "make_replica_task_mask must be uniform within each atomic group; group_index={} offset={} len={}", + group_index, offset, group_len + )); + } + offset = end; + } + let enabled = !write_through && make_replica_task; + Ok(requested_mask + .into_iter() + .map(|requested| enabled && requested) + .collect()) +} + +fn normalize_put_atomic_group_lens( + key_count: usize, + requested_group_lens: Option>, +) -> Result, String> { + let Some(group_lens) = requested_group_lens else { + return Ok(vec![1; key_count]); + }; + if group_lens.is_empty() { + return Err("atomic_group_lens must be non-empty when provided".to_string()); + } + if let Some((index, _)) = group_lens + .iter() + .enumerate() + .find(|(_, length)| **length == 0) + { + return Err(format!( + "atomic_group_lens entries must be > 0; index={} got=0", + index + )); + } + let group_sum = group_lens + .iter() + .try_fold(0usize, |sum, length| sum.checked_add(*length)) + .ok_or_else(|| "atomic_group_lens sum overflowed usize".to_string())?; + if group_sum != key_count { + return Err(format!( + "atomic_group_lens must sum to keys length; sum={} keys={}", + group_sum, key_count + )); + } + Ok(group_lens) } #[pymethods] @@ -2751,6 +3287,9 @@ impl KvClient { framework: Some(framework), runtime: Some(runtime), config: final_config, + registered_buffers: RwLock::new(RegisteredBufferRegistry::default()), + plan_registry: Arc::new(RwLock::new(HashMap::new())), + locality_counters: Arc::new(FluxonLocalityCounters::default()), }; match Py::new(py, client) { @@ -2800,28 +3339,9 @@ impl KvClient { self.config.cluster_name.clone() } - /// Allocate a fluxon-kv lease id synchronously. - /// Always allocate a new lease id (no reuse by requested id). - /// Allocate with the provided TTL seconds (must be >= MIN_CLIENT_TTL_SECONDS). - #[pyo3(signature = (ttl_seconds))] - fn allocate_lease(&self, ttl_seconds: u64, py: Python) -> PyObject { - fn allocate_lease_inner( - client: &KvClient, - ttl_seconds: u64, - py: Python, - ) -> ApiResult { - // Enforce minimum TTL at the PyO3 boundary so obvious mistakes fail fast. - if ttl_seconds - < fluxon_kv::master_lease_manager::MasterLeaseManager::MIN_CLIENT_TTL_SECONDS - { - return ApiResult::new_error(new_invalid_argument_error( - py, - &format!( - "allocate_lease(ttl_seconds) requires ttl_seconds >= {} seconds", - fluxon_kv::master_lease_manager::MasterLeaseManager::MIN_CLIENT_TTL_SECONDS, - ), - )); - } + /// Wait until the local segment mapping is ready for direct hostless access. + fn wait_local_segments_ready(&self, py: Python) -> PyObject { + fn wait_local_segments_ready_inner(client: &KvClient, py: Python) -> ApiResult { let framework = match require_kv_framework_api(client, py) { Ok(v) => v, Err(e) => return ApiResult::new_error(e), @@ -2835,64 +3355,121 @@ impl KvClient { )); } }; - // Blocking call on the client's runtime; simple and predictable for callers. - let r: Result = match py.allow_threads(|| { + + if framework.is_external_mode() { + let framework_for_wait = framework.clone(); + let ready_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_wait + .external_client_api_view() + .external_client_api() + .inner() + .wait_current_owner_mapped_range() + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + return match ready_result { + Ok((node_id, owner_start_time, write_ptr, read_ptr, len)) => { + let item = PyDict::new_bound(py); + item.set_item("segment_label", "external_owner:0") + .expect("set segment_label"); + item.set_item("write_ptr", write_ptr) + .expect("set write_ptr"); + item.set_item("read_ptr", read_ptr).expect("set read_ptr"); + item.set_item("len", len).expect("set len"); + item.set_item("generation", owner_start_time) + .expect("set generation"); + item.set_item("node_id", node_id).expect("set node_id"); + let out = PyList::empty_bound(py); + out.append(item).expect("append external owner segment"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "wait_local_segments_ready failed", + )), + }; + } + + let self_info = framework + .cluster_manager_view() + .cluster_manager() + .get_self_info(); + let framework_for_wait = framework.clone(); + let ready_result = match py.allow_threads(|| { runtime.run_async_from_sync(async move { - framework - .kv_allocate_lease(ttl_seconds) + framework_for_wait + .client_seg_pool_view() + .client_seg_pool() + .mapped_range() .await - .map_err(|e| e.to_string()) + .ok_or_else(|| { + CoreKvError::Api(CoreApiError::Unknown { + detail: "local owner segment is not mounted".to_string(), + }) + }) }) }) { Ok(v) => v, - Err(e) => Err(format!("runtime bridge failed: {}", e)), + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } }; - match r { - Ok(id) => ApiResult::new_success(Python::with_gil(|py| id.into_py(py))), - Err(e) => ApiResult::new_error(Python::with_gil(|py| { - new_network_error(py, &format!("Allocate lease failed: {}", e), None) - })), + match ready_result { + Ok((write_ptr, read_ptr, len)) => { + let item = PyDict::new_bound(py); + item.set_item("segment_label", "cpu:0") + .expect("set segment_label"); + item.set_item("write_ptr", write_ptr) + .expect("set write_ptr"); + item.set_item("read_ptr", read_ptr).expect("set read_ptr"); + item.set_item("len", len).expect("set len"); + item.set_item("generation", self_info.node_start_time) + .expect("set generation"); + item.set_item("node_id", self_info.id).expect("set node_id"); + let out = PyList::empty_bound(py); + out.append(item).expect("append local segment"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "wait_local_segments_ready failed", + )), } } - allocate_lease_inner(self, ttl_seconds, py).into_py_object(py) + + wait_local_segments_ready_inner(self, py).into_py_object(py) } - /// Keepalive a lease synchronously. Type must be specified to avoid ambiguity. - /// This uses the lease's existing TTL on the master. - #[pyo3(signature = (lease_id, lease_type))] - fn keepalive_lease( - &self, - lease_id: u64, - lease_type: &Bound<'_, PyAny>, - py: Python, - ) -> PyObject { - fn keepalive_lease_inner( + #[pyo3(signature = ())] + fn reserve_local_grant_blocking(&self, py: Python) -> PyObject { + fn reserve_local_grant_blocking_inner( client: &KvClient, - lease_id: u64, - lease_type: &Bound<'_, PyAny>, py: Python, ) -> ApiResult { - // Accept simple enum-like strings: "kvclient" | "etcd" - let lease_type_str = match lease_type.extract::() { - Ok(s) => s.to_ascii_lowercase(), - Err(_) => { - return ApiResult::new_error(new_invalid_argument_error( - py, - "lease_type must be 'kvclient' or 'etcd'", - )); - } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), }; - - if lease_type_str != "kvclient" { + if framework.is_external_mode() { return ApiResult::new_error(new_invalid_argument_error( py, - "keepalive_lease(type=etcd) is not supported in fluxon_pyo3; use fluxon_mq.LeaseManagerHandle for etcd leases", + "reserve_local_grant_blocking currently only supports owner mode", )); } - let framework = match require_kv_framework_api(client, py) { - Ok(v) => v, - Err(e) => return ApiResult::new_error(e), - }; let runtime = match client.runtime.as_ref() { Some(v) => v, None => { @@ -2902,59 +3479,2110 @@ impl KvClient { )); } }; - let r: Result<(), String> = match py.allow_threads(|| { + let framework_for_reserve = framework.clone(); + let reserve_resp = match py.allow_threads(|| { runtime.run_async_from_sync(async move { - framework - .kv_keepalive_lease(lease_id) + framework_for_reserve + .client_kv_api_view() + .client_kv_api() + .inner() + .reserve_local_grant() .await - .map_err(|e| e.to_string()) }) }) { Ok(v) => v, - Err(e) => Err(format!("runtime bridge failed: {}", e)), + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } }; - match r { - Ok(_) => { - ApiResult::new_success(Python::with_gil(|py| new_none_success_instance(py))) + match reserve_resp { + Ok(ReserveLocalGrantOutcome::Granted { + grant_id, + node_id: _, + addr, + base_addr, + len, + }) => ApiResult::new_success((grant_id, addr, base_addr, len).into_py(py)), + Ok(ReserveLocalGrantOutcome::None) => { + unreachable!("reserve_local_grant filters empty successful outcomes") } - Err(e) => ApiResult::new_error(Python::with_gil(|py| { - new_network_error(py, &format!("Keepalive lease failed: {}", e), None) - })), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + &format!( + "reserve_local_grant_blocking failed for fixed quantum={}B", + OWNER_LOCAL_RESERVE_GRANT_QUANTUM_BYTES + ), + )), + } + } + + reserve_local_grant_blocking_inner(self, py).into_py_object(py) + } + + #[pyo3(signature = (grant_id))] + fn release_local_grant_blocking(&self, grant_id: u64, py: Python) -> PyObject { + fn release_local_grant_blocking_inner( + client: &KvClient, + grant_id: u64, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "release_local_grant_blocking currently only supports owner mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework_for_release = framework.clone(); + let release_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_release + .client_kv_api_view() + .client_kv_api() + .inner() + .release_local_grant(grant_id) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match release_result { + Ok(()) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "release_local_grant_blocking failed", + )), + } + } + + release_local_grant_blocking_inner(self, grant_id, py).into_py_object(py) + } + + #[pyo3(signature = (keys, value_len, reject_if_inflight_same_key=false, reject_if_exist_same_key=false, write_through=true, make_replica_task=true, make_replica_task_mask=None, atomic_group_lens=None))] + fn local_fast_put_start( + &self, + keys: Vec, + value_len: u64, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, + make_replica_task_mask: Option>, + atomic_group_lens: Option>, + py: Python, + ) -> PyObject { + fn local_fast_put_start_inner( + client: &KvClient, + keys: Vec, + value_len: u64, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, + make_replica_task_mask: Option>, + atomic_group_lens: Option>, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_start requires at least one key", + )); + } + if value_len == 0 { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_start requires value_len > 0", + )); + } + if value_len > u32::MAX as u64 { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_start value_len exceeds u32 limit", + )); + } + let atomic_group_lens = + match normalize_put_atomic_group_lens(keys.len(), atomic_group_lens) { + Ok(group_lens) => group_lens, + Err(detail) => { + return ApiResult::new_error(new_invalid_argument_error(py, &detail)); + } + }; + let make_replica_task_mask = match resolve_make_replica_task_mask( + keys.len(), + write_through, + make_replica_task, + make_replica_task_mask, + &atomic_group_lens, + ) { + Ok(mask) => mask, + Err(detail) => { + return ApiResult::new_error(new_invalid_argument_error(py, &detail)); + } + }; + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + if framework.is_external_mode() { + let framework_for_start = framework.clone(); + let keys_for_req = keys.clone(); + let make_replica_task_mask_for_req = make_replica_task_mask.clone(); + let atomic_group_lens_for_req = atomic_group_lens.clone(); + let external_start_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + let external_api_view = framework_for_start.external_client_api_view(); + let inner = external_api_view.external_client_api().inner(); + let started_time = inner.current_owner_start_time().await; + let base_addr = inner.base_ptr().await?; + let short_circuit_payload = inner.short_circuit_put_payload_path_enabled(); + let resp = inner + .external_batch_put_start_rpc(ExternalBatchPutStartReq { + items: keys_for_req + .iter() + .zip(make_replica_task_mask_for_req.iter()) + .map(|(key, make_replica_task)| ExternalBatchPutStartItemReq { + key: key.clone(), + len: value_len, + reject_if_inflight_same_key, + reject_if_exist_same_key, + make_replica_task: *make_replica_task, + preferred_sub_cluster: None, + }) + .collect(), + atomic_group_lens: Some(atomic_group_lens_for_req), + started_time, + }) + .await?; + Ok::<_, CoreKvError>((base_addr, started_time, short_circuit_payload, resp)) + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + let (base_addr, started_time, short_circuit_payload, resp) = + match external_start_result { + Ok(v) => v, + Err(err) => { + return ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "local_fast_put_start failed", + )); + } + }; + let revoke_started_puts = |started_items: Vec<( + String, + fluxon_kv::master_kv_router::put::PutIDForAKey, + )>| + -> Option { + if started_items.is_empty() { + return None; + } + let framework_for_revoke = framework.clone(); + match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + let external_api_view = framework_for_revoke.external_client_api_view(); + let inner = external_api_view.external_client_api().inner(); + for (key, put_id) in started_items { + inner + .external_put_revoke_rpc(ExternalPutRevokeReq { + key, + put_id: Some(put_id), + started_time, + }) + .await?; + } + Ok::<(), CoreKvError>(()) + }) + }) { + Ok(Ok(())) => None, + Ok(Err(err)) => Some(err), + Err(err) => Some(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "runtime bridge failed during external local_fast_put_start cleanup: {}", + err + ), + })), + } + }; + let local_fast_put_start_error = |primary_err: CoreKvError, + cleanup_err: Option| + -> ApiResult { + let err = match cleanup_err { + Some(cleanup_err) => combine_kv_errors( + "local_fast_put_start external cleanup failed", + primary_err, + cleanup_err, + ), + None => primary_err, + }; + ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "local_fast_put_start failed", + )) + }; + if resp.items.len() != keys.len() { + let started_items = keys + .iter() + .zip(resp.items.iter()) + .filter_map(|(key, item)| { + if item.error_code == OK { + item.put_id.map(|put_id| (key.clone(), put_id)) + } else { + None + } + }) + .collect::>(); + let err = CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_start external response length mismatch: expected={} got={}", + keys.len(), + resp.items.len() + ), + }); + return local_fast_put_start_error(err, revoke_started_puts(started_items)); + } + let mut value_ptrs = Vec::with_capacity(keys.len()); + let mut items = Vec::with_capacity(keys.len()); + let mut started_items = Vec::with_capacity(keys.len()); + for (key, item) in keys.iter().zip(resp.items.into_iter()) { + if item.error_code != OK { + let err = CoreKvError::from_json(item.error_code, &item.error_json); + return local_fast_put_start_error(err, revoke_started_puts(started_items)); + } + let Some(put_id) = item.put_id else { + let err = CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_start external missing put_id in success response for key={}", + key + ), + }); + return local_fast_put_start_error(err, revoke_started_puts(started_items)); + }; + started_items.push((key.clone(), put_id)); + let value_ptr = (base_addr as u64) + .checked_add(item.target_offset) + .ok_or_else(|| { + CoreKvError::Api(CoreApiError::Unknown { + detail: "local_fast_put_start external value ptr overflow" + .to_string(), + }) + }); + let value_ptr = match value_ptr { + Ok(v) => v, + Err(err) => { + return local_fast_put_start_error( + err, + revoke_started_puts(started_items), + ); + } + }; + value_ptrs.push(value_ptr); + items.push(StagedExternalPutItem { + put_id, + src_offset: item.src_offset, + target_offset: item.target_offset, + transfer_target_offset: item.transfer_target_offset, + peer_id: item.peer_id, + target_base_addr: item.target_base_addr, + }); + } + let plan = Arc::new(FluxonPutPlan { + blob: build_plan_blob(&value_ptrs), + state: Mutex::new(FluxonPutPlanState::Prepared(StagedPutPlanData::External( + StagedExternalPutPlanData { + keys, + value_len, + started_time, + items, + short_circuit_payload, + }, + ))), + }); + let plan_ptr = plan_ptr_from_blob(plan.blob.as_ref()); + client + .plan_registry + .write() + .expect("plan_registry poisoned") + .insert(plan_ptr, FluxonPlanRegistryEntry::Put(plan)); + return ApiResult::new_success((plan_ptr as u64).into_py(py)); + } + + if write_through { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_start write_through=true requires external mode", + )); + } + + let key_count = keys.len(); + let framework_for_start = framework.clone(); + let prepare_items = keys + .iter() + .map(|key| BatchPreparePutKeyItemReq { + key: key.clone(), + reject_if_inflight_same_key, + reject_if_exist_same_key, + }) + .collect::>(); + + let prepare_and_slot_lease_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + let client_kv_api_view = framework_for_start.client_kv_api_view(); + let inner = client_kv_api_view.client_kv_api().inner(); + let prepared = inner.batch_prepare_put_keys(prepare_items).await?; + if prepared.reservation_ids.len() != key_count { + let cleanup_result = inner + .batch_release_put_key_reservations(prepared.reservation_ids.clone()) + .await; + return match cleanup_result { + Ok(_) => Err(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_start reservation length mismatch: expected={} got={}", + key_count, + prepared.reservation_ids.len() + ), + })), + Err(cleanup_err) => Err(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_start reservation length mismatch and cleanup failed: expected={} got={} cleanup_err={}", + key_count, + prepared.reservation_ids.len(), + cleanup_err + ), + })), + }; + } + let slot_lease = match inner + .owner_claim_local_reserve_slot_lease(value_len, key_count) + .await + { + Ok(slot_lease) => slot_lease, + Err(err) => { + let cleanup_result = inner + .batch_release_put_key_reservations(prepared.reservation_ids.clone()) + .await; + return match cleanup_result { + Ok(_) => Err(err), + Err(cleanup_err) => Err(CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_start reserve_local_grant failed: {} cleanup_err={}", + err, cleanup_err + ), + })), + }; + } + }; + Ok((prepared.reservation_ids, slot_lease)) + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + let (key_reservation_ids, slot_lease) = match prepare_and_slot_lease_result { + Ok(v) => v, + Err(err) => { + return ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "local_fast_put_start failed", + )); + } + }; + + let value_ptrs = slot_lease.value_ptrs(); + + let plan = Arc::new(FluxonPutPlan { + blob: build_plan_blob(&value_ptrs), + state: Mutex::new(FluxonPutPlanState::Prepared(StagedPutPlanData::Owner( + StagedOwnerPutPlanData { + keys, + value_len, + key_reservation_ids, + slot_lease, + make_replica_task_mask, + atomic_group_lens, + }, + ))), + }); + let plan_ptr = plan_ptr_from_blob(plan.blob.as_ref()); + client + .plan_registry + .write() + .expect("plan_registry poisoned") + .insert(plan_ptr, FluxonPlanRegistryEntry::Put(plan)); + ApiResult::new_success((plan_ptr as u64).into_py(py)) + } + + local_fast_put_start_inner( + self, + keys, + value_len, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + make_replica_task_mask, + atomic_group_lens, + py, + ) + .into_py_object(py) + } + + #[pyo3(signature = (plan_ptr))] + fn local_fast_put_commit(&self, plan_ptr: u64, py: Python) -> PyObject { + fn local_fast_put_commit_inner( + client: &KvClient, + plan_ptr: u64, + py: Python, + ) -> ApiResult { + let plan_ptr_usize = match plan_ptr_u64_to_usize(plan_ptr, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime_handle = match client.runtime.as_ref() { + Some(v) => v.handle().clone(), + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let registry_entry = { + let guard = client.plan_registry.read().expect("plan_registry poisoned"); + guard.get(&plan_ptr_usize).cloned() + }; + let Some(FluxonPlanRegistryEntry::Put(plan)) = registry_entry else { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_commit requires a live put plan_ptr", + )); + }; + let plan_data = { + let mut state_guard = plan.state.lock().expect("put plan state poisoned"); + match std::mem::replace(&mut *state_guard, FluxonPutPlanState::Committing) { + FluxonPutPlanState::Prepared(data) => data, + FluxonPutPlanState::Committing => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "local_fast_put_commit cannot be called twice on the same plan_ptr", + )); + } + } + }; + let framework_for_commit = framework.clone(); + let plan_registry = client.plan_registry.clone(); + let transfer_concurrency = 128usize; + let future = async move { + let ret_codes = match plan_data { + StagedPutPlanData::External(external_data) => { + let external_api_view = framework_for_commit.external_client_api_view(); + let inner = external_api_view.external_client_api().inner(); + let keys = external_data.keys; + let items = external_data.items; + let result = if external_data.short_circuit_payload { + inner + .external_batch_put_commit_rpc(ExternalBatchPutCommitReq { + items: keys + .iter() + .zip(items.iter()) + .map(|(key, item)| ExternalBatchPutCommitItemReq { + key: key.clone(), + len: external_data.value_len, + src_offset: item.src_offset, + remote_target: item.peer_id.is_some(), + put_id: Some(item.put_id), + lease_id: None, + }) + .collect(), + started_time: external_data.started_time, + }) + .await + .map(|resp| { + resp.items + .into_iter() + .map(|item| { + rpc_item_error_to_ret_code( + item.error_code, + &item.error_json, + ) + }) + .collect::>() + }) + } else { + inner + .external_batch_put_transfer_end_rpc( + ExternalBatchPutTransferEndReq { + items: keys + .iter() + .zip(items.iter()) + .map(|(key, item)| ExternalBatchPutTransferEndItemReq { + key: key.clone(), + len: external_data.value_len, + src_offset: item.src_offset, + target_offset: item + .transfer_target_offset + .unwrap_or(item.target_offset), + peer_id: item.peer_id.clone(), + target_base_addr: if item.peer_id.is_some() { + Some(item.target_base_addr) + } else { + None + }, + put_id: Some(item.put_id), + lease_id: None, + }) + .collect(), + started_time: external_data.started_time, + transfer_concurrency, + }, + ) + .await + .map(|resp| { + resp.items + .into_iter() + .map(|item| { + rpc_item_error_to_ret_code( + item.error_code, + &item.error_json, + ) + }) + .collect::>() + }) + }; + match result { + Ok(ret_codes) => ret_codes, + Err(err) => vec![kv_error_to_ret_code(&err); keys.len()], + } + } + StagedPutPlanData::Owner(owner_data) => { + let client_kv_api_view = framework_for_commit.client_kv_api_view(); + let inner = client_kv_api_view.client_kv_api().inner(); + let slot_refs = owner_data.slot_lease.slots.clone(); + let slot_size = owner_data.slot_lease.slot_size; + if slot_refs.len() != owner_data.keys.len() + || owner_data.make_replica_task_mask.len() != owner_data.keys.len() + { + let err = CoreKvError::Api(CoreApiError::Unknown { + detail: format!( + "local_fast_put_commit staged length mismatch: keys={} slots={} replica_mask={}", + owner_data.keys.len(), + slot_refs.len(), + owner_data.make_replica_task_mask.len() + ), + }); + let _ = release_staged_put_resources( + inner, + owner_data.key_reservation_ids, + owner_data.slot_lease, + ) + .await; + return ApiResult::new_success(vec![ + kv_error_to_ret_code(&err); + owner_data.keys.len() + ]); + } + let value_len_u32 = match u32::try_from(owner_data.value_len) { + Ok(value) => value, + Err(_) => { + let err = CoreKvError::Api(CoreApiError::InvalidArgument { + detail: "local_fast_put_commit value_len does not fit into u32" + .to_string(), + }); + let _ = release_staged_put_resources( + inner, + owner_data.key_reservation_ids, + owner_data.slot_lease, + ) + .await; + return ApiResult::new_success(vec![ + kv_error_to_ret_code(&err); + owner_data.keys.len() + ]); + } + }; + let mut publish_items = Vec::with_capacity(owner_data.keys.len()); + let mut ret_codes = Vec::with_capacity(owner_data.keys.len()); + let put_ids = owner_data + .keys + .iter() + .map(|_| inner.next_owner_local_first_put_id()) + .collect::>(); + let keys_and_put_ids = owner_data + .keys + .iter() + .cloned() + .zip(put_ids.iter().copied()) + .collect::>(); + let atomic_groups = match build_put_atomic_group_assignments( + &keys_and_put_ids, + &owner_data.atomic_group_lens, + ) { + Ok(groups) => groups, + Err(detail) => { + let err = + CoreKvError::Api(CoreApiError::InvalidArgument { detail }); + let _ = release_staged_put_resources( + inner, + owner_data.key_reservation_ids, + owner_data.slot_lease, + ) + .await; + return ApiResult::new_success(vec![ + kv_error_to_ret_code(&err); + owner_data.keys.len() + ]); + } + }; + for (idx, ((key, slot_ref), make_replica_task)) in owner_data + .keys + .iter() + .zip(slot_refs.iter()) + .zip(owner_data.make_replica_task_mask.iter()) + .enumerate() + { + let memory_info = inner + .build_local_reserve_resident_memory_info( + key, + slot_ref.ptr, + value_len_u32, + slot_size, + slot_ref.grant_id, + slot_ref.slot_index, + ) + .await; + inner.install_precommit_local_visible_memory_info(key, memory_info); + let put_id = put_ids[idx]; + inner.record_put_locality(false, owner_data.value_len, 0); + publish_items.push(OwnerLocalPublishItem { + key: key.clone(), + put_id, + value_len: owner_data.value_len, + lease_id: None, + committed_slot: PutDoneCommittedSlot { + grant_id: slot_ref.grant_id, + slot_index: slot_ref.slot_index, + slot_size, + addr: slot_ref.ptr, + base_addr: slot_ref.base_addr, + len: owner_data.value_len, + }, + make_replica_task: *make_replica_task, + // This owner-local path has no earlier remote + // reservation. Selected items are admitted to + // run append Start, where allocation occurs. + remote_replica_admitted: *make_replica_task, + preferred_sub_cluster: None, + atomic_group: atomic_groups[idx].clone(), + }); + ret_codes.push(0); + } + if !publish_items.is_empty() { + let publish_keys = publish_items + .iter() + .map(|item| item.key.clone()) + .collect::>(); + if let Err(err) = inner + .enqueue_owner_local_publish(OwnerLocalPublishJob { + items: publish_items, + key_reservation_ids: owner_data.key_reservation_ids.clone(), + external_pending_contexts: Vec::new(), + }) + .await + { + tracing::warn!( + "local_fast_put_commit owner local publish enqueue failed after local commit: {}", + err + ); + for key in publish_keys { + if let Some(memory_info) = + inner.precommit_local_visible_memory_info(&key) + { + let _ = inner + .remove_precommit_local_reserve_resident_slot_if_same( + &key, + &memory_info, + ); + } + } + if let Err(cleanup_err) = release_put_key_reservations_only( + inner, + owner_data.key_reservation_ids, + ) + .await + { + tracing::warn!( + "local_fast_put_commit key reservation cleanup failed after owner local publish enqueue error: {}", + cleanup_err + ); + } + ret_codes = vec![kv_error_to_ret_code(&err); ret_codes.len()]; + } + } else if let Err(err) = + release_put_key_reservations_only(inner, owner_data.key_reservation_ids) + .await + { + tracing::warn!( + "local_fast_put_commit key reservation cleanup failed after partial local commit: {}", + err + ); + } + ret_codes + } + }; + cleanup_plan_registry_entry(&plan_registry, plan_ptr_usize); + ApiResult::new_success(ret_codes) + }; + match KvFuture::new(future, runtime_handle, py) { + Ok(v) => ApiResult::new_success(v.into_any()), + Err(e) => ApiResult::new_error(new_general_error( + py, + &format!("Failed to create future: {}", e), + )), + } + } + + local_fast_put_commit_inner(self, plan_ptr, py).into_py_object(py) + } + + #[pyo3(signature = (plan_ptr))] + fn put_abort(&self, plan_ptr: u64, py: Python) -> PyObject { + fn put_abort_inner(client: &KvClient, plan_ptr: u64, py: Python) -> ApiResult { + let plan_ptr_usize = match plan_ptr_u64_to_usize(plan_ptr, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let removed = { + let mut guard = client + .plan_registry + .write() + .expect("plan_registry poisoned"); + guard.remove(&plan_ptr_usize) + }; + let Some(FluxonPlanRegistryEntry::Put(plan)) = removed else { + return ApiResult::new_error(new_invalid_argument_error( + py, + "put_abort requires a live put plan_ptr", + )); + }; + let plan_data = { + let mut state_guard = plan.state.lock().expect("put plan state poisoned"); + match std::mem::replace(&mut *state_guard, FluxonPutPlanState::Committing) { + FluxonPutPlanState::Prepared(data) => data, + FluxonPutPlanState::Committing => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "put_abort cannot be used after local_fast_put_commit", + )); + } + } + }; + let abort_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + match plan_data { + StagedPutPlanData::External(external_data) => { + let external_api_view = framework.external_client_api_view(); + let inner = external_api_view.external_client_api().inner(); + for (key, item) in + external_data.keys.iter().zip(external_data.items.iter()) + { + inner + .external_put_revoke_rpc(ExternalPutRevokeReq { + key: key.clone(), + put_id: Some(item.put_id), + started_time: external_data.started_time, + }) + .await?; + } + Ok::<(), CoreKvError>(()) + } + StagedPutPlanData::Owner(owner_data) => { + let client_kv_api_view = framework.client_kv_api_view(); + let inner = client_kv_api_view.client_kv_api().inner(); + release_staged_put_resources( + inner, + owner_data.key_reservation_ids, + owner_data.slot_lease, + ) + .await + } + } + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match abort_result { + Ok(()) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "put_abort failed", + )), + } + } + + put_abort_inner(self, plan_ptr, py).into_py_object(py) + } + + #[pyo3(signature = (keys, concurrency=None))] + fn get_views(&self, keys: Vec, concurrency: Option, py: Python) -> PyObject { + fn get_views_inner( + client: &KvClient, + keys: Vec, + concurrency: Option, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_views requires at least one key", + )); + } + if matches!(concurrency, Some(0)) { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_concurrency must be > 0", + )); + } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let keys_for_get = keys.clone(); + let locality_counters = client.locality_counters.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + let mut holders = Vec::with_capacity(keys_for_get.len()); + let mut value_ptrs = Vec::with_capacity(keys_for_get.len()); + for key in &keys_for_get { + match framework.kv_get(key).await? { + KvGetResult::Owner(Some(holder)) => { + let bytes = holder.bytes().len() as u64; + locality_counters.record_l2_hit(false, bytes); + locality_counters.record_get(false, bytes, 0); + value_ptrs.push(holder.bytes().as_ptr() as u64); + holders.push(StagedGetViewHolder::Owner { _holder: holder }); + } + KvGetResult::External(Some(holder)) => { + let bytes = holder.bytes().len() as u64; + locality_counters.record_l2_hit(true, bytes); + locality_counters.record_get(true, bytes, 0); + value_ptrs.push(holder.bytes().as_ptr() as u64); + holders.push(StagedGetViewHolder::External { _holder: holder }); + } + KvGetResult::Owner(None) | KvGetResult::External(None) => { + return Err(CoreKvError::Api(CoreApiError::KeyNotFound { + key: key.clone(), + })); + } + } + } + Ok::<_, CoreKvError>((value_ptrs, holders)) + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + let (value_ptrs, holders) = match result { + Ok(v) => v, + Err(err) => { + return ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_views failed", + )); + } + }; + let plan = Arc::new(FluxonGetViewsPlan { + blob: build_plan_blob(&value_ptrs), + _holders: holders, + }); + let plan_ptr = plan_ptr_from_blob(plan.blob.as_ref()); + client + .plan_registry + .write() + .expect("plan_registry poisoned") + .insert(plan_ptr, FluxonPlanRegistryEntry::Get(plan)); + ApiResult::new_success((plan_ptr as u64).into_py(py)) + } + + get_views_inner(self, keys, concurrency, py).into_py_object(py) + } + + #[pyo3(signature = (plan_ptr))] + fn release_views(&self, plan_ptr: u64, py: Python) -> PyObject { + fn release_views_inner( + client: &KvClient, + plan_ptr: u64, + py: Python, + ) -> ApiResult { + let plan_ptr_usize = match plan_ptr_u64_to_usize(plan_ptr, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let removed = { + let mut guard = client + .plan_registry + .write() + .expect("plan_registry poisoned"); + guard.remove(&plan_ptr_usize) + }; + match removed { + Some(FluxonPlanRegistryEntry::Get(plan)) => { + defer_drop_to_runtime(client.runtime.as_ref(), plan); + ApiResult::new_success(0i32.into_py(py)) + } + Some(FluxonPlanRegistryEntry::Put(_)) => ApiResult::new_error( + new_invalid_argument_error(py, "release_views requires a get-views plan_ptr"), + ), + None => ApiResult::new_error(new_invalid_argument_error( + py, + "release_views requires a live get-views plan_ptr", + )), + } + } + + release_views_inner(self, plan_ptr, py).into_py_object(py) + } + + #[pyo3(signature = (keys, prefix_best_effort=true, atomic_group_lens=None, concurrency=None))] + fn get_start( + &self, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + concurrency: Option, + py: Python, + ) -> PyObject { + fn get_start_inner( + client: &KvClient, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + concurrency: Option, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start requires at least one key", + )); + } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let batch_concurrency = match concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY) { + 0 => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_concurrency must be > 0", + )); + } + v => v, + }; + let framework_for_start = framework.clone(); + let started = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_start + .external_client_api_view() + .external_client_api() + .inner() + .get_start( + keys, + prefix_best_effort, + atomic_group_lens, + batch_concurrency, + ) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match started { + Ok(v) => ApiResult::new_success(external_get_start_result_to_py(&v, py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_start failed", + )), + } + } + + get_start_inner( + self, + keys, + prefix_best_effort, + atomic_group_lens, + concurrency, + py, + ) + .into_py_object(py) + } + + #[pyo3(signature = (handle, consume_prefix_len=None))] + fn get_transfer(&self, handle: u64, consume_prefix_len: Option, py: Python) -> PyObject { + fn get_transfer_inner( + client: &KvClient, + handle: u64, + consume_prefix_len: Option, + py: Python, + ) -> ApiResult { + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_transfer is supported only in external-client mode", + )); + } + let framework_for_transfer = framework.clone(); + let transfer_results = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_transfer + .external_client_api_view() + .external_client_api() + .inner() + .get_transfer(handle, consume_prefix_len) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + let transfer_results = match transfer_results { + Ok(v) => v, + Err(err) => { + return ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_transfer failed", + )); + } + }; + let mut value_ptrs = Vec::with_capacity(transfer_results.len()); + let mut holders = Vec::with_capacity(transfer_results.len()); + for (idx, item) in transfer_results.into_iter().enumerate() { + match item { + Ok(Some(holder)) => { + value_ptrs.push(holder.bytes().as_ptr() as u64); + holders.push(StagedGetViewHolder::External { _holder: holder }); + } + Ok(None) => { + return ApiResult::new_error(new_key_not_found_error( + py, + &format!("get_transfer external transfer missed key index {}", idx), + None, + )); + } + Err(err) => { + return ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_transfer external transfer item failed", + )); + } + } + } + let view_plan = Arc::new(FluxonGetViewsPlan { + blob: build_plan_blob(&value_ptrs), + _holders: holders, + }); + let plan_ptr = plan_ptr_from_blob(view_plan.blob.as_ref()); + let mut registry_guard = client + .plan_registry + .write() + .expect("plan_registry poisoned"); + registry_guard.insert(plan_ptr, FluxonPlanRegistryEntry::Get(view_plan)); + drop(registry_guard); + ApiResult::new_success((plan_ptr as u64).into_py(py)) + } + + get_transfer_inner(self, handle, consume_prefix_len, py).into_py_object(py) + } + + #[pyo3(signature = (handle))] + fn cancel_get_transfer(&self, handle: u64, py: Python) -> PyObject { + fn cancel_get_transfer_inner( + client: &KvClient, + handle: u64, + py: Python, + ) -> ApiResult { + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "cancel_get_transfer is supported only in external-client mode", + )); + } + let framework_for_cancel = framework.clone(); + let cancel_res = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_cancel + .external_client_api_view() + .external_client_api() + .inner() + .cancel_get_transfer(handle) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match cancel_res { + Ok(()) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "cancel_get_transfer failed", + )), + } + } + + cancel_get_transfer_inner(self, handle, py).into_py_object(py) + } + + /// Plan one generation-fenced Get prefix without allocating a destination. + #[pyo3(signature = (keys, prefix_best_effort=true, atomic_group_lens=None))] + fn get_plan( + &self, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + py: Python, + ) -> PyObject { + fn get_plan_inner( + client: &KvClient, + keys: Vec, + prefix_best_effort: bool, + atomic_group_lens: Option>, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_plan requires at least one key", + )); + } + let framework = match require_kv_framework_api(client, py) { + Ok(value) => value, + Err(err) => return ApiResult::new_error(err), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_plan is supported only in external-client mode", + )); + } + let Some(runtime) = client.runtime.as_ref() else { + return ApiResult::new_error(new_general_error(py, "Client runtime is missing")); + }; + let framework_for_plan = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_plan + .external_client_api_view() + .external_client_api() + .inner() + .get_plan(keys, prefix_best_effort, atomic_group_lens) + .await + }) + }) { + Ok(result) => result, + Err(err) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )); + } + }; + match result { + Ok(plan) => { + let out = PyDict::new_bound(py); + out.set_item("handle", plan.handle).expect("set handle"); + out.set_item("raw_prefix_hit_len", plan.raw_prefix_hit_len as u64) + .expect("set raw_prefix_hit_len"); + out.set_item("gpu_raw_prefix_hit_len", plan.gpu_raw_prefix_hit_len as u64) + .expect("set gpu_raw_prefix_hit_len"); + out.set_item("gpu_remote_indices", plan.gpu_remote_indices) + .expect("set gpu_remote_indices"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_plan failed", + )), + } + } + + get_plan_inner(self, keys, prefix_best_effort, atomic_group_lens, py).into_py_object(py) + } + + #[pyo3(signature = (handle, consume_prefix_len, concurrency=None))] + fn execute_get_plan_cpu( + &self, + handle: u64, + consume_prefix_len: usize, + concurrency: Option, + py: Python, + ) -> PyObject { + let batch_concurrency = concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY); + let framework = match require_kv_framework_api(self, py) { + Ok(value) => value, + Err(err) => return ApiResult::::new_error(err).into_py_object(py), + }; + if !framework.is_external_mode() || batch_concurrency == 0 { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "execute_get_plan_cpu requires external mode and positive concurrency", + )) + .into_py_object(py); + } + let Some(runtime) = self.runtime.as_ref() else { + return ApiResult::::new_error(new_general_error( + py, + "Client runtime is missing", + )) + .into_py_object(py); + }; + let framework_for_execute = framework.clone(); + let result = py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_execute + .external_client_api_view() + .external_client_api() + .inner() + .execute_get_plan_cpu(handle, consume_prefix_len, batch_concurrency) + .await + }) + }); + match result { + Ok(Ok(())) => ApiResult::new_success(0i32.into_py(py)), + Ok(Err(err)) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "execute_get_plan_cpu failed", + )), + Err(err) => ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )), + } + .into_py_object(py) + } + + #[pyo3(signature = (handle, destinations, consume_prefix_len, concurrency=None))] + fn execute_get_plan_gpu( + &self, + handle: u64, + destinations: Vec<(u64, u64, u64)>, + consume_prefix_len: usize, + concurrency: Option, + py: Python, + ) -> PyObject { + let batch_concurrency = concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY); + let framework = match require_kv_framework_api(self, py) { + Ok(value) => value, + Err(err) => return ApiResult::::new_error(err).into_py_object(py), + }; + if !framework.is_external_mode() || batch_concurrency == 0 { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "execute_get_plan_gpu requires external mode and positive concurrency", + )) + .into_py_object(py); + } + let Some(runtime) = self.runtime.as_ref() else { + return ApiResult::::new_error(new_general_error( + py, + "Client runtime is missing", + )) + .into_py_object(py); + }; + let destinations = destinations + .into_iter() + .map(|(registration_id, addr, capacity)| { + fluxon_kv::external_client_api::ExternalGpuDestination { + registration_id, + addr, + capacity, + } + }) + .collect(); + let framework_for_execute = framework.clone(); + let result = py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_execute + .external_client_api_view() + .external_client_api() + .inner() + .execute_get_plan_gpu( + handle, + destinations, + consume_prefix_len, + batch_concurrency, + ) + .await + }) + }); + match result { + Ok(Ok(())) => ApiResult::new_success(0i32.into_py(py)), + Ok(Err(err)) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "execute_get_plan_gpu failed", + )), + Err(err) => ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )), + } + .into_py_object(py) + } + + #[pyo3(signature = (handle))] + fn cancel_get_plan(&self, handle: u64, py: Python) -> PyObject { + let framework = match require_kv_framework_api(self, py) { + Ok(value) => value, + Err(err) => return ApiResult::::new_error(err).into_py_object(py), + }; + let Some(runtime) = self.runtime.as_ref() else { + return ApiResult::::new_error(new_general_error( + py, + "Client runtime is missing", + )) + .into_py_object(py); + }; + let framework_for_cancel = framework.clone(); + let result = py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_cancel + .external_client_api_view() + .external_client_api() + .inner() + .cancel_get_plan(handle) + .await + }) + }); + match result { + Ok(Ok(())) => ApiResult::new_success(0i32.into_py(py)), + Ok(Err(err)) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "cancel_get_plan failed", + )), + Err(err) => ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )), + } + .into_py_object(py) + } + + /// Start background RDMA pulls directly into caller-owned GPU destinations. + #[pyo3(signature = (keys, destinations, prefix_best_effort=true, atomic_group_lens=None, concurrency=None))] + fn get_start_gpu( + &self, + keys: Vec, + destinations: Vec<(u64, u64, u64)>, + prefix_best_effort: bool, + atomic_group_lens: Option>, + concurrency: Option, + py: Python, + ) -> PyObject { + fn get_start_gpu_inner( + client: &KvClient, + keys: Vec, + destinations: Vec<(u64, u64, u64)>, + prefix_best_effort: bool, + atomic_group_lens: Option>, + concurrency: Option, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start_gpu requires at least one key", + )); + } + if keys.len() != destinations.len() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start_gpu requires one destination per key", + )); + } + let batch_concurrency = match concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY) { + 0 => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start_gpu concurrency must be > 0", + )); + } + value => value, + }; + let framework = match require_kv_framework_api(client, py) { + Ok(value) => value, + Err(err) => return ApiResult::new_error(err), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_start_gpu is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(value) => value, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let destinations = destinations + .into_iter() + .map(|(registration_id, addr, capacity)| { + fluxon_kv::external_client_api::ExternalGpuDestination { + registration_id, + addr, + capacity, + } + }) + .collect(); + let framework_for_start = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_start + .external_client_api_view() + .external_client_api() + .inner() + .get_start_gpu( + keys, + destinations, + prefix_best_effort, + atomic_group_lens, + batch_concurrency, + ) + .await + }) + }) { + Ok(value) => value, + Err(err) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )); + } + }; + match result { + Ok(started) => { + let out = PyDict::new_bound(py); + out.set_item("handle", started.handle).expect("set handle"); + out.set_item("raw_prefix_hit_len", started.raw_prefix_hit_len as u64) + .expect("set raw_prefix_hit_len"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_start_gpu failed", + )), + } + } + + get_start_gpu_inner( + self, + keys, + destinations, + prefix_best_effort, + atomic_group_lens, + concurrency, + py, + ) + .into_py_object(py) + } + + /// Wait for a GPU Get terminal and consume an atomic prefix. + #[pyo3(signature = (handle, consume_prefix_len=None))] + fn get_transfer_gpu( + &self, + handle: u64, + consume_prefix_len: Option, + py: Python, + ) -> PyObject { + fn get_transfer_gpu_inner( + client: &KvClient, + handle: u64, + consume_prefix_len: Option, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(value) => value, + Err(err) => return ApiResult::new_error(err), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "get_transfer_gpu is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(value) => value, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework_for_transfer = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_transfer + .external_client_api_view() + .external_client_api() + .inner() + .get_transfer_gpu(handle, consume_prefix_len) + .await + }) + }) { + Ok(value) => value, + Err(err) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )); + } + }; + match result { + Ok(terminal) => { + if terminal.value_ptrs.len() != terminal.consumed_prefix_len { + return ApiResult::new_error(new_general_error( + py, + &format!( + "get_transfer_gpu source plan length mismatch: values={} consumed={}", + terminal.value_ptrs.len(), + terminal.consumed_prefix_len + ), + )); + } + let view_plan = Arc::new(FluxonGetViewsPlan { + blob: build_plan_blob(&terminal.value_ptrs), + _holders: terminal + .local_holders + .into_iter() + .map(|holder| StagedGetViewHolder::External { _holder: holder }) + .collect(), + }); + let plan_ptr = plan_ptr_from_blob(view_plan.blob.as_ref()); + client + .plan_registry + .write() + .expect("plan_registry poisoned") + .insert(plan_ptr, FluxonPlanRegistryEntry::Get(view_plan)); + let out = PyDict::new_bound(py); + out.set_item( + "transferred_prefix_len", + terminal.transferred_prefix_len as u64, + ) + .expect("set transferred_prefix_len"); + out.set_item("consumed_prefix_len", terminal.consumed_prefix_len as u64) + .expect("set consumed_prefix_len"); + out.set_item("transfer_wall_us", terminal.transfer_wall_us) + .expect("set transfer_wall_us"); + out.set_item("finish_wait_us", terminal.finish_wait_us) + .expect("set finish_wait_us"); + out.set_item("terminal_before_consume", terminal.terminal_before_consume) + .expect("set terminal_before_consume"); + out.set_item("terminal_to_consume_us", terminal.terminal_to_consume_us) + .expect("set terminal_to_consume_us"); + out.set_item("plan_ptr", plan_ptr as u64) + .expect("set plan_ptr"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "get_transfer_gpu failed", + )), + } + } + + get_transfer_gpu_inner(self, handle, consume_prefix_len, py).into_py_object(py) + } + + /// Cancel a live GPU Get handle and wait until master cleanup is terminal. + #[pyo3(signature = (handle))] + fn cancel_get_transfer_gpu(&self, handle: u64, py: Python) -> PyObject { + fn cancel_get_transfer_gpu_inner( + client: &KvClient, + handle: u64, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(value) => value, + Err(err) => return ApiResult::new_error(err), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "cancel_get_transfer_gpu is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(value) => value, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework_for_cancel = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_cancel + .external_client_api_view() + .external_client_api() + .inner() + .cancel_get_transfer_gpu(handle) + .await + }) + }) { + Ok(value) => value, + Err(err) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {err}"), + )); + } + }; + match result { + Ok(()) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "cancel_get_transfer_gpu failed", + )), + } + } + + cancel_get_transfer_gpu_inner(self, handle, py).into_py_object(py) + } + + /// Allocate a fluxon-kv lease id synchronously. + /// Always allocate a new lease id (no reuse by requested id). + /// Allocate with the provided TTL seconds (must be >= MIN_CLIENT_TTL_SECONDS). + #[pyo3(signature = (ttl_seconds))] + fn allocate_lease(&self, ttl_seconds: u64, py: Python) -> PyObject { + fn allocate_lease_inner( + client: &KvClient, + ttl_seconds: u64, + py: Python, + ) -> ApiResult { + // Enforce minimum TTL at the PyO3 boundary so obvious mistakes fail fast. + if ttl_seconds + < fluxon_kv::master_lease_manager::MasterLeaseManager::MIN_CLIENT_TTL_SECONDS + { + return ApiResult::new_error(new_invalid_argument_error( + py, + &format!( + "allocate_lease(ttl_seconds) requires ttl_seconds >= {} seconds", + fluxon_kv::master_lease_manager::MasterLeaseManager::MIN_CLIENT_TTL_SECONDS, + ), + )); + } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + // Blocking call on the client's runtime; simple and predictable for callers. + let r: Result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework + .kv_allocate_lease(ttl_seconds) + .await + .map_err(|e| e.to_string()) + }) + }) { + Ok(v) => v, + Err(e) => Err(format!("runtime bridge failed: {}", e)), + }; + match r { + Ok(id) => ApiResult::new_success(Python::with_gil(|py| id.into_py(py))), + Err(e) => ApiResult::new_error(Python::with_gil(|py| { + new_network_error(py, &format!("Allocate lease failed: {}", e), None) + })), + } + } + allocate_lease_inner(self, ttl_seconds, py).into_py_object(py) + } + + /// Keepalive a lease synchronously. Type must be specified to avoid ambiguity. + /// This uses the lease's existing TTL on the master. + #[pyo3(signature = (lease_id, lease_type))] + fn keepalive_lease( + &self, + lease_id: u64, + lease_type: &Bound<'_, PyAny>, + py: Python, + ) -> PyObject { + fn keepalive_lease_inner( + client: &KvClient, + lease_id: u64, + lease_type: &Bound<'_, PyAny>, + py: Python, + ) -> ApiResult { + // Accept simple enum-like strings: "kvclient" | "etcd" + let lease_type_str = match lease_type.extract::() { + Ok(s) => s.to_ascii_lowercase(), + Err(_) => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "lease_type must be 'kvclient' or 'etcd'", + )); + } + }; + + if lease_type_str != "kvclient" { + return ApiResult::new_error(new_invalid_argument_error( + py, + "keepalive_lease(type=etcd) is not supported in fluxon_pyo3; use fluxon_mq.LeaseManagerHandle for etcd leases", + )); + } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let r: Result<(), String> = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework + .kv_keepalive_lease(lease_id) + .await + .map_err(|e| e.to_string()) + }) + }) { + Ok(v) => v, + Err(e) => Err(format!("runtime bridge failed: {}", e)), + }; + match r { + Ok(_) => { + ApiResult::new_success(Python::with_gil(|py| new_none_success_instance(py))) + } + Err(e) => ApiResult::new_error(Python::with_gil(|py| { + new_network_error(py, &format!("Keepalive lease failed: {}", e), None) + })), } } keepalive_lease_inner(self, lease_id, lease_type, py).into_py_object(py) } - /// Put a key-value pair (non-blocking) by encoding a flat dict from raw entries. - /// - /// `ptrs` is a list of `(type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra)`: - /// - `dict_key_ptr/dict_key_len`: UTF-8 bytes of the dict field key. - /// - For scalar types (bool/int64/float64), `val_u64` stores raw bits and `val_len` is fixed. - /// - For bytes-like types (string/bytes), `val_u64` stores a pointer and `val_len` is the byte length. - /// - `extra`: reserved for future use. - /// - /// Note: dict field keys cannot be passed as `&str` across async; this function must be able to - /// move all inputs into a Rust future. Therefore we accept pointers for keys and values and rely - /// on the caller to keep the pointed-to memory alive until the async call completes. - /// - /// The backend encoding/copy runs on the Rust runtime without holding the Python GIL. - #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, callback=None))] - fn put( + /// Put a key-value pair (non-blocking) by encoding a flat dict from raw entries. + /// + /// `ptrs` is a list of `(type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra)`: + /// - `dict_key_ptr/dict_key_len`: UTF-8 bytes of the dict field key. + /// - For scalar types (bool/int64/float64), `val_u64` stores raw bits and `val_len` is fixed. + /// - For bytes-like types (string/bytes), `val_u64` stores a pointer and `val_len` is the byte length. + /// - `extra`: reserved for future use. + /// + /// Note: dict field keys cannot be passed as `&str` across async; this function must be able to + /// move all inputs into a Rust future. Therefore we accept pointers for keys and values and rely + /// on the caller to keep the pointed-to memory alive until the async call completes. + /// + /// The backend encoding/copy runs on the Rust runtime without holding the Python GIL. + #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, reject_if_exist_same_key=false, callback=None, write_through=true, make_replica_task=true))] + fn put( + &self, + key: &str, + ptrs: Vec<(u8, u64, u32, u64, u32, Option)>, + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + callback: Option, + write_through: bool, + make_replica_task: bool, + py: Python, + ) -> PyObject { + fn put_inner( + client: &KvClient, + key: String, + ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + callback: Option, + write_through: bool, + make_replica_task: bool, + py: Python, + ) -> ApiResult { + if ptrs.len() > (u32::MAX as usize) { + return ApiResult::new_error(new_invalid_argument_error(py, "flat dict too large")); + } + + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime_handle = match client.runtime.as_ref() { + Some(v) => v.handle().clone(), + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let put_opts = build_put_optional_args( + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + ); + + let future = async move { + let result = unsafe { framework.kv_put_ptrs(&key, ptrs, put_opts).await }; + match result { + Ok(_) => Python::with_gil(|py| { + if let Some(cb) = callback { + let args = PyTuple::new_bound(py, &[new_none_success_instance(py)]); + let _ = cb.call1(py, args); + } + ApiResult::new_success(new_none_success_instance(py)) + }), + Err(e) => Python::with_gil(|py| { + let err_obj = crate::error::py_error_from_kv_error(py, &e, "Put failed"); + ApiResult::new_error(err_obj) + }), + } + }; + + let kv_future = KvFuture::new(future, runtime_handle, py); + match kv_future { + Ok(py_future) => ApiResult::new_success(py_future.into_any()), + Err(e) => ApiResult::new_error(new_general_error( + py, + &format!("Failed to create future: {}", e), + )), + } + } + + let key_owned = key.to_string(); + let mut ptrs_owned: Vec<(u8, usize, u32, u64, u32, Option)> = + Vec::with_capacity(ptrs.len()); + for (type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra) in ptrs.into_iter() { + let dict_key_ptr_usize: usize = match usize::try_from(dict_key_ptr) { + Ok(v) => v, + Err(_) => { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "dict_key_ptr out of range", + )) + .into_py_object(py); + } + }; + ptrs_owned.push(( + type_id, + dict_key_ptr_usize, + dict_key_len, + val_u64, + val_len, + extra, + )); + } + put_inner( + self, + key_owned, + ptrs_owned, + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + callback, + write_through, + make_replica_task, + py, + ) + .into_py_object(py) + } + + /// Put a key-value pair and wait for completion before returning. + #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false, reject_if_exist_same_key=false, write_through=true, make_replica_task=true))] + fn put_blocking( &self, key: &str, ptrs: Vec<(u8, u64, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, - callback: Option, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, py: Python, ) -> PyObject { - fn put_inner( + fn put_blocking_inner( client: &KvClient, key: String, ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, lease_id: Option, reject_if_inflight_same_key: bool, - callback: Option, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, py: Python, ) -> ApiResult { if ptrs.len() > (u32::MAX as usize) { @@ -2965,8 +5593,8 @@ impl KvClient { Ok(v) => v, Err(e) => return ApiResult::new_error(e), }; - let runtime_handle = match client.runtime.as_ref() { - Some(v) => v.handle().clone(), + let runtime = match client.runtime.as_ref() { + Some(v) => v, None => { return ApiResult::new_error(new_general_error( py, @@ -2974,99 +5602,561 @@ impl KvClient { )); } }; - let put_opts = { - let mut o = fluxon_kv::client_kv_api::PutOptionalArgs::new(); - if let Some(id) = lease_id { - o.0.push(fluxon_kv::client_kv_api::PutOptionalArg::LeaseId(id)); + let framework = borrow_stable_owner(&framework); + let put_opts = build_put_optional_args( + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + ); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async { + unsafe { framework.kv_put_ptrs(&key, ptrs, put_opts).await } + }) + }) { + Ok(v) => v, + Err(e) => Err(anyhow::anyhow!("runtime bridge failed: {}", e).into()), + }; + + match result { + Ok(_) => ApiResult::new_success(new_none_success_instance(py)), + Err(e) => { + let err_obj = crate::error::py_error_from_kv_error(py, &e, "Put failed"); + ApiResult::new_error(err_obj) + } + } + } + + let key_owned = key.to_string(); + let mut ptrs_owned: Vec<(u8, usize, u32, u64, u32, Option)> = + Vec::with_capacity(ptrs.len()); + for (type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra) in ptrs.into_iter() { + let dict_key_ptr_usize: usize = match usize::try_from(dict_key_ptr) { + Ok(v) => v, + Err(_) => { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "dict_key_ptr out of range", + )) + .into_py_object(py); + } + }; + ptrs_owned.push(( + type_id, + dict_key_ptr_usize, + dict_key_len, + val_u64, + val_len, + extra, + )); + } + put_blocking_inner( + self, + key_owned, + ptrs_owned, + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + py, + ) + .into_py_object(py) + } + + /// Put a batch of key-value pairs and wait for completion before returning. + /// + /// Each item uses the same flatdict encoding contract as `put_blocking()`. + #[pyo3(signature = (keys, ptrs_groups, lease_id=None, reject_if_inflight_same_key=false, reject_if_exist_same_key=false, concurrency=None, write_through=true, make_replica_task=true))] + fn batch_put_blocking( + &self, + keys: Vec, + ptrs_groups: Vec)>>, + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + concurrency: Option, + write_through: bool, + make_replica_task: bool, + py: Python, + ) -> PyObject { + fn batch_put_blocking_inner( + client: &KvClient, + keys: Vec, + ptrs_groups: Vec)>>, + lease_id: Option, + reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + concurrency: Option, + write_through: bool, + make_replica_task: bool, + py: Python, + ) -> ApiResult { + if keys.len() != ptrs_groups.len() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_put_blocking requires keys and ptrs_groups to have the same length", + )); + } + if keys.is_empty() { + return ApiResult::new_success(Vec::::new().into_py(py)); + } + + let batch_concurrency = match concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY) { + 0 => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_concurrency must be > 0", + )); } - if reject_if_inflight_same_key { - o.0.push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfInflightSameKey); + v => v, + }; + + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); } - o }; + let framework = borrow_stable_owner(&framework); + let keys_for_results = keys; + + let put_opts = build_put_optional_args( + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + ); - let future = async move { - let result = unsafe { framework.kv_put_ptrs(&key, ptrs, put_opts).await }; - match result { - Ok(_) => Python::with_gil(|py| { - if let Some(cb) = callback { - let args = PyTuple::new_bound(py, &[new_none_success_instance(py)]); - let _ = cb.call1(py, args); + if !framework.is_external_mode() { + let framework_for_backend = framework.clone(); + let keys_for_backend = keys_for_results.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + unsafe { + framework_for_backend + .client_kv_api_view() + .client_kv_api() + .inner() + .batch_put_flat_dict_ptrs( + keys_for_backend, + ptrs_groups, + put_opts, + batch_concurrency, + ) + .await } - ApiResult::new_success(new_none_success_instance(py)) - }), - Err(e) => Python::with_gil(|py| { - let err_obj = crate::error::py_error_from_kv_error(py, &e, "Put failed"); - ApiResult::new_error(err_obj) - }), + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + return Python::with_gil(|py| match backend_result { + Ok(per_item_results) => { + let out = PyList::empty_bound(py); + for (key, item_result) in keys_for_results + .into_iter() + .zip(per_item_results.into_iter()) + { + let item = match item_result { + Ok(()) => crate::error::new_result_success( + py, + new_none_success_instance(py), + ), + Err(err) => crate::error::new_result_error( + py, + crate::error::py_error_from_kv_error( + py, + &err, + &format!("Put failed for key {}", key), + ), + ), + }; + out.append(item).expect("append batch_put result"); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch put failed", + )), + }); + } + + let framework_for_backend = framework.clone(); + let keys_for_backend = keys_for_results.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + unsafe { + framework_for_backend + .external_client_api_view() + .external_client_api() + .inner() + .batch_put_flat_dict_ptrs( + keys_for_backend, + ptrs_groups, + put_opts, + batch_concurrency, + ) + .await + } + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + Python::with_gil(|py| match backend_result { + Ok(per_item_results) => { + let out = PyList::empty_bound(py); + for (key, item_result) in keys_for_results + .into_iter() + .zip(per_item_results.into_iter()) + { + let item = match item_result { + Ok(()) => { + crate::error::new_result_success(py, new_none_success_instance(py)) + } + Err(err) => { + let err_obj = crate::error::py_error_from_kv_error( + py, + &err, + &format!("Put failed for key {}", key), + ); + crate::error::new_result_error(py, err_obj) + } + }; + out.append(item).expect("append batch_put result"); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch put failed", + )), + }) + } + + let mut ptrs_groups_owned: Vec)>> = + Vec::with_capacity(ptrs_groups.len()); + for ptrs in ptrs_groups.into_iter() { + let mut ptrs_owned: Vec<(u8, usize, u32, u64, u32, Option)> = + Vec::with_capacity(ptrs.len()); + for (type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra) in ptrs.into_iter() { + let dict_key_ptr_usize: usize = match usize::try_from(dict_key_ptr) { + Ok(v) => v, + Err(_) => { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "dict_key_ptr out of range", + )) + .into_py_object(py); + } + }; + ptrs_owned.push(( + type_id, + dict_key_ptr_usize, + dict_key_len, + val_u64, + val_len, + extra, + )); + } + if ptrs_owned.len() > (u32::MAX as usize) { + return ApiResult::::new_error(new_invalid_argument_error( + py, + "flat dict too large", + )) + .into_py_object(py); + } + ptrs_groups_owned.push(ptrs_owned); + } + + batch_put_blocking_inner( + self, + keys, + ptrs_groups_owned, + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + concurrency, + write_through, + make_replica_task, + py, + ) + .into_py_object(py) + } + + /// Register one caller-owned GPU staging range with the transfer engine. + #[pyo3(signature = (ptr, len, device_id))] + fn register_gpu_buffer(&self, ptr: u64, len: u64, device_id: u32, py: Python) -> PyObject { + fn register_gpu_buffer_inner( + client: &KvClient, + ptr: u64, + len: u64, + device_id: u32, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "register_gpu_buffer is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework_for_register = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_register + .client_seg_pool_view() + .client_transfer_engine() + .register_gpu_memory(ptr, len, device_id) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match result { + Ok(registration) => { + let out = PyDict::new_bound(py); + out.set_item("registration_id", registration.registration_id) + .expect("set registration_id"); + out.set_item("ptr", registration.addr).expect("set ptr"); + out.set_item("len", registration.len).expect("set len"); + out.set_item("device_id", registration.device_id) + .expect("set device_id"); + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "register_gpu_buffer failed", + )), + } + } + + register_gpu_buffer_inner(self, ptr, len, device_id, py).into_py_object(py) + } + + /// Unregister a GPU staging range after every destination guard is released. + #[pyo3(signature = (registration_id))] + fn unregister_gpu_buffer(&self, registration_id: u64, py: Python) -> PyObject { + fn unregister_gpu_buffer_inner( + client: &KvClient, + registration_id: u64, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "unregister_gpu_buffer is supported only in external-client mode", + )); + } + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework_for_unregister = framework.clone(); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_unregister + .client_seg_pool_view() + .client_transfer_engine() + .unregister_gpu_memory(registration_id) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); } }; + match result { + Ok(()) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "unregister_gpu_buffer failed", + )), + } + } - let kv_future = KvFuture::new(future, runtime_handle, py); - match kv_future { - Ok(py_future) => ApiResult::new_success(py_future.into_any()), - Err(e) => ApiResult::new_error(new_general_error( + unregister_gpu_buffer_inner(self, registration_id, py).into_py_object(py) + } + + /// Validate one destination against the exact GPU registration generation. + #[pyo3(signature = (registration_id, ptr, capacity))] + fn validate_gpu_destination( + &self, + registration_id: u64, + ptr: u64, + capacity: u64, + py: Python, + ) -> PyObject { + fn validate_gpu_destination_inner( + client: &KvClient, + registration_id: u64, + ptr: u64, + capacity: u64, + py: Python, + ) -> ApiResult { + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + if !framework.is_external_mode() { + return ApiResult::new_error(new_invalid_argument_error( py, - &format!("Failed to create future: {}", e), + "validate_gpu_destination is supported only in external-client mode", + )); + } + match framework + .client_seg_pool_view() + .client_transfer_engine() + .validate_gpu_destination(registration_id, ptr, capacity) + { + Ok(_guard) => ApiResult::new_success(0i32.into_py(py)), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "validate_gpu_destination failed", )), } } - let key_owned = key.to_string(); - let mut ptrs_owned: Vec<(u8, usize, u32, u64, u32, Option)> = - Vec::with_capacity(ptrs.len()); - for (type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra) in ptrs.into_iter() { - let dict_key_ptr_usize: usize = match usize::try_from(dict_key_ptr) { + validate_gpu_destination_inner(self, registration_id, ptr, capacity, py).into_py_object(py) + } + + /// Register a caller-owned host buffer range used by batch zero-copy paths. + #[pyo3(signature = (ptr, len))] + fn register_buffer(&self, ptr: u64, len: u64, py: Python) -> PyObject { + fn register_buffer_inner( + client: &KvClient, + ptr: u64, + len: u64, + py: Python, + ) -> ApiResult { + let ptr_usize = match usize::try_from(ptr) { Ok(v) => v, Err(_) => { - return ApiResult::::new_error(new_invalid_argument_error( + return ApiResult::new_error(new_invalid_argument_error( py, - "dict_key_ptr out of range", - )) - .into_py_object(py); + "register_buffer ptr out of range", + )); } }; - ptrs_owned.push(( - type_id, - dict_key_ptr_usize, - dict_key_len, - val_u64, - val_len, - extra, - )); + let len_usize = match usize::try_from(len) { + Ok(v) => v, + Err(_) => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "register_buffer len out of range", + )); + } + }; + if ptr_usize == 0 && len_usize > 0 { + return ApiResult::new_error(new_invalid_argument_error( + py, + "register_buffer ptr must be non-zero when len > 0", + )); + } + let mut guard = client + .registered_buffers + .write() + .expect("registered_buffers poisoned"); + guard.register(ptr_usize, len_usize); + ApiResult::new_success(0i32.into_py(py)) } - put_inner( - self, - key_owned, - ptrs_owned, - lease_id, - reject_if_inflight_same_key, - callback, - py, - ) - .into_py_object(py) + + register_buffer_inner(self, ptr, len, py).into_py_object(py) } - /// Put a key-value pair and wait for completion before returning. - #[pyo3(signature = (key, ptrs, lease_id=None, reject_if_inflight_same_key=false))] - fn put_blocking( + /// Native blocking batch put path for payload pointers. + #[pyo3(signature = (keys, payload_ptrs, payload_sizes, lease_id=None, reject_if_inflight_same_key=false, reject_if_exist_same_key=false, write_through=true, make_replica_task=true))] + fn batch_put_from( &self, - key: &str, - ptrs: Vec<(u8, u64, u32, u64, u32, Option)>, + keys: Vec, + payload_ptrs: Vec, + payload_sizes: Vec, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, py: Python, ) -> PyObject { - fn put_blocking_inner( + fn batch_put_from_inner( client: &KvClient, - key: String, - ptrs: Vec<(u8, usize, u32, u64, u32, Option)>, + keys: Vec, + payload_ptrs: Vec, + payload_sizes: Vec, lease_id: Option, reject_if_inflight_same_key: bool, + reject_if_exist_same_key: bool, + write_through: bool, + make_replica_task: bool, py: Python, ) -> ApiResult { - if ptrs.len() > (u32::MAX as usize) { - return ApiResult::new_error(new_invalid_argument_error(py, "flat dict too large")); + if keys.len() != payload_ptrs.len() || keys.len() != payload_sizes.len() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_put_from requires keys, payload_ptrs, and payload_sizes to have the same length", + )); + } + if keys.is_empty() { + return ApiResult::new_success(Vec::::new().into_py(py)); } let framework = match require_kv_framework_api(client, py) { @@ -3082,65 +6172,168 @@ impl KvClient { )); } }; - let framework = borrow_stable_owner(&framework); - let mut put_opts = fluxon_kv::client_kv_api::PutOptionalArgs::new(); - if let Some(id) = lease_id { - put_opts - .0 - .push(fluxon_kv::client_kv_api::PutOptionalArg::LeaseId(id)); + + let registered = client + .registered_buffers + .read() + .expect("registered_buffers poisoned"); + let mut payload_ptrs_usize = Vec::with_capacity(payload_ptrs.len()); + for (&ptr, &size_u32) in payload_ptrs.iter().zip(payload_sizes.iter()) { + let ptr_usize = match usize::try_from(ptr) { + Ok(v) => v, + Err(_) => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_put_from payload_ptr out of range", + )); + } + }; + let size = size_u32 as usize; + if ptr_usize == 0 && size > 0 { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_put_from payload_ptr must be non-zero when payload_size > 0", + )); + } + if !registered.ranges.is_empty() && !registered.contains(ptr_usize, size) { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_put_from payload range is not inside any registered buffer", + )); + } + payload_ptrs_usize.push(ptr_usize); } - if reject_if_inflight_same_key { - put_opts - .0 - .push(fluxon_kv::client_kv_api::PutOptionalArg::RejectIfInflightSameKey); + drop(registered); + + let framework = borrow_stable_owner(&framework); + let put_opts = build_put_optional_args( + lease_id, + reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, + ); + + let keys_for_backend = keys.clone(); + let ptrs_groups = keys_for_backend + .iter() + .zip(payload_ptrs_usize.iter()) + .zip(payload_sizes.iter()) + .map(|((_key, payload_ptr), payload_size)| { + vec![( + 5u8, + b"payload".as_ptr() as usize, + 7u32, + *payload_ptr as u64, + *payload_size, + None, + )] + }) + .collect::>(); + let batch_concurrency = DEFAULT_PYO3_BATCH_CONCURRENCY; + + if !framework.is_external_mode() { + let framework_for_backend = framework.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + unsafe { + framework_for_backend + .client_kv_api_view() + .client_kv_api() + .inner() + .batch_put_flat_dict_ptrs( + keys_for_backend, + ptrs_groups, + put_opts, + batch_concurrency, + ) + .await + } + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + return match backend_result { + Ok(per_item_results) => ApiResult::new_success( + per_item_results + .into_iter() + .map(|item| match item { + Ok(()) => 0, + Err(err) => kv_error_to_ret_code(&err), + }) + .collect::>() + .into_py(py), + ), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch put failed", + )), + }; } - let result = match py.allow_threads(|| { - runtime.run_async_from_sync(async { - unsafe { framework.kv_put_ptrs(&key, ptrs, put_opts).await } + + let framework_for_backend = framework.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + unsafe { + framework_for_backend + .external_client_api_view() + .external_client_api() + .inner() + .batch_put_flat_dict_ptrs( + keys_for_backend, + ptrs_groups, + put_opts, + batch_concurrency, + ) + .await + } }) }) { Ok(v) => v, - Err(e) => Err(anyhow::anyhow!("runtime bridge failed: {}", e).into()), - }; - - match result { - Ok(_) => ApiResult::new_success(new_none_success_instance(py)), Err(e) => { - let err_obj = crate::error::py_error_from_kv_error(py, &e, "Put failed"); - ApiResult::new_error(err_obj) - } - } - } - - let key_owned = key.to_string(); - let mut ptrs_owned: Vec<(u8, usize, u32, u64, u32, Option)> = - Vec::with_capacity(ptrs.len()); - for (type_id, dict_key_ptr, dict_key_len, val_u64, val_len, extra) in ptrs.into_iter() { - let dict_key_ptr_usize: usize = match usize::try_from(dict_key_ptr) { - Ok(v) => v, - Err(_) => { - return ApiResult::::new_error(new_invalid_argument_error( + return ApiResult::new_error(new_general_error( py, - "dict_key_ptr out of range", - )) - .into_py_object(py); + &format!("runtime bridge failed: {}", e), + )); } }; - ptrs_owned.push(( - type_id, - dict_key_ptr_usize, - dict_key_len, - val_u64, - val_len, - extra, - )); + + match backend_result { + Ok(per_item_results) => ApiResult::new_success( + per_item_results + .into_iter() + .map(|item| match item { + Ok(()) => 0, + Err(err) => kv_error_to_ret_code(&err), + }) + .collect::>() + .into_py(py), + ), + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch put failed", + )), + } } - put_blocking_inner( + + batch_put_from_inner( self, - key_owned, - ptrs_owned, + keys, + payload_ptrs, + payload_sizes, lease_id, reject_if_inflight_same_key, + reject_if_exist_same_key, + write_through, + make_replica_task, py, ) .into_py_object(py) @@ -3307,19 +6500,208 @@ impl KvClient { let pyo3_external = ExternalMemHolder::new(external_mem_holder); pyo3_external.into_py_mem_holder(py) } - Ok(KvGetResult::External(None)) => ApiResult::new_error(new_key_not_found_error( + Ok(KvGetResult::External(None)) => ApiResult::new_error(new_key_not_found_error( + py, + &format!("Key not found: {}", key), + Some(&key), + )), + Err(e) => { + let err_obj = crate::error::py_error_from_kv_error(py, &e, "Get failed"); + ApiResult::new_error(err_obj) + } + } + } + + get_blocking_inner(self, key, py).into_py_object(py) + } + + /// Get a batch of keys and wait for completion before returning. + /// + /// Returns a Python list of per-item `Result` objects, aligned with + /// `get_blocking()` semantics. + #[pyo3(signature = (keys, concurrency=None))] + fn batch_get_blocking( + &self, + keys: Vec, + concurrency: Option, + py: Python, + ) -> PyObject { + fn batch_get_blocking_inner( + client: &KvClient, + keys: Vec, + concurrency: Option, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_success(Vec::::new().into_py(py)); + } + + let batch_concurrency = match concurrency.unwrap_or(DEFAULT_PYO3_BATCH_CONCURRENCY) { + 0 => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_concurrency must be > 0", + )); + } + v => v, + }; + + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework = borrow_stable_owner(&framework); + let keys_for_results = keys; + + if !framework.is_external_mode() { + let framework_for_backend = framework.clone(); + let keys_for_backend = keys_for_results.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_backend + .client_kv_api_view() + .client_kv_api() + .inner() + .batch_get(keys_for_backend, batch_concurrency) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + return Python::with_gil(|py| match backend_result { + Ok(per_item_results) => { + let out = PyList::empty_bound(py); + for (key, item_result) in keys_for_results + .into_iter() + .zip(per_item_results.into_iter()) + { + let item = match item_result { + Ok(Some((rust_holder, _remote_info))) => { + let mem_holder = MemHolder::new(rust_holder); + match mem_holder.into_py_mem_holder(py) { + ApiResult::Success(py_holder) => { + crate::error::new_result_success(py, py_holder) + } + ApiResult::Error(err) => { + crate::error::new_result_error(py, err) + } + } + } + Ok(None) => crate::error::new_result_error( + py, + new_key_not_found_error( + py, + &format!("Key not found: {}", key), + Some(&key), + ), + ), + Err(err) => crate::error::new_result_error( + py, + crate::error::py_error_from_kv_error( + py, + &err, + &format!("Get failed for key {}", key), + ), + ), + }; + out.append(item).expect("append batch_get result"); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch get failed", + )), + }); + } + + let framework_for_backend = framework.clone(); + let keys_for_backend = keys_for_results.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_backend + .external_client_api_view() + .external_client_api() + .inner() + .batch_get(keys_for_backend, batch_concurrency) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + Python::with_gil(|py| match backend_result { + Ok(per_item_results) => { + let out = PyList::empty_bound(py); + for (key, item_result) in keys_for_results + .into_iter() + .zip(per_item_results.into_iter()) + { + let item = match item_result { + Ok(Some(external_mem_holder)) => { + let pyo3_external = ExternalMemHolder::new(external_mem_holder); + match pyo3_external.into_py_mem_holder(py) { + ApiResult::Success(py_holder) => { + crate::error::new_result_success(py, py_holder) + } + ApiResult::Error(err) => { + crate::error::new_result_error(py, err) + } + } + } + Ok(None) => crate::error::new_result_error( + py, + new_key_not_found_error( + py, + &format!("Key not found: {}", key), + Some(&key), + ), + ), + Err(err) => crate::error::new_result_error( + py, + crate::error::py_error_from_kv_error( + py, + &err, + &format!("Get failed for key {}", key), + ), + ), + }; + out.append(item).expect("append batch_get result"); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( py, - &format!("Key not found: {}", key), - Some(&key), + &err, + "Batch get failed", )), - Err(e) => { - let err_obj = crate::error::py_error_from_kv_error(py, &e, "Get failed"); - ApiResult::new_error(err_obj) - } - } + }) } - get_blocking_inner(self, key, py).into_py_object(py) + batch_get_blocking_inner(self, keys, concurrency, py).into_py_object(py) } /// Delete a key (synchronous from Python; only put/get use KvFuture) @@ -3416,9 +6798,84 @@ impl KvClient { metrics_snapshot_inner(self, py).into_py_object(py) } + /// Async snapshot for Fluxon KV locality and IO counters. + fn observability_snapshot_async(&self, py: Python) -> PyObject { + fn observability_snapshot_async_inner( + client: &KvClient, + py: Python, + ) -> ApiResult { + let _framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime_handle = match client.runtime.as_ref() { + Some(v) => v.handle().clone(), + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let locality_counters = client.locality_counters.clone(); + let future = async move { + let metrics = locality_counters.snapshot(); + Python::with_gil(|py| { + let out = PyDict::new_bound(py); + out.set_item("l2_local_hit_pages", metrics.l2_local_hit_pages) + .expect("set l2_local_hit_pages"); + out.set_item("l2_local_hit_bytes", metrics.l2_local_hit_bytes) + .expect("set l2_local_hit_bytes"); + out.set_item("l2_remote_hit_pages", metrics.l2_remote_hit_pages) + .expect("set l2_remote_hit_pages"); + out.set_item("l2_remote_hit_bytes", metrics.l2_remote_hit_bytes) + .expect("set l2_remote_hit_bytes"); + let io = PyDict::new_bound(py); + for (key, item) in [ + ("put_local", metrics.put_local.clone()), + ("put_remote", metrics.put_remote.clone()), + ("get_local", metrics.get_local.clone()), + ("get_remote", metrics.get_remote.clone()), + ] { + let entry = PyDict::new_bound(py); + entry + .set_item("op_count", item.op_count) + .expect("set op_count"); + entry.set_item("bytes", item.bytes).expect("set bytes"); + entry + .set_item("transfer_us", item.transfer_us) + .expect("set transfer_us"); + entry + .set_item("bandwidth_gbps", item.bandwidth_gbps()) + .expect("set bandwidth_gbps"); + io.set_item(key, entry).expect("set io entry"); + } + out.set_item("io", io).expect("set io"); + ApiResult::new_success(out.into_py(py)) + }) + }; + match KvFuture::new(future, runtime_handle, py) { + Ok(py_future) => ApiResult::new_success(py_future.into_any()), + Err(e) => ApiResult::new_error(new_general_error( + py, + &format!("Failed to create future: {}", e), + )), + } + } + + observability_snapshot_async_inner(self, py).into_py_object(py) + } + /// Check if a key exists (synchronous; returns bool wrapped in Result) - fn is_exist(&self, key: String, py: Python) -> PyObject { - fn is_exist_inner(client: &KvClient, key: String, py: Python) -> ApiResult { + #[pyo3(signature = (key, allow_local_snapshot=false))] + fn is_exist(&self, key: String, allow_local_snapshot: bool, py: Python) -> PyObject { + fn is_exist_inner( + client: &KvClient, + key: String, + allow_local_snapshot: bool, + py: Python, + ) -> ApiResult { + let _ = allow_local_snapshot; let framework = match require_kv_framework_api(client, py) { Ok(v) => v, Err(e) => return ApiResult::new_error(e), @@ -3450,7 +6907,298 @@ impl KvClient { }), } } - is_exist_inner(self, key, py).into_py_object(py) + is_exist_inner(self, key, allow_local_snapshot, py).into_py_object(py) + } + + /// Native blocking batch get path writing payload bytes into caller-provided pointers. + #[pyo3(signature = (keys, payload_ptrs, payload_capacities))] + fn batch_get_into( + &self, + keys: Vec, + payload_ptrs: Vec, + payload_capacities: Vec, + py: Python, + ) -> PyObject { + fn batch_get_into_inner( + client: &KvClient, + keys: Vec, + payload_ptrs: Vec, + payload_capacities: Vec, + py: Python, + ) -> ApiResult { + if keys.len() != payload_ptrs.len() || keys.len() != payload_capacities.len() { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_get_into requires keys, payload_ptrs, and payload_capacities to have the same length", + )); + } + if keys.is_empty() { + return ApiResult::new_success(Vec::::new().into_py(py)); + } + + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + + let registered = client + .registered_buffers + .read() + .expect("registered_buffers poisoned"); + let mut payload_ptrs_usize = Vec::with_capacity(payload_ptrs.len()); + for (&ptr, &cap_u32) in payload_ptrs.iter().zip(payload_capacities.iter()) { + let ptr_usize = match usize::try_from(ptr) { + Ok(v) => v, + Err(_) => { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_get_into payload_ptr out of range", + )); + } + }; + if ptr_usize == 0 && cap_u32 > 0 { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_get_into payload_ptr must be non-zero when payload_capacity > 0", + )); + } + let cap = cap_u32 as usize; + if !registered.ranges.is_empty() && !registered.contains(ptr_usize, cap) { + return ApiResult::new_error(new_invalid_argument_error( + py, + "batch_get_into destination range is not inside any registered buffer", + )); + } + payload_ptrs_usize.push(ptr_usize); + } + drop(registered); + + let framework = borrow_stable_owner(&framework); + let batch_concurrency = DEFAULT_PYO3_BATCH_CONCURRENCY; + + if !framework.is_external_mode() { + let framework_for_backend = framework.clone(); + let keys_for_backend = keys.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_backend + .client_kv_api_view() + .client_kv_api() + .inner() + .batch_get(keys_for_backend, batch_concurrency) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + return match backend_result { + Ok(per_item_results) => { + let mut out = Vec::with_capacity(per_item_results.len()); + for (((key, payload_ptr), payload_capacity), item_result) in keys + .into_iter() + .zip(payload_ptrs_usize.into_iter()) + .zip(payload_capacities.into_iter()) + .zip(per_item_results.into_iter()) + { + let code = match item_result { + Ok(Some((holder, _remote_info))) => { + let holder_result = KvGetResult::Owner(Some(holder)); + let payload = + match payload_slice_from_get_result(&key, &holder_result) { + Ok(payload) => payload, + Err(err) => { + out.push(kv_error_to_ret_code(&err)); + continue; + } + }; + if payload.len() > payload_capacity as usize { + -(4003i32) + } else { + if !payload.is_empty() { + unsafe { + std::ptr::copy_nonoverlapping( + payload.as_ptr(), + payload_ptr as *mut u8, + payload.len(), + ); + } + } + payload.len() as i32 + } + } + Ok(None) => kv_error_to_ret_code(&CoreKvError::Api( + CoreApiError::KeyNotFound { key }, + )), + Err(err) => kv_error_to_ret_code(&err), + }; + out.push(code); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch get failed", + )), + }; + } + + let framework_for_backend = framework.clone(); + let keys_for_backend = keys.clone(); + let backend_result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework_for_backend + .external_client_api_view() + .external_client_api() + .inner() + .batch_get(keys_for_backend, batch_concurrency) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + + match backend_result { + Ok(per_item_results) => { + let mut out = Vec::with_capacity(per_item_results.len()); + for (((key, payload_ptr), payload_capacity), item_result) in keys + .into_iter() + .zip(payload_ptrs_usize.into_iter()) + .zip(payload_capacities.into_iter()) + .zip(per_item_results.into_iter()) + { + let code = + match item_result { + Ok(Some(holder)) => { + let holder_result = KvGetResult::External(Some(holder)); + let payload = + match payload_slice_from_get_result(&key, &holder_result) { + Ok(payload) => payload, + Err(err) => { + out.push(kv_error_to_ret_code(&err)); + continue; + } + }; + if payload.len() > payload_capacity as usize { + -(4003i32) + } else { + if !payload.is_empty() { + unsafe { + std::ptr::copy_nonoverlapping( + payload.as_ptr(), + payload_ptr as *mut u8, + payload.len(), + ); + } + } + payload.len() as i32 + } + } + Ok(None) => kv_error_to_ret_code(&CoreKvError::Api( + CoreApiError::KeyNotFound { key }, + )), + Err(err) => kv_error_to_ret_code(&err), + }; + out.push(code); + } + ApiResult::new_success(out.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch get failed", + )), + } + } + + batch_get_into_inner(self, keys, payload_ptrs, payload_capacities, py).into_py_object(py) + } + + /// Native blocking batch existence check. + #[pyo3(signature = (keys, allow_local_snapshot=false))] + fn batch_is_exist( + &self, + keys: Vec, + allow_local_snapshot: bool, + py: Python, + ) -> PyObject { + fn batch_is_exist_inner( + client: &KvClient, + keys: Vec, + allow_local_snapshot: bool, + py: Python, + ) -> ApiResult { + if keys.is_empty() { + return ApiResult::new_success(Vec::::new().into_py(py)); + } + let framework = match require_kv_framework_api(client, py) { + Ok(v) => v, + Err(e) => return ApiResult::new_error(e), + }; + let runtime = match client.runtime.as_ref() { + Some(v) => v, + None => { + return ApiResult::new_error(new_general_error( + py, + "Client runtime is missing", + )); + } + }; + let framework = borrow_stable_owner(&framework); + let result = match py.allow_threads(|| { + runtime.run_async_from_sync(async move { + framework + .kv_batch_is_exist(keys, allow_local_snapshot) + .await + }) + }) { + Ok(v) => v, + Err(e) => { + return ApiResult::new_error(new_general_error( + py, + &format!("runtime bridge failed: {}", e), + )); + } + }; + match result { + Ok(exists_list) => { + let ints = exists_list + .into_iter() + .map(|exists| if exists { 1i32 } else { 0i32 }) + .collect::>(); + ApiResult::new_success(ints.into_py(py)) + } + Err(err) => ApiResult::new_error(crate::error::py_error_from_kv_error( + py, + &err, + "Batch existence check failed", + )), + } + } + + batch_is_exist_inner(self, keys, allow_local_snapshot, py).into_py_object(py) } /// Count number of keys whose name starts with the given prefix. @@ -3526,14 +7274,11 @@ impl KvClient { timeout_ms: u64, py: Python, ) -> ApiResult { - if timeout_ms < fluxon_kv::user_rpc::USER_RPC_MIN_TIMEOUT_MS { - return ApiResult::new_error(new_invalid_argument_error( + if let Err(err) = fluxon_kv::user_rpc::validate_timeout_ms(timeout_ms) { + return ApiResult::new_error(crate::error::py_error_from_kv_error( py, - &format!( - "timeout_ms must be >= {} (got {})", - fluxon_kv::user_rpc::USER_RPC_MIN_TIMEOUT_MS, - timeout_ms - ), + &err, + "Invalid RPC timeout", )); } @@ -3649,14 +7394,13 @@ impl KvClient { timeout_ms: u64, py: Python, ) -> ApiResult { - const MIN_TIMEOUT_MS: u64 = 10_000; - if timeout_ms < MIN_TIMEOUT_MS { - return ApiResult::new_error(new_invalid_argument_error( + if let Err(err) = fluxon_kv::p2p::msg_pack::validate_explicit_rpc_timeout_ms(timeout_ms) + { + let err = CoreKvError::from(err); + return ApiResult::new_error(crate::error::py_error_from_kv_error( py, - &format!( - "timeout_ms must be >= {} (got {})", - MIN_TIMEOUT_MS, timeout_ms - ), + &err, + "Invalid RPC timeout", )); } if target_instance_key.trim().is_empty() { @@ -4018,7 +7762,7 @@ impl Drop for KvMaster { /// This function creates a master, runs it until Ctrl+C, then shuts down #[pyfunction] #[pyo3(signature = (config=None))] -fn run_master_blocking(config: Option<&Bound<'_, PyAny>>, py: Python) -> PyObject { +fn run_master_blocking(config: Option<&Bound<'_, PyAny>>, py: Python) -> PyResult { fn run_master_inner(config: Option<&Bound<'_, PyAny>>, py: Python) -> ApiResult { // Debug config println!("🛠️ Master init configuration: {:?}", config); @@ -4140,7 +7884,20 @@ fn run_master_blocking(config: Option<&Bound<'_, PyAny>>, py: Python) -> PyObjec out } - run_master_inner(config, py).into_py_object(py) + let result = run_master_inner(config, py); + // The same SIGINT wakes Tokio's ctrl_c future and becomes a pending Python + // KeyboardInterrupt while the GIL is released. Once the framework has + // completed its graceful shutdown, consume only that expected interrupt + // before constructing the Python Result object. Other Python exceptions + // remain visible to the caller. + if let Err(err) = py.check_signals() { + if err.is_instance_of::(py) { + tracing::debug!("consumed KeyboardInterrupt after graceful KV master shutdown"); + } else { + return Err(err); + } + } + Ok(result.into_py_object(py)) } /// Python module definition @@ -4148,11 +7905,16 @@ fn run_master_blocking(config: Option<&Bound<'_, PyAny>>, py: Python) -> PyObjec #[pyo3(name = "fluxon_pyo3")] fn fluxon_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { init_dynamic_libraries()?; + m.add( + "MIN_EXPLICIT_RPC_TIMEOUT_MS", + fluxon_kv::p2p::msg_pack::MIN_EXPLICIT_RPC_TIMEOUT_MS, + )?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -4164,6 +7926,7 @@ fn fluxon_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(run_master_blocking, m)?)?; + m.add_function(wrap_pyfunction!(decode_flat_dict_payload, m)?)?; m.add_function(wrap_pyfunction!(monitor_render_cli, m)?)?; m.add_function(wrap_pyfunction!(monitor_render_web, m)?)?; m.add_function(wrap_pyfunction!(fluxon_ops_controller_blocking, m)?)?; @@ -4180,18 +7943,151 @@ fn fluxon_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { #[cfg(test)] mod tests { use super::{ - bundled_driver_names_from_entries, configure_bundled_rdmav_driver_env, + build_put_optional_args, bundled_driver_names_from_entries, + configure_bundled_rdmav_driver_env, defer_drop_to_runtime, discover_bundled_ibverbs_driver_config, extract_fluxon_pyo3_libs_root_from_loaded_library_line, loaded_fluxon_pyo3_libs_roots, - parse_bundled_ibverbs_driver_name, sanitize_bundled_ld_library_path_entries, + normalize_put_atomic_group_lens, parse_bundled_ibverbs_driver_name, + resolve_make_replica_task_mask, sanitize_bundled_ld_library_path_entries, set_authoritative_bundled_ld_library_path, validate_single_fluxon_pyo3_libs_root, }; use std::path::{Path, PathBuf}; - use std::sync::{Mutex, OnceLock}; - use std::time::{SystemTime, UNIX_EPOCH}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex, OnceLock}; + use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; static ENV_LOCK: OnceLock> = OnceLock::new(); + struct DropProbe(Arc); + + impl Drop for DropProbe { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn defer_drop_to_runtime_eventually_drops_value() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let dropped = Arc::new(AtomicUsize::new(0)); + defer_drop_to_runtime(Some(&runtime), DropProbe(dropped.clone())); + + let deadline = Instant::now() + Duration::from_secs(2); + while dropped.load(Ordering::SeqCst) == 0 && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(1)); + } + assert_eq!(dropped.load(Ordering::SeqCst), 1); + } + + #[test] + fn defer_drop_without_runtime_is_synchronous() { + let dropped = Arc::new(AtomicUsize::new(0)); + defer_drop_to_runtime(None, DropProbe(dropped.clone())); + assert_eq!(dropped.load(Ordering::SeqCst), 1); + } + + #[test] + fn put_options_control_async_replica_task() { + let write_through = build_put_optional_args(None, false, false, true, true); + let write_back = build_put_optional_args(None, false, false, false, true); + let local_only = build_put_optional_args(None, false, false, false, false); + + assert!(!write_through.make_replica_task()); + assert!(write_back.make_replica_task()); + assert!(!local_only.make_replica_task()); + } + + #[test] + fn local_fast_put_replica_mask_broadcasts_scalar() { + assert_eq!( + resolve_make_replica_task_mask(3, false, true, None, &[1, 1, 1]).unwrap(), + vec![true, true, true] + ); + assert_eq!( + resolve_make_replica_task_mask(3, false, false, None, &[1, 1, 1]).unwrap(), + vec![false, false, false] + ); + } + + #[test] + fn local_fast_put_replica_mask_selects_items() { + assert_eq!( + resolve_make_replica_task_mask( + 4, + false, + true, + Some(vec![true, false, true, false]), + &[1, 1, 1, 1], + ) + .unwrap(), + vec![true, false, true, false] + ); + } + + #[test] + fn local_fast_put_replica_mask_respects_global_gates() { + let requested = Some(vec![true, false, true]); + assert_eq!( + resolve_make_replica_task_mask(3, true, true, requested.clone(), &[1, 1, 1]).unwrap(), + vec![false, false, false] + ); + assert_eq!( + resolve_make_replica_task_mask(3, false, false, requested, &[1, 1, 1]).unwrap(), + vec![false, false, false] + ); + } + + #[test] + fn local_fast_put_replica_mask_rejects_length_mismatch() { + let err = + resolve_make_replica_task_mask(3, false, true, Some(vec![true, false]), &[1, 1, 1]) + .unwrap_err(); + assert_eq!( + err, + "make_replica_task_mask length must match keys length; keys=3 mask=2" + ); + } + + #[test] + fn local_fast_put_atomic_groups_accept_valid_partition() { + assert_eq!( + normalize_put_atomic_group_lens(5, Some(vec![2, 3])).unwrap(), + vec![2, 3] + ); + assert_eq!( + normalize_put_atomic_group_lens(3, None).unwrap(), + vec![1, 1, 1] + ); + } + + #[test] + fn local_fast_put_atomic_groups_reject_zero_and_sum_mismatch() { + assert_eq!( + normalize_put_atomic_group_lens(3, Some(vec![1, 0, 2])).unwrap_err(), + "atomic_group_lens entries must be > 0; index=1 got=0" + ); + assert_eq!( + normalize_put_atomic_group_lens(3, Some(vec![1, 1])).unwrap_err(), + "atomic_group_lens must sum to keys length; sum=2 keys=3" + ); + } + + #[test] + fn local_fast_put_replica_mask_rejects_split_atomic_group() { + let err = resolve_make_replica_task_mask( + 5, + false, + true, + Some(vec![true, true, false, true, true]), + &[2, 3], + ) + .unwrap_err(); + assert_eq!( + err, + "make_replica_task_mask must be uniform within each atomic group; group_index=1 offset=2 len=3" + ); + } + struct EnvSnapshot { key: &'static str, value: Option, diff --git a/fluxon_rs/fluxon_util/src/dev_config.rs b/fluxon_rs/fluxon_util/src/dev_config.rs index c860910..d3f5953 100644 --- a/fluxon_rs/fluxon_util/src/dev_config.rs +++ b/fluxon_rs/fluxon_util/src/dev_config.rs @@ -1,4 +1,4 @@ -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde_yaml::Value; use std::fs; use std::path::{Path, PathBuf}; diff --git a/fluxon_rs/fluxon_util/src/fixed_slab_allocator.rs b/fluxon_rs/fluxon_util/src/fixed_slab_allocator.rs new file mode 100644 index 0000000..fe9c6ce --- /dev/null +++ b/fluxon_rs/fluxon_util/src/fixed_slab_allocator.rs @@ -0,0 +1,245 @@ +use parking_lot::Mutex; +use std::error::Error; +use std::fmt::{Display, Formatter}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FixedSlabAllocatorError { + ZeroCapacity, + DuplicateSlot { slot: u32 }, + SlotOutOfBounds { slot: u32, capacity: u32 }, + SlotNotAllocated { slot: u32 }, +} + +impl Display for FixedSlabAllocatorError { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::ZeroCapacity => write!(formatter, "fixed slab capacity must be positive"), + Self::DuplicateSlot { slot } => { + write!( + formatter, + "fixed slab release contains duplicate slot {slot}" + ) + } + Self::SlotOutOfBounds { slot, capacity } => write!( + formatter, + "fixed slab slot {slot} is outside capacity {capacity}" + ), + Self::SlotNotAllocated { slot } => { + write!( + formatter, + "fixed slab slot {slot} is not currently allocated" + ) + } + } + } +} + +impl Error for FixedSlabAllocatorError {} + +#[derive(Debug)] +struct FixedSlabState { + free_slots: Vec, + allocated: Vec, + validation_marks: Vec, + validation_epoch: u64, +} + +/// Thread-safe allocator for a fixed number of equally sized slab slots. +/// +/// Reservations are all-or-none. Releases validate every supplied slot before +/// changing allocator state, so an invalid release cannot partially free slots. +#[derive(Debug)] +pub struct FixedSlabAllocator { + capacity: u32, + state: Mutex, +} + +impl FixedSlabAllocator { + pub fn new(capacity: u32) -> Result { + if capacity == 0 { + return Err(FixedSlabAllocatorError::ZeroCapacity); + } + Ok(Self { + capacity, + state: Mutex::new(FixedSlabState { + free_slots: (0..capacity).rev().collect(), + allocated: vec![false; capacity as usize], + validation_marks: vec![0; capacity as usize], + validation_epoch: 0, + }), + }) + } + + /// Reserves exactly `count` slots, or returns `None` without changing state. + pub fn try_reserve(&self, count: u32) -> Option> { + let mut state = self.state.lock(); + let count = count as usize; + if count > state.free_slots.len() { + return None; + } + + let mut slots = Vec::with_capacity(count); + for _ in 0..count { + let slot = state + .free_slots + .pop() + .expect("fixed slab free count was checked before reservation"); + let was_allocated = std::mem::replace(&mut state.allocated[slot as usize], true); + assert!(!was_allocated, "fixed slab freelist contained a live slot"); + slots.push(slot); + } + Some(slots) + } + + /// Atomically releases all supplied slots after validating the full input. + pub fn release(&self, slots: &[u32]) -> Result<(), FixedSlabAllocatorError> { + let mut state = self.state.lock(); + state.validation_epoch = state.validation_epoch.wrapping_add(1); + if state.validation_epoch == 0 { + state.validation_marks.fill(0); + state.validation_epoch = 1; + } + let validation_epoch = state.validation_epoch; + + for &slot in slots { + if slot >= self.capacity { + return Err(FixedSlabAllocatorError::SlotOutOfBounds { + slot, + capacity: self.capacity, + }); + } + let slot_index = slot as usize; + if state.validation_marks[slot_index] == validation_epoch { + return Err(FixedSlabAllocatorError::DuplicateSlot { slot }); + } + state.validation_marks[slot_index] = validation_epoch; + } + for &slot in slots { + if !state.allocated[slot as usize] { + return Err(FixedSlabAllocatorError::SlotNotAllocated { slot }); + } + } + + for &slot in slots { + state.allocated[slot as usize] = false; + state.free_slots.push(slot); + } + Ok(()) + } + + pub fn capacity(&self) -> u32 { + self.capacity + } + + pub fn free_count(&self) -> u32 { + self.state.lock().free_slots.len() as u32 + } + + pub fn live_count(&self) -> u32 { + self.capacity - self.free_count() + } + + pub fn is_empty(&self) -> bool { + self.free_count() == self.capacity + } +} + +#[cfg(test)] +mod tests { + use super::{FixedSlabAllocator, FixedSlabAllocatorError}; + use std::collections::HashSet; + use std::sync::{Arc, Mutex}; + use std::thread; + + #[test] + fn rejects_zero_capacity() { + assert!(matches!( + FixedSlabAllocator::new(0), + Err(FixedSlabAllocatorError::ZeroCapacity) + )); + } + + #[test] + fn reserve_is_ordered_and_all_or_none() { + let allocator = FixedSlabAllocator::new(4).unwrap(); + assert_eq!(allocator.try_reserve(3), Some(vec![0, 1, 2])); + assert_eq!(allocator.free_count(), 1); + assert_eq!(allocator.live_count(), 3); + + assert_eq!(allocator.try_reserve(2), None); + assert_eq!(allocator.free_count(), 1); + assert_eq!(allocator.try_reserve(1), Some(vec![3])); + assert_eq!(allocator.try_reserve(0), Some(Vec::new())); + assert!(!allocator.is_empty()); + } + + #[test] + fn released_slots_are_reused_and_counts_close() { + let allocator = FixedSlabAllocator::new(4).unwrap(); + let slots = allocator.try_reserve(3).unwrap(); + allocator.release(&slots[2..]).unwrap(); + assert_eq!(allocator.try_reserve(1), Some(vec![2])); + allocator.release(&slots[..2]).unwrap(); + allocator.release(&[2]).unwrap(); + + assert_eq!(allocator.capacity(), 4); + assert_eq!(allocator.free_count(), 4); + assert_eq!(allocator.live_count(), 0); + assert!(allocator.is_empty()); + } + + #[test] + fn invalid_release_is_atomic() { + let allocator = FixedSlabAllocator::new(4).unwrap(); + assert_eq!(allocator.try_reserve(2), Some(vec![0, 1])); + + assert_eq!( + allocator.release(&[0, 0]), + Err(FixedSlabAllocatorError::DuplicateSlot { slot: 0 }) + ); + assert_eq!(allocator.live_count(), 2); + assert_eq!( + allocator.release(&[0, 4]), + Err(FixedSlabAllocatorError::SlotOutOfBounds { + slot: 4, + capacity: 4, + }) + ); + assert_eq!(allocator.live_count(), 2); + + allocator.release(&[0]).unwrap(); + assert_eq!( + allocator.release(&[0]), + Err(FixedSlabAllocatorError::SlotNotAllocated { slot: 0 }) + ); + assert_eq!(allocator.live_count(), 1); + allocator.release(&[1]).unwrap(); + assert!(allocator.is_empty()); + } + + #[test] + fn concurrent_single_slot_reservations_are_unique() { + let allocator = Arc::new(FixedSlabAllocator::new(64).unwrap()); + let reserved = Arc::new(Mutex::new(Vec::new())); + let workers = (0..8) + .map(|_| { + let allocator = Arc::clone(&allocator); + let reserved = Arc::clone(&reserved); + thread::spawn(move || { + while let Some(mut slots) = allocator.try_reserve(1) { + reserved.lock().unwrap().append(&mut slots); + } + }) + }) + .collect::>(); + for worker in workers { + worker.join().unwrap(); + } + + let slots = reserved.lock().unwrap(); + assert_eq!(slots.len(), 64); + assert_eq!(slots.iter().copied().collect::>().len(), 64); + allocator.release(&slots).unwrap(); + assert!(allocator.is_empty()); + } +} diff --git a/fluxon_rs/fluxon_util/src/lib.rs b/fluxon_rs/fluxon_util/src/lib.rs index a85aed0..b95ff96 100644 --- a/fluxon_rs/fluxon_util/src/lib.rs +++ b/fluxon_rs/fluxon_util/src/lib.rs @@ -1,5 +1,6 @@ pub mod dag_viz_html; pub mod dev_config; +pub mod fixed_slab_allocator; pub mod fs_statvfs; pub mod fs_watch; pub mod git_version_build_record; @@ -7,6 +8,7 @@ pub mod hash; pub mod init_dag_compiler; pub mod map_lock; pub mod merge_recent_async_notifies; +pub mod pin_aware_moka; pub mod prefix_scan; pub mod prom_remote_write; pub mod scoped_future_set; @@ -37,10 +39,9 @@ pub mod limitrate; pub mod pyo3; // Re-export for stable public API: existing call sites can keep using `fluxon_util::init_log`. pub use log::{ - current_daily_sharded_log_path, current_log_file_path, daily_sharded_log_path, - display_runtime_log_path, init_log, init_log_test, init_log_with_extra_layer, - latest_existing_daily_sharded_log_path, resolve_readable_log_path, - DEFAULT_DAILY_LOG_RETENTION_DAYS, + DEFAULT_DAILY_LOG_RETENTION_DAYS, current_daily_sharded_log_path, current_log_file_path, + daily_sharded_log_path, display_runtime_log_path, init_log, init_log_test, + init_log_with_extra_layer, latest_existing_daily_sharded_log_path, resolve_readable_log_path, }; #[cfg(test)] mod test_util_test; @@ -251,7 +252,12 @@ mod tests { ); assert_logged_text( &active_log_path, - &["debug message", "info message", "warning message", "error message"], + &[ + "debug message", + "info message", + "warning message", + "error message", + ], ); } } diff --git a/fluxon_rs/fluxon_util/src/log.rs b/fluxon_rs/fluxon_util/src/log.rs index fc6066f..b941ed5 100644 --- a/fluxon_rs/fluxon_util/src/log.rs +++ b/fluxon_rs/fluxon_util/src/log.rs @@ -61,9 +61,7 @@ fn read_test_log_shard_window_config() -> anyhow::Result 0"); @@ -85,7 +83,9 @@ fn read_test_log_shard_window_config() -> anyhow::Result) -> anyhow::Result { +fn resolve_shard_date_from_datetime( + now: chrono::DateTime, +) -> anyhow::Result { let Some(config) = read_test_log_shard_window_config()? else { return Ok(now.date_naive()); }; @@ -99,8 +99,8 @@ fn resolve_shard_date_from_datetime(now: chrono::DateTime) -> anyho ); } let bucket_index = delta_seconds / config.window_seconds; - let base_date = chrono::NaiveDate::from_ymd_opt(2026, 1, 1) - .expect("valid hard-coded synthetic base date"); + let base_date = + chrono::NaiveDate::from_ymd_opt(2026, 1, 1).expect("valid hard-coded synthetic base date"); Ok(base_date + chrono::Days::new(bucket_index as u64)) } @@ -108,10 +108,7 @@ fn current_shard_date() -> anyhow::Result { resolve_shard_date_from_datetime(chrono::Utc::now()) } -fn cleanup_old_daily_sharded_logs( - base_path: &Path, - retention_days: usize, -) -> anyhow::Result<()> { +fn cleanup_old_daily_sharded_logs(base_path: &Path, retention_days: usize) -> anyhow::Result<()> { let parent = match base_path.parent() { Some(parent) => parent, None => return Ok(()), @@ -124,7 +121,8 @@ fn cleanup_old_daily_sharded_logs( return Ok(()); }; fs::create_dir_all(parent)?; - let keep_since = current_shard_date()? - chrono::Days::new(retention_days.saturating_sub(1) as u64); + let keep_since = + current_shard_date()? - chrono::Days::new(retention_days.saturating_sub(1) as u64); let prefix = format!("{stem}."); for entry in std::fs::read_dir(parent)? { let entry = entry?; @@ -180,10 +178,7 @@ impl DailyShardedFileWriter { current_daily_sharded_log_path(&self.base_path) } - fn rotate_if_needed( - &self, - state: &mut DailyShardedFileWriterState, - ) -> io::Result<()> { + fn rotate_if_needed(&self, state: &mut DailyShardedFileWriterState) -> io::Result<()> { let next_path = self .current_path() .map_err(|err| io::Error::new(io::ErrorKind::Other, err.to_string()))?; @@ -231,20 +226,6 @@ fn setup_global_log_guards(file_guard: WorkerGuard, console_guard: WorkerGuard) let _ = GLOBAL_CONSOLE_LOG_GUARD.set(console_guard); } -fn workspace_targets_filter( - workspace_level: filter::LevelFilter, - non_workspace_default: filter::LevelFilter, -) -> filter::Targets { - let mut targets = filter::Targets::new().with_default(non_workspace_default); - for c in generated_crates::OUR_CRATES { - targets = targets.with_target(*c, workspace_level); - } - for c in RDMA_DEBUG_TARGETS { - targets = targets.with_target(*c, workspace_level); - } - targets -} - fn third_party_log_target_overrides( enable_iceoryx_logs: bool, default_level: filter::LevelFilter, @@ -314,20 +295,19 @@ pub fn daily_sharded_log_path( base_path: &Path, date: chrono::NaiveDate, ) -> anyhow::Result { - let file_name = base_path.file_name().and_then(|v| v.to_str()).ok_or_else(|| { - anyhow::anyhow!( - "log path must end with a valid utf-8 filename: {}", - base_path.display() - ) - })?; + let file_name = base_path + .file_name() + .and_then(|v| v.to_str()) + .ok_or_else(|| { + anyhow::anyhow!( + "log path must end with a valid utf-8 filename: {}", + base_path.display() + ) + })?; let stem = file_name .strip_suffix(".log") .ok_or_else(|| anyhow::anyhow!("log path must end with .log: {}", base_path.display()))?; - Ok(base_path.with_file_name(format!( - "{}.{}.log", - stem, - date.format("%Y-%m-%d") - ))) + Ok(base_path.with_file_name(format!("{}.{}.log", stem, date.format("%Y-%m-%d")))) } pub fn current_daily_sharded_log_path(base_path: &Path) -> anyhow::Result { @@ -486,11 +466,16 @@ where } if let Some(level) = crate_level { - // Use build-time generated workspace member list. - let our_crates: &[&str] = generated_crates::OUR_CRATES; + // Use build-time generated workspace members plus the explicit RDMA + // diagnostic targets. A bare level remains one canonical setting for + // console, file, and exporter layers. + let our_crates = generated_crates::OUR_CRATES + .iter() + .copied() + .chain(RDMA_DEBUG_TARGETS.iter().copied()); let mut changed = removed_bare_debug; for c in our_crates { - if !existing_targets.contains(&c.to_string()) { + if !existing_targets.contains(c) { new_parts.push(format!("{}={}", c, level)); changed = true; } @@ -522,8 +507,9 @@ where } } - // File log keeps workspace crates at DEBUG; non-workspace crates default to WARN. - // This avoids dumping verbose dependency debug logs (e.g. h2/tower) into file output. + // All sinks follow the canonical RUST_LOG/FLUXON_LOG filter. Keeping a fixed DEBUG + // file layer here makes an `info` or `error` runtime silently continue formatting + // and writing high-volume debug events. let file_path = current_daily_log_file_path(log_path, instance_key); // Keep a copy for the whole process lifetime; collectors can clone it. if let Some(prev) = GLOBAL_LOG_FILE_PATH.get() { @@ -548,18 +534,17 @@ where .map(|v| v.trim().to_ascii_lowercase()), Some(v) if !matches!(v.as_str(), "" | "0" | "false" | "no") ); - let file_filter = - workspace_targets_filter(filter::LevelFilter::DEBUG, filter::LevelFilter::WARN); + let file_filter = EnvFilter::from_default_env(); let file_layer = tracing_subscriber::fmt::layer() .with_timer(UtcSecondTimer) .with_writer(file_writer) .with_ansi(false) - .with_filter(file_filter.clone().or(third_party_log_target_overrides( + .with_filter(file_filter.or(third_party_log_target_overrides( enable_iceoryx_logs, filter::LevelFilter::WARN, ))); - // Console logging follows user config (RUST_LOG/FLUXON_LOG); file logging ignores it. + // Console logging follows the same user config as file output. let (console_writer, console_guard) = non_blocking(io::stdout()); let console_env_filter = EnvFilter::from_default_env(); let console_layer = tracing_subscriber::fmt::layer() @@ -567,9 +552,8 @@ where .with_filter(console_env_filter); // Register layers. - // `extra_layer` follows the same target filtering as file output (workspace=DEBUG, deps=WARN) - // to avoid exporting noisy dependency debug logs by default. - let extra_layer = extra_layer.with_filter(file_filter.clone().or( + // Exporters follow the same canonical user filter as local sinks. + let extra_layer = extra_layer.with_filter(EnvFilter::from_default_env().or( third_party_log_target_overrides(enable_iceoryx_logs, filter::LevelFilter::WARN), )); let _ = tracing_subscriber::registry() diff --git a/fluxon_rs/fluxon_util/src/pin_aware_moka.rs b/fluxon_rs/fluxon_util/src/pin_aware_moka.rs new file mode 100644 index 0000000..b518fac --- /dev/null +++ b/fluxon_rs/fluxon_util/src/pin_aware_moka.rs @@ -0,0 +1,671 @@ +//! Explicit pin management around Moka without changing Moka's entry key. +//! +//! The first pin removes one entry from Moka. The final unpin inserts the same +//! generation again. A pin alias can be reserved before its entry is inserted. + +use moka::notification::RemovalCause; +use moka::ops::compute::{CompResult, Op}; +use moka::sync::Cache; +use parking_lot::Mutex; +use std::collections::{HashMap, HashSet}; +use std::hash::Hash; +use std::sync::{Arc, Weak}; + +type EvictionListener = dyn Fn(Arc, Value, RemovalCause) + Send + Sync + 'static; +type Weigher = dyn Fn(&Key, &Value) -> u32 + Send + Sync + 'static; + +/// Builder for [`PinAwareMoka`]. +pub struct PinAwareMokaBuilder { + max_capacity: u64, + weigher: Arc>, + eviction_listener: Option>>, + _pin_alias: std::marker::PhantomData, +} + +impl PinAwareMokaBuilder { + fn new(max_capacity: u64) -> Self { + Self { + max_capacity, + weigher: Arc::new(|_, _| 1), + eviction_listener: None, + _pin_alias: std::marker::PhantomData, + } + } + + pub fn weigher( + mut self, + weigher: impl Fn(&Key, &Value) -> u32 + Send + Sync + 'static, + ) -> Self { + self.weigher = Arc::new(weigher); + self + } + + /// Set the listener for Moka removals. + /// + /// Pin/unpin removals are internal and are not sent to this listener. + pub fn eviction_listener( + mut self, + listener: impl Fn(Arc, Value, RemovalCause) + Send + Sync + 'static, + ) -> Self { + self.eviction_listener = Some(Arc::new(listener)); + self + } +} + +impl PinAwareMokaBuilder +where + Key: Clone + Eq + Hash + Send + Sync + 'static, + PinAlias: Clone + Eq + Hash + Send + Sync + 'static, + Value: Clone + Send + Sync + 'static, +{ + pub fn build(self) -> PinAwareMoka { + assert!( + self.max_capacity > 0, + "pin-aware Moka capacity must be positive" + ); + let weigher = self.weigher; + let eviction_listener = self.eviction_listener; + let inner = Arc::new_cyclic(|weak: &Weak>| { + let listener_inner = weak.clone(); + let cache = Cache::builder() + .max_capacity(self.max_capacity) + .weigher(move |_key: &Key, entry: &MokaEntry| entry.weight) + .eviction_listener(move |key, entry, cause| { + if let Some(inner) = listener_inner.upgrade() { + inner.handle_moka_removal(key, entry, cause); + } + }) + .build(); + Inner { + cache, + state: Mutex::new(State::default()), + mutation: Mutex::new(()), + weigher, + eviction_listener, + } + }); + PinAwareMoka { inner } + } +} + +/// A key-preserving Moka wrapper with explicit pin aliases. +pub struct PinAwareMoka { + inner: Arc>, +} + +impl Clone for PinAwareMoka { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +#[derive(Clone)] +struct MokaEntry { + generation: u64, + value: Value, + weight: u32, + _pin_alias: std::marker::PhantomData, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Residency { + InMoka, + Pinned, + Selected, +} + +struct EntryState { + generation: u64, + pin_aliases: Arc<[PinAlias]>, + value: Value, + weight: u32, + pin_count: usize, + residency: Residency, +} + +enum PinLease { + Pending(PinAlias), + Assigned { key: Key, generation: u64 }, +} + +struct State { + next_generation: u64, + next_pin_lease: u64, + entries: HashMap>, + aliases: HashMap, + pin_leases: HashMap>, + pending_pins: HashMap>, +} + +impl Default for State { + fn default() -> Self { + Self { + next_generation: 1, + next_pin_lease: 1, + entries: HashMap::new(), + aliases: HashMap::new(), + pin_leases: HashMap::new(), + pending_pins: HashMap::new(), + } + } +} + +struct Inner { + cache: Cache>, + state: Mutex>, + // State transitions and Moka mutations are serialized, but the state lock + // is always released before calling Moka and its synchronous listener. + mutation: Mutex<()>, + weigher: Arc>, + eviction_listener: Option>>, +} + +impl Inner +where + Key: Clone + Eq + Hash + Send + Sync + 'static, + PinAlias: Clone + Eq + Hash + Send + Sync + 'static, + Value: Clone + Send + Sync + 'static, +{ + fn remove_moka_generation(&self, key: &Key, generation: u64) -> bool { + matches!( + self.cache.entry(key.clone()).and_compute_with(|current| { + if current + .as_ref() + .is_some_and(|entry| entry.value().generation == generation) + { + Op::Remove + } else { + Op::Nop + } + }), + CompResult::Removed(_) + ) + } + + fn handle_moka_removal( + &self, + key: Arc, + entry: MokaEntry, + cause: RemovalCause, + ) { + if cause != RemovalCause::Size { + return; + } + let selected = { + let mut state = self.state.lock(); + let Some(current) = state.entries.get_mut(key.as_ref()) else { + return; + }; + if current.generation != entry.generation || current.residency != Residency::InMoka { + return; + } + debug_assert_eq!(current.pin_count, 0); + current.residency = Residency::Selected; + true + }; + if selected && let Some(listener) = self.eviction_listener.as_ref() { + listener(key, entry.value, cause); + } + } + + fn release_pin(self: &Arc, lease_id: u64) { + let _mutation = self.mutation.lock(); + let entry_to_admit = { + let mut state = self.state.lock(); + let Some(lease) = state.pin_leases.remove(&lease_id) else { + return; + }; + match lease { + PinLease::Pending(alias) => { + if let Some(leases) = state.pending_pins.get_mut(&alias) { + leases.remove(&lease_id); + if leases.is_empty() { + state.pending_pins.remove(&alias); + } + } + None + } + PinLease::Assigned { key, generation } => { + let Some(entry) = state.entries.get_mut(&key) else { + return; + }; + if entry.generation != generation { + return; + } + entry.pin_count = entry + .pin_count + .checked_sub(1) + .expect("pin-aware Moka pin count underflow"); + if entry.pin_count == 0 && entry.residency == Residency::Pinned { + entry.residency = Residency::InMoka; + Some(( + key, + MokaEntry { + generation, + value: entry.value.clone(), + weight: entry.weight, + _pin_alias: std::marker::PhantomData, + }, + )) + } else { + None + } + } + } + }; + if let Some((key, entry)) = entry_to_admit { + self.cache.insert(key, entry); + } + } +} + +impl PinAwareMoka +where + Key: Clone + Eq + Hash + Send + Sync + 'static, + PinAlias: Clone + Eq + Hash + Send + Sync + 'static, + Value: Clone + Send + Sync + 'static, +{ + pub fn builder(max_capacity: u64) -> PinAwareMokaBuilder { + PinAwareMokaBuilder::new(max_capacity) + } + + /// Insert one normal Moka entry with one or more explicit pin identities. + pub fn insert( + &self, + key: Key, + pin_aliases: impl IntoIterator, + value: Value, + ) -> u64 { + let pin_aliases = pin_aliases.into_iter().collect::>(); + assert!( + !pin_aliases.is_empty(), + "a pin-aware Moka entry must have at least one pin alias" + ); + assert_eq!( + pin_aliases.iter().collect::>().len(), + pin_aliases.len(), + "a pin-aware Moka entry cannot contain duplicate pin aliases" + ); + let weight = (self.inner.weigher)(&key, &value); + let pin_aliases: Arc<[PinAlias]> = pin_aliases.into(); + let _mutation = self.inner.mutation.lock(); + let (generation, admit, displaced) = { + let mut state = self.inner.state.lock(); + let generation = state.next_generation; + state.next_generation = state + .next_generation + .checked_add(1) + .expect("pin-aware Moka generation space exhausted"); + + let mut displaced_keys = HashSet::new(); + if state.entries.contains_key(&key) { + displaced_keys.insert(key.clone()); + } + for alias in pin_aliases.iter() { + if let Some((old_key, _)) = state.aliases.get(alias) { + displaced_keys.insert(old_key.clone()); + } + } + let mut displaced = Vec::new(); + for old_key in displaced_keys { + if let Some(old) = state.entries.remove(&old_key) { + for alias in old.pin_aliases.iter() { + if state.aliases.get(alias) == Some(&(old_key.clone(), old.generation)) { + state.aliases.remove(alias); + } + } + displaced.push((old_key, old.generation)); + } + } + + let mut pin_count = 0usize; + for alias in pin_aliases.iter() { + let pending = state.pending_pins.remove(alias).unwrap_or_default(); + pin_count = pin_count + .checked_add(pending.len()) + .expect("pin-aware Moka pin count overflow"); + for lease_id in pending { + *state + .pin_leases + .get_mut(&lease_id) + .expect("pending pin lease must exist") = PinLease::Assigned { + key: key.clone(), + generation, + }; + } + state + .aliases + .insert(alias.clone(), (key.clone(), generation)); + } + let residency = if pin_count == 0 { + Residency::InMoka + } else { + Residency::Pinned + }; + state.entries.insert( + key.clone(), + EntryState { + generation, + pin_aliases: pin_aliases.clone(), + value: value.clone(), + weight, + pin_count, + residency, + }, + ); + let admit = (residency == Residency::InMoka).then(|| MokaEntry { + generation, + value, + weight, + _pin_alias: std::marker::PhantomData, + }); + (generation, admit, displaced) + }; + for (old_key, old_generation) in displaced { + self.inner.remove_moka_generation(&old_key, old_generation); + } + if let Some(entry) = admit { + self.inner.cache.insert(key, entry); + } + generation + } + + /// Pin the entry identified by `alias`, or reserve a pin before insertion. + /// + /// Returns `None` when Moka selection already owns that entry. + pub fn try_pin_alias(&self, alias: PinAlias) -> Option { + self.try_pin_alias_if(alias, |_| true) + } + + /// Pin only if an existing value satisfies `predicate`. + pub fn try_pin_alias_if( + &self, + alias: PinAlias, + predicate: impl FnOnce(&Value) -> bool, + ) -> Option { + let _mutation = self.inner.mutation.lock(); + let mut state = self.inner.state.lock(); + let lease_id = state.next_pin_lease; + state.next_pin_lease = state + .next_pin_lease + .checked_add(1) + .expect("pin-aware Moka pin lease space exhausted"); + let mut remove = None; + if let Some((key, generation)) = state.aliases.get(&alias).cloned() { + let entry = state + .entries + .get_mut(&key) + .expect("pin alias must reference a live entry"); + assert_eq!(entry.generation, generation); + if !predicate(&entry.value) || entry.residency == Residency::Selected { + return None; + } + entry.pin_count = entry + .pin_count + .checked_add(1) + .expect("pin-aware Moka pin count overflow"); + if entry.residency == Residency::InMoka { + entry.residency = Residency::Pinned; + remove = Some((key.clone(), generation)); + } + state + .pin_leases + .insert(lease_id, PinLease::Assigned { key, generation }); + } else { + state + .pending_pins + .entry(alias.clone()) + .or_default() + .insert(lease_id); + state.pin_leases.insert(lease_id, PinLease::Pending(alias)); + } + drop(state); + if let Some((key, generation)) = remove { + self.inner.remove_moka_generation(&key, generation); + } + Some(PinGuard { + _lease: Arc::new(PinGuardLease { + release: Box::new({ + let inner = Arc::downgrade(&self.inner); + move || { + if let Some(inner) = inner.upgrade() { + inner.release_pin(lease_id); + } + } + }), + }), + }) + } + + /// Read and touch an unselected entry. Pinned entries remain readable here. + pub fn get(&self, key: &Key) -> Option { + let _mutation = self.inner.mutation.lock(); + let (generation, value, in_moka) = { + let state = self.inner.state.lock(); + let entry = state.entries.get(key)?; + if entry.residency == Residency::Selected { + return None; + } + ( + entry.generation, + entry.value.clone(), + entry.residency == Residency::InMoka, + ) + }; + if in_moka { + let _ = self + .inner + .cache + .get(key) + .filter(|entry| entry.generation == generation); + } + Some(value) + } + + pub fn contains_key(&self, key: &Key) -> bool { + self.inner + .state + .lock() + .entries + .get(key) + .is_some_and(|entry| entry.residency != Residency::Selected) + } + + /// Remove an entry generation from both the wrapper and Moka. + pub fn invalidate(&self, key: &Key) -> bool { + self.invalidate_if(key, |_| true) + } + + pub fn invalidate_if(&self, key: &Key, predicate: impl FnOnce(&Value) -> bool) -> bool { + self.take_if(key, predicate).is_some() + } + + pub fn take_if(&self, key: &Key, predicate: impl FnOnce(&Value) -> bool) -> Option { + let _mutation = self.inner.mutation.lock(); + let removed = { + let mut state = self.inner.state.lock(); + let Some(entry) = state.entries.get(key) else { + return None; + }; + if !predicate(&entry.value) { + return None; + } + let entry = state + .entries + .remove(key) + .expect("validated pin-aware entry disappeared under lock"); + for alias in entry.pin_aliases.iter() { + if state.aliases.get(alias) == Some(&(key.clone(), entry.generation)) { + state.aliases.remove(alias); + } + } + entry + }; + self.inner.remove_moka_generation(key, removed.generation); + Some(removed.value) + } + + /// Pop LRU entries until at least `weight_to_evict` weight is selected. + pub fn evict_some(&self, weight_to_evict: u64) -> u64 { + let _mutation = self.inner.mutation.lock(); + self.inner.cache.evict_some(weight_to_evict) + } + + pub fn run_pending_tasks(&self) { + let _mutation = self.inner.mutation.lock(); + self.inner.cache.run_pending_tasks(); + } + + pub fn set_max_capacity(&self, capacity: u64) -> Result<(), moka::CapacityError> { + self.inner.cache.set_max_capacity(capacity) + } + + pub fn max_capacity(&self) -> Option { + self.inner.cache.policy().max_capacity() + } + + pub fn weighted_size(&self) -> u64 { + self.inner.cache.weighted_size() + } + + pub fn entry_count(&self) -> u64 { + self.inner.cache.entry_count() + } +} + +/// A cloneable pin. The final clone performs the unpin transition. +#[derive(Clone)] +pub struct PinGuard { + _lease: Arc, +} + +struct PinGuardLease { + release: Box, +} + +impl Drop for PinGuardLease { + fn drop(&mut self) { + (self.release)(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + use std::thread; + + type TestCache = PinAwareMoka; + + fn cache(capacity: u64, evicted: Arc>>) -> TestCache { + TestCache::builder(capacity) + .weigher(|_, weight| *weight) + .eviction_listener(move |key, _, cause| { + assert_eq!(cause, RemovalCause::Size); + evicted.lock().push((*key).clone()); + }) + .build() + } + + fn insert(cache: &TestCache, key: &str, alias: &str, weight: u32) { + cache.insert(key.to_string(), [alias.to_string()], weight); + cache.run_pending_tasks(); + } + + #[test] + fn first_pin_removes_and_final_unpin_readmits_one_key() { + let cache = cache(100, Arc::new(Mutex::new(Vec::new()))); + insert(&cache, "key", "alias", 10); + let first = cache.try_pin_alias("alias".to_string()).unwrap(); + let second = cache.try_pin_alias("alias".to_string()).unwrap(); + cache.run_pending_tasks(); + assert_eq!(cache.entry_count(), 0); + assert_eq!(cache.get(&"key".to_string()), Some(10)); + + drop(first); + cache.run_pending_tasks(); + assert_eq!(cache.entry_count(), 0); + drop(second); + cache.run_pending_tasks(); + assert_eq!(cache.entry_count(), 1); + } + + #[test] + fn pin_before_insert_defers_moka_admission() { + let cache = cache(100, Arc::new(Mutex::new(Vec::new()))); + let pin = cache.try_pin_alias("alias".to_string()).unwrap(); + insert(&cache, "key", "alias", 10); + assert_eq!(cache.entry_count(), 0); + drop(pin); + cache.run_pending_tasks(); + assert_eq!(cache.entry_count(), 1); + } + + #[test] + fn pin_and_evict_race_has_one_winner() { + for iteration in 0..100 { + let evicted = Arc::new(Mutex::new(Vec::new())); + let cache = cache(100, evicted.clone()); + insert(&cache, "key", "alias", 10); + let barrier = Arc::new(Barrier::new(3)); + let pin_cache = cache.clone(); + let pin_barrier = barrier.clone(); + let pin = thread::spawn(move || { + pin_barrier.wait(); + pin_cache.try_pin_alias("alias".to_string()) + }); + let pop_cache = cache.clone(); + let pop_barrier = barrier.clone(); + let pop = thread::spawn(move || { + pop_barrier.wait(); + pop_cache.evict_some(1) + }); + barrier.wait(); + let pin = pin.join().unwrap(); + let selected_weight = pop.join().unwrap(); + assert_ne!(pin.is_some(), selected_weight != 0, "iteration {iteration}"); + assert_eq!(evicted.lock().len(), usize::from(selected_weight != 0)); + drop(pin); + } + } + + #[test] + fn stale_guard_does_not_readmit_a_new_generation() { + let cache = cache(100, Arc::new(Mutex::new(Vec::new()))); + let stale = cache.try_pin_alias("old-alias".to_string()).unwrap(); + insert(&cache, "key", "old-alias", 10); + insert(&cache, "key", "new-alias", 20); + drop(stale); + cache.run_pending_tasks(); + assert_eq!(cache.get(&"key".to_string()), Some(20)); + assert_eq!(cache.entry_count(), 1); + } + + #[test] + fn selected_entry_rejects_pin_until_restore_or_invalidate() { + let evicted = Arc::new(Mutex::new(Vec::new())); + let cache = cache(100, evicted.clone()); + insert(&cache, "key", "alias", 10); + assert_eq!(cache.evict_some(1), 10); + assert_eq!(&*evicted.lock(), &["key"]); + assert!(cache.try_pin_alias("alias".to_string()).is_none()); + + insert(&cache, "key", "alias", 10); + let pin = cache.try_pin_alias("alias".to_string()).unwrap(); + drop(pin); + assert!(cache.invalidate(&"key".to_string())); + assert!(!cache.contains_key(&"key".to_string())); + } + + #[test] + fn weighted_pop_keeps_key_granularity_and_can_overshoot() { + let evicted = Arc::new(Mutex::new(Vec::new())); + let cache = cache(100, evicted.clone()); + insert(&cache, "first", "a", 12); + insert(&cache, "second", "b", 7); + assert_eq!(cache.evict_some(13), 19); + assert_eq!(evicted.lock().len(), 2); + } +} diff --git a/fluxon_rs/fluxon_util/src/vallocator/mod.rs b/fluxon_rs/fluxon_util/src/vallocator/mod.rs index 2f56285..14d7dfd 100644 --- a/fluxon_rs/fluxon_util/src/vallocator/mod.rs +++ b/fluxon_rs/fluxon_util/src/vallocator/mod.rs @@ -4,6 +4,7 @@ pub use frame::*; pub mod test; use parking_lot::RwLock; +use std::collections::{BTreeMap, BTreeSet}; crate::define_error_code_enum_with_from! { #[repr(i32)] @@ -42,6 +43,7 @@ impl std::fmt::Display for AllocError { impl std::error::Error for AllocError {} pub const ORDER: usize = 64; +const CONTIGUOUS_ALLOCATION_ALIGNMENT: u64 = 4096; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AllocRegion { @@ -49,8 +51,138 @@ pub struct AllocRegion { pub size: u64, } +#[derive(Debug)] +struct ContiguousRangeAllocator { + free_by_start: BTreeMap, + free_by_size: BTreeSet<(u64, u64)>, + allocated: u64, + total: u64, +} + +impl ContiguousRangeAllocator { + fn new(total: u64) -> Self { + let mut allocator = Self { + free_by_start: BTreeMap::new(), + free_by_size: BTreeSet::new(), + allocated: 0, + total, + }; + if total != 0 { + allocator.insert_free_range(0, total); + } + allocator + } + + fn aligned_capacity(size: u64) -> Option { + size.checked_add(CONTIGUOUS_ALLOCATION_ALIGNMENT - 1) + .map(|value| value / CONTIGUOUS_ALLOCATION_ALIGNMENT * CONTIGUOUS_ALLOCATION_ALIGNMENT) + } + + fn insert_free_range(&mut self, start: u64, len: u64) { + if len == 0 { + return; + } + assert!(self.free_by_start.insert(start, len).is_none()); + assert!(self.free_by_size.insert((len, start))); + } + + fn remove_free_range(&mut self, start: u64, len: u64) { + assert_eq!(self.free_by_start.remove(&start), Some(len)); + assert!(self.free_by_size.remove(&(len, start))); + } + + fn alloc(&mut self, requested: u64) -> Option<(u64, u64)> { + let capacity = Self::aligned_capacity(requested)?; + let (range_len, range_start) = self.free_by_size.range((capacity, 0)..).next().copied()?; + self.remove_free_range(range_start, range_len); + let remaining = range_len - capacity; + if remaining != 0 { + self.insert_free_range(range_start + capacity, remaining); + } + self.allocated = self.allocated.checked_add(capacity)?; + Some((range_start, capacity)) + } + + fn free(&mut self, start: u64, capacity: u64) -> Result { + let end = start.checked_add(capacity).ok_or_else(|| AllocError { + code: AllocErrorCode::InvalidAddress, + message: format!("Allocation range overflows: start={start} capacity={capacity}"), + })?; + if capacity == 0 + || capacity % CONTIGUOUS_ALLOCATION_ALIGNMENT != 0 + || start % CONTIGUOUS_ALLOCATION_ALIGNMENT != 0 + || end > self.total + { + return Err(AllocError { + code: AllocErrorCode::InvalidAddress, + message: format!( + "Invalid aligned allocation range: start={start} capacity={capacity} total={}", + self.total + ), + }); + } + + let previous = self + .free_by_start + .range(..start) + .next_back() + .map(|(&range_start, &range_len)| (range_start, range_len)); + if previous + .is_some_and(|(range_start, range_len)| range_start.saturating_add(range_len) > start) + { + return Err(AllocError { + code: AllocErrorCode::DoubleFree, + message: format!( + "Range overlaps previous free range: start={start} capacity={capacity}" + ), + }); + } + let next = self + .free_by_start + .range(start..) + .next() + .map(|(&range_start, &range_len)| (range_start, range_len)); + if next.is_some_and(|(range_start, _)| range_start < end) { + return Err(AllocError { + code: AllocErrorCode::DoubleFree, + message: format!( + "Range overlaps next free range: start={start} capacity={capacity}" + ), + }); + } + + let mut merged_start = start; + let mut merged_len = capacity; + if let Some((previous_start, previous_len)) = previous + && previous_start + previous_len == start + { + self.remove_free_range(previous_start, previous_len); + merged_start = previous_start; + merged_len += previous_len; + } + if let Some((next_start, next_len)) = next + && end == next_start + { + self.remove_free_range(next_start, next_len); + merged_len += next_len; + } + self.insert_free_range(merged_start, merged_len); + self.allocated = self + .allocated + .checked_sub(capacity) + .ok_or_else(|| AllocError { + code: AllocErrorCode::DoubleFree, + message: format!( + "Allocated byte counter underflow: allocated={} capacity={capacity}", + self.allocated + ), + })?; + Ok(capacity) + } +} + pub struct VirtualAllocator { - inner: RwLock>, + inner: RwLock, } impl std::fmt::Debug for VirtualAllocator { @@ -65,11 +197,8 @@ impl std::fmt::Debug for VirtualAllocator { impl VirtualAllocator { pub fn new(total_size: u64) -> Result { - let mut allocator = FrameAllocator::::new(); - allocator.add_frame(0, total_size); - Ok(Self { - inner: RwLock::new(allocator), + inner: RwLock::new(ContiguousRangeAllocator::new(total_size)), }) } @@ -95,8 +224,7 @@ impl VirtualAllocator { } let mut allocator = self.inner.write(); - let (addr_opt, actual_size) = allocator.alloc(size); - if let Some(addr) = addr_opt { + if let Some((addr, actual_size)) = allocator.alloc(size) { if actual_size < size { return Err(AllocError { code: AllocErrorCode::AllocationFailed, @@ -127,7 +255,7 @@ impl VirtualAllocator { } let mut allocator = self.inner.write(); - let freed_size = allocator.dealloc(ptr, size); + let freed_size = allocator.free(ptr, size)?; if freed_size != size { return Err(AllocError { code: AllocErrorCode::DeallocationException, diff --git a/fluxon_rs/fluxon_util/src/vallocator/test.rs b/fluxon_rs/fluxon_util/src/vallocator/test.rs index 132879f..07a92b1 100644 --- a/fluxon_rs/fluxon_util/src/vallocator/test.rs +++ b/fluxon_rs/fluxon_util/src/vallocator/test.rs @@ -1,5 +1,40 @@ use crate::vallocator::VirtualAllocator; +#[test] +fn virtual_allocator_keeps_sglang_pages_exact_instead_of_power_of_two_rounding() { + const PAGE_BYTES: u64 = 4_718_592; + let allocator = VirtualAllocator::new(PAGE_BYTES * 3).unwrap(); + let first = allocator.alloc(PAGE_BYTES).unwrap(); + let second = allocator.alloc(PAGE_BYTES).unwrap(); + let third = allocator.alloc(PAGE_BYTES).unwrap(); + + assert_eq!(first.size, PAGE_BYTES); + assert_eq!(second.size, PAGE_BYTES); + assert_eq!(third.size, PAGE_BYTES); + assert_eq!(allocator.get_allocated_size(), PAGE_BYTES * 3); + assert!(allocator.alloc(PAGE_BYTES).is_err()); + + allocator.free(second.start_addr, second.size).unwrap(); + allocator.free(first.start_addr, first.size).unwrap(); + allocator.free(third.start_addr, third.size).unwrap(); + assert_eq!(allocator.get_allocated_size(), 0); + + let whole = allocator.alloc(PAGE_BYTES * 3).unwrap(); + assert_eq!(whole.start_addr, 0); + assert_eq!(whole.size, PAGE_BYTES * 3); +} + +#[test] +fn virtual_allocator_rounds_only_to_transfer_page_alignment() { + let allocator = VirtualAllocator::new(16 * 1024).unwrap(); + let region = allocator.alloc(4097).unwrap(); + assert_eq!(region.start_addr, 0); + assert_eq!(region.size, 8192); + assert_eq!(allocator.get_allocated_size(), 8192); + allocator.free(region.start_addr, region.size).unwrap(); + assert_eq!(allocator.get_allocated_size(), 0); +} + #[test] fn test_empty_frame_allocator() { let mut frame = crate::vallocator::frame::FrameAllocator::<32>::new(); diff --git a/fluxon_rs/fluxon_util/tests/log_mgmt.rs b/fluxon_rs/fluxon_util/tests/log_mgmt.rs index 431c5fc..a458c47 100644 --- a/fluxon_rs/fluxon_util/tests/log_mgmt.rs +++ b/fluxon_rs/fluxon_util/tests/log_mgmt.rs @@ -59,9 +59,15 @@ fn kv_log_shards_roll_and_cleanup_with_test_window() { .expect("unix epoch") .as_secs() as i64; let _window_guard = EnvVarGuard::set(TEST_LOG_SHARD_WINDOW_SECONDS_ENV, "10"); - let _anchor_guard = EnvVarGuard::set(TEST_LOG_SHARD_ANCHOR_UNIX_SECONDS_ENV, (now - 2).to_string()); + let _anchor_guard = EnvVarGuard::set( + TEST_LOG_SHARD_ANCHOR_UNIX_SECONDS_ENV, + (now - 2).to_string(), + ); + let _rust_log_guard = EnvVarGuard::set("RUST_LOG", "info"); + let _fluxon_log_guard = EnvVarGuard::set("FLUXON_LOG", "info"); fluxon_util::init_log(log_path, instance_key); + tracing::debug!(target: "fluxon_util", "[kv-log-mgmt][debug-must-be-filtered]"); tracing::info!(target: "fluxon_util", "[kv-log-mgmt][phase=before] ts={}", now); std::thread::sleep(Duration::from_millis(300)); std::thread::sleep(Duration::from_secs(11)); @@ -104,6 +110,11 @@ fn kv_log_shards_roll_and_cleanup_with_test_window() { !shard_2_text.contains("[kv-log-mgmt][phase=before]"), "second shard should not contain the before marker" ); + assert!( + !shard_1_text.contains("[kv-log-mgmt][debug-must-be-filtered]") + && !shard_2_text.contains("[kv-log-mgmt][debug-must-be-filtered]"), + "file output must honor the canonical info filter" + ); assert_eq!(DEFAULT_DAILY_LOG_RETENTION_DAYS, 31); } @@ -115,7 +126,8 @@ fn resolve_readable_log_path_ignores_plain_base_log_when_daily_shards_exist() { let shard_path = temp_dir.path().join("startup.2026-06-21.log"); fs::write(&shard_path, "shard\n").expect("write shard log"); - let resolved = fluxon_util::resolve_readable_log_path(&base_path).expect("resolve readable log path"); + let resolved = + fluxon_util::resolve_readable_log_path(&base_path).expect("resolve readable log path"); assert_eq!(resolved, shard_path); } @@ -128,7 +140,7 @@ fn latest_existing_daily_sharded_log_path_skips_invalid_candidates() { fs::write(&invalid_shard_path, "invalid\n").expect("write invalid shard"); fs::write(&valid_shard_path, "valid\n").expect("write valid shard"); - let resolved = - fluxon_util::latest_existing_daily_sharded_log_path(&base_path).expect("resolve latest shard"); + let resolved = fluxon_util::latest_existing_daily_sharded_log_path(&base_path) + .expect("resolve latest shard"); assert_eq!(resolved, valid_shard_path); } diff --git a/pegaflow b/pegaflow new file mode 160000 index 0000000..939363f --- /dev/null +++ b/pegaflow @@ -0,0 +1 @@ +Subproject commit 939363f198a0ffded9e3f30d8af9cdd74439f16c diff --git a/scripts/preview_doc_html.py b/scripts/preview_doc_html.py new file mode 100755 index 0000000..18f5093 --- /dev/null +++ b/scripts/preview_doc_html.py @@ -0,0 +1,547 @@ +#!/usr/bin/env python3 +"""Generate a standalone HTML preview for a Markdown document.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import subprocess +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_OUTPUT_ROOT = REPO_ROOT / ".dever/html_preview" + + +HTML_TEMPLATE = """ + + + + + {title} + + + + + + +
+
{title}
+
+ Waiting + + +
+
+ +
+
+
+ + + + + + + + + + + +""" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate a standalone HTML preview for a Markdown document with Mermaid fullscreen zoom/pan support." + ) + parser.add_argument( + "markdown_path", + type=Path, + help="Markdown input path.", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help=f"HTML output path. Default: {DEFAULT_OUTPUT_ROOT.relative_to(REPO_ROOT)}/.html", + ) + parser.add_argument( + "--open", + action="store_true", + help="Open the generated HTML file with the system browser.", + ) + return parser.parse_args() + + +def default_output_path(input_path: Path) -> Path: + try: + relative_input = input_path.relative_to(REPO_ROOT) + except ValueError: + digest = hashlib.sha256(str(input_path).encode("utf-8")).hexdigest()[:12] + relative_input = Path("external") / f"{input_path.stem}-{digest}{input_path.suffix}" + return (DEFAULT_OUTPUT_ROOT / relative_input).with_suffix(".html") + + +def infer_html_lang(input_path: Path, markdown: str) -> str: + parts = set(input_path.parts) + if "fluxon_doc_cn" in parts or input_path.name == "README_CN.md": + return "zh-CN" + if "fluxon_doc_en" in parts or input_path.name == "README.md": + return "en" + return "zh-CN" if re.search(r"[\u4e00-\u9fff]", markdown) else "en" + + +def first_markdown_heading(markdown: str, fallback: str) -> str: + for line in markdown.splitlines(): + stripped = line.strip() + if stripped.startswith("# "): + return stripped[2:].strip() or fallback + return fallback + + +def open_in_browser(path: Path) -> None: + if sys.platform == "darwin": + subprocess.run(["open", str(path)], check=False) + elif os.name == "nt": + os.startfile(str(path)) # type: ignore[attr-defined] + else: + subprocess.run(["xdg-open", str(path)], check=False) + + +def main() -> int: + args = parse_args() + input_path = args.markdown_path.resolve() + + if not input_path.exists(): + print(f"Input markdown does not exist: {input_path}", file=sys.stderr) + return 1 + if not input_path.is_file(): + print(f"Input path is not a file: {input_path}", file=sys.stderr) + return 1 + + output_path = args.output.resolve() if args.output else default_output_path(input_path) + + markdown = input_path.read_text(encoding="utf-8") + title = first_markdown_heading(markdown, input_path.stem) + html = HTML_TEMPLATE.format( + title=title, + html_lang=infer_html_lang(input_path, markdown), + markdown_json=json.dumps(markdown, ensure_ascii=False), + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(html, encoding="utf-8") + + print(f"Generated HTML preview: {output_path}") + print("Open it in a browser. Mermaid fullscreen zoom/pan uses CDN assets, so internet access is required.") + + if args.open: + open_in_browser(output_path) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/setup_and_pack/nix/pack_fluxonkv_pylib.py b/setup_and_pack/nix/pack_fluxonkv_pylib.py index e12f8fe..cb0af6d 100644 --- a/setup_and_pack/nix/pack_fluxonkv_pylib.py +++ b/setup_and_pack/nix/pack_fluxonkv_pylib.py @@ -198,6 +198,11 @@ def _cleanup_temp_workspace_mount_dirs() -> None: "libstdc++.so.6", "libgomp.so.1", ) +EXTERNAL_DRIVER_RUNTIME_LIBRARY_NAMES = ( + # libcuda is supplied by the NVIDIA kernel-driver installation on the target GPU node. + # Packaging the build-time stub would make the wheel non-functional at runtime. + "libcuda.so.1", +) def _dedupe_relative_paths(relative_paths: tuple[str, ...]) -> tuple[str, ...]: ordered_relative_paths: list[str] = [] @@ -2273,6 +2278,10 @@ def _is_core_system_runtime_lib(lib_name: str) -> bool: return any(lib_name.startswith(prefix) for prefix in CORE_SYSTEM_RUNTIME_LIB_PREFIXES) +def _is_external_driver_runtime_lib(lib_name: str) -> bool: + return lib_name in EXTERNAL_DRIVER_RUNTIME_LIBRARY_NAMES + + def _read_runtime_dependency_entries( path: Path, *, @@ -2314,6 +2323,12 @@ def _read_runtime_dependency_entries( if "=>" in line: needed_name, raw_target = (part.strip() for part in line.split("=>", 1)) if raw_target == "not found": + if _is_external_driver_runtime_lib(needed_name): + print( + "wheel finalize: leaving target-driver dependency external: " + f"{path} -> {needed_name}" + ) + continue raise RuntimeError(f"runtime dependency not found for {path}: {needed_name}") dep_path_str = raw_target.split(" ", 1)[0].strip() else: diff --git a/setup_and_pack/rather_no_git_submodule.yaml b/setup_and_pack/rather_no_git_submodule.yaml index 9f00b04..453a0e3 100644 --- a/setup_and_pack/rather_no_git_submodule.yaml +++ b/setup_and_pack/rather_no_git_submodule.yaml @@ -1,4 +1,4 @@ modules: - path: fluxon_rs/moka repo: https://github.com/ActivePeter/moka - checkout: main + checkout: 91c97076e7b1bbac306d7ace7b74f9e994440b2d diff --git a/setup_and_pack/tests/test_vendor_runtime_layout.py b/setup_and_pack/tests/test_vendor_runtime_layout.py index 48f87be..52eac3d 100644 --- a/setup_and_pack/tests/test_vendor_runtime_layout.py +++ b/setup_and_pack/tests/test_vendor_runtime_layout.py @@ -33,6 +33,15 @@ def _load_module(module_path: Path, module_name: str): class VendorRuntimeLayoutTest(unittest.TestCase): + def test_cuda_driver_is_external_but_cuda_runtime_is_packaged(self) -> None: + for module_path in PACKER_MODULE_PATHS: + mod = _load_module( + module_path=module_path, + module_name=f"vendor_runtime_cuda_external_{module_path.stem}", + ) + self.assertTrue(mod._is_external_driver_runtime_lib("libcuda.so.1")) + self.assertFalse(mod._is_external_driver_runtime_lib("libcudart.so.12")) + def test_ensure_vendor_runtime_soname_aliases_materializes_missing_alias_files(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: root = Path(tmpdir)