channel.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/context.hpp"
#include "condy/intrusive.hpp"
#include "condy/invoker.hpp"
#include "condy/runtime.hpp"
#include "condy/utils.hpp"
#include <bit>
#include <cerrno>
#include <coroutine>
#include <cstddef>
#include <new>
#include <optional>
#include <type_traits>
 
namespace condy {

template <typename T, size_t N = 2> class Channel {
public:
    Channel(size_t capacity)
        : buffer_(capacity ? std::bit_ceil(capacity) : 0) {}
    ~Channel() {
        std::lock_guard<std::mutex> lock(mutex_);
        push_close_inner_();
        destruct_all_();
    }
 
    Channel(const Channel &) = delete;
    Channel &operator=(const Channel &) = delete;
    Channel(Channel &&) = delete;
    Channel &operator=(Channel &&) = delete;
 
public:
    template <typename U>
        requires std::is_same_v<std::decay_t<U>, T>
    int try_push(U &&item) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        if (closed_) {
            return -EPIPE;
        }
        if (try_push_inner_(std::forward<U>(item))) {
            return 0;
        }
        return -EAGAIN;
    }

    std::pair<int, T> try_pop() noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        auto item = try_pop_inner_();
        if (item.has_value()) {
            return {0, std::move(item.value())};
        } else if (closed_) {
            return {-EPIPE, T()};
        } else {
            return {-EAGAIN, T()};
        }
    }
 
    void force_push(T item) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        if (closed_) [[unlikely]] {
            panic_on("Push to closed channel");
        }
        if (try_push_inner_(std::move(item))) [[likely]] {
            return;
        }
        // This is safe because if try_push_inner_ returns false, the item has
        // not been moved into the channel.
        // NOLINTBEGIN(bugprone-use-after-move)
        auto *fake_handle =
            new (std::nothrow) PushFinishHandle(std::move(item));
        // NOLINTEND(bugprone-use-after-move)
        if (!fake_handle) {
            panic_on("Allocation failed for PushFinishHandle");
        }
        assert(pop_awaiters_.empty());
        push_awaiters_.push_back(fake_handle);
    }
 
    struct [[nodiscard]] PushAwaiter;
    PushAwaiter push(T item) noexcept { return {*this, std::move(item)}; }
 
    struct [[nodiscard]] PopAwaiter;
    PopAwaiter pop() noexcept { return {*this}; }

    size_t capacity() const noexcept { return buffer_.capacity(); }

    size_t size() const noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return size_;
    }

    bool empty() const noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return size_ == 0;
    }

    bool is_closed() const noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return closed_;
    }

    void push_close() noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        push_close_inner_();
    }
 
private:
    class PushFinishHandle;
    class PopFinishHandle;
 
    int request_push_(PushFinishHandle *finish_handle) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        if (closed_) {
            return -EPIPE;
        }
        if (try_push_inner_(std::move(finish_handle->get_item()))) {
            return 0;
        }
        assert(pop_awaiters_.empty());
        push_awaiters_.push_back(finish_handle);
        detail::Context::current().runtime()->pend_work();
        return -EAGAIN;
    }
 
    bool cancel_push_(PushFinishHandle *finish_handle) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return push_awaiters_.remove(finish_handle);
    }
 
    std::pair<int, T> request_pop_(PopFinishHandle *finish_handle) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        auto result = try_pop_inner_();
        if (result.has_value()) {
            return {0, std::move(result.value())};
        }
        assert(push_awaiters_.empty());
        if (closed_) {
            return {-EPIPE, T()};
        }
        pop_awaiters_.push_back(finish_handle);
        detail::Context::current().runtime()->pend_work();
        return {-EAGAIN, T()};
    }
 
    bool cancel_pop_(PopFinishHandle *finish_handle) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return pop_awaiters_.remove(finish_handle);
    }
 
