op_states.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/concepts.hpp"
#include "condy/condy_uring.hpp"
#include "condy/finish_handles.hpp"
#include "condy/type_traits.hpp"
#include "condy/utils.hpp"
#include <array>
#include <cstddef>
#include <optional>
#include <stop_token>
#include <tuple>
#include <vector>
 
namespace condy {
namespace detail {
 
template <typename Handle, PrepFuncLike Func> class OpSenderOperationState {
public:
    template <typename... HandleArgs>
    OpSenderOperationState(Func prep_func, HandleArgs &&...handle_args)
        : prep_func_(std::move(prep_func)),
          finish_handle_(std::forward<HandleArgs>(handle_args)...) {}
 
    CONDY_DELETE_COPY_MOVE(OpSenderOperationState);
 
public:
    void start(unsigned int flags) noexcept {
        auto &context = detail::Context::current();
        auto &ring = context.runtime()->ring();
        context.runtime()->pend_work();
        io_uring_sqe *sqe = prep_func_(&ring);
        assert(sqe && "prep_func must return a valid sqe");
        io_uring_sqe_set_flags(sqe, sqe->flags | flags);
        auto work = encode_work(&finish_handle_.get(), WorkType::Common);
        io_uring_sqe_set_data64(sqe, work);
        ring.maybe_submit();
 
        finish_handle_.get().maybe_set_cancel(context.runtime());
    }
 
private:
    Func prep_func_;
    HandleBox<Handle> finish_handle_;
};
 
template <unsigned int Flags, typename Sender, typename Receiver>
class FlaggedOpState {
public:
    FlaggedOpState(Sender sender, Receiver receiver)
        : op_state_(sender.connect_impl(std::move(receiver))) {}
 
    CONDY_DELETE_COPY_MOVE(FlaggedOpState);
 
    void start(unsigned int flags) noexcept { op_state_.start(flags | Flags); }
 
private:
    using OperationState = operation_state_t<Sender, Receiver>;
    OperationState op_state_;
};
 
template <typename TokenType> class WhenAnyCanceller {
public:
    auto chain_token(TokenType token) noexcept {
        if (token.stop_possible()) {
            stop_callback_.emplace(std::move(token), Cancellation{this});
        }
        return stop_source_.get_token();
    }
 
    template <typename R> void operator()(const R &) noexcept {
        stop_source_.request_stop();
    }
 
    void reset() noexcept { stop_callback_.reset(); }
 
private:
    struct Cancellation {
        WhenAnyCanceller *self;
        void operator()() noexcept { self->cancel_(); }
    };
    void cancel_() noexcept { stop_source_.request_stop(); }
 
    using StopCallbackType = stop_callback_t<TokenType, Cancellation>;
 
    std::stop_source stop_source_;
    std::optional<StopCallbackType> stop_callback_;
};
 
template <typename TokenType> class WhenAllCanceller {
public:
    auto chain_token(TokenType token) noexcept { return token; }
 
    template <typename R> void operator()(const R &) noexcept {}
 
    void reset() noexcept {}
};
 
template <typename Receiver, typename Canceller, typename... Senders>
class ParallelOperationState {
public:
    ParallelOperationState(std::tuple<Senders...> senders, Receiver receiver)
        : receiver_(std::move(receiver)) {
        auto next_token = canceller_.chain_token(receiver_.get_stop_token());
        connect_senders_(senders, next_token);
    }
 
    CONDY_DELETE_COPY_MOVE(ParallelOperationState);
 
    ~ParallelOperationState() {
        std::apply([](auto &&...states) { (states.destroy(), ...); },
                   op_states_);
    }
 
    void start(unsigned int flags) noexcept {
        if constexpr (sizeof...(Senders) == 0) {
            std::move(receiver_)(
                std::make_pair(std::move(order_), std::move(results_)));
        } else {
            std::apply(
                [&](auto &&...states) { (states.get().start(flags), ...); },
                op_states_);
        }
    }
 
private:
    using TokenType =
        std::remove_cvref_t<decltype(std::declval<Canceller &>().chain_token(
            std::declval<stop_token_t<Receiver>>()))>;
 
