[Bugfix][TENT] Drain RDMA work before publishing software timeout

docs/source/api-reference/cpp/tent.md

@@ -334,6 +334,16 @@ Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
 
 Queries the status of transfer requests.
 
+For RDMA, a software timeout cancels further attempts of the expired task and
+retires its affected endpoints. The task remains `PENDING` until its posted
+work has completed or flushed, or the owning QPs have been successfully
+destroyed. It then reports `TIMEOUT`, including for completions that succeed
+after timeout cancellation. Other tasks sharing a QP can retry after their
+flush completions. A failed QP destruction does not itself finish a transfer;
+if neither completion nor destruction succeeds, the buffers must remain alive.
+This is not a hard completion deadline and does not change Python or Store
+wait deadlines.
+
 - `batch_id`: The batch to query.
 - `task_id`: Index of the specific task (for single-task query).
 - `status` / `status_list` / `overall_status`: Output parameter(s) for status.

mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h

@@ -113,6 +113,11 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     void beginDestroy();
     bool finishDestroy();
 
+    // Retire only if this posted slice still belongs to the endpoint. Leave
+    // its queue/accounting intact for CQ drain or successful destruction.
+    // Returns true when this call starts retirement.
+    bool retireForTimeout(RdmaSlice* slice);
+
     Status connect(const std::string& peer_server_name,
                    const std::string& peer_nic_name,
                    const std::string& peer_rpc_server_addr = "");
@@ -245,7 +250,7 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     void beginDestroyNoLock();
     int deconstructUnlocked();
 
-    void resetInflightSlices();
+    void resetInflightSlices(int qp_index);
 
     void postNotifyRecv(size_t idx);
     // postNotifyRecv() for a slot that has just been consumed, subject to

mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h

@@ -113,8 +113,8 @@ struct RdmaTask {
     std::atomic<int> resolved_slices{0};
     volatile TransferStatusEnum first_error = PENDING;
 
-    // Set by the control thread. Workers observe this flag before posting or
-    // retrying a slice. Already-posted WRs are allowed to drain normally.
+    // Set by the control thread or a worker detecting a software timeout.
+    // Workers observe it before posting/retrying; posted WRs must still drain.
     std::atomic<bool> cancel_requested{false};
 
     // Reference counting for UAF protection
@@ -229,6 +229,13 @@ static inline void updateSliceStatus(RdmaSlice* slice,
              task->num_slices)
                 ? COMPLETED
                 : task->first_error;
+        // Timeout cancellation prevents further attempts but is not itself
+        // completion evidence. Only publish it after all slices resolve,
+        // even when every CQE reports a late success. Keep task access after
+        // the slice CAS: an already-terminal slice may outlive its task.
+        if (task->cancel_requested.load(std::memory_order_acquire) &&
+            __atomic_load_n(&task->first_error, __ATOMIC_ACQUIRE) == TIMEOUT)
+            final_st = TIMEOUT;
         if (final_st == PENDING) final_st = FAILED;
         __sync_bool_compare_and_swap(&task->status_word, PENDING, final_st);
     }

mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp

