From 00e3df37c97c0981adc35f60bd1ce82daf6e9891 Mon Sep 17 00:00:00 2001 From: BGQ99 <1132767344@qq.com> Date: Tue, 20 Oct 2026 16:28:29 +0800 Subject: [PATCH 1/3] Fix URMA polling shutdown lifetime --- src/brpc/urma/urma_endpoint.cpp | 58 ++++++++---- src/brpc/urma/urma_endpoint.h | 4 +- src/brpc/urma/urma_helper.cpp | 21 +++-- test/brpc_urma_unittest.cpp | 154 ++++++++++++++++++++++++++++++++ 4 files changed, 215 insertions(+), 22 deletions(-) diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 26c7a1b9dc..b24714ba13 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -38,6 +38,7 @@ #include "butil/iobuf.h" #include "butil/logging.h" #include "butil/macros.h" +#include "butil/memory/singleton_on_pthread_once.h" #include "butil/sys_byteorder.h" #include "butil/time.h" #include "bthread/bthread.h" @@ -112,7 +113,12 @@ static int PreparedJettyCount() { return static_cast(max_prepared); } -std::vector UrmaEndpoint::_poller_groups; +std::vector& UrmaEndpoint::PollerGroups() { + // Endpoints owned by process-wide Socket state may be destroyed after + // ordinary static objects. Keep the poller queues alive until process + // termination so late endpoint cleanup never enqueues into freed storage. + return *butil::get_leaky_singleton>(); +} // ============================================================================ // UrmaResource lifecycle. @@ -1616,12 +1622,13 @@ int UrmaEndpoint::PollingModeInitialize( if (!FLAGS_urma_use_polling) { return 0; } - if (tag >= _poller_groups.size() || - _poller_groups[tag].pollers.empty()) { + auto& poller_groups = PollerGroups(); + if (tag >= poller_groups.size() || + poller_groups[tag].pollers.empty()) { errno = EINVAL; return -1; } - auto& group = _poller_groups[tag]; + auto& group = poller_groups[tag]; bool expected = false; if (!group.running.compare_exchange_strong(expected, true)) { return 0; @@ -1705,11 +1712,14 @@ int UrmaEndpoint::PollingModeInitialize( } void UrmaEndpoint::PollingModeRelease(bthread_tag_t tag) { - if (!FLAGS_urma_use_polling || tag >= _poller_groups.size()) { + auto& poller_groups = PollerGroups(); + if (!FLAGS_urma_use_polling || tag >= poller_groups.size()) { + return; + } + auto& group = poller_groups[tag]; + if (!group.running.exchange(false, butil::memory_order_acq_rel)) { return; } - auto& group = _poller_groups[tag]; - group.running.store(false, butil::memory_order_relaxed); for (auto& poller : group.pollers) { if (poller.tid != INVALID_BTHREAD) { bthread_join(poller.tid, nullptr); @@ -1719,14 +1729,19 @@ void UrmaEndpoint::PollingModeRelease(bthread_tag_t tag) { } void UrmaEndpoint::PollerAddCqSid() { - if (_cq_sid == INVALID_SOCKET_ID || _poller_groups.empty()) { + auto& poller_groups = PollerGroups(); + if (_cq_sid == INVALID_SOCKET_ID || poller_groups.empty()) { return; } _poller_tag = bthread_self_tag(); - if (_poller_tag >= _poller_groups.size()) { + if (_poller_tag >= poller_groups.size()) { return; } - auto& pollers = _poller_groups[_poller_tag].pollers; + auto& group = poller_groups[_poller_tag]; + if (!group.running.load(butil::memory_order_acquire)) { + return; + } + auto& pollers = group.pollers; if (pollers.empty()) { return; } @@ -1737,11 +1752,16 @@ void UrmaEndpoint::PollerAddCqSid() { } void UrmaEndpoint::PollerRemoveCqSid() { - if (_cq_sid == INVALID_SOCKET_ID || _poller_groups.empty() || - _poller_tag >= _poller_groups.size()) { + auto& poller_groups = PollerGroups(); + if (_cq_sid == INVALID_SOCKET_ID || poller_groups.empty() || + _poller_tag >= poller_groups.size()) { return; } - auto& pollers = _poller_groups[_poller_tag].pollers; + auto& group = poller_groups[_poller_tag]; + if (!group.running.load(butil::memory_order_acquire)) { + return; + } + auto& pollers = group.pollers; if (pollers.empty()) { return; } @@ -1754,7 +1774,8 @@ void UrmaEndpoint::PollerRemoveCqSid() { int UrmaEndpoint::GlobalInitialize() { // Pre-allocate the prepared jetty pool. Skipped if URMA init is skipped // (unit-test mode). - if (FLAGS_urma_use_polling && _poller_groups.empty()) { + auto& poller_groups = PollerGroups(); + if (FLAGS_urma_use_polling && poller_groups.empty()) { if (FLAGS_urma_poller_num <= 0) { LOG(ERROR) << "urma_poller_num must be positive"; errno = EINVAL; @@ -1764,7 +1785,7 @@ int UrmaEndpoint::GlobalInitialize() { if (ntags == 0) { ntags = 1; } - _poller_groups = std::vector(ntags); + poller_groups = std::vector(ntags); } if (g_prepared_cnt > 0) { return 0; @@ -1838,7 +1859,12 @@ void UrmaEndpoint::GlobalRelease() { } g_prepared_cnt = 0; } - for (size_t tag = 0; tag < _poller_groups.size(); ++tag) { + GlobalPollingModeRelease(); +} + +void UrmaEndpoint::GlobalPollingModeRelease() { + auto& poller_groups = PollerGroups(); + for (size_t tag = 0; tag < poller_groups.size(); ++tag) { PollingModeRelease(static_cast(tag)); } } diff --git a/src/brpc/urma/urma_endpoint.h b/src/brpc/urma/urma_endpoint.h index e5aed5ac66..5e99e95f2d 100644 --- a/src/brpc/urma/urma_endpoint.h +++ b/src/brpc/urma/urma_endpoint.h @@ -157,6 +157,8 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { std::function init_fn, std::function release_fn); static void PollingModeRelease(bthread_tag_t tag); + // Stop every polling group without releasing the global IOBuf pool. + static void GlobalPollingModeRelease(); // ---- Handshake IO helpers (also used by urma_handshake.cpp) ---- // Read at most @len bytes from the TCP fd into @data; waits on _read_butex @@ -338,7 +340,7 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { std::vector pollers; butil::atomic running; }; - static std::vector _poller_groups; + static std::vector& PollerGroups(); bthread_tag_t _poller_tag{0}; DISALLOW_COPY_AND_ASSIGN(UrmaEndpoint); diff --git a/src/brpc/urma/urma_helper.cpp b/src/brpc/urma/urma_helper.cpp index ded662fb22..7af22aa6fe 100644 --- a/src/brpc/urma/urma_helper.cpp +++ b/src/brpc/urma/urma_helper.cpp @@ -460,6 +460,13 @@ static void GlobalRelease() { g_device = nullptr; } +static void ReleasePollingModeAtExit() { + // Match RDMA's shutdown ordering for pollers, but deliberately leave the + // URMA IOBuf pool mapped because thread-local IOBuf blocks may outlive + // atexit handlers. + UrmaEndpoint::GlobalPollingModeRelease(); +} + // ============================================================================ // Global initialization. // ============================================================================ @@ -620,12 +627,16 @@ static bool GlobalUrmaInitializeImpl() { return false; } + if (FLAGS_urma_use_polling) { + atexit(ReleasePollingModeAtExit); + } + g_urma_available.store(true, butil::memory_order_release); - // Do not register GlobalRelease with atexit. IOBuf keeps blocks in - // thread-local chains whose destructors may run after atexit handlers. - // Unmapping the registered pool here would leave those TLS chains - // pointing into unmapped memory. The process reclaims global URMA - // resources on exit; GlobalRelease remains available for init rollback. + // Register only the polling shutdown above, not GlobalRelease. IOBuf keeps + // blocks in thread-local chains whose destructors may run after atexit + // handlers. Unmapping the registered pool here would leave those TLS + // chains pointing into unmapped memory. The process reclaims the remaining + // global URMA resources; GlobalRelease remains available for init rollback. LOG(INFO) << "URMA initialized: device=" << device_name << " bonding=" << g_is_bonding_device << " max_sge=" << g_max_sge diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index d27367a504..c8c9f614e5 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -21,9 +21,12 @@ #include #if BRPC_WITH_URMA +#include +#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" @@ -31,16 +34,167 @@ using namespace brpc; +DECLARE_int32(task_group_ntags); + namespace brpc { namespace urma { DECLARE_int32(urma_client_handshake_version); +DECLARE_int32(urma_buffer_size); +DECLARE_int32(urma_buffer_count); +DECLARE_int32(urma_sq_size); +DECLARE_int32(urma_rq_size); +DECLARE_int32(urma_prepared_jetty_cnt); extern bool g_skip_urma_init; extern butil::atomic g_urma_available; } // namespace urma } // namespace brpc +class UrmaPollingTest : public ::testing::TestWithParam { +protected: + void SetUp() override { + urma::FLAGS_urma_use_polling = true; + urma::FLAGS_urma_poller_num = 2; + urma::FLAGS_urma_disable_bthread = GetParam(); + ASSERT_EQ(0, urma::UrmaEndpoint::GlobalInitialize()); + } + + void TearDown() override { + urma::UrmaEndpoint::PollingModeRelease(0); + } + + GFLAGS_NAMESPACE::FlagSaver _flags; +}; + +TEST_P(UrmaPollingTest, release_joins_and_can_be_repeated) { + std::atomic initialized{0}; + std::atomic released{0}; + ASSERT_EQ(0, urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, [&] { ++initialized; }, [&] { ++released; })); + urma::UrmaEndpoint::PollingModeRelease(0); + EXPECT_EQ(2, initialized.load()); + EXPECT_EQ(2, released.load()); + auto& group = urma::UrmaEndpoint::PollerGroups()[0]; + EXPECT_FALSE(group.running.load()); + for (const auto& poller : group.pollers) { + EXPECT_EQ(INVALID_BTHREAD, poller.tid); + } + urma::UrmaEndpoint::PollingModeRelease(0); + EXPECT_EQ(2, released.load()); + + ASSERT_EQ(0, urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, nullptr, [&] { ++released; })); + urma::UrmaEndpoint::PollingModeRelease(0); + EXPECT_EQ(4, released.load()); +} + +TEST_P(UrmaPollingTest, late_cq_operations_skip_stopped_pollers) { + ASSERT_EQ(0, urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, nullptr, nullptr)); + urma::UrmaEndpoint::PollingModeRelease(0); + urma::UrmaEndpoint endpoint(nullptr); + endpoint._cq_sid = 7; + endpoint.PollerAddCqSid(); + endpoint.PollerRemoveCqSid(); + for (auto& poller : urma::UrmaEndpoint::PollerGroups()[0].pollers) { + urma::UrmaEndpoint::CqSidOp op; + EXPECT_FALSE(poller.op_queue.Dequeue(op)); + } +} + +INSTANTIATE_TEST_SUITE_P(BthreadAndPthread, UrmaPollingTest, + ::testing::Bool()); + +namespace { +std::atomic g_exit_pollers_released{0}; +urma::UrmaEndpoint* g_exit_endpoint = nullptr; + +void CheckPollingExit() { + if (g_exit_pollers_released.load() != 2) { + std::_Exit(91); + } + // This handler was registered before URMA initialization, so it runs + // after the polling exit handler, like a late Socket/endpoint destructor. + delete g_exit_endpoint; + g_exit_endpoint = nullptr; +} +} // namespace + +TEST(UrmaPollingDeathTest, global_shutdown_joins_all_tags) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + EXPECT_EXIT({ + urma::FLAGS_urma_use_polling = true; + urma::FLAGS_urma_poller_num = 2; + FLAGS_task_group_ntags = 2; + if (urma::UrmaEndpoint::GlobalInitialize() != 0 || + urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, nullptr, + [] { ++g_exit_pollers_released; }) != 0) { + std::_Exit(92); + } + if (urma::UrmaEndpoint::PollingModeInitialize( + 1, nullptr, nullptr, + [] { ++g_exit_pollers_released; }) != 0) { + std::_Exit(92); + } + urma::UrmaEndpoint::GlobalPollingModeRelease(); + urma::UrmaEndpoint::GlobalPollingModeRelease(); + if (g_exit_pollers_released.load() != 4) { + std::_Exit(93); + } + std::_Exit(0); + }, ::testing::ExitedWithCode(0), ""); +} + +TEST(UrmaPollingDeathTest, exit_stops_pollers_before_late_endpoint_cleanup) { + // Only the link-time mock can provide device resources independently of + // the host. Release the probe initialization before starting the child. + urma_init_attr_t init_attr{}; + const urma_status_t status = urma_init(&init_attr); + if (status != URMA_SUCCESS && status != URMA_EEXIST) { + GTEST_SKIP() << "Polling exit test requires the URMA link-time mock"; + } + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + const bool using_mock = devices && num_devices > 0 && + strcmp(devices[0]->name, "mock_urma_device") == 0; + if (devices) { + urma_free_device_list(devices); + } + if (status == URMA_SUCCESS) { + ASSERT_EQ(URMA_SUCCESS, urma_uninit()); + } + if (!using_mock) { + GTEST_SKIP() << "Polling exit test requires the URMA link-time mock"; + } + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + EXPECT_EXIT({ + urma::FLAGS_urma_use_polling = true; + urma::FLAGS_urma_poller_num = 2; + urma::FLAGS_urma_buffer_size = 4096; + urma::FLAGS_urma_buffer_count = 512; + urma::FLAGS_urma_sq_size = 16; + urma::FLAGS_urma_rq_size = 16; + urma::FLAGS_urma_prepared_jetty_cnt = 0; + urma::g_skip_urma_init = false; + if (std::atexit(CheckPollingExit) != 0) { + std::_Exit(94); + } + urma::GlobalUrmaInitializeOrDie(); + if (!urma::IsUrmaAvailable() || + urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, nullptr, + [] { ++g_exit_pollers_released; }) != 0) { + std::_Exit(95); + } + g_exit_endpoint = new urma::UrmaEndpoint(nullptr); + g_exit_endpoint->_resource = new urma::UrmaResource; + g_exit_endpoint->_cq_sid = 7; + std::exit(0); + }, ::testing::ExitedWithCode(0), ""); +} + // --------------------------------------------------------------------------- // v2 binary HelloMessage: serialize + deserialize round-trips. // --------------------------------------------------------------------------- From 0ef609d68d51827a508c8e3201c6ccefd8754099 Mon Sep 17 00:00:00 2001 From: BGQ99 <1132767344@qq.com> Date: Thu, 22 Oct 2026 20:09:49 +0800 Subject: [PATCH 2/3] Check URMA polling shutdown handler registration --- src/brpc/urma/urma_helper.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/brpc/urma/urma_helper.cpp b/src/brpc/urma/urma_helper.cpp index 7af22aa6fe..7515c9b92e 100644 --- a/src/brpc/urma/urma_helper.cpp +++ b/src/brpc/urma/urma_helper.cpp @@ -627,8 +627,10 @@ static bool GlobalUrmaInitializeImpl() { return false; } - if (FLAGS_urma_use_polling) { - atexit(ReleasePollingModeAtExit); + if (FLAGS_urma_use_polling && + atexit(ReleasePollingModeAtExit) != 0) { + LOG(ERROR) << "Fail to register URMA polling shutdown handler"; + return false; } g_urma_available.store(true, butil::memory_order_release); From 3ae8b7456e7d7217c1d5f45ce7fd3828c16bdd17 Mon Sep 17 00:00:00 2001 From: BGQ99 <1132767344@qq.com> Date: Thu, 22 Oct 2026 22:06:27 +0800 Subject: [PATCH 3/3] Stop URMA pollers regardless of polling flag --- src/brpc/urma/urma_endpoint.cpp | 2 +- test/brpc_urma_unittest.cpp | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index b24714ba13..507125e94f 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -1713,7 +1713,7 @@ int UrmaEndpoint::PollingModeInitialize( void UrmaEndpoint::PollingModeRelease(bthread_tag_t tag) { auto& poller_groups = PollerGroups(); - if (!FLAGS_urma_use_polling || tag >= poller_groups.size()) { + if (tag >= poller_groups.size()) { return; } auto& group = poller_groups[tag]; diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index c8c9f614e5..57183d0cc8 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -138,6 +138,7 @@ TEST(UrmaPollingDeathTest, global_shutdown_joins_all_tags) { [] { ++g_exit_pollers_released; }) != 0) { std::_Exit(92); } + urma::FLAGS_urma_use_polling = false; urma::UrmaEndpoint::GlobalPollingModeRelease(); urma::UrmaEndpoint::GlobalPollingModeRelease(); if (g_exit_pollers_released.load() != 4) { @@ -191,6 +192,7 @@ TEST(UrmaPollingDeathTest, exit_stops_pollers_before_late_endpoint_cleanup) { g_exit_endpoint = new urma::UrmaEndpoint(nullptr); g_exit_endpoint->_resource = new urma::UrmaResource; g_exit_endpoint->_cq_sid = 7; + urma::FLAGS_urma_use_polling = false; std::exit(0); }, ::testing::ExitedWithCode(0), ""); }