cobalt/detail/join.hpp

0001 //
0002 // Copyright (c) 2022 Klemens Morgenstern (klemens.morgenstern@gmx.net)
0003 //
0004 // Distributed under the Boost Software License, Version 1.0. (See accompanying
0005 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
0006 //
0007
0008 #ifndef BOOST_COBALT_DETAIL_JOIN_HPP
0009 #define BOOST_COBALT_DETAIL_JOIN_HPP
0010
0011 #include <boost/cobalt/detail/await_result_helper.hpp>
0012 #include <boost/cobalt/detail/exception.hpp>
0013 #include <boost/cobalt/detail/fork.hpp>
0014 #include <boost/cobalt/detail/forward_cancellation.hpp>
0015 #include <boost/cobalt/detail/util.hpp>
0016 #include <boost/cobalt/detail/wrapper.hpp>
0017 #include <boost/cobalt/task.hpp>
0018 #include <boost/cobalt/this_thread.hpp>
0019
0020 #include <boost/asio/associated_cancellation_slot.hpp>
0021 #include <boost/asio/bind_cancellation_slot.hpp>
0022 #include <boost/asio/cancellation_signal.hpp>
0023
0024
0025 #include <boost/core/ignore_unused.hpp>
0026 #include <boost/intrusive_ptr.hpp>
0027 #include <boost/system/result.hpp>
0028 #include <boost/variant2/variant.hpp>
0029
0030 #include <array>
0031 #include <coroutine>
0032 #include <algorithm>
0033
0034 namespace boost::cobalt::detail
0035 {
0036
0037 template<typename ... Args>
0038 struct join_variadic_impl
0039 {
0040   using tuple_type = std::tuple<decltype(get_awaitable_type(std::declval<Args&&>()))...>;
0041
0042   join_variadic_impl(Args && ... args)
0043       : args{std::forward<Args>(args)...}
0044   {
0045   }
0046
0047   std::tuple<Args...> args;
0048
0049   constexpr static std::size_t tuple_size = sizeof...(Args);
0050
0051   struct awaitable : fork::static_shared_state<256 * tuple_size>
0052   {
0053     template<std::size_t ... Idx>
0054     awaitable(std::tuple<Args...> & args, std::index_sequence<Idx...>) :
0055         aws(awaitable_type_getter<Args>(std::get<Idx>(args))...)
0056     {
0057     }
0058
0059     tuple_type aws;
0060
0061     std::array<asio::cancellation_signal, tuple_size> cancel_;
0062     template<typename > constexpr static auto make_null() {return nullptr;};
0063     std::array<asio::cancellation_signal*, tuple_size> cancel = {make_null<Args>()...};
0064
0065     constexpr static bool all_void = (std::is_void_v<co_await_result_t<Args>> && ...);
0066     template<typename T>
0067     using result_store_part =
0068         std::optional<void_as_monostate<co_await_result_t<T>>>;
0069
0070     std::conditional_t<all_void,
0071                        variant2::monostate,
0072                        std::tuple<result_store_part<Args>...>> result;
0073     std::exception_ptr error;
0074
0075     template<std::size_t Idx>
0076     void cancel_step()
0077     {
0078       auto &r = cancel[Idx];
0079       if (r)
0080         std::exchange(r, nullptr)->emit(asio::cancellation_type::all);
0081     }
0082
0083     void cancel_all()
0084     {
0085       mp11::mp_for_each<mp11::mp_iota_c<sizeof...(Args)>>
0086           ([&](auto idx)
0087            {
0088              cancel_step<idx>();
0089            });
0090     }
0091
0092
0093
0094     template<std::size_t Idx>
0095     void interrupt_await_step()
0096     {
0097       using type = std::tuple_element_t<Idx, tuple_type>;
0098       using t = std::conditional_t<std::is_reference_v<std::tuple_element_t<Idx, std::tuple<Args...>>>,
0099           type &,
0100           type &&>;
0101
0102       if constexpr (interruptible<t>)
0103         if (this->cancel[Idx] != nullptr)
0104           static_cast<t>(std::get<Idx>(aws)).interrupt_await();
0105     }
0106
0107     void interrupt_await()
0108     {
0109       mp11::mp_for_each<mp11::mp_iota_c<sizeof...(Args)>>
0110           ([&](auto idx)
0111            {
0112              interrupt_await_step<idx>();
0113            });
0114     }
0115
0116
0117     // GCC doesn't like member funs
0118     template<std::size_t Idx>
0119     static detail::fork await_impl(awaitable & this_)
0120     try
0121     {
0122       auto & aw = std::get<Idx>(this_.aws);
0123       // check manually if we're ready
0124       auto rd = aw.await_ready();
0125       if (!rd)
0126       {
0127         this_.cancel[Idx] = &this_.cancel_[Idx];
0128         co_await this_.cancel[Idx]->slot();
0129         // make sure the executor is set
0130         co_await detail::fork::wired_up;
0131         // do the await - this doesn't call await-ready again
0132
0133         if constexpr (std::is_void_v<decltype(aw.await_resume())>)
0134         {
0135           co_await aw;
0136           if constexpr (!all_void)
0137             std::get<Idx>(this_.result).emplace();
0138         }
0139         else
0140           std::get<Idx>(this_.result).emplace(co_await aw);
0141       }
0142       else
0143       {
0144         if constexpr (std::is_void_v<decltype(aw.await_resume())>)
0145         {
0146           aw.await_resume();
0147           if constexpr (!all_void)
0148             std::get<Idx>(this_.result).emplace();
0149         }
0150         else
0151           std::get<Idx>(this_.result).emplace(aw.await_resume());
0152       }
0153
0154     }
0155     catch(...)
0156     {
0157       if (!this_.error)
0158            this_.error = std::current_exception();
0159       this_.cancel_all();
0160     }
0161
0162     std::array<detail::fork(*)(awaitable&), tuple_size> impls {
0163         []<std::size_t ... Idx>(std::index_sequence<Idx...>)
0164         {
0165           return std::array<detail::fork(*)(awaitable&), tuple_size>{&await_impl<Idx>...};
0166         }(std::make_index_sequence<tuple_size>{})
0167     };
0168
0169     detail::fork last_forked;
0170     std::size_t last_index = 0u;
0171
0172     bool await_ready()
0173     {
0174       while (last_index < tuple_size)
0175       {
0176         last_forked = impls[last_index++](*this);
0177         if (!last_forked.done())
0178           return false; // one coro didn't immediately complete!
0179       }
0180       last_forked.release();
0181       return true;
0182     }
0183
0184     template<typename H>
0185     auto await_suspend(
0186           std::coroutine_handle<H> h
0187 #if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
0188         , const boost::source_location & loc = BOOST_CURRENT_LOCATION
0189 #endif
0190     )
0191     {
0192 #if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
0193       this->loc = loc;
0194 #endif
0195       this->exec = &detail::get_executor(h);
0196       last_forked.release().resume();
0197       while (last_index < tuple_size)
0198         impls[last_index++](*this).release();
0199
0200       if (error)
0201         cancel_all();
0202
0203       if (!this->outstanding_work()) // already done, resume rightaway.
0204         return false;
0205
0206       // arm the cancel
0207       assign_cancellation(
0208           h,
0209           [&](asio::cancellation_type ct)
0210           {
0211             for (auto cs : cancel)
0212               if (cs)
0213                 cs->emit(ct);
0214           });
0215
0216       this->coro.reset(h.address());
0217       return true;
0218     }
0219
0220 #if _MSC_VER
0221     BOOST_NOINLINE
0222 #endif
0223     auto await_resume()
0224     {
0225       if (error)
0226         std::rethrow_exception(error);
0227       if constexpr(!all_void)
0228         return mp11::tuple_transform(
0229             []<typename T>(std::optional<T> & var)
0230                 -> T
0231             {
0232               BOOST_ASSERT(var.has_value());
0233               return std::move(*var);
0234             }, result);
0235     }
0236
0237     auto await_resume(const as_tuple_tag &)
0238     {
0239       using t = decltype(await_resume());
0240       if constexpr(!all_void)
0241       {
0242         if (error)
0243           return std::make_tuple(error, t{});
0244         else
0245           return std::make_tuple(std::current_exception(),
0246                                  mp11::tuple_transform(
0247               []<typename T>(std::optional<T> & var)
0248                   -> T
0249               {
0250                 BOOST_ASSERT(var.has_value());
0251                 return std::move(*var);
0252               }, result));
0253       }
0254       else
0255         return std::make_tuple(error);
0256     }
0257
0258     auto await_resume(const as_result_tag &)
0259     {
0260       using t = decltype(await_resume());
0261       using rt = system::result<t, std::exception_ptr>;
0262       if (error)
0263         return rt(system::in_place_error, error);
0264       if constexpr(!all_void)
0265         return mp11::tuple_transform(
0266             []<typename T>(std::optional<T> & var)
0267                 -> T
0268             {
0269               BOOST_ASSERT(var.has_value());
0270               return std::move(*var);
0271             }, result);
0272       else
0273         return system::in_place_value;
0274     }
0275   };
0276   awaitable operator co_await() &&
0277   {
0278     return awaitable(args, std::make_index_sequence<sizeof...(Args)>{});
0279   }
0280 };
0281
0282 template<typename Range>
0283 struct join_ranged_impl
0284 {
0285   Range aws;
0286
0287   using result_type = co_await_result_t<std::decay_t<decltype(*std::begin(std::declval<Range>()))>>;
0288
0289   constexpr static std::size_t result_size =
0290       sizeof(std::conditional_t<std::is_void_v<result_type>, variant2::monostate, result_type>);
0291
0292   struct awaitable : fork::shared_state
0293   {
0294     struct dummy
0295     {
0296       template<typename ... Args>
0297       dummy(Args && ...) {}
0298     };
0299
0300     using type = std::decay_t<decltype(*std::begin(std::declval<Range>()))>;
0301 #if !defined(BOOST_COBALT_NO_PMR)
0302     pmr::polymorphic_allocator<void> alloc{&resource};
0303
0304     std::conditional_t<awaitable_type<type>, Range &,
0305                        pmr::vector<co_awaitable_type<type>>> aws;
0306
0307     pmr::vector<bool> ready{std::size(aws), alloc};
0308     pmr::vector<asio::cancellation_signal> cancel_{std::size(aws), alloc};
0309     pmr::vector<asio::cancellation_signal*> cancel{std::size(aws), alloc};
0310
0311
0312
0313     std::conditional_t<
0314         std::is_void_v<result_type>,
0315         dummy,
0316         pmr::vector<std::optional<void_as_monostate<result_type>>>>
0317           result{
0318           cancel.size(),
0319           alloc};
0320 #else
0321     std::allocator<void> alloc;
0322     std::conditional_t<awaitable_type<type>, Range &,  std::vector<co_awaitable_type<type>>> aws;
0323
0324     std::vector<bool> ready{std::size(aws), alloc};
0325     std::vector<asio::cancellation_signal> cancel_{std::size(aws), alloc};
0326     std::vector<asio::cancellation_signal*> cancel{std::size(aws), alloc};
0327
0328     std::conditional_t<
0329         std::is_void_v<result_type>,
0330         dummy,
0331         std::vector<std::optional<void_as_monostate<result_type>>>>
0332           result{
0333           cancel.size(),
0334           alloc};
0335 #endif
0336     std::exception_ptr error;
0337
0338     awaitable(Range & aws_, std::false_type /* needs  operator co_await */)
0339       :  fork::shared_state((512 + sizeof(co_awaitable_type<type>) + result_size) * std::size(aws_))
0340       , aws{alloc}
0341       , ready{std::size(aws_), alloc}
0342       , cancel_{std::size(aws_), alloc}
0343       , cancel{std::size(aws_), alloc}
0344     {
0345       aws.reserve(std::size(aws_));
0346       for (auto && a : aws_)
0347       {
0348         using a_0 = std::decay_t<decltype(a)>;
0349         using a_t = std::conditional_t<
0350             std::is_lvalue_reference_v<Range>, a_0 &, a_0 &&>;
0351         aws.emplace_back(awaitable_type_getter<a_t>(static_cast<a_t>(a)));
0352       }
0353
0354       std::transform(std::begin(this->aws),
0355                      std::end(this->aws),
0356                      std::begin(ready),
0357                      [](auto & aw) {return aw.await_ready();});
0358     }
0359     awaitable(Range & aws, std::true_type /* needs operator co_await */)
0360         : fork::shared_state((512 + sizeof(co_awaitable_type<type>) + result_size) * std::size(aws))
0361         , aws(aws)
0362     {
0363       std::transform(std::begin(aws), std::end(aws), std::begin(ready), [](auto & aw) {return aw.await_ready();});
0364     }
0365
0366     awaitable(Range & aws)
0367       : awaitable(aws, std::bool_constant<awaitable_type<type>>{})
0368     {
0369     }
0370
0371     void cancel_all()
0372     {
0373       for (auto & r : cancel)
0374         if (r)
0375           std::exchange(r, nullptr)->emit(asio::cancellation_type::all);
0376     }
0377
0378     void interrupt_await()
0379     {
0380       using t = std::conditional_t<std::is_reference_v<Range>,
0381           co_awaitable_type<type> &,
0382           co_awaitable_type<type> &&>;
0383
0384       if constexpr (interruptible<t>)
0385       {
0386         std::size_t idx = 0u;
0387         for (auto & aw : aws)
0388           if (cancel[idx])
0389             static_cast<t>(aw).interrupt_await();
0390       }
0391     }
0392
0393
0394     static detail::fork await_impl(awaitable & this_, std::size_t idx)
0395     try
0396     {
0397       auto & aw = *std::next(std::begin(this_.aws), idx);
0398       auto rd = aw.await_ready();
0399       if (!rd)
0400       {
0401         this_.cancel[idx] = &this_.cancel_[idx];
0402         co_await this_.cancel[idx]->slot();
0403         co_await detail::fork::wired_up;
0404         if constexpr (std::is_void_v<decltype(aw.await_resume())>)
0405           co_await aw;
0406         else
0407           this_.result[idx].emplace(co_await aw);
0408       }
0409       else
0410       {
0411         if constexpr (std::is_void_v<decltype(aw.await_resume())>)
0412           aw.await_resume();
0413         else
0414           this_.result[idx].emplace(aw.await_resume());
0415       }
0416     }
0417     catch(...)
0418     {
0419       if (!this_.error)
0420         this_.error = std::current_exception();
0421       this_.cancel_all();
0422     }
0423
0424     detail::fork last_forked;
0425     std::size_t last_index = 0u;
0426
0427     bool await_ready()
0428     {
0429       while (last_index < cancel.size())
0430       {
0431         last_forked = await_impl(*this, last_index++);
0432         if (!last_forked.done())
0433           return false; // one coro didn't immediately complete!
0434       }
0435       last_forked.release();
0436       return true;
0437     }
0438
0439
0440     template<typename H>
0441     auto await_suspend(
0442         std::coroutine_handle<H> h
0443 #if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
0444         , const boost::source_location & loc = BOOST_CURRENT_LOCATION
0445 #endif
0446     )
0447     {
0448 #if defined(BOOST_ASIO_ENABLE_HANDLER_TRACKING)
0449       this->loc = loc;
0450 #endif
0451       exec = &detail::get_executor(h);
0452
0453       last_forked.release().resume();
0454       while (last_index < cancel.size())
0455         await_impl(*this, last_index++).release();
0456
0457       if (error)
0458         cancel_all();
0459
0460       if (!this->outstanding_work()) // already done, resume right away.
0461         return false;
0462
0463       // arm the cancel
0464       assign_cancellation(
0465           h,
0466           [&](asio::cancellation_type ct)
0467           {
0468             for (auto cs : cancel)
0469               if (cs)
0470                 cs->emit(ct);
0471           });
0472
0473
0474       this->coro.reset(h.address());
0475       return true;
0476     }
0477
0478     auto await_resume(const as_tuple_tag & )
0479     {
0480 #if defined(BOOST_COBALT_NO_PMR)
0481       std::vector<result_type> rr;
0482 #else
0483       pmr::vector<result_type> rr{this_thread::get_allocator()};
0484 #endif
0485
0486       if (error)
0487         return std::make_tuple(error, rr);
0488       if constexpr (!std::is_void_v<result_type>)
0489       {
0490         rr.reserve(result.size());
0491         for (auto & t : result)
0492           rr.push_back(*std::move(t));
0493         return std::make_tuple(std::exception_ptr(), std::move(rr));
0494       }
0495     }
0496
0497     auto await_resume(const as_result_tag & )
0498     {
0499 #if defined(BOOST_COBALT_NO_PMR)
0500       std::vector<result_type> rr;
0501 #else
0502       pmr::vector<result_type> rr{this_thread::get_allocator()};
0503 #endif
0504
0505       if (error)
0506         return system::result<decltype(rr), std::exception_ptr>(error);
0507       if constexpr (!std::is_void_v<result_type>)
0508       {
0509         rr.reserve(result.size());
0510         for (auto & t : result)
0511           rr.push_back(*std::move(t));
0512         return rr;
0513       }
0514     }
0515
0516 #if _MSC_VER
0517     BOOST_NOINLINE
0518 #endif
0519     auto await_resume()
0520     {
0521       if (error)
0522         std::rethrow_exception(error);
0523       if constexpr (!std::is_void_v<result_type>)
0524       {
0525 #if defined(BOOST_COBALT_NO_PMR)
0526         std::vector<result_type> rr;
0527 #else
0528         pmr::vector<result_type> rr{this_thread::get_allocator()};
0529 #endif
0530         rr.reserve(result.size());
0531         for (auto & t : result)
0532           rr.push_back(*std::move(t));
0533         return rr;
0534       }
0535     }
0536   };
0537   awaitable operator co_await() && {return awaitable{aws};}
0538 };
0539
0540 }
0541
0542
0543 #endif //BOOST_COBALT_DETAIL_JOIN_HPP