rdma support polling mode

Author: yanglimingcn

docs/cn/rdma.md

@@ -47,6 +47,8 @@ RDMA要求数据收发所使用的内存空间必须被注册（memory register
 
 RDMA是硬件相关的通信技术，有很多独特的概念，比如device、port、GID、LID、MaxSge等。这些参数在初始化时会从对应的网卡中读取出来，并且做出默认的选择（参见src/brpc/rdma/rdma_helper.cpp）。有时默认的选择并非用户的期望，则可以通过flag参数方式指定。
 
+RDMA支持事件驱动和轮询两种模式，默认是事件驱动模式，通过设置rdma_use_polling可以开启轮询模式。轮询模式下还可以设置轮询器数目（rdma_poller_num），以及是否主动放弃CPU（rdma_poller_yield）。轮询模式下还可以设置一个回调函数，在每次轮询时调用，可以配合io_uring/spdk等使用。
+
 # 参数
 
 可配置参数说明：
@@ -68,3 +70,6 @@ RDMA是硬件相关的通信技术，有很多独特的概念，比如device、p
 * rdma_memory_pool_max_regions: 最大的内存池块数，默认16
 * rdma_memory_pool_buckets: 内存池中为避免竞争采用的bucket数目，默认为4
 * rdma_memory_pool_tls_cache_num: 内存池中thread local的缓存block数目，默认为128
+* rdma_use_polling: 是否使用RDMA的轮询模式，默认false
+* rdma_poller_num: 轮询模式下的poller数目，默认1
+* rdma_poller_yield: 轮询模式下的poller是否主动放弃CPU，默认是false

src/brpc/rdma/rdma_endpoint.cpp

@@ -58,6 +58,9 @@ DEFINE_int32(rdma_prepared_qp_size, 128, "SQ and RQ size for prepared QP.");
 DEFINE_int32(rdma_prepared_qp_cnt, 1024, "Initial count of prepared QP.");
 DEFINE_bool(rdma_trace_verbose, false, "Print log message verbosely");
 BRPC_VALIDATE_GFLAG(rdma_trace_verbose, brpc::PassValidate);
+DEFINE_bool(rdma_use_polling, false, "Use polling mode for RDMA.");
+DEFINE_int32(rdma_poller_num, 1, "Poller number in RDMA polling mode.");
+DEFINE_bool(rdma_poller_yield, false, "Yield thread in RDMA polling mode.");
 
 static const size_t IOBUF_BLOCK_HEADER_LEN = 32; // implementation-dependent
 
@@ -699,7 +702,6 @@ void* RdmaEndpoint::ProcessHandshakeAtServer(void* arg) {
         LOG_IF(INFO, FLAGS_rdma_trace_verbose) 
             << "Handshake ends (use tcp) on " << s->description();
     }
- 
     ep->TryReadOnTcp();
 
     return NULL;
@@ -1041,26 +1043,36 @@ static RdmaResource* AllocateQpCq(uint16_t sq_size, uint16_t rq_size) {
         return NULL;
     }
 
