From d911f4c948d3ce1fee97e4d35df6bbd8a4596fe0 Mon Sep 17 00:00:00 2001 From: shengtiedan Date: Sun, 23 Aug 2026 20:58:45 +0800 Subject: [PATCH] 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. --- src/brpc/input_messenger.cpp | 25 +++++++++-- src/brpc/rdma/rdma_endpoint.cpp | 79 ++++++++++++++++++++++++--------- src/brpc/rdma/rdma_endpoint.h | 7 +++ src/brpc/rdma_transport.cpp | 9 ++++ src/brpc/rdma_transport.h | 1 + src/brpc/transport.h | 5 +++ 6 files changed, 101 insertions(+), 25 deletions(-) diff --git a/src/brpc/input_messenger.cpp b/src/brpc/input_messenger.cpp index 81154be134..7dbdd3a75a 100644 --- a/src/brpc/input_messenger.cpp +++ b/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) { diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index fd70fb5c2b..c37f2ed02f 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/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,6 +629,7 @@ 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); } @@ -640,9 +637,25 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s << "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(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(s->user()); + if (messenger->ProcessNewMessage( + s.get(), bytes, false, received_us, base_realtime, last_msg) < 0) { + return; + } } } } diff --git a/src/brpc/rdma/rdma_endpoint.h b/src/brpc/rdma/rdma_endpoint.h index 388e31d78e..de36194a35 100644 --- a/src/brpc/rdma/rdma_endpoint.h +++ b/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; diff --git a/src/brpc/rdma_transport.cpp b/src/brpc/rdma_transport.cpp index ee5151c3a5..69b3eaf13a 100644 --- a/src/brpc/rdma_transport.cpp +++ b/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 RdmaTransport::Connect() { if (_default_connect == nullptr) { return std::make_shared(); diff --git a/src/brpc/rdma_transport.h b/src/brpc/rdma_transport.h index 1d78fbb430..2c76603796 100644 --- a/src/brpc/rdma_transport.h +++ b/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; diff --git a/src/brpc/transport.h b/src/brpc/transport.h index edef24879c..671e83b1b1 100644 --- a/src/brpc/transport.h +++ b/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; }