fix: avoid cq burn out

Author: alogfans

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

@@ -39,6 +39,12 @@ class RdmaTransport;
 class WorkerPool;
 class EndpointStore;
 
+struct RdmaCq {
+    RdmaCq() : native(nullptr), outstanding(0) {}
+    ibv_cq *native;
+    volatile int outstanding;
+};
+
 // RdmaContext represents the set of resources controlled by each local NIC,
 // including Memory Region, CQ, EndPoint (QPs), etc.
 class RdmaContext {
@@ -111,6 +117,10 @@ class RdmaContext {
 
     ibv_cq *cq();
 
+    volatile int *cqOutstandingCount(int cq_index) {
+        return &cq_list_[cq_index].outstanding;
+    }
+
     int cqCount() const { return cq_list_.size(); }
 
     int poll(int num_entries, ibv_wc *wc, int cq_index = 0);
@@ -150,7 +160,8 @@ class RdmaContext {
 
     RWSpinlock memory_regions_lock_;
     std::vector<ibv_mr *> memory_region_list_;
-    std::vector<ibv_cq *> cq_list_;
+    // std::vector<ibv_cq *> cq_list_;
+    std::vector<RdmaCq> cq_list_;
 
     std::shared_ptr<EndpointStore> endpoint_store_;
 

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

@@ -125,6 +125,7 @@ class RdmaEndPoint {
     int max_wr_depth_;
 
     volatile bool active_;
+    volatile int *cq_outstanding_;
 };
 
 }  // namespace mooncake

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

@@ -108,13 +108,16 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels,
 
     cq_list_.resize(num_cq_list);
     for (size_t i = 0; i < num_cq_list; ++i) {
-        cq_list_[i] = ibv_create_cq(context_, max_cqe, this /* CQ context */,
-                                    compChannel(), compVector());
-        if (!cq_list_[i]) {
+        auto cq =
+            ibv_create_cq(context_, max_cqe,
+                          (void *)&cq_list_[i].outstanding /* CQ context */,
+                          compChannel(), compVector());
+        if (!cq) {
             PLOG(ERROR) << "Failed to create completion queue";
             close(event_fd_);
             return ERR_CONTEXT;
         }
+        cq_list_[i].native = cq;
     }
 
     worker_pool_ = std::make_shared<WorkerPool>(*this, socketId());
@@ -153,7 +156,7 @@ int RdmaContext::deconstruct() {
     memory_region_list_.clear();
 
     for (size_t i = 0; i < cq_list_.size(); ++i) {
-        int ret = ibv_destroy_cq(cq_list_[i]);
+        int ret = ibv_destroy_cq(cq_list_[i].native);
         if (ret) {
             PLOG(ERROR) << "Failed to destroy completion queue";
         }
@@ -310,7 +313,7 @@ std::string RdmaContext::gid() const {
 
 ibv_cq *RdmaContext::cq() {
     int index = (next_cq_list_index_++) % cq_list_.size();
-    return cq_list_[index];
+    return cq_list_[index].native;
 }
 
 ibv_comp_channel *RdmaContext::compChannel() {
@@ -498,7 +501,7 @@ int RdmaContext::joinNonblockingPollList(int event_fd, int data_fd) {
 }
 
 int RdmaContext::poll(int num_entries, ibv_wc *wc, int cq_index) {
-    int nr_poll = ibv_poll_cq(cq_list_[cq_index], num_entries, wc);
+    int nr_poll = ibv_poll_cq(cq_list_[cq_index].native, num_entries, wc);
     if (nr_poll < 0) {
         LOG(ERROR) << "Failed to poll CQ " << cq_index << " of device "
                    << device_name_;

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

@@ -27,7 +27,10 @@ const static uint8_t TIMEOUT = 14;
 const static uint8_t RETRY_CNT = 7;
 
 RdmaEndPoint::RdmaEndPoint(RdmaContext &context)
-    : context_(context), status_(INITIALIZING), active_(true) {}
+    : context_(context),
+      status_(INITIALIZING),
+      active_(true),
+      cq_outstanding_(nullptr) {}
 
 RdmaEndPoint::~RdmaEndPoint() {
     if (!qp_list_.empty()) deconstruct();
@@ -42,6 +45,7 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list,
     }
 
     qp_list_.resize(num_qp_list);
+    cq_outstanding_ = (volatile int *)cq->cq_context;
 
     max_wr_depth_ = (int)max_wr_depth;
     wr_depth_list_ = new volatile int[num_qp_list];
@@ -224,7 +228,9 @@ int RdmaEndPoint::submitPostSend(
     int qp_index = SimpleRandom::Get().next(qp_list_.size());
     int wr_count = std::min(max_wr_depth_ - wr_depth_list_[qp_index],
                             (int)slice_list.size());
-    if (wr_count == 0) return 0;
+    wr_count =
+        std::min(int(globalConfig().max_cqe) - *cq_outstanding_, wr_count);
+    if (wr_count <= 0) return 0;
 
     ibv_send_wr wr_list[wr_count], *bad_wr = nullptr;
     ibv_sge sge_list[wr_count];
@@ -258,13 +264,15 @@ int RdmaEndPoint::submitPostSend(
         // }
     }
     __sync_fetch_and_add(&wr_depth_list_[qp_index], wr_count);
+    __sync_fetch_and_add(cq_outstanding_, wr_count);
     int rc = ibv_post_send(qp_list_[qp_index], wr_list, &bad_wr);
     if (rc) {
         PLOG(ERROR) << "Failed to ibv_post_send";
         while (bad_wr) {
             int i = bad_wr - wr_list;
             failed_slice_list.push_back(slice_list[i]);
             __sync_fetch_and_sub(&wr_depth_list_[qp_index], 1);
+            __sync_fetch_and_sub(cq_outstanding_, 1);
             bad_wr = bad_wr->next;
         }
     }

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

@@ -305,6 +305,9 @@ void WorkerPool::performPollCq(int thread_id) {
                 processed_slice_count++;
             }
         }
+        if (nr_poll)
+            __sync_fetch_and_sub(context_.cqOutstandingCount(cq_index),
+                                 nr_poll);
     }
 
     for (auto &entry : qp_depth_set)
@@ -393,10 +396,12 @@ int WorkerPool::doProcessContextEvents() {
         event.event_type == IBV_EVENT_LID_CHANGE) {
         context_.set_active(false);
         context_.disconnectAllEndpoints();
-        LOG(INFO) << "Worker: Context " << context_.deviceName() << " is now inactive";
+        LOG(INFO) << "Worker: Context " << context_.deviceName()
+                  << " is now inactive";
     } else if (event.event_type == IBV_EVENT_PORT_ACTIVE) {
         context_.set_active(true);
-        LOG(INFO) << "Worker: Context " << context_.deviceName() << " is now active";
+        LOG(INFO) << "Worker: Context " << context_.deviceName()
+                  << " is now active";
     }
     ibv_ack_async_event(&event);
     return 0;