BasicExecutionContext.hpp

#pragma once
 
#include "util/Assert.hpp"
#include "util/async/Concepts.hpp"
#include "util/async/Error.hpp"
#include "util/async/Operation.hpp"
#include "util/async/Outcome.hpp"
#include "util/async/context/impl/Cancellation.hpp"
#include "util/async/context/impl/Execution.hpp"
#include "util/async/context/impl/Strand.hpp"
#include "util/async/context/impl/Timer.hpp"
#include "util/async/context/impl/Utils.hpp"
#include "util/async/impl/ErrorHandling.hpp"
 
#include <boost/asio.hpp>
#include <boost/asio/error.hpp>
#include <boost/asio/strand.hpp>
#include <boost/asio/thread_pool.hpp>
 
#include <chrono>
#include <cstddef>
#include <expected>
#include <memory>
#include <optional>
#include <type_traits>
#include <utility>

namespace util::async {
namespace impl {
 
struct AsioPoolStrandContext {
    using Executor = boost::asio::strand<boost::asio::thread_pool::executor_type>;
    using Timer = SteadyTimer<Executor>;
 
    Executor&
    getExecutor()
    {
        return executor;
    }
 
    Executor executor;
};
 
struct AsioPoolContext {
    using Executor = boost::asio::thread_pool;
    using Timer = SteadyTimer<Executor>;
    using Strand = AsioPoolStrandContext;
 
    AsioPoolContext(std::size_t numThreads) : executor(std::make_unique<Executor>(numThreads))
    {
    }
 
    AsioPoolContext(AsioPoolContext const&) = delete;
    AsioPoolContext(AsioPoolContext&&) = default;
 
    Strand
    makeStrand() const
    {
        ASSERT(executor, "Called after executor was moved from.");
        return {boost::asio::make_strand(*executor)};
    }
 
    void
    stop() const
    {
        if (executor)  // don't call if executor was moved from
            executor->stop();
    }
 
    void
    join() const
    {
        if (executor)  // don't call if executor was moved from
            executor->join();
    }
 
    Executor&
    getExecutor() const
    {
        ASSERT(executor, "Called after executor was moved from.");
        return *executor;
    }
 
    std::unique_ptr<Executor> executor;
};
 
}  // namespace impl

template <
    typename ContextType,
    typename StopSourceType,
    typename DispatcherType,
    typename TimerContextProvider = impl::SelfContextProvider,
    typename ErrorHandlerType = impl::DefaultErrorHandler>
class BasicExecutionContext : public ExecutionContextTag {
    ContextType context_;

    friend impl::AssociatedExecutorExtractor;
 
public:
    static constexpr bool kIS_NOEXCEPT =
        noexcept(ErrorHandlerType::wrap([](auto&) { throw 0; })) and
        noexcept(ErrorHandlerType::catchAndAssert([] { throw 0; }));
 
    using ContextHolderType = ContextType;
 
    using ExecutorType = typename ContextHolderType::Executor;
 
    template <typename T>
    using ValueType = std::expected<T, ExecutionError>;
 
    using StopSource = StopSourceType;
 
    using StopToken = typename StopSourceType::Token;
 
    template <typename T>
    using StoppableOperation = StoppableOperation<ValueType<T>, StopSourceType>;
 
    template <typename T>
    using Operation = Operation<ValueType<T>>;
 
    using Strand = impl::BasicStrand<
        BasicExecutionContext,
        StopSourceType,
        DispatcherType,
        TimerContextProvider,
        ErrorHandlerType>;
 
    using Timer = typename ContextHolderType::Timer;
 
    // note: scheduled operations are always stoppable
    template <typename T>
    using ScheduledOperation = ScheduledOperation<BasicExecutionContext, StoppableOperation<T>>;
 
    // note: repeating operations are always stoppable and must return void
    using RepeatedOperation = RepeatingOperation<BasicExecutionContext>;

    explicit BasicExecutionContext(std::size_t numThreads = 1) noexcept : context_{numThreads}
    {
    }

    ~BasicExecutionContext() override
    {
        stop();
    }
 
    BasicExecutionContext(BasicExecutionContext&&) = default;
    BasicExecutionContext(BasicExecutionContext const&) = delete;

    [[nodiscard]] auto
    scheduleAfter(
        SomeStdDuration auto delay,
        SomeHandlerWith<StopToken> auto&& fn,
        std::optional<std::chrono::milliseconds> timeout = std::nullopt
    ) noexcept(kIS_NOEXCEPT)
    {
        if constexpr (not std::is_same_v<
                          decltype(TimerContextProvider::getContext(*this)),
                          decltype(*this)>) {
            return TimerContextProvider::getContext(*this).scheduleAfter(
                delay, std::forward<decltype(fn)>(fn), timeout
            );
        } else {
            using FnRetType = std::decay_t<std::invoke_result_t<decltype(fn), StopToken>>;
            return ScheduledOperation<FnRetType>(
                impl::extractAssociatedExecutor(*this),
                delay,
                [this, timeout, fn = std::forward<decltype(fn)>(fn)](auto) mutable {
                    return this->execute(
                        [fn = std::forward<decltype(fn)>(fn)](auto stopToken) mutable {
                            if constexpr (std::is_void_v<FnRetType>) {
                                std::invoke(std::forward<decltype(fn)>(fn), std::move(stopToken));
                            } else {
                                return std::invoke(
                                    std::forward<decltype(fn)>(fn), std::move(stopToken)
                                );
                            }
                        },
                        timeout
                    );
                }
            );
        }
    }

