[TE] Add ProgressWorker skeleton (#2199)

* [TE] Add explicit progressBatch API

* [TE] Add progress worker skeleton

Author: ZhenyuePan

mooncake-transfer-engine/tent/config/transfer-engine.json

@@ -12,6 +12,7 @@
     "log_level": "warning",
     "max_failover_attempts": 3,
     "enable_auto_failover_on_poll": true,
+    "enable_progress_worker": false,
     "metrics": {
         "enabled": true,
         "http_port": 9100,

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

@@ -0,0 +1,74 @@
+// Copyright 2024 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.
+
+#ifndef PROGRESS_WORKER_H_
+#define PROGRESS_WORKER_H_
+
+#include <atomic>
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <thread>
+#include <unordered_set>
+
+#include "tent/common/types.h"
+
+namespace mooncake {
+namespace tent {
+
+class TransferEngineImpl;
+
+// Event-driven progress worker for issue #2116. When the engine is configured
+// with enable_progress_worker=true, transports (or test hooks) call
+// notifyBatchMaybeReady to wake this worker, which then drives one
+// progressBatch step per notification. This decouples failover/resubmit from
+// the caller polling loop, so integrators that turn off
+// enable_auto_failover_on_poll do not need to spin a polling thread of their
+// own to keep failover progressing.
+class ProgressWorker {
+   public:
+    explicit ProgressWorker(TransferEngineImpl* impl);
+    ~ProgressWorker();
+
+    ProgressWorker(const ProgressWorker&) = delete;
+    ProgressWorker& operator=(const ProgressWorker&) = delete;
+
+    void start();
+
+    // Idempotent. Signals the worker thread to exit and joins it. After stop()
+    // returns, notifyBatchMaybeReady becomes a no-op.
+    void stop();
+
+    // Safe from any thread. De-duplicates: enqueueing a batch that is already
+    // queued is a no-op. No-op if the worker has been stopped or never
+    // started.
+    void notifyBatchMaybeReady(BatchID batch_id);
+
+   private:
+    void runner();
+
+    TransferEngineImpl* impl_;
+    std::atomic<bool> running_{false};
+    std::thread thread_;
+
+    std::mutex mu_;
+    std::condition_variable cv_;
+    std::unordered_set<BatchID> queued_;
+    std::deque<BatchID> order_;
+};
+
+}  // namespace tent
+}  // namespace mooncake
+
+#endif  // PROGRESS_WORKER_H_

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

@@ -46,6 +46,7 @@ class ControlService;
 class SegmentTracker;
 class Platform;
 class ProxyManager;
+class ProgressWorker;
 
 struct TaskInfo {
     TransportType type{UNSPEC};
@@ -152,6 +153,8 @@ class TransferEngineImpl {
 
     Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
 
+    Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
+
     Status waitTransferCompletion(BatchID batch_id);
 
     Status transferSync(const std::vector<Request>& request_list);
@@ -172,6 +175,11 @@ class TransferEngineImpl {
         }
     }
 
+    // Wake the optional event-driven progress worker for `batch_id`. No-op if
+    // enable_progress_worker is false. Currently used by test/integration
+    // hooks; transports will be migrated to call this in a follow-up PR.
+    void notifyBatchMaybeReady(BatchID batch_id);
+
    private:
     Status construct();
 
@@ -189,8 +197,15 @@ class TransferEngineImpl {
 
     Status resubmitTransferTask(Batch* batch, size_t task_id);
 
-    void updateTaskStatusFromPoll(Batch* batch, size_t task_id,
-                                  TransferStatus& task_status);
+    Status pollTaskStatus(Batch* batch, size_t task_id,
+                          TransferStatus& task_status);
+
+    void updateTaskStatusAfterPoll(Batch* batch, size_t task_id,
+                                   TransferStatus& task_status,
+                                   bool allow_failover);
+
+    Status getBatchStatus(BatchID batch_id, TransferStatus& overall_status,
+                          bool allow_failover);
 
     SelectionResult resolveTransport(const Request& req, int transport_index,
                                      bool invalidate_on_fail = true);
@@ -244,6 +259,15 @@ class TransferEngineImpl {
     bool merge_requests_;
     int max_failover_attempts_{3};
     bool enable_auto_failover_on_poll_{true};
+    bool enable_progress_worker_{false};
+
+    // Guards alive_batches_ and serializes pollTaskStatus /
+    // updateTaskStatusAfterPoll / lazyFreeBatch against the optional
+    // ProgressWorker thread. Recursive because freeBatch -> lazyFreeBatch ->
+    // getTransferStatus can re-enter on the same thread. See issue #2116.
+    std::recursive_mutex progress_mutex_;
+    std::unordered_set<BatchID> alive_batches_;
+    std::unique_ptr<ProgressWorker> progress_worker_;
 };
 }  // namespace tent
 }  // namespace mooncake

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

@@ -308,6 +308,14 @@ class TransferEngine {
 
     Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
 
+    // Drive one progress step on a batch and return its aggregated status.
+    // Unlike getTransferStatus, this always allows internal failover/resubmit
+    // regardless of enable_auto_failover_on_poll. The call is non-blocking and
+    // performs at most one state-machine step per task; callers that want to
+    // wait for completion must invoke it in a loop. PENDING means "make
+    // progress later"; terminal states (COMPLETED/FAILED) will not be revived.
+    Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
+
    private:
     std::unique_ptr<TransferEngineImpl> impl_;
 };

mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp

@@ -0,0 +1,81 @@
+// Copyright 2024 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.
+
+#include "tent/runtime/progress_worker.h"
+
+#include "tent/common/status.h"
+#include "tent/runtime/transfer_engine_impl.h"
+
+namespace mooncake {
+namespace tent {
+
+ProgressWorker::ProgressWorker(TransferEngineImpl* impl) : impl_(impl) {}
+
+ProgressWorker::~ProgressWorker() { stop(); }
+
+void ProgressWorker::start() {
+    if (running_.exchange(true, std::memory_order_acq_rel)) return;
+    thread_ = std::thread(&ProgressWorker::runner, this);
+}
+
+void ProgressWorker::stop() {
+    if (!running_.exchange(false, std::memory_order_acq_rel)) return;
+    {
+        std::lock_guard<std::mutex> lk(mu_);
+        // Drop pending work; outstanding batches will be reaped via the
+        // user thread's freeBatch path.
+        order_.clear();
+        queued_.clear();
+    }
+    cv_.notify_all();
+    if (thread_.joinable()) thread_.join();
+}
+
+void ProgressWorker::notifyBatchMaybeReady(BatchID batch_id) {
+    if (!batch_id) return;
+    if (!running_.load(std::memory_order_acquire)) return;
+    {
+        std::lock_guard<std::mutex> lk(mu_);
+        if (!queued_.insert(batch_id).second) return;
+        order_.push_back(batch_id);
+    }
+    cv_.notify_one();
+}
+
+void ProgressWorker::runner() {
+    while (true) {
+        BatchID batch_id = 0;
+        {
+            std::unique_lock<std::mutex> lk(mu_);
+            cv_.wait(lk, [&] {
+                return !running_.load(std::memory_order_acquire) ||
+                       !order_.empty();
+            });
+            if (!running_.load(std::memory_order_acquire)) return;
+            batch_id = order_.front();
+            order_.pop_front();
+            queued_.erase(batch_id);
+        }
+        // progressBatch acquires the engine's progress_mutex_ and silently
+        // returns InvalidArgument if the batch was freed before we got here.
+        // PENDING means "kick again later"; the next notify wakes us up.
+        // Terminal states leave the batch alone — freeBatch on the user
+        // thread is responsible for reclamation.
+        TransferStatus s;
+        (void)impl_->progressBatch(batch_id, s);
+    }
+}
+
+}  // namespace tent
+}  // namespace mooncake

mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp

@@ -385,8 +385,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
 
             case StageState::INFLIGHT: {
                 TransferStatus xfer_status;
-                CHECK_STATUS(
-                    impl_->getTransferStatus(chunk.batch, xfer_status));
+                CHECK_STATUS(impl_->progressBatch(chunk.batch, xfer_status));
                 if (xfer_status.s == PENDING) {
                     event_queue.push(id);
                     break;