AnyOf: pass correct executor instance and cancelation flag

Author: DanglingPointer

src/include/crhandle/taskhandle.hpp

@@ -148,6 +148,10 @@ struct Promise
    const bool * parentCanceled = nullptr;
    stdcr::coroutine_handle<> parentHandle = nullptr;
 
+   const bool & CancelationFlag() const noexcept
+   {
+      return parentCanceled ? *parentCanceled : canceled;
+   }
    const E & Executor() const noexcept { return static_cast<const E &>(*this); }
    E & Executor() noexcept { return static_cast<E &>(*this); }
 
@@ -168,9 +172,7 @@ struct Promise
    auto await_transform(TaskHandle<R, E> && innerTask) const
    {
       using InnerAwaiter = std::remove_reference_t<decltype(innerTask.Run())>;
-      return CancelingAwaiter<InnerAwaiter>{
-         innerTask.Run(Executor(), parentCanceled ? parentCanceled : &canceled),
-         *this};
+      return CancelingAwaiter<InnerAwaiter>{innerTask.Run(Executor(), &CancelationFlag()), *this};
    }
 
    auto initial_suspend() noexcept

src/include/crhandle/taskutils.hpp

@@ -11,17 +11,19 @@ namespace cr {
 
 namespace internal {
 
+template <typename P = void>
 struct CurrentHandleRetriever
 {
-   stdcr::coroutine_handle<> handle;
+   stdcr::coroutine_handle<P> handle;
 
    bool await_ready() const noexcept { return false; }
-   bool await_suspend(stdcr::coroutine_handle<> h) noexcept
+   template <typename T = P>
+   bool await_suspend(stdcr::coroutine_handle<T> h) noexcept
    {
       handle = h;
       return false;
    }
-   stdcr::coroutine_handle<> await_resume() const noexcept { return handle; }
+   stdcr::coroutine_handle<P> await_resume() const noexcept { return handle; }
 };
 
 template <typename F, typename T, size_t... Is>
@@ -53,15 +55,16 @@ struct NonVoid<C<Ts...>>
 
 } // namespace internal
 
-
 template <typename T>
 using NonVoid = typename internal::NonVoid<T>::Type;
 
 template <Executor E, TaskResult... Rs>
 TaskHandle<NonVoid<std::variant<Rs...>>, E> AnyOf(TaskHandle<Rs, E>... ts)
 {
+   using ThisTaskHandle = TaskHandle<NonVoid<std::variant<Rs...>>, E>;
+
    std::optional<NonVoid<std::variant<Rs...>>> ret;
-   stdcr::coroutine_handle<> thisHandle = nullptr;
+   stdcr::coroutine_handle<> continuation = nullptr;
 
    auto TaskWrapper = [&]<size_t I, typename R>(std::in_place_index_t<I> i,
                                                 TaskHandle<R, E> task) -> TaskHandle<void, E> {
@@ -75,19 +78,23 @@ TaskHandle<NonVoid<std::variant<Rs...>>, E> AnyOf(TaskHandle<Rs, E>... ts)
          R tmp = co_await std::move(task);
          ret.emplace(i, std::move(tmp));
       }
-      if (thisHandle)
-         thisHandle.resume();
+      if (continuation)
+         continuation.resume();
    };
 
    auto handles = internal::CreateArray(std::make_index_sequence<sizeof...(Rs)>{},
                                         std::forward_as_tuple(std::move(ts)...),
                                         TaskWrapper);
 
+   auto thisHandle =
+      co_await internal::CurrentHandleRetriever<typename ThisTaskHandle::promise_type>{};
+   const auto & thisPromise = thisHandle.promise();
+
    for (auto & h : handles)
-      h.Run();
+      h.Run(thisPromise.Executor(), &thisPromise.CancelationFlag());
 
    if (!ret.has_value()) {
-      thisHandle = co_await internal::CurrentHandleRetriever{};
+      continuation = thisHandle;
       co_await stdcr::suspend_always{};
    }
 

test/dispatcher.hpp

@@ -55,6 +55,12 @@ struct ManualDispatcher
       return true;
    }
 