-    res->comp_channel = IbvCreateCompChannel(GetRdmaContext());
-    if (!res->comp_channel) {
-        PLOG(WARNING) << "Fail to create comp channel for CQ";
-        delete res;
-        return NULL;
-    }
+    if (!FLAGS_rdma_use_polling) {
+        res->comp_channel = IbvCreateCompChannel(GetRdmaContext());
+        if (!res->comp_channel) {
+            PLOG(WARNING) << "Fail to create comp channel for CQ";
+            delete res;
+            return NULL;
+        }
 
-    butil::make_close_on_exec(res->comp_channel->fd);
-    if (butil::make_non_blocking(res->comp_channel->fd) < 0) {
-        PLOG(WARNING) << "Fail to set comp channel nonblocking";
-        delete res;
-        return NULL;
-    }
+        butil::make_close_on_exec(res->comp_channel->fd);
+        if (butil::make_non_blocking(res->comp_channel->fd) < 0) {
+            PLOG(WARNING) << "Fail to set comp channel nonblocking";
+            delete res;
+            return NULL;
+        }
 
-    res->cq = IbvCreateCq(GetRdmaContext(), 2 * FLAGS_rdma_prepared_qp_size,
-            NULL, res->comp_channel, GetRdmaCompVector());
-    if (!res->cq) {
-        PLOG(WARNING) << "Fail to create CQ";
-        delete res;
-        return NULL;
+        res->cq = IbvCreateCq(GetRdmaContext(), 2 * FLAGS_rdma_prepared_qp_size,
+                              NULL, res->comp_channel, GetRdmaCompVector());
+        if (!res->cq) {
+            PLOG(WARNING) << "Fail to create CQ";
+            delete res;
+            return NULL;
+        }
+    } else {
+        res->cq = IbvCreateCq(GetRdmaContext(), 2 * FLAGS_rdma_prepared_qp_size,
+                              NULL, NULL, 0);
+        if (!res->cq) {
+            PLOG(WARNING) << "Fail to create CQ";
+            delete res;
+            return NULL;
+        }
     }
 
     ibv_qp_init_attr attr;
@@ -1117,19 +1129,30 @@ int RdmaEndpoint::AllocateResources() {
         return -1;
     }
 
-    SocketOptions options;
-    options.user = this;
-    options.keytable_pool = _socket->_keytable_pool;
-    options.fd = _resource->comp_channel->fd;
-    options.on_edge_triggered_events = PollCq;
-    if (Socket::Create(options, &_cq_sid) < 0) {
-        PLOG(WARNING) << "Fail to create socket for cq";
-        return -1;
-    }
+    if (!FLAGS_rdma_use_polling) {
+        SocketOptions options;
+        options.user = this;
+        options.keytable_pool = _socket->_keytable_pool;
+        options.fd = _resource->comp_channel->fd;
+        options.on_edge_triggered_events = PollCq;
+        if (Socket::Create(options, &_cq_sid) < 0) {
+            PLOG(WARNING) << "Fail to create socket for cq";
+            return -1;
+        }
 
-    if (ibv_req_notify_cq(_resource->cq, 1) < 0) {
-        PLOG(WARNING) << "Fail to arm CQ comp channel";
-        return -1;
+        if (ibv_req_notify_cq(_resource->cq, 1) < 0) {
+            PLOG(WARNING) << "Fail to arm CQ comp channel";
+            return -1;
+        }
+    } else {
+        SocketOptions options;
+        options.user = this;
+        options.keytable_pool = _socket->_keytable_pool;
+        if (Socket::Create(options, &_cq_sid) < 0) {
+            PLOG(WARNING) << "Fail to create socket for cq";
+            return -1;
+        }
+        PollerAddCqSocket();
     }
 
     _sbuf.resize(_sq_size - RESERVED_WR_NUM);
@@ -1227,6 +1250,9 @@ void RdmaEndpoint::DeallocateResources() {
     if (!_resource) {
         return;
     }
+    if (FLAGS_rdma_use_polling) {
+        PollerRemoveCqSocket();
+    }
     bool move_to_rdma_resource_list = false;
     if (_sq_size <= FLAGS_rdma_prepared_qp_size &&
         _rq_size <= FLAGS_rdma_prepared_qp_size &&
@@ -1237,7 +1263,10 @@ void RdmaEndpoint::DeallocateResources() {
             move_to_rdma_resource_list = true;
         }
     }
-    int fd = _resource->comp_channel->fd;
+    int fd = -1;
+    if (_resource->comp_channel) {
+        fd = _resource->comp_channel->fd;
+    }
     if (!move_to_rdma_resource_list) {
         if (_resource->qp) {
             if (IbvDestroyQp(_resource->qp) < 0) {
@@ -1327,12 +1356,14 @@ void RdmaEndpoint::PollCq(Socket* m) {
     }
     CHECK(ep == s->_rdma_ep);
 
-    if (ep->GetAndAckEvents() < 0) {
-        const int saved_errno = errno;
-        PLOG(ERROR) << "Fail to get cq event: " << s->description();
-        s->SetFailed(saved_errno, "Fail to get cq event from %s: %s",
-                s->description().c_str(), berror(saved_errno));
-        return;
+    if (!FLAGS_rdma_use_polling) {
+        if (ep->GetAndAckEvents() < 0) {
+            const int saved_errno = errno;
+            PLOG(ERROR) << "Fail to get cq event: " << s->description();
+            s->SetFailed(saved_errno, "Fail to get cq event from %s: %s",
+                         s->description().c_str(), berror(saved_errno));
+            return;
+        }
     }
 
     int progress = Socket::PROGRESS_INIT;
@@ -1349,6 +1380,9 @@ void RdmaEndpoint::PollCq(Socket* m) {
             return;
         }
         if (cnt == 0) {
+            if (FLAGS_rdma_use_polling) {
+                return;
+            }
             if (!notified) {
                 // Since RDMA only provides one shot event, we have to call the
                 // notify function every time. Because there is a possibility
@@ -1370,7 +1404,7 @@ void RdmaEndpoint::PollCq(Socket* m) {
             }
             if (ep->GetAndAckEvents() < 0) {
                 s->SetFailed(errno, "Fail to ack CQ event on %s",
-                        s->description().c_str());
+                    s->description().c_str());
                 return;
             }
             notified = false;
@@ -1474,6 +1508,13 @@ int RdmaEndpoint::GlobalInitialize() {
         g_rdma_resource_list = res;
     }
 
+    if (FLAGS_rdma_use_polling) {
+        auto rc = PollingModeInitialize();
+        if (rc != 0) {
+            return -1;
+        }
+    }
+
     return 0;
 }
 
@@ -1486,6 +1527,84 @@ void RdmaEndpoint::GlobalRelease() {
             delete res;
         }
     }
+
+    if (FLAGS_rdma_use_polling) {
+        PollingModeRelease();
+    }
+}
+
+std::vector<RdmaEndpoint::Poller> RdmaEndpoint::_pollers;
+std::atomic<bool> RdmaEndpoint::_running(false);
+std::function<void()> RdmaEndpoint::_callback(nullptr);
+butil::Mutex RdmaEndpoint::_cb_mutex;
+
+void RdmaEndpoint::SetCallbackFn(std::function<void()> cb) { _callback = cb; }
+
+int RdmaEndpoint::PollingModeInitialize() {
+    auto fn = [](void* args) -> void* {
+        auto poller = static_cast<Poller*>(args);
+        while (_running.load(butil::memory_order_relaxed)) {
+            std::list<Socket*> sockets;
+            {
+                std::unique_lock<butil::Mutex> lk(poller->mutex);
+                sockets = poller->sockets;  // copy all sockets is not good
+            }
+            for (auto m : sockets) {
+                PollCq(m);
+            }
+            {
+                std::unique_lock<butil::Mutex> lk(_cb_mutex);
+                if (_callback) {
+                    _callback();
+                }
+            }
+            if (FLAGS_rdma_poller_yield) {
+                bthread_yield();
+            }
+        }
+        return nullptr;
+    };
+    _pollers = std::vector<Poller>(FLAGS_rdma_poller_num);
+    _running.store(true, butil::memory_order_relaxed);
+    for (int i = 0; i < FLAGS_rdma_poller_num; ++i) {
+        auto rc = bthread_start_background(
+            &_pollers[i].tid, &BTHREAD_ATTR_NORMAL, fn, &_pollers[i]);
+        if (rc != 0) {
+            LOG(ERROR) << "Fail to start rdma polling bthread";
+            return -1;
+        }
+    }
+    return 0;
+}
+
+void RdmaEndpoint::PollingModeRelease() {
+    _running.store(false, butil::memory_order_relaxed);
+    for (int i = 0; i < FLAGS_rdma_poller_num; ++i) {
+        bthread_join(_pollers[i].tid, nullptr);
+    }
+}
+
+// Add socket to poller
+void RdmaEndpoint::PollerAddCqSocket() {
+    SocketUniquePtr s;
+    if (Socket::Address(_cq_sid, &s) != 0) {
+        LOG(ERROR) << "Fail to address socket id: " << _cq_sid;
+        return;
+    }
+    auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num;
+    auto& poller = _pollers[index];
+    std::unique_lock<butil::Mutex> lk(poller.mutex);
+    _cq_iter = poller.sockets.insert(poller.sockets.end(), s.get());
+}
+
+// Remove socket from poller
+void RdmaEndpoint::PollerRemoveCqSocket() {
+    auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num;
+    auto& poller = _pollers[index];
+    std::unique_lock<butil::Mutex> lk(poller.mutex);
+    if (_cq_iter != poller.sockets.end()) {
+        poller.sockets.erase(_cq_iter);
+    }
 }
 
 }  // namespace rdma

src/brpc/rdma/rdma_endpoint.h

@@ -23,6 +23,8 @@
 #include <iostream>
 #include <string>
 #include <vector>
+#include <list>
+#include <functional>
 #include <infiniband/verbs.h>
 #include "butil/atomicops.h"
 #include "butil/iobuf.h"
@@ -34,6 +36,8 @@ namespace brpc {
 class Socket;
 namespace rdma {
 
+DECLARE_bool(rdma_use_polling);
+
 class RdmaConnect : public AppConnect {
 public:
     void StartConnect(const Socket* socket, 
@@ -90,6 +94,9 @@ friend class brpc::Socket;
     // Callback when there is new epollin event on TCP fd
     static void OnNewDataFromTcp(Socket* m);
 
+    // Set callback function
+    static void SetCallbackFn(std::function<void()> cb);
+
 private:
     enum State {
         UNINIT = 0x0,
@@ -191,6 +198,17 @@ friend class brpc::Socket;
     // Try to read data on TCP fd in _socket
     inline void TryReadOnTcp();
 
+    // Initialize polling mode
+    static int PollingModeInitialize();
+
+    static void PollingModeRelease();
+
+    // Add cq socket to poller
+    void PollerAddCqSocket();
+
+    // Remove cq socket to poller
+    void PollerRemoveCqSocket();
+
     // Not owner
     Socket* _socket;
 
@@ -206,6 +224,9 @@ friend class brpc::Socket;
     // the SocketId which wrap the comp channel of CQ
     SocketId _cq_sid;
 
+    // the Socket pointer's iterator in list
+    std::list<Socket*>::iterator _cq_iter;
+
     // Capacity of local Send Queue and local Recv Queue
     uint16_t _sq_size;
     uint16_t _rq_size;
@@ -245,6 +266,18 @@ friend class brpc::Socket;
     butil::atomic<int> *_read_butex;
 
     DISALLOW_COPY_AND_ASSIGN(RdmaEndpoint);
+
+    struct Poller {
+        bthread_t tid;
+        // TODO(yangliming): lockless queue is better
+        std::list<Socket*> sockets;
+        butil::Mutex mutex;
+    };
+    static std::vector<Poller> _pollers;
+    static std::atomic<bool> _running;
+    // Callback used for io_uring/spdk async engine 
+    static std::function<void()> _callback;
+    static butil::Mutex _cb_mutex;
 };
 
 }  // namespace rdma

src/brpc/socket.cpp

@@ -2273,6 +2273,14 @@ int Socket::OnInputEvent(void* user_data, uint32_t events,
         attr.tag = bthread_self_tag();
         if (FLAGS_usercode_in_coroutine) {
             ProcessEvent(p);
+#if BRPC_WITH_RDMA
+        } else if (rdma::FLAGS_rdma_use_polling && p->_rdma_state == RDMA_ON) {
+            auto rc = bthread_start_background(&tid, &attr, ProcessEvent, p);
+            if (rc != 0) {
+                LOG(FATAL) << "Fail to start ProcessEvent";
+                ProcessEvent(p);
+            }
+#endif
         } else if (bthread_start_urgent(&tid, &attr, ProcessEvent, p) != 0) {
             LOG(FATAL) << "Fail to start ProcessEvent";
             ProcessEvent(p);

src/brpc/span.cpp

@@ -17,6 +17,7 @@
 
 
 #include <netinet/in.h>
+#include <functional>
 #include <gflags/gflags.h>
 #include <leveldb/db.h>
 #include <leveldb/comparator.h>