@@ -226,7 +226,6 @@ int RdmaEndPoint::deconstructUnlocked() {
     auto current_status = status_.load(std::memory_order_relaxed);
     // Idempotent, including delayed destruction through an external shared_ptr.
     if (current_status == EP_DESTROYED) return 0;
-    resetInflightSlices();
     peer_qp_num_list_.clear();
 
     // A default-constructed endpoint owns no verbs resources. A partially
@@ -282,18 +281,22 @@ int RdmaEndPoint::deconstructUnlocked() {
 
     bool all_qps_destroyed = true;
     for (size_t i = 0; i < qp_list_.size(); ++i) {
-        if (wr_depth_list_) {
-            const int outstanding =
-                wr_depth_list_[i].value.load(std::memory_order_relaxed);
-            if (outstanding != 0) cancelQuota(i, outstanding);
-        }
-        if (!qp_list_[i]) continue;
-        if (context_->verbs_.ibv_destroy_qp(qp_list_[i])) {
+        if (qp_list_[i] && context_->verbs_.ibv_destroy_qp(qp_list_[i])) {
             PLOG(ERROR) << "Failed to destroy data QP[" << i << "]";
             result = -1;
             all_qps_destroyed = false;
         } else {
             qp_list_[i] = nullptr;
+            if (wr_depth_list_) {
+                const int outstanding =
+                    wr_depth_list_[i].value.load(std::memory_order_relaxed);
+                if (outstanding != 0) cancelQuota(i, outstanding);
+            }
+            // Detach this QP's slices in the same critical section as its
+            // quota release. Another QP may fail destruction; a later CQE
+            // must not find these entries and release their quota again.
+            // Only this successfully destroyed QP can publish cancellation.
+            resetInflightSlices(i);
         }
     }
     if (all_qps_destroyed) {
@@ -322,6 +325,25 @@ void RdmaEndPoint::beginDestroy() {
     beginDestroyNoLock();
 }
 
+bool RdmaEndPoint::retireForTimeout(RdmaSlice* slice) {
+    {
+        RWSpinlock::WriteGuard guard(lock_);
+        const int qp_index = slice->qp_index;
+        if (qp_index < 0 || qp_index >= (int)slice_queue_.size() ||
+            !slice_queue_[qp_index].contains(slice) ||
+            status_.load(std::memory_order_relaxed) != EP_READY)
+            return false;
+        // A stale worker-set entry must not retire unrelated work. Once
+        // validated, stop new posts and request flushes without resolving
+        // any WR here. Other tasks on the QP retain their normal retry path.
+        beginDestroyNoLock();
+    }
+    // Match resetConnection(): remove outside the endpoint lock so the
+    // monitor can reclaim it even when there are no future submissions.
+    context_->endpointStore()->remove(this);
+    return true;
+}
+
 void RdmaEndPoint::beginDestroyNoLock() {
     auto current_status = status_.load(std::memory_order_relaxed);
     if (current_status == EP_DESTROYING || current_status == EP_DESTROYED)
@@ -884,16 +906,15 @@ int RdmaEndPoint::submitRecvImmDataRequest(int qp_index, uint64_t id) {
     return 1;
 }
 
-void RdmaEndPoint::resetInflightSlices() {
+void RdmaEndPoint::resetInflightSlices(int qp_index) {
     // qp_list_ is sized up front, while slice_queue_ grows after each
     // successful QP creation. Use the latter during partial-construction
     // cleanup to avoid indexing queues that were never created.
-    for (int qp_index = 0; qp_index < (int)slice_queue_.size(); ++qp_index) {
-        auto& queue = slice_queue_[qp_index];
-        while (!queue.empty()) {
-            auto current = queue.pop();
-            updateSliceStatus(current, TransferStatusEnum::CANCELED);
-        }
+    if (qp_index >= (int)slice_queue_.size()) return;
+    auto& queue = slice_queue_[qp_index];
+    while (!queue.empty()) {
+        auto current = queue.pop();
+        updateSliceStatus(current, TransferStatusEnum::CANCELED);
     }
 }
 

mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp

@@ -965,30 +965,32 @@ void Workers::expireTimedOutSlices(WorkerContext& worker, uint64_t now_ns) {
             retire(slice);
             continue;
         }
-        if (now_ns - slice->enqueue_ts > slice_timeout_ns_) {
+        const bool task_timed_out =
+            slice->task->cancel_requested.load(std::memory_order_acquire) &&
+            __atomic_load_n(&slice->task->first_error, __ATOMIC_ACQUIRE) ==
+                TIMEOUT;
+        if (task_timed_out || now_ns - slice->enqueue_ts > slice_timeout_ns_) {
+            // Stop queued siblings and prevent a flush completion from
+            // resubmitting this task while its QP is being aborted.
+            if (__sync_bool_compare_and_swap(&slice->task->first_error, PENDING,
+                                             TIMEOUT))
+                LOG(WARNING) << "Slice " << slice
+                             << " failed: transfer timeout (software)";
+            slice->task->cancel_requested.store(true,
+                                                std::memory_order_release);
             auto ep = slice->ep_weak_ptr.lock();
-            LOG(WARNING) << "Slice " << slice
-                         << " failed: transfer timeout (software)";
-            // The slice turns terminal here, not on the CQ. Its flush
-            // completion may never be polled (acknowledge() zeroes wr_depth,
-            // so the endpoint destroys the QP and its unpolled CQEs), so
-            // return the selector charge now for it and every slice swept
-            // with it; releaseSliceQuota() is idempotent.
-            if (!ep) {
+            // A software deadline does not stop DMA. Only a completed WR or
+            // successful endpoint destruction permits terminal publication.
+            if (!ep || ep->status() == RdmaEndPoint::EP_DESTROYED ||
+                slice->completions_owed.load(std::memory_order_acquire) == 0) {
                 retire(slice);
                 updateSliceStatus(slice, TIMEOUT);
                 continue;
             }
-            auto num_slices = ep->acknowledge(slice, TIMEOUT, retire);
-            disableEndpoint(slice);
-            if (num_slices == 0) {
-                // A neighbour's retry already popped it off the queue pair
-                // and discounted it there, so only its charge and its set
-                // entry are still outstanding.
-                releaseSliceQuota(slice, now_ns);
-                if (routeSetRemoval(worker, slice))
-                    slice_to_remove.push_back(slice);
-                updateSliceStatus(slice, TIMEOUT);
+            if (ep->retireForTimeout(slice)) {
+                if (auto* rail = slice->rail_monitor)
+                    rail->markFailed(slice->source_dev_id,
+                                     slice->target_dev_id);
             }
         }
     }
@@ -1075,6 +1077,19 @@ void Workers::handleCompletion(WorkerContext& worker, RdmaContext& context,
         discountFromOwner(worker, slice);
         return;
     }
+    if (wc.status != IBV_WC_SUCCESS &&
+        slice->task->cancel_requested.load(std::memory_order_acquire) &&
+        __atomic_load_n(&slice->task->first_error, __ATOMIC_ACQUIRE) ==
+            TIMEOUT) {
+        // This CQE proves only this slice is done. Do not fail other tasks
+        // in the QP prefix or spend a retry on the canceled task.
+        ep->acknowledge(slice, PENDING, [&](RdmaSlice* swept) {
+            discountFromOwner(worker, swept);
+        });
+        updateSliceStatus(slice, TIMEOUT);
+        discountFromOwner(worker, slice);
+        return;
+    }
     if (wc.status != IBV_WC_SUCCESS) {
         if (wc.status != IBV_WC_WR_FLUSH_ERR) {
             // TE handles them automatically