Allow disabling send size zerocopy for rdma

example/rdma_performance/client.cpp

@@ -38,6 +38,7 @@ DEFINE_int32(expected_qps, 0, "The expected QPS");
 DEFINE_int32(max_thread_num, 16, "The max number of threads are used");
 DEFINE_int32(attachment_size, -1, "Attachment size is used (in Bytes)");
 DEFINE_bool(echo_attachment, false, "Select whether attachment should be echo");
+DEFINE_bool(attachment_as_userdata, false, "Append attachment as user_data");
 DEFINE_string(connection_type, "single", "Connection type of the channel");
 DEFINE_string(protocol, "baidu_std", "Protocol type.");
 DEFINE_string(servers, "0.0.0.0:8002+0.0.0.0:8002", "IP Address of servers");
@@ -86,7 +87,12 @@ class PerformanceTest {
  if (attachment_size > 0) {
  _addr = malloc(attachment_size);
  butil::fast_rand_bytes(_addr, attachment_size);
- _attachment.append(_addr, attachment_size);
+ if (FLAGS_attachment_as_userdata) {
+ brpc::rdma::RegisterMemoryForRdma(_addr, (size_t)attachment_size);
+ _attachment.append_user_data(_addr, attachment_size, NULL);
+ } else {
+ _attachment.append(_addr, attachment_size);
+ }
  }
  _echo_attachment = echo_attachment;
  }

src/brpc/rdma/rdma_endpoint.cpp

@@ -49,6 +49,7 @@ extern bool g_skip_rdma_init;
 
 DEFINE_int32(rdma_sq_size, 128, "SQ size for RDMA");
 DEFINE_int32(rdma_rq_size, 128, "RQ size for RDMA");
+DEFINE_bool(rdma_send_zerocopy, true, "Enable zerocopy for send side");
 DEFINE_bool(rdma_recv_zerocopy, true, "Enable zerocopy for receive side");
 DEFINE_int32(rdma_zerocopy_min_size, 512, "The minimal size for receive zerocopy");
 DEFINE_string(rdma_recv_block_type, "default", "Default size type for recv WR: "
@@ -801,29 +802,45 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) {
  wr.sg_list = sglist;
  wr.opcode = IBV_WR_SEND_WITH_IMM;
 
- RdmaIOBuf* data = (RdmaIOBuf*)from[current];
  size_t sge_index = 0;
  while (sge_index < (uint32_t)max_sge &&
  this_len < _remote_recv_block_size) {
- if (data->size() == 0) {
+ if (from[current]->size() == 0) {
  // The current IOBuf is empty, find next one
  ++current;
  if (current == ndata) {
  break;
  }
- data = (RdmaIOBuf*)from[current];
  continue;
  }
 
- ssize_t len = data->cut_into_sglist_and_iobuf(
- sglist, &sge_index, to, max_sge,
- _remote_recv_block_size - this_len);
- if (len < 0) {
- return -1;
+ ssize_t len = 0;
+ if (FLAGS_rdma_send_zerocopy) {
+ ssize_t len = ((RdmaIOBuf*)from[current])->cut_into_sglist_and_iobuf(
+ sglist, &sge_index, to, max_sge,
+ _remote_recv_block_size - this_len);
+ if (len < 0) {
+ return -1;
+ }
+ this_len += len;
+ total_len += len;
+ } else {
+ len = _remote_recv_block_size - this_len;
+ void* buf = AllocBlock(len);
+ if (!buf) {
+ return -1;
+ }
+ len = from[current]->copy_to(buf, len);
+ from[current]->cutn(to, len);
+ sglist[sge_index].length = len;
+ sglist[sge_index].addr = (uint64_t)buf;
+ sglist[sge_index].lkey = GetLKey(buf);
+ ++sge_index;
+ this_len += len;
+ total_len += len;
+ _sbuf_data[_sq_current] = buf;
+ break;
  }
- CHECK(len > 0);
- this_len += len;
- total_len += len;
  }
  if (this_len == 0) {
  continue;
@@ -951,6 +968,9 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
  uint32_t acks = butil::NetToHost32(wc.imm_data);
  uint32_t num = acks;
  while (num > 0) {
+ if (!FLAGS_rdma_send_zerocopy) {
+ DeallocBlock(_sbuf_data[_sq_sent]);
+ }
  _sbuf[_sq_sent++].clear();
  if (_sq_sent == _sq_size - RESERVED_WR_NUM) {
  _sq_sent = 0;
@@ -1139,6 +1159,10 @@ int RdmaEndpoint::AllocateResources() {
  if (_rbuf.size() != _rq_size) {
  return -1;
  }
+ _sbuf_data.resize(_sq_size, NULL);
+ if (_sbuf_data.size() != _sq_size) {
+ return -1;
+ }
  _rbuf_data.resize(_rq_size, NULL);
  if (_rbuf_data.size() != _rq_size) {
  return -1;

src/brpc/rdma/rdma_endpoint.h

@@ -213,6 +213,8 @@ friend class brpc::Socket;
  // Act as sendbuf and recvbuf, but requires no memcpy
  std::vector<butil::IOBuf> _sbuf;
  std::vector<butil::IOBuf> _rbuf;
+ // Data address of _sbuf
+ std::vector<void*> _sbuf_data;
  // Data address of _rbuf
  std::vector<void*> _rbuf_data;
  // Remote block size for receiving

src/brpc/rdma/rdma_helper.cpp

@@ -643,12 +643,15 @@ ibv_pd* GetRdmaPd() {
 }
 
 uint32_t GetLKey(void* buf) {
- BAIDU_SCOPED_LOCK(*g_user_mrs_lock);
- ibv_mr** mr_ptr = g_user_mrs->seek(buf);
- if (mr_ptr) {
- return (*mr_ptr)->lkey;
+ uint32_t lkey = GetRegionId(buf);
+ if (lkey == 0) {
+ BAIDU_SCOPED_LOCK(*g_user_mrs_lock);
+ ibv_mr** mr_ptr = g_user_mrs->seek(buf);
+ if (mr_ptr) {
+ return (*mr_ptr)->lkey;
+ }
  }
- return 0;
+ return lkey;
 }
 
 ibv_gid GetRdmaGid() {

test/brpc_rdma_unittest.cpp

@@ -68,6 +68,7 @@ struct HelloMessage {
 };
 
 DECLARE_bool(rdma_trace_verbose);
+DECLARE_bool(rdma_send_zerocopy);
 DECLARE_int32(rdma_memory_pool_max_regions);
 extern ibv_cq* (*IbvCreateCq)(ibv_context*, int, void*, ibv_comp_channel*, int);
 extern int (*IbvDestroyCq)(ibv_cq*);
@@ -1873,7 +1874,11 @@ TEST_F(RdmaTest, send_rpcs_with_user_defined_iobuf) {
  google::protobuf::Closure* done = DoNothing();
  ::test::EchoService::Stub(&channel).Echo(&cntl[0], &req[0], &res[0], done);
  bthread_id_join(cntl[0].call_id());
- ASSERT_EQ(ERDMAMEM, cntl[0].ErrorCode());
+ if (rdma::FLAGS_rdma_send_zerocopy) {
+ ASSERT_EQ(ERDMAMEM, cntl[0].ErrorCode());
+ } else {
+ ASSERT_EQ(0, cntl[0].ErrorCode());
+ }
  attach.clear();
  sleep(2); // wait for client recover from EHOSTDOWN
  cntl[0].Reset();