add pending signaled wrs checking

Author: yanglimingcn

src/brpc/rdma/rdma_endpoint.cpp

@@ -19,6 +19,7 @@
 
 #include <gflags/gflags.h>
 #include "butil/fd_utility.h"
+#include "butil/fast_rand.h"
 #include "butil/logging.h"                   // CHECK, LOG
 #include "butil/sys_byteorder.h"             // HostToNet,NetToHost
 #include "bthread/bthread.h"
@@ -191,6 +192,7 @@ RdmaEndpoint::RdmaEndpoint(Socket* s)
     , _remote_window_capacity(0)
     , _window_size(0)
     , _new_rq_wrs(0)
+    , _pending_signaled_wrs(0)
 {
     if (_sq_size < MIN_QP_SIZE) {
         _sq_size = MIN_QP_SIZE;
@@ -875,6 +877,7 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) {
             // http::www.rdmamojo.com/2014/06/30/working-unsignaled-completions/
             wr.send_flags |= IBV_SEND_SIGNALED;
             _sq_unsignaled = 0;
+            _pending_signaled_wrs.fetch_add(1, butil::memory_order_relaxed);
         }
 
         ibv_send_wr* bad = NULL;
@@ -904,6 +907,35 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) {
     return total_len;
 }
 
+struct WrIdGen {
+    bool init;
+    uint32_t seq;
+    uint64_t current_random;
+    butil::FastRandSeed seed;
+};
+
+static __thread WrIdGen tls_wr_id_gen = { false, 0, 0, { { 0, 0 } } };
+
+inline uint64_t UpdateTLSRandom64(IdGen* g) {
+    if (!g->init) {
+        g->init = true;
+        init_fast_rand_seed(&g->seed);
+    }
+    const uint64_t val = fast_rand(&g->seed);
+    g->current_random = val;
+    return val;
+}
+
+inline uint64_t GenerateWrId() {
+    // 0 is an invalid Id
+    IdGen* g = &tls_wr_id_gen;
+    if (g->seq == 0) {
+        UpdateTLSRandom64(g);
+        g->seq = 1;
+    }
+    return (g->current_random & 0xFFFFFFFF00000000ULL) | g->seq++;
+}
+
 int RdmaEndpoint::SendAck(int num) {
     if (_new_rq_wrs.fetch_add(num, butil::memory_order_relaxed) > _remote_window_capacity / 2) {
         return SendImm(_new_rq_wrs.exchange(0, butil::memory_order_relaxed));
@@ -918,6 +950,7 @@ int RdmaEndpoint::SendImm(uint32_t imm) {
 
     ibv_send_wr wr;
     memset(&wr, 0, sizeof(wr));
+    wr.wr_id = GenerateWrId();
     wr.opcode = IBV_WR_SEND_WITH_IMM;
     wr.imm_data = butil::HostToNet32(imm);
     wr.send_flags |= IBV_SEND_SOLICITED;
@@ -938,7 +971,9 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
     bool zerocopy = FLAGS_rdma_recv_zerocopy;
     switch (wc.opcode) {
     case IBV_WC_SEND: {  // send completion
-        // Do nothing
+        if (wc.wr_id == 0) {
+            _pending_signaled_wrs.fetch_sub(1, butil::memory_order_relaxed);
+        }
         break;
     }
     case IBV_WC_RECV: {  // recv completion
@@ -972,10 +1007,12 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
 
             // Update window
             uint32_t wnd_thresh = _local_window_capacity / 8;
-            if (_window_size.fetch_add(acks, butil::memory_order_relaxed) >= wnd_thresh
-                    || acks >= wnd_thresh) {
+            if (_pending_signaled_wrs.load(butil::memory_order_relaxed) <= 1 &&
+                (_window_size.fetch_add(acks, butil::memory_order_relaxed) >= wnd_thresh ||
+                 acks >= wnd_thresh)) {
                 // Do not wake up writing thread right after _window_size > 0.
-                // Otherwise the writing thread may switch to background too quickly.
+                // Otherwise the writing thread may switch to background too
+                // quickly.
                 _socket->WakeAsEpollOut();
             }
         }

src/brpc/rdma/rdma_endpoint.h

@@ -261,6 +261,8 @@ friend class brpc::Socket;
     butil::atomic<uint16_t> _window_size;
     // The number of new WRs posted in the local Recv Queue
     butil::atomic<uint16_t> _new_rq_wrs;
+    // The number of pending signaled send WRs
+    butil::atomic<uint16_t> _pending_signaled_wrs;
 
     // butex for inform read events on TCP fd during handshake
     butil::atomic<int> *_read_butex;