include/cppcoro/round_robin_scheduler.hpp

0001 ///////////////////////////////////////////////////////////////////////////////
0002 // Copyright (c) Lewis Baker
0003 // Licenced under MIT license. See LICENSE.txt for details.
0004 ///////////////////////////////////////////////////////////////////////////////
0005 #ifndef CPPCORO_ROUND_ROBIN_SCHEDULER_HPP_INCLUDED
0006 #define CPPCORO_ROUND_ROBIN_SCHEDULER_HPP_INCLUDED
0007
0008 #include <cppcoro/config.hpp>
0009
0010 #include <cppcoro/coroutine.hpp>
0011 #include <array>
0012 #include <cassert>
0013 #include <algorithm>
0014 #include <utility>
0015
0016 namespace cppcoro
0017 {
0018 #if CPPCORO_COMPILER_SUPPORTS_SYMMETRIC_TRANSFER
0019     /// This is a scheduler class that schedules coroutines in a round-robin
0020     /// fashion once N coroutines have been scheduled to it.
0021     ///
0022     /// Only supports access from a single thread at a time so
0023     ///
0024     /// This implementation was inspired by Gor Nishanov's CppCon 2018 talk
0025     /// about nano-coroutines.
0026     ///
0027     /// The implementation relies on symmetric transfer and noop_coroutine()
0028     /// and so only works with a relatively recent version of Clang and does
0029     /// not yet work with MSVC.
0030     template<size_t N>
0031     class round_robin_scheduler
0032     {
0033         static_assert(
0034             N >= 2,
0035             "Round robin scheduler must be configured to support at least two coroutines");
0036
0037         class schedule_operation
0038         {
0039         public:
0040             explicit schedule_operation(round_robin_scheduler& s) noexcept : m_scheduler(s) {}
0041
0042             bool await_ready() noexcept
0043             {
0044                 return false;
0045             }
0046
0047             cppcoro::coroutine_handle<> await_suspend(
0048                 cppcoro::coroutine_handle<> awaitingCoroutine) noexcept
0049             {
0050                 return m_scheduler.exchange_next(awaitingCoroutine);
0051             }
0052
0053             void await_resume() noexcept {}
0054
0055         private:
0056             round_robin_scheduler& m_scheduler;
0057         };
0058
0059         friend class schedule_operation;
0060
0061     public:
0062         round_robin_scheduler() noexcept
0063             : m_index(0)
0064             , m_noop(cppcoro::noop_coroutine())
0065         {
0066             for (size_t i = 0; i < N - 1; ++i)
0067             {
0068                 m_coroutines[i] = m_noop();
0069             }
0070         }
0071
0072         ~round_robin_scheduler()
0073         {
0074             // All tasks should have been joined before calling destructor.
0075             assert(std::all_of(
0076                 m_coroutines.begin(),
0077                 m_coroutines.end(),
0078                 [&](auto h) { return h == m_noop; }));
0079         }
0080
0081         schedule_operation schedule() noexcept
0082         {
0083             return schedule_operation{ *this };
0084         }
0085
0086         /// Resume any queued coroutines until there are no more coroutines.
0087         void drain() noexcept
0088         {
0089             size_t countRemaining = N - 1;
0090             do
0091             {
0092                 auto nextToResume = exchange_next(m_noop);
0093                 if (nextToResume != m_noop)
0094                 {
0095                     nextToResume.resume();
0096                     countRemaining = N - 1;
0097                 }
0098                 else
0099                 {
0100                     --countRemaining;
0101                 }
0102             } while (countRemaining > 0);
0103         }
0104
0105     private:
0106
0107         cppcoro::coroutine_handle exchange_next(
0108             cppcoro::coroutine_handle<> coroutine) noexcept
0109         {
0110             auto coroutineToResume = std::exchange(
0111                 m_scheduler.m_coroutines[m_scheduler.m_index],
0112                 awaitingCoroutine);
0113             m_scheduler.m_index = m_scheduler.m_index < (N - 2) ? m_scheduler.m_index + 1 : 0;
0114             return coroutineToResume;
0115         }
0116
0117         size_t m_index;
0118         const cppcoro::coroutine_handle<> m_noop;
0119         std::array<cppcoro::coroutine_handle<>, N - 1> m_coroutines;
0120     };
0121 #endif
0122 }
0123
0124 #endif