    [[nodiscard]] auto
    scheduleAfter(
        SomeStdDuration auto delay,
        SomeHandlerWith<StopToken, bool> auto&& fn,
        std::optional<std::chrono::milliseconds> timeout = std::nullopt
    ) noexcept(kIS_NOEXCEPT)
    {
        if constexpr (not std::is_same_v<
                          decltype(TimerContextProvider::getContext(*this)),
                          decltype(*this)>) {
            return TimerContextProvider::getContext(*this).scheduleAfter(
                delay, std::forward<decltype(fn)>(fn), timeout
            );
        } else {
            using FnRetType = std::decay_t<std::invoke_result_t<decltype(fn), StopToken, bool>>;
            return ScheduledOperation<FnRetType>(
                impl::extractAssociatedExecutor(*this),
                delay,
                [this, timeout, fn = std::forward<decltype(fn)>(fn)](auto ec) mutable {
                    return this->execute(
                        [fn = std::forward<decltype(fn)>(fn),
                         isAborted = (ec == boost::asio::error::operation_aborted)](
                            auto stopToken
                        ) mutable {
                            if constexpr (std::is_void_v<FnRetType>) {
                                std::invoke(
                                    std::forward<decltype(fn)>(fn), std::move(stopToken), isAborted
                                );
                            } else {
                                return std::invoke(
                                    std::forward<decltype(fn)>(fn), std::move(stopToken), isAborted
                                );
                            }
                        },
                        timeout
                    );
                }
            );
        }
    }

    [[nodiscard]] auto
    executeRepeatedly(
        SomeStdDuration auto interval,
        SomeHandlerWithoutStopToken auto&& fn
    ) noexcept(kIS_NOEXCEPT)
    {
        if constexpr (not std::is_same_v<
                          decltype(TimerContextProvider::getContext(*this)),
                          decltype(*this)>) {
            return TimerContextProvider::getContext(*this).executeRepeatedly(
                interval, std::forward<decltype(fn)>(fn)
            );
        } else {
            return RepeatedOperation(
                impl::extractAssociatedExecutor(*this), interval, std::forward<decltype(fn)>(fn)
            );
        }
    }

    [[nodiscard]] auto
    execute(
        SomeHandlerWith<StopToken> auto&& fn,
        std::optional<std::chrono::milliseconds> timeout = std::nullopt
    ) noexcept(kIS_NOEXCEPT)
    {
        return DispatcherType::dispatch(
            context_,
            impl::outcomeForHandler<StopSourceType>(fn),
            ErrorHandlerType::wrap([this, timeout, fn = std::forward<decltype(fn)>(fn)](
                                       auto& outcome, auto& stopSource, auto stopToken
                                   ) mutable {
                [[maybe_unused]] auto timeoutHandler = impl::getTimeoutHandleIfNeeded(
                    TimerContextProvider::getContext(*this), timeout, stopSource
                );
 
                using FnRetType = std::decay_t<std::invoke_result_t<decltype(fn), StopToken>>;
                if constexpr (std::is_void_v<FnRetType>) {
                    std::invoke(std::forward<decltype(fn)>(fn), std::move(stopToken));
                    outcome.setValue();
                } else {
                    outcome.setValue(
                        std::invoke(std::forward<decltype(fn)>(fn), std::move(stopToken))
                    );
                }
            })
        );
    }

    [[nodiscard]] auto
    execute(SomeHandlerWith<StopToken> auto&& fn, SomeStdDuration auto timeout) noexcept(
        kIS_NOEXCEPT
    )
    {
        return execute(
            std::forward<decltype(fn)>(fn),
            std::make_optional(std::chrono::duration_cast<std::chrono::milliseconds>(timeout))
        );
    }

    [[nodiscard]] auto
    execute(SomeHandlerWithoutStopToken auto&& fn) noexcept(kIS_NOEXCEPT)
    {
        return DispatcherType::dispatch(
            context_,
            impl::outcomeForHandler<StopSourceType>(fn),
            ErrorHandlerType::wrap([fn = std::forward<decltype(fn)>(fn)](auto& outcome) mutable {
                using FnRetType = std::decay_t<std::invoke_result_t<decltype(fn)>>;
                if constexpr (std::is_void_v<FnRetType>) {
                    std::invoke(std::forward<decltype(fn)>(fn));
                    outcome.setValue();
                } else {
                    outcome.setValue(std::invoke(std::forward<decltype(fn)>(fn)));
                }
            })
        );
    }

    void
    submit(SomeHandlerWithoutStopToken auto&& fn) noexcept(kIS_NOEXCEPT)
    {
        DispatcherType::post(context_, ErrorHandlerType::catchAndAssert(fn));
    }

    [[nodiscard]] Strand
    makeStrand()
    {
        return Strand(*this, context_.makeStrand());
    }

    void
    stop() const noexcept
    {
        context_.stop();
    }

    void
    join() const noexcept
    {
        context_.join();
    }

    typename ContextType::Executor&
    getExecutor()
    {
        return context_.getExecutor();
    }
};

using CoroExecutionContext = BasicExecutionContext<
    impl::AsioPoolContext,
    impl::YieldContextStopSource,
    impl::SpawnDispatchStrategy>;

using PoolExecutionContext =
    BasicExecutionContext<impl::AsioPoolContext, impl::BasicStopSource, impl::PostDispatchStrategy>;
 
}  // namespace util::async