fix _read_buf race between PollCq and OnNewMessages in RDMA server - #3480
fix _read_buf race between PollCq and OnNewMessages in RDMA server#3480bzs1118 wants to merge 1 commit into
Conversation
Sorry, something went wrong.
There was a problem hiding this comment.
Pull request overview
This PR addresses an RDMA server crash caused by concurrent access to Socket::_read_buf from two bthreads (PollCq via HandleCompletion() and the TCP-driven OnNewMessages handshake/fallback path). It does so by switching the edge-trigger handler after handshake, adding a state gate in RDMA recv completion handling, and relaxing overly-strict ACK parsing so TCP-fallback clients can coalesce ACK + first request.
Changes:
- Switch server-side edge trigger to RdmaEndpoint::OnNewDataFromTcp on all handshake terminal paths, and change the ESTABLISHED return to PARSE_ERROR_NOT_ENOUGH_DATA to stop OnNewMessages from further parsing.
- In HandleCompletion(IBV_WC_RECV), skip touching _read_buf unless state is ESTABLISHED.
- In RdmaTransport::Reset(), restore the edge-trigger handler (added in this PR).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| src/brpc/rdma/rdma_endpoint.cpp | Adjusts server handshake terminal behavior/returns and adds a state gate in RDMA recv completion to avoid concurrent _read_buf mutation. |
| src/brpc/rdma_transport.cpp | Resets the transport edge-trigger callback during Reset() to support re-handshake behavior. |
src/brpc/rdma/rdma_endpoint.cpp:935
- This branch intentionally drops RDMA payloads received before the server reaches ESTABLISHED. The issue description notes this can happen after BringUpQp (client may start sending RDMA data before server finishes processing the ACK), so dropping here risks losing the first real RPC bytes and causing timeouts/hangs. Consider buffering these bytes until the handshake finalizes, or preventing the main socket from touching _read_buf during S_ACK_WAIT so PollCq can safely append without loss.
// 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).
if (_state.load(butil::memory_order_acquire) != ESTABLISHED) {
LOG(WARNING) << "RDMA recv completion in non-ESTABLISHED state "
<< GetStateStr() << ", drop "
<< wc.byte_len << " bytes from "
<< _socket->description();
src/brpc/rdma/rdma_endpoint.cpp:936
- In the non-ESTABLISHED recv-completion path, PostRecv() failures are ignored, which can leave the receive queue un-posted and stall the connection. Also, a plain LOG(WARNING) here can spam logs if the peer sends early/invalid RDMA traffic. It’s safer to (a) check PostRecv() and propagate failure and (b) rate-limit the warning log.
LOG(WARNING) << "RDMA recv completion in non-ESTABLISHED state "
<< GetStateStr() << ", drop "
<< wc.byte_len << " bytes from "
<< _socket->description();
PostRecv(1, zerocopy);
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Sorry, something went wrong.
| rdma_transport->_on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp; | ||
| // Return NOT_ENOUGH_DATA (not TRY_OTHERS) so that OnNewMessages stops | ||
| // processing _read_buf immediately, before PollCq starts writing RDMA | ||
| // data into _read_buf. | ||
| return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); |
There was a problem hiding this comment.
addressed this by adding a Transport::ShouldStopReading() virtual mechanism
Sorry, something went wrong.
| if (_state.load(butil::memory_order_acquire) != ESTABLISHED) { | ||
| LOG(WARNING) << "RDMA recv completion in non-ESTABLISHED state " | ||
| << GetStateStr() << ", drop " | ||
| << wc.byte_len << " bytes from " | ||
| << _socket->description(); | ||
| PostRecv(1, zerocopy); | ||
| return 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; } |
There was a problem hiding this comment.
If ABI stability becomes a requirement in the future, this can be refactored to use a non-virtual capability query, but for now I think it's OK
Sorry, something went wrong.
| bool RdmaTransport::ShouldStopReading() const { | ||
| return _rdma_state == RDMA_ON; | ||
| } |
There was a problem hiding this comment.
Added RdmaEndpoint::IsEstablished() — return _state.load(acquire) == ESTABLISHED
Sorry, something went wrong.
| if (_state.load(butil::memory_order_acquire) != ESTABLISHED) { | ||
| LOG(WARNING) << "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; | ||
| } |
There was a problem hiding this comment.
The race window is between BringUpQp (QP→RTS) and state.store(ESTABLISHED) — the server only needs to process the 4-byte ACK (one cutn call), So at most 0–1 RDMA messages arrive(or lost). Buffer the data outside _read_buf could add significant complexity.
Plus, baidu_std is request-response: the client detects the missing response via timeout and retries. The retry succeeds because the server is now ESTABLISHED.
Sorry, something went wrong.
| if (wc.byte_len > 0) { | ||
| SendAck(1); | ||
| } | ||
| return wc.byte_len; | ||
| return bytes_written; |
There was a problem hiding this comment.
same
Sorry, something went wrong.
| LOG(WARNING) << "RDMA recv completion in non-ESTABLISHED state " | ||
| << GetStateStr() << ", drop " | ||
| << wc.byte_len << " bytes from " | ||
| << _socket->description(); |
| if (_state.load(butil::memory_order_acquire) != ESTABLISHED) { | ||
| LOG(WARNING) << "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; | ||
| } |
There was a problem hiding this comment.
This is already handled by PostRecv. For zerocopy mode, PostRecv unconditionally calls _rbuf[_rq_received].clear() before allocating a new block. For non-zerocopy mode, _rbuf_data[_rq_received] is a raw pointer to a fixed-size buffer — the next recv completion simply overwrites it, so no cleanup is needed.
Sorry, something went wrong.
| if (m->_transport->ShouldStopReading()) { | ||
| return; | ||
| } |
There was a problem hiding this comment.
done
Sorry, something went wrong.
| if (m->_transport->ShouldStopReading()) { | ||
| while (m->MoreReadEvents(&progress)) {} | ||
| return; | ||
| } |
| 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 OnNewMessages stops | ||
| // processing _read_buf immediately, before PollCq starts writing RDMA | ||
| // data into _read_buf. | ||
| return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); |
There was a problem hiding this comment.
The preferred_index issue is self-correcting: the first PollCq-driven ProcessNewMessage tries the handshake parser, gets a magic mismatch → TRY_OTHERS, and switches to the correct parser. This is a one-time cost, and should be ok
Sorry, something went wrong.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Suppressed comments (1)src/brpc/rdma/rdma_endpoint.cpp:615
- The handshake parsing behavior here was changed to (a) accept ACK coalesced with the first RPC bytes by draining exactly 4 bytes, and (b) transition to ESTABLISHED while switching the edge trigger + relying on stop-reading logic. There are existing RDMA unit tests in test/brpc_rdma_unittest.cpp, but no test appears to cover these coalescing / transition cases, which makes regressions hard to catch.
Please add a unit test that simulates a TCP-fallback client sending the 4-byte ACK followed immediately by a valid baidu_std request in the same write/read, and asserts the server does not drop the connection and the request is parsed/handled correctly (matching the intended behavior described in the PR).
// Phase 2: drain the 4B ACK and finalize.
if (source->size() < HELLO_ACK_LEN) {
return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
}
uint32_t flags_be = 0;
CHECK_EQ(source->cutn(&flags_be, HELLO_ACK_LEN), HELLO_ACK_LEN);
uint32_t flags = butil::NetToHost32(flags_be);
Sorry, something went wrong.
| 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; |
There was a problem hiding this comment.
Added a ShouldStopReading() check before the length() read in ProcessNewMessage to skip it
Sorry, something went wrong.
| // 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 |
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.
chenBright
left a comment
There was a problem hiding this comment.
The overall direction doesn't look right to me. According to the design of #3350, what this needs is a change in the flow that removes the concurrency and the contention in the first place, rather than arbitrating it. Absent that, reverting to plain OnNewDataFromTcp would be preferable to landing this.
Sorry, something went wrong.
| if (m->_transport->ShouldStopReading()) { | ||
| while (m->MoreReadEvents(&progress)) {} | ||
| if (read_eof) { | ||
| m->SetEOF(); | ||
| } | ||
| return; | ||
| } |
There was a problem hiding this comment.
-
ShouldStopReading() is more like a patch than an abstraction, which makes this hard to maintain going forward.
- "stop the read loop immediately" leaves undefined whether the caller must drain _nevent, call SetEOF(), flush last_msg, or skip the _read_buf.length() read.
- Introducing any new state later (a reconnect/migration state, or an intermediate one between the ACK and ESTABLISHED) means remembering to update the gate in HandleCompletion, the condition here, and every terminal path that assigns _on_edge_trigger. The five scattered _on_edge_trigger assignments are already easy to forget: omitting one on a newly added terminal path silently reintroduces the race.
-
Draining _nevent discards events under EPOLLET, the server never learns the peer is gone. .
Sorry, something went wrong.
| 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(); |
There was a problem hiding this comment.
Dropping RDMA payload is not acceptable.
Sorry, something went wrong.
|
@bzs1118 Thank you for your feedback and PR. To quickly fix this issue and avoid blocking the release of a new version, I submitted a version without concurrency. Please try #3505 to see if it resolves the problem. If you have any questions about the PR, please let me know. |
Sorry, something went wrong.
The server-side RDMA socket's _read_buf is accessed by two independent bthreads: PollCq (CQ socket) writes RDMA data via HandleCompletion, 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:
Switch edge trigger to OnNewDataFromTcpin ALL ExecuteServerHandshake end paths (ESTABLISHED + 5 failure paths). OnNewDataFromTcpchecks 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.
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).
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 OnNewMessages stops processing _read_buf before PollCq starts writing.
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.
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 _rdma_ep->IsEstablished() — a new RdmaEndpoint method that performs 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 while OnNewMessages is driving the handshake on the same _read_buf.
What problem does this PR solve?
Issue Number: #3479
Problem Summary:
What is changed and the side effects?
Changed:
Side effects:
Performance effects: HandleCompletion adds one atomic load per recv completion (negligible)
ShouldStopReading() adds one virtual call per ProcessNewMessage iteration (negligible);
Breaking backward compatibility: No. The OnNewDataFromTcp behavior in FALLBACK_TCP state is unchanged (delegates to OnNewMessages). The removed > HELLO_ACK_LEN check aligns with FallbackServerHandshake's existing behavior. ShouldStopReading() defaults to false for non-RDMA transports.
Check List: