[coro_io] add support for cancellation

include/ylt/coro_io/io_context_pool.hpp

@@ -31,6 +31,7 @@
 #include <vector>
 
 #include "asio/dispatch.hpp"
+#include "async_simple/Signal.h"
 #ifdef __linux__
 #include <pthread.h>
 #include <sched.h>
@@ -104,11 +105,36 @@ 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_unique<asio::steady_timer>(executor_, dur);
-    auto tm = timer.get();
-    tm->async_wait([fn = std::move(func), timer = std::move(timer)](auto ec) {
-      fn();
-    });
+    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_acq_rel)) {
+                       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_acq_rel)) {
+          timer->first.cancel();
+        }
+      }
+    }
   }
 };
 

include/ylt/metric/summary_impl.hpp

@@ -135,6 +135,8 @@ class summary_impl {
       if (piece) {
         if constexpr (inc_order) {
           for (int j = 0; j < piece->size(); ++j) {
+            // tsan check data race here is expected. stat dont need to be very
+            // strict. we allow old value.
             auto value = (*piece)[j].load(std::memory_order_relaxed);
             if (value) {
               result.emplace_back(get_ordered_index(i * piece_size + j), value);

include/ylt/thirdparty/async_simple/Signal.h

@@ -17,7 +17,6 @@
 #define ASYNC_SIMPLE_SIGNAL_H
 
 #ifndef ASYNC_SIMPLE_USE_MODULES
-
 #include <assert.h>
 #include <any>
 #include <atomic>
@@ -205,14 +204,16 @@ class Slot {
                       "we dont allow emplace an empty signal handler");
         logicAssert(std::popcount(static_cast<uint64_t>(type)) == 1,
                     "It's not allow to emplace for multiple signals");
-        // trigger-once signal has already been triggered
+        auto handler = std::make_unique<detail::SignalSlotSharedState::Handler>(
+            std::forward<Args>(args)...);
+        auto oldHandlerPtr = loadHandler<true>(type);
+        // check trigger-once signal has already been triggered
+        // if signal has already been triggered, return false
         if (!detail::SignalSlotSharedState::isMultiTriggerSignal(type) &&
             (signal()->state() & type)) {
             return false;
         }
-        auto handler = std::make_unique<detail::SignalSlotSharedState::Handler>(
-            std::forward<Args>(args)...);
-        auto oldHandlerPtr = loadHandler<true>(type);
+        // if signal triggered later, we will found it by cas failed.
         auto oldHandler = oldHandlerPtr->load(std::memory_order_acquire);
         if (oldHandler ==
             &detail::SignalSlotSharedState::HandlerManager::emittedTag) {

include/ylt/thirdparty/async_simple/coro/Collect.h

@@ -166,8 +166,8 @@ struct CollectAnyAwaiter {
                 _slot, [c = continuation, e = event, size = input.size()](
                            SignalType type, Signal*) mutable {
                     auto count = e->downCount();
-                    if (count > size + 1) {
-                        c.resume();
+                    if (count == size + 1) {
+                      c.resume();
                     }
                 })) {  // has canceled
             return false;
@@ -186,14 +186,14 @@ struct CollectAnyAwaiter {
                     assert(e != nullptr);
                     auto count = e->downCount();
                     // n+1: n coro + 1 cancel handler
-                    if (count > size + 1) {
-                        _result = std::make_unique<ResultType>();
-                        _result->_idx = i;
-                        _result->_value = std::move(result);
-                        if (auto ptr = local->getSlot(); ptr) {
-                            ptr->signal()->emit(_SignalType);
-                        }
-                        c.resume();
+                    if (count == size + 1) {
+                      _result = std::make_unique<ResultType>();
+                      _result->_idx = i;
+                      _result->_value = std::move(result);
+                      if (auto ptr = local->getSlot(); ptr) {
+                        ptr->signal()->emit(_SignalType);
+                      }
+                      c.resume();
                     }
                 });
         }  // end for
@@ -268,8 +268,8 @@ struct CollectAnyVariadicAwaiter {
                 _slot, [c = continuation, e = event](SignalType type,
                                                      Signal*) mutable {
                     auto count = e->downCount();
-                    if (count > std::tuple_size<InputType>() + 1) {
-                        c.resume();
+                    if (count == std::tuple_size<InputType>() + 1) {
+                      c.resume();
                     }
                 })) {  // has canceled
             return false;
@@ -290,13 +290,13 @@ struct CollectAnyVariadicAwaiter {
                             res) mutable {
                         auto count = e->downCount();
                         // n+1: n coro + 1 cancel handler
-                        if (count > std::tuple_size<InputType>() + 1) {
-                            _result = std::make_unique<ResultType>(
-                                std::in_place_index_t<index>(), std::move(res));
-                            if (auto ptr = local->getSlot(); ptr) {
-                                ptr->signal()->emit(_SignalType);
-                            }
-                            c.resume();
+                        if (count == std::tuple_size<InputType>() + 1) {
+                          _result = std::make_unique<ResultType>(
+                              std::in_place_index_t<index>(), std::move(res));
+                          if (auto ptr = local->getSlot(); ptr) {
+                            ptr->signal()->emit(_SignalType);
+                          }
+                          c.resume();
                         }
                     });
             }(),
@@ -388,15 +388,19 @@ struct CollectAllAwaiter {
             _slot->chainedSignal(_signal.get());
 
         auto executor = promise_type._executor;
-        for (size_t i = 0; i < _input.size(); ++i) {
-            auto& exec = _input[i]._coro.promise()._executor;
-            if (exec == nullptr) {
-                exec = executor;
-            }
-            std::unique_ptr<LazyLocalBase> local;
-            local = std::make_unique<LazyLocalBase>(_signal.get());
-            _input[i]._coro.promise()._lazy_local = local.get();
-            auto&& func = [this, i, local = std::move(local)]() mutable {
+
+        _event.setAwaitingCoro(continuation);
+        auto size = _input.size();
+        for (size_t i = 0; i < size; ++i) {
+          auto& exec = _input[i]._coro.promise()._executor;
+          if (exec == nullptr) {
+            exec = executor;
+          }
+          std::unique_ptr<LazyLocalBase> local;
+          local = std::make_unique<LazyLocalBase>(_signal.get());
+          _input[i]._coro.promise()._lazy_local = local.get();
+          auto&& func =
+              [this, i, local = std::move(local)]() mutable {
                 _input[i].start([this, i, local = std::move(local)](
                                     Try<ValueType>&& result) {
                     _output[i] = std::move(result);
@@ -412,20 +416,15 @@ struct CollectAllAwaiter {
                         awaitingCoro.resume();
                     }
                 });
-            };
-            if (Para == true && _input.size() > 1) {
-                if (exec != nullptr)
-                    AS_LIKELY {
-                        exec->schedule_move_only(std::move(func));
-                        continue;
-                    }
-            }
-            func();
-        }
-        _event.setAwaitingCoro(continuation);
-        auto awaitingCoro = _event.down();
-        if (awaitingCoro) {
-            awaitingCoro.resume();
+              };
+          if (Para == true && _input.size() > 1) {
+            if (exec != nullptr)
+              AS_LIKELY {
+                exec->schedule_move_only(std::move(func));
+                continue;
+              }
+          }
+          func();
         }
     }
     inline auto await_resume() { return std::move(_output); }
@@ -602,10 +601,6 @@ struct CollectAllVariadicAwaiter {
                 }
             }(std::get<index>(_inputs), std::get<index>(_results)),
             ...);
-
-        if (auto awaitingCoro = _event.down(); awaitingCoro) {
-            awaitingCoro.resume();
-        }
     }
 
     void await_suspend(std::coroutine_handle<> continuation) {

include/ylt/thirdparty/async_simple/coro/CountEvent.h

@@ -34,15 +34,15 @@ namespace detail {
 // The last 'down' will resume the awaiting coroutine on this event.
 class CountEvent {
 public:
-    CountEvent(size_t count) : _count(count + 1) {}
-    CountEvent(const CountEvent&) = delete;
-    CountEvent(CountEvent&& other)
-        : _count(other._count.exchange(0, std::memory_order_relaxed)),
-          _awaitingCoro(std::exchange(other._awaitingCoro, nullptr)) {}
+ CountEvent(size_t count) : _count(count) {}
+ CountEvent(const CountEvent&) = delete;
+ CountEvent(CountEvent&& other)
+     : _count(other._count.exchange(0, std::memory_order_relaxed)),
+       _awaitingCoro(std::exchange(other._awaitingCoro, nullptr)) {}
 
-    [[nodiscard]] CoroHandle<> down(size_t n = 1) {
-        std::size_t oldCount;
-        return down(oldCount, n);
+ [[nodiscard]] CoroHandle<> down(size_t n = 1) {
+   std::size_t oldCount;
+   return down(oldCount, n);
     }
     [[nodiscard]] CoroHandle<> down(size_t& oldCount, std::size_t n) {
         // read acquire and write release, _awaitingCoro store can not be

src/coro_io/tests/CMakeLists.txt

@@ -5,6 +5,7 @@ add_executable(coro_io_test
         test_client_pool.cpp
         test_rate_limiter.cpp
         test_coro_channel.cpp
+        test_cancel.cpp
         main.cpp
         )
 if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_SYSTEM_NAME MATCHES "Windows") # mingw-w64