From b8538e6ee2bff41819d7c044e6f9d287c5503222 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Wed, 26 Aug 2026 17:58:06 +0800
Subject: [PATCH 01/10] [TransferEngine] Rebuild single Jetty on ACK timeout
 (status=9)

Avoid deleting the whole endpoint after a transient ACK timeout by draining
the faulty Jetty, recreating it locally, and rebinding to the existing peer
id. Falls back to deleteEndpoint on flush-done timeout or rebuild failure.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../jetty-ack-timeout-rebuild.md              |  15 +
 .../jetty-single-rebuild-plan.md              | 211 ++++++++++
 .../kunpeng_transport/urma/urma_endpoint.h    |  44 ++
 .../kunpeng_transport/urma/mock_urma.cpp      |  13 +-
 .../kunpeng_transport/urma/urma_endpoint.cpp  | 375 +++++++++++++++++-
 5 files changed, 647 insertions(+), 11 deletions(-)
 create mode 100644 docs/source/design/transfer-engine/jetty-ack-timeout-rebuild.md
 create mode 100644 docs/source/design/transfer-engine/jetty-single-rebuild-plan.md

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-single-rebuild-plan.md b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
new file mode 100644
index 0000000000..7fd20d67f8
--- /dev/null
+++ b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
@@ -0,0 +1,211 @@
+# 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 + 更新槽位 ──┘
+                     │
+                     └── 失败 / 超时 ──► 回退：deleteEndpointByPtr
+```
+
+- `ACTIVE`：正常收发
+- `DRAINING`：已 `modify(ERROR)`，禁止 post，等待 `FLUSH_ERR_DONE`
+- `REBUILDING`：排空完成，正在 flush / 删建 / rebind
+- 同 EP 任意时刻最多一个 jetty 处于 DRAINING/REBUILDING（串行）
+
+---
+
+## 4. 关键设计
+
+### 4.1 per-jetty 状态
+
+```cpp
+enum JettyState { ACTIVE, DRAINING, REBUILDING };
+std::vector<JettyState> jetty_state_;             // 与 jetty_list_ 平行
+std::unordered_map<uint32_t, int> jetty_id_map_;  // jetty_id → slot index
+std::vector<uint32_t> peer_jetty_id_;             // 每槽对端 id，delete 前保留
+```
+
+`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 则不再 `modify`）。
+
+```
+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 入队
+- 段 C 替换前若 depth 仍非 0：打日志并在持锁下归零该槽与对应 JFC 计数（与今日
+  `deconstruct` 对 outstanding 的处理同思路），避免泄漏
+
+### 4.7 段 C 完整顺序（持锁）
+
+```
+1. urma_flush_jetty（回收残余 WR；注意与 poll 可能重复，见 §7）
+2. unbind → unimport（旧本地 jetty 上的对端视图）
+3. delete_jetty；从 jetty_id_map_ 删旧 id
+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
+```
+
+任一步失败 → 回退 `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 行为标为硬件验证项
+
+---
+
+## 6. 改动文件
+
+| 文件 | 改动 |
+|---|---|
+| `urma_endpoint.h` | `JettyState`、`jetty_state_`、`jetty_id_map_`、`peer_jetty_id_`、`onJettyError()` / `onFlushDone()` |
+| `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。
+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/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
index 6b6fa7076a..d98b50993b 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 <memory>
 #include <string>
 #include <thread>
+#include <unordered_map>
+#include <unordered_set>
 #include <utility>
+#include <vector>
 #include "common.h"
 #include "config.h"
 #include "urma_api.h"
@@ -47,6 +50,8 @@ static urma_import_seg_flag_t import_flag = {
            .reserved = 0}};
 
 // define the UrmaContext class
+class UrmaEndpoint;
+
 class UrmaContext : public UbContext {
     friend class UrmaEndpoint;
 
@@ -81,6 +86,15 @@ 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();
+
    private:
     int construct(GlobalConfig& config) override;
     int deconstruct() override;
@@ -146,11 +160,21 @@ class UrmaContext : public UbContext {
 
     urma_import_seg_flag_t import_flag_ = mooncake::import_flag;
     std::unordered_map<std::string, urma_target_seg_t*> import_tseg_map;
+
+    RWSpinlock jetty_owner_lock_;
+    struct JettyOwner {
+        UrmaEndpoint* endpoint = nullptr;
+        int slot = -1;
+    };
+    std::unordered_map<uint32_t, JettyOwner> jetty_owner_map_;
+    std::unordered_set<UrmaEndpoint*> draining_endpoints_;
 };
 
 // define the UrmaEndpoint class
 class UrmaEndpoint : public UbEndPoint {
    public:
+    enum JettyState { ACTIVE = 0, DRAINING = 1, REBUILDING = 2 };
+
     UrmaEndpoint(UrmaContext* context)
         : context_(context), jfc_outstanding_(nullptr) {}
 
@@ -173,6 +197,13 @@ 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);
+    void onFlushDone(int slot);
+    void checkDrainTimeout();
+
+    int findSlotByDepth(volatile int* depth) const;
+
    private:
     void disconnectUnlocked() override;
 
@@ -187,6 +218,12 @@ class UrmaEndpoint : public UbEndPoint {
                           uint32_t peer_jetty_num,
                           std::string* reply_msg = nullptr);
 
+    bool hasNonActiveJettyUnlocked() const;
+    int selectActiveJettyUnlocked();
+    int rebuildJettyUnlocked(int slot);
+
+    static constexpr uint64_t kJettyDrainTimeoutNs = 3000000000ull;  // 3s
+
    private:
     UrmaContext* context_;
     urma_token_t urma_token = {.token = 0xACFE};
@@ -195,6 +232,13 @@ class UrmaEndpoint : public UbEndPoint {
     int max_wr_depth_;
     volatile int* jfc_outstanding_;
     std::unordered_map<urma_jetty_t*, urma_target_jetty_t*> imported_jetty_map_;
+
+    std::vector<JettyState> jetty_state_;
+    std::unordered_map<uint32_t, int> jetty_id_map_;
+    std::vector<uint32_t> peer_jetty_id_;
+    std::string peer_eid_;
+    uint64_t drain_start_ns_ = 0;
+    int draining_slot_ = -1;
 };
 }  // namespace mooncake
 #endif  // URMA_ENDPOINT_H
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..a05e04280c 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
@@ -346,11 +346,12 @@ urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) {
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
     }
+    static std::atomic<uint32_t> 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;
@@ -420,6 +421,16 @@ 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) {
+    (void)cr_cnt;
+    (void)cr;
+    std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
+    if (!jetty || jetty_map.find(jetty) == jetty_map.end()) {
+        return -1;
+    }
+    return 0;
+}
+
 urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
                                       urma_jfs_wr_t **bad_wr) {
     {
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..7a9cda3cf6 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,55 @@ bool UrmaContext::transEidFromString(const std::string& eid_str,
     return index == URMA_EID_SIZE;
 }
 
+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::vector<UrmaEndpoint*> endpoints;
+    {
+        RWSpinlock::ReadGuard guard(jetty_owner_lock_);
+        endpoints.assign(draining_endpoints_.begin(),
+                         draining_endpoints_.end());
+    }
+    for (auto* endpoint : endpoints) {
+        if (endpoint) endpoint->checkDrainTimeout();
+    }
+}
+
 int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
                       int& num_failed,
                       std::unordered_map<volatile int*, int>& jetty_depth_set,
                       int jfc_index) {
     num_failed = 0;
+    checkJettyDrainTimeouts();
     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,11 +574,26 @@ 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);
+            } 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;
         }
+        ++wr_completions;
 
         // All deref of `slice` (including the jetty_depth aggregation below)
         // MUST happen before markSuccess(): once that publishes completion,
@@ -553,6 +612,22 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
             continue;
         }
 
+        if (cr[i].status == URMA_CR_ACK_TIMEOUT_ERR) {
+            auto* endpoint = static_cast<UrmaEndpoint*>(slice->ub.endpoint);
+            if (endpoint) {
+                int slot = endpoint->findSlotByDepth(depth);
+                if (slot < 0) {
+                    UrmaEndpoint* mapped = nullptr;
+                    int mapped_slot = -1;
+                    if (findJettyOwner(cr[i].local_id, &mapped, &mapped_slot) &&
+                        mapped == endpoint) {
+                        slot = mapped_slot;
+                    }
+                }
+                if (slot >= 0) endpoint->onJettyError(slot);
+            }
+        }
+
         if (cr[i].status != URMA_CR_WR_FLUSH_ERR ||
             show_work_request_flushed_error_)
             LOG(ERROR) << "Worker: Process failed for slice (opcode: "
@@ -576,7 +651,8 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
         // safely deref it.
         failed_slices[num_failed++] = slice;
     }
-    return nr_poll;
+    // Exclude FLUSH_ERR_DONE from outstanding accounting (it was never posted).
+    return wr_completions;
 }
 
 volatile int* UrmaContext::outstandingCount(int jfc_index) {
@@ -626,6 +702,12 @@ int UrmaEndpoint::construct(GlobalConfig& config) {
     }
 
     jetty_list_.resize(num_jetty_list);
+    jetty_state_.assign(num_jetty_list, ACTIVE);
+    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;
 
@@ -660,8 +742,11 @@ int UrmaEndpoint::construct(GlobalConfig& config) {
             PLOG(ERROR) << "Failed to create jetty";
             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<int>(i);
+        context_->registerJettyOwner(jetty_id, this, static_cast<int>(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 +757,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 +789,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,8 +889,12 @@ 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";
@@ -812,7 +914,11 @@ void UrmaEndpoint::disconnectUnlocked() {
             __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]);
             wr_depth_list_[i] = 0;
         }
+        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 +979,8 @@ int UrmaEndpoint::submitPostSend(
     std::vector<Transport::Slice*>& 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 =
@@ -929,8 +1036,6 @@ int UrmaEndpoint::submitPostSend(
     }
     __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 +1073,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 +1117,261 @@ 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<int>(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<int>(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;
+}
+
+void UrmaEndpoint::onJettyError(int slot) {
+    bool delete_ep = false;
+    {
+        RWSpinlock::WriteGuard guard(lock_);
+        if (slot < 0 || slot >= static_cast<int>(jetty_list_.size())) return;
+        if (jetty_state_[slot] == DRAINING || jetty_state_[slot] == REBUILDING) {
+            return;  // idempotent
+        }
+        // Serial rebuild: at most one non-ACTIVE jetty per endpoint.
+        if (hasNonActiveJettyUnlocked()) {
+            LOG(INFO) << "Skip jetty rebuild for slot " << slot
+                      << ": another jetty is already draining/rebuilding on "
+                      << toString();
+            return;
+        }
+        if (!jetty_list_[slot]) return;
+
+        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;
+            delete_ep = true;
+        } else {
+            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();
+        }
+    }
+    if (delete_ep) {
+        LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: "
+                   << "modify_jetty(ERROR) failed on " << toString();
+        context_->deleteEndpointByPtr(this);
+    }
+}
+
+void UrmaEndpoint::onFlushDone(int slot) {
+    bool delete_ep = false;
+    {
+        RWSpinlock::WriteGuard guard(lock_);
+        if (slot < 0 || slot >= static_cast<int>(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)) {
+            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();
+        context_->deleteEndpointByPtr(this);
+    }
+}
+
+void UrmaEndpoint::checkDrainTimeout() {
+    bool delete_ep = false;
+    {
+        RWSpinlock::WriteGuard guard(lock_);
+        if (draining_slot_ < 0) return;
+        int slot = draining_slot_;
+        if (slot >= static_cast<int>(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();
+        context_->removeDrainingEndpoint(this);
+        draining_slot_ = -1;
+        drain_start_ns_ = 0;
+        delete_ep = true;
+    }
+    if (delete_ep) {
+        LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: "
+                   << "flush-done timeout on " << toString();
+        context_->deleteEndpointByPtr(this);
+    }
+}
+
+int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
+    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 (may overlap with already-polled CRs).
+    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;
+        // Completions for these WRs should already have been (or will be)
+        // accounted via poll; do not touch slice pointers from flush CRs.
+    }
+
+    // 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";
+        jetty_list_[slot] = nullptr;
+        return ERR_ENDPOINT;
+    }
+    jetty_list_[slot] = nullptr;
+    if (wr_depth_list_[slot] != 0) {
+        __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[slot]);
+        wr_depth_list_[slot] = 0;
+    }
+
+    // 4) Create replacement jetty with the same JFC/JFR config.
+    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;
+    }
+
+    // 5) Re-import peer and bind locally (no peer protocol).
+    if (peer_eid_.empty()) {
+        LOG(ERROR) << "Missing peer eid during jetty rebuild";
+        urma_delete_jetty(new_jetty);
+        return ERR_ENDPOINT;
+    }
+    urma_eid_t eid;
+    if (!context_->transEidFromString(peer_eid_, eid)) {
+        LOG(ERROR) << "Invalid peer eid during jetty rebuild: " << peer_eid_;
+        urma_delete_jetty(new_jetty);
+        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";
+        urma_delete_jetty(new_jetty);
+        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);
+        urma_delete_jetty(new_jetty);
+        return ERR_ENDPOINT;
+    }
+
+    jetty_list_[slot] = new_jetty;
+    imported_jetty_map_[new_jetty] = imported;
+    const uint32_t new_id = new_jetty->jetty_id.id;
+    jetty_id_map_[new_id] = slot;
+    context_->registerJettyOwner(new_id, this, slot);
+    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_id
+                 << " peer_id=" << peer_id << " elapsed_ms="
+                 << ((getCurrentTimeInNano() - started_ns) / 1000000ull)
+                 << " on " << toString();
+    return 0;
+}
+
 std::shared_ptr<UbEndPoint> UrmaContext::makeEndpoint() {
     return std::make_shared<UrmaEndpoint>(this);
 }

From 115191f24b8c1f307209daff93b5b3a55294ca4e Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Wed, 26 Aug 2026 19:20:53 +0800
Subject: [PATCH 02/10] [TransferEngine] Fix CI format and docs orphan metadata

Mark the jetty rebuild design doc as Sphinx orphan and apply clang-format
to the UrmaEndpoint changes so PR checks pass.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../design/transfer-engine/jetty-single-rebuild-plan.md     | 4 ++++
 .../transport/kunpeng_transport/urma/urma_endpoint.h        | 6 +++---
 .../src/transport/kunpeng_transport/urma/urma_endpoint.cpp  | 6 +++---
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
index 7fd20d67f8..85ec8b0e0b 100644
--- a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
+++ b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
@@ -1,3 +1,7 @@
+---
+orphan: true
+---
+
 # Jetty ACK Timeout（status=9）单 Jetty 重建方案
 
 状态：实现中（分支 `feat/jetty-ack-timeout-single-rebuild`）— **权威方案**（取代已废弃的 `jetty-ack-timeout-rebuild.md`）
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 d98b50993b..af456e88e0 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
@@ -89,8 +89,7 @@ class UrmaContext : public UbContext {
     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);
+    bool findJettyOwner(uint32_t jetty_id, UrmaEndpoint** endpoint, int* slot);
     void addDrainingEndpoint(UrmaEndpoint* endpoint);
     void removeDrainingEndpoint(UrmaEndpoint* endpoint);
     void checkJettyDrainTimeouts();
@@ -197,7 +196,8 @@ class UrmaEndpoint : public UbEndPoint {
 
     const std::string toString() const override;
 
-    // Called from UrmaContext::poll on ACK timeout / flush-done / drain timeout.
+    // Called from UrmaContext::poll on ACK timeout / flush-done / drain
+    // timeout.
     void onJettyError(int slot);
     void onFlushDone(int slot);
     void checkDrainTimeout();
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 7a9cda3cf6..14a317f3d9 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
@@ -1156,7 +1156,8 @@ void UrmaEndpoint::onJettyError(int slot) {
     {
         RWSpinlock::WriteGuard guard(lock_);
         if (slot < 0 || slot >= static_cast<int>(jetty_list_.size())) return;
-        if (jetty_state_[slot] == DRAINING || jetty_state_[slot] == REBUILDING) {
+        if (jetty_state_[slot] == DRAINING ||
+            jetty_state_[slot] == REBUILDING) {
             return;  // idempotent
         }
         // Serial rebuild: at most one non-ACTIVE jetty per endpoint.
@@ -1313,8 +1314,7 @@ int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
     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);
+    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;

From 60e947da6cdc52dfd09192e5d587ae295c750214 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 00:41:18 +0800
Subject: [PATCH 03/10] [TransferEngine] Fix jetty rebuild safety issues on ACK
 timeout path

Defer endpoint deletion until after poll depth accounting to avoid UAF,
deliver flush completions during rebuild, and isolate stale CQEs with
per-slot jetty epochs. Align disconnect teardown with ERROR flush fence
constraints and document the updated rebuild accounting model.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../jetty-single-rebuild-plan.md              |  48 ++-
 .../transport/kunpeng_transport/ub_context.h  |  22 +-
 .../kunpeng_transport/urma/urma_endpoint.h    |  62 ++-
 .../include/transport/transport.h             |   1 +
 .../kunpeng_transport/ub_context.cpp          |  28 +-
 .../kunpeng_transport/urma/urma_endpoint.cpp  | 395 ++++++++++++------
 6 files changed, 382 insertions(+), 174 deletions(-)

diff --git a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
index 85ec8b0e0b..956cb201ce 100644
--- a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
+++ b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
@@ -60,16 +60,21 @@ kunpeng 数据面是单边 `URMA_OPC_READ/WRITE`，DMA 目标是对端 **segment
 每个 jetty 槽位有独立状态：
 
 ```
