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 1/2] 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 2/2] 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), "");
 }