Fix scheduler affinity (#405)

* Fix scheduler affinity

We have been storing a `task<>`'s scheduler as an `any_scheduler_ref`,
which has proven to be a source of use-after-free bugs.  This change
switches all the `any_scheduler_ref`s to `any_scheduler`s, fixing the
lifetime issues.

include/unifex/at_coroutine_exit.hpp

@@ -122,14 +122,14 @@ struct _cleanup_promise_base {
     return unstoppable_token{};
   }
 
-  friend any_scheduler_ref
+  friend any_scheduler
   tag_invoke(tag_t<get_scheduler>, const _cleanup_promise_base& p) noexcept {
     return p.sched_;
   }
 
   inline static constexpr inline_scheduler _default_scheduler{};
   continuation_handle<> continuation_{};
-  any_scheduler_ref sched_{_default_scheduler};
+  any_scheduler sched_{_default_scheduler};
   bool isUnhandledDone_{false};
 };
 

include/unifex/task.hpp

@@ -107,7 +107,7 @@ struct _promise_base {
     }
   };
 
-  void transform_schedule_sender_impl_(any_scheduler_ref newSched);
+  void transform_schedule_sender_impl_(any_scheduler newSched);
 
   coro::suspend_always initial_suspend() noexcept {
     return {};
@@ -127,7 +127,7 @@ struct _promise_base {
     return p.stoken_;
   }
 
-  friend any_scheduler_ref tag_invoke(tag_t<get_scheduler>, const _promise_base& p) noexcept {
+  friend any_scheduler tag_invoke(tag_t<get_scheduler>, const _promise_base& p) noexcept {
     return p.sched_;
   }
 
@@ -140,7 +140,7 @@ struct _promise_base {
 
   continuation_handle<> continuation_;
   inplace_stop_token stoken_;
-  any_scheduler_ref sched_{_default_scheduler};
+  any_scheduler sched_{_default_scheduler};
   bool rescheduled_ = false;
 };
 
@@ -368,7 +368,7 @@ struct _awaiter {
     using scheduler_t = remove_cvref_t<get_scheduler_result_t<OtherPromise&>>;
     using stop_token_t = remove_cvref_t<stop_token_type_t<OtherPromise>>;
     using needs_scheduler_t =
-        bool_constant<!same_as<scheduler_t, any_scheduler_ref>>;
+        bool_constant<!same_as<scheduler_t, any_scheduler>>;
     using needs_stop_token_t =
         bool_constant<!same_as<stop_token_t, inplace_stop_token>>;
 

source/task.cpp

@@ -21,7 +21,7 @@
 #include <unifex/at_coroutine_exit.hpp>
 
 namespace unifex::_task {
-void _promise_base::transform_schedule_sender_impl_(any_scheduler_ref newSched) {
+void _promise_base::transform_schedule_sender_impl_(any_scheduler newSched) {
   // If we haven't already inserted a cleanup action to take us back to the correct
   // scheduler, do so now:
   if (!std::exchange(this->rescheduled_, true)) {
@@ -37,7 +37,7 @@ void _promise_base::transform_schedule_sender_impl_(any_scheduler_ref newSched)
   // Update the current scheduler. (Don't do this before we have inserted the
   // cleanup action because the insertion of the cleanup action reads this task's
   // current scheduler.)
-  this->sched_ = newSched;
+  this->sched_ = std::move(newSched);
 }
 } // unifex::_task
 

test/task_scheduler_affinity_test.cpp

@@ -82,8 +82,8 @@ task<bool> test_current_scheduler(Scheduler s) {
 UNIFEX_TEMPLATE(typename Scheduler)
   (requires scheduler<Scheduler>)
 task<std::pair<bool, std::thread::id>> test_current_scheduler_is_inherited_impl(Scheduler s) {
-  any_scheduler_ref s2 = co_await current_scheduler();
-  bool sameScheduler = s2.equal_to(s);
+  any_scheduler s2 = co_await current_scheduler();
+  bool sameScheduler = (s2 == s);
   co_return std::make_pair(sameScheduler, std::this_thread::get_id());
 }
 UNIFEX_TEMPLATE(typename Scheduler)