feat: bridge TE wait via poll executor for coroutine-friendly paths

Author: shichangzhang064

mooncake-integration/store/store_py.cpp

@@ -1045,6 +1045,7 @@ PYBIND11_MODULE(store, m) {
                size_t async_sender_thread_count = 0,
                size_t async_max_batch_size = 2000,
                size_t async_route_queue_size = 0,
+               size_t te_async_poll_worker_num = 32,
                const py::object& engine = py::none()) {
                 auto& resource_tracker = ResourceTracker::getInstance();
                 self.use_dummy_client_ = false;
@@ -1070,7 +1071,7 @@ PYBIND11_MODULE(store, m) {
                     local_transfer_mode, local_memcpy_async_worker_num,
                     metrics_port, enable_metrics_http, {},
                     async_sender_thread_count, async_max_batch_size,
-                    async_route_queue_size);
+                    async_route_queue_size, te_async_poll_worker_num);
 
                 auto ret = real_client->setup(config);
                 self.store_ = real_client;
@@ -1091,6 +1092,7 @@ PYBIND11_MODULE(store, m) {
             py::arg("async_sender_thread_count") = 0,
             py::arg("async_max_batch_size") = 2000,
             py::arg("async_route_queue_size") = 0,
+            py::arg("te_async_poll_worker_num") = 32,
             py::arg("engine") = py::none(),
             "Setup the store in P2P architecture.")
         .def(

mooncake-store/include/client_config_builder.h

@@ -171,6 +171,11 @@ struct P2PClientConfig : RealClientConfigBase {
     // When local_transfer_mode == MEMCPY, the following parameter is used:
     // 0 means forbid async memcpy (fall back to synchronous).
     size_t local_memcpy_async_worker_num = 32;
+
+    // Worker threads for offloading TransferEngine batch polling (WaitAllTransferBatches)
+    // in DataManager: local TE Put path and remote forward TE (co_await) paths. Independent
+    // of local_transfer_mode. 0 means synchronous TE wait on the caller/coroutine thread.
+    size_t te_async_poll_worker_num = 32;
 };
 
 // ============================================================================
@@ -238,7 +243,8 @@ class ClientConfigBuilder {
         bool enable_metrics_http = true,
         const std::map<std::string, std::string>& labels = {},
         size_t async_sender_thread_count = 0,
-        size_t async_max_batch_size = 2000, size_t async_route_queue_size = 0) {
+        size_t async_max_batch_size = 2000, size_t async_route_queue_size = 0,
+        size_t te_async_poll_worker_num = 32) {
         P2PClientConfig config;
         fill_real_client_config_base(
             config, local_hostname, metadata_connstring, protocol, rdma_devices,
@@ -251,6 +257,7 @@ class ClientConfigBuilder {
         config.route_cache_ttl_ms = route_cache_ttl_ms;
         config.local_transfer_mode =
             parse_p2p_local_transfer_mode(local_transfer_mode);
+        config.te_async_poll_worker_num = te_async_poll_worker_num;
         if (config.local_transfer_mode == LocalTransferMode::MEMCPY) {
             config.local_memcpy_async_worker_num =
                 local_memcpy_async_worker_num;

mooncake-store/include/data_manager.h

@@ -13,6 +13,7 @@
 #include <thread>
 #include <unordered_map>
 #include <vector>
+#include <async_simple/coro/Lazy.h>
 #include <ylt/util/tl/expected.hpp>
 #include "async_memcpy_executor.h"
 #include "client_buffer.hpp"
@@ -51,6 +52,11 @@ struct LocalTransferConfig {
     // When mode == MEMCPY, the following parameters are used:
     // 0 means forbid async memcpy (fall back to synchronous).
     size_t local_memcpy_async_worker_num = 32;
+
+    // Dedicated worker threads to offload TransferEngine batch polling
+    // (WaitAllTransferBatches) for local TE Put and remote forward TE paths.
+    // Independent of `mode`. 0 keeps synchronous TE wait on the caller thread.
+    size_t te_async_poll_worker_num = 32;
 };
 
 /**
@@ -211,15 +217,17 @@ class DataManager {
                                            const UUID& pin_token);
 
     /**
-     * @brief TE transfer without tier DRAM staging (PrepareDRAM*).
+     * @brief TE transfer without tier DRAM staging (PrepareDRAM*), with TE
+     * completion polling offloaded when `te_async_poll_worker_num > 0`.
      *
      * Caller guarantees `local_transfer_base` covers a contiguous layout of
-     * `total_size` bytes that is valid for TransferEngine (typically registered
-     * DRAM). Used by forward RDMA paths where buffers are already TE-ready.
+     * `total_size` bytes valid for TransferEngine (typically registered DRAM).
+     * Used by forward RDMA paths where buffers are already TE-ready.
      *
      * @param opcode WRITE: local -> peer_buffers; READ: peer_buffers -> local
      */
-    tl::expected<void, ErrorCode> TransferWithTeNoTierStaging(
+    async_simple::coro::Lazy<tl::expected<void, ErrorCode>>
+    TransferWithTeNoTierStagingAsync(
         void* local_transfer_base, size_t total_size,
         const std::vector<RemoteBufferDesc>& peer_buffers,
         Transport::TransferRequest::OpCode opcode);
@@ -350,6 +358,19 @@ class DataManager {
         const std::vector<RemoteBufferDesc>& remote_buffers,
         Transport::TransferRequest::OpCode opcode);
 
+    tl::expected<std::vector<std::tuple<Transport::BatchID, size_t, std::string>>,
+                 ErrorCode>
+    SubmitTeNoTierStagingBatches(
+        void* local_transfer_base, size_t total_size,
+        const std::vector<RemoteBufferDesc>& peer_buffers,
+        Transport::TransferRequest::OpCode opcode);
+
+    /** Synchronous TE wait on caller thread (used when no poll executor). */
+    tl::expected<void, ErrorCode> TransferWithTeNoTierStaging(
+        void* local_transfer_base, size_t total_size,
+        const std::vector<RemoteBufferDesc>& peer_buffers,
+        Transport::TransferRequest::OpCode opcode);
+
     tl::expected<
         std::vector<std::tuple<Transport::BatchID, size_t, std::string>>,
         ErrorCode>
@@ -517,6 +538,7 @@ class DataManager {
 
     LocalTransferConfig local_transfer_config_;
     std::unique_ptr<AsyncMemcpyExecutor> async_memcpy_executor_;
+    std::unique_ptr<AsyncMemcpyExecutor> te_poll_executor_;
     std::chrono::milliseconds lease_duration_;
     std::chrono::milliseconds lease_scan_interval_;
     std::atomic<bool> lease_scanner_stop_requested_{false};

mooncake-store/src/data_manager.cpp

@@ -18,6 +18,8 @@
 #include "utils/scoped_vlog_timer.h"
 #include "utils.h"
 
+#include <async_simple/coro/FutureAwaiter.h>
+
 namespace mooncake {
 
 namespace {
@@ -148,6 +150,10 @@ DataManager::DataManager(std::unique_ptr<TieredBackend> tiered_backend,
         async_memcpy_executor_ = std::make_unique<AsyncMemcpyExecutor>(
             local_transfer_config_.local_memcpy_async_worker_num);
     }
+    if (local_transfer_config_.te_async_poll_worker_num > 0) {
+        te_poll_executor_ = std::make_unique<AsyncMemcpyExecutor>(
+            local_transfer_config_.te_async_poll_worker_num);
+    }
 
     lease_duration_ = std::chrono::milliseconds(GetEnvOr<uint32_t>(
         "P2P_RPC_LEASE_DURATION_MS", kDefaultLeaseDurationMs));
@@ -162,8 +168,10 @@ DataManager::DataManager(std::unique_ptr<TieredBackend> tiered_backend,
                       ? "TE"
                       : "MEMCPY")
               << ", te_endpoint=" << local_transfer_config_.te_endpoint
-              << ", async_memcpy_workers="
+              << ", memcpy_async_workers="
               << local_transfer_config_.local_memcpy_async_worker_num
+              << ", te_async_poll_workers="
+              << local_transfer_config_.te_async_poll_worker_num
               << ", lease_duration_ms=" << lease_duration_.count()
               << ", lease_scan_interval_ms=" << lease_scan_interval_.count();
 }
@@ -176,6 +184,9 @@ void DataManager::Stop() {
     if (async_memcpy_executor_) {
         async_memcpy_executor_->Shutdown();
     }
+    if (te_poll_executor_) {
+        te_poll_executor_->Shutdown();
+    }
     if (tiered_backend_) {
         tiered_backend_->Stop();
     }
@@ -485,16 +496,6 @@ tl::expected<std::unique_ptr<TaskHandle<void>>, ErrorCode> DataManager::Put(
     return tl::unexpected(ErrorCode::INTERNAL_ERROR);
 }
 
-// TODO: The returned CallableTaskHandle's WaitAsync() falls back to a
-// synchronous Wait() on the coroutine's current thread, because the
-// WaitAllTransferBatches() is a loop with no async completion notification.
-// Possible optimizations:
-//   (1) run a polling coroutine on yalantinglibs coro_io's io_context (via
-//       co_await coro_io::sleep_for(100us) + getTransferStatus), no new thread;
-//   (2) introduce a lightweight timer service to bridge cv-poll to
-//       async_simple::Promise;
-//   (3) introduce a completion callback from transfer_engine itself.
-// Once any of these lands, switch the return type to FutureHandle.
 tl::expected<std::unique_ptr<TaskHandle<void>>, ErrorCode>
 DataManager::PutViaTe(std::string_view key, std::vector<Slice>& slices) {
     // using Te, treat local memory as remote memory
@@ -536,46 +537,51 @@ DataManager::PutViaTe(std::string_view key, std::vector<Slice>& slices) {
         return tl::unexpected(submit_result.error());
     }
 
-    return CallableTaskHandle<void>::Create(
-        [this, ctx = std::move(*submit_result), alloc_handle, kctx,
-         pending_write_token]() mutable -> tl::expected<void, ErrorCode> {
-            ScopedVLogTimer timer(1, "DataManager::PutViaTe");
-            timer.LogRequest("key=", kctx.key);
+    auto te_phase = [this, ctx = std::move(*submit_result), alloc_handle, kctx,
+                     pending_write_token]() mutable -> tl::expected<void, ErrorCode> {
+        ScopedVLogTimer timer(1, "DataManager::PutViaTe");
+        timer.LogRequest("key=", kctx.key);
+
+        auto wait_result = WaitAllTransferBatches(ctx.transfer_batches);
+        if (!wait_result) {
+            LOG(ERROR) << "WaitAllTransferBatches failed"
+                       << ", key=" << kctx.key
+                       << ", error_code=" << toString(wait_result.error());
+            (void)WriteRevokeInternal(kctx, pending_write_token);
+            return tl::unexpected(wait_result.error());
+        }
 
-            auto wait_result = WaitAllTransferBatches(ctx.transfer_batches);
-            if (!wait_result) {
-                LOG(ERROR) << "WaitAllTransferBatches failed"
+        if (ctx.handle->loc.data.type != MemoryType::DRAM && ctx.temp_buffer) {
+            auto& loc_data = ctx.handle->loc.data;
+            auto copy_result = CopyFromDRAMBuffer(
+                ctx.temp_buffer.get(),
+                reinterpret_cast<void*>(loc_data.buffer->data()),
+                loc_data.type, loc_data.buffer->size(), ctx.handle->backend);
+            if (!copy_result) {
+                LOG(ERROR) << "CopyFromDRAMBuffer failed"
                            << ", key=" << kctx.key
-                           << ", error_code=" << toString(wait_result.error());
+                           << ", error_code=" << toString(copy_result.error());
                 (void)WriteRevokeInternal(kctx, pending_write_token);
-                return tl::unexpected(wait_result.error());
+                return tl::unexpected(copy_result.error());
             }
+        }
 
-            if (ctx.handle->loc.data.type != MemoryType::DRAM &&
-                ctx.temp_buffer) {
-                auto& loc_data = ctx.handle->loc.data;
-                auto copy_result = CopyFromDRAMBuffer(
-                    ctx.temp_buffer.get(),
-                    reinterpret_cast<void*>(loc_data.buffer->data()),
-                    loc_data.type, loc_data.buffer->size(),
-                    ctx.handle->backend);
-                if (!copy_result) {
-                    LOG(ERROR)
-                        << "CopyFromDRAMBuffer failed"
-                        << ", key=" << kctx.key
-                        << ", error_code=" << toString(copy_result.error());
-                    (void)WriteRevokeInternal(kctx, pending_write_token);
-                    return tl::unexpected(copy_result.error());
-                }
-            }
+        auto commit_result = WriteCommitInternal(kctx, pending_write_token);
+        if (!commit_result) {
+            return tl::unexpected(commit_result.error());
+        }
+        timer.LogResponse("error_code=", ErrorCode::OK);
+        return {};
+    };
 
-            auto commit_result = WriteCommitInternal(kctx, pending_write_token);
-            if (!commit_result) {
-                return tl::unexpected(commit_result.error());
-            }
-            timer.LogResponse("error_code=", ErrorCode::OK);
-            return {};
-        });
+    if (te_poll_executor_) {
+        auto future = te_poll_executor_
+                          ->SubmitSingleTask<tl::expected<void, ErrorCode>>(
+                              std::move(te_phase));
+        return FutureHandle<void>::Create(std::shared_ptr<void>{},
+                                          std::move(future));
+    }
+    return CallableTaskHandle<void>::Create(std::move(te_phase));
 }
 
 tl::expected<std::unique_ptr<TaskHandle<void>>, ErrorCode>
@@ -756,19 +762,27 @@ tl::expected<ReadTaskHandle, ErrorCode> DataManager::BuildDataCopierViaTe(
 
     ReadTaskHandle res;
     res.data_size = static_cast<int64_t>(source_size);
-    res.task_handle = CallableTaskHandle<void>::Create(
-        [this, ctx = std::move(submit_result.value()),
-         h = handle]() mutable -> tl::expected<void, ErrorCode> {
-            ScopedVLogTimer timer(1, "DataManager::BuildDataCopierViaTe");
-            auto wait_result = WaitAllTransferBatches(ctx.transfer_batches);
-            if (!wait_result) {
-                LOG(ERROR) << "Failed to wait TE read transfer, error_code="
-                           << wait_result.error();
-                return tl::unexpected(wait_result.error());
-            }
-            timer.LogResponse("error_code=", ErrorCode::OK);
-            return {};
-        });
+    auto te_wait = [this, ctx = std::move(submit_result.value()),
+                    h = handle]() mutable -> tl::expected<void, ErrorCode> {
+        ScopedVLogTimer timer(1, "DataManager::BuildDataCopierViaTe");
+        auto wait_result = WaitAllTransferBatches(ctx.transfer_batches);
+        if (!wait_result) {
+            LOG(ERROR) << "Failed to wait TE read transfer, error_code="
+                       << wait_result.error();
+            return tl::unexpected(wait_result.error());
+        }
+        timer.LogResponse("error_code=", ErrorCode::OK);
+        return {};
+    };
+    if (te_poll_executor_) {
+        auto future = te_poll_executor_
+                          ->SubmitSingleTask<tl::expected<void, ErrorCode>>(
+                              std::move(te_wait));
+        res.task_handle = FutureHandle<void>::Create(std::shared_ptr<void>{},
+                                                     std::move(future));
+    } else {
+        res.task_handle = CallableTaskHandle<void>::Create(std::move(te_wait));
+    }
     return res;
 }
 
@@ -1375,7 +1389,9 @@ DataManager::SubmitTeTransferBatches(
     return submitted_batches;
 }
 
-tl::expected<void, ErrorCode> DataManager::TransferWithTeNoTierStaging(
+tl::expected<std::vector<std::tuple<Transport::BatchID, size_t, std::string>>,
+             ErrorCode>
+DataManager::SubmitTeNoTierStagingBatches(
     void* local_transfer_base, size_t total_size,
     const std::vector<RemoteBufferDesc>& peer_buffers,
     Transport::TransferRequest::OpCode opcode) {
@@ -1394,8 +1410,16 @@ tl::expected<void, ErrorCode> DataManager::TransferWithTeNoTierStaging(
         LOG(ERROR) << "TransferEngine not initialized";
         return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
     }
-    auto batches = SubmitTeTransferBatches(local_transfer_base, total_size,
-                                           peer_buffers, opcode);
+    return SubmitTeTransferBatches(local_transfer_base, total_size, peer_buffers,
+                                 opcode);
+}
+
+tl::expected<void, ErrorCode> DataManager::TransferWithTeNoTierStaging(
+    void* local_transfer_base, size_t total_size,
+    const std::vector<RemoteBufferDesc>& peer_buffers,
+    Transport::TransferRequest::OpCode opcode) {
+    auto batches = SubmitTeNoTierStagingBatches(
+        local_transfer_base, total_size, peer_buffers, opcode);
     if (!batches) {
         return tl::unexpected(batches.error());
     }
@@ -1409,6 +1433,29 @@ tl::expected<void, ErrorCode> DataManager::TransferWithTeNoTierStaging(
     return {};
 }
 
+async_simple::coro::Lazy<tl::expected<void, ErrorCode>>
+DataManager::TransferWithTeNoTierStagingAsync(
+    void* local_transfer_base, size_t total_size,
+    const std::vector<RemoteBufferDesc>& peer_buffers,
+    Transport::TransferRequest::OpCode opcode) {
+    if (!te_poll_executor_) {
+        co_return TransferWithTeNoTierStaging(local_transfer_base, total_size,
+                                              peer_buffers, opcode);
+    }
+    auto batches = SubmitTeNoTierStagingBatches(
+        local_transfer_base, total_size, peer_buffers, opcode);
+    if (!batches) {
+        co_return tl::make_unexpected(batches.error());
+    }
+    auto batch_vec = std::move(batches.value());
+    auto fut =
+        te_poll_executor_->SubmitSingleTask<tl::expected<void, ErrorCode>>(
+            [this, batch_vec = std::move(batch_vec)]() mutable {
+                return WaitAllTransferBatches(batch_vec);
+            });
+    co_return co_await std::move(fut);
+}
+
 tl::expected<void, ErrorCode> DataManager::ValidateRemoteBuffers(
     const std::vector<RemoteBufferDesc>& buffers) {
     if (buffers.empty()) {