Skip to content

Fix URMA send buffer reclamation order #3626

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
BGQ99 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Fix URMA send buffer reclamation order #3626

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6c05373
Fix URMA send buffer reclamation order
BGQ99 Oct 19, 2026
45a7cfb
Simplify URMA send completion credit update
BGQ99 Oct 21, 2026
d05cf78
Clear errno before URMA completion error assertion
BGQ99 Oct 22, 2026
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
28 changes: 12 additions & 16 deletions src/brpc/urma/urma_endpoint.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -952,28 +952,24 @@ ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) {
uint16_t wnd = 1; // We signal every WR (complete_enable=1).
uint16_t old =
_sq_window_size.load(butil::memory_order_relaxed);
while (true) {
if (old >= _local_window_capacity) {
LOG(WARNING)
<< "URMA send completion exceeds SQ window: old=" << old
<< " increment=" << wnd
<< " capacity=" << _local_window_capacity
<< " user_ctx=" << cr.user_ctx
<< " on " << _socket->description();
errno = EPROTO;
return -1;
}
if (_sq_window_size.compare_exchange_weak(
old, static_cast<uint16_t>(old + wnd),
butil::memory_order_relaxed)) {
break;
}
if (old >= _local_window_capacity) {

@wwbmmm wwbmmm Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] This capacity check plus its warning is now duplicated: the same check is still inside the while (true) CAS loop below. The in-loop branch is now only reachable after the buffer has already been cleared and _sq_sent advanced, in which case the credit is dropped while the state was already mutated. Since _sq_window_size is only incremented here (writers only decrement it), consider doing the single pre-check and then publishing with _sq_window_size.fetch_add(wnd, butil::memory_order_relaxed) after the barrier, removing the now-redundant retry loop so reclamation and credit publication stay in one place.

@BGQ99 BGQ99 Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Updated in 45a7cfb. Kept the capacity check before buffer reclamation and replaced the subsequent CAS loop with fetch_add, preserving the memory barrier before credit publication.

Comment on lines 952 to +955
LOG(WARNING)
<< "URMA send completion exceeds SQ window: old=" << old
<< " increment=" << wnd
<< " capacity=" << _local_window_capacity
<< " user_ctx=" << cr.user_ctx
<< " on " << _socket->description();
errno = EPROTO;
return -1;
}
for (uint16_t i = 0; i < wnd; ++i) {
_sbuf[_sq_sent].clear();
_sq_sent = (_sq_sent + 1) % (_sq_size - RESERVED_WR_NUM);
}
butil::subtle::MemoryBarrier();
// Reclaim the buffer before publishing the SQ credit. Otherwise the
// writer may reuse this slot while the completion thread clears it.
_sq_window_size.fetch_add(wnd, butil::memory_order_relaxed);
Comment on lines 952 to +972
Comment on lines +955 to +972
if (_remote_rq_window_size.load(butil::memory_order_relaxed) >=
_local_window_capacity / 8) {
_socket->WakeAsEpollOut();
Expand Down
88 changes: 88 additions & 0 deletions test/brpc_urma_unittest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,11 @@
#include <gflags/gflags.h>

#if BRPC_WITH_URMA
#include <cerrno>
#include "butil/atomicops.h"
#include "butil/sys_byteorder.h"
#include "urma_api.h"
#include "brpc/urma/urma_endpoint.h"
#include "brpc/urma/urma_handshake.h"
#include "brpc/urma/urma_handshake.pb.h"
#include "brpc/urma/urma_helper.h"
Expand All @@ -41,6 +43,92 @@ extern butil::atomic<bool> g_urma_available;
} // namespace urma
} // namespace brpc

class UrmaSendCompletionTest : public ::testing::Test {
protected:
void SetUp() override {
SocketOptions options;
SocketId id;
ASSERT_EQ(0, Socket::Create(options, &id));
ASSERT_EQ(0, Socket::Address(id, &_socket));
}

void TearDown() override {
if (_socket) {
_socket->SetFailed();
}
}

static urma_cr_t SendCompletion() {
urma_cr_t cr{};
cr.status = URMA_CR_SUCCESS;
cr.user_ctx = 1;
return cr;
}

SocketUniquePtr _socket;
};

