awaiters.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/concepts.hpp"
#include "condy/condy_uring.hpp"
#include "condy/context.hpp"
#include "condy/finish_handles.hpp"
#include "condy/ring.hpp"
#include "condy/runtime.hpp"
#include "condy/work_type.hpp"
#include <coroutine>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <tuple>
 
namespace condy {
 
template <OpFinishHandleLike Handle> class HandleBox {
public:
    HandleBox(Handle h) : handle_(std::move(h)) {}
 
    Handle &get() { return handle_; }
 
    void maybe_release() { /* No-op */ }
 
private:
    Handle handle_;
};
 
template <typename Func, OpFinishHandleLike HandleBase>
class HandleBox<ZeroCopyMixin<Func, HandleBase>> {
public:
    using Handle = ZeroCopyMixin<Func, HandleBase>;
    HandleBox(Handle h) : handle_ptr_(std::make_unique<Handle>(std::move(h))) {}
    HandleBox(const HandleBox &other) // Deep copy
        : handle_ptr_(std::make_unique<Handle>(*other.handle_ptr_)) {}
 
    Handle &get() { return *handle_ptr_; }
 
    void maybe_release() { handle_ptr_.release(); }
 
private:
    std::unique_ptr<Handle> handle_ptr_;
};
 
template <OpFinishHandleLike Handle, PrepFuncLike Func> class OpAwaiterBase {
public:
    using HandleType = Handle;
 
    OpAwaiterBase(HandleBox<Handle> handle, Func func)
        : prep_func_(func), finish_handle_(std::move(handle)) {}
 
public:
    HandleType *get_handle() { return &finish_handle_.get(); }
 
    void init_finish_handle() { /* Leaf node, no-op */ }
 
    void register_operation(unsigned int flags) noexcept {
        auto &context = detail::Context::current();
        auto *ring = context.ring();
 
        context.runtime()->pend_work();
 
        io_uring_sqe *sqe = prep_func_(ring);
        assert(sqe && "prep_func must return a valid sqe");
        sqe->flags |= static_cast<uint8_t>(flags);
        io_uring_sqe_set_data(
            sqe, encode_work(&finish_handle_.get(), WorkType::Common));
    }
 
public:
    bool await_ready() const noexcept { return false; }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
        init_finish_handle();
        finish_handle_.get().set_invoker(&h.promise());
        register_operation(0);
    }
 
    auto await_resume() {
        auto result = finish_handle_.get().extract_result();
        finish_handle_.maybe_release();
        return result;
    }
 
protected:
    Func prep_func_;
    HandleBox<Handle> finish_handle_;
};
 
template <PrepFuncLike PrepFunc, CQEHandlerLike CQEHandler>
class [[nodiscard]] OpAwaiter
    : public OpAwaiterBase<OpFinishHandle<CQEHandler>, PrepFunc> {
public:
    using Base = OpAwaiterBase<OpFinishHandle<CQEHandler>, PrepFunc>;
    template <typename... Args>
    OpAwaiter(PrepFunc func, Args &&...args)
        : Base(HandleBox(
                   OpFinishHandle<CQEHandler>(std::forward<Args>(args)...)),
               std::move(func)) {}
};
 
template <PrepFuncLike PrepFunc, CQEHandlerLike CQEHandler,
          typename MultiShotFunc>
class [[nodiscard]] MultiShotOpAwaiter
    : public OpAwaiterBase<MultiShotOpFinishHandle<CQEHandler, MultiShotFunc>,
                           PrepFunc> {
public:
    using Base =
        OpAwaiterBase<MultiShotOpFinishHandle<CQEHandler, MultiShotFunc>,
                      PrepFunc>;
    template <typename... Args>
    MultiShotOpAwaiter(PrepFunc func, MultiShotFunc multishot_func,
                       Args &&...args)
        : Base(HandleBox(MultiShotOpFinishHandle<CQEHandler, MultiShotFunc>(
                   std::move(multishot_func), std::forward<Args>(args)...)),
               std::move(func)) {}
};
 
