runtime.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/condy_uring.hpp"
#include "condy/context.hpp"
#include "condy/finish_handles.hpp"
#include "condy/intrusive.hpp"
#include "condy/invoker.hpp"
#include "condy/ring.hpp"
#include "condy/runtime_options.hpp"
#include "condy/singleton.hpp"
#include "condy/utils.hpp"
#include "condy/work_type.hpp"
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <mutex>
 
namespace condy {
 
namespace detail {
 
class ThreadLocalRing : public ThreadLocalSingleton<ThreadLocalRing> {
public:
    Ring *ring() { return &ring_; }
 
    ThreadLocalRing() {
        io_uring_params params = {};
        params.flags |= IORING_SETUP_CLAMP;
        params.flags |= IORING_SETUP_SINGLE_ISSUER;
        params.flags |= IORING_SETUP_SUBMIT_ALL;
        [[maybe_unused]] int r = ring_.init(8, &params);
        assert(r == 0);
    }
 
private:
    Ring ring_;
};
 
inline int sync_msg_ring(io_uring_sqe *sqe_data) noexcept {
#if !IO_URING_CHECK_VERSION(2, 12) // >= 2.12
    return io_uring_register_sync_msg(sqe_data);
#else
    auto *ring = ThreadLocalRing::current().ring();
    auto *sqe = ring->get_sqe();
    *sqe = *sqe_data;
    int r;
    [[maybe_unused]] auto n =
        ring->reap_completions_wait([&](io_uring_cqe *cqe) { r = cqe->res; });
    assert(n == 1);
    return r;
#endif
}
 
} // namespace detail

class Runtime {
public:
    Runtime(const RuntimeOptions &options = {}) {
        io_uring_params params;
        std::memset(&params, 0, sizeof(params));
 
        params.flags |= IORING_SETUP_CLAMP;
        params.flags |= IORING_SETUP_SINGLE_ISSUER;
        params.flags |= IORING_SETUP_SUBMIT_ALL;
        params.flags |= IORING_SETUP_R_DISABLED;
 
        size_t ring_entries = options.sq_size_;
        if (options.cq_size_ != 0) { // 0 means default
            params.flags |= IORING_SETUP_CQSIZE;
            params.cq_entries = options.cq_size_;
        }
 
        if (options.enable_iopoll_) {
            params.flags |= IORING_SETUP_IOPOLL;
#if !IO_URING_CHECK_VERSION(2, 9) // >= 2.9
            if (options.enable_hybrid_iopoll_) {
                params.flags |= IORING_SETUP_HYBRID_IOPOLL;
            }
#endif
        }
 
        if (options.enable_sqpoll_) {
            params.flags |= IORING_SETUP_SQPOLL;
            params.sq_thread_idle = options.sqpoll_idle_time_ms_;
            if (options.sqpoll_thread_cpu_.has_value()) {
                params.flags |= IORING_SETUP_SQ_AFF;
                params.sq_thread_cpu = *options.sqpoll_thread_cpu_;
            }
        }
 
        if (options.attach_wq_target_ != nullptr) {
            params.flags |= IORING_SETUP_ATTACH_WQ;
            params.wq_fd = options.attach_wq_target_->ring_.ring()->ring_fd;
        }
 
        if (options.enable_defer_taskrun_) {
            params.flags |= IORING_SETUP_DEFER_TASKRUN;
            params.flags |= IORING_SETUP_TASKRUN_FLAG;
        }
 
        if (options.enable_coop_taskrun_) {
            params.flags |= IORING_SETUP_COOP_TASKRUN;
            params.flags |= IORING_SETUP_TASKRUN_FLAG;
        }
 
        if (options.enable_sqe128_) {
            params.flags |= IORING_SETUP_SQE128;
        }
 
        if (options.enable_cqe32_) {
            params.flags |= IORING_SETUP_CQE32;
        }
 
#if !IO_URING_CHECK_VERSION(2, 13) // >= 2.13
        if (options.enable_sqe_mixed_) {
            params.flags |= IORING_SETUP_SQE_MIXED;
        }
#endif
 
#if !IO_URING_CHECK_VERSION(2, 13) // >= 2.13
        if (options.enable_cqe_mixed_) {
            params.flags |= IORING_SETUP_CQE_MIXED;
        }
#endif
 
        void *buf = nullptr;
        size_t buf_size = 0;
#if !IO_URING_CHECK_VERSION(2, 5) // >= 2.5
        if (options.enable_no_mmap_) {
            params.flags |= IORING_SETUP_NO_MMAP;
            buf = options.no_mmap_buf_;
            buf_size = options.no_mmap_buf_size_;
        }
#endif
 
        int r = ring_.init(ring_entries, &params, buf, buf_size);
        if (r < 0) {
            throw make_system_error("io_uring_queue_init_params", -r);
        }
 
        event_interval_ = options.event_interval_;
        disable_register_ring_fd_ = options.disable_register_ring_fd_;
    }
 
