[TENT] Add best-effort RDMA task cancellation

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

@@ -75,6 +75,7 @@ For users migrating from Transfer Engine, the following table shows how TE APIs
 | `allocateBatchID(batch_size)` | `allocateBatch(batch_size)` | Renamed |
 | `freeBatchID(batch_id)` | `freeBatch(batch_id)` | Renamed |
 | `submitTransfer(batch_id, entries)` | `submitTransfer(batch_id, request_list)` | `TransferRequest` → `Request` |
+| *Not available* | `cancelTransfer(batch_id, task_id)` | TENT-only: best-effort cancellation for queued and RDMA tasks |
 | `submitTransferWithNotify(batch_id, entries, notify_msg)` | `submitTransfer(batch_id, request_list, notifi)` | Unified API with optional notification |
 | `getTransferStatus(batch_id, task_id, status)` | `getTransferStatus(batch_id, task_id, status)` | Same |
 | `getBatchTransferStatus(batch_id, status)` | `getTransferStatus(batch_id, status)` | Overloaded; single `TransferStatus` output = overall status |
@@ -274,6 +275,24 @@ Queries the status of transfer requests.
 - `status` / `status_list` / `overall_status`: Output parameter(s) for status.
 - Return value: `Status::OK()` on success; otherwise a non-OK status.
 
