[Bugfix] Preserve deferred batch cleanup ownership after query errors

mooncake-transfer-engine/include/multi_transport.h

@@ -15,6 +15,7 @@
 #ifndef MULTI_TRANSPORT_H_
 #define MULTI_TRANSPORT_H_
 
+#include <chrono>
 #include <condition_variable>
 #include <functional>
 #include <map>
@@ -55,10 +56,18 @@ class MultiTransport {
 
     BatchID allocateBatchID(size_t batch_size);
 
-    // An OK or BatchCleanupDeferred return invalidates batch_id. A deferred
-    // cleanup means in-flight transport work may still own transfer buffers.
+    // An OK or BatchCleanupDeferred return invalidates batch_id. Deferred
+    // transfers batch cleanup responsibility, not ownership of user buffers.
+    // Cleanup status is independent of transfer success: do not turn an
+    // otherwise successful transfer into a failure solely due to deferral.
+    // In-flight work may still access borrowed buffers; keep them alive until
+    // retirement is independently established, even after handle handoff.
+    // A pre-delete hook (when installed by TransferEngineImpl) releases
+    // descriptor side state; its invocation is not a buffer-retirement signal.
     Status freeBatchID(BatchID batch_id);
 
+    Transport::DeferredCleanupStats getDeferredCleanupStats();
+
     Status submitTransfer(BatchID batch_id,
                           const std::vector<TransferRequest> &entries);
 
@@ -80,6 +89,11 @@ class MultiTransport {
 
     Status getBatchTransferStatus(BatchID batch_id, TransferStatus &status);
 
+    // Observe retirement without releasing or transferring batch ownership.
+    // A non-OK result reports a query/transfer failure; only retired permits
+    // the caller to stop retaining borrowed requests and transfer buffers.
+    Status getBatchRetirementStatus(BatchID batch_id, bool &retired);
+
     Transport *installTransport(const std::string &proto,
                                 std::shared_ptr<Topology> topo);
 
@@ -134,8 +148,16 @@ class MultiTransport {
     Status selectTransport(const TransferRequest &entry, Transport *&transport,
                            bool *allows_reuse = nullptr);
 
+    // The hook cleans descriptor-scoped state before deletion. It is not a
+    // user-buffer retirement callback; observe retirement before free when
+    // borrowed requests/buffers must be reused.
+    Transport::DeferredCleanupStats collectDeferredCleanupStats(
+        std::chrono::steady_clock::time_point now) const;
+
     Status tryFreeBatchID(BatchID batch_id,
-                          const std::function<void()> &before_delete);
+                          const std::function<void()> &before_delete,
+                          bool allow_completion_fallback = false,
+                          bool poll_transport = true);
 
     void deferredCleanupLoop();
 
@@ -154,8 +176,20 @@ class MultiTransport {
 
     std::mutex deferred_cleanup_mutex_;
     std::condition_variable deferred_cleanup_cv_;
-    std::unordered_map<BatchID, std::function<void()>>
-        deferred_cleanup_batches_;
+    struct DeferredCleanupEntry {
+        explicit DeferredCleanupEntry(const std::function<void()> &callback)
+            : before_delete(callback) {}
+
+        std::function<void()> before_delete;
+        std::chrono::steady_clock::time_point enqueued_at =
+            std::chrono::steady_clock::now();
+        Status last_error;
+        bool completion_only = false;
+        std::chrono::steady_clock::time_point next_poll{};
+        std::chrono::milliseconds retry_delay{100};
+    };
+    std::unordered_map<BatchID, DeferredCleanupEntry> deferred_cleanup_batches_;
+    std::chrono::steady_clock::time_point next_deferred_cleanup_warning_{};
     bool stop_deferred_cleanup_ = false;
     std::thread deferred_cleanup_thread_;
 };

mooncake-transfer-engine/include/transfer_engine.h

@@ -268,8 +268,12 @@ class TransferEngine {
 
     BatchID allocateBatchID(size_t batch_size);
 
-    // An OK or BatchCleanupDeferred return invalidates batch_id. A deferred
-    // cleanup means in-flight transport work may still own transfer buffers.
+    // An OK or BatchCleanupDeferred return invalidates batch_id. Deferred
+    // transfers batch cleanup responsibility, not ownership of user buffers.
+    // Cleanup status is independent of transfer success: do not turn an
+    // otherwise successful transfer into a failure solely due to deferral.
+    // In-flight work may still access borrowed buffers; keep them alive until
+    // retirement is independently established, even after handle handoff.
     Status freeBatchID(BatchID batch_id);
 
     int getNotifies(std::vector<TransferMetadata::NotifyDesc>& notifies);
@@ -289,6 +293,27 @@ class TransferEngine {
 
     Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status);
 
+#ifndef USE_TENT
+    // After submission returns, poll a caller-owned, valid batch until retired
+    // is true. All tasks and completion callbacks must have stopped accessing
+    // borrowed requests/buffers. This does not free or invalidate batch_id.
+    // A non-OK result reports an observed query/transfer failure independently
+    // of retirement. TIMEOUT, CANCELED, INVALID, or another error alone never
+    // authorizes buffer reuse; rely on retired for the independent proof.
+    // Do not submit more work, free, or query this batch concurrently. As with
+    // transfer completion, transports must publish terminal state after I/O
+    // retirement (including failure paths). TIMEOUT is not terminal: tasks may
+    // subsequently complete/fail, but the timeout itself never permits reuse.
+    // Never query after free returns OK or BatchCleanupDeferred: the handle is
+    // invalid, possibly already deleted or reused. This API does not validate
+    // expired pointer-derived IDs.
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired);
+
+    // Current reaper-owned backlog. This is not a count of all live batches,
+    // leaked allocations, historical IDs, or bytes in borrowed device buffers.
+    Transport::DeferredCleanupStats getDeferredCleanupStats();
+#endif
+
     Status getNicLoadStats(std::vector<NicLoadStats>& stats) const;
 
     Transport* getTransport(const std::string& proto);

mooncake-transfer-engine/include/transfer_engine_impl.h

@@ -276,9 +276,7 @@ class TransferEngineImpl {
 
             // Only record metrics for successful completions
             if (status.s == TransferStatusEnum::COMPLETED) {
-                if (status.transferred_bytes > 0) {
-                    transferred_bytes_counter_.inc(status.transferred_bytes);
-                }
+                recordTransferredBytes(task, status.transferred_bytes);
                 auto now = std::chrono::steady_clock::now();
                 auto duration =
                     std::chrono::duration_cast<std::chrono::microseconds>(
@@ -319,8 +317,14 @@ class TransferEngineImpl {
 #ifdef WITH_METRICS
         if (metrics_enabled_ && !skip_metrics && result.ok() &&
             status.s == TransferStatusEnum::COMPLETED) {
-            if (status.transferred_bytes > 0) {
-                transferred_bytes_counter_.inc(status.transferred_bytes);
+            // The cached aggregate is built from these task byte counts.
+            // Share the per-task claim with getTransferStatus so mixed and
+            // repeated observers cannot count the same successful bytes twice.
+            auto& batch = Transport::toBatchDesc(batch_id);
+            for (auto& task : batch.task_list) {
+                recordTransferredBytes(
+                    task,
+                    __atomic_load_n(&task.transferred_bytes, __ATOMIC_ACQUIRE));
             }
         }
 #endif
@@ -342,6 +346,24 @@ class TransferEngineImpl {
         return result;
     }
 
+    Status getBatchRetirementStatus(BatchID batch_id, bool& retired) {
+        Status result =
+            multi_transports_->getBatchRetirementStatus(batch_id, retired);
+        if (retired && result.ok()) {
+            // Preserve completion metrics/notifications previously driven by
+            // the aggregate status polling path before the caller frees.
+            TransferStatus status;
+            Status completed = getBatchTransferStatus(batch_id, status);
+            if (!completed.ok()) return completed;
+        }
+        return result;
+    }
+
+    Transport::DeferredCleanupStats getDeferredCleanupStats() {
+        if (!multi_transports_) return {};
+        return multi_transports_->getDeferredCleanupStats();
+    }
+
     Transport* getTransport(const std::string& proto) {
         if (!multi_transports_) return nullptr;
         return multi_transports_->getTransport(proto);
@@ -527,6 +549,15 @@ class TransferEngineImpl {
     bool metrics_enabled_{false};
     uint64_t metrics_interval_seconds_{5};
 
+    void recordTransferredBytes(Transport::TransferTask& task,
+                                uint64_t transferred_bytes) {
+        if (transferred_bytes > 0 &&
+            !__atomic_exchange_n(&task.log_bytes_recorded, true,
+                                 __ATOMIC_ACQ_REL)) {
+            transferred_bytes_counter_.inc(transferred_bytes);
+        }
+    }
+
     // Helper methods for metrics reporting thread management
     void InitializeMetricsConfig();
     void StartMetricsReportingThread();

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

@@ -504,6 +504,9 @@ class Transport {
         volatile uint64_t completion_state = Slice::kCompletionSealBit;
         volatile uint64_t transferred_bytes = 0;
         volatile bool is_finished = false;
+        // Set by MultiTransport only after a successful terminal status query.
+        // Legacy is_finished alone is not proof after an API query error.
+        volatile bool retirement_observed = false;
         uint64_t total_bytes = 0;
         BatchID batch_id = 0;
 
@@ -522,6 +525,9 @@ class Transport {
         // make the task non-movable).
         volatile bool export_metrics_recorded = false;
         volatile bool log_metrics_recorded = false;
+        // Task and aggregate/retirement observers share byte accounting.
+        // Keep latency publication independent from the first byte observer.
+        volatile bool log_bytes_recorded = false;
         // Set by the event-driven completion callback when it finishes the
         // task: steady_clock ticks (else 0) and COMPLETED or FAILED (else
         // WAITING). Atomic: a poll may read them while the callback writes.
@@ -546,6 +552,13 @@ class Transport {
         }
     };
 
+    // A live snapshot of reaper-owned descriptors, not historical BatchIDs.
+    struct DeferredCleanupStats {
+        uint64_t pending_batches = 0;
+        uint64_t query_error_batches = 0;
+        uint64_t oldest_age_ms = 0;
+    };
+
     struct BatchDesc {
         BatchID id;
         size_t batch_size;