From b960c5ac7e65e8f4bcb3c5d293cfffc10d882382 Mon Sep 17 00:00:00 2001
From: BGQ99 <1132767344@qq.com>
Date: Wed, 28 Oct 2026 15:35:57 +0800
Subject: [PATCH 1/2] Restore late CQ removal coverage in URMA exit test

---
 test/brpc_urma_unittest.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp
index 57183d0cc8..d477d95026 100644
--- a/test/brpc_urma_unittest.cpp
+++ b/test/brpc_urma_unittest.cpp
@@ -116,6 +116,7 @@ void CheckPollingExit() {
     }
     // 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;
 }

From 9391cdc9c9e7be1397294bdd8eeae0424a3436d9 Mon Sep 17 00:00:00 2001
From: BGQ99 <1132767344@qq.com>
Date: Wed, 28 Oct 2026 16:58:13 +0800
Subject: [PATCH 2/2] Preserve late URMA CQ removals across poller restarts

---
 src/brpc/urma/urma_endpoint.cpp |  5 ++-
 test/brpc_urma_unittest.cpp     | 63 ++++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp
index 507125e94f..21032fe729 100644
--- a/src/brpc/urma/urma_endpoint.cpp
+++ b/src/brpc/urma/urma_endpoint.cpp
@@ -1758,9 +1758,8 @@ void UrmaEndpoint::PollerRemoveCqSid() {
         return;
     }
     auto& group = poller_groups[_poller_tag];
-    if (!group.running.load(butil::memory_order_acquire)) {
-        return;
-    }
+    // 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;
diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp
index d477d95026..5bdfbd23ad 100644
--- a/test/brpc_urma_unittest.cpp
+++ b/test/brpc_urma_unittest.cpp
@@ -23,6 +23,7 @@
 #if BRPC_WITH_URMA
 #include <atomic>
 #include <cstdlib>
+#include "bthread/bthread.h"
 #include "butil/atomicops.h"
 #include "butil/sys_byteorder.h"
 #include "urma_api.h"
@@ -89,7 +90,7 @@ TEST_P(UrmaPollingTest, release_joins_and_can_be_repeated) {
     EXPECT_EQ(4, released.load());
 }
 
-TEST_P(UrmaPollingTest, late_cq_operations_skip_stopped_pollers) {
+TEST_P(UrmaPollingTest, stopped_pollers_skip_add_and_preserve_remove) {
     ASSERT_EQ(0, urma::UrmaEndpoint::PollingModeInitialize(
         0, nullptr, nullptr, nullptr));
     urma::UrmaEndpoint::PollingModeRelease(0);
@@ -97,10 +98,17 @@ TEST_P(UrmaPollingTest, late_cq_operations_skip_stopped_pollers) {
     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,
@@ -122,6 +130,59 @@ void CheckPollingExit() {
 }
 }  // 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<bool> 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({