template <PrepFuncLike PrepFunc, CQEHandlerLike CQEHandler, typename FreeFunc>
class [[nodiscard]] ZeroCopyOpAwaiter
    : public OpAwaiterBase<ZeroCopyOpFinishHandle<CQEHandler, FreeFunc>,
                           PrepFunc> {
public:
    using Base =
        OpAwaiterBase<ZeroCopyOpFinishHandle<CQEHandler, FreeFunc>, PrepFunc>;
    template <typename... Args>
    ZeroCopyOpAwaiter(PrepFunc func, FreeFunc free_func, Args &&...args)
        : Base(HandleBox(ZeroCopyOpFinishHandle<CQEHandler, FreeFunc>(
                   std::move(free_func), std::forward<Args>(args)...)),
               std::move(func)) {}
};
 
template <unsigned int Flags, AwaiterLike Awaiter>
class [[nodiscard]] FlaggedOpAwaiter : public Awaiter {
public:
    using Base = Awaiter;
    FlaggedOpAwaiter(Awaiter awaiter) : Base(std::move(awaiter)) {}
 
    void register_operation(unsigned int flags) {
        Base::register_operation(flags | Flags);
    }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) {
        Base::init_finish_handle();
        Base::get_handle()->set_invoker(&h.promise());
        register_operation(0);
    }
};
 
template <HandleLike Handle, AwaiterLike Awaiter>
class [[nodiscard]] RangedParallelAwaiterBase {
public:
    using HandleType = Handle;
 
    RangedParallelAwaiterBase(std::vector<Awaiter> awaiters)
        : awaiters_(std::move(awaiters)) {}
 
public:
    HandleType *get_handle() { return &finish_handle_; }
 
    void init_finish_handle() {
        using ChildHandle = typename Awaiter::HandleType;
        std::vector<ChildHandle *> handles;
        handles.reserve(awaiters_.size());
        for (auto &awaiter : awaiters_) {
            awaiter.init_finish_handle();
            handles.push_back(awaiter.get_handle());
        }
        finish_handle_.init(std::move(handles));
    }
 
    void register_operation(unsigned int flags) {
        for (auto &awaiter : awaiters_) {
            awaiter.register_operation(flags);
        }
    }
 
public:
    bool await_ready() const noexcept { return awaiters_.empty(); }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) {
        init_finish_handle();
        finish_handle_.set_invoker(&h.promise());
        register_operation(0);
    }
 
    typename Handle::ReturnType await_resume() {
        return finish_handle_.extract_result();
    }
 
public:
    void push(Awaiter awaiter) { awaiters_.push_back(std::move(awaiter)); }
 
protected:
    HandleType finish_handle_;
    std::vector<Awaiter> awaiters_;
};

template <typename Awaiter>
using RangedParallelAllAwaiter = RangedParallelAwaiterBase<
    RangedParallelAllFinishHandle<typename Awaiter::HandleType>, Awaiter>;

template <typename Awaiter>
using RangedParallelAnyAwaiter = RangedParallelAwaiterBase<
    RangedParallelAnyFinishHandle<typename Awaiter::HandleType>, Awaiter>;

template <typename Awaiter>
using RangedWhenAllAwaiter = RangedParallelAwaiterBase<
    RangedWhenAllFinishHandle<typename Awaiter::HandleType>, Awaiter>;

template <typename Awaiter>
using RangedWhenAnyAwaiter = RangedParallelAwaiterBase<
    RangedWhenAnyFinishHandle<typename Awaiter::HandleType>, Awaiter>;
 
template <unsigned int Flags, AwaiterLike Awaiter>
class [[nodiscard]] RangedLinkAwaiterBase
    : public RangedWhenAllAwaiter<Awaiter> {
public:
    using Base = RangedWhenAllAwaiter<Awaiter>;
    using Base::Base;
 
    void register_operation(unsigned int flags) {
        auto *ring = detail::Context::current().ring();
        ring->reserve_space(Base::awaiters_.size());
        for (int i = 0; i < Base::awaiters_.size() - 1; ++i) {
            Base::awaiters_[i].register_operation(flags | Flags);
        }
        Base::awaiters_.back().register_operation(flags);
    }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) {
        Base::init_finish_handle();
        Base::finish_handle_.set_invoker(&h.promise());
        register_operation(0);
    }
};

template <typename Awaiter>
using RangedLinkAwaiter = RangedLinkAwaiterBase<IOSQE_IO_LINK, Awaiter>;

template <typename Awaiter>
using RangedHardLinkAwaiter = RangedLinkAwaiterBase<IOSQE_IO_HARDLINK, Awaiter>;
 
template <HandleLike Handle, AwaiterLike... Awaiters>
class [[nodiscard]] ParallelAwaiterBase {
public:
    using HandleType = Handle;
 