+   void ProcessAll()
+   {
+      while (ProcessOneTask())
+         ;
+   }
+
    std::deque<Task> queue;
 };
 

test/test_taskutils.cpp

@@ -4,6 +4,7 @@
 #include "crhandle/detachedhandle.hpp"
 #include "crhandle/taskhandle.hpp"
 #include "crhandle/taskutils.hpp"
+#include "dispatcher.hpp"
 
 #include <optional>
 
@@ -152,4 +153,83 @@ TEST_F(TaskUtilsFixture, anyof_handles_immediate_task_and_short_circuits)
    EXPECT_FALSE(stringResult2.has_value());
 }
 
+TEST_F(TaskUtilsFixture, anyof_uses_provided_executor_instance)
+{
+   ManualDispatcher dispatcher;
+
+   struct State
+   {
+      stdcr::coroutine_handle<> handle = nullptr;
+   } state1, state2;
+   size_t index = std::numeric_limits<size_t>::max();
+
+   static auto VoidTask = [](State & s) -> cr::TaskHandle<void, ManualDispatcher::Executor> {
+      co_await Awaitable<State>{s};
+   };
+   static auto OuterTask = [](State & state1,
+                              State & state2,
+                              size_t & index) -> cr::TaskHandle<void, ManualDispatcher::Executor> {
+      auto result = co_await cr::AnyOf(VoidTask(state1), VoidTask(state2));
+      index = result.index();
+   };
+
+   auto handle = OuterTask(state1, state2, index);
+   handle.Run(ManualDispatcher::Executor{&dispatcher});
+   EXPECT_FALSE(state1.handle);
+   EXPECT_FALSE(state2.handle);
+
+   dispatcher.ProcessAll();
+   EXPECT_TRUE(state1.handle);
+   EXPECT_TRUE(state2.handle);
+   EXPECT_EQ(std::numeric_limits<size_t>::max(), index);
+
+   state2.handle.resume();
+   EXPECT_EQ(std::numeric_limits<size_t>::max(), index);
+
+   dispatcher.ProcessAll();
+   EXPECT_EQ(1u, index);
+
+   state1.handle.resume();
+   dispatcher.ProcessAll();
+   EXPECT_EQ(1u, index);
+}
+
+TEST_F(TaskUtilsFixture, anyof_cancels_inner_tasks)
+{
+   ManualDispatcher dispatcher;
+
+   struct State
+   {
+      stdcr::coroutine_handle<> handle = nullptr;
+      bool done = false;
+   } state1, state2;
+   size_t index = std::numeric_limits<size_t>::max();
+
+   static auto VoidTask = [](State & s) -> cr::TaskHandle<void, ManualDispatcher::Executor> {
+      co_await Awaitable<State>{s};
+      s.done = true;
+   };
+   static auto OuterTask = [](State & state1,
+                              State & state2,
+                              size_t & index) -> cr::TaskHandle<void, ManualDispatcher::Executor> {
+      auto result = co_await cr::AnyOf(VoidTask(state1), VoidTask(state2));
+      index = result.index();
+   };
+
+   auto handle = OuterTask(state1, state2, index);
+   handle.Run(ManualDispatcher::Executor{&dispatcher});
+   dispatcher.ProcessAll();
+   EXPECT_TRUE(state1.handle);
+   EXPECT_TRUE(state2.handle);
+   EXPECT_EQ(std::numeric_limits<size_t>::max(), index);
+
+   handle = {};
+   state1.handle.resume();
+   state2.handle.resume();
+   dispatcher.ProcessAll();
+   EXPECT_EQ(std::numeric_limits<size_t>::max(), index);
+   EXPECT_FALSE(state1.done);
+   EXPECT_FALSE(state2.done);
+}
+
 } // namespace