    template <size_t I = 0>
    void connect_senders_(std::tuple<Senders...> &senders,
                          const TokenType &token) noexcept {
        if constexpr (I < sizeof...(Senders)) {
            std::get<I>(op_states_).accept([&] {
                return std::move(std::get<I>(senders))
                    .connect_impl(ChildReceiver<I>{this, token});
            });
            connect_senders_<I + 1>(senders, token);
        }
    }
 
    template <size_t I, typename R> void receive_(R &&result) noexcept {
        canceller_(result);
        auto no = completed_count_++;
        order_[no] = I;
        std::get<I>(results_) = std::forward<R>(result);
        if (no + 1 == sizeof...(Senders)) {
            canceller_.reset();
            std::move(receiver_)(
                std::make_pair(std::move(order_), std::move(results_)));
        }
    }
 
    template <size_t I> struct ChildReceiver {
        ParallelOperationState *self;
        TokenType stop_token;
        template <typename R> void operator()(R &&result) noexcept {
            self->receive_<I>(std::forward<R>(result));
        }
        auto get_stop_token() const noexcept { return stop_token; }
    };
 
    template <typename T> struct operation_state_traits;
    template <size_t... Is>
    struct operation_state_traits<std::index_sequence<Is...>> {
        using type = std::tuple<
            RawStorage<operation_state_t<Senders, ChildReceiver<Is>>>...>;
    };
    using OperationStates = typename operation_state_traits<
        std::make_index_sequence<sizeof...(Senders)>>::type;
 
protected:
    OperationStates op_states_;
    std::array<size_t, sizeof...(Senders)> order_;
    std::tuple<typename Senders::ReturnType...> results_;
    size_t completed_count_ = 0;
    Receiver receiver_;
    Canceller canceller_;
};
 
template <typename Receiver, typename... Senders>
using ParallelAnyOperationState =
    ParallelOperationState<Receiver, WhenAnyCanceller<stop_token_t<Receiver>>,
                           Senders...>;
 
template <typename Receiver, typename... Senders>
using ParallelAllOperationState =
    ParallelOperationState<Receiver, WhenAllCanceller<stop_token_t<Receiver>>,
                           Senders...>;
 
template <typename Receiver> struct ReceiverAllWrapper {
    Receiver receiver;
    ReceiverAllWrapper(Receiver receiver) : receiver(std::move(receiver)) {}
    template <typename R> void operator()(R &&result) noexcept {
        auto &[order, results] = result;
        std::move(receiver)(std::move(results));
    }
    auto get_stop_token() const noexcept { return receiver.get_stop_token(); }
};
 
template <typename Receiver> struct ReceiverAnyWrapper {
    Receiver receiver;
    ReceiverAnyWrapper(Receiver receiver) : receiver(std::move(receiver)) {}
    template <typename R> void operator()(R &&result) noexcept {
        auto &[order, results] = result;
        size_t index = order[0];
        std::move(receiver)(tuple_at(results, index));
    }
    auto get_stop_token() const noexcept { return receiver.get_stop_token(); }
};
 
template <typename Receiver, typename... Senders>
using WhenAnyOperationState =
    ParallelAnyOperationState<ReceiverAnyWrapper<Receiver>, Senders...>;
 
template <typename Receiver, typename... Senders>
using WhenAllOperationState =
    ParallelAllOperationState<ReceiverAllWrapper<Receiver>, Senders...>;
 
template <typename Receiver, unsigned int Flags, typename... Senders>
class LinkOperationState : public WhenAllOperationState<Receiver, Senders...> {
public:
    using Base = WhenAllOperationState<Receiver, Senders...>;
    using Base::Base;
 
    void start(unsigned int flags) noexcept {
        auto &ring = detail::Context::current().runtime()->ring();
        ring.reserve_space(sizeof...(Senders));
        start_linked_operations_(flags);
    }
 
private:
    template <size_t I = 0>
    void start_linked_operations_(unsigned int flags) noexcept {
        if constexpr (I < sizeof...(Senders)) {
            auto &state = std::get<I>(Base::op_states_);
            if constexpr (I < sizeof...(Senders) - 1) {
                state.get().start(flags | Flags);
            } else {
                state.get().start(flags);
            }
            start_linked_operations_<I + 1>(flags);
        }
    }
};
 
template <typename Receiver, typename Canceller, typename Sender>
class RangedParallelOperationState {
public:
    RangedParallelOperationState(std::vector<Sender> senders, Receiver receiver)
        : op_states_(senders.size()), order_(senders.size()),
          results_(senders.size()), receiver_(std::move(receiver)) {
        auto next_token = canceller_.chain_token(receiver_.get_stop_token());
        for (size_t i = 0; i < senders.size(); ++i) {
            op_states_[i].accept([&] {
                return std::move(senders[i])
                    .connect_impl(ChildReceiver{this, i, next_token});
            });
        }
    }
 