-ACTIVE ──(status=9)──► DRAINING ──(FLUSH_ERR_DONE)──► REBUILDING
-                                                          │
-                     ◄── create + 本端 rebind + 更新槽位 ──┘
-                     │
-                     └── 失败 / 超时 ──► 回退：deleteEndpointByPtr
+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（串行）
 
 ---
@@ -79,10 +84,11 @@ ACTIVE ──(status=9)──► DRAINING ──(FLUSH_ERR_DONE)──► REBUIL
 ### 4.1 per-jetty 状态
 
 ```cpp
-enum JettyState { ACTIVE, DRAINING, REBUILDING };
+enum JettyState { ACTIVE, DRAINING, REBUILDING, PENDING_DRAIN };
 std::vector<JettyState> jetty_state_;             // 与 jetty_list_ 平行
 std::unordered_map<uint32_t, int> jetty_id_map_;  // jetty_id → slot index
 std::vector<uint32_t> peer_jetty_id_;             // 每槽对端 id，delete 前保留
+std::vector<uint64_t> jetty_epoch_;               // 每槽重建代次，丢弃旧代次的迟到 CR
 ```
 
 `peer_jetty_id_` 在握手 `doSetupConnection` 时写入；段 C 在 delete 本地 jetty
@@ -119,7 +125,9 @@ if cr.status == FLUSH_ERR_DONE:
 ### 4.5 status=9 分流与触发时机
 
 **首次** poll 到 `ACK_TIMEOUT_ERR (9)` 即触发排空（不等重试耗尽）；`onJettyError`