    ~Runtime() { ring_.destroy(); }
 
    Runtime(const Runtime &) = delete;
    Runtime &operator=(const Runtime &) = delete;
    Runtime(Runtime &&) = delete;
    Runtime &operator=(Runtime &&) = delete;
 
public:
    void allow_exit() noexcept {
        pending_works_--;
        wakeup_();
    }
 
    void schedule(WorkInvoker *work) noexcept {
        auto *curr_runtime = detail::Context::current().runtime();
        if (curr_runtime == this) {
            local_queue_.push_back(work);
            return;
        }
 
        auto state = state_.load();
        if (state == State::Enabled) {
            // Fast path: if the ring is enabled, we can directly schedule the
            // work
            tsan_release(work);
            schedule_msg_ring_(curr_runtime, work, WorkType::Schedule);
        } else {
            // Slow path: if the ring is not enabled, we need to acquire the
            // mutex to ensure the work is scheduled before the ring is enabled
            std::unique_lock<std::mutex> lock(mutex_);
            state = state_.load();
            if (state == State::Enabled) {
                lock.unlock();
                tsan_release(work);
                schedule_msg_ring_(curr_runtime, work, WorkType::Schedule);
            } else {
                global_queue_.push_back(work);
            }
        }
    }
 
    void pend_work() noexcept { pending_works_++; }
 
    void resume_work() noexcept { pending_works_--; }

    void run() noexcept {
        State expected = State::Idle;
        [[maybe_unused]] bool success =
            state_.compare_exchange_strong(expected, State::Running);
        assert(success && "Runtime is already running or stopped");
        auto d1 = defer([this]() { state_.store(State::Stopped); });
 
        [[maybe_unused]] int r;
        r = io_uring_enable_rings(ring_.ring());
        assert(r == 0);
 
        {
            std::lock_guard<std::mutex> lock(mutex_);
            flush_global_queue_();
            // Now that the ring is enabled and all pending works are scheduled,
            // we can set the state to Enabled.
            state_.store(State::Enabled);
        }
 
        if (!disable_register_ring_fd_) {
            r = io_uring_register_ring_fd(ring_.ring());
            assert(r == 1); // 1 indicates success for this call
        }
 
        detail::Context::current().init(&ring_, this);
        auto d2 = defer([]() { detail::Context::current().reset(); });
 
        while (true) {
            tick_count_++;
 
            if (tick_count_ % event_interval_ == 0) {
                flush_ring_();
            }
 
            if (auto *work = local_queue_.pop_front()) {
                (*work)();
                continue;
            }
 
            if (pending_works_ == 0) {
                break;
            }
            flush_ring_wait_();
        }
    }

    auto &fd_table() noexcept { return ring_.fd_table(); }

    auto &buffer_table() noexcept { return ring_.buffer_table(); }

    auto &settings() noexcept { return ring_.settings(); }
 
private:
    void schedule_msg_ring_(Runtime *curr_runtime, WorkInvoker *work,
                            WorkType type) noexcept {
        if (curr_runtime != nullptr) {
            io_uring_sqe *sqe = curr_runtime->ring_.get_sqe();
            prep_msg_ring_(sqe, work, type);
            curr_runtime->pend_work();
        } else {
            io_uring_sqe sqe = {};
            prep_msg_ring_(&sqe, work, type);
            [[maybe_unused]] int r = detail::sync_msg_ring(&sqe);
            assert(r == 0);
        }
    }
 
    // Wakeup the runtime if it's blocked in Ring::reap_completions_wait()
    void wakeup_() noexcept {
        auto *curr_runtime = detail::Context::current().runtime();
        if (curr_runtime == this) {
            return;
        }
 
        auto state = state_.load();
        if (state != State::Enabled) {
            return;
        }
 
        schedule_msg_ring_(curr_runtime, nullptr, WorkType::Ignore);
    }
 
    void flush_global_queue_() noexcept {
        local_queue_.push_back(std::move(global_queue_));
    }
 
    void prep_msg_ring_(io_uring_sqe *sqe, WorkInvoker *work,
                        WorkType type) noexcept {
        auto data = encode_work(work, type);
        io_uring_prep_msg_ring(sqe, this->ring_.ring()->ring_fd, 0,
                               reinterpret_cast<uint64_t>(data), 0);
        io_uring_sqe_set_data(sqe, encode_work(nullptr, WorkType::Schedule));
    }
 
    size_t flush_ring_() noexcept {
        return ring_.reap_completions(
            [this](io_uring_cqe *cqe) { process_cqe_(cqe); });
    }
 
    size_t flush_ring_wait_() noexcept {
        return ring_.reap_completions_wait(
            [this](io_uring_cqe *cqe) { process_cqe_(cqe); });
    }
 
    void process_cqe_(io_uring_cqe *cqe) noexcept {
        auto *data_raw = io_uring_cqe_get_data(cqe);
        auto [data, type] = decode_work(data_raw);
 
        if (type == WorkType::Ignore) {
            // No-op
            assert(cqe->res != -EINVAL); // If EINVAL, something is wrong
        } else if (type == WorkType::SendFd) {
            auto &fd_table = ring_.fd_table();
            if (fd_table.fd_accepter_ == nullptr) [[unlikely]] {
                std::cerr << "[Deprecated Warning] Received a file "
                             "descriptor but no accepter is set.\n";
            } else {
                uint64_t payload = reinterpret_cast<uint64_t>(data) >> 3;
                if (payload == 0) { // Auto-allocate
                    fd_table.fd_accepter_(cqe->res);
                } else {
                    int target_fd = static_cast<int>(payload - 1);
                    fd_table.fd_accepter_(target_fd);
                }
            }
        } else if (type == WorkType::Schedule) {
            if (data == nullptr) {
                assert(cqe->res == 0);
                pending_works_--;
            } else {
                auto *work = static_cast<WorkInvoker *>(data);
                tsan_acquire(data);
                local_queue_.push_back(work);
            }
        } else if (type == WorkType::Common) {
            auto *handle = static_cast<OpFinishHandleBase *>(data);
            auto action = handle->handle_cqe(cqe);
            if (action.op_finish) {
                pending_works_--;
            }
            if (action.queue_work) {
                local_queue_.push_back(handle);
            }
        } else {
            assert(false && "Invalid work type");
        }
    }
 
private:
    enum class State : uint8_t {
        Idle,    // Not running
        Running, // Started running
        Enabled, // Running and ring enabled
        Stopped, // Stopped
    };
    static_assert(std::atomic<State>::is_always_lock_free);
 
    using WorkListQueue =
        IntrusiveSingleList<WorkInvoker, &WorkInvoker::work_queue_entry_>;
 
    // Global state
    std::mutex mutex_;
    WorkListQueue global_queue_;
    std::atomic_size_t pending_works_ = 1;
    std::atomic<State> state_ = State::Idle;
 
    // Local state
    WorkListQueue local_queue_;
    Ring ring_;
    size_t tick_count_ = 0;
 
    // Configurable parameters
    size_t event_interval_ = 61;
    bool disable_register_ring_fd_ = false;
};

inline auto &current_runtime() { return *detail::Context::current().runtime(); }
 
} // namespace condy