[Bugfix][TransferEngine] Serialize RDMA endpoint lifecycle by peer owner

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

@@ -41,7 +41,7 @@ class EndpointStore {
     virtual std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr) = 0;
     virtual std::shared_ptr<RdmaEndPoint> insertEndpoint(
-        const std::string &peer_nic_path, RdmaContext *context) = 0;
+        const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) = 0;
     virtual int deleteEndpoint(const std::string &peer_nic_path) = 0;
     // Deletes the endpoint matching endpoint_ptr (by pointer identity, under
     // the store lock -- the pointer is never dereferenced, so a stale/freed
@@ -83,7 +83,8 @@ class FIFOEndpointStore : public EndpointStore {
     std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr) override;
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
-        const std::string &peer_nic_path, RdmaContext *context) override;
+        const std::string &peer_nic_path, RdmaContext *context,
+        ibv_cq *cq) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;
     int deleteEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr,
@@ -125,7 +126,8 @@ class SIEVEEndpointStore : public EndpointStore {
     std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr) override;
     std::shared_ptr<RdmaEndPoint> insertEndpoint(
-        const std::string &peer_nic_path, RdmaContext *context) override;
+        const std::string &peer_nic_path, RdmaContext *context,
+        ibv_cq *cq) override;
     int deleteEndpoint(const std::string &peer_nic_path) override;
     int deleteEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr,
@@ -143,6 +145,11 @@ class SIEVEEndpointStore : public EndpointStore {
     }
 
     void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep) override;
+    // Test-only: push a pre-constructed endpoint into the active map so
+    // pointer-identity lookup/delete paths can be exercised without standing up
+    // an RDMA device.
+    void testOnlyInsertEndpoint(const std::string &peer_nic_path,
+                                std::shared_ptr<RdmaEndPoint> ep);
 
    private:
     RWSpinlock endpoint_map_lock_;

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

@@ -173,6 +173,10 @@ class RdmaContext {
    public:
     // EndPoint Management
     std::shared_ptr<RdmaEndPoint> endpoint(const std::string &peer_nic_path);
+    std::shared_ptr<RdmaEndPoint> endpoint(const std::string &peer_nic_path,
+                                           int cq_index);
+    std::shared_ptr<RdmaEndPoint> findEndpoint(
+        const std::string &peer_nic_path);
 
     std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
         const RdmaEndPoint *endpoint_ptr);
