diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 26c7a1b9dc..21032fe729 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 (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& group = poller_groups[_poller_tag]; + if (!group.running.load(butil::memory_order_acquire)) { return; } - auto& pollers = _poller_groups[_poller_tag].pollers; + auto& pollers = group.pollers; if (pollers.empty()) { return; } @@ -1737,11 +1752,15 @@ 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]; + // Preserve removals for ADDs left queued when a worker stopped. The queues + // outlive endpoints, and a restarted worker must consume the matching pair. + auto& pollers = group.pollers; if (pollers.empty()) { return; } @@ -1754,7 +1773,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 +1784,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 +1858,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..7515c9b92e 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,18 @@ static bool GlobalUrmaInitializeImpl() { return false; } + 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); - // 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..5bdfbd23ad 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -21,9 +21,13 @@ #include #if BRPC_WITH_URMA +#include +#include +#include "bthread/bthread.h" #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 +35,230 @@ 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, stopped_pollers_skip_add_and_preserve_remove) { + 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(); + int removals = 0; + for (auto& poller : urma::UrmaEndpoint::PollerGroups()[0].pollers) { + urma::UrmaEndpoint::CqSidOp op; + if (poller.op_queue.Dequeue(op)) { + EXPECT_EQ(urma::UrmaEndpoint::CqSidOp::REMOVE, op.type); + EXPECT_EQ(7, op.sid); + ++removals; + } + EXPECT_FALSE(poller.op_queue.Dequeue(op)); + } + EXPECT_EQ(1, removals); +} + +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. + urma::FLAGS_urma_use_polling = true; + delete g_exit_endpoint; + g_exit_endpoint = nullptr; +} +} // namespace + +TEST(UrmaPollingDeathTest, restart_consumes_pending_add_and_late_remove) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (bool use_pthread : {false, true}) { + EXPECT_EXIT({ + urma::FLAGS_urma_use_polling = true; + urma::FLAGS_urma_poller_num = 1; + urma::FLAGS_urma_disable_bthread = use_pthread; + if (urma::UrmaEndpoint::GlobalInitialize() != 0 || + urma::UrmaEndpoint::PollingModeInitialize( + 0, nullptr, nullptr, nullptr) != 0) { + std::_Exit(92); + } + urma::UrmaEndpoint::PollingModeRelease(0); + auto& poller = urma::UrmaEndpoint::PollerGroups()[0].pollers[0]; + auto& queue = poller.op_queue; + // Seed the stopped queue with an ADD left unconsumed by a worker. + urma::UrmaEndpoint::CqSidOp add; + add.type = urma::UrmaEndpoint::CqSidOp::ADD; + add.sid = 7; + queue.Enqueue(add); + urma::UrmaEndpoint endpoint(nullptr); + endpoint._cq_sid = 7; + endpoint.PollerRemoveCqSid(); + urma::UrmaEndpoint::CqSidOp op; + if (!queue.Dequeue(op) || op.type != add.type || op.sid != add.sid) { + std::_Exit(96); + } + if (!queue.Dequeue(op) || + op.type != urma::UrmaEndpoint::CqSidOp::REMOVE || + op.sid != add.sid) { + std::_Exit(97); + } + // Replay the verified pair through a restarted worker. Its callback + // runs after draining the queue, so observing it bounds the wait. + queue.Enqueue(add); + queue.Enqueue(op); + std::atomic polled{false}; + if (urma::UrmaEndpoint::PollingModeInitialize( + 0, [&] { polled.store(true); }, nullptr, nullptr) != 0) { + std::_Exit(92); + } + for (int i = 0; i < 5000 && !polled.load(); ++i) { + bthread_usleep(1000); + } + urma::UrmaEndpoint::PollingModeRelease(0); + if (!polled.load() || queue.Dequeue(op)) { + std::_Exit(98); + } + std::_Exit(0); + }, ::testing::ExitedWithCode(0), ""); + } +} + +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::FLAGS_urma_use_polling = false; + 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; + urma::FLAGS_urma_use_polling = false; + std::exit(0); + }, ::testing::ExitedWithCode(0), ""); +} + // --------------------------------------------------------------------------- // v2 binary HelloMessage: serialize + deserialize round-trips. // ---------------------------------------------------------------------------