private:
    template <typename U>
        requires std::is_same_v<std::decay_t<U>, T>
    bool try_push_inner_(U &&item) noexcept {
        if (!pop_awaiters_.empty()) {
            assert(empty_inner_());
            auto *pop_handle = pop_awaiters_.pop_front();
            pop_handle->set_result({0, std::forward<U>(item)});
            pop_handle->schedule();
            return true;
        }
        if (!full_inner_()) {
            push_inner_(std::forward<U>(item));
            return true;
        }
        return false;
    }
 
    std::optional<T> try_pop_inner_() noexcept {
        if (!push_awaiters_.empty()) {
            assert(full_inner_());
            auto *push_handle = push_awaiters_.pop_front();
            T item = std::move(push_handle->get_item());
            push_handle->set_result(0);
            push_handle->schedule();
            return pop_and_push_(std::move(item));
        }
        if (!empty_inner_()) {
            T result = pop_inner_();
            return result;
        }
        return std::nullopt;
    }
 
    T pop_and_push_(T item) noexcept {
        if (no_buffer_()) {
            return item;
        } else {
            T result = pop_inner_();
            push_inner_(std::move(item));
            return result;
        }
    }
 
    template <typename U>
        requires std::is_same_v<std::decay_t<U>, T>
    void push_inner_(U &&item) noexcept {
        assert(!full_inner_());
        auto mask = buffer_.capacity() - 1;
        buffer_[tail_ & mask].construct(std::forward<U>(item));
        tail_++;
        size_++;
    }
 
    T pop_inner_() noexcept {
        assert(!empty_inner_());
        auto mask = buffer_.capacity() - 1;
        T item = std::move(buffer_[head_ & mask].get());
        buffer_[head_ & mask].destroy();
        head_++;
        size_--;
        return item;
    }
 
    bool no_buffer_() const noexcept { return buffer_.capacity() == 0; }
 
    bool empty_inner_() const noexcept {
        if (no_buffer_()) {
            return true;
        }
        return size_ == 0;
    }
 
    bool full_inner_() const noexcept {
        if (no_buffer_()) {
            return true;
        }
        return size_ == buffer_.capacity();
    }
 
    void push_close_inner_() noexcept {
        if (closed_) {
            return;
        }
        closed_ = true;
        // Cancel all pending pop awaiters
        PopFinishHandle *pop_handle = nullptr;
        while ((pop_handle = pop_awaiters_.pop_front()) != nullptr) {
            assert(empty_inner_());
            pop_handle->set_result({-EPIPE, T()});
            pop_handle->schedule();
        }
        // Cancel all pending push awaiters
        PushFinishHandle *push_handle = nullptr;
        while ((push_handle = push_awaiters_.pop_front()) != nullptr) {
            assert(full_inner_());
            push_handle->set_result(-EPIPE);
            push_handle->schedule();
        }
    }
 
    void destruct_all_() noexcept {
        while (!empty_inner_()) {
            pop_inner_();
        }
        assert(size_ == 0);
        assert(head_ == tail_);
    }
 
private:
    template <typename Handle>
    using HandleList = IntrusiveDoubleList<Handle, &Handle::link_entry_>;
 
    mutable std::mutex mutex_;
    HandleList<PushFinishHandle> push_awaiters_;
    HandleList<PopFinishHandle> pop_awaiters_;
    size_t head_ = 0;
    size_t tail_ = 0;
    size_t size_ = 0;
    SmallArray<RawStorage<T>, N> buffer_;
    bool closed_ = false;
};
 
template <typename T, size_t N>
class Channel<T, N>::PushFinishHandle
    : public InvokerAdapter<PushFinishHandle, WorkInvoker> {
public:
    using ReturnType = int;
 
    PushFinishHandle(T item) : item_(std::move(item)) {}
 
    void cancel([[maybe_unused]] Runtime *runtime) noexcept {
        assert(runtime == runtime_);
        if (channel_->cancel_push_(this)) {
            // Successfully canceled
            assert(result_ == -ENOTRECOVERABLE);
            result_ = -ECANCELED;
            need_resume_ = true;
            runtime_->schedule(this);
        }
    }
 
    ReturnType extract_result() noexcept { return result_; }
 
    void set_invoker(Invoker *invoker) noexcept { invoker_ = invoker; }
 
    void invoke() noexcept {
        if (need_resume_) {
            runtime_->resume_work();
        }
        (*invoker_)();
    }
 
public:
    void init(Channel *channel, Runtime *runtime) noexcept {
        channel_ = channel;
        runtime_ = runtime;
    }
 
    T &get_item() noexcept { return item_; }
 
    void schedule() noexcept {
        if (runtime_ == nullptr) [[unlikely]] {
            // Fake handle, no need to schedule
            delete this;
        } else {
            need_resume_ = true;
            runtime_->schedule(this);
        }
    }
 
    void set_result(int result) noexcept { result_ = result; }
 
public:
    DoubleLinkEntry link_entry_;
 
private:
    Invoker *invoker_ = nullptr;
    Channel *channel_ = nullptr;
    Runtime *runtime_ = nullptr;
    T item_;
    bool need_resume_ = false;
    int result_ = -ENOTRECOVERABLE; // Internal error if not set
};
 
