bugfix: safe batch retirement and bounded abandon.

mooncake-transfer-engine/include/multi_transport.h

@@ -30,6 +30,7 @@ class MultiTransport {
     using TransferRequest = Transport::TransferRequest;
     using TransferStatus = Transport::TransferStatus;
     using BatchDesc = Transport::BatchDesc;
+    using AbandonOutcome = Transport::AbandonOutcome;
 
     MultiTransport(std::shared_ptr<TransferMetadata> metadata,
                    std::string &local_server_name);
@@ -54,6 +55,24 @@ class MultiTransport {
 
     Status getBatchTransferStatus(BatchID batch_id, TransferStatus &status);
 
+    // Drives all unfinished tasks and reports whether borrowed buffers and
+    // callback state can be released. retired is valid even on an error Status;
+    // transfer/query failure alone never authorizes resource reuse.
+    Status getBatchRetirementStatus(BatchID batch_id, bool &retired);
+
+    // Same as above plus a safety verdict: `quiescent` is false when the
+    // terminal state came from a give-up (timeout/abort not confirmed), so the
+    // caller must not release the batch's borrowed buffers.
+    Status getBatchRetirementStatus(BatchID batch_id, bool &retired,
+                                    bool &quiescent);
+
+    // Bounded attempt to stop a stalled batch. Only kQuiescentRecoverable
+    // publishes the terminal counts that let the caller free the batch;
+    // kQuiescentDeviceFault keeps the batch pending on purpose so the caller
+    // shuts the engine down instead of releasing buffers it cannot trust.
+    Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout,
+                        AbandonOutcome &outcome);
+
     Transport *installTransport(const std::string &proto,
                                 std::shared_ptr<Topology> topo);
 

mooncake-transfer-engine/include/transfer_engine.h

@@ -165,6 +165,21 @@ class TransferEngine {
 
     Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status);
 
+    // Drives all unfinished tasks and reports whether borrowed buffers and
+    // callback state can be released. retired is valid even on an error Status;
+    // transfer/query failure alone never authorizes resource reuse.
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired);
+
+    // Same as above plus a safety verdict: `quiescent` is false when the
+    // terminal state came from a give-up, so borrowed buffers must not be
+    // released. Callers must not free a retired but non-quiescent batch.
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired,
+                                    bool& quiescent);
+
+    // Bounded attempt to stop a stalled batch. See Transport::AbandonOutcome.
+    Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout,
+                        Transport::AbandonOutcome& outcome);
+
     Status getNicLoadStats(std::vector<NicLoadStats>& stats) const;
 
     Transport* getTransport(const std::string& proto);

mooncake-transfer-engine/include/transfer_engine_impl.h

@@ -52,6 +52,7 @@ using SegmentHandle = Transport::SegmentHandle;
 using SegmentID = Transport::SegmentID;
 using BatchID = Transport::BatchID;
 using BufferEntry = Transport::BufferEntry;
+using AbandonOutcome = Transport::AbandonOutcome;
 
 #ifdef ENABLE_MULTI_PROTOCOL
 using RegisteredBuffer = TransferEngine::RegisteredBuffer;
@@ -318,6 +319,25 @@ class TransferEngineImpl {
         return result;
     }
 
