support gpu direct rdma

BUILD.bazel

@@ -54,6 +54,9 @@ COPTS = [
 }) + select({
     "//bazel/config:brpc_with_asan": ["-fsanitize=address"],
     "//conditions:default": [""],
+}) + select({
+    ":brpc_with_gdr": ["-DBRPC_WITH_GDR=1"],
+    "//conditions:default": [""],
 }) + select({
     "//bazel/config:brpc_with_no_pthread_mutex_hook": ["-DNO_PTHREAD_MUTEX_HOOK"],
     "//conditions:default": [""],
@@ -232,6 +235,7 @@ BUTIL_SRCS = [
     "src/butil/iobuf.cpp",
     "src/butil/single_iobuf.cpp",
     "src/butil/iobuf_profiler.cpp",
+    "src/butil/gpu/gpu_block_pool.cpp",
     "src/butil/binary_printer.cpp",
     "src/butil/recordio.cc",
     "src/butil/popen.cpp",
@@ -337,6 +341,9 @@ cc_library(
             "-DUNIT_TEST",
         ],
         "//conditions:default": [],
+    }) + select({
+        ":brpc_with_gdr": ["@local_config_cuda//cuda:cuda_headers"],
+        "//conditions:default": [],
     }),
     includes = [
         "src/",
@@ -356,6 +363,9 @@ cc_library(
     }) + select({
         "//bazel/config:brpc_with_boringssl": ["@boringssl//:ssl", "@boringssl//:crypto"],
         "//conditions:default": ["@openssl//:ssl", "@openssl//:crypto"],
+    }) + select({
+        ":brpc_with_gdr": ["@local_config_cuda//cuda:cuda_headers"],
+        "//conditions:default": [],
     }),
 )
 
@@ -573,6 +583,9 @@ cc_library(
             "@org_apache_thrift//:thrift",
         ],
         "//conditions:default": [],
+    }) + select({
+        ":brpc_with_gdr": ["@local_config_cuda//cuda:cuda_headers"],
+        "//conditions:default": [],
     }),
 )
 

bazel/config/BUILD.bazel

@@ -104,6 +104,12 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+config_setting(
+    name = "brpc_with_gdr",
+    define_values = {"BRPC_WITH_GDR": "true"},
+    visibility = ["//visibility:public"],
+)
+
 config_setting(
     name = "brpc_with_boringssl",
     define_values = {"BRPC_WITH_BORINGSSL": "true"},
@@ -148,4 +154,4 @@ config_setting(
     name = "with_babylon_counter",
     define_values = {"with_babylon_counter": "true"},
     visibility = ["//visibility:public"],
-)
+)

src/brpc/acceptor.cpp

@@ -41,6 +41,7 @@ Acceptor::Acceptor(bthread_keytable_pool_t* pool)
     , _force_ssl(false)
     , _ssl_ctx(NULL) 
     , _use_rdma(false)
+    , _use_gdr(false)
     , _bthread_tag(BTHREAD_TAG_DEFAULT) {
 }
 

src/brpc/acceptor.h

@@ -113,6 +113,9 @@ friend class Server;
     // Whether to use rdma or not
     bool _use_rdma;
 
+    // Whether to use gdr or not
+    bool _use_gdr;
+
     // Acceptor belongs to this tag
     bthread_tag_t _bthread_tag;
 };

src/brpc/channel.cpp

@@ -61,6 +61,7 @@ ChannelOptions::ChannelOptions()
     , succeed_without_server(true)
     , log_succeed_without_server(true)
     , use_rdma(false)
+    , use_gdr(false)
     , auth(NULL)
     , backup_request_policy(NULL)
     , retry_policy(NULL)
@@ -123,6 +124,9 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) {
         if (opt.use_rdma) {
             buf.append("|rdma");
         }
+        if (opt.use_gdr) {
+            buf.append("|gdr");
+        }
         butil::MurmurHash3_x64_128_Update(&mm_ctx, buf.data(), buf.size());
         buf.clear();
 
@@ -197,6 +201,11 @@ int Channel::InitChannelOptions(const ChannelOptions* options) {
             return -1;
         }
         rdma::GlobalRdmaInitializeOrDie();
+#if BRPC_WITH_GDR
+        if (_options.use_gdr) {
+            rdma::GlobalGdrInitializeOrDie();
+        }
+#endif // BRPC_WITH_GDR
         if (!rdma::InitPollingModeWithTag(bthread_self_tag())) {
             return -1;
         }
@@ -369,7 +378,8 @@ int Channel::InitSingle(const butil::EndPoint& server_addr_and_port,
         return -1;
     }
     if (SocketMapInsert(SocketMapKey(server_addr_and_port, sig),
-                        &_server_id, ssl_ctx, _options.use_rdma, _options.hc_option) != 0) {
+                        &_server_id, ssl_ctx, _options.use_rdma, _options.use_gdr,
+                        _options.hc_option) != 0) {
         LOG(ERROR) << "Fail to insert into SocketMap";
         return -1;
     }
@@ -407,6 +417,7 @@ int Channel::Init(const char* ns_url,
     ns_opt.succeed_without_server = _options.succeed_without_server;
     ns_opt.log_succeed_without_server = _options.log_succeed_without_server;
     ns_opt.use_rdma = _options.use_rdma;
+    ns_opt.use_gdr = _options.use_gdr;
     ns_opt.channel_signature = ComputeChannelSignature(_options);
     ns_opt.hc_option =  _options.hc_option;
     if (CreateSocketSSLContext(_options, &ns_opt.ssl_ctx) != 0) {

src/brpc/channel.h

@@ -109,6 +109,10 @@ struct ChannelOptions {
     // Default: false
     bool use_rdma;
 
+    // Let this channel use gdu direct rdma.
+    // Default: false
+    bool use_gdr;
+
     // Turn on authentication for this channel if `auth' is not NULL.
     // Note `auth' will not be deleted by channel and must remain valid when
     // the channel is being used.

src/brpc/details/naming_service_thread.cpp

@@ -126,7 +126,8 @@ void NamingServiceThread::Actions::ResetServers(
         //       to pick those Sockets with the right settings during OnAddedServers
         const SocketMapKey key(_added[i], _owner->_options.channel_signature);
         CHECK_EQ(0, SocketMapInsert(key, &tagged_id.id, _owner->_options.ssl_ctx,
-                                    _owner->_options.use_rdma, _owner->_options.hc_option));
+                                    _owner->_options.use_rdma, _owner->_options.use_gdr,
+                                    _owner->_options.hc_option));
         _added_sockets.push_back(tagged_id);
     }
 

src/brpc/details/naming_service_thread.h

@@ -45,11 +45,13 @@ struct GetNamingServiceThreadOptions {
     GetNamingServiceThreadOptions()
         : succeed_without_server(false)
         , log_succeed_without_server(true)
-        , use_rdma(false) {}
+        , use_rdma(false)
+        , use_gdr(false) {}
 
     bool succeed_without_server;
     bool log_succeed_without_server;
     bool use_rdma;
+    bool use_gdr;
     HealthCheckOption hc_option;
     ChannelSignature channel_signature;
     std::shared_ptr<SocketSSLContext> ssl_ctx;