fix _read_buf race between PollCq and OnNewMessages in RDMA server

src/brpc/input_messenger.cpp

@@ -210,8 +210,15 @@ int InputMessenger::ProcessNewMessage(
             if (pr.error() == PARSE_ERROR_NOT_ENOUGH_DATA) {
                 // incomplete message, re-read.
                 // However, some buffer may have been consumed
-                // under protocols like HTTP. Record this size
-                m->_last_msg_size += (last_size - m->_read_buf.length());
+                // under protocols like HTTP. Record this size.
+                // Skip the length() read when the transport is stopping
+                // (e.g., RDMA handshake just completed and ESTABLISHED is
+                // published), because HandleCompletion may be appending
+                // to _read_buf concurrently. The last_msg_size stat is
+                // metrics-only and safe to skip.
+                if (!m->_transport->ShouldStopReading()) {
+                    m->_last_msg_size += (last_size - m->_read_buf.length());
+                }
                 break;
             } else if (pr.error() == PARSE_ERROR_TRY_OTHERS) {
                 LOG(WARNING)
@@ -369,7 +376,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;
     }