fix _read_buf race between PollCq and OnNewMessages in RDMA server

The server-side RDMA socket's _read_buf is accessed by two independent
bthreads: PollCq (CQ socket) writes RDMA data via HandleCompletion and
calls ProcessNewMessage (which reads _read_buf via CutInputMessage), and
OnNewMessages (main socket) reads TCP data for handshake / fallback.
Since IOBuf is not thread-safe, concurrent access corrupts internal
state and causes intermittent core dumps.

Three fixes:

1. Switch edge trigger to OnNewDataFromTcp in ALL ExecuteServerHandshake
   end paths (ESTABLISHED + 5 failure paths). OnNewDataFromTcp checks the
   RDMA state: in ESTABLISHED it only reads 1 byte for EOF detection
   without touching _read_buf; in FALLBACK_TCP it delegates to
   OnNewMessages for TCP data. This prevents post-handshake races.

2. Guard HandleCompletion (IBV_WC_RECV) with a state check: skip writing
   to _read_buf if the state is not ESTABLISHED, but still handle imm
   data, re-post the recv WR (with failure check), and send ack. This
   prevents races during the handshake (after BringUpQp puts the QP into
   RTS, the client may start sending RDMA data before the server
   finishes processing the ACK).

3. Remove the source->size() > HELLO_ACK_LEN check in Phase 2. When a
   client falls back to TCP, the 4-byte ACK and the first RPC request
   may arrive in the same readv() call. Use cutn() to drain the 4-byte
   ACK and let remaining data be processed by other parsers, matching
   FallbackServerHandshake's behavior.

Additionally:

- Return NOT_ENOUGH_DATA (not TRY_OTHERS) from the ESTABLISHED path so
  CutInputMessage returns immediately without reading _read_buf,
  minimizing the race window with PollCq.
- Clear _read_buf before transitioning to ESTABLISHED so that residual
  TCP data cannot become a prefix of the RDMA recv stream (HandleCompletion
  appends to _read_buf, not overwrites). The clear is safe because
  HandleCompletion only writes after seeing ESTABLISHED (acquire), which
  is stored (release) strictly after the clear.
- Use memory_order_release for ESTABLISHED stores (both client and server)
  to properly pair with the acquire load in HandleCompletion.
- Restore edge trigger in RdmaTransport::Reset() based on CreatedByConnect():
  OnNewDataFromTcp for client-side sockets, OnNewMessages for server-side
  sockets, matching the logic in Init().
- Add Transport::ShouldStopReading() virtual method (default false),
  overridden by RdmaTransport to return true when the RDMA endpoint has
  reached ESTABLISHED (via RdmaEndpoint::IsEstablished(), which uses an
  acquire load on the already-atomic _state, pairing with the release
  store in ExecuteServerHandshake). OnNewMessages checks this after
  ProcessNewMessage returns and exits immediately, preventing it from
  calling DoRead again on _read_buf after the edge trigger has been
  switched.
- Guard ProcessNewMessage in PollCq with bytes > 0: when bytes == 0
  (IBV_WC_SEND completions, or IBV_WC_RECV dropped during handshake),
  skip ProcessNewMessage entirely. This prevents PollCq from calling
  CutInputMessage on _read_buf (via ProcessNewMessage) while
  OnNewMessages is driving the handshake on the same _read_buf.

Author: shengtiedan

src/brpc/input_messenger.cpp