    CONDY_DELETE_COPY_MOVE(RangedParallelOperationState);
 
    ~RangedParallelOperationState() {
        for (auto &op_state : op_states_) {
            op_state.destroy();
        }
    }
 
    void start(unsigned int flags) noexcept {
        if (op_states_.empty()) {
            std::move(receiver_)(
                std::make_pair(std::move(order_), std::move(results_)));
        } else {
            for (auto &op_state : op_states_) {
                op_state.get().start(flags);
            }
        }
    }
 
private:
    using TokenType =
        std::remove_cvref_t<decltype(std::declval<Canceller &>().chain_token(
            std::declval<stop_token_t<Receiver>>()))>;
 
    template <typename R> void receive_(size_t index, R &&result) noexcept {
        canceller_(result);
        size_t no = completed_count_++;
        order_[no] = index;
        results_[index] = std::forward<R>(result);
        if (no + 1 == op_states_.size()) {
            std::move(receiver_)(
                std::make_pair(std::move(order_), std::move(results_)));
        }
    }
 
    struct ChildReceiver {
        RangedParallelOperationState *self;
        size_t index;
        TokenType stop_token;
        template <typename R> void operator()(R &&result) noexcept {
            self->receive_(index, std::forward<R>(result));
        }
        auto get_stop_token() const noexcept { return stop_token; }
    };
 
    using OperationStates =
        std::vector<RawStorage<operation_state_t<Sender, ChildReceiver>>>;
 
protected:
    OperationStates op_states_;
    std::vector<size_t> order_;
    std::vector<typename Sender::ReturnType> results_;
    size_t completed_count_ = 0;
    Receiver receiver_;
    Canceller canceller_;
};
 
template <typename Receiver, typename Sender>
using RangedParallelAllOperationState = RangedParallelOperationState<
    Receiver, WhenAllCanceller<stop_token_t<Receiver>>, Sender>;
 
template <typename Receiver, typename Sender>
using RangedParallelAnyOperationState = RangedParallelOperationState<
    Receiver, WhenAnyCanceller<stop_token_t<Receiver>>, Sender>;
 
template <typename Receiver>
using ReceiverRangedAllWrapper = ReceiverAllWrapper<Receiver>;
 
template <typename Receiver> struct ReceiverRangedAnyWrapper {
    Receiver receiver;
    ReceiverRangedAnyWrapper(Receiver receiver)
        : receiver(std::move(receiver)) {}
    template <typename R> void operator()(R &&result) noexcept {
        auto &[order, results] = result;
        size_t index = order[0];
        std::move(receiver)(std::make_pair(index, std::move(results[index])));
    }
    auto get_stop_token() const noexcept { return receiver.get_stop_token(); }
};
 
template <typename Receiver, typename Sender>
using RangedWhenAllOperationState =
    RangedParallelAllOperationState<ReceiverRangedAllWrapper<Receiver>, Sender>;
 
template <typename Receiver, typename Sender>
using RangedWhenAnyOperationState =
    RangedParallelAnyOperationState<ReceiverRangedAnyWrapper<Receiver>, Sender>;
 
template <typename Receiver, unsigned int Flags, typename Sender>
class RangedLinkOperationState
    : public RangedWhenAllOperationState<Receiver, Sender> {
public:
    using Base = RangedWhenAllOperationState<Receiver, Sender>;
    using Base::Base;
 
    void start(unsigned int flags) noexcept {
        auto &ring = detail::Context::current().runtime()->ring();
        ring.reserve_space(Base::op_states_.size());
        for (size_t i = 0; i < Base::op_states_.size(); ++i) {
            auto &op_state = Base::op_states_[i];
            if (i < Base::op_states_.size() - 1) {
                op_state.get().start(flags | Flags);
            } else {
                op_state.get().start(flags);
            }
        }
    }
};
 
} // namespace detail
 
} // namespace condy