include/cppcoro/async_scope.hpp

0001 ///////////////////////////////////////////////////////////////////////////////
0002 // Copyright (c) Lewis Baker
0003 // Licenced under MIT license. See LICENSE.txt for details.
0004 ///////////////////////////////////////////////////////////////////////////////
0005 #ifndef CPPCORO_ASYNC_SCOPE_HPP_INCLUDED
0006 #define CPPCORO_ASYNC_SCOPE_HPP_INCLUDED
0007
0008 #include <cppcoro/on_scope_exit.hpp>
0009
0010 #include <atomic>
0011 #include <cppcoro/coroutine.hpp>
0012 #include <type_traits>
0013 #include <cassert>
0014
0015 namespace cppcoro
0016 {
0017     class async_scope
0018     {
0019     public:
0020
0021         async_scope() noexcept
0022             : m_count(1u)
0023         {}
0024
0025         ~async_scope()
0026         {
0027             // scope must be co_awaited before it destructs.
0028             assert(m_continuation);
0029         }
0030
0031         template<typename AWAITABLE>
0032         void spawn(AWAITABLE&& awaitable)
0033         {
0034             [](async_scope* scope, std::decay_t<AWAITABLE> awaitable) -> oneway_task
0035             {
0036                 scope->on_work_started();
0037                 auto decrementOnCompletion = on_scope_exit([scope] { scope->on_work_finished(); });
0038                 co_await std::move(awaitable);
0039             }(this, std::forward<AWAITABLE>(awaitable));
0040         }
0041
0042         [[nodiscard]] auto join() noexcept
0043         {
0044             class awaiter
0045             {
0046                 async_scope* m_scope;
0047             public:
0048                 awaiter(async_scope* scope) noexcept : m_scope(scope) {}
0049
0050                 bool await_ready() noexcept
0051                 {
0052                     return m_scope->m_count.load(std::memory_order_acquire) == 0;
0053                 }
0054
0055                 bool await_suspend(cppcoro::coroutine_handle<> continuation) noexcept
0056                 {
0057                     m_scope->m_continuation = continuation;
0058                     return m_scope->m_count.fetch_sub(1u, std::memory_order_acq_rel) > 1u;
0059                 }
0060
0061                 void await_resume() noexcept
0062                 {}
0063             };
0064
0065             return awaiter{ this };
0066         }
0067
0068     private:
0069
0070         void on_work_finished() noexcept
0071         {
0072             if (m_count.fetch_sub(1u, std::memory_order_acq_rel) == 1)
0073             {
0074                 m_continuation.resume();
0075             }
0076         }
0077
0078         void on_work_started() noexcept
0079         {
0080             assert(m_count.load(std::memory_order_relaxed) != 0);
0081             m_count.fetch_add(1, std::memory_order_relaxed);
0082         }
0083
0084         struct oneway_task
0085         {
0086             struct promise_type
0087             {
0088                 cppcoro::suspend_never initial_suspend() { return {}; }
0089                 cppcoro::suspend_never final_suspend() noexcept { return {}; }
0090                 void unhandled_exception() { std::terminate(); }
0091                 oneway_task get_return_object() { return {}; }
0092                 void return_void() {}
0093             };
0094         };
0095
0096         std::atomic<size_t> m_count;
0097         cppcoro::coroutine_handle<> m_continuation;
0098
0099     };
0100 }
0101
0102 #endif