TEST_F(UrmaSendCompletionTest, reclaim_buffer_before_publishing_credit) {
char payload = 'x';
bool released = false;
urma::UrmaEndpoint endpoint(_socket.get());
endpoint._sbuf.resize(endpoint._sq_size -

@wwbmmm wwbmmm Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[replied by brpc-oncall robot] These tests cannot compile when the tree is built with WITH_URMA=ON: _sbuf, _sq_size, _sq_sent, _sq_window_size, _local_window_capacity and HandleCompletion() are all private in UrmaEndpoint (see src/brpc/urma/urma_endpoint.h: everything after the private: section that starts at the State enum), and the fixture is not declared as a friend. CI does not catch this because test/brpc_urma_unittest.cpp compiles to an empty translation unit unless BRPC_WITH_URMA is 1. Please grant access explicitly, e.g. add #include <gtest/gtest_prod.h> plus FRIEND_TEST(UrmaSendCompletionTest, reclaim_buffer_before_publishing_credit); (and the other two test names) in UrmaEndpoint, or expose a small test-only accessor, so that the URMA unit-test build actually compiles.

@BGQ99 BGQ99 Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks for checking. The existing test build disables C++ access checks with -fno-access-control. In CMake, this is set on brpc_test_config, which the URMA unit-test target links against; test/Makefile includes the same flag in CXXFLAGS. Therefore, these private-member accesses are permitted by the existing CMake and Make test configurations, including when URMA is enabled.

urma::UrmaEndpoint::RESERVED_WR_NUM);
endpoint._local_window_capacity = endpoint._sbuf.size();
Comment on lines +75 to +77
endpoint._sq_window_size.store(0, butil::memory_order_relaxed);
Comment on lines +74 to +78
ASSERT_EQ(0, endpoint._sbuf[0].append_user_data(
&payload, sizeof(payload), [&](void*) {
released = true;
// Observe the credit at the exact point the old buffer is freed.
// Publishing first lets the writer reuse a slot still being cleared.
EXPECT_EQ(0, endpoint._sq_window_size.load(
butil::memory_order_relaxed));
}));
Comment on lines +75 to +86

@BGQ99 BGQ99 Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

These tests use the existing test build’s -fno-access-control flag, without changing production member visibility. Direct access lets the ordering test observe the SQ credit from the buffer deleter at the exact point of reclamation. Adding FRIEND_TEST would make access explicit, but would not remove the tests’ dependence on these internals. I’ve kept the existing test convention for this focused regression fix.


ASSERT_EQ(0, endpoint.HandleCompletion(SendCompletion()));
EXPECT_TRUE(released);
EXPECT_TRUE(endpoint._sbuf[0].empty());
EXPECT_EQ(1, endpoint._sq_sent);

@BGQ99 BGQ99 Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Same rationale as above; this duplicates the other comment.

EXPECT_EQ(1, endpoint._sq_window_size.load(butil::memory_order_relaxed));
}

TEST_F(UrmaSendCompletionTest, reclaim_last_slot_and_wrap_cursor) {
urma::UrmaEndpoint endpoint(_socket.get());
endpoint._sbuf.resize(endpoint._sq_size -
urma::UrmaEndpoint::RESERVED_WR_NUM);
endpoint._local_window_capacity = endpoint._sbuf.size();
Comment on lines +97 to +99
endpoint._sq_sent = endpoint._sbuf.size() - 1;
const uint16_t old_credit = endpoint._local_window_capacity - 1;
endpoint._sq_window_size.store(old_credit, butil::memory_order_relaxed);
ASSERT_EQ(0, endpoint._sbuf[endpoint._sq_sent].append("last"));
ASSERT_EQ(0, endpoint._sbuf[0].append("next"));

ASSERT_EQ(0, endpoint.HandleCompletion(SendCompletion()));
EXPECT_EQ(0, endpoint._sq_sent);
EXPECT_TRUE(endpoint._sbuf.back().empty());
EXPECT_EQ("next", endpoint._sbuf[0].to_string());
EXPECT_EQ(endpoint._local_window_capacity,
endpoint._sq_window_size.load(butil::memory_order_relaxed));
}

TEST_F(UrmaSendCompletionTest, reject_full_window_before_reclaiming_buffer) {
urma::UrmaEndpoint endpoint(_socket.get());
endpoint._sbuf.resize(endpoint._sq_size -
urma::UrmaEndpoint::RESERVED_WR_NUM);
endpoint._local_window_capacity = endpoint._sbuf.size();
Comment on lines +116 to +118
endpoint._sq_window_size.store(endpoint._local_window_capacity,
butil::memory_order_relaxed);
ASSERT_EQ(0, endpoint._sbuf[0].append("keep"));

errno = 0;
EXPECT_EQ(-1, endpoint.HandleCompletion(SendCompletion()));
EXPECT_EQ(EPROTO, errno);
Comment on lines +124 to +125
EXPECT_EQ("keep", endpoint._sbuf[0].to_string());
EXPECT_EQ(0, endpoint._sq_sent);
EXPECT_EQ(endpoint._local_window_capacity,
endpoint._sq_window_size.load(butil::memory_order_relaxed));
}

// ---------------------------------------------------------------------------
// v2 binary HelloMessage: serialize + deserialize round-trips.
// ---------------------------------------------------------------------------
Expand Down
Loading