@@ -369,7 +369,19 @@ void InputMessenger::OnNewMessages(Socket* m) {
         if (messenger->ProcessNewMessage(m, nr, read_eof, received_us,
                                          base_realtime, last_msg) < 0) {
             return;
-        } 
+        }
+        // If the transport switched its edge trigger during parsing (e.g.,
+        // RDMA handshake completed and edge trigger changed to
+        // OnNewDataFromTcp), stop reading to avoid racing with the new
+        // edge trigger handler on _read_buf. Drain _nevent so future
+        // epoll events can schedule the new edge trigger handler.
+        if (m->_transport->ShouldStopReading()) {
+            while (m->MoreReadEvents(&progress)) {}
+            if (read_eof) {
+                m->SetEOF();
+            }
+            return;
+        }
     }
 
     if (read_eof) {

src/brpc/rdma/rdma_endpoint.cpp

@@ -500,7 +500,7 @@ void* RdmaEndpoint::ProcessHandshakeAtClient(void* arg) {
     }
 
     if (rdma_transport->_rdma_state == RdmaTransport::RDMA_ON) {
-        ep->_state.store(ESTABLISHED, butil::memory_order_relaxed);
+        ep->_state.store(ESTABLISHED, butil::memory_order_release);
         LOG_IF(INFO, FLAGS_rdma_trace_verbose)
             << "Client handshake ends (use rdma v" << ep->_handshake_version
             << ") on " << s->description();
@@ -560,6 +560,7 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s
         }
         if (r == RemoteHelloResult::ERROR) {
             ep->_state.store(FAILED, butil::memory_order_relaxed);
+            rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
             return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
         }
 
@@ -592,6 +593,7 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s
         if (hs->SendLocalHello() < 0) {
             PLOG(WARNING) << "Fail to send server hello to " << s->description();
             ep->_state.store(FAILED, butil::memory_order_relaxed);
+            rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
             return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
         }
 
@@ -607,13 +609,6 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s
     if (source->size() < HELLO_ACK_LEN) {
         return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
     }
-    if (source->size() > HELLO_ACK_LEN) {
-        LOG(WARNING) << "Too many bytes in handshake ACK, drop connection: "
-                     << s->description();
-        ep->_state.store(FAILED, butil::memory_order_relaxed);
-        s->reset_parsing_context(nullptr);
-        return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
-    }
 
     uint32_t flags_be = 0;
     CHECK_EQ(source->cutn(&flags_be, HELLO_ACK_LEN), HELLO_ACK_LEN);
@@ -625,6 +620,7 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s
         rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF;
         ep->_state.store(FALLBACK_TCP, butil::memory_order_release);
         s->reset_parsing_context(nullptr);
+        rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
         return MakeParseError(PARSE_ERROR_TRY_OTHERS);
     }
 
@@ -633,16 +629,33 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s
                      << s->description();
         ep->_state.store(FAILED, butil::memory_order_relaxed);
         s->reset_parsing_context(nullptr);
+        rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
         return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
     }
 
     LOG_IF(INFO, FLAGS_rdma_trace_verbose)
         << "Server handshake ends (use rdma v" << ep->_handshake_version
         << ") on " << s->description();
     rdma_transport->_rdma_state = RdmaTransport::RDMA_ON;
-    ep->_state.store(ESTABLISHED, butil::memory_order_relaxed);
+    // Clear any residual TCP data so it cannot pollute the RDMA recv
+    // stream. HandleCompletion appends (not overwrites) to _read_buf,
+    // so leftover bytes would become a prefix to RDMA data and break
+    // parsing. This clear is safe because HandleCompletion only writes
+    // _read_buf after seeing ESTABLISHED (acquire), which is stored
+    // below (release) — strictly after this clear.
+    source->clear();
+    ep->_state.store(ESTABLISHED, butil::memory_order_release);
     s->reset_parsing_context(nullptr);
