include/cppcoro/resume_on.hpp

0001 ///////////////////////////////////////////////////////////////////////////////
0002 // Copyright (c) Lewis Baker
0003 // Licenced under MIT license. See LICENSE.txt for details.
0004 ///////////////////////////////////////////////////////////////////////////////
0005 #ifndef CPPCORO_RESUME_ON_HPP_INCLUDED
0006 #define CPPCORO_RESUME_ON_HPP_INCLUDED
0007
0008 #include <cppcoro/task.hpp>
0009 #include <cppcoro/async_generator.hpp>
0010 #include <cppcoro/awaitable_traits.hpp>
0011 #include <cppcoro/detail/get_awaiter.hpp>
0012
0013 #include <exception>
0014 #include <type_traits>
0015
0016 namespace cppcoro
0017 {
0018     template<typename SCHEDULER>
0019     struct resume_on_transform
0020     {
0021         explicit resume_on_transform(SCHEDULER& s) noexcept
0022             : scheduler(s)
0023         {}
0024
0025         SCHEDULER& scheduler;
0026     };
0027
0028     template<typename SCHEDULER>
0029     resume_on_transform<SCHEDULER> resume_on(SCHEDULER& scheduler) noexcept
0030     {
0031         return resume_on_transform<SCHEDULER>(scheduler);
0032     }
0033
0034     template<typename T, typename SCHEDULER>
0035     decltype(auto) operator|(T&& value, resume_on_transform<SCHEDULER> transform)
0036     {
0037         return resume_on(transform.scheduler, std::forward<T>(value));
0038     }
0039
0040     template<
0041         typename SCHEDULER,
0042         typename AWAITABLE,
0043         typename AWAIT_RESULT = detail::remove_rvalue_reference_t<typename awaitable_traits<AWAITABLE>::await_result_t>,
0044         std::enable_if_t<!std::is_void_v<AWAIT_RESULT>, int> = 0>
0045     auto resume_on(SCHEDULER& scheduler, AWAITABLE awaitable)
0046         -> task<AWAIT_RESULT>
0047     {
0048         bool rescheduled = false;
0049         std::exception_ptr ex;
0050         try
0051         {
0052             // We manually get the awaiter here so that we can keep
0053             // it alive across the call to `scheduler.schedule()`
0054             // just in case the result is a reference to a value
0055             // in the awaiter that would otherwise be a temporary
0056             // and destructed before the value could be returned.
0057
0058             auto&& awaiter = detail::get_awaiter(static_cast<AWAITABLE&&>(awaitable));
0059
0060             auto&& result = co_await static_cast<decltype(awaiter)>(awaiter);
0061
0062             // Flag as rescheduled before scheduling in case it is the
0063             // schedule() operation that throws an exception as we don't
0064             // want to attempt to schedule twice if scheduling fails.
0065             rescheduled = true;
0066
0067             co_await scheduler.schedule();
0068
0069             co_return static_cast<decltype(result)>(result);
0070         }
0071         catch (...)
0072         {
0073             ex = std::current_exception();
0074         }
0075
0076         // We still want to resume on the scheduler even in the presence
0077         // of an exception.
0078         if (!rescheduled)
0079         {
0080             co_await scheduler.schedule();
0081         }
0082
0083         std::rethrow_exception(ex);
0084     }
0085
0086     template<
0087         typename SCHEDULER,
0088         typename AWAITABLE,
0089         typename AWAIT_RESULT = detail::remove_rvalue_reference_t<typename awaitable_traits<AWAITABLE>::await_result_t>,
0090         std::enable_if_t<std::is_void_v<AWAIT_RESULT>, int> = 0>
0091     auto resume_on(SCHEDULER& scheduler, AWAITABLE awaitable)
0092         -> task<>
0093     {
0094         std::exception_ptr ex;
0095         try
0096         {
0097             co_await static_cast<AWAITABLE&&>(awaitable);
0098         }
0099         catch (...)
0100         {
0101             ex = std::current_exception();
0102         }
0103
0104         // NOTE: We're assuming that `schedule()` operation is noexcept
0105         // here. If it were to throw what would we do if 'ex' was non-null?
0106         // Presumably we'd treat it the same as throwing an exception while
0107         // unwinding and call std::terminate()?
0108
0109         co_await scheduler.schedule();
0110
0111         if (ex)
0112         {
0113             std::rethrow_exception(ex);
0114         }
0115     }
0116
0117     template<typename SCHEDULER, typename T>
0118     async_generator<T> resume_on(SCHEDULER& scheduler, async_generator<T> source)
0119     {
0120         for (auto iter = co_await source.begin(); iter != source.end(); co_await ++iter)
0121         {
0122             auto& value = *iter;
0123             co_await scheduler.schedule();
0124             co_yield value;
0125         }
0126     }
0127 }
0128
0129 #endif