Fix URMA polling shutdown lifetime(CI)

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<int>(max_prepared);
 }
 
-std::vector<UrmaEndpoint::PollerGroup> UrmaEndpoint::_poller_groups;
+std::vector<UrmaEndpoint::PollerGroup>& 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<std::vector<PollerGroup>>();
+}
 
 // ============================================================================
 // 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& 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<PollerGroup>(ntags);
+        poller_groups = std::vector<PollerGroup>(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<bthread_tag_t>(tag));
     }
 }

src/brpc/urma/urma_endpoint.h

@@ -157,6 +157,8 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser {
                                      std::function<void()> init_fn,
                                      std::function<void()> 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<Poller> pollers;
         butil::atomic<bool> running;
     };
-    static std::vector<PollerGroup> _poller_groups;
+    static std::vector<PollerGroup>& PollerGroups();
     bthread_tag_t _poller_tag{0};
 
     DISALLOW_COPY_AND_ASSIGN(UrmaEndpoint);

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

test/brpc_urma_unittest.cpp

@@ -21,26 +21,182 @@
 #include <gflags/gflags.h>
 
 #if BRPC_WITH_URMA
+#include <atomic>
+#include <cstdlib>
 #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"
 #include "urma_types.h"
 
 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<bool> g_urma_available;
 
 }  // namespace urma
 }  // namespace brpc
 
+class UrmaPollingTest : public ::testing::TestWithParam<bool> {
+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<int> initialized{0};
+    std::atomic<int> 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<int> 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::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.
 // ---------------------------------------------------------------------------