+    // Drives unfinished classic tasks and reports whether borrowed buffers
+    // and completion callbacks may be released. The public TransferEngine
+    // facade forwards here, so this middle layer must not be skipped even
+    // though MultiTransport already implements the check.
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired) {
+        return multi_transports_->getBatchRetirementStatus(batch_id, retired);
+    }
+
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired,
+                                    bool& quiescent) {
+        return multi_transports_->getBatchRetirementStatus(batch_id, retired,
+                                                           quiescent);
+    }
+
+    Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout,
+                        AbandonOutcome& outcome) {
+        return multi_transports_->abandonBatch(batch_id, timeout, outcome);
+    }
+
     Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status,
                                   bool skip_metrics = false) {
         Status result =

mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h

@@ -50,6 +50,10 @@ class AscendDirectTransport : public Transport {
     Status getTransferStatus(BatchID batch_id, size_t task_id,
                              TransferStatus &status) override;
 
+    // False when any slice of the batch was failed by a give-up path whose
+    // in-flight ADXL operation could not be cancelled or drained.
+    bool isBatchQuiescent(BatchID batch_id) override;
+
     int install(std::string &local_server_name,
                 std::shared_ptr<TransferMetadata> meta,
                 std::shared_ptr<Topology> topo) override;

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

@@ -27,6 +27,7 @@
 #include <queue>
 #include <string>
 #include <atomic>
+#include <chrono>
 #include <functional>
 #include <mutex>
 #include <condition_variable>
@@ -80,6 +81,16 @@ class Transport {
         FAILED
     };
 
+    // Result of a bounded attempt to stop a stalled batch. kQuiescent* means
+    // the transport confirmed that it will not access the batch's borrowed
+    // buffers again; kNotQuiescent means no such proof exists and the caller
+    // must keep waiting (returning without it reintroduces a use-after-free).
+    enum class AbandonOutcome {
+        kQuiescentRecoverable,
+        kQuiescentDeviceFault,
+        kNotQuiescent,
+    };
+
     struct TransferStatus {
         TransferStatusEnum s;
         size_t transferred_bytes;
@@ -178,6 +189,10 @@ class Transport {
                 void *handle;
                 int64_t start_time;
                 int32_t engine_id;
+                // Set when the terminal state came from a timeout/abort that
+                // could not prove the transfer stopped. Such a slice must never
+                // be treated as safe to release.
+                bool unsafe_terminal;
             } ascend_direct;
             struct {
                 uint64_t dest_addr;
@@ -186,24 +201,42 @@ class Transport {
 
        public:
         void markSuccess() {
+            CompletionCallbackPin pin(task);
             status = Slice::SUCCESS;
             __atomic_fetch_add(&task->transferred_bytes, length,
                                __ATOMIC_RELAXED);
-            __atomic_fetch_add(&task->success_slice_count, 1, __ATOMIC_RELAXED);
+            __atomic_fetch_add(&task->success_slice_count, 1, __ATOMIC_ACQ_REL);
 
             check_batch_completion(false);
         }
 
         void markFailed() {
+            CompletionCallbackPin pin(task);
             status = Slice::FAILED;
-            __atomic_fetch_add(&task->failed_slice_count, 1, __ATOMIC_RELAXED);
+            __atomic_fetch_add(&task->failed_slice_count, 1, __ATOMIC_ACQ_REL);
 
             check_batch_completion(true);
         }
 
         volatile int64_t ts;
 
        private:
+        class CompletionCallbackPin {
+           public:
+            explicit CompletionCallbackPin(TransferTask *task)
+                : batch_(toBatchDesc(task->batch_id)) {
+                batch_.active_completion_callbacks.fetch_add(
+                    1, std::memory_order_acq_rel);
+            }
+            ~CompletionCallbackPin() {
+                batch_.active_completion_callbacks.fetch_sub(
+                    1, std::memory_order_release);
+            }
+
+           private:
+            BatchDesc &batch_;
+        };
+
         inline void check_batch_completion(bool is_failed) {
 #ifdef USE_EVENT_DRIVEN_COMPLETION
             auto &batch_desc = toBatchDesc(task->batch_id);
@@ -346,6 +379,9 @@ class Transport {
         void *context;  // for transport implementers.
         int64_t start_timestamp;
 
+        // A terminal count does not mean its completion callback has returned.
+        std::atomic<uint64_t> active_completion_callbacks{0};
+
         // Track batch progress and notifies waiters
         std::atomic<bool> has_failure{false};
         std::atomic<bool> is_finished{
@@ -371,6 +407,18 @@ class Transport {
     /// @brief Free an allocated batch.
     virtual Status freeBatchID(BatchID batch_id);
 
+    /// @brief Stop a stalled batch and report whether its buffers are safe to
+    /// release. Implementations must only report kQuiescent* when they can
+    /// prove that no transfer access is in flight any more.
+    virtual Status abandonBatch(BatchID batch_id,
+                                std::chrono::milliseconds timeout,
+                                AbandonOutcome& outcome);
+
+    /// @brief Whether a retired batch is provably idle (no in-flight access).
+    /// Transports that publish terminal state on give-up/timeout must override
+    /// this and return false for those batches.
+    virtual bool isBatchQuiescent(BatchID batch_id);
+
     /// @brief Submit a batch of transfer requests to the batch.
     /// @return The number of successfully submitted transfers on success. If
     /// that number is less than nr, errno is set.

mooncake-transfer-engine/src/multi_transport.cpp

@@ -96,6 +96,13 @@ MultiTransport::BatchID MultiTransport::allocateBatchID(size_t batch_size) {
 }
 
 Status MultiTransport::freeBatchID(BatchID batch_id) {
+    bool retired = false;
+    auto status = getBatchRetirementStatus(batch_id, retired);
+    if (!retired) {
+        return Status::BatchBusy("Batch tasks or callbacks are pending: " +
+                                 status.ToString());
+    }
+    // A transfer/query error does not prevent cleanup after retirement.
     auto& batch_desc = *((BatchDesc*)(batch_id));
     const size_t task_count = batch_desc.task_list.size();
     for (size_t task_id = 0; task_id < task_count; task_id++) {
@@ -254,7 +261,7 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
         } else {
             status.s = Transport::TransferStatusEnum::COMPLETED;
         }
-        task.is_finished = true;
+        __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE);
     } else {
         if (checkSliceTimeout(task)) {
             status.s = Transport::TransferStatusEnum::TIMEOUT;
@@ -265,6 +272,136 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
     return Status::OK();
 }
 
+Status MultiTransport::getBatchRetirementStatus(BatchID batch_id,
+                                                bool& retired) {
+    auto& batch = Transport::toBatchDesc(batch_id);
+    retired = true;
+    Status result = Status::OK();
+    for (size_t index = 0; index < batch.task_list.size(); ++index) {
+        auto& task = batch.task_list[index];
+        const uint64_t count =
+            __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE);
+        uint64_t succeeded =
+            __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE);
+        uint64_t failed =
+            __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE);
+        if (succeeded > count || failed != count - succeeded) {
+            TransferStatus status;
+            auto queried = getTransferStatus(batch_id, index, status);
+            if (result.ok() && !queried.ok()) {
+                result = queried;
+            } else if (result.ok() &&
+                       (status.s == Transport::TransferStatusEnum::FAILED ||
+                        status.s == Transport::TransferStatusEnum::TIMEOUT ||
+                        status.s == Transport::TransferStatusEnum::CANCELED)) {
+                result =
+                    Status::Context("Transfer task failed, state=" +
+                                    std::to_string(static_cast<int>(status.s)));
+            }
+            succeeded =
+                __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE);
+            failed =
+                __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE);
+        }
+        if (succeeded > count || failed != count - succeeded) retired = false;
+        if (result.ok() && failed != 0) {
+            result = Status::Context("Transfer task failed: " +
+                                     std::to_string(index));
+        }
+    }
+    retired = retired && batch.active_completion_callbacks.load(
+                             std::memory_order_acquire) == 0;
+    if (retired) {
+        for (auto& task : batch.task_list) {
+            __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE);
+        }
+    }
+    return result;
+}
+
+Status MultiTransport::getBatchRetirementStatus(BatchID batch_id,
+                                                bool &retired,
+                                                bool &quiescent) {
+    const Status status = getBatchRetirementStatus(batch_id, retired);
+    quiescent = retired;
+    if (retired) {
+        auto &batch = Transport::toBatchDesc(batch_id);
+        for (auto &task : batch.task_list) {
+            if (task.transport_ == nullptr) continue;
+            if (!task.transport_->isBatchQuiescent(batch_id)) {
+                quiescent = false;
+                break;
+            }
+        }
+    }
+    return status;
+}
+
+Status MultiTransport::abandonBatch(BatchID batch_id,
+                                    std::chrono::milliseconds timeout,
+                                    AbandonOutcome &outcome) {
+    auto &batch = Transport::toBatchDesc(batch_id);
+    outcome = AbandonOutcome::kQuiescentRecoverable;
+
+    // Collect the transports that still own unfinished tasks of this batch.
+    std::vector<Transport *> pending;
+    for (auto &task : batch.task_list) {
+        const uint64_t count =
+            __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE);
+        const uint64_t succeeded =
+            __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE);
+        const uint64_t failed =
+            __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE);
+        if (succeeded + failed == count) continue;
+        if (task.transport_ == nullptr) {
+            // No transport to ask: the batch can never be proven quiescent.
+            outcome = AbandonOutcome::kNotQuiescent;
+            return Status::OK();
+        }
+        if (std::find(pending.begin(), pending.end(), task.transport_) ==
+            pending.end()) {
+            pending.push_back(task.transport_);
+        }
+    }
+
+    Status result = Status::OK();
+    for (auto *transport : pending) {
+        AbandonOutcome local = AbandonOutcome::kNotQuiescent;
+        const Status status = transport->abandonBatch(batch_id, timeout, local);
+        if (!status.ok()) {
+            // The transport cannot stop this batch: no safety proof.
+            outcome = AbandonOutcome::kNotQuiescent;
+            return status;
+        }
+        if (local == AbandonOutcome::kQuiescentDeviceFault) {
+            outcome = local;
+        } else if (local == AbandonOutcome::kNotQuiescent &&
+                   outcome == AbandonOutcome::kQuiescentRecoverable) {
+            outcome = local;
+        }
+    }
+
+    if (outcome == AbandonOutcome::kQuiescentRecoverable) {
+        // The transport proved that the batch is stopped, so publishing the
+        // terminal counts is safe and lets the caller free the batch.
+        for (auto &task : batch.task_list) {
+            const uint64_t count =
+                __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE);
+            const uint64_t succeeded =
+                __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE);
+            const uint64_t failed =
+                __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE);
+            if (succeeded + failed < count) {
+                __atomic_add_fetch(&task.failed_slice_count,
+                                   count - succeeded - failed,
+                                   __ATOMIC_ACQ_REL);
+            }
+            __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE);
+        }
+    }
+    return result;
+}
+
 Status MultiTransport::getBatchTransferStatus(BatchID batch_id,
                                               TransferStatus& status) {
     auto& batch_desc = *((BatchDesc*)(batch_id));