@@ -260,13 +264,20 @@ class RdmaContext {
 
     int eventFd() const { return event_fd_; }
 
-    ibv_cq *cq();
+    ibv_cq *cq(int cq_index);
 
     std::atomic<int> *cqOutstandingCount(int cq_index) {
         return &cq_list_[cq_index].outstanding;
     }
 
     int cqCount() const { return cq_list_.size(); }
+    int postingThreadForPeer(const std::string &peer_nic_path) const;
+    int cqIndexForPostingThread(int thread_id) const;
+    int cqIndexForPeer(const std::string &peer_nic_path) const;
+    int transferWorkerCount() const { return transfer_worker_count_; }
+    std::unique_lock<std::mutex> lockEndpointLifecycle(
+        const std::string &peer_nic_path) const;
+    std::vector<std::unique_lock<std::mutex>> lockAllEndpointLifecycles() const;
 
     int poll(int num_entries, ibv_wc *wc, int cq_index = 0);
 
@@ -329,7 +340,9 @@ class RdmaContext {
 
     std::atomic<int> next_comp_channel_index_;
     std::atomic<int> next_comp_vector_index_;
-    std::atomic<int> next_cq_list_index_;
+
+    int transfer_worker_count_ = 0;
+    std::vector<std::unique_ptr<std::mutex>> endpoint_lifecycle_locks_;
 
     std::shared_ptr<WorkerPool> worker_pool_;
 

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

@@ -71,6 +71,7 @@ class RdmaEndPoint {
 
    public:
     void setPeerNicPath(const std::string &peer_nic_path);
+    std::string peerNicPath() const;
 
     int setupConnectionsByActive();
 
@@ -197,9 +198,10 @@ class RdmaEndPoint {
     static constexpr uint32_t kWaitExistingHandshakeInitialSleepUs = 50;
     static constexpr uint32_t kWaitExistingHandshakeMaxSleepUs = 2000;
 
-    // Maximum time (in seconds) to wait for outstanding WRs to drain in
-    // finishDestroy before forcing QP destruction. This guards against
-    // ibv_modify_qp-to-ERR failures that prevent WR flushing.
+    // Maximum time (in seconds) to wait for outstanding WRs to drain before
+    // treating the endpoint as leaked. Timed-out endpoints stay in the
+    // EndpointStore waiting list so stale in-flight references cannot turn
+    // into UAF.
     static constexpr double kFinishDestroyTimeoutSec = 30.0;
 
     // Maximum number of deconstructLocked retries in finishDestroy before
@@ -210,7 +212,7 @@ class RdmaEndPoint {
     RdmaContext &context_;
     std::atomic<Status> status_;
 
-    RWSpinlock lock_;
+    mutable RWSpinlock lock_;
     std::vector<ibv_qp *> qp_list_;
     uint64_t qp_generation_;
 
@@ -225,8 +227,10 @@ class RdmaEndPoint {
     size_t max_inline_bytes_;
 
     std::atomic<bool> active_;
+    ibv_cq *cq_;
     std::atomic<int> *cq_outstanding_;
     std::atomic<uint64_t> inactive_time_;
+    bool finish_destroy_timeout_logged_ = false;
     int finish_destroy_retries_ = 0;
 };
 

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

@@ -114,9 +114,9 @@ class RdmaTransport : public Transport {
     virtual int onSetupRdmaConnections(const HandShakeDesc &peer_desc,
                                        HandShakeDesc &local_desc);
 
-    int sendHandshake(const std::string &peer_server_name,
-                      const HandShakeDesc &local_desc,
-                      HandShakeDesc &peer_desc) {
+    virtual int sendHandshake(const std::string &peer_server_name,
+                              const HandShakeDesc &local_desc,
+                              HandShakeDesc &peer_desc) {
         return metadata_->sendHandshake(peer_server_name, local_desc,
                                         peer_desc);
     }

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

@@ -15,8 +15,13 @@
 #ifndef WORKER_H
 #define WORKER_H
 
-#include <queue>
+#include <atomic>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+#include <unordered_map>
 #include <unordered_set>
+#include <vector>
 
 #include "config.h"
 #include "rdma_context.h"
@@ -41,20 +46,23 @@ class WorkerPool {
 
    private:
     using SliceList = std::vector<Transport::Slice *>;
-    const static int kShardCount = 8;
 
     // Enqueue slices that were prepared by another WorkerPool. Used for
     // local-NIC failure handoff: the original worker keeps the remote path
     // fixed, updates the local lkey, and pushes the slice to this context's
     // worker queue.
     int submitPreparedPostSend(
         const std::vector<Transport::Slice *> &slice_list);
-    void enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount],
+    void enqueuePreparedSlices(const SliceList &slice_list,
                                uint64_t submitted_slice_count);
+    void enqueueSliceToOwner(Transport::Slice *slice);
+    int postingThreadForPeer(const std::string &peer_nic_path) const;
+    int cqIndexForPostingThread(int thread_id) const;
 
     void performPostSend(int thread_id);
 
     void performPollCq(int thread_id);
+    void processCompletions(int thread_id, const std::vector<ibv_wc> &wc_list);
 
     void redispatch(std::vector<Transport::Slice *> &slice_list, int thread_id,
                     bool handoff_to_local_worker = false);
@@ -130,6 +138,7 @@ class WorkerPool {
    private:
     RdmaContext &context_;
     const int numa_socket_id_;
+    const int worker_count_;
 
     std::vector<std::thread> worker_thread_;
     std::atomic<bool> workers_running_;
@@ -151,12 +160,10 @@ class WorkerPool {
     std::mutex cond_mutex_;
     std::condition_variable cond_var_;
 
-    std::unordered_map<std::string, SliceList> slice_queue_[kShardCount];
-    std::atomic<uint64_t> slice_queue_count_[kShardCount];
-    TicketLock slice_queue_lock_[kShardCount];
-
     std::vector<std::unordered_map<std::string, SliceList>>
         collective_slice_queue_;
+    std::vector<std::unordered_map<std::string, SliceList>> worker_slice_queue_;
+    std::vector<std::mutex> worker_slice_queue_lock_;
 
     std::atomic<uint64_t> submitted_slice_count_, processed_slice_count_;
     std::atomic<uint64_t> recovery_activate_after_ns_{0};

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

@@ -47,22 +47,26 @@ std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::getEndpointByPtr(
 }
 
 std::shared_ptr<RdmaEndPoint> FIFOEndpointStore::insertEndpoint(
-    const std::string &peer_nic_path, RdmaContext *context) {
+    const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     if (endpoint_map_.find(peer_nic_path) != endpoint_map_.end()) {
         LOG(INFO) << "Endpoint " << peer_nic_path
                   << " already exists in FIFOEndpointStore";
         return endpoint_map_[peer_nic_path];
     }
+    if (!cq) {
+        LOG(ERROR) << "Cannot insert endpoint " << peer_nic_path
+                   << ": completion queue is null";
+        return nullptr;
+    }
     auto endpoint = std::make_shared<RdmaEndPoint>(*context);
     if (!endpoint) {
         LOG(ERROR) << "Failed to allocate memory for RdmaEndPoint";
         return nullptr;
     }
     auto &config = globalConfig();
-    int ret =
-        endpoint->construct(context->cq(), config.num_qp_per_ep, config.max_sge,
-                            config.max_wr, config.max_inline);
+    int ret = endpoint->construct(cq, config.num_qp_per_ep, config.max_sge,
+                                  config.max_wr, config.max_inline);
     if (ret) return nullptr;
 
     while (this->getSize() >= max_size_) evictEndpoint();
@@ -80,8 +84,9 @@ int FIFOEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     auto iter = endpoint_map_.find(peer_nic_path);
     // Begin two-phase destruction: mark endpoint as destroying and move QPs
     // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved
-    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once
-    // all outstanding WRs have been drained.
+    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() only
+    // after all outstanding WRs have been drained. Timed-out endpoints remain
+    // retired in waiting_list_ rather than being force-freed.
     if (iter != endpoint_map_.end()) {
         waiting_list_len_++;
         iter->second->beginDestroy();
@@ -220,22 +225,26 @@ std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::getEndpointByPtr(
 }
 
 std::shared_ptr<RdmaEndPoint> SIEVEEndpointStore::insertEndpoint(
-    const std::string &peer_nic_path, RdmaContext *context) {
+    const std::string &peer_nic_path, RdmaContext *context, ibv_cq *cq) {
     RWSpinlock::WriteGuard guard(endpoint_map_lock_);
     if (endpoint_map_.find(peer_nic_path) != endpoint_map_.end()) {
         LOG(INFO) << "Endpoint " << peer_nic_path
                   << " already exists in SIEVEEndpointStore";
         return endpoint_map_[peer_nic_path].first;
     }
+    if (!cq) {
+        LOG(ERROR) << "Cannot insert endpoint " << peer_nic_path
+                   << ": completion queue is null";
+        return nullptr;
+    }
     auto endpoint = std::make_shared<RdmaEndPoint>(*context);
     if (!endpoint) {
         LOG(ERROR) << "Failed to allocate memory for RdmaEndPoint";
         return nullptr;
     }
     auto &config = globalConfig();
-    int ret =
-        endpoint->construct(context->cq(), config.num_qp_per_ep, config.max_sge,
-                            config.max_wr, config.max_inline);
+    int ret = endpoint->construct(cq, config.num_qp_per_ep, config.max_sge,
+                                  config.max_wr, config.max_inline);
     if (ret) return nullptr;
 
     while (this->getSize() >= max_size_) evictEndpoint();
@@ -252,8 +261,9 @@ int SIEVEEndpointStore::deleteEndpoint(const std::string &peer_nic_path) {
     auto iter = endpoint_map_.find(peer_nic_path);
     // Begin two-phase destruction: mark endpoint as destroying and move QPs
     // to ERR state so inflight WRs are flushed to CQ. The endpoint is moved
-    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() once
-    // all outstanding WRs have been drained.
+    // to waiting_list_ and will be fully destroyed by reclaimEndpoint() only
+    // after all outstanding WRs have been drained. Timed-out endpoints remain
+    // retired in waiting_list_ rather than being force-freed.
     if (iter != endpoint_map_.end()) {
         iter->second.first->beginDestroy();
         waiting_list_len_++;
@@ -377,6 +387,14 @@ void SIEVEEndpointStore::testOnlyInsertWaiting(
     waiting_list_len_++;
 }
 
+void SIEVEEndpointStore::testOnlyInsertEndpoint(
+    const std::string &peer_nic_path, std::shared_ptr<RdmaEndPoint> ep) {
+    RWSpinlock::WriteGuard guard(endpoint_map_lock_);
+    endpoint_map_[peer_nic_path] = std::make_pair(ep, true);
+    fifo_list_.push_front(peer_nic_path);
+    fifo_map_[peer_nic_path] = fifo_list_.begin();
+}
+
 size_t SIEVEEndpointStore::getTotalQPNumber() {
     RWSpinlock::ReadGuard guard(endpoint_map_lock_);
     size_t total_qps = 0;