-必须幂等（已在 DRAINING/REBUILDING 则不再 `modify`）。
+必须幂等（已在 DRAINING/REBUILDING/PENDING_DRAIN 则不再 `modify`）。
+若同 EP 已有 jetty 在重建，则把本槽标记为 `PENDING_DRAIN` 排队：选槽立即跳过
+它，等前一个重建完成后由其尾部调用 `startDrainUnlocked` 串行启动排空。
 
 ```
 if cr.status == ACK_TIMEOUT_ERR (9):
@@ -136,19 +144,29 @@ if cr.status == ACK_TIMEOUT_ERR (9):
 - 带有效 `user_ctx` 的失败 CQE：仍走现有 `jetty_depth_set` / retry 路径扣
   `wr_depth_list_` 与 JFC outstanding
 - `FLUSH_ERR_DONE`：不当作 slice；不得 `markSuccess` / 不得当失败 slice 入队
-- 段 C 替换前若 depth 仍非 0：打日志并在持锁下归零该槽与对应 JFC 计数（与今日
-  `deconstruct` 对 outstanding 的处理同思路），避免泄漏
+- **每个 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；注意与 poll 可能重复，见 §7）
+1. urma_flush_jetty（回收残余 WR，逐条经 processWrCompletion 交付；
+   注意与 poll 可能重复，见 §7）
 2. unbind → unimport（旧本地 jetty 上的对端视图）
-3. delete_jetty；从 jetty_id_map_ 删旧 id
+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`。
@@ -200,7 +218,7 @@ if cr.status == ACK_TIMEOUT_ERR (9):
 
 | 文件 | 改动 |
 |---|---|
-| `urma_endpoint.h` | `JettyState`、`jetty_state_`、`jetty_id_map_`、`peer_jetty_id_`、`onJettyError()` / `onFlushDone()` |
+| `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`） |
 
@@ -208,7 +226,11 @@ if cr.status == ACK_TIMEOUT_ERR (9):
 
 ## 7. 风险与开放问题
 
-1. **urma_flush_jetty 与 poll 重复**：flush 返回的 WR 级 CR 是否已在 JFC poll 中出现过，需实测确认，避免 double-complete。
+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 降级路径。
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<Transport::Slice*>& failed_slices,
                      std::unordered_map<volatile int*, int>& jetty_depth_set,
+                     std::vector<UbEndPoint*>& 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 af456e88e0..25d3e71a2d 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
@@ -63,8 +63,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<Transport::Slice*>& failed_slices,
              std::unordered_map<volatile int*, int>& jetty_depth_set,
