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 1 commit 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
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
22 changes: 17 additions & 5 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,6 +952,23 @@ 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);
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.

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.
while (true) {
if (old >= _local_window_capacity) {
LOG(WARNING)
Expand All @@ -969,11 +986,6 @@ ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) {
break;
}
}
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();
if (_remote_rq_window_size.load(butil::memory_order_relaxed) >=
_local_window_capacity / 8) {
_socket->WakeAsEpollOut();
Expand Down
87 changes: 87 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,91 @@ 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();
endpoint._sq_window_size.store(0, butil::memory_order_relaxed);
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

ASSERT_EQ(0, endpoint.HandleCompletion(SendCompletion()));
EXPECT_TRUE(released);
EXPECT_TRUE(endpoint._sbuf[0].empty());
EXPECT_EQ(1, endpoint._sq_sent);
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();
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();
endpoint._sq_window_size.store(endpoint._local_window_capacity,
butil::memory_order_relaxed);
ASSERT_EQ(0, endpoint._sbuf[0].append("keep"));

EXPECT_EQ(-1, endpoint.HandleCompletion(SendCompletion()));
EXPECT_EQ(EPROTO, errno);
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