yield_to.h

// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
 
#pragma once
 
#include <boost/asio/executor_work_guard.hpp>
#include <boost/asio/io_context.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/optional.hpp>
#include <boost/thread/csbl/memory/allocator_arg.hpp>
 
#include <condition_variable>
#include <mutex>
#include <thread>
#include <vector>
 
namespace beast::test {

class EnableYieldTo
{
protected:
    boost::asio::io_context ios_;
 
private:
    boost::optional<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> work_;
    std::vector<std::thread> threads_;
    std::mutex m_;
    std::condition_variable cv_;
    std::size_t running_ = 0;
 
public:
    using yield_context = boost::asio::yield_context;
 
    explicit EnableYieldTo(std::size_t concurrency = 1) : work_(boost::asio::make_work_guard(ios_))
    {
        threads_.reserve(concurrency);
        for (std::size_t i = 0; i < concurrency; ++i)
        {
            threads_.emplace_back([&] { ios_.run(); });
        }
    }
 
    ~EnableYieldTo()
    {
        work_ = boost::none;
        for (auto& t : threads_)
            t.join();
    }

    boost::asio::io_context&
    getIoContext()
    {
        return ios_;
    }

#if BEAST_DOXYGEN
    template <class... FN>
    void
    yield_to(FN&&... fn);
#else
    template <class F0, class... FN>
    void
    yieldTo(F0&& f0, FN&&... fn);
#endif
 
private:
    void
    spawn()
    {
    }
 
    template <class F0, class... FN>
    void
    spawn(F0&& f, FN&&... fn);
};
 
template <class F0, class... FN>
void
EnableYieldTo::yieldTo(F0&& f0, FN&&... fn)
{
    running_ = 1 + sizeof...(FN);
    spawn(f0, fn...);
    std::unique_lock<std::mutex> lock{m_};
    cv_.wait(lock, [&] { return running_ == 0; });
}
 
template <class F0, class... FN>
inline void
EnableYieldTo::spawn(F0&& f, FN&&... fn)
{
    boost::asio::spawn(
        ios_,
        boost::allocator_arg,
        boost::context::fixedsize_stack(2 * 1024 * 1024),
        [&](yield_context yield) {
            f(yield);
            std::scoped_lock const lock{m_};
            if (--running_ == 0)
                cv_.notify_all();
        },
        [](std::exception_ptr e) {
            if (e)
                std::rethrow_exception(e);
        });
    spawn(fn...);
}
 
}  // namespace beast::test