+#### TransferEngine::cancelTransfer
+
+```cpp
+Status cancelTransfer(BatchID batch_id, size_t task_id);
+```
+
+Requests best-effort cancellation of one public task. A task still waiting in
+the TENT admission queue becomes `CANCELED` without being dispatched. For RDMA,
+workers suppress slices they observe before `ibv_post_send`; work already
+posted to a QP is allowed to drain and may complete successfully. Consequently,
+the API returning `OK` means the cancellation request was accepted, not that
+the task is already terminal. Continue polling `getTransferStatus` before
+calling `freeBatch`.
+
+Cancellation is idempotent. Merged public tasks share one physical transfer,
+so canceling any alias cancels the shared task. Direct cancellation of staging
+or non-RDMA transport work currently returns `Status::NotImplemented`.
+
 #### TransferEngine::freeBatch
 
 ```cpp

mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h

@@ -134,6 +134,11 @@ class LocalTransferAdmissionQueue {
 
     Status complete(QueueOwnerId owner_id, TransferStatusEnum terminal_status);
 
+    // Cancel an owner that has not entered the dispatch window. Idempotent for
+    // an owner already canceled; dispatching owners must be canceled through
+    // their selected transport instead.
+    Status cancel(QueueOwnerId owner_id);
+
     Status retireBatch(uint64_t batch_token);
 
     Status resolveOwner(uint64_t batch_token, size_t public_task_id,

mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h

@@ -58,6 +58,7 @@ struct TaskInfo {
     std::string qp_pool;          // Named QP pool (RFC #2568 step 3), "" = none
     Request request;
     bool staging{false};
+    bool cancel_requested{false};
     TransferStatusEnum status{TransferStatusEnum::PENDING};
     volatile TransferStatusEnum staging_status{TransferStatusEnum::PENDING};
     std::chrono::steady_clock::time_point start_time{};  // For latency tracking
@@ -140,6 +141,8 @@ class TransferEngineImpl {
                           const std::vector<Request>& request_list,
                           const Notification& notifi);
 
+    Status cancelTransfer(BatchID batch_id, size_t task_id);
+
     Status sendNotification(SegmentID target_id, const Notification& notifi);
 
     Status receiveNotification(std::vector<Notification>& notifi_list);
@@ -246,6 +249,8 @@ class TransferEngineImpl {
     Status finishQueuedOwner(QueueOwnerId owner_id,
                              TransferStatusEnum terminal_status);
 
+    Status cancelQueuedOwner(QueueOwnerId owner_id);
+
     Status retireQueueForBatch(Batch* batch);
 
     Status pollTaskStatus(Batch* batch, size_t task_id,

mooncake-transfer-engine/tent/include/tent/runtime/transport.h

@@ -63,14 +63,14 @@ class Transport {
         std::function<void(BatchID)> notify_progress;
     };
 
-    using SubBatchRef = SubBatch *;
+    using SubBatchRef = SubBatch*;
 
    public:
     Transport() = default;
 
     virtual ~Transport() = default;
 
-    virtual Status install(std::string &local_segment_name,
+    virtual Status install(std::string& local_segment_name,
                            std::shared_ptr<ControlService> metadata,
                            std::shared_ptr<Topology> local_topology,
                            std::shared_ptr<Config> conf = nullptr) {
@@ -81,74 +81,85 @@ class Transport {
 
     virtual const Capabilities capabilities() const { return caps; }
 
-    virtual Status allocateSubBatch(SubBatchRef &batch, size_t max_size) {
+    virtual Status allocateSubBatch(SubBatchRef& batch, size_t max_size) {
         return Status::NotImplemented(
             "allocateSubBatch not implemented" LOC_MARK);
     }
 
-    virtual Status freeSubBatch(SubBatchRef &batch) {
+    virtual Status freeSubBatch(SubBatchRef& batch) {
         return Status::NotImplemented("freeSubBatch not implemented" LOC_MARK);
     }
 
     virtual Status submitTransferTasks(
-        SubBatchRef batch, const std::vector<Request> &request_list) {
+        SubBatchRef batch, const std::vector<Request>& request_list) {
         return Status::NotImplemented(
             "submitTransferTasks not implemented" LOC_MARK);
     }
 
     virtual Status getTransferStatus(SubBatchRef batch, int task_id,
-                                     TransferStatus &status) {
+                                     TransferStatus& status) {
         return Status::NotImplemented(
             "getTransferStatus not implemented" LOC_MARK);
     }
 
-    virtual Status allocateLocalMemory(void **addr, size_t size,
-                                       MemoryOptions &options) {
+    // Cancellation is best effort: implementations must prevent work that has
+    // not reached the device from being submitted, but work already posted to
+    // a device may still complete. Callers must continue polling until the
+    // task reaches a terminal state.
+    virtual bool supportsCancellation() const { return false; }
+
+    virtual Status cancelTransferTask(SubBatchRef batch, int task_id) {
+        return Status::NotImplemented(
+            "cancelTransferTask not implemented" LOC_MARK);
+    }
+
+    virtual Status allocateLocalMemory(void** addr, size_t size,
+                                       MemoryOptions& options) {
         return Platform::getLoader().allocate(addr, size, options);
     }
 
-    virtual Status freeLocalMemory(void *addr, size_t size) {
+    virtual Status freeLocalMemory(void* addr, size_t size) {
         return Platform::getLoader().free(addr, size);
     }
 
     // Pre-registration warm-up that pins pages before NUMA probing.
     // Returns true if pages were successfully pinned (caller may skip
     // prefault). Default: no-op, returns false.
-    virtual bool warmupMemory(void *addr, size_t length) { return false; }
+    virtual bool warmupMemory(void* addr, size_t length) { return false; }
 
-    virtual Status addMemoryBuffer(BufferDesc &desc,
-                                   const MemoryOptions &options) {
+    virtual Status addMemoryBuffer(BufferDesc& desc,
+                                   const MemoryOptions& options) {
         return Status::NotImplemented(
             "addMemoryBuffer not implemented" LOC_MARK);
     }
 
-    virtual Status addMemoryBuffer(std::vector<BufferDesc> &desc_list,
-                                   const MemoryOptions &options) {
-        for (auto &desc : desc_list) {
+    virtual Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
+                                   const MemoryOptions& options) {
+        for (auto& desc : desc_list) {
             CHECK_STATUS(addMemoryBuffer(desc, options));
         }
         return Status::OK();
     }
 
-    virtual Status removeMemoryBuffer(BufferDesc &desc) {
+    virtual Status removeMemoryBuffer(BufferDesc& desc) {
         return Status::NotImplemented(
             "removeMemoryBuffer not implemented" LOC_MARK);
     }
 
     virtual bool supportNotification() const { return false; }
 
     virtual Status sendNotification(SegmentID target_id,
-                                    const Notification &notify) {
+                                    const Notification& notify) {
         return Status::NotImplemented(
             "sendNotification not implemented" LOC_MARK);
     }
 
-    virtual Status receiveNotification(std::vector<Notification> &notify_list) {
+    virtual Status receiveNotification(std::vector<Notification>& notify_list) {
         return Status::NotImplemented(
             "receiveNotification not implemented" LOC_MARK);
     }
 
-    virtual const char *getName() const { return "<generic>"; }
+    virtual const char* getName() const { return "<generic>"; }
 
    protected:
     Capabilities caps;

mooncake-transfer-engine/tent/include/tent/transfer_engine.h

@@ -174,6 +174,9 @@ void tent_free_notifs(tent_notifi_info* info);
 int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id,
                      size_t task_id, tent_status_t* status);
 
+int tent_cancel_task(tent_engine_t engine, tent_batch_id_t batch_id,
+                     size_t task_id);
+
 int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id,
                         tent_status_t* status);
 
@@ -299,6 +302,11 @@ class TransferEngine {
                           const std::vector<Request>& request_list,
                           const Notification& notifi);
 
+    // Best-effort task cancellation. Work that has not reached the transport
+    // is prevented from being submitted. Device work already posted may still
+    // complete, so callers must continue polling for a terminal status.
+    Status cancelTransfer(BatchID batch_id, size_t task_id);
+
     Status sendNotification(SegmentID target_id, const Notification& notifi);
 
     Status receiveNotification(std::vector<Notification>& notifi_list);
@@ -328,4 +336,4 @@ class TransferEngine {
 }  // namespace mooncake
 #endif
 
-#endif
+#endif

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

@@ -79,6 +79,10 @@ class RdmaTransport : public Transport {
     virtual Status getTransferStatus(SubBatchRef batch, int task_id,
                                      TransferStatus& status);
 
+    bool supportsCancellation() const override { return true; }
+
+    Status cancelTransferTask(SubBatchRef batch, int task_id) override;
+
     virtual Status addMemoryBuffer(BufferDesc& desc,
                                    const MemoryOptions& options);
 

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

@@ -57,10 +57,14 @@ struct RdmaTask {
     std::string qp_pool;
     volatile TransferStatusEnum status_word;
     volatile size_t transferred_bytes;
-    volatile int success_slices;
-    volatile int resolved_slices;
+    std::atomic<int> success_slices{0};
+    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.
+    std::atomic<bool> cancel_requested{false};
+
     // Reference counting for UAF protection
     std::atomic<int> ref_count{0};
 
@@ -92,6 +96,9 @@ struct RdmaSlice {
     int qp_index = 0;
     int retry_count = 0;
     bool failed = false;
+    // True while DeviceSelector accounts this slice against source_dev_id.
+    // The worker clears it exactly once on completion, failure, or cancel.
+    bool quota_charged = false;
     uint64_t enqueue_ts = 0;
     uint64_t submit_ts = 0;
     // Non-owning pointer to the per-worker RailMonitor for this slice's
@@ -111,15 +118,18 @@ static inline void updateSliceStatus(RdmaSlice* slice,
     if (!__sync_bool_compare_and_swap(&slice->word, PENDING, status)) return;
     if (status == COMPLETED) {
         __sync_fetch_and_add(&task->transferred_bytes, slice->length);
-        __sync_fetch_and_add(&task->success_slices, 1);
+        task->success_slices.fetch_add(1, std::memory_order_acq_rel);
     } else {
         __sync_bool_compare_and_swap(&task->first_error, PENDING, status);
     }
-    int resolved = __sync_add_and_fetch(&task->resolved_slices, 1);
+    int resolved =
+        task->resolved_slices.fetch_add(1, std::memory_order_acq_rel) + 1;
     if (resolved >= task->num_slices) {
-        TransferStatusEnum final_st = (task->success_slices == task->num_slices)
-                                          ? COMPLETED
-                                          : task->first_error;
+        TransferStatusEnum final_st =
+            (task->success_slices.load(std::memory_order_acquire) ==
+             task->num_slices)
+                ? COMPLETED
+                : task->first_error;
         if (final_st == PENDING) final_st = FAILED;
         __sync_bool_compare_and_swap(&task->status_word, PENDING, final_st);
     }
@@ -129,4 +139,4 @@ static inline void updateSliceStatus(RdmaSlice* slice,
 }  // namespace tent
 }  // namespace mooncake
 
-#endif  // TENT_SLICE_H
+#endif  // TENT_SLICE_H