-    return MakeParseError(PARSE_ERROR_TRY_OTHERS);
+    rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
+    // Return NOT_ENOUGH_DATA (not TRY_OTHERS) so that CutInputMessage
+    // returns immediately without invoking other parsers on _read_buf.
+    // With TRY_OTHERS, each parser would call source->size() on _read_buf,
+    // racing with HandleCompletion which may be appending RDMA data to
+    // _read_buf concurrently. The preferred_index is left as the handshake
+    // parser, but this is self-correcting: the first PollCq-driven
+    // ProcessNewMessage will try the handshake parser, get TRY_OTHERS
+    // (magic mismatch), and switch to the correct parser.
+    return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
 }
 
 bool RdmaEndpoint::IsWritable() const {
@@ -916,16 +929,31 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
     }
     case IBV_WC_RECV: {  // recv completion
         // Please note that only the first wc.byte_len bytes is valid
+        ssize_t bytes_written = 0;
         if (wc.byte_len > 0) {
             if (wc.byte_len < (uint32_t)FLAGS_rdma_zerocopy_min_size) {
                 zerocopy = false;
             }
-            CHECK_NE(_state.load(butil::memory_order_relaxed), FALLBACK_TCP);
-            if (zerocopy) {
-                _rbuf[_rq_received].cutn(&_socket->_read_buf, wc.byte_len);
+            // Don't write to _read_buf until the handshake is fully done
+            // (ESTABLISHED). During the handshake (S_ACK_WAIT etc.), the
+            // main socket's OnNewMessages is driving the handshake via
+            // _read_buf; PollCq writing to _read_buf concurrently corrupts
+            // the IOBuf (non-thread-safe). Fall through to handle imm
+            // data, re-post recv WR, and send ack normally.
+            if (_state.load(butil::memory_order_acquire) != ESTABLISHED) {
+                LOG_EVERY_N(WARNING, 100)
+                    << "RDMA recv completion in non-ESTABLISHED state "
+                    << GetStateStr() << ", drop "
+                    << wc.byte_len << " bytes from "
+                    << _socket->description();
             } else {
-                // Copy data when the receive data is really small
-                _socket->_read_buf.append(_rbuf_data[_rq_received], wc.byte_len);
+                if (zerocopy) {
+                    _rbuf[_rq_received].cutn(&_socket->_read_buf, wc.byte_len);
+                } else {
+                    // Copy data when the receive data is really small
+                    _socket->_read_buf.append(_rbuf_data[_rq_received], wc.byte_len);
+                }
+                bytes_written = wc.byte_len;
             }
         }
         if (0 != (wc.wc_flags & IBV_WC_WITH_IMM) && wc.imm_data > 0) {
@@ -948,7 +976,7 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
         if (wc.byte_len > 0) {
             SendAck(1);
         }
-        return wc.byte_len;
+        return bytes_written;
     }
     default:
         // Some driver bugs may lead to unexpected completion opcode.
@@ -1593,12 +1621,19 @@ void RdmaEndpoint::PollCq(Socket* m) {
 
         // Just call PrcessNewMessage once for all of these CQEs.
         // Otherwise it may call too many bthread_flush to affect performance.
-        const int64_t received_us = butil::cpuwide_time_us();
-        const int64_t base_realtime = butil::gettimeofday_us() - received_us;
-        InputMessenger* messenger = static_cast<InputMessenger*>(s->user());
-        if (messenger->ProcessNewMessage(
-                    s.get(), bytes, false, received_us, base_realtime, last_msg) < 0) {
-            return;
+        // Only call when bytes > 0: when bytes == 0, HandleCompletion wrote
+        // nothing to _read_buf (e.g., IBV_WC_SEND completions, or IBV_WC_RECV
+        // dropped during handshake). Calling ProcessNewMessage with bytes == 0
+        // would still invoke CutInputMessage on _read_buf, racing with
+        // OnNewMessages which is driving the handshake via the same _read_buf.
+        if (bytes > 0) {
+            const int64_t received_us = butil::cpuwide_time_us();
+            const int64_t base_realtime = butil::gettimeofday_us() - received_us;
+            InputMessenger* messenger = static_cast<InputMessenger*>(s->user());
+            if (messenger->ProcessNewMessage(
+                        s.get(), bytes, false, received_us, base_realtime, last_msg) < 0) {
+                return;
+            }
         }
     }
 }

src/brpc/rdma/rdma_endpoint.h

@@ -124,6 +124,13 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&);
     // Whether the endpoint can send more data
     bool IsWritable() const;
 
+    // Whether the RDMA handshake has reached ESTABLISHED.
+    // Uses acquire load to pair with the release store in
+    // ExecuteServerHandshake / ProcessHandshakeAtClient.
+    bool IsEstablished() const {
+        return _state.load(butil::memory_order_acquire) == ESTABLISHED;
+    }
+
     // For debug
     void DebugInfo(std::ostream& os,
                    butil::StringPiece connector = "\n") const;

src/brpc/rdma_transport.cpp

@@ -70,10 +70,19 @@ int RdmaTransport::Reset(int32_t expected_nref) {
     if (_rdma_ep) {
         _rdma_ep->Reset();
         _rdma_state = RDMA_UNKNOWN;
+        if (_socket->CreatedByConnect()) {
+            _on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp;
+        } else {
+            _on_edge_trigger = InputMessenger::OnNewMessages;
+        }
     }
     return 0;
 }
 
+bool RdmaTransport::ShouldStopReading() const {
+    return _rdma_ep && _rdma_ep->IsEstablished();
+}
+
 std::shared_ptr<AppConnect> RdmaTransport::Connect() {
     if (_default_connect == nullptr) {
         return  std::make_shared<rdma::RdmaConnect>();

src/brpc/rdma_transport.h

@@ -41,6 +41,7 @@ friend class rdma::RdmaHandshakeServerV3;
     void ProcessEvent(bthread_attr_t attr) override;
     void QueueMessage(InputMessageClosure& inputMsg, int* num_bthread_created, bool last_msg) override;
     void Debug(std::ostream &os) override;
+    bool ShouldStopReading() const override;
     rdma::RdmaEndpoint* GetRdmaEp() {
         CHECK(_rdma_ep != nullptr);
         return _rdma_ep;

src/brpc/transport.h

@@ -51,6 +51,11 @@ class Transport {
     virtual void QueueMessage(InputMessageClosure& input_msg, int* num_bthread_created, bool last_msg) = 0;
     virtual void Debug(std::ostream &os) = 0;
 
+    // Returns true if OnNewMessages should stop its read loop immediately
+    // (e.g., RDMA transport after handshake completes and edge trigger
+    // is switched to OnNewDataFromTcp). Default: never stop.
+    virtual bool ShouldStopReading() const { return false; }
+
     bool HasOnEdgeTrigger() {
         return _on_edge_trigger != nullptr;
     }