template <typename T, size_t N>
class Channel<T, N>::PopFinishHandle
    : public InvokerAdapter<PopFinishHandle, WorkInvoker> {
public:
    using ReturnType = std::pair<int, T>;
 
    void cancel([[maybe_unused]] Runtime *runtime) noexcept {
        assert(runtime == runtime_);
        if (channel_->cancel_pop_(this)) {
            // Successfully canceled
            assert(result_.first == -ENOTRECOVERABLE);
            result_.first = -ECANCELED;
            need_resume_ = true;
            runtime_->schedule(this);
        }
    }
 
    ReturnType extract_result() noexcept { return std::move(result_); }
 
    void set_invoker(Invoker *invoker) noexcept { invoker_ = invoker; }
 
    void invoke() noexcept {
        if (need_resume_) {
            runtime_->resume_work();
        }
        (*invoker_)();
    }
 
public:
    void init(Channel *channel, Runtime *runtime) noexcept {
        channel_ = channel;
        runtime_ = runtime;
    }
 
    void set_result(ReturnType result) noexcept { result_ = std::move(result); }
 
    void schedule() noexcept {
        assert(runtime_ != nullptr);
        need_resume_ = true;
        runtime_->schedule(this);
    }
 
public:
    DoubleLinkEntry link_entry_;
 
private:
    Invoker *invoker_ = nullptr;
    Channel *channel_ = nullptr;
    Runtime *runtime_ = nullptr;
    ReturnType result_ = {-ENOTRECOVERABLE, T()}; // Internal error if not set
    bool need_resume_ = false;
};

template <typename T, size_t N> struct Channel<T, N>::PushAwaiter {
public:
    using HandleType = PushFinishHandle;
 
    PushAwaiter(Channel &channel, T item)
        : channel_(channel), finish_handle_(std::move(item)) {}
 
public:
    HandleType *get_handle() noexcept { return &finish_handle_; }
 
    void init_finish_handle() noexcept { /* Leaf node, no-op */ }
 
    void register_operation(unsigned int /*flags*/) noexcept {
        auto *runtime = detail::Context::current().runtime();
        finish_handle_.init(&channel_, runtime);
        int r = channel_.request_push_(&finish_handle_);
        if (r != -EAGAIN) {
            // Operation completed immediately
            finish_handle_.set_result(r);
            runtime->schedule(&finish_handle_);
        }
    }
 
public:
    bool await_ready() const noexcept { return false; }
 
    template <typename PromiseType>
    bool await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
        init_finish_handle();
        finish_handle_.set_invoker(&h.promise());
        finish_handle_.init(&channel_, detail::Context::current().runtime());
        int r = channel_.request_push_(&finish_handle_);
        if (r != -EAGAIN) {
            // Operation completed immediately
            finish_handle_.set_result(r);
            return false; // Do not suspend
        }
        return true; // Suspend
    }
 
    auto await_resume() noexcept { return finish_handle_.extract_result(); }
 
private:
    Channel &channel_;
    PushFinishHandle finish_handle_;
};

template <typename T, size_t N> struct Channel<T, N>::PopAwaiter {
public:
    using HandleType = PopFinishHandle;
 
    PopAwaiter(Channel &channel) : channel_(channel) {}
 
public:
    HandleType *get_handle() noexcept { return &finish_handle_; }
 
    void init_finish_handle() noexcept { /* Leaf node, no-op */ }
 
    void register_operation(unsigned int /*flags*/) noexcept {
        auto *runtime = detail::Context::current().runtime();
        finish_handle_.init(&channel_, runtime);
        auto item = channel_.request_pop_(&finish_handle_);
        auto r = item.first;
        if (r != -EAGAIN) {
            finish_handle_.set_result(std::move(item));
            runtime->schedule(&finish_handle_);
        }
    }
 
public:
    bool await_ready() const noexcept { return false; }
 
    template <typename PromiseType>
    bool await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
        init_finish_handle();
        finish_handle_.set_invoker(&h.promise());
        finish_handle_.init(&channel_, detail::Context::current().runtime());
        auto item = channel_.request_pop_(&finish_handle_);
        auto r = item.first;
        if (r != -EAGAIN) {
            finish_handle_.set_result(std::move(item));
            return false; // Do not suspend
        }
        return true; // Suspend
    }
 
    auto await_resume() noexcept { return finish_handle_.extract_result(); }
 
private:
    Channel &channel_;
    PopFinishHandle finish_handle_;
};
 
} // namespace condy