From 0b0823ace6ede3d8132e5fe1baf5f256bd180924 Mon Sep 17 00:00:00 2001 From: BGQ99 <1132767344@qq.com> Date: Thu, 15 Oct 2026 10:34:51 +0800 Subject: [PATCH] Fix URMA send buffer reclamation order --- src/brpc/urma/urma_endpoint.cpp | 22 +++++++-- test/brpc_urma_unittest.cpp | 87 +++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 5 deletions(-) diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 26c7a1b9dc..c7c2f05ef6 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -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) { + 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) @@ -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(); diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index d27367a504..87e8dfb77c 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -21,9 +21,11 @@ #include #if BRPC_WITH_URMA +#include #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" @@ -41,6 +43,91 @@ extern butil::atomic 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 - + 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)); + })); + + 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. // ---------------------------------------------------------------------------