+             std::vector<UbEndPoint*>& deferred_deletes,
              int jfc_index) override;
     volatile int* outstandingCount(int jfc_index) override;
     int submitPostSend(
@@ -92,7 +93,10 @@ class UrmaContext : public UbContext {
     bool findJettyOwner(uint32_t jetty_id, UrmaEndpoint** endpoint, int* slot);
     void addDrainingEndpoint(UrmaEndpoint* endpoint);
     void removeDrainingEndpoint(UrmaEndpoint* endpoint);
-    void checkJettyDrainTimeouts();
+    void checkJettyDrainTimeouts(
+        std::unordered_map<volatile int*, int>& jetty_depth_set,
+        std::vector<Transport::Slice*>& failed_slices,
+        std::vector<UbEndPoint*>& deferred_deletes);
 
    private:
     int construct(GlobalConfig& config) override;
@@ -171,8 +175,20 @@ class UrmaContext : public UbContext {
 
 // define the UrmaEndpoint class
 class UrmaEndpoint : public UbEndPoint {
+    // UrmaContext::poll drives the jetty state machine via processWrCompletion.
+    friend class UrmaContext;
+
    public:
-    enum JettyState { ACTIVE = 0, DRAINING = 1, REBUILDING = 2 };
+    // 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
+    };
 
     UrmaEndpoint(UrmaContext* context)
         : context_(context), jfc_outstanding_(nullptr) {}
@@ -198,9 +214,16 @@ class UrmaEndpoint : public UbEndPoint {
 
     // Called from UrmaContext::poll on ACK timeout / flush-done / drain
     // timeout.
-    void onJettyError(int slot);
-    void onFlushDone(int slot);
-    void checkDrainTimeout();
+    void onJettyError(int slot, std::vector<UbEndPoint*>& deferred_deletes);
+    void onFlushDone(int slot,
+                     std::unordered_map<volatile int*, int>& jetty_depth_set,
+                     std::vector<Transport::Slice*>& failed_slices,
+                     std::vector<UbEndPoint*>& deferred_deletes,
+                     int& resolved_wr_count);
+    void checkDrainTimeout(
+        std::unordered_map<volatile int*, int>& jetty_depth_set,
+        std::vector<Transport::Slice*>& failed_slices,
+        std::vector<UbEndPoint*>& deferred_deletes);
 
     int findSlotByDepth(volatile int* depth) const;
 
@@ -220,7 +243,31 @@ class UrmaEndpoint : public UbEndPoint {
 
     bool hasNonActiveJettyUnlocked() const;
     int selectActiveJettyUnlocked();
-    int rebuildJettyUnlocked(int slot);
+    // 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<volatile int*, int>& jetty_depth_set,
+        std::vector<Transport::Slice*>& failed_slices,
+        std::vector<UbEndPoint*>& 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<volatile int*, int>& jetty_depth_set,
+        std::vector<Transport::Slice*>& failed_slices,
+        std::vector<UbEndPoint*>& 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
 
@@ -239,6 +286,7 @@ class UrmaEndpoint : public UbEndPoint {
     std::string peer_eid_;
     uint64_t drain_start_ns_ = 0;
     int draining_slot_ = -1;
+    std::vector<uint64_t> 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<volatile int*, int> jetty_depth_set;
+    std::vector<UbEndPoint*> deferred_deletes;
     std::vector<UbTransport::Slice*> 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<UbTransport::Slice*> 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<int>(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/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
index 14a317f3d9..c63ae43540 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
@@ -549,7 +549,10 @@ void UrmaContext::removeDrainingEndpoint(UrmaEndpoint* endpoint) {
     draining_endpoints_.erase(endpoint);
 }
 
-void UrmaContext::checkJettyDrainTimeouts() {
+void UrmaContext::checkJettyDrainTimeouts(
+    std::unordered_map<volatile int*, int>& jetty_depth_set,
+    std::vector<Transport::Slice*>& failed_slices,
+    std::vector<UbEndPoint*>& deferred_deletes) {
     std::vector<UrmaEndpoint*> endpoints;
     {
         RWSpinlock::ReadGuard guard(jetty_owner_lock_);
@@ -557,16 +560,27 @@ void UrmaContext::checkJettyDrainTimeouts() {
                          draining_endpoints_.end());
     }
     for (auto* endpoint : endpoints) {
-        if (endpoint) endpoint->checkDrainTimeout();
+        if (endpoint)
+            endpoint->checkDrainTimeout(jetty_depth_set, failed_slices,
+                                        deferred_deletes);
     }
 }
 
-int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
-                      int& num_failed,
+namespace {
+void deferEndpointDelete(UrmaEndpoint* endpoint,
+                         std::vector<UbEndPoint*>& 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<Transport::Slice*>& failed_slices,
                       std::unordered_map<volatile int*, int>& jetty_depth_set,
+                      std::vector<UbEndPoint*>& deferred_deletes,
                       int jfc_index) {
-    num_failed = 0;
-    checkJettyDrainTimeouts();
     urma_cr_t cr[num_entries];
     int nr_poll = urma_poll_jfc(jfc_list_[jfc_index].native, num_entries, cr);
     if (nr_poll < 0) {
@@ -581,7 +595,8 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
             UrmaEndpoint* endpoint = nullptr;
             int slot = -1;
             if (findJettyOwner(cr[i].local_id, &endpoint, &slot) && endpoint) {
-                endpoint->onFlushDone(slot);
+                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_;
@@ -593,65 +608,22 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices,
         if (!slice) {
             continue;
         }
-        ++wr_completions;
-
-        // 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<UrmaEndpoint*>(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_ACK_TIMEOUT_ERR) {
-            auto* endpoint = static_cast<UrmaEndpoint*>(slice->ub.endpoint);
-            if (endpoint) {
-                int slot = endpoint->findSlotByDepth(depth);
-                if (slot < 0) {
-                    UrmaEndpoint* mapped = nullptr;
-                    int mapped_slot = -1;
-                    if (findJettyOwner(cr[i].local_id, &mapped, &mapped_slot) &&
-                        mapped == endpoint) {
-                        slot = mapped_slot;
-                    }
-                }
-                if (slot >= 0) endpoint->onJettyError(slot);
-            }
+        if (endpoint->processWrCompletion(cr[i], jetty_depth_set, failed_slices,
+                                          deferred_deletes, jfc_index,
+                                          /*allow_error_trigger=*/true)) {
+            ++wr_completions;
         }
-
-        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<urma_target_seg_t*>(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;
-    }
-    // Exclude FLUSH_ERR_DONE from outstanding accounting (it was never posted).
+    }
+    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;
 }
 
@@ -703,6 +675,7 @@ 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();
@@ -740,6 +713,19 @@ 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;
         }
         uint32_t jetty_id = jetty_list_[i]->jetty_id.id;
@@ -896,10 +882,21 @@ void UrmaEndpoint::disconnectUnlocked() {
     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.
+        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";
@@ -914,7 +911,7 @@ void UrmaEndpoint::disconnectUnlocked() {
             __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]);
             wr_depth_list_[i] = 0;
         }
-        jetty_state_[i] = ACTIVE;
+        if (!reset_ret) jetty_state_[i] = ACTIVE;
     }
     imported_jetty_map_.clear();
     peer_jetty_id_.assign(jetty_list_.size(), 0);
@@ -1031,6 +1028,7 @@ 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;
     }
@@ -1151,53 +1149,62 @@ int UrmaEndpoint::selectActiveJettyUnlocked() {
     return -1;
 }
 
-void UrmaEndpoint::onJettyError(int slot) {
+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<UbEndPoint*>& deferred_deletes) {
     bool delete_ep = false;
     {
         RWSpinlock::WriteGuard guard(lock_);
         if (slot < 0 || slot >= static_cast<int>(jetty_list_.size())) return;
-        if (jetty_state_[slot] == DRAINING ||
-            jetty_state_[slot] == REBUILDING) {
-            return;  // idempotent
-        }
-        // Serial rebuild: at most one non-ACTIVE jetty per endpoint.
-        if (hasNonActiveJettyUnlocked()) {
-            LOG(INFO) << "Skip jetty rebuild for slot " << slot
-                      << ": another jetty is already draining/rebuilding on "
-                      << toString();
-            return;
-        }
+        if (jetty_state_[slot] != ACTIVE) return;  // idempotent
         if (!jetty_list_[slot]) return;
-
-        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;
-            delete_ep = true;
-        } else {
-            jetty_state_[slot] = DRAINING;
-            draining_slot_ = slot;
-            drain_start_ns_ = getCurrentTimeInNano();
-            context_->addDrainingEndpoint(this);
-            LOG(WARNING) << "Jetty ACK timeout: start drain slot=" << slot
+        // 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
-                         << " on " << toString();
+                         << ": 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();
-        context_->deleteEndpointByPtr(this);
+        deferEndpointDelete(this, deferred_deletes);
     }
 }
 
-void UrmaEndpoint::onFlushDone(int slot) {
+void UrmaEndpoint::onFlushDone(
+    int slot, std::unordered_map<volatile int*, int>& jetty_depth_set,
+    std::vector<Transport::Slice*>& failed_slices,
+    std::vector<UbEndPoint*>& deferred_deletes, int& resolved_wr_count) {
     bool delete_ep = false;
     {
         RWSpinlock::WriteGuard guard(lock_);
@@ -1206,7 +1213,8 @@ void UrmaEndpoint::onFlushDone(int slot) {
         jetty_state_[slot] = REBUILDING;
         LOG(INFO) << "Jetty flush-done: rebuild slot=" << slot << " on "
                   << toString();
-        if (rebuildJettyUnlocked(slot)) {
+        if (rebuildJettyUnlocked(slot, jetty_depth_set, failed_slices,
+                                 deferred_deletes, resolved_wr_count)) {
             context_->removeDrainingEndpoint(this);
             draining_slot_ = -1;
             drain_start_ns_ = 0;
@@ -1216,11 +1224,14 @@ void UrmaEndpoint::onFlushDone(int slot) {
     if (delete_ep) {
         LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: "
                    << "rebuildJetty failed on " << toString();
-        context_->deleteEndpointByPtr(this);
+        deferEndpointDelete(this, deferred_deletes);
     }
 }
 
-void UrmaEndpoint::checkDrainTimeout() {
+void UrmaEndpoint::checkDrainTimeout(
+    std::unordered_map<volatile int*, int>& /*jetty_depth_set*/,
+    std::vector<Transport::Slice*>& /*failed_slices*/,
+    std::vector<UbEndPoint*>& deferred_deletes) {
     bool delete_ep = false;
     {
         RWSpinlock::WriteGuard guard(lock_);
@@ -1243,11 +1254,75 @@ void UrmaEndpoint::checkDrainTimeout() {
     if (delete_ep) {
         LOG(ERROR) << "Jetty rebuild fallback to deleteEndpoint: "
                    << "flush-done timeout on " << toString();
-        context_->deleteEndpointByPtr(this);
+        deferEndpointDelete(this, deferred_deletes);
     }
 }
 
-int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
+bool UrmaEndpoint::processWrCompletion(
+    urma_cr_t& cr, std::unordered_map<volatile int*, int>& jetty_depth_set,
+    std::vector<Transport::Slice*>& failed_slices,
+    std::vector<UbEndPoint*>& deferred_deletes, int jfc_index,
+    bool allow_error_trigger) {
+    auto slice = reinterpret_cast<Transport::Slice*>(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<int>(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<urma_target_seg_t*>(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<volatile int*, int>& jetty_depth_set,
+    std::vector<Transport::Slice*>& failed_slices,
+    std::vector<UbEndPoint*>& 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;
@@ -1256,7 +1331,7 @@ int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
     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 (may overlap with already-polled CRs).
+    // 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);
@@ -1265,8 +1340,14 @@ int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
             return ERR_ENDPOINT;
         }
         if (flushed == 0) break;
-        // Completions for these WRs should already have been (or will be)
-        // accounted via poll; do not touch slice pointers from flush CRs.
+        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.
@@ -1292,44 +1373,36 @@ int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
         return ERR_ENDPOINT;
     }
     jetty_list_[slot] = nullptr;
-    if (wr_depth_list_[slot] != 0) {
-        __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[slot]);
-        wr_depth_list_[slot] = 0;
-    }
-
-    // 4) Create replacement jetty with the same JFC/JFR config.
-    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";
+    // 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 = {};
@@ -1343,32 +1416,78 @@ int UrmaEndpoint::rebuildJettyUnlocked(int slot) {
         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;
     }
 
-    jetty_list_[slot] = new_jetty;
     imported_jetty_map_[new_jetty] = imported;
-    const uint32_t new_id = new_jetty->jetty_id.id;
-    jetty_id_map_[new_id] = slot;
-    context_->registerJettyOwner(new_id, this, slot);
     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_id
+                 << " 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<int>(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;
 }
 

From 1112d30f63c7ddade6fee85e264198f3bc62aa00 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 00:47:41 +0800
Subject: [PATCH 04/10] [TransferEngine] Apply clang-format to jetty rebuild
 safety fixes

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../kunpeng_transport/urma/urma_endpoint.cpp  | 28 +++++++++----------
 1 file changed, 13 insertions(+), 15 deletions(-)

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 c63ae43540..3a7e18e401 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
@@ -1303,17 +1303,16 @@ bool UrmaEndpoint::processWrCompletion(
 
     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<urma_target_seg_t*>(slice->ub.r_seg)
-                          ->seg.token_id
-                   << ", retry_cnt: " << slice->ub.retry_cnt
-                   << "): " << cr.status << ", jfc idx : " << jfc_index;
+        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<urma_target_seg_t*>(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;
@@ -1342,8 +1341,8 @@ int UrmaEndpoint::rebuildJettyUnlocked(
         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,
+            if (processWrCompletion(flush_crs[j], jetty_depth_set,
+                                    failed_slices, deferred_deletes, -1,
                                     /*allow_error_trigger=*/false)) {
                 ++resolved_wr_count;
             }
@@ -1477,8 +1476,7 @@ int UrmaEndpoint::recreateJettyUnlocked(int slot, urma_jfc_t* reuse_jfc,
     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);
+    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;

From 4b7bc8cd9793aa726156d98ae8bcf74341d983cd Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 11:40:58 +0800
Subject: [PATCH 05/10] [TransferEngine] Add mock-injected tests for jetty ACK
 timeout rebuild

Extend the mock URMA provider so CI can drive the status=9 rebuild path
without real hardware: record each posted WR's jetty local id, script the
next poll status (URMA_CR_ACK_TIMEOUT_ERR), inject a FLUSH_ERR_DONE fence,
return WR_FLUSH_ERR from urma_flush_jetty, and optionally fail the next
urma_create_jetty. Add gtest cases TC-1..TC-3 covering the happy-path
rebuild, flush-CR delivery, and stale-epoch drop, and wire them into the
tent-ci (ub-mock) job as a scoped ctest.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .github/workflows/ci.yml                      |  12 +
 .../jetty-rebuild-mock-test-plan.md           | 341 ++++++++++++++++++
 .../jetty-single-rebuild-plan.md              |   2 +
 .../kunpeng_transport/urma/urma_endpoint.h    |   9 +
 .../kunpeng_transport/urma/mock_urma.cpp      | 175 ++++++++-
 .../urma/mock_urma_test_ctrl.h                |  55 +++
 mooncake-transfer-engine/tests/CMakeLists.txt |  13 +
 .../tests/urma_jetty_rebuild_test.cpp         | 334 +++++++++++++++++
 8 files changed, 925 insertions(+), 16 deletions(-)
 create mode 100644 docs/source/design/transfer-engine/jetty-rebuild-mock-test-plan.md
 create mode 100644 mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h
 create mode 100644 mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp

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-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<PendingWr> 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
index 956cb201ce..24c6e527e9 100644
--- a/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
+++ b/docs/source/design/transfer-engine/jetty-single-rebuild-plan.md
@@ -211,6 +211,8 @@ if cr.status == ACK_TIMEOUT_ERR (9):
 - 单测：状态机、选槽跳过、假 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)`）
 
 ---
 
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 25d3e71a2d..d78e58f23f 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
@@ -51,9 +51,13 @@ static urma_import_seg_flag_t import_flag = {
 
 // 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,
@@ -173,10 +177,15 @@ class UrmaContext : public UbContext {
     std::unordered_set<UrmaEndpoint*> 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
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 a05e04280c..f2dae63133 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 <algorithm>
 #include <atomic>
 #include <cstring>
@@ -10,11 +11,39 @@
 
 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;
+};
+
 struct JfcState {
     std::mutex mutex;
-    std::deque<uint64_t> pending_ctx;
+    std::deque<PendingWr> 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<uint32_t> 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;
 };
 
+std::mutex g_script_mutex;
+MockScript g_script;
+
 std::shared_mutex g_rw_mutex;
 bool initialized = false;
 std::vector<urma_device_t *> device_list;
@@ -342,6 +371,13 @@ 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<std::mutex> script_lock(g_script_mutex);
+        if (g_script.fail_next_create_jetty) {
+            g_script.fail_next_create_jetty = false;
+            return nullptr;
+        }
+    }
     std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
     if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) {
         return nullptr;
@@ -422,13 +458,51 @@ urma_status_t urma_modify_jetty(urma_jetty_t *jetty, urma_jetty_attr_t *attr) {
 }
 
 int urma_flush_jetty(urma_jetty_t *jetty, int cr_cnt, urma_cr_t *cr) {
-    (void)cr_cnt;
-    (void)cr;
-    std::shared_lock<std::shared_mutex> lock(g_rw_mutex);
-    if (!jetty || jetty_map.find(jetty) == jetty_map.end()) {
-        return -1;
+    urma_jfc_t *jfc = nullptr;
+    uint32_t local_id = 0;
+    {
+        std::shared_lock<std::shared_mutex> 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<std::mutex> 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<std::shared_mutex> 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<std::mutex> 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 0;
+    return produced;
 }
 
 urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
@@ -461,7 +535,8 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
         std::lock_guard<std::mutex> jfc_lock(state->mutex);
         urma_jfs_wr_t *current_wr = wr;
         while (current_wr) {
-            state->pending_ctx.push_back(current_wr->user_ctx);
+            state->pending.push_back(
+                PendingWr{current_wr->user_ctx, jetty->jetty_id.id});
             current_wr = current_wr->next;
         }
     }
@@ -483,17 +558,85 @@ int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) {
         state = it->second;
     }
 
+    std::lock_guard<std::mutex> 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<std::mutex> jfc_lock(state->mutex);
-        int available = static_cast<int>(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<std::mutex> 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;
+            }
+        }
+    }
+
+    int available = static_cast<int>(state->pending.size());
+    int wr_completed = std::min(num_entries - num_completed, available);
+    for (int i = 0; i < wr_completed; ++i) {
+        urma_cr_status_t status = URMA_CR_SUCCESS;
+        {
+            std::lock_guard<std::mutex> script_lock(g_script_mutex);
+            if (g_script.poll_status_count > 0) {
+                status = g_script.poll_status;
+                --g_script.poll_status_count;
+            }
         }
-        state->pending_ctx.erase(state->pending_ctx.begin(),
-                                 state->pending_ctx.begin() + num_completed);
+        cr_list[num_completed].status = status;
+        cr_list[num_completed].user_ctx = state->pending[i].user_ctx;
+        cr_list[num_completed].local_id = state->pending[i].jetty_local_id;
+        ++num_completed;
     }
+    state->pending.erase(state->pending.begin(),
+                         state->pending.begin() + wr_completed);
     return num_completed;
 }
+
+// ---------------------------------------------------------------------------
+// Test-only scripted control API. See mock_urma_test_ctrl.h.
+// ---------------------------------------------------------------------------
+
+void mock_urma_test_reset(void) {
+    std::lock_guard<std::mutex> script_lock(g_script_mutex);
+    g_script = MockScript{};
+}
+
+void mock_urma_set_next_poll_status(int status, int count) {
+    std::lock_guard<std::mutex> script_lock(g_script_mutex);
+    g_script.poll_status = static_cast<urma_cr_status_t>(status);
+    g_script.poll_status_count = count;
+}
+
+void mock_urma_enqueue_flush_done(uint32_t jetty_local_id) {
+    std::lock_guard<std::mutex> 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<std::mutex> script_lock(g_script_mutex);
+    g_script.flush_err_count = count;
+}
+
+void mock_urma_fail_next_create_jetty(void) {
+    std::lock_guard<std::mutex> 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..78fb089efd
--- /dev/null
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma_test_ctrl.h
@@ -0,0 +1,55 @@
+// 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 <stdint.h>
+
+#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);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MOCK_URMA_TEST_CTRL_H
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..f9cc2e25d5
--- /dev/null
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -0,0 +1,334 @@
+// 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 <gtest/gtest.h>
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#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 <sanitizer/lsan_interface.h>
+#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<UbSIEVEEndpointStore>(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; }
+
+    // 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<UrmaContext>(*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<UrmaEndpoint>(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());
+    }
+
+    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<void *>(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 = nullptr;
+        slice->ub.l_seg = nullptr;
+        slice->ub.retry_cnt = 0;
+        slice->ub.max_retry_cnt = 0;
+
+        std::vector<Transport::Slice *> slices{slice};
+        std::vector<Transport::Slice *> 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<Transport::Slice *> &failed_slices,
+                 std::unordered_map<volatile int *, int> &jetty_depth_set,
+                 std::vector<UbEndPoint *> &deferred_deletes) {
+        return context_->poll(16, failed_slices, jetty_depth_set,
+                              deferred_deletes, 0);
+    }
+
+    TestUbTransport *transport_ = nullptr;
+    std::unique_ptr<UrmaContext> context_;
+    std::unique_ptr<UrmaEndpoint> endpoint_;
+};
+
+// 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<Transport::Slice *> failed;
+    std::unordered_map<volatile int *, int> depth_set;
+    std::vector<UbEndPoint *> 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);
+    Transport::Slice *s2 = postOneSlice(&task);
+    const uint32_t old_id = UrmaEndpointTestPeer::jettyId(*endpoint_, 0);
+    EXPECT_EQ(2, UrmaEndpointTestPeer::wrDepth(*endpoint_, 0));
+
+    // Drive into DRAINING without consuming the two WRs.
+    mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1);
+    // Move the two real WRs aside so the ACK-timeout completion can refer to
+    // the first one; the remaining one stays outstanding for the flush loop.
+    std::vector<Transport::Slice *> failed;
+    std::unordered_map<volatile int *, int> depth_set;
+    std::vector<UbEndPoint *> deferred;
+    // Re-post nothing; mark both WRs with ACK timeout would fail both. Instead
+    // fail only the first (count=1), leaving s2 outstanding.
+    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);
+    EXPECT_EQ(0, resolved);  // fence not counted; flush accounted internally
+
+    // 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<Transport::Slice *> failed;
+    std::unordered_map<volatile int *, int> depth_set;
+    std::vector<UbEndPoint *> 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;
+}
+
+}  // namespace

From b806242d60353a56b09a4dd68a9efacc37fc2724 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 14:17:46 +0800
Subject: [PATCH 06/10] [TransferEngine] Fix segfault in jetty rebuild mock
 test fixture

TC-1 segfaulted on CI because postOneSlice left slice->ub.r_seg null while
processWrCompletion's failure log dereferences it. Import a real mock target
segment in the fixture and attach it to each posted slice.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../transport/kunpeng_transport/urma/mock_urma.cpp   |  2 +-
 .../tests/urma_jetty_rebuild_test.cpp                | 12 +++++++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

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 f2dae63133..69cf5c3da6 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
@@ -339,7 +339,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;
diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
index f9cc2e25d5..28da8b330b 100644
--- a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -162,6 +162,15 @@ class UrmaJettyRebuildTest : public ::testing::Test {
         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 {
@@ -185,7 +194,7 @@ class UrmaJettyRebuildTest : public ::testing::Test {
         slice->task = task;
         slice->from_cache = false;
         slice->ub.dest_addr = 0x2000;
-        slice->ub.r_seg = nullptr;
+        slice->ub.r_seg = imported_seg_;
         slice->ub.l_seg = nullptr;
         slice->ub.retry_cnt = 0;
         slice->ub.max_retry_cnt = 0;
@@ -209,6 +218,7 @@ class UrmaJettyRebuildTest : public ::testing::Test {
     TestUbTransport *transport_ = nullptr;
     std::unique_ptr<UrmaContext> context_;
     std::unique_ptr<UrmaEndpoint> endpoint_;
+    urma_target_seg_t *imported_seg_ = nullptr;
 };
 
 // TC-1: status=9 -> DRAINING -> FLUSH_ERR_DONE -> rebuild -> slice resolved.

From 31d0205375ec5e85fb1ae650e1dbd3c910c69d5d Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 15:46:13 +0800
Subject: [PATCH 07/10] [TransferEngine] Keep a residual WR outstanding in
 jetty rebuild flush test

TC-2 failed on CI because the single-JFC/single-slot fixture drained both
posted WRs in the first poll (s1 failed with status 9, s2 succeeded), leaving
nothing for the rebuild flush loop to deliver. Add a withhold hook to the mock
so a posted WR stays outstanding until urma_flush_jetty completes it, and post
s2 with that hook so the flush path is genuinely exercised.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../kunpeng_transport/urma/mock_urma.cpp      | 41 +++++++++++++++----
 .../urma/mock_urma_test_ctrl.h                |  6 +++
 .../tests/urma_jetty_rebuild_test.cpp         |  9 ++--
 3 files changed, 43 insertions(+), 13 deletions(-)

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 69cf5c3da6..61dc5bbc5b 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
@@ -18,6 +18,10 @@ namespace {
 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 {
@@ -39,6 +43,9 @@ struct MockScript {
     int flush_err_count = 0;
     // When true, the next urma_create_jetty returns NULL.
     bool fail_next_create_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;
@@ -535,8 +542,16 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
         std::lock_guard<std::mutex> jfc_lock(state->mutex);
         urma_jfs_wr_t *current_wr = wr;
         while (current_wr) {
+            bool withhold = false;
+            {
+                std::lock_guard<std::mutex> 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});
+                PendingWr{current_wr->user_ctx, jetty->jetty_id.id, withhold});
             current_wr = current_wr->next;
         }
     }
@@ -590,9 +605,14 @@ int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) {
         }
     }
 
-    int available = static_cast<int>(state->pending.size());
-    int wr_completed = std::min(num_entries - num_completed, available);
-    for (int i = 0; i < wr_completed; ++i) {
+    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<std::mutex> script_lock(g_script_mutex);
@@ -602,12 +622,12 @@ int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) {
             }
         }
         cr_list[num_completed].status = status;
-        cr_list[num_completed].user_ctx = state->pending[i].user_ctx;
-        cr_list[num_completed].local_id = state->pending[i].jetty_local_id;
+        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);
     }
-    state->pending.erase(state->pending.begin(),
-                         state->pending.begin() + wr_completed);
     return num_completed;
 }
 
@@ -636,6 +656,11 @@ void mock_urma_set_flush_returns_errors(int count) {
     g_script.flush_err_count = count;
 }
 
+void mock_urma_withhold_next_post(int count) {
+    std::lock_guard<std::mutex> script_lock(g_script_mutex);
+    g_script.withhold_next_post_count = count;
+}
+
 void mock_urma_fail_next_create_jetty(void) {
     std::lock_guard<std::mutex> 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
index 78fb089efd..f40acee8f2 100644
--- 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
@@ -48,6 +48,12 @@ void mock_urma_set_flush_returns_errors(int count);
 // rebuild-failure -> deferred-delete fallback.
 void mock_urma_fail_next_create_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
diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
index 28da8b330b..cbbe79160b 100644
--- a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -269,19 +269,18 @@ TEST_F(UrmaJettyRebuildTest, AckTimeoutTriggersRebuildHappyPath) {
 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 without consuming the two WRs.
+    // Drive into DRAINING: only s1 is polled (status 9); s2 stays outstanding.
     mock_urma_set_next_poll_status(URMA_CR_ACK_TIMEOUT_ERR, 1);
-    // Move the two real WRs aside so the ACK-timeout completion can refer to
-    // the first one; the remaining one stays outstanding for the flush loop.
     std::vector<Transport::Slice *> failed;
     std::unordered_map<volatile int *, int> depth_set;
     std::vector<UbEndPoint *> deferred;
-    // Re-post nothing; mark both WRs with ACK timeout would fail both. Instead
-    // fail only the first (count=1), leaving s2 outstanding.
     pollOnce(failed, depth_set, deferred);
     ASSERT_EQ(UrmaEndpoint::DRAINING,
               UrmaEndpointTestPeer::jettyState(*endpoint_, 0));

From 36e65a112b59f074ace8cafed29b79c2ae8cdd87 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 16:25:35 +0800
Subject: [PATCH 08/10] [TransferEngine] Correct resolved-count assertion in
 rebuild flush test

The rebuild flush loop delivers each outstanding WR through
processWrCompletion and increments poll's resolved_wr_count, so delivering s2
via flush yields resolved == 1 (the fence marker itself is not counted). Fix
the test to assert that instead of 0.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
index cbbe79160b..098a2397fa 100644
--- a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -293,7 +293,9 @@ TEST_F(UrmaJettyRebuildTest, FlushCompletionsDeliveredOnRebuild) {
     depth_set.clear();
     deferred.clear();
     int resolved = pollOnce(failed, depth_set, deferred);
-    EXPECT_EQ(0, resolved);  // fence not counted; flush accounted internally
+    // 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;

From 600f5063b5e0d2dcf3078dfbb20e43c3f33276a6 Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 19:29:18 +0800
Subject: [PATCH 09/10] [TransferEngine] Harden jetty rebuild failure and
 drain-timeout paths

Rebuild failure (urma_endpoint.cpp, urma_endpoint.h): when urma_delete_jetty
fails mid-rebuild, keep the old jetty handle and mark the slot
REBUILDING_FAILED instead of nulling it, so deconstruct can retry the delete
rather than leak the handle. The new state is never selected for posting.

Drain timeout (urma_endpoint.cpp): when the flush-done fence never arrives,
checkDrainTimeout now flushes the jetty and delivers each residual WR through
processWrCompletion (endpoint still alive) so stuck slices are failed/retried
and depth accounting returns to zero before the endpoint is deferred for
deletion. disconnectUnlocked keeps its drain-only flush (endpoint is being
torn down), with a comment clarifying why CRs are not delivered there.

Add TC-4 (rebuild failure keeps the handle, deferred delete) and TC-5 (drain
timeout delivers residual WRs) to urma_jetty_rebuild_test.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../kunpeng_transport/urma/urma_endpoint.h    |  7 +-
 .../kunpeng_transport/urma/urma_endpoint.cpp  | 51 ++++++++++-
 .../tests/urma_jetty_rebuild_test.cpp         | 90 +++++++++++++++++++
 3 files changed, 143 insertions(+), 5 deletions(-)

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 d78e58f23f..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
@@ -196,7 +196,12 @@ class UrmaEndpoint : public UbEndPoint {
         ACTIVE = 0,
         DRAINING = 1,
         REBUILDING = 2,
-        PENDING_DRAIN = 3
+        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)
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 3a7e18e401..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
@@ -885,7 +885,13 @@ void UrmaEndpoint::disconnectUnlocked() {
         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.
+        // 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) {
@@ -1229,10 +1235,11 @@ void UrmaEndpoint::onFlushDone(
 }
 
 void UrmaEndpoint::checkDrainTimeout(
-    std::unordered_map<volatile int*, int>& /*jetty_depth_set*/,
-    std::vector<Transport::Slice*>& /*failed_slices*/,
+    std::unordered_map<volatile int*, int>& jetty_depth_set,
+    std::vector<Transport::Slice*>& failed_slices,
     std::vector<UbEndPoint*>& deferred_deletes) {
     bool delete_ep = false;
+    int resolved_wr_count = 0;
     {
         RWSpinlock::WriteGuard guard(lock_);
         if (draining_slot_ < 0) return;
@@ -1246,11 +1253,44 @@ void UrmaEndpoint::checkDrainTimeout(
         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();
@@ -1368,7 +1408,10 @@ int UrmaEndpoint::rebuildJettyUnlocked(
     ret = urma_delete_jetty(old_jetty);
     if (ret) {
         PLOG(ERROR) << "Failed to delete jetty during rebuild";
-        jetty_list_[slot] = nullptr;
+        // 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;
diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
index 098a2397fa..53949ea262 100644
--- a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -93,6 +93,14 @@ class UrmaEndpointTestPeer {
     }
     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) {
@@ -342,4 +350,86 @@ TEST_F(UrmaJettyRebuildTest, StaleEpochCompletionDropped) {
     delete slice;
 }
 
+// TC-4: a failed rebuild keeps the old jetty handle (marked REBUILDING_FAILED)
+// so deconstruct can retry urma_delete_jetty 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<Transport::Slice *> failed;
+    std::unordered_map<volatile int *, int> depth_set;
+    std::vector<UbEndPoint *> deferred;
+    pollOnce(failed, depth_set, deferred);
+    ASSERT_EQ(UrmaEndpoint::DRAINING,
+              UrmaEndpointTestPeer::jettyState(*endpoint_, 0));
+
+    // Make rebuild's recreate step fail (urma_create_jetty returns NULL), then
+    // inject the fence so onFlushDone runs rebuildJettyUnlocked to failure.
+    mock_urma_fail_next_create_jetty();
+    mock_urma_enqueue_flush_done(old_id);
+    failed.clear();
+    depth_set.clear();
+    deferred.clear();
+    pollOnce(failed, depth_set, deferred);
+
+    // The rebuild failed: the endpoint is deferred for deletion, and the old
+    // jetty handle is preserved (not nulled) so it can be cleaned up later.
+    ASSERT_FALSE(deferred.empty());
+    EXPECT_EQ(static_cast<UbEndPoint *>(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<Transport::Slice *> failed;
+    std::unordered_map<volatile int *, int> depth_set;
+    std::vector<UbEndPoint *> 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

From 63acf9aef7d411284d9a50aa2fbc11daa03e180d Mon Sep 17 00:00:00 2001
From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Date: Thu, 27 Aug 2026 19:50:44 +0800
Subject: [PATCH 10/10] [TransferEngine] Test jetty delete-failure path, not
 create-failure
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

TC-4 was driving rebuild failure via fail_next_create_jetty, which fails in
recreateJettyUnlocked after the old jetty was already deleted — so the handle
is legitimately null and the slot stays REBUILDING. The handle-preservation
fix targets urma_delete_jetty failure instead. Add a fail_next_delete_jetty
hook to the mock and point TC-4 at that branch, which is where the old handle
must be kept (REBUILDING_FAILED) for deconstruct to retry.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../kunpeng_transport/urma/mock_urma.cpp       | 15 +++++++++++++++
 .../urma/mock_urma_test_ctrl.h                 |  5 +++++
 .../tests/urma_jetty_rebuild_test.cpp          | 18 ++++++++++--------
 3 files changed, 30 insertions(+), 8 deletions(-)

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 61dc5bbc5b..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
@@ -43,6 +43,9 @@ struct MockScript {
     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;
@@ -402,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<std::mutex> 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<std::shared_mutex> lock(g_rw_mutex);
     if (!jetty || jetty_map.find(jetty) == jetty_map.end()) {
         return URMA_EINVAL;
@@ -661,6 +671,11 @@ void mock_urma_withhold_next_post(int count) {
     g_script.withhold_next_post_count = count;
 }
 
+void mock_urma_fail_next_delete_jetty(void) {
+    std::lock_guard<std::mutex> script_lock(g_script_mutex);
+    g_script.fail_next_delete_jetty = true;
+}
+
 void mock_urma_fail_next_create_jetty(void) {
     std::lock_guard<std::mutex> 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
index f40acee8f2..3f42167657 100644
--- 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
@@ -48,6 +48,11 @@ void mock_urma_set_flush_returns_errors(int count);
 // 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
diff --git a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
index 53949ea262..06458f2a1d 100644
--- a/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
+++ b/mooncake-transfer-engine/tests/urma_jetty_rebuild_test.cpp
@@ -350,9 +350,10 @@ TEST_F(UrmaJettyRebuildTest, StaleEpochCompletionDropped) {
     delete slice;
 }
 
-// TC-4: a failed rebuild keeps the old jetty handle (marked REBUILDING_FAILED)
-// so deconstruct can retry urma_delete_jetty instead of leaking it, and the
-// endpoint is deferred for deletion rather than torn down inside poll.
+// 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);
@@ -369,17 +370,18 @@ TEST_F(UrmaJettyRebuildTest, RebuildFailureKeepsJettyHandle) {
     ASSERT_EQ(UrmaEndpoint::DRAINING,
               UrmaEndpointTestPeer::jettyState(*endpoint_, 0));
 
-    // Make rebuild's recreate step fail (urma_create_jetty returns NULL), then
-    // inject the fence so onFlushDone runs rebuildJettyUnlocked to failure.
-    mock_urma_fail_next_create_jetty();
+    // 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: the endpoint is deferred for deletion, and the old
-    // jetty handle is preserved (not nulled) so it can be cleaned up later.
+    // 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<UbEndPoint *>(endpoint_.get()), deferred.back());
     EXPECT_EQ(UrmaEndpoint::REBUILDING_FAILED,