diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 26c7a1b9dc..fca97fbe86 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -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(old + wnd), - butil::memory_order_relaxed)) { - break; - } + 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. + _sq_window_size.fetch_add(wnd, butil::memory_order_relaxed); 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. // ---------------------------------------------------------------------------