fix timer user-after-free

alibaba · Jan 23, 2025 · 802a570 · 802a570
1 parent a1ebf51
commit 802a570
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 24 deletions.
diff --git a/include/ylt/coro_io/coro_io.hpp b/include/ylt/coro_io/coro_io.hpp
@@ -110,18 +110,18 @@ class callback_awaitor_base {
       set_value(std::forward<Args>(args)...);
       resume();
     }
-    template <typename... Args>
-    void set_value(Args &&...args) const {
-      if constexpr (!std::is_void_v<Arg>) {
-        obj->arg_ = {std::forward<Args>(args)...};
-      }
+    template <typename RetType1, typename RetType2>
+    void set_value(RetType1 &&t1, RetType2 &&t2) const {
+      std::get<0>(obj->arg_) = std::forward<RetType1>(t1);
+      std::get<1>(obj->arg_) = std::forward<RetType2>(t2);
     }
-    void set_value(std::error_code &&ec) const {
-      if constexpr (std::is_same_v<Arg, std::error_code>) {
-        obj->arg_ = std::move(ec);
+    template <typename RetType>
+    void set_value(RetType &&ec) const {
+      if constexpr (!requires { typename Arg::first_type; }) {
+        obj->arg_ = std::forward<RetType>(ec);
       }
       else {
-        std::get<0>(obj->arg_) = std::move(ec);
+        std::get<0>(obj->arg_) = std::forward<RetType>(ec);
       }
     }
     void resume() const { obj->coro_.resume(); }

diff --git a/include/ylt/coro_io/io_context_pool.hpp b/include/ylt/coro_io/io_context_pool.hpp
@@ -105,17 +105,35 @@ class ExecutorWrapper : public async_simple::Executor {
   }
   void schedule(Func func, Duration dur, uint64_t hint,
                 async_simple::Slot *slot = nullptr) override {
-    auto timer = std::make_shared<asio::steady_timer>(executor_, dur);
-    timer->async_wait([fn = std::move(func), timer](const auto &ec) {
-      fn();
-    });
-    if (!async_simple::signalHelper{async_simple::SignalType::Terminate}
-             .tryEmplace(slot, [timer](auto signalType, auto *signal) mutable {
-               asio::dispatch(timer->get_executor(), [timer]() {
-                 timer->cancel();
-               });
-             })) {
-      timer->cancel();
+    auto timer =
+        std::make_shared<std::pair<asio::steady_timer, std::atomic<bool>>>(
+            asio::steady_timer{executor_, dur}, false);
+    if (!slot) {
+      timer->first.async_wait([fn = std::move(func), timer](const auto &ec) {
+        fn();
+      });
+    }
+    else {
+      if (!async_simple::signalHelper{async_simple::SignalType::Terminate}
+               .tryEmplace(
+                   slot, [timer](auto signalType, auto *signal) mutable {
+                     if (bool expected = false;
+                         !timer->second.compare_exchange_strong(
+                             expected, true, std::memory_order_release)) {
+                       timer->first.cancel();
+                     }
+                   })) {
+        asio::dispatch(timer->first.get_executor(), func);
+      }
+      else {
+        timer->first.async_wait([fn = std::move(func), timer](const auto &ec) {
+          fn();
+        });
+        if (bool expected = false; !timer->second.compare_exchange_strong(
+                expected, true, std::memory_order_release)) {
+          timer->first.cancel();
+        }
+      }
     }
   }
 };

diff --git a/src/coro_io/tests/test_client_pool.cpp b/src/coro_io/tests/test_client_pool.cpp
@@ -91,8 +91,8 @@ TEST_CASE("test client pool") {
     REQUIRE(is_started.hasResult() == false);
     auto pool = coro_io::client_pool<coro_rpc::coro_rpc_client>::create(
         "127.0.0.1:8801", {.max_connection = 100,
-                           .idle_timeout = 300ms,
-                           .short_connect_idle_timeout = 50ms});
+                           .idle_timeout = 700ms,
+                           .short_connect_idle_timeout = 200ms});
     SpinLock lock;
     ConditionVariable<SpinLock> cv;
     auto res = co_await event(20, *pool, cv, lock);
@@ -102,10 +102,10 @@ TEST_CASE("test client pool") {
     CHECK(res);
     auto sz = pool->free_client_count();
     CHECK(sz == 200);
-    co_await coro_io::sleep_for(150ms);
+    co_await coro_io::sleep_for(500ms);
     sz = pool->free_client_count();
     CHECK((sz >= 100 && sz <= 105));
-    co_await coro_io::sleep_for(550ms);
+    co_await coro_io::sleep_for(1000ms);
     CHECK(pool->free_client_count() == 0);
     server.stop();
   }());