    ParallelAwaiterBase(Awaiters... awaiters)
        : awaiters_(std::move(awaiters)...) {}
    template <typename ParallelAwaiter, AwaiterLike New>
    ParallelAwaiterBase(ParallelAwaiter &&aws, New new_awaiter)
        : awaiters_(std::tuple_cat(std::move(aws.awaiters_),
                                   std::make_tuple(std::move(new_awaiter)))) {}
 
public:
    HandleType *get_handle() { return &finish_handle_; }
 
    void init_finish_handle() {
        auto handles = foreach_init_finish_handle_();
        static_assert(std::tuple_size<decltype(handles)>::value ==
                          sizeof...(Awaiters),
                      "Number of handles must match number of awaiters");
        std::apply(
            [this](auto &&...handle_ptrs) {
                finish_handle_.init(handle_ptrs...);
            },
            handles);
    }
 
    void register_operation(unsigned int flags) {
        foreach_register_operation_(flags);
    }
 
public:
    bool await_ready() const noexcept { return false; }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) {
        init_finish_handle();
        finish_handle_.set_invoker(&h.promise());
        register_operation(0);
    }
 
    typename Handle::ReturnType await_resume() {
        return finish_handle_.extract_result();
    }
 
private:
    template <size_t Idx = 0> auto foreach_init_finish_handle_() {
        if constexpr (Idx < sizeof...(Awaiters)) {
            std::get<Idx>(awaiters_).init_finish_handle();
            return std::tuple_cat(
                std::make_tuple(std::get<Idx>(awaiters_).get_handle()),
                foreach_init_finish_handle_<Idx + 1>());
        } else {
            return std::tuple<>();
        }
    }
 
    template <size_t Idx = 0>
    void foreach_register_operation_(unsigned int flags) {
        if constexpr (Idx < sizeof...(Awaiters)) {
            std::get<Idx>(awaiters_).register_operation(flags);
            foreach_register_operation_<Idx + 1>(flags);
        }
    }
 
protected:
    HandleType finish_handle_;
    std::tuple<Awaiters...> awaiters_;
 
    // Make awaiters_ accessible to all template instantiations
    template <HandleLike, AwaiterLike...> friend class ParallelAwaiterBase;
};

template <typename... Awaiter>
using ParallelAllAwaiter = ParallelAwaiterBase<
    ParallelAllFinishHandle<typename Awaiter::HandleType...>, Awaiter...>;

template <typename... Awaiter>
using ParallelAnyAwaiter = ParallelAwaiterBase<
    ParallelAnyFinishHandle<typename Awaiter::HandleType...>, Awaiter...>;

template <typename... Awaiter>
using WhenAllAwaiter =
    ParallelAwaiterBase<WhenAllFinishHandle<typename Awaiter::HandleType...>,
                        Awaiter...>;

template <typename... Awaiter>
using WhenAnyAwaiter =
    ParallelAwaiterBase<WhenAnyFinishHandle<typename Awaiter::HandleType...>,
                        Awaiter...>;
 
template <unsigned int Flags, AwaiterLike... Awaiter>
class [[nodiscard]] LinkAwaiterBase : public WhenAllAwaiter<Awaiter...> {
public:
    using Base = WhenAllAwaiter<Awaiter...>;
    using Base::Base;
 
    void register_operation(unsigned int flags) {
        auto *ring = detail::Context::current().ring();
        ring->reserve_space(sizeof...(Awaiter));
        foreach_register_operation_(flags);
    }
 
    template <typename PromiseType>
    void await_suspend(std::coroutine_handle<PromiseType> h) {
        Base::init_finish_handle();
        Base::finish_handle_.set_invoker(&h.promise());
        register_operation(0);
    }
 
private:
    template <size_t Idx = 0>
    void foreach_register_operation_(unsigned int flags) {
        if constexpr (Idx < sizeof...(Awaiter)) {
            std::get<Idx>(Base::awaiters_)
                .register_operation(Idx < sizeof...(Awaiter) - 1 ? flags | Flags
                                                                 : flags);
            foreach_register_operation_<Idx + 1>(flags);
        }
    }
};

template <typename... Awaiter>
using LinkAwaiter = LinkAwaiterBase<IOSQE_IO_LINK, Awaiter...>;

template <typename... Awaiter>
using HardLinkAwaiter = LinkAwaiterBase<IOSQE_IO_HARDLINK, Awaiter...>;
 
} // namespace condy