Skip to content

fix _read_buf race between PollCq and OnNewMessages in RDMA server #3480

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bzs1118 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix _read_buf race between PollCq and OnNewMessages in RDMA server #3480

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion src/brpc/input_messenger.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,14 @@ 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.
if (m->_transport->ShouldStopReading()) {
return;
}
Comment on lines +377 to +379
}

if (read_eof) {
Expand Down
74 changes: 52 additions & 22 deletions src/brpc/rdma/rdma_endpoint.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
}

Expand All @@ -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);
Expand All @@ -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);
}

Expand All @@ -633,16 +629,28 @@ 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 OnNewMessages stops
// processing _read_buf immediately, before PollCq starts writing RDMA
// data into _read_buf.
return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
Comment on lines +649 to +653

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed this by adding a Transport::ShouldStopReading() virtual mechanism

}

bool RdmaEndpoint::IsWritable() const {
Expand Down Expand Up @@ -916,16 +924,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;
}
Comment on lines +938 to 952

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +938 to 952

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

}
if (0 != (wc.wc_flags & IBV_WC_WITH_IMM) && wc.imm_data > 0) {
Expand All @@ -948,7 +971,7 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) {
if (wc.byte_len > 0) {
SendAck(1);
}
return wc.byte_len;
return bytes_written;
Comment on lines 971 to +974

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same

}
default:
// Some driver bugs may lead to unexpected completion opcode.
Expand Down Expand Up @@ -1593,12 +1616,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;
}
}
}
}
Expand Down
7 changes: 7 additions & 0 deletions src/brpc/rdma/rdma_endpoint.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
9 changes: 9 additions & 0 deletions src/brpc/rdma_transport.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Comment on lines +82 to +84

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added RdmaEndpoint::IsEstablished() — return _state.load(acquire) == ESTABLISHED


std::shared_ptr<AppConnect> RdmaTransport::Connect() {
if (_default_connect == nullptr) {
return std::make_shared<rdma::RdmaConnect>();
Expand Down
1 change: 1 addition & 0 deletions src/brpc/rdma_transport.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
5 changes: 5 additions & 0 deletions src/brpc/transport.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Comment on lines +54 to +57

@bzs1118 bzs1118 Aug 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


bool HasOnEdgeTrigger() {
return _on_edge_trigger != nullptr;
}
Expand Down
Loading