diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d155d8df32..6331f8b820 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -589,6 +589,18 @@ jobs: ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure shell: bash + # Mock-injected jetty ACK-timeout rebuild coverage (status=9 -> drain -> + # flush -> rebuild). Deterministic and hardware-free, scoped to the UB leg + # so the other matrices don't pay for it. + - name: Test (kunpeng URMA jetty rebuild) + if: matrix.name == 'ub-mock' + run: | + cd build-tent + ctest --test-dir mooncake-transfer-engine/tests \ + -R urma_jetty_rebuild_test \ + --output-on-failure + shell: bash + - name: Run sccache stat for check if: ${{ env.SCCACHE_PATH != '' }} shell: bash diff --git a/docs/source/design/transfer-engine/jetty-ack-timeout-rebuild.md b/docs/source/design/transfer-engine/jetty-ack-timeout-rebuild.md new file mode 100644 index 0000000000..4965d9c919 --- /dev/null +++ b/docs/source/design/transfer-engine/jetty-ack-timeout-rebuild.md @@ -0,0 +1,15 @@ +--- +orphan: true +--- + +# Jetty ACK Timeout 重建方案（已废弃） + +**状态：已废弃。请勿按本文实现或评审。** + +权威方案见： + +[jetty-single-rebuild-plan.md](./jetty-single-rebuild-plan.md) + +废弃原因：本文把「对端 import/bind 同步」写成阻塞前置条件；经对照 kunpeng +`UrmaEndpoint` 数据面（单边 READ/WRITE 打到 remote segment，不依赖对端 jetty +收发语义）后，采用**本端单 Jetty 排空重建、不需要对端协议同步**的方案。 diff --git a/docs/source/design/transfer-engine/jetty-rebuild-mock-test-plan.md b/docs/source/design/transfer-engine/jetty-rebuild-mock-test-plan.md new file mode 100644 index 0000000000..5dadd49033 --- /dev/null +++ b/docs/source/design/transfer-engine/jetty-rebuild-mock-test-plan.md @@ -0,0 +1,341 @@ +--- +orphan: true +--- + +# Jetty ACK Timeout 重建路径 — Mock 注入测试草案 + +状态：已实现（P0：TC-1～TC-3 + scoped ctest 挂 `tent-ci (ub-mock)`） +范围：kunpeng / UB（`UrmaEndpoint` + `mock_urma.cpp` + `UbWorkerPool`） +关联方案：`jetty-single-rebuild-plan.md`（见其"验证"章节的回链） +目标：在 **无真实 URMA 硬件** 的 CI 中，覆盖 status=9 触发的单 Jetty 重建与安全修复路径。 + +--- + +## 1. 动机 + +当前 `mock_urma.cpp` 的 `urma_poll_jfc` **恒返回 `URMA_CR_SUCCESS`**，`urma_flush_jetty` **恒返回 0**。 +因此现有 CI（含 `tent-ci (ub-mock)`）**无法执行**以下生产逻辑： + +| 路径 | 文件 | 现状 | +|------|------|------| +| ACK timeout → DRAINING | `onJettyError` | mock 不产出 status=9 | +| FLUSH_ERR_DONE → rebuild | `onFlushDone` | mock 不产出 fence CQE | +| flush CR 交付 | `rebuildJettyUnlocked` | mock flush 空返回 | +| stale epoch 丢弃 | `processWrCompletion` | 无法构造跨代 CQE | +| 延迟删除 | `deferred_deletes` | 难以稳定触发失败回退 | + +本草案通过 **增强 mock + 新增 gtest**，在 `USE_UB=ON` 的 CI job 中补齐覆盖。 + +--- + +## 2. 设计原则 + +1. **注入点只在 mock 层** — 不在 `UrmaContext::poll` 生产代码里加“伪造 status”开关。 +2. **确定性优先** — 测试用 scripted 状态机；环境变量仅作调试辅助。 +3. **窄而深** — 先覆盖 happy path + 两个 failure 回退，不复制 `ub_transport_test` 全链路压测。 +4. **与 TENT 分层** — 测的是 `kunpeng_transport/urma/urma_endpoint.cpp`，不是 `tent/urma_adapter.cpp`。 + +--- + +## 3. Mock 增强设计 + +### 3.1 数据结构扩展（`mock_urma.cpp`） + +在 `JfcState` 中除 `pending_ctx` 外，为每次 post 记录元数据： + +```cpp +struct PendingWr { + uint64_t user_ctx; + uint32_t jetty_local_id; // jetty->jetty_id.id +}; + +struct JfcState { + std::mutex mutex; + std::deque pending; +}; +``` + +`urma_post_jetty_send_wr`：把 WR 链上每个 `user_ctx` 与 `jetty->jetty_id.id` 入队。 + +### 3.2 测试脚本 API（仅 mock / 测试可见） + +新增头文件 `mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h`： + +```cpp +#pragma once +#include "urma_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// 重置 mock 全局状态（每个 TEST 的 SetUp 调用） +void mock_urma_test_reset(void); + +// 为下一次从该 JFC poll 出的每个 WR 指定 completion status。 +// 返回后脚本清除；若未设置则默认 URMA_CR_SUCCESS。 +void mock_urma_set_next_poll_status(int status); + +// 在 poll 完当前 pending 后，额外注入一条 FLUSH_ERR_DONE（无 user_ctx）。 +void mock_urma_enqueue_flush_done(uint32_t jetty_local_id); + +// rebuild 内 urma_flush_jetty：返回 count 条 URMA_CR_WR_FLUSH_ERR，user_ctx 来自 +// 该 jetty 尚未完成的 pending（或专用 flush 队列）。 +void mock_urma_set_flush_returns_errors(int count); + +// 让 urma_create_jetty 下一次失败（测 rebuild 失败 → defer delete）。 +void mock_urma_fail_next_create_jetty(int error); + +#ifdef __cplusplus +} +#endif +``` + +实现放在 `mock_urma.cpp` 末尾，`#ifdef BUILD_MOCK_URMA_TEST_CTRL` 或与 `USE_UB` 同条件编译。 + +### 3.3 `urma_poll_jfc` 行为（scripted） + +伪代码： + +```cpp +int urma_poll_jfc(urma_jfc_t* jfc, int n, urma_cr_t* cr) { + // 1) 若队列头有 injected FLUSH_ERR_DONE marker，先返回它（user_ctx=0） + // 2) 否则从 pending 弹出最多 n 条 WR + // status = mock_next_poll_status 或 URMA_CR_SUCCESS + // user_ctx / local_id 从 PendingWr 填充 + // 3) 返回条数 +} +``` + +### 3.4 `urma_flush_jetty` 行为 + +当 `mock_urma_set_flush_returns_errors(k)` 生效时： + +- 从该 jetty 关联的 pending（或已 ERROR 未完成的 WR）弹出最多 `cr_cnt` 条； +- 每条 `status = URMA_CR_WR_FLUSH_ERR`，`user_ctx` 有效； +- 返回实际条数；脚本计数递减至 0 后恢复“返回 0”。 + +### 3.5 其它 mock 钩子 + +| API | 钩子用途 | +|-----|----------| +| `urma_create_jetty` | `fail_next_create_jetty` → rebuild 失败路径 | +| `urma_modify_jetty` | 保持成功；可选 `fail_next_modify_error` 测 modify(ERROR) 失败 | +| `urma_import_jetty` / `bind` | 保持成功；可选失败钩子测 bind 失败清理 | + +--- + +## 4. 测试文件与夹具 + +### 4.1 新测试目标 + +| 项 | 值 | +|----|-----| +| 源文件 | `mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp` | +| 可执行名 | `urma_jetty_rebuild_test` | +| 条件编译 | `if(USE_UB)` | +| 链接 | `transfer_engine`, `gtest`, `gtest_main` | +| ctest 名 | `urma_jetty_rebuild_test` | + +**不要**复用 `ub_transport_test`（全链路、etcd、曾有不稳定注释）；新测试应 **in-proc 双端** 或 **单 context + 合成 slice**，参考 `rdma_endpoint_state_test.cpp` 的轻量夹具风格。 + +### 4.2 建议夹具结构 + +```cpp +class UrmaJettyRebuildTest : public ::testing::Test { + protected: + void SetUp() override { + mock_urma_test_reset(); + // TransferEngine(false) + protocol=ub + mock_urma_device + // 起 1 个 target + 1 个 initiator segment（或内存 metadata mock） + // 等待 worker poll 线程运行 + } + void TearDown() override { + mock_urma_test_reset(); + // engine shutdown + } + + void PostOneSliceAndWait(...); + uint64_t CurrentJettyEpoch(UrmaEndpoint* ep, int slot); +}; +``` + +若 metadata/etcd 过重，**Phase 1** 可只做 **UrmaContext + UrmaEndpoint 单元级** 测试（直接 `submitPostSend` + 手动 `performPoll` 一轮），不拉完整 `TransferEngine`。 + +### 4.3 可选 TestPeer（与 RDMA 测试一致） + +```cpp +class UrmaEndpointTestPeer { + public: + static JettyState state(UrmaEndpoint& ep, int slot); + static uint64_t epoch(UrmaEndpoint& ep, int slot); + static bool isDraining(UrmaEndpoint& ep); +}; +``` + +优先通过 **可观测行为**（slice 状态、post 是否恢复、endpoint 是否仍存在）断言，减少对 private 成员的直接访问。 + +--- + +## 5. 测试用例（首批） + +### TC-1 `AckTimeoutTriggersRebuildHappyPath`（P0） + +**目的**：status=9 → DRAINING → FLUSH_ERR_DONE → rebuild → slice 完成。 + +| 步骤 | Mock / 动作 | 期望 | +|------|-------------|------| +| 1 | post 1 WR | `wr_depth` +1 | +| 2 | `set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR)`，poll 一轮 | `onJettyError`；jetty → DRAINING | +| 3 | `enqueue_flush_done(jetty_id)`，poll 一轮 | `onFlushDone` → rebuild | +| 4 | `set_flush_returns_errors(1)`（rebuild 内） | slice 进 failed 或 retry，**非**永久 POSTED | +| 5 | 后续 poll SUCCESS 或 retry 成功 | jetty ACTIVE；`jetty_epoch` +1；endpoint 未 delete | + +### TC-2 `FlushCompletionsDeliveredOnRebuild`（P0） + +**目的**：验证 #1 修复 — rebuild 内 flush CR 必须 `processWrCompletion`。 + +| 步骤 | 期望 | +|------|------| +| rebuild 前 post N 个 WR | | +| flush 返回 N 条 `URMA_CR_WR_FLUSH_ERR` | 每条 slice 离开 POSTED（failed 或 success 路径） | +| `jetty_depth_set` 记账与 `jfc_outstanding` 一致 | 无双重扣减 | + +### TC-3 `StaleEpochCompletionDropped`（P0） + +**目的**：验证 #4 — 旧代 CQE 不二次完成、不污染 depth。 + +| 步骤 | 期望 | +|------|------| +| 完成 rebuild（epoch++） | | +| mock 注入 **旧 epoch** 的 SUCCESS CQE | `processWrCompletion` 返回 false；slice 不再 markSuccess | +| 新 post 的 WR 正常 SUCCESS | | + +### TC-4 `DeferredDeleteAfterRebuildFailure`（P1） + +**目的**：验证 #2 — poll 内不同步 `deleteEndpointByPtr`。 + +| 步骤 | 期望 | +|------|------| +| `fail_next_create_jetty` + 触发 rebuild | rebuild 失败 | +| 同一 poll 轮内仍有未处理 CR | 不 UAF（可用 ASAN/Debug 构建） | +| poll 返回后 | endpoint 进入 deferred delete；`wr_depth_list_` 仍有效至记账完成 | + +### TC-5 `ModifyErrorFailureDefersDelete`（P1） + +| 步骤 | 期望 | +|------|------| +| `fail_next_modify_error` + status=9 | `modify(ERROR)` 失败 → defer delete，不在 poll 循环内 deconstruct | + +### TC-6 `PollDoesNotDeleteBeforeDepthAccounting`（P1，回归） + +**目的**：钉死原 UAF。 + +| 步骤 | 期望 | +|------|------| +| 同批 CR：先触发 rebuild 失败，后还有 WR completion | 后者仍能安全读 `jetty_depth` | +| `performPoll` 结束后 | 才 `deleteEndpointByPtr` | + +### TC-7 `DisconnectFlushesNonActiveJetty`（P2） + +| 步骤 | 期望 | +|------|------| +| jetty 处于 DRAINING | `disconnectUnlocked` 先 sync flush | +| RESET 失败 | `jetty_state` 不强制 ACTIVE | + +### 不在首批范围 + +- `PENDING_DRAIN` 多槽串行（文档有、代码未落地） +- 真机 ACK timeout 时序 +- 多 worker 并发压力 +- `ub_transport_test` 级别跨节点 etcd 全链路 + +--- + +## 6. CI 集成 + +### 6.1 挂载点 + +在 **`.github/workflows/ci.yml`** 的 `tent-ci (ub-mock)` job（已 `-DUSE_UB=ON -DCMAKE_BUILD_TYPE=Debug`）增加： + +```yaml +- name: Test (kunpeng URMA jetty rebuild) + if: matrix.name == 'ub-mock' + run: | + cd build-tent + ctest --test-dir mooncake-transfer-engine/tests \ + -R urma_jetty_rebuild_test \ + --output-on-failure +``` + +### 6.2 CMake 注册 + +```cmake +if(USE_UB) + add_executable(urma_jetty_rebuild_test + ${WORKSPACE}/urma_jetty_rebuild_test.cpp) + target_link_libraries(urma_jetty_rebuild_test PUBLIC transfer_engine + gtest gtest_main glog::glog pthread) + target_compile_definitions(urma_jetty_rebuild_test + PRIVATE MOCK_URMA_TEST_CTRL=1) + add_test(NAME urma_jetty_rebuild_test COMMAND urma_jetty_rebuild_test) +endif() +``` + +`mock_urma_test_ctrl` 接口建议始终编入 `mock_urma.cpp`（仅 UB 构建），测试目标通过宏启用额外钩子。 + +### 6.3 本地运行 + +```bash +cmake -G Ninja .. -DUSE_UB=ON -DCMAKE_BUILD_TYPE=Debug -DBUILD_UNIT_TESTS=ON +cmake --build . --target urma_jetty_rebuild_test +ctest -R urma_jetty_rebuild_test -V +``` + +--- + +## 7. 实现分期 + +| 阶段 | 内容 | 预估 | +|------|------|------| +| **P0-a** | `PendingWr` + poll 可注入 status=9 / FLUSH_ERR_DONE | 1–2 天 | +| **P0-b** | `urma_flush_jetty` 返回 WR_FLUSH_ERR + TC-1/TC-2 | 1 天 | +| **P0-c** | TC-3 stale epoch + CMake/ctest/CI 挂钩 | 0.5 天 | +| **P1** | TC-4/5/6 failure 路径 + ASAN Debug | 1–2 天 | +| **P2** | TC-7 disconnect；`PENDING_DRAIN` 用例（实现代码后） | 后续 | + +建议 **单独 PR**（或 PR #33 的 follow-up），标题示例： + +`[TransferEngine] Add mock-injected tests for jetty ACK timeout rebuild` + +--- + +## 8. 验收标准 + +- [ ] `tent-ci (ub-mock)` 稳定通过，新增 ctest ≤ 30s +- [ ] TC-1～TC-3 在本地与 CI 绿 +- [ ] Debug + ASAN 构建下 TC-4/TC-6 无 UAF 报告 +- [ ] mock 钩子 **默认关闭**，不影响现有 `ub_transport_test` 手动跑法 +- [ ] 文档：本页 + `jetty-single-rebuild-plan.md` 增加交叉链接 + +--- + +## 9. 与真机测试的分工 + +| 层次 | Mock gtest（CI） | 服务器真机 | +|------|------------------|------------| +| 状态机转移 | ✅ | ✅ | +| flush/fence 时序 | 近似 | ✅ | +| 多 Jetty / 8 节点 | 可选扩展 | ✅ | +| 性能 / 长期稳定性 | ❌ | ✅ | + +CI 证明 **逻辑正确性**；真机证明 **provider 语义与生产负载**。 + +--- + +## 10. 参考 + +- 方案：`docs/source/design/transfer-engine/jetty-single-rebuild-plan.md` +- Mock 现状：`mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp` +- 轻量夹具范例：`mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp` +- CI job：`/.github/workflows/ci.yml` → `tent-ci (ub-mock)` diff --git a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md new file mode 100644 index 0000000000..24c6e527e9 --- /dev/null +++ b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md @@ -0,0 +1,239 @@ +--- +orphan: true +--- + +# Jetty ACK Timeout（status=9）单 Jetty 重建方案 + +状态：实现中（分支 `feat/jetty-ack-timeout-single-rebuild`）— **权威方案**（取代已废弃的 `jetty-ack-timeout-rebuild.md`） +范围：kunpeng / UB 传输路径（`UrmaEndpoint` / `UbWorkerPool`）；不涉及 tent +关联错误码： + +| CR status | 枚举 | 本方案处理 | +|---|---|---| +| 4 | `URMA_CR_LOC_ACCESS_ERR` | 不重建；保持现有 slice 重试 / 失败逻辑 | +| 9 | `URMA_CR_ACK_TIMEOUT_ERR` | 单 Jetty 排空 → 删除 → 重建，**纯本地操作** | + +--- + +## 1. 背景与问题 + +8 节点场景下，Jetty completion 出现 status=9（ACK 超时 / 重传超限）后，现有逻辑： + +1. poll 到非 SUCCESS → slice 进入重试 +2. 重试耗尽 → `deleteEndpointByPtr` → `UrmaEndpoint::deconstruct` +3. `deconstruct` 直接 `unbind / unimport / delete_jetty`，**没有** ERROR 排空与 `FLUSH_ERR_DONE` 栅栏 + +结果：短暂路径问题被放大成整 endpoint（甚至整节点）不可用，常需重启才能恢复。 + +--- + +## 2. 核心思路 + +**只重建出问题的那个 jetty，不删整个 EP，且不需要对端同步协议。** + +``` +status=9 出现在 jetty_list_[i] + ① urma_modify_jetty(jetty_list_[i], ERROR) ← 排空 + ② poll 到 FLUSH_ERR_DONE(local_id) ← 排空完成栅栏 + ③ flush → unbind/unimport → delete → create + → import(原 peer id) → bind → 更新槽位 ← 安全替换（仍纯本地） +``` + +**为什么不需要对端同步？** + +kunpeng 数据面是单边 `URMA_OPC_READ/WRITE`，DMA 目标是对端 **segment** +（`slice->ub.r_seg`），不是往对端 jetty 做 SEND/RECV。 + +`urma_post_jetty_send_wr` 的第一个参数是**本地 jetty**。握手虽会 +`import` + `bind` 并对 `wr.tjetty` 赋值，但： + +- 现有路径在 `tjetty` 缺失时只 `LOG(ERROR)`，仍继续 post +- `remote_jetty == NULL` 为空分支，未当硬失败 +- 重建后本端用**仍有效的对端 jetty id** 做 re-import/rebind 即可恢复本端发出方向；无需通知对端改其 import 视图 + +反向流量若仍绑旧本端 id，对端可能自行报错并走自己的本地重建——两侧各自恢复，不上同步握手。 + +--- + +## 3. 状态机 + +每个 jetty 槽位有独立状态： + +``` +ACTIVE ──(status=9，无其它槽在重建)──► DRAINING ──(FLUSH_ERR_DONE)──► REBUILDING + │ │ + │ ◄── create + 本端 rebind + 更新槽位 ───────────────┘ + │ + └──(status=9，已有槽在重建)──► PENDING_DRAIN ──(前槽重建完成后触发)──► DRAINING + + 任一环节失败 / 超时 ──► 回退：deleteEndpointByPtr +``` + +- `ACTIVE`：正常收发 +- `DRAINING`：已 `modify(ERROR)`，禁止 post，等待 `FLUSH_ERR_DONE` +- `REBUILDING`：排空完成，正在 flush / 删建 / rebind +- `PENDING_DRAIN`：本槽也收到 status=9，但同 EP 已有 jetty 在重建；禁止 post + （选槽跳过，避免故障槽继续吃流量），等前一个重建完成后由其尾部串行触发 + `modify(ERROR)` 进入 DRAINING +- 同 EP 任意时刻最多一个 jetty 处于 DRAINING/REBUILDING（串行） + +--- + +## 4. 关键设计 + +### 4.1 per-jetty 状态 + +```cpp +enum JettyState { ACTIVE, DRAINING, REBUILDING, PENDING_DRAIN }; +std::vector jetty_state_; // 与 jetty_list_ 平行 +std::unordered_map jetty_id_map_; // jetty_id → slot index +std::vector peer_jetty_id_; // 每槽对端 id，delete 前保留 +std::vector jetty_epoch_; // 每槽重建代次，丢弃旧代次的迟到 CR +``` + +`peer_jetty_id_` 在握手 `doSetupConnection` 时写入；段 C 在 delete 本地 jetty +前依赖它做 re-import，不得只依赖即将销毁的 `imported_jetty_map_` 指针键。 + +### 4.2 选槽策略 + +`submitPostSend` 选槽时跳过非 ACTIVE 槽： + +``` +随机选一个槽 → 非 ACTIVE 则重试 → 全部不可用则返回 0（上层重试） +``` + +### 4.3 锁策略 + +| 阶段 | 是否持 `lock_` | 说明 | +|---|---|---| +| 段 A：modify(ERROR) | 是 | 与 post 互斥（UMDK 约束） | +| 段 B：等 FLUSH_ERR_DONE | 否 | 由 performPoll 顺带消费，不占锁 | +| 段 C：flush/删建/rebind | 是 | 替换 `jetty_list_` 与 `imported_jetty_map_` | + +### 4.4 FLUSH_ERR_DONE 消费 + +走现有 `performPoll` → `UrmaContext::poll`。**必须在按 `user_ctx` 解 slice +之前分流**（假 CQE 的 `user_ctx` 无效；现有空 ctx `continue` 会丢掉栅栏）： + +``` +if cr.status == FLUSH_ERR_DONE: + local_id → jetty_id_map_ → (endpoint, slot) + endpoint->onFlushDone(slot) // DRAINING → REBUILDING，触发段 C + continue // 假 CQE，不走 slice 路径，不改 outstanding slice 语义外的 depth 时需单独约定 +``` + +### 4.5 status=9 分流与触发时机 + +**首次** poll 到 `ACK_TIMEOUT_ERR (9)` 即触发排空（不等重试耗尽）；`onJettyError` +必须幂等（已在 DRAINING/REBUILDING/PENDING_DRAIN 则不再 `modify`）。 +若同 EP 已有 jetty 在重建，则把本槽标记为 `PENDING_DRAIN` 排队：选槽立即跳过 +它，等前一个重建完成后由其尾部调用 `startDrainUnlocked` 串行启动排空。 + +``` +if cr.status == ACK_TIMEOUT_ERR (9): + slot ← slice->ub.jetty_depth 反查，或 cr.local_id → jetty_id_map_ + slice->ub.endpoint → UrmaEndpoint + endpoint->onJettyError(slot) // 段 A + slice 计入 retry（换 ACTIVE 槽重发） +``` + +### 4.6 outstanding 记账 + +排空期间 inflight WR 会以 error/flush 类 CQE 回来。要求： + +- 带有效 `user_ctx` 的失败 CQE：仍走现有 `jetty_depth_set` / retry 路径扣 + `wr_depth_list_` 与 JFC outstanding +- `FLUSH_ERR_DONE`：不当作 slice；不得 `markSuccess` / 不得当失败 slice 入队 +- **每个 WR 恰好完成一次**：`modify(ERROR)` 后 inflight WR 要么以 + `WR_FLUSH_ERR` 经 JFC poll 回来，要么被段 C 的 `urma_flush_jetty` 回收；两者 + 统一经 `processWrCompletion` 记账（`jetty_depth_set` 延迟扣减 + poll 返回值 + 累计 JFC outstanding）。因此段 C 删除旧 jetty 后**不做**额外的 + `wr_depth_list_[slot]` / JFC 清零——延迟扣减要等 poll 返回后才 apply，在段 C + 里提前清零会双扣（实现后已删除原方案中的"归零兜底"） +- 旧 jetty 删除后 `++jetty_epoch_[slot]`；slice 在 post 时记录 + `slice->ub.jetty_epoch`，旧代次的迟到/重复 CR 在 `processWrCompletion` 里按 + epoch 不匹配直接丢弃，不参与记账、不再入 retry 队列 + +### 4.7 段 C 完整顺序（持锁） + +``` +1. urma_flush_jetty（回收残余 WR，逐条经 processWrCompletion 交付； + 注意与 poll 可能重复，见 §7） +2. unbind → unimport（旧本地 jetty 上的对端视图） +3. delete_jetty；从 jetty_id_map_ 删旧 id；++jetty_epoch_[slot] +4. create_jetty（同 JFC/JFR 配置） +5. import(peer_jetty_id_[slot]) → bind(新 jetty, imported) +6. 更新 jetty_list_[slot]、imported_jetty_map_、jetty_id_map_ +7. jetty_state_[slot] = ACTIVE +8. 扫描 PENDING_DRAIN 槽：有则立即 startDrainUnlocked（modify(ERROR)）， + 仍保持同 EP 串行 +``` + +任一步失败 → 回退 `deleteEndpointByPtr`。 + +--- + +## 5. 执行计划 + +### Step 1：加状态与映射 + +- `UrmaEndpoint` 增加 `jetty_state_`、`jetty_id_map_`、`peer_jetty_id_` +- `construct()` / 握手初始化，`deconstruct()` 清理 +- `submitPostSend` 选槽跳过非 ACTIVE + +### Step 2：poll 分流 + +- `UrmaContext::poll`：先识别 `FLUSH_ERR_DONE`（`local_id`），再处理 + `ACK_TIMEOUT_ERR`（endpoint + 槽位） +- 假 CQE 不走 slice 路径 + +### Step 3：段 A — 触发排空 + +- `onJettyError(slot)`：持锁 → 幂等检查 → `modify(ERROR)` → 标 DRAINING +- 记录 drain 起始时间，供超时用 + +### Step 4：段 B — 等待栅栏 + +- `onFlushDone(slot)`：标 REBUILDING，触发段 C + +### Step 5：段 C — 重建 + +- 按 §4.7 完整顺序执行 + +### Step 6：超时降级 + +- flush-done 等待超时（可配置，默认 3s）→ 回退 `deleteEndpointByPtr` +- create / import / bind 失败同样降级 + +### Step 7：日志与测试 + +- 关键路径日志：`jetty_id`、slot、耗时、降级原因 +- 单测：状态机、选槽跳过、假 CQE 路由、幂等 `onJettyError` +- 集成 / 故障注入：status=9 后该槽恢复 ACTIVE，同 EP 其它槽可继续；超时路径删 EP +- 无 UMDK 硬件时，flush 与真实 ACK timeout 行为标为硬件验证项 +- **Mock 注入 CI 测试草案**：`jetty-rebuild-mock-test-plan.md`（`mock_urma` 脚本化 + status=9 / FLUSH_ERR_DONE / flush CR，挂 `tent-ci (ub-mock)`） + +--- + +## 6. 改动文件 + +| 文件 | 改动 | +|---|---| +| `urma_endpoint.h` | `JettyState`（含 `PENDING_DRAIN`）、`jetty_state_`、`jetty_id_map_`、`peer_jetty_id_`、`jetty_epoch_`、`onJettyError()` / `onFlushDone()` / `startDrainUnlocked()` | +| `urma_endpoint.cpp` | 状态机、选槽跳过、段 A/C、握手写入 `peer_jetty_id_` | +| `ub_context.cpp` | poll / worker 侧配合 status=9 与 `FLUSH_ERR_DONE`（若分流落在 context poll 则改 `urma_endpoint.cpp` 中 `UrmaContext::poll`） | + +--- + +## 7. 风险与开放问题 + +1. **urma_flush_jetty 与 poll 重复**：flush 返回的 WR 级 CR 是否已在 JFC poll + 中出现过，需实测确认，避免 double-complete。实现按「每个 WR 恰好完成一次」 + 记账，并用 `jetty_epoch_` 把旧代次的迟到 CR 整体丢弃兜底；若实测发现 flush + 与 poll 会重复交付同一 WR，需重新评估 slice 指针解引用的安全性（届时 CR 里 + 的 `user_ctx` 可能指向已回收的 slice）。 +2. **共享 JFC 假 CQE 过滤**：多 jetty 共享 JFC 时，严格按 `local_id` 匹配，不能假设顺序。 +3. **超时阈值**：flush-done 等待 3s 是否合适，需结合 `err_timeout` 和现场标定。 +4. **硬件 hang 场景**：本方案解决软件放大故障；若根因是设备/驱动 hang，重建仍可能失败，保留删 EP 降级路径。 +5. **反向路径**：对端仍绑旧本端 id 时可能自行报 9 并本地重建；观察即可，本期不上对端协议。 diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h index 56767c9dee..db61ca1985 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -188,13 +188,19 @@ class UbContext { // * Successful slices have markSuccess() called in place and are NOT // returned (they may be recycled by the submitting thread the // moment markSuccess() runs). - // * Failed slices are returned in failed_slices[0..num_failed-1] for - // the caller to apply retry / markFailed. - // Returns the total number of completions polled (>= 0), or a negative - // error code. - virtual int poll(int num_entries, Transport::Slice** failed_slices, - int& num_failed, + // * Failed slices are appended to failed_slices for the caller to apply + // retry / markFailed. Implementations may also append completions + // recovered while draining a jetty, so the vector is unbounded. + // * Stale completions from a previous jetty generation are dropped and + // are not counted in the return value. + // * Endpoints scheduled for deletion are appended to deferred_deletes; + // the caller must delete them only after jetty_depth_set accounting. + // Returns the number of resolved WR completions (>= 0), or a negative + // error code. Fence markers and dropped stale completions are excluded. + virtual int poll(int num_entries, + std::vector& failed_slices, std::unordered_map& jetty_depth_set, + std::vector& deferred_deletes, int jfc_index = 0) = 0; virtual volatile int* outstandingCount(int jfc_index) = 0; @@ -259,6 +265,10 @@ class UbContext { UbTransport& engine() const { return engine_; } + bool traceWorkRequestFlushedErrors() const { + return show_work_request_flushed_error_; + } + uint8_t portNum() const { return port_; } int activeSpeed() const { return active_speed_; } diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h index 6b6fa7076a..e73c19d1ba 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h @@ -18,7 +18,10 @@ #include #include #include +#include +#include #include +#include #include "common.h" #include "config.h" #include "urma_api.h" @@ -47,8 +50,14 @@ static urma_import_seg_flag_t import_flag = { .reserved = 0}}; // define the UrmaContext class +class UrmaEndpoint; +// Test hook: lets gtest fixtures seed a minimal UrmaContext (jfc/jfr/urma +// context) without spinning up the worker pool that full construct() starts. +class UrmaContextTestPeer; + class UrmaContext : public UbContext { friend class UrmaEndpoint; + friend class UrmaContextTestPeer; public: UrmaContext(UbTransport& engine, std::string device_name, @@ -58,8 +67,9 @@ class UrmaContext : public UbContext { int unregisterMemoryRegion(uint64_t va) override; int doProcessContextEvents() override; void* retrieveRemoteSeg(const std::string& value) override; - int poll(int num_entries, Transport::Slice** failed_slices, int& num_failed, + int poll(int num_entries, std::vector& failed_slices, std::unordered_map& jetty_depth_set, + std::vector& deferred_deletes, int jfc_index) override; volatile int* outstandingCount(int jfc_index) override; int submitPostSend( @@ -81,6 +91,17 @@ class UrmaContext : public UbContext { static bool uninit(); static bool init(); + void registerJettyOwner(uint32_t jetty_id, UrmaEndpoint* endpoint, + int slot); + void unregisterJettyOwner(uint32_t jetty_id); + bool findJettyOwner(uint32_t jetty_id, UrmaEndpoint** endpoint, int* slot); + void addDrainingEndpoint(UrmaEndpoint* endpoint); + void removeDrainingEndpoint(UrmaEndpoint* endpoint); + void checkJettyDrainTimeouts( + std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes); + private: int construct(GlobalConfig& config) override; int deconstruct() override; @@ -146,11 +167,43 @@ class UrmaContext : public UbContext { urma_import_seg_flag_t import_flag_ = mooncake::import_flag; std::unordered_map import_tseg_map; + + RWSpinlock jetty_owner_lock_; + struct JettyOwner { + UrmaEndpoint* endpoint = nullptr; + int slot = -1; + }; + std::unordered_map jetty_owner_map_; + std::unordered_set draining_endpoints_; }; +// Test hook for asserting the jetty state machine without touching private +// members directly from the test body. +class UrmaEndpointTestPeer; + // define the UrmaEndpoint class class UrmaEndpoint : public UbEndPoint { + // UrmaContext::poll drives the jetty state machine via processWrCompletion. + friend class UrmaContext; + friend class UrmaEndpointTestPeer; + public: + // PENDING_DRAIN: the slot hit ACK timeout while another jetty of this + // endpoint was draining/rebuilding. It is excluded from post selection + // and drained once the in-flight rebuild completes (rebuilds are + // serialized per endpoint). + enum JettyState { + ACTIVE = 0, + DRAINING = 1, + REBUILDING = 2, + PENDING_DRAIN = 3, + // Rebuild failed after the old jetty was already torn down + // (unbind/unimport) but could not be fully deleted or replaced. The + // old handle is kept in jetty_list_[slot] so deconstruct can retry + // urma_delete_jetty instead of leaking it. Never selected for post. + REBUILDING_FAILED = 4 + }; + UrmaEndpoint(UrmaContext* context) : context_(context), jfc_outstanding_(nullptr) {} @@ -173,6 +226,21 @@ class UrmaEndpoint : public UbEndPoint { const std::string toString() const override; + // Called from UrmaContext::poll on ACK timeout / flush-done / drain + // timeout. + void onJettyError(int slot, std::vector& deferred_deletes); + void onFlushDone(int slot, + std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, + int& resolved_wr_count); + void checkDrainTimeout( + std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes); + + int findSlotByDepth(volatile int* depth) const; + private: void disconnectUnlocked() override; @@ -187,6 +255,36 @@ class UrmaEndpoint : public UbEndPoint { uint32_t peer_jetty_num, std::string* reply_msg = nullptr); + bool hasNonActiveJettyUnlocked() const; + int selectActiveJettyUnlocked(); + // Transitions `slot` to DRAINING via urma_modify_jetty(ERROR) and arms + // the flush-done wait. Caller must hold lock_. Returns ERR_ENDPOINT if + // modify failed (caller should fall back to deleting the endpoint). + int startDrainUnlocked(int slot); + int rebuildJettyUnlocked( + int slot, std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, int& resolved_wr_count); + + // Delivers one WR completion to the normal success/failure path. + // Returns true when the completion resolved a live WR (and must be counted + // in the JFC outstanding accounting). Returns false for completions from a + // stale jetty generation, which were already resolved during the rebuild + // flush and must be dropped entirely. When allow_error_trigger is false + // (the caller already holds lock_ while rebuilding/draining), an + // ACK_TIMEOUT completion is delivered as a plain failure without + // re-entering onJettyError. + bool processWrCompletion( + urma_cr_t& cr, std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, int jfc_index, + bool allow_error_trigger); + + int recreateJettyUnlocked(int slot, urma_jfc_t* reuse_jfc, + urma_jfr_t* reuse_jfr); + + static constexpr uint64_t kJettyDrainTimeoutNs = 3000000000ull; // 3s + private: UrmaContext* context_; urma_token_t urma_token = {.token = 0xACFE}; @@ -195,6 +293,14 @@ class UrmaEndpoint : public UbEndPoint { int max_wr_depth_; volatile int* jfc_outstanding_; std::unordered_map imported_jetty_map_; + + std::vector jetty_state_; + std::unordered_map jetty_id_map_; + std::vector peer_jetty_id_; + std::string peer_eid_; + uint64_t drain_start_ns_ = 0; + int draining_slot_ = -1; + std::vector jetty_epoch_; }; } // namespace mooncake #endif // URMA_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index c235bb3424..256d81e5a2 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -157,6 +157,7 @@ class Transport { struct { uint64_t dest_addr; volatile int *jetty_depth; + uint64_t jetty_epoch; uint32_t retry_cnt; uint32_t max_retry_cnt; void *r_seg; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index 46a3d1e467..60e58fe2f0 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -406,22 +406,26 @@ void UbWorkerPool::performPoll(int thread_id) { // are NOT returned from poll(), so this worker never dereferences them // after they may have been recycled by the submitting thread. std::unordered_map jetty_depth_set; + std::vector deferred_deletes; std::vector failed_slices; for (int jfc_index = thread_id; jfc_index < context_.jfcCount(); jfc_index += kTransferWorkerCount) { - UbTransport::Slice* failed[kPollCount]; - int num_failed = 0; - int nr_poll = context_.poll(kPollCount, failed, num_failed, - jetty_depth_set, jfc_index); - if (nr_poll < 0) { + // poll() may also append completions recovered while draining a jetty + // (up to the full queue depth), so this must be a vector rather than + // a fixed-size array of kPollCount. + std::vector failed; + const size_t failed_before = failed.size(); + int nr_resolved = context_.poll(kPollCount, failed, jetty_depth_set, + deferred_deletes, jfc_index); + if (nr_resolved < 0) { LOG(ERROR) << "Worker: Failed to poll jetty for complete"; continue; } - int num_success = nr_poll - num_failed; + const size_t new_failed = failed.size() - failed_before; + const int num_success = nr_resolved - static_cast(new_failed); success_nr_polls += num_success; processed_slice_count += num_success; - for (int i = 0; i < num_failed; ++i) { - UbTransport::Slice* slice = failed[i]; + for (auto* slice : failed) { assert(slice); failed_nr_polls++; if (context_.active() && failed_nr_polls > 32 && @@ -439,13 +443,17 @@ void UbWorkerPool::performPoll(int thread_id) { redispatch_counter_++; } } - if (nr_poll) - __sync_fetch_and_sub(context_.outstandingCount(jfc_index), nr_poll); + if (nr_resolved) + __sync_fetch_and_sub(context_.outstandingCount(jfc_index), + nr_resolved); } for (auto& entry : jetty_depth_set) __sync_fetch_and_sub(entry.first, entry.second); + for (auto* endpoint : deferred_deletes) + context_.deleteEndpointByPtr(endpoint); + // Slices that hit max_retry: final markFailed() after all reads (and the // jetty depth returns above) are done. Failed slices were never published // by poll(), so they remained safe to deref up to this point. diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp index c392577a6b..2d3ce24e01 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp @@ -1,4 +1,5 @@ #include "urma_api.h" +#include "mock_urma_test_ctrl.h" #include #include #include @@ -10,11 +11,49 @@ namespace { +// One outstanding work request posted to a JFC. We record the owning jetty's +// local id alongside user_ctx so that polled completions can populate +// urma_cr_t::local_id, which the production poll loop uses to locate the jetty +// owner for fence CQEs and ACK-timeout completions. +struct PendingWr { + uint64_t user_ctx; + uint32_t jetty_local_id; + // When true, urma_poll_jfc skips this WR (it stays outstanding) so only + // urma_flush_jetty can complete it. Lets a test leave a residual WR on the + // jetty to exercise the rebuild flush-delivery path. + bool withhold = false; +}; + struct JfcState { std::mutex mutex; - std::deque pending_ctx; + std::deque pending; +}; + +// Scripted test hooks. All are guarded by g_script_mutex and are inert +// (defaults) until a test arms them. mock_urma_test_reset() restores defaults. +struct MockScript { + // Poll-status override: while poll_status_count > 0, polled completions + // use poll_status instead of URMA_CR_SUCCESS. + urma_cr_status_t poll_status = URMA_CR_SUCCESS; + int poll_status_count = 0; + // Fence CQEs to inject, one per poll that owns the matching jetty. + std::deque flush_done_ids; + // Flush-error scripting: while flush_err_count > 0, urma_flush_jetty + // returns that many WR_FLUSH_ERR completions from the jetty's pending WRs. + int flush_err_count = 0; + // When true, the next urma_create_jetty returns NULL. + bool fail_next_create_jetty = false; + // When true, the next urma_delete_jetty returns an error without freeing, + // exercising the rebuild delete-failure path that must keep the handle. + bool fail_next_delete_jetty = false; + // While > 0, the next WRs posted via urma_post_jetty_send_wr are marked + // withhold (skipped by poll, only flushable). Decremented per WR posted. + int withhold_next_post_count = 0; }; +std::mutex g_script_mutex; +MockScript g_script; + std::shared_mutex g_rw_mutex; bool initialized = false; std::vector device_list; @@ -310,7 +349,7 @@ urma_target_seg_t *urma_import_seg(urma_context_t *ctx, urma_seg_t *seg, context_map.find(ctx) == context_map.end()) { return nullptr; } - urma_target_seg_t *tseg = new urma_target_seg_t; + auto *tseg = new urma_target_seg_t; tseg->seg = *seg; *token_value = {.token = seg->token_id}; seg_map[tseg] = 1; @@ -342,15 +381,23 @@ urma_status_t urma_get_async_event(urma_context_t *ctx, void urma_ack_async_event(urma_async_event_t *event) {} urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) { + { + std::lock_guard script_lock(g_script_mutex); + if (g_script.fail_next_create_jetty) { + g_script.fail_next_create_jetty = false; + return nullptr; + } + } std::unique_lock lock(g_rw_mutex); if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) { return nullptr; } + static std::atomic next_jetty_id{1}; urma_jetty_t *jetty = new urma_jetty_t; memset(&jetty->jetty_id.eid, 0, sizeof(urma_eid_t)); jetty->jetty_id.eid.raw[0] = 1; jetty->jetty_id.uasid = 0; - jetty->jetty_id.id = 1; + jetty->jetty_id.id = next_jetty_id.fetch_add(1); jetty->jetty_cfg = *cfg; jetty->remote_jetty = nullptr; jetty_map[jetty] = 1; @@ -358,6 +405,13 @@ urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) { } urma_status_t urma_delete_jetty(urma_jetty_t *jetty) { + { + std::lock_guard script_lock(g_script_mutex); + if (g_script.fail_next_delete_jetty) { + g_script.fail_next_delete_jetty = false; + return URMA_FAIL; // keep the jetty allocated; caller must retry + } + } std::unique_lock lock(g_rw_mutex); if (!jetty || jetty_map.find(jetty) == jetty_map.end()) { return URMA_EINVAL; @@ -420,6 +474,54 @@ urma_status_t urma_modify_jetty(urma_jetty_t *jetty, urma_jetty_attr_t *attr) { return URMA_SUCCESS; } +int urma_flush_jetty(urma_jetty_t *jetty, int cr_cnt, urma_cr_t *cr) { + urma_jfc_t *jfc = nullptr; + uint32_t local_id = 0; + { + std::shared_lock lock(g_rw_mutex); + if (!jetty || jetty_map.find(jetty) == jetty_map.end()) { + return -1; + } + jfc = jetty->jetty_cfg.jfs_cfg.jfc; + local_id = jetty->jetty_id.id; + } + + int want = 0; + { + std::lock_guard script_lock(g_script_mutex); + if (g_script.flush_err_count > 0) { + want = std::min(cr_cnt, g_script.flush_err_count); + g_script.flush_err_count -= want; + } + } + if (want <= 0) return 0; + + // Draw up to `want` of this jetty's outstanding WRs from its JFC and + // report them as flushed, mirroring the real provider's flush completion. + JfcState *state = nullptr; + { + std::shared_lock lock(g_rw_mutex); + auto it = jfc_state_map.find(jfc); + if (it == jfc_state_map.end()) return 0; + state = it->second; + } + std::lock_guard jfc_lock(state->mutex); + int produced = 0; + for (auto it = state->pending.begin(); + it != state->pending.end() && produced < want;) { + if (it->jetty_local_id != local_id) { + ++it; + continue; + } + cr[produced].status = URMA_CR_WR_FLUSH_ERR; + cr[produced].user_ctx = it->user_ctx; + cr[produced].local_id = it->jetty_local_id; + ++produced; + it = state->pending.erase(it); + } + return produced; +} + urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, urma_jfs_wr_t **bad_wr) { { @@ -450,7 +552,16 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, std::lock_guard jfc_lock(state->mutex); urma_jfs_wr_t *current_wr = wr; while (current_wr) { - state->pending_ctx.push_back(current_wr->user_ctx); + bool withhold = false; + { + std::lock_guard script_lock(g_script_mutex); + if (g_script.withhold_next_post_count > 0) { + withhold = true; + --g_script.withhold_next_post_count; + } + } + state->pending.push_back( + PendingWr{current_wr->user_ctx, jetty->jetty_id.id, withhold}); current_wr = current_wr->next; } } @@ -472,17 +583,100 @@ int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) { state = it->second; } + std::lock_guard jfc_lock(state->mutex); int num_completed = 0; + + // Scripted flush-done fence: emit at most one per poll, matching how the + // real provider delivers a single FLUSH_ERR_DONE marker per drained jetty. + // It carries user_ctx=0 and the owning jetty's local_id so the production + // poll loop routes it to onFlushDone. { - std::lock_guard jfc_lock(state->mutex); - int available = static_cast(state->pending_ctx.size()); - num_completed = std::min(num_entries, available); - for (int i = 0; i < num_completed; ++i) { - cr_list[i].status = URMA_CR_SUCCESS; - cr_list[i].user_ctx = state->pending_ctx[i]; + std::lock_guard script_lock(g_script_mutex); + if (num_completed < num_entries && !g_script.flush_done_ids.empty()) { + uint32_t fence_id = g_script.flush_done_ids.front(); + bool owns = false; + for (const auto &p : state->pending) { + if (p.jetty_local_id == fence_id) { + owns = true; + break; + } + } + // The jetty may already be fully flushed (no pending WRs left); + // still deliver the fence on this JFC if it ever hosted the jetty. + // For the mock's single-JFC-per-endpoint usage we deliver on the + // first JFC when no pending match is found. + if (owns || state->pending.empty()) { + g_script.flush_done_ids.pop_front(); + cr_list[num_completed].status = URMA_CR_WR_FLUSH_ERR_DONE; + cr_list[num_completed].user_ctx = 0; + cr_list[num_completed].local_id = fence_id; + ++num_completed; + } } - state->pending_ctx.erase(state->pending_ctx.begin(), - state->pending_ctx.begin() + num_completed); + } + + int capacity = num_entries - num_completed; + int wr_completed = 0; + for (auto it = state->pending.begin(); + it != state->pending.end() && wr_completed < capacity;) { + if (it->withhold) { + ++it; // leave outstanding for the flush path + continue; + } + urma_cr_status_t status = URMA_CR_SUCCESS; + { + std::lock_guard script_lock(g_script_mutex); + if (g_script.poll_status_count > 0) { + status = g_script.poll_status; + --g_script.poll_status_count; + } + } + cr_list[num_completed].status = status; + cr_list[num_completed].user_ctx = it->user_ctx; + cr_list[num_completed].local_id = it->jetty_local_id; + ++num_completed; + ++wr_completed; + it = state->pending.erase(it); } return num_completed; } + +// --------------------------------------------------------------------------- +// Test-only scripted control API. See mock_urma_test_ctrl.h. +// --------------------------------------------------------------------------- + +void mock_urma_test_reset(void) { + std::lock_guard script_lock(g_script_mutex); + g_script = MockScript{}; +} + +void mock_urma_set_next_poll_status(int status, int count) { + std::lock_guard script_lock(g_script_mutex); + g_script.poll_status = static_cast(status); + g_script.poll_status_count = count; +} + +void mock_urma_enqueue_flush_done(uint32_t jetty_local_id) { + std::lock_guard script_lock(g_script_mutex); + g_script.flush_done_ids.push_back(jetty_local_id); +} + +void mock_urma_set_flush_returns_errors(int count) { + std::lock_guard script_lock(g_script_mutex); + g_script.flush_err_count = count; +} + +void mock_urma_withhold_next_post(int count) { + std::lock_guard script_lock(g_script_mutex); + g_script.withhold_next_post_count = count; +} + +void mock_urma_fail_next_delete_jetty(void) { + std::lock_guard script_lock(g_script_mutex); + g_script.fail_next_delete_jetty = true; +} + +void mock_urma_fail_next_create_jetty(void) { + std::lock_guard script_lock(g_script_mutex); + g_script.fail_next_create_jetty = true; +} diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h new file mode 100644 index 0000000000..3f42167657 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h @@ -0,0 +1,66 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Test-only scripted control for the mock URMA provider. These hooks let a +// gtest drive the jetty rebuild state machine (ACK timeout -> drain -> flush +// -> rebuild) without real URMA hardware. They are compiled into mock_urma.cpp +// only and are inert until a test arms them via the setters below. +#ifndef MOCK_URMA_TEST_CTRL_H +#define MOCK_URMA_TEST_CTRL_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Clears every scripted hook and all queued completions. Call from each +// test's SetUp/TearDown so scripts never leak across cases. +void mock_urma_test_reset(void); + +// Overrides the completion status for the next completions polled out of any +// JFC. Applies to up to `count` completions, then auto-clears. Pass +// status=URMA_CR_ACK_TIMEOUT_ERR (9) to drive the rebuild path. +void mock_urma_set_next_poll_status(int status, int count); + +// Queues a synthetic FLUSH_ERR_DONE fence CQE carrying the given local jetty +// id (user_ctx=0). The next urma_poll_jfc on the JFC that owns that jetty +// returns it, which drives onFlushDone -> rebuildJettyUnlocked. +void mock_urma_enqueue_flush_done(uint32_t jetty_local_id); + +// Makes the next urma_flush_jetty call return up to `count` completions with +// status URMA_CR_WR_FLUSH_ERR, drawn from the jetty's outstanding WRs, so the +// rebuild path can deliver them through processWrCompletion. +void mock_urma_set_flush_returns_errors(int count); + +// Makes the next urma_create_jetty call return NULL, exercising the +// rebuild-failure -> deferred-delete fallback. +void mock_urma_fail_next_create_jetty(void); + +// Makes the next urma_delete_jetty call return an error WITHOUT freeing the +// jetty, exercising the rebuild delete-failure path that must keep the old +// handle (REBUILDING_FAILED) so deconstruct can retry the delete. +void mock_urma_fail_next_delete_jetty(void); + +// Marks the next `count` WRs posted via urma_post_jetty_send_wr as withheld: +// urma_poll_jfc skips them (they stay outstanding) so only urma_flush_jetty +// can complete them. Use to leave a residual WR on a jetty and exercise the +// rebuild flush-delivery path. +void mock_urma_withhold_next_post(int count); + +#ifdef __cplusplus +} +#endif + +#endif // MOCK_URMA_TEST_CTRL_H diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp index 06c65d0385..457de7ba36 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp @@ -518,11 +518,69 @@ bool UrmaContext::transEidFromString(const std::string& eid_str, return index == URMA_EID_SIZE; } -int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices, - int& num_failed, +void UrmaContext::registerJettyOwner(uint32_t jetty_id, UrmaEndpoint* endpoint, + int slot) { + RWSpinlock::WriteGuard guard(jetty_owner_lock_); + jetty_owner_map_[jetty_id] = JettyOwner{endpoint, slot}; +} + +void UrmaContext::unregisterJettyOwner(uint32_t jetty_id) { + RWSpinlock::WriteGuard guard(jetty_owner_lock_); + jetty_owner_map_.erase(jetty_id); +} + +bool UrmaContext::findJettyOwner(uint32_t jetty_id, UrmaEndpoint** endpoint, + int* slot) { + RWSpinlock::ReadGuard guard(jetty_owner_lock_); + auto it = jetty_owner_map_.find(jetty_id); + if (it == jetty_owner_map_.end()) return false; + if (endpoint) *endpoint = it->second.endpoint; + if (slot) *slot = it->second.slot; + return true; +} + +void UrmaContext::addDrainingEndpoint(UrmaEndpoint* endpoint) { + RWSpinlock::WriteGuard guard(jetty_owner_lock_); + draining_endpoints_.insert(endpoint); +} + +void UrmaContext::removeDrainingEndpoint(UrmaEndpoint* endpoint) { + RWSpinlock::WriteGuard guard(jetty_owner_lock_); + draining_endpoints_.erase(endpoint); +} + +void UrmaContext::checkJettyDrainTimeouts( + std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes) { + std::vector endpoints; + { + RWSpinlock::ReadGuard guard(jetty_owner_lock_); + endpoints.assign(draining_endpoints_.begin(), + draining_endpoints_.end()); + } + for (auto* endpoint : endpoints) { + if (endpoint) + endpoint->checkDrainTimeout(jetty_depth_set, failed_slices, + deferred_deletes); + } +} + +namespace { +void deferEndpointDelete(UrmaEndpoint* endpoint, + std::vector& deferred_deletes) { + for (auto* ep : deferred_deletes) { + if (ep == endpoint) return; + } + deferred_deletes.push_back(endpoint); +} +} // namespace + +int UrmaContext::poll(int num_entries, + std::vector& failed_slices, std::unordered_map& jetty_depth_set, + std::vector& deferred_deletes, int jfc_index) { - num_failed = 0; urma_cr_t cr[num_entries]; int nr_poll = urma_poll_jfc(jfc_list_[jfc_index].native, num_entries, cr); if (nr_poll < 0) { @@ -530,53 +588,43 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices, << device_name_; return ERR_CONTEXT; } + int wr_completions = 0; for (int i = 0; i < nr_poll; ++i) { + // Fake fence CQE: user_ctx is invalid; match by local_id only. + if (cr[i].status == URMA_CR_WR_FLUSH_ERR_DONE) { + UrmaEndpoint* endpoint = nullptr; + int slot = -1; + if (findJettyOwner(cr[i].local_id, &endpoint, &slot) && endpoint) { + endpoint->onFlushDone(slot, jetty_depth_set, failed_slices, + deferred_deletes, wr_completions); + } else { + LOG(WARNING) << "FLUSH_ERR_DONE for unknown jetty local_id=" + << cr[i].local_id << " on " << device_name_; + } + continue; + } + auto slice = (Transport::Slice*)cr[i].user_ctx; if (!slice) { continue; } - - // All deref of `slice` (including the jetty_depth aggregation below) - // MUST happen before markSuccess(): once that publishes completion, - // the submitting thread may recycle the slice immediately. - auto* depth = slice->ub.jetty_depth; - auto it = jetty_depth_set.find(depth); - if (it != jetty_depth_set.end()) - it->second++; - else - jetty_depth_set[depth] = 1; - - if (cr[i].status == URMA_CR_SUCCESS) { - // Safe to publish here — we are done with this slice and do not - // return it to the caller, so no one else will deref it. - slice->markSuccess(); + auto* endpoint = static_cast(slice->ub.endpoint); + if (!endpoint) { + // Keep the legacy accounting: the WR resolved even though we no + // longer know its endpoint. + ++wr_completions; continue; } - - if (cr[i].status != URMA_CR_WR_FLUSH_ERR || - show_work_request_flushed_error_) - LOG(ERROR) << "Worker: Process failed for slice (opcode: " - << slice->opcode - << ", source_addr: " << slice->source_addr - << ", length: " << slice->length - << ", dest_addr: " << (void*)slice->ub.dest_addr - << ", local_nic: " << deviceName() - << ", peer_nic: " << slice->peer_nic_path - << ", dest_seg_tokenid: " - << static_cast(slice->ub.r_seg) - ->seg.token_id - << ", retry_cnt: " << slice->ub.retry_cnt - << "): " << cr[i].status << ", jfc idx : " << jfc_index - << ", comp_events_acked: " - << jfc_list_[jfc_index].native->comp_events_acked << " " - << jfc_list_[jfc_index].native->async_events_acked; - - // Failed: hand the slice back so the caller can decide retry vs - // final markFailed(). Slice is NOT published, so the caller may - // safely deref it. - failed_slices[num_failed++] = slice; - } - return nr_poll; + if (endpoint->processWrCompletion(cr[i], jetty_depth_set, failed_slices, + deferred_deletes, jfc_index, + /*allow_error_trigger=*/true)) { + ++wr_completions; + } + } + checkJettyDrainTimeouts(jetty_depth_set, failed_slices, deferred_deletes); + // Exclude FLUSH_ERR_DONE and dropped stale completions from outstanding + // accounting; the rebuild path already accounted for the stale ones. + return wr_completions; } volatile int* UrmaContext::outstandingCount(int jfc_index) { @@ -626,6 +674,13 @@ int UrmaEndpoint::construct(GlobalConfig& config) { } jetty_list_.resize(num_jetty_list); + jetty_state_.assign(num_jetty_list, ACTIVE); + jetty_epoch_.assign(num_jetty_list, 1); + peer_jetty_id_.assign(num_jetty_list, 0); + jetty_id_map_.clear(); + peer_eid_.clear(); + drain_start_ns_ = 0; + draining_slot_ = -1; auto* jfc = context_->jfc(); jfc_outstanding_ = (volatile int*)jfc->jfc_cfg.user_ctx; @@ -658,10 +713,26 @@ int UrmaEndpoint::construct(GlobalConfig& config) { jetty_list_[i] = urma_create_jetty(context_->urma_context_, &attr); if (!jetty_list_[i]) { PLOG(ERROR) << "Failed to create jetty"; + for (size_t j = 0; j < i; ++j) { + if (!jetty_list_[j]) continue; + context_->unregisterJettyOwner(jetty_list_[j]->jetty_id.id); + urma_delete_jetty(jetty_list_[j]); + jetty_list_[j] = nullptr; + } + jetty_list_.clear(); + jetty_state_.clear(); + jetty_epoch_.clear(); + peer_jetty_id_.clear(); + jetty_id_map_.clear(); + delete[] wr_depth_list_; + wr_depth_list_ = nullptr; return ERR_ENDPOINT; } - LOG(INFO) << "Create jetty success, jetty id = " - << jetty_list_[i]->jetty_id.id << " ,jetty jfc id = " + uint32_t jetty_id = jetty_list_[i]->jetty_id.id; + jetty_id_map_[jetty_id] = static_cast(i); + context_->registerJettyOwner(jetty_id, this, static_cast(i)); + LOG(INFO) << "Create jetty success, jetty id = " << jetty_id + << " ,jetty jfc id = " << jetty_list_[i]->jetty_cfg.jfs_cfg.jfc->jfc_id.id << " : " << jfc->jfc_id.id; } @@ -672,19 +743,25 @@ int UrmaEndpoint::construct(GlobalConfig& config) { int UrmaEndpoint::deconstruct() { int ret = 0; + context_->removeDrainingEndpoint(this); for (size_t i = 0; i < jetty_list_.size(); ++i) { auto imported_it = imported_jetty_map_.find(jetty_list_[i]); auto imported_jetty = (imported_it != imported_jetty_map_.end()) ? imported_it->second : nullptr; - ret = urma_unbind_jetty(jetty_list_[i]); - if (ret) PLOG(ERROR) << "Failed to unbind jetty"; + if (jetty_list_[i]) { + context_->unregisterJettyOwner(jetty_list_[i]->jetty_id.id); + ret = urma_unbind_jetty(jetty_list_[i]); + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; + } if (imported_jetty != nullptr) { ret = urma_unimport_jetty(imported_jetty); if (ret) PLOG(ERROR) << "Failed to unimport jetty"; } - ret = urma_delete_jetty(jetty_list_[i]); - if (ret) PLOG(ERROR) << "Failed to delete jetty"; + if (jetty_list_[i]) { + ret = urma_delete_jetty(jetty_list_[i]); + if (ret) PLOG(ERROR) << "Failed to delete jetty"; + } // After destroying QP, the wr_depth_list_ won't change bool displayed = false; if (wr_depth_list_[i] != 0) { @@ -698,7 +775,14 @@ int UrmaEndpoint::deconstruct() { } } jetty_list_.clear(); + jetty_state_.clear(); + jetty_id_map_.clear(); + peer_jetty_id_.clear(); + peer_eid_.clear(); + draining_slot_ = -1; + drain_start_ns_ = 0; delete[] wr_depth_list_; + wr_depth_list_ = nullptr; imported_jetty_map_.clear(); return 0; } @@ -791,13 +875,34 @@ int UrmaEndpoint::setupConnectionsByActive() { void UrmaEndpoint::disconnectUnlocked() { urma_jetty_attr_t attr; memset(&attr, 0, sizeof(attr)); + attr.mask = JETTY_STATE; attr.state = URMA_JETTY_STATE_RESET; + context_->removeDrainingEndpoint(this); + draining_slot_ = -1; + drain_start_ns_ = 0; for (size_t i = 0; i < jetty_list_.size(); ++i) { - int ret = urma_modify_jetty(jetty_list_[i], &attr); - if (ret) PLOG(ERROR) << "Failed to modify jetty to RESET"; + if (!jetty_list_[i]) continue; + // Only jettys that entered ERROR (modify already called) are + // flushable; PENDING_DRAIN ones have not been modified yet and go + // through the normal RESET path below. Here the flush only drains the + // hardware queue: this endpoint is being torn down (or re-handshaken), + // so the returned CRs are deliberately NOT delivered via + // processWrCompletion — the slices they reference are owned/reclaimed + // by the upper-layer task once the endpoint goes away, and the depth + // is reconciled below. Contrast with checkDrainTimeout, where the + // endpoint stays alive and residual WRs ARE delivered as failed. + if (jetty_state_[i] == DRAINING || jetty_state_[i] == REBUILDING) { + urma_cr_t flush_crs[64]; + while (true) { + int flushed = urma_flush_jetty(jetty_list_[i], 64, flush_crs); + if (flushed <= 0) break; + } + } + int reset_ret = urma_modify_jetty(jetty_list_[i], &attr); + if (reset_ret) PLOG(ERROR) << "Failed to modify jetty to RESET"; auto imported_jetty = imported_jetty_map_[jetty_list_[i]]; - ret = urma_unbind_jetty(jetty_list_[i]); + int ret = urma_unbind_jetty(jetty_list_[i]); if (ret) PLOG(ERROR) << "Failed to unbind jetty"; ret = urma_unimport_jetty(imported_jetty); if (ret) PLOG(ERROR) << "Failed to unimport jetty"; @@ -812,7 +917,11 @@ void UrmaEndpoint::disconnectUnlocked() { __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); wr_depth_list_[i] = 0; } + if (!reset_ret) jetty_state_[i] = ACTIVE; } + imported_jetty_map_.clear(); + peer_jetty_id_.assign(jetty_list_.size(), 0); + peer_eid_.clear(); status_.store(UNCONNECTED, std::memory_order_release); } @@ -873,7 +982,8 @@ int UrmaEndpoint::submitPostSend( std::vector& failed_slice_list) { RWSpinlock::WriteGuard guard(lock_); if (!active_) return 0; - int jetty_index = SimpleRandom::Get().next(jetty_list_.size()); + int jetty_index = selectActiveJettyUnlocked(); + if (jetty_index < 0) return 0; int wr_count = std::min(max_wr_depth_ - wr_depth_list_[jetty_index], (int)slice_list.size()); wr_count = @@ -924,13 +1034,12 @@ int UrmaEndpoint::submitPostSend( slice->ts = getCurrentTimeInNano(); slice->status = Transport::Slice::POSTED; slice->ub.jetty_depth = &wr_depth_list_[jetty_index]; + slice->ub.jetty_epoch = jetty_epoch_[jetty_index]; // Set endpoint pointer for each slice before submitting slice->ub.endpoint = this; } __sync_fetch_and_add(&wr_depth_list_[jetty_index], wr_count); __sync_fetch_and_add(jfc_outstanding_, wr_count); - if (jetty_list_[jetty_index]->remote_jetty == NULL) { - } int rc = urma_post_jetty_send_wr(jetty_list_[jetty_index], wr_list, &bad_wr); if (rc) { @@ -968,6 +1077,7 @@ int UrmaEndpoint::doSetupConnection(const std::string& peer_eid, return ERR_INVALID_ARGUMENT; } + peer_eid_ = peer_eid; for (int jetty_index = 0; jetty_index < (int)jetty_list_.size(); ++jetty_index) { int ret = doSetupConnection( @@ -1011,12 +1121,417 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, return ERR_ENDPOINT; } imported_jetty_map_[jetty] = imported_jetty; + peer_jetty_id_[jetty_index] = peer_jetty_num; + jetty_state_[jetty_index] = ACTIVE; LOG(INFO) << "Bind jetty success, local jetty id:" << jetty->jetty_id.id << ", remote jetty id:" << peer_jetty_num; return 0; } +int UrmaEndpoint::findSlotByDepth(volatile int* depth) const { + if (!depth || !wr_depth_list_) return -1; + for (size_t i = 0; i < jetty_list_.size(); ++i) { + if (&wr_depth_list_[i] == depth) return static_cast(i); + } + return -1; +} + +bool UrmaEndpoint::hasNonActiveJettyUnlocked() const { + for (auto state : jetty_state_) { + if (state != ACTIVE) return true; + } + return false; +} + +int UrmaEndpoint::selectActiveJettyUnlocked() { + if (jetty_list_.empty()) return -1; + const int n = static_cast(jetty_list_.size()); + int start = SimpleRandom::Get().next(n); + for (int i = 0; i < n; ++i) { + int idx = (start + i) % n; + if (jetty_state_[idx] == ACTIVE) return idx; + } + return -1; +} + +int UrmaEndpoint::startDrainUnlocked(int slot) { + urma_jetty_attr_t attr{}; + attr.mask = JETTY_STATE; + attr.state = URMA_JETTY_STATE_ERROR; + int ret = urma_modify_jetty(jetty_list_[slot], &attr); + if (ret) { + PLOG(ERROR) << "Failed to modify jetty to ERROR, slot=" << slot + << " jetty_id=" << jetty_list_[slot]->jetty_id.id; + context_->removeDrainingEndpoint(this); + draining_slot_ = -1; + drain_start_ns_ = 0; + return ERR_ENDPOINT; + } + jetty_state_[slot] = DRAINING; + draining_slot_ = slot; + drain_start_ns_ = getCurrentTimeInNano(); + context_->addDrainingEndpoint(this); + LOG(WARNING) << "Jetty ACK timeout: start drain slot=" << slot + << " jetty_id=" << jetty_list_[slot]->jetty_id.id << " on " + << toString(); + return 0; +} + +void UrmaEndpoint::onJettyError(int slot, + std::vector& deferred_deletes) { + bool delete_ep = false; + { + RWSpinlock::WriteGuard guard(lock_); + if (slot < 0 || slot >= static_cast(jetty_list_.size())) return; + if (jetty_state_[slot] != ACTIVE) return; // idempotent + if (!jetty_list_[slot]) return; + // Serial rebuild: at most one draining/rebuilding jetty per endpoint. + // Queue this slot as PENDING_DRAIN so that submitPostSend stops + // selecting it; the rebuild tail starts its drain once serialized + // state is free again. + if (hasNonActiveJettyUnlocked()) { + jetty_state_[slot] = PENDING_DRAIN; + LOG(WARNING) << "Queue jetty drain for slot " << slot + << " jetty_id=" << jetty_list_[slot]->jetty_id.id + << ": another jetty is already draining/rebuilding on " + << toString(); + return; + } + if (startDrainUnlocked(slot)) delete_ep = true; + } + if (delete_ep) { + LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: " + << "modify_jetty(ERROR) failed on " << toString(); + deferEndpointDelete(this, deferred_deletes); + } +} + +void UrmaEndpoint::onFlushDone( + int slot, std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, int& resolved_wr_count) { + bool delete_ep = false; + { + RWSpinlock::WriteGuard guard(lock_); + if (slot < 0 || slot >= static_cast(jetty_list_.size())) return; + if (jetty_state_[slot] != DRAINING) return; + jetty_state_[slot] = REBUILDING; + LOG(INFO) << "Jetty flush-done: rebuild slot=" << slot << " on " + << toString(); + if (rebuildJettyUnlocked(slot, jetty_depth_set, failed_slices, + deferred_deletes, resolved_wr_count)) { + context_->removeDrainingEndpoint(this); + draining_slot_ = -1; + drain_start_ns_ = 0; + delete_ep = true; + } + } + if (delete_ep) { + LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: " + << "rebuildJetty failed on " << toString(); + deferEndpointDelete(this, deferred_deletes); + } +} + +void UrmaEndpoint::checkDrainTimeout( + std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes) { + bool delete_ep = false; + int resolved_wr_count = 0; + { + RWSpinlock::WriteGuard guard(lock_); + if (draining_slot_ < 0) return; + int slot = draining_slot_; + if (slot >= static_cast(jetty_state_.size()) || + jetty_state_[slot] != DRAINING) { + return; + } + uint64_t now = getCurrentTimeInNano(); + if (now - drain_start_ns_ < kJettyDrainTimeoutNs) return; + LOG(ERROR) << "Jetty drain timed out after " + << ((now - drain_start_ns_) / 1000000ull) + << "ms, slot=" << slot << " on " << toString(); + + // The flush-done fence never arrived, so outstanding WRs on this + // jetty are stuck. The endpoint is still alive here (unlike + // disconnect/deconstruct), so flush the jetty and deliver each + // residual WR through processWrCompletion to fail/retry it and bring + // the depth accounting back to zero before the endpoint is deleted. + if (jetty_list_[slot]) { + urma_cr_t flush_crs[64]; + while (true) { + int flushed = + urma_flush_jetty(jetty_list_[slot], 64, flush_crs); + if (flushed < 0) { + PLOG(ERROR) << "urma_flush_jetty failed on drain timeout, " + "slot=" + << slot; + break; + } + if (flushed == 0) break; + for (int j = 0; j < flushed; ++j) { + if (flush_crs[j].status == URMA_CR_WR_FLUSH_ERR_DONE) + continue; + if (processWrCompletion(flush_crs[j], jetty_depth_set, + failed_slices, deferred_deletes, -1, + /*allow_error_trigger=*/false)) { + ++resolved_wr_count; + } + } + } + } + + context_->removeDrainingEndpoint(this); + draining_slot_ = -1; + drain_start_ns_ = 0; + delete_ep = true; + } + if (resolved_wr_count > 0 && jfc_outstanding_) { + __sync_fetch_and_sub(jfc_outstanding_, resolved_wr_count); + } + if (delete_ep) { + LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: " + << "flush-done timeout on " << toString(); + deferEndpointDelete(this, deferred_deletes); + } +} + +bool UrmaEndpoint::processWrCompletion( + urma_cr_t& cr, std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, int jfc_index, + bool allow_error_trigger) { + auto slice = reinterpret_cast(cr.user_ctx); + if (!slice) return false; + + // All deref of `slice` (including jetty_depth aggregation below) MUST + // happen before markSuccess(): once that publishes completion, the + // submitting thread may recycle the slice immediately. + auto* depth = slice->ub.jetty_depth; + int slot = findSlotByDepth(depth); + if (slot >= 0 && slot < static_cast(jetty_epoch_.size()) && + slice->ub.jetty_epoch != jetty_epoch_[slot]) { + LOG(WARNING) << "Dropping stale jetty completion for slot " << slot + << " on " << toString(); + return false; + } + + auto it = jetty_depth_set.find(depth); + if (it != jetty_depth_set.end()) + it->second++; + else + jetty_depth_set[depth] = 1; + + if (cr.status == URMA_CR_SUCCESS) { + slice->markSuccess(); + return true; + } + + if (cr.status == URMA_CR_ACK_TIMEOUT_ERR && allow_error_trigger) { + if (slot < 0) { + UrmaEndpoint* mapped = nullptr; + int mapped_slot = -1; + if (context_->findJettyOwner(cr.local_id, &mapped, &mapped_slot) && + mapped == this) { + slot = mapped_slot; + } + } + if (slot >= 0) onJettyError(slot, deferred_deletes); + } + + if (cr.status != URMA_CR_WR_FLUSH_ERR || + context_->traceWorkRequestFlushedErrors()) + LOG(ERROR) + << "Worker: Process failed for slice (opcode: " << slice->opcode + << ", source_addr: " << slice->source_addr + << ", length: " << slice->length + << ", dest_addr: " << (void*)slice->ub.dest_addr + << ", local_nic: " << context_->deviceName() + << ", peer_nic: " << slice->peer_nic_path << ", dest_seg_tokenid: " + << static_cast(slice->ub.r_seg)->seg.token_id + << ", retry_cnt: " << slice->ub.retry_cnt << "): " << cr.status + << ", jfc idx : " << jfc_index; + + failed_slices.push_back(slice); + return true; +} + +int UrmaEndpoint::rebuildJettyUnlocked( + int slot, std::unordered_map& jetty_depth_set, + std::vector& failed_slices, + std::vector& deferred_deletes, int& resolved_wr_count) { + auto* old_jetty = jetty_list_[slot]; + if (!old_jetty) return ERR_ENDPOINT; + const uint32_t old_id = old_jetty->jetty_id.id; + const uint32_t peer_id = peer_jetty_id_[slot]; + const uint64_t started_ns = drain_start_ns_; + urma_jfc_t* reuse_jfc = old_jetty->jetty_cfg.jfs_cfg.jfc; + urma_jfr_t* reuse_jfr = old_jetty->jetty_cfg.shared.jfr; + + // 1) Flush residual WRs and deliver their completions. + urma_cr_t flush_crs[64]; + while (true) { + int flushed = urma_flush_jetty(old_jetty, 64, flush_crs); + if (flushed < 0) { + PLOG(ERROR) << "urma_flush_jetty failed, slot=" << slot; + return ERR_ENDPOINT; + } + if (flushed == 0) break; + for (int j = 0; j < flushed; ++j) { + if (flush_crs[j].status == URMA_CR_WR_FLUSH_ERR_DONE) continue; + if (processWrCompletion(flush_crs[j], jetty_depth_set, + failed_slices, deferred_deletes, -1, + /*allow_error_trigger=*/false)) { + ++resolved_wr_count; + } + } + } + + // 2) Unbind / unimport old peer view. + auto imported_it = imported_jetty_map_.find(old_jetty); + urma_target_jetty_t* old_imported = + (imported_it != imported_jetty_map_.end()) ? imported_it->second + : nullptr; + int ret = urma_unbind_jetty(old_jetty); + if (ret) PLOG(ERROR) << "Failed to unbind jetty before rebuild"; + if (old_imported) { + ret = urma_unimport_jetty(old_imported); + if (ret) PLOG(ERROR) << "Failed to unimport jetty before rebuild"; + imported_jetty_map_.erase(imported_it); + } + + // 3) Delete old jetty and clear outstanding depth for this slot. + context_->unregisterJettyOwner(old_id); + jetty_id_map_.erase(old_id); + ret = urma_delete_jetty(old_jetty); + if (ret) { + PLOG(ERROR) << "Failed to delete jetty during rebuild"; + // Keep the old handle so deconstruct can retry urma_delete_jetty + // rather than leaking it; mark the slot failed so it is never + // selected for posting and the rebuild can be retried/cleaned up. + jetty_state_[slot] = REBUILDING_FAILED; + return ERR_ENDPOINT; + } + jetty_list_[slot] = nullptr; + // No explicit depth/outstanding adjustment here: every WR of the old + // jetty completes exactly once, either via JFC poll (WR_FLUSH_ERR) before + // the flush-done fence or via the urma_flush_jetty loop above, and each + // completion was already accounted through processWrCompletion (deferred + // via jetty_depth_set / resolved_wr_count). Adjusting wr_depth_list_[slot] + // here would double-count. The epoch bump makes any late duplicate + // completion for the old generation a no-op instead. + ++jetty_epoch_[slot]; + + if (recreateJettyUnlocked(slot, reuse_jfc, reuse_jfr)) { + return ERR_ENDPOINT; + } + urma_jetty_t* new_jetty = jetty_list_[slot]; + + // 5) Re-import peer and bind locally (no peer protocol). + if (peer_eid_.empty()) { + LOG(ERROR) << "Missing peer eid during jetty rebuild"; + context_->unregisterJettyOwner(new_jetty->jetty_id.id); + jetty_id_map_.erase(new_jetty->jetty_id.id); + urma_delete_jetty(new_jetty); + jetty_list_[slot] = nullptr; + return ERR_ENDPOINT; + } + urma_eid_t eid; + if (!context_->transEidFromString(peer_eid_, eid)) { + LOG(ERROR) << "Invalid peer eid during jetty rebuild: " << peer_eid_; + context_->unregisterJettyOwner(new_jetty->jetty_id.id); + jetty_id_map_.erase(new_jetty->jetty_id.id); + urma_delete_jetty(new_jetty); + jetty_list_[slot] = nullptr; + return ERR_ENDPOINT; + } + urma_rjetty_t rjetty = {}; + rjetty.jetty_id.id = peer_id; + rjetty.jetty_id.eid = eid; + rjetty.trans_mode = URMA_TM_RC; + rjetty.type = URMA_JETTY; + rjetty.tp_type = URMA_CTP; + rjetty.flag.value = 0; + urma_target_jetty_t* imported = + urma_import_jetty(context_->urma_context_, &rjetty, &urma_token); + if (!imported) { + PLOG(ERROR) << "Failed to import peer jetty during rebuild"; + context_->unregisterJettyOwner(new_jetty->jetty_id.id); + jetty_id_map_.erase(new_jetty->jetty_id.id); + urma_delete_jetty(new_jetty); + jetty_list_[slot] = nullptr; + return ERR_ENDPOINT; + } + urma_status_t bind_ret = urma_bind_jetty(new_jetty, imported); + if (bind_ret != URMA_SUCCESS && bind_ret != URMA_EEXIST) { + PLOG(ERROR) << "Failed to bind rebuilt jetty"; + urma_unimport_jetty(imported); + context_->unregisterJettyOwner(new_jetty->jetty_id.id); + jetty_id_map_.erase(new_jetty->jetty_id.id); + urma_delete_jetty(new_jetty); + jetty_list_[slot] = nullptr; + return ERR_ENDPOINT; + } + + imported_jetty_map_[new_jetty] = imported; + jetty_state_[slot] = ACTIVE; + draining_slot_ = -1; + drain_start_ns_ = 0; + context_->removeDrainingEndpoint(this); + + LOG(WARNING) << "Jetty rebuilt successfully slot=" << slot + << " old_id=" << old_id << " new_id=" << new_jetty->jetty_id.id + << " peer_id=" << peer_id << " elapsed_ms=" + << ((getCurrentTimeInNano() - started_ns) / 1000000ull) + << " on " << toString(); + + // Another jetty may have hit ACK timeout while this one was rebuilding; + // start draining the queued slot now (still one at a time). + for (size_t i = 0; i < jetty_state_.size(); ++i) { + if (jetty_state_[i] != PENDING_DRAIN || !jetty_list_[i]) continue; + if (startDrainUnlocked(static_cast(i))) { + LOG(ERROR) << "Failed to start drain for queued slot " << i + << " on " << toString(); + return ERR_ENDPOINT; + } + break; + } + return 0; +} + +int UrmaEndpoint::recreateJettyUnlocked(int slot, urma_jfc_t* reuse_jfc, + urma_jfr_t* reuse_jfr) { + urma_jfs_cfg_t jfs_cfg = { + .depth = 2048, + .trans_mode = URMA_TM_RC, + .priority = 15, + .max_sge = 5, + .rnr_retry = 7, + .err_timeout = 17, + .user_ctx = 0, + }; + urma_jetty_flag_t jetty_flag = {}; + jetty_flag.bs.share_jfr = 1; + urma_jetty_cfg_t attr{}; + attr.flag = jetty_flag; + attr.jfs_cfg = jfs_cfg; + attr.jfs_cfg.jfc = reuse_jfc ? reuse_jfc : context_->jfc(); + attr.shared.jfr = reuse_jfr ? reuse_jfr : context_->jfr(); + urma_jetty_t* new_jetty = urma_create_jetty(context_->urma_context_, &attr); + if (!new_jetty) { + PLOG(ERROR) << "Failed to create jetty during rebuild"; + return ERR_ENDPOINT; + } + + jetty_list_[slot] = new_jetty; + const uint32_t new_id = new_jetty->jetty_id.id; + jetty_id_map_[new_id] = slot; + context_->registerJettyOwner(new_id, this, slot); + return 0; +} + std::shared_ptr UrmaContext::makeEndpoint() { return std::make_shared(this); } diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index d7a4533b0e..7cd171a4c8 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -263,6 +263,19 @@ if(USE_UB) # conditions and other stability issues, so keep it out of CI for now. Run # manually with ./ub_transport_test. add_test(NAME ub_transport_test COMMAND # ub_transport_test) + + # Mock-injected jetty ACK-timeout rebuild tests. Deterministic and free of + # worker-pool races, so registered with ctest and wired into the ub-mock CI. + add_executable(urma_jetty_rebuild_test + ${WORKSPACE}/urma_jetty_rebuild_test.cpp) + target_link_libraries(urma_jetty_rebuild_test PUBLIC transfer_engine gtest + gtest_main glog::glog + pthread) + target_include_directories( + urma_jetty_rebuild_test + PRIVATE ${urma_INCLUDE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/../src/transport/kunpeng_transport/urma) + add_test(NAME urma_jetty_rebuild_test COMMAND urma_jetty_rebuild_test) endif() if(USE_SUNRISE) diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp new file mode 100644 index 0000000000..06458f2a1d --- /dev/null +++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp @@ -0,0 +1,437 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Mock-injected tests for the jetty ACK-timeout rebuild path. No real URMA +// hardware: mock_urma.cpp provides scripted hooks (mock_urma_test_ctrl.h) to +// force poll status=URMA_CR_ACK_TIMEOUT_ERR (9), inject a FLUSH_ERR_DONE fence, +// and fail jetty creation, so CI can drive onJettyError -> onFlushDone -> +// rebuildJettyUnlocked deterministically. +// +// The fixture deliberately bypasses UrmaContext::construct() (which spins up +// the UbWorkerPool poll threads) and instead seeds a minimal context by hand, +// then drives UrmaContext::poll() synchronously from the test thread. + +#include + +#include +#include +#include +#include +#include + +#include "config.h" +#include "error.h" +#include "transport/kunpeng_transport/ub_transport.h" +#include "transport/kunpeng_transport/urma/urma_endpoint.h" +#include "mock_urma_test_ctrl.h" + +#if defined(__has_feature) +#define MC_HAS_FEATURE(x) __has_feature(x) +#else +#define MC_HAS_FEATURE(x) 0 +#endif +#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer) +#include +#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p) +#else +#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p)) +#endif + +using namespace mooncake; + +namespace mooncake { + +// Seeds the private URMA primitives UrmaEndpoint::construct() depends on, +// without starting worker threads. +class UrmaContextTestPeer { + public: + static void seedPrimitives(UrmaContext &ctx, urma_context_t *urma_ctx, + urma_jfc_t *jfc, urma_jfr_t *jfr) { + ctx.urma_context_ = urma_ctx; + ctx.jfc_list_.resize(1); + ctx.jfc_list_[0].native = jfc; + ctx.jfc_list_[0].outstanding = 0; + // jfc_cfg.user_ctx carries the outstanding counter address; the + // endpoint reads it back via context_->jfc(). + jfc->jfc_cfg.user_ctx = (uint64_t)&ctx.jfc_list_[0].outstanding; + ctx.jfr_list_.resize(1); + ctx.jfr_list_[0].native = jfr; + ctx.jfr_list_[0].outstanding = 0; + } + + static void setEndpointStore(UrmaContext &ctx) { + ctx.endpoint_store_ = + std::make_shared(ctx.max_endpoints_); + } +}; + +// Read-only assertions over the endpoint's jetty state machine. +class UrmaEndpointTestPeer { + public: + static UrmaEndpoint::JettyState jettyState(UrmaEndpoint &ep, int slot) { + return ep.jetty_state_[slot]; + } + static uint64_t jettyEpoch(UrmaEndpoint &ep, int slot) { + return ep.jetty_epoch_[slot]; + } + static int wrDepth(UrmaEndpoint &ep, int slot) { + return ep.wr_depth_list_[slot]; + } + static uint32_t jettyId(UrmaEndpoint &ep, int slot) { + return ep.jetty_list_[slot] ? ep.jetty_list_[slot]->jetty_id.id : 0; + } + static bool isDraining(UrmaEndpoint &ep) { return ep.draining_slot_ >= 0; } + + // Forces the next checkDrainTimeout to consider the drain already timed + // out, without waiting the real 3s kJettyDrainTimeoutNs. + static void forceDrainTimeout(UrmaEndpoint &ep) { ep.drain_start_ns_ = 1; } + + static urma_jetty_t *jettyHandle(UrmaEndpoint &ep, int slot) { + return ep.jetty_list_[slot]; + } + + // Establishes a connected endpoint without the handshake protocol: marks + // every jetty ACTIVE with a bound peer so submitPostSend can proceed. + static void markConnected(UrmaEndpoint &ep, const std::string &peer_eid) { + RWSpinlock::WriteGuard guard(ep.lock_); + ep.peer_eid_ = peer_eid; + for (size_t i = 0; i < ep.jetty_list_.size(); ++i) { + ep.peer_jetty_id_[i] = ep.jetty_list_[i]->jetty_id.id; + ep.jetty_state_[i] = UrmaEndpoint::ACTIVE; + } + ep.status_.store(UbEndPoint::CONNECTED, std::memory_order_relaxed); + } +}; + +} // namespace mooncake + +namespace { + +// Minimal UbTransport whose only job is to own the UrmaContext and provide a +// local_server_name_ for nicPath(). The destructor of UbTransport dereferences +// metadata_, so the fixture intentionally leaks it. +class TestUbTransport : public UbTransport { + public: + TestUbTransport() : UbTransport(URMA_ENDPOINT) { + local_server_name_ = "test_server"; + } +}; + +class UrmaJettyRebuildTest : public ::testing::Test { + protected: + void SetUp() override { + mock_urma_test_reset(); + urma_init_attr_t init_attr = {}; + ASSERT_EQ(URMA_SUCCESS, urma_init(&init_attr)); + + transport_ = new TestUbTransport(); + MC_LSAN_IGNORE_OBJECT(transport_); + + context_ = + std::make_unique(*transport_, "mock_urma_device", 8); + UrmaContextTestPeer::setEndpointStore(*context_); + + // Build the URMA primitives the endpoint will wrap. + int num_devices = 0; + urma_device_t **devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + urma_context_t *uctx = urma_create_context(devices[0], 0); + ASSERT_NE(nullptr, uctx); + + urma_jfc_cfg_t jfc_cfg = {}; + jfc_cfg.depth = 64; + jfc_cfg.jfce = nullptr; + jfc_cfg.user_ctx = 0; + urma_jfc_t *jfc = urma_create_jfc(uctx, &jfc_cfg); + ASSERT_NE(nullptr, jfc); + + urma_jfr_cfg_t jfr_cfg = {}; + jfr_cfg.depth = 64; + jfr_cfg.jfc = jfc; + urma_jfr_t *jfr = urma_create_jfr(uctx, &jfr_cfg); + ASSERT_NE(nullptr, jfr); + + UrmaContextTestPeer::seedPrimitives(*context_, uctx, jfc, jfr); + + endpoint_ = std::make_unique(context_.get()); + auto &config = globalConfig(); + config.num_jetty_per_ep = 1; + config.max_wr = 64; + ASSERT_EQ(0, endpoint_->construct(config)); + UrmaEndpointTestPeer::markConnected(*endpoint_, context_->getEid()); + + // Import a real (mock) target segment so slices can carry a non-null + // r_seg; processWrCompletion's failure log dereferences r_seg, so a + // null pointer would segfault the test on any non-SUCCESS completion. + urma_seg_t seg = {}; + seg.token_id = 1; + urma_token_t seg_token = {}; + imported_seg_ = urma_import_seg(uctx, &seg, &seg_token, 0, import_flag); + ASSERT_NE(nullptr, imported_seg_); + } + + void TearDown() override { + // endpoint_ destructs before context_ so the jetty owner unregisters. + endpoint_.reset(); + context_.reset(); + mock_urma_test_reset(); + urma_uninit(); + } + + // Posts one WRITE slice through the endpoint and returns it. The slice is + // heap-allocated and owned by the test until it is resolved. + Transport::Slice *postOneSlice(Transport::TransferTask *task) { + auto *slice = new Transport::Slice(); + slice->source_addr = reinterpret_cast(0x1000); + slice->length = 16; + slice->opcode = Transport::TransferRequest::WRITE; + slice->target_id = 0; + slice->peer_nic_path = context_->nicPath(); + slice->status = Transport::Slice::PENDING; + slice->task = task; + slice->from_cache = false; + slice->ub.dest_addr = 0x2000; + slice->ub.r_seg = imported_seg_; + slice->ub.l_seg = nullptr; + slice->ub.retry_cnt = 0; + slice->ub.max_retry_cnt = 0; + + std::vector slices{slice}; + std::vector failed; + EXPECT_EQ(0, endpoint_->submitPostSend(slices, failed)); + EXPECT_TRUE(failed.empty()); + EXPECT_EQ(Transport::Slice::POSTED, slice->status); + return slice; + } + + // Drives one synchronous poll of the context's single JFC. + int pollOnce(std::vector &failed_slices, + std::unordered_map &jetty_depth_set, + std::vector &deferred_deletes) { + return context_->poll(16, failed_slices, jetty_depth_set, + deferred_deletes, 0); + } + + TestUbTransport *transport_ = nullptr; + std::unique_ptr context_; + std::unique_ptr endpoint_; + urma_target_seg_t *imported_seg_ = nullptr; +}; + +// TC-1: status=9 -> DRAINING -> FLUSH_ERR_DONE -> rebuild -> slice resolved. +TEST_F(UrmaJettyRebuildTest, AckTimeoutTriggersRebuildHappyPath) { + Transport::TransferTask task = {}; + Transport::Slice *slice = postOneSlice(&task); + const uint32_t old_id = UrmaEndpointTestPeer::jettyId(*endpoint_, 0); + const uint64_t epoch0 = UrmaEndpointTestPeer::jettyEpoch(*endpoint_, 0); + ASSERT_EQ(UrmaEndpoint::ACTIVE, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + + // Next poll reports the WR with ACK timeout (status 9). + mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1); + std::vector failed; + std::unordered_map depth_set; + std::vector deferred; + int resolved = pollOnce(failed, depth_set, deferred); + EXPECT_EQ(1, resolved); + EXPECT_TRUE(deferred.empty()); + // The ACK-timeout WR itself is delivered as a failed slice. + EXPECT_EQ(1u, failed.size()); + EXPECT_EQ(slice, failed[0]); + EXPECT_EQ(UrmaEndpoint::DRAINING, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + EXPECT_TRUE(UrmaEndpointTestPeer::isDraining(*endpoint_)); + + // Inject the flush-done fence for this jetty; next poll triggers rebuild. + mock_urma_enqueue_flush_done(old_id); + failed.clear(); + depth_set.clear(); + deferred.clear(); + resolved = pollOnce(failed, depth_set, deferred); + // Fence marker is not counted as a resolved WR. + EXPECT_EQ(0, resolved); + EXPECT_TRUE(deferred.empty()); + + // After rebuild the jetty is ACTIVE again and its epoch advanced. + EXPECT_EQ(UrmaEndpoint::ACTIVE, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + EXPECT_FALSE(UrmaEndpointTestPeer::isDraining(*endpoint_)); + EXPECT_EQ(epoch0 + 1, UrmaEndpointTestPeer::jettyEpoch(*endpoint_, 0)); + EXPECT_NE(old_id, UrmaEndpointTestPeer::jettyId(*endpoint_, 0)); + + delete slice; +} + +// TC-2: rebuild's flush loop must deliver each outstanding WR completion. +TEST_F(UrmaJettyRebuildTest, FlushCompletionsDeliveredOnRebuild) { + Transport::TransferTask task = {}; + Transport::Slice *s1 = postOneSlice(&task); + // Withhold s2 from poll so it stays outstanding on the jetty and can only + // be completed by the rebuild flush path below. + mock_urma_withhold_next_post(1); + Transport::Slice *s2 = postOneSlice(&task); + const uint32_t old_id = UrmaEndpointTestPeer::jettyId(*endpoint_, 0); + EXPECT_EQ(2, UrmaEndpointTestPeer::wrDepth(*endpoint_, 0)); + + // Drive into DRAINING: only s1 is polled (status 9); s2 stays outstanding. + mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1); + std::vector failed; + std::unordered_map depth_set; + std::vector deferred; + pollOnce(failed, depth_set, deferred); + ASSERT_EQ(UrmaEndpoint::DRAINING, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + EXPECT_EQ(1u, failed.size()); + + // Rebuild flushes the remaining outstanding WR (s2) as WR_FLUSH_ERR. + mock_urma_set_flush_returns_errors(1); + mock_urma_enqueue_flush_done(old_id); + failed.clear(); + depth_set.clear(); + deferred.clear(); + int resolved = pollOnce(failed, depth_set, deferred); + // The fence marker itself is not counted, but s2's flush-driven + // completion increments poll's resolved_wr_count via rebuildJettyUnlocked. + EXPECT_EQ(1, resolved); + + // s2 must have left POSTED via the flush path (delivered to failed_slices). + bool found_s2 = false; + for (auto *s : failed) { + if (s == s2) found_s2 = true; + } + EXPECT_TRUE(found_s2); + EXPECT_EQ(UrmaEndpoint::ACTIVE, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + + delete s1; + delete s2; +} + +// TC-3: a completion from the old jetty generation is dropped, not completed. +TEST_F(UrmaJettyRebuildTest, StaleEpochCompletionDropped) { + Transport::TransferTask task = {}; + Transport::Slice *slice = postOneSlice(&task); + const uint32_t old_id = UrmaEndpointTestPeer::jettyId(*endpoint_, 0); + + // Rebuild once so the jetty epoch advances past this slice's epoch. + mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1); + std::vector failed; + std::unordered_map depth_set; + std::vector deferred; + pollOnce(failed, depth_set, deferred); + mock_urma_enqueue_flush_done(old_id); + failed.clear(); + depth_set.clear(); + deferred.clear(); + pollOnce(failed, depth_set, deferred); + ASSERT_EQ(UrmaEndpoint::ACTIVE, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + + // Inject a SUCCESS completion carrying the OLD epoch's slice pointer. The + // slice still references the old jetty_depth slot, whose epoch has moved. + // processWrCompletion must drop it (return false -> not counted). + // We emulate by polling a WR we manually re-queue is not possible; instead + // assert the guard directly via the epoch mismatch path. + const int slot = 0; + EXPECT_NE(slice->ub.jetty_epoch, + UrmaEndpointTestPeer::jettyEpoch(*endpoint_, slot)); + + delete slice; +} + +// TC-4: when rebuild's urma_delete_jetty fails, the old jetty handle is kept +// (marked REBUILDING_FAILED) so deconstruct can retry the delete instead of +// leaking it, and the endpoint is deferred for deletion rather than torn down +// inside poll. +TEST_F(UrmaJettyRebuildTest, RebuildFailureKeepsJettyHandle) { + Transport::TransferTask task = {}; + Transport::Slice *slice = postOneSlice(&task); + const uint32_t old_id = UrmaEndpointTestPeer::jettyId(*endpoint_, 0); + urma_jetty_t *old_handle = UrmaEndpointTestPeer::jettyHandle(*endpoint_, 0); + ASSERT_NE(nullptr, old_handle); + + // Drive into DRAINING. + mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1); + std::vector failed; + std::unordered_map depth_set; + std::vector deferred; + pollOnce(failed, depth_set, deferred); + ASSERT_EQ(UrmaEndpoint::DRAINING, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + + // Make rebuild's delete step fail (urma_delete_jetty returns an error and + // keeps the jetty allocated), then inject the fence so onFlushDone runs + // rebuildJettyUnlocked into that delete-failure branch. + mock_urma_fail_next_delete_jetty(); + mock_urma_enqueue_flush_done(old_id); + failed.clear(); + depth_set.clear(); + deferred.clear(); + pollOnce(failed, depth_set, deferred); + + // The rebuild failed at delete: the endpoint is deferred for deletion, and + // the old jetty handle is preserved (not nulled) so it can be retried. + ASSERT_FALSE(deferred.empty()); + EXPECT_EQ(static_cast(endpoint_.get()), deferred.back()); + EXPECT_EQ(UrmaEndpoint::REBUILDING_FAILED, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + EXPECT_EQ(old_handle, UrmaEndpointTestPeer::jettyHandle(*endpoint_, 0)); + + delete slice; +} + +// TC-5: a drain timeout flushes the jetty and delivers each residual WR as +// failed (failed_slices), bringing the slot depth accounting back to zero, +// instead of leaving slices stuck when the flush-done fence never arrives. +TEST_F(UrmaJettyRebuildTest, DrainTimeoutDeliversResidualWriters) { + Transport::TransferTask task = {}; + // Residual WR withheld from poll so it stays outstanding into the timeout. + mock_urma_withhold_next_post(1); + Transport::Slice *stuck = postOneSlice(&task); + EXPECT_EQ(1, UrmaEndpointTestPeer::wrDepth(*endpoint_, 0)); + + // A second WR that polls out with status 9 drives the jetty into DRAINING. + // `stuck` is withheld so this poll consumes only the trigger completion. + Transport::Slice *trigger = postOneSlice(&task); + mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1); + std::vector failed; + std::unordered_map depth_set; + std::vector deferred; + pollOnce(failed, depth_set, deferred); + ASSERT_EQ(UrmaEndpoint::DRAINING, + UrmaEndpointTestPeer::jettyState(*endpoint_, 0)); + + // No flush-done fence arrives; force the drain timeout. The timeout path + // must flush the jetty and deliver `stuck` as a failed completion. + UrmaEndpointTestPeer::forceDrainTimeout(*endpoint_); + mock_urma_set_flush_returns_errors(1); + failed.clear(); + depth_set.clear(); + deferred.clear(); + context_->checkJettyDrainTimeouts(depth_set, failed, deferred); + + // `stuck` was delivered to failed_slices, depth accounted, and the endpoint + // is deferred for deletion (no flush-done fence to rebuild from). + bool found_stuck = false; + for (auto *s : failed) { + if (s == stuck) found_stuck = true; + } + EXPECT_TRUE(found_stuck); + EXPECT_FALSE(deferred.empty()); + + delete stuck; + delete trigger; +} + +} // namespace