[TransferEngine] Keep RDMA owner CQ polled during active handshake

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

@@ -61,11 +61,13 @@ class WorkerPool {
 
     void performPostSend(int thread_id);
 
-    void performPollCq(int thread_id);
-    void processCompletions(int thread_id, const std::vector<ibv_wc> &wc_list);
+    int performPollCq(int thread_id, bool defer_local_redispatch = false);
+    void processCompletions(int thread_id, const std::vector<ibv_wc> &wc_list,
+                            bool defer_local_redispatch = false);
 
     void redispatch(std::vector<Transport::Slice *> &slice_list, int thread_id,
-                    bool handoff_to_local_worker = false);
+                    bool handoff_to_local_worker = false,
+                    bool defer_local_redispatch = false);
 
     void transferWorker(int thread_id);
 

mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp

@@ -16,8 +16,15 @@
 
 #include <sys/epoll.h>
 
+#include <algorithm>
 #include <cassert>
+#include <chrono>
+#include <condition_variable>
+#include <exception>
 #include <functional>
+#include <future>
+#include <queue>
+#include <thread>
 
 #include "config.h"
 #include "memory_location.h"
@@ -52,14 +59,116 @@ struct ActiveEndpointSetupResult {
     bool endpoint_current = false;
 };
 
+class ActiveEndpointSetupExecutor {
+   public:
+    explicit ActiveEndpointSetupExecutor(size_t thread_count) : running_(true) {
+        workers_.reserve(thread_count);
+        for (size_t i = 0; i < thread_count; ++i) {
+            workers_.emplace_back([this] { run(); });
+        }
+    }
+
+    ~ActiveEndpointSetupExecutor() {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            running_ = false;
+        }
+        cv_.notify_all();
+        for (auto &worker : workers_) {
+            if (worker.joinable()) worker.join();
+        }
+    }
+
+    std::future<int> submit(const std::shared_ptr<RdmaEndPoint> &endpoint) {
+        std::packaged_task<int()> task(
+            [endpoint] { return endpoint->setupConnectionsByActive(); });
+        auto future = task.get_future();
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            if (!running_) {
+                task();
+                return future;
+            }
+            tasks_.emplace(std::move(task));
+        }
+        cv_.notify_one();
+        return future;
+    }
+
+   private:
+    void run() {
+        while (true) {
+            std::packaged_task<int()> task;
+            {
+                std::unique_lock<std::mutex> lock(mutex_);
+                cv_.wait(lock, [this] { return !running_ || !tasks_.empty(); });
+                if (!running_ && tasks_.empty()) return;
+                task = std::move(tasks_.front());
+                tasks_.pop();
+            }
+            task();
+        }
+    }
+
+    std::vector<std::thread> workers_;
+    std::queue<std::packaged_task<int()>> tasks_;
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    bool running_;
+};
+
+static ActiveEndpointSetupExecutor &activeEndpointSetupExecutor() {
+    constexpr size_t kActiveEndpointSetupThreadCount = 8;
+    static ActiveEndpointSetupExecutor executor(
+        kActiveEndpointSetupThreadCount);
+    return executor;
+}
+
 static ActiveEndpointSetupResult setupEndpointByActiveOutsideLifecycleGate(
     RdmaContext &context, const std::string &peer_nic_path,
     const std::shared_ptr<RdmaEndPoint> &endpoint,
-    std::unique_lock<std::mutex> &endpoint_lifecycle_lock) {
+    std::unique_lock<std::mutex> &endpoint_lifecycle_lock,
+    const std::function<int()> &drain_cq_while_waiting) {
     const bool had_lifecycle_gate = endpoint_lifecycle_lock.owns_lock();
     if (had_lifecycle_gate) endpoint_lifecycle_lock.unlock();
 
-    int ret = endpoint->setupConnectionsByActive();
+    // The owner worker also polls this peer's CQ. Keep it draining while the
+    // active handshake can block on TCP connect/read timeouts.
+    int ret = ERR_ENDPOINT;
+    try {
+        auto setup_future = activeEndpointSetupExecutor().submit(endpoint);
+        auto wait_period = std::chrono::milliseconds(1);
+        constexpr auto kMaxWaitPeriod = std::chrono::milliseconds(25);
+        while (setup_future.wait_for(std::chrono::milliseconds(0)) !=
+               std::future_status::ready) {
+            bool made_progress = false;
+            while (setup_future.wait_for(std::chrono::milliseconds(0)) !=
+                   std::future_status::ready) {
+                const int drained = drain_cq_while_waiting();
+                if (drained <= 0) break;
+                made_progress = true;
+            }
+            if (setup_future.wait_for(std::chrono::milliseconds(0)) ==
+                std::future_status::ready) {
+                break;
+            }
+            if (made_progress) {
+                wait_period = std::chrono::milliseconds(1);
+            } else {
+                if (setup_future.wait_for(wait_period) ==
+                    std::future_status::ready) {
+                    break;
+                }
+                wait_period = std::min(wait_period * 2, kMaxWaitPeriod);
+            }
+        }
+        ret = setup_future.get();
+    } catch (const std::exception &ex) {
+        LOG(ERROR) << "Worker: Active endpoint setup threw: " << ex.what();
+    } catch (...) {
+        LOG(ERROR)
+            << "Worker: Active endpoint setup threw an unknown exception";
+    }
 
     if (had_lifecycle_gate) endpoint_lifecycle_lock.lock();
     auto current_endpoint = context.findEndpoint(peer_nic_path);
@@ -455,7 +564,11 @@ void WorkerPool::performPostSend(int thread_id) {
         }
         if (!endpoint->connected()) {
             auto setup_result = setupEndpointByActiveOutsideLifecycleGate(
-                context_, entry.first, endpoint, endpoint_lifecycle_lock);
+                context_, entry.first, endpoint, endpoint_lifecycle_lock,
+                [this, thread_id] {
+                    if (!hasOutstandingCq(thread_id)) return 0;
+                    return performPollCq(thread_id, true);
+                });
             if (!setup_result.endpoint_current) {
                 LOG(WARNING)
                     << "Worker: Endpoint changed while active handshake was "
@@ -559,13 +672,14 @@ void WorkerPool::performPostSend(int thread_id) {
     }
 }
 
-void WorkerPool::performPollCq(int thread_id) {
-    if (context_.cqCount() <= 0) return;
-    if (thread_id < 0 || thread_id >= worker_count_) return;
+int WorkerPool::performPollCq(int thread_id, bool defer_local_redispatch) {
+    if (context_.cqCount() <= 0) return 0;
+    if (thread_id < 0 || thread_id >= worker_count_) return 0;
 
     const uint64_t poll_ts = getCurrentTimeInNano();
     const uint64_t previous_poll_ts =
-        last_poll_ts_ns_.exchange(poll_ts, std::memory_order_relaxed);
+        last_poll_ts_ns_.exchange(poll_ts, std::memory_order_release);
+    last_poll_ts_ns_.notify_all();
     if (previous_poll_ts > 0 && poll_ts > previous_poll_ts) {
         const uint64_t interval = poll_ts - previous_poll_ts;
         last_poll_interval_ns_.store(interval, std::memory_order_relaxed);
@@ -581,12 +695,13 @@ void WorkerPool::performPollCq(int thread_id) {
     std::unordered_map<std::atomic<int> *, int> qp_depth_set;
     std::vector<ibv_wc> wc_list;
     const int cq_index = cqIndexForPostingThread(thread_id);
-    if (cq_index < 0) return;
+    if (cq_index < 0) return 0;
+    if (!context_.cq(cq_index)) return 0;
     ibv_wc wc[kPollCount];
     int nr_poll = context_.poll(kPollCount, wc, cq_index);
     if (nr_poll < 0) {
         LOG(ERROR) << "Worker: Failed to poll completion queues";
-        return;
+        return nr_poll;
     }
 
     if (nr_poll > 0 && globalConfig().track_rdma_posted_slices) {
@@ -614,11 +729,14 @@ void WorkerPool::performPollCq(int thread_id) {
     for (auto &entry : qp_depth_set)
         entry.first->fetch_sub(entry.second, std::memory_order_acq_rel);
 
-    if (!wc_list.empty()) processCompletions(thread_id, wc_list);
+    if (!wc_list.empty())
+        processCompletions(thread_id, wc_list, defer_local_redispatch);
+    return nr_poll;
 }
 
 void WorkerPool::processCompletions(int thread_id,
-                                    const std::vector<ibv_wc> &wc_list) {
+                                    const std::vector<ibv_wc> &wc_list,
+                                    bool defer_local_redispatch) {
     // Slices this call drove to a terminal state, successes and
     // retry-exhausted failures alike; folded into processed_slice_count_,
     // which gates worker parking. Every terminal outcome below goes through
@@ -756,15 +874,17 @@ void WorkerPool::processCompletions(int thread_id,
     if (succeeded_slice_count) markContextSuccess();
 
     if (!local_failed_slice_list.empty()) {
-        redispatch(local_failed_slice_list, thread_id, true);
+        redispatch(local_failed_slice_list, thread_id, true,
+                   defer_local_redispatch);
     }
     if (!failed_slice_list.empty()) {
-        redispatch(failed_slice_list, thread_id);
+        redispatch(failed_slice_list, thread_id, false, defer_local_redispatch);
     }
 }
 
 void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
-                            int thread_id, bool handoff_to_local_worker) {
+                            int thread_id, bool handoff_to_local_worker,
+                            bool defer_local_redispatch) {
     std::unordered_map<SegmentID, std::shared_ptr<Transport::SegmentDesc>>
         segment_desc_map;
     int shared_redispatch_count = 0;
@@ -878,7 +998,7 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
             }
             slice->ts = 0;
             const int owner_thread = postingThreadForPeer(peer_nic_path);
-            if (owner_thread == thread_id) {
+            if (owner_thread == thread_id && !defer_local_redispatch) {
                 collective_slice_queue_[thread_id][peer_nic_path].push_back(
                     slice);
             } else {

mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp

@@ -19,9 +19,11 @@
 #include <condition_variable>
 #include <cstdint>
 #include <future>
+#include <limits>
 #include <memory>
 #include <mutex>
 #include <string>
+#include <thread>
 #include <utility>
 #include <vector>
 
@@ -118,6 +120,12 @@ class RdmaContextTestPeer {
     static void clearCompletionQueues(RdmaContext &context) {
         context.cq_list_.clear();
     }
+
+    static void clearNativeCompletionQueue(RdmaContext &context, int cq_index) {
+        ASSERT_GE(cq_index, 0);
+        ASSERT_LT(static_cast<size_t>(cq_index), context.cq_list_.size());
+        context.cq_list_[cq_index].native = nullptr;
+    }
 };
 
 class WorkerPoolTestPeer {
@@ -145,6 +153,43 @@ class WorkerPoolTestPeer {
         for (auto &queue : pool.collective_slice_queue_) queue.clear();
         for (auto &queue : pool.worker_slice_queue_) queue.clear();
     }
+
+    static uint64_t lastPollTs(const WorkerPool &pool) {
+        return pool.last_poll_ts_ns_.load(std::memory_order_acquire);
+    }
+
+    static bool waitForPollTs(WorkerPool &pool,
+                              std::chrono::milliseconds timeout) {
+        if (lastPollTs(pool) != 0) return true;
+
+        std::mutex mutex;
+        std::condition_variable cv;
+        bool observed = false;
+        std::thread waiter([&] {
+            uint64_t poll_ts =
+                pool.last_poll_ts_ns_.load(std::memory_order_acquire);
+            while (poll_ts == 0) {
+                pool.last_poll_ts_ns_.wait(poll_ts, std::memory_order_acquire);
+                poll_ts = pool.last_poll_ts_ns_.load(std::memory_order_acquire);
+            }
+            {
+                std::lock_guard<std::mutex> lock(mutex);
+                observed = true;
+            }
+            cv.notify_one();
+        });
+
+        std::unique_lock<std::mutex> lock(mutex);
+        const bool ok = cv.wait_for(lock, timeout, [&] { return observed; });
+        if (!ok) {
+            pool.last_poll_ts_ns_.store(std::numeric_limits<uint64_t>::max(),
+                                        std::memory_order_release);
+            pool.last_poll_ts_ns_.notify_all();
+        }
+        lock.unlock();
+        waiter.join();
+        return ok;
+    }
 };
 
 class RdmaEndPointTestPeer {
@@ -193,10 +238,12 @@ class InProcessRdmaTransport : public RdmaTransport {
             ++barrier_->arrivals;
             barrier_->cv.notify_all();
             if (!barrier_->cv.wait_for(lock, std::chrono::seconds(5), [&] {
-                    return barrier_->arrivals >= 2;
+                    return barrier_->arrivals >= barrier_->required_arrivals ||
+                           barrier_->released;
                 })) {
                 return ERR_ENDPOINT;
             }
+            if (fail_after_barrier_) return ERR_ENDPOINT;
         }
 
         return peer_->onSetupRdmaConnections(local_desc, peer_desc);
@@ -210,11 +257,14 @@ class InProcessRdmaTransport : public RdmaTransport {
         std::mutex mutex;
         std::condition_variable cv;
         int arrivals = 0;
+        int required_arrivals = 2;
+        bool released = false;
     };
 
     RdmaContext *local_context_ = nullptr;
     InProcessRdmaTransport *peer_ = nullptr;
     std::shared_ptr<Barrier> barrier_;
+    bool fail_after_barrier_ = false;
 
    private:
     std::atomic<bool> lifecycle_gate_was_free_{true};
@@ -512,4 +562,54 @@ TEST(RdmaEndpointLifecycleGateTest,
     WorkerPoolTestPeer::clearQueuedSlices(*peer_b.worker_pool);
 }
 
+TEST(RdmaEndpointLifecycleGateTest, ActiveHandshakeWaitDrainsOwnerCq) {
+    auto barrier = std::make_shared<InProcessRdmaTransport::Barrier>();
+    FakeRdmaPeer peer;
+
+    ASSERT_NO_FATAL_FAILURE(
+        initFakePeer(peer, "rdma-drain-a:10000", "mlx5_drain_a", barrier));
+    peer.transport->fail_after_barrier_ = true;
+
+    const std::string peer_nic_path =
+        MakeNicPath("rdma-drain-b:10000", "mlx5_drain_b");
+    ASSERT_NO_FATAL_FAILURE(installZeroQpEndpoint(peer, peer_nic_path));
+
+    // The fake CQ is not a real verbs object. Clearing the native pointer keeps
+    // the regression hardware-free while still proving the owner worker enters
+    // its CQ poll path during the blocked active handshake.
+    RdmaContextTestPeer::clearNativeCompletionQueue(*peer.context, 0);
+    peer.context->cqOutstandingCount(0)->store(1, std::memory_order_relaxed);
+
+    Transport::Slice slice;
+    Transport::TransferTask task;
+    queueHandshakeSlice(peer, peer_nic_path, slice, task);
+    slice.rdma.max_retry_cnt = 1;
+
+    auto active = std::async(std::launch::async, [&] {
+        WorkerPoolTestPeer::performPostSend(*peer.worker_pool, 0);
+    });
+
+    {
+        std::unique_lock<std::mutex> lock(barrier->mutex);
+        ASSERT_TRUE(barrier->cv.wait_for(lock, std::chrono::seconds(5), [&] {
+            return barrier->arrivals >= 1;
+        }));
+    }
+
+    EXPECT_TRUE(WorkerPoolTestPeer::waitForPollTs(*peer.worker_pool,
+                                                  std::chrono::seconds(5)));
+
+    {
+        std::lock_guard<std::mutex> lock(barrier->mutex);
+        barrier->released = true;
+    }
+    barrier->cv.notify_all();
+
+    ASSERT_EQ(active.wait_for(std::chrono::seconds(5)),
+              std::future_status::ready);
+    active.get();
+
+    WorkerPoolTestPeer::clearQueuedSlices(*peer.worker_pool);
+}
+
 }  // namespace