runtime.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/condy_uring.hpp"
#include "condy/context.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 <atomic>
#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);
        // If we can construct Runtime, we should be able to construct this
        // thread-local ring.
        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 = 0;
    auto n =
        ring->reap_completions_wait([&](io_uring_cqe *cqe) { r = cqe->res; });
    if (n < 0) {
        return static_cast<int>(n);
    }
    assert(n == 1);
    return r;
#endif
}
 
class CancelRequest {
public:
    CancelRequest(uintptr_t data) : data_(data) {}
 
    void wait() noexcept {
        while (!finished_.load(std::memory_order_acquire)) {
            finished_.wait(false, std::memory_order_relaxed);
        }
    }
 
    void notify() noexcept {
        finished_.store(true, std::memory_order_release);
        finished_.notify_one();
    }
 
    uintptr_t data() const noexcept { return data_; }
 
private:
    uintptr_t data_;
    std::atomic_bool finished_ = false;
};
 
} // namespace detail
 
class OpFinishHandleBase {
public:
    using HandleFunc = bool (*)(void *, io_uring_cqe *) noexcept;
 
    bool handle(io_uring_cqe *cqe) noexcept {
        assert(handle_func_ != nullptr);
        return handle_func_(this, cqe);
    }
 
protected:
    OpFinishHandleBase() = default;
 
protected:
    HandleFunc handle_func_ = nullptr;
};

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
 
#if !IO_URING_CHECK_VERSION(2, 14) // >= 2.14
        if (options.enable_sq_rewind_) {
            params.flags |= IORING_SETUP_SQ_REWIND;
        }
#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,
                               encode_work(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,
                                   encode_work(work, WorkType::Schedule));
            } else {
                global_queue_.push_back(work);
            }
        }
    }
 
    // Internal use only. Schedule a cancel request for the given data.
    void cancel(uintptr_t data) noexcept {
        auto *curr_runtime = detail::Context::current().runtime();
        if (curr_runtime == this) {
            io_uring_sqe *sqe = ring_.get_sqe();
            prep_cancel_(sqe, data);
            return;
        }
 
        auto state = state_.load();
        if (state != State::Enabled) {
            return;
        }
 
        detail::CancelRequest request(data);
        tsan_release(&request);
        schedule_msg_ring_(curr_runtime,
                           encode_work(&request, WorkType::Cancel));
        if (curr_runtime != nullptr) {
            // Ensure the cancel msg is submitted.
            curr_runtime->ring_.submit();
        }
        // Block until the runtime thread has submitted the cancel SQE. This is
        // important to prevent address reuse of the same data pointer, which
        // can lead to incorrect cancellation or other bugs.
        request.wait();
    }
 
    void pend_work() noexcept { pending_works_++; }
 
    void resume_work() noexcept { pending_works_--; }

    void run() {
        State expected = State::Idle;
        bool success = state_.compare_exchange_strong(expected, State::Running);
        if (!success) {
            throw std::runtime_error(
                "Runtime is already running or has been 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, uintptr_t data) noexcept {
        int ring_fd = this->ring_.ring()->ring_fd;
        if (curr_runtime != nullptr) {
            io_uring_sqe *sqe = curr_runtime->ring_.get_sqe();
            prep_msg_ring_(ring_fd, sqe, data);
            curr_runtime->pend_work();
        } else {
            io_uring_sqe sqe = {};
            prep_msg_ring_(ring_fd, &sqe, data);
            int r = detail::sync_msg_ring(&sqe);
            if (r < 0) {
                panic_on(std::format("sync_msg_ring: {}", std::strerror(-r)));
            }
        }
    }
 
    // 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,
                           encode_work(nullptr, WorkType::Ignore));
    }
 
    void flush_global_queue_() noexcept {
        local_queue_.push_back(std::move(global_queue_));
    }
 
    static void prep_msg_ring_(int ring_fd, io_uring_sqe *sqe,
                               uintptr_t data) noexcept {
        io_uring_prep_msg_ring(sqe, ring_fd, 0, data, 0);
        io_uring_sqe_set_data64(sqe, encode_work(nullptr, WorkType::Schedule));
    }
 
    static void prep_cancel_(io_uring_sqe *sqe, uintptr_t data) noexcept {
        io_uring_prep_cancel64(sqe, data, 0);
        io_uring_sqe_set_data64(sqe, encode_work(nullptr, WorkType::Ignore));
        io_uring_sqe_set_flags(sqe, IOSQE_CQE_SKIP_SUCCESS);
    }
 
    void flush_ring_() noexcept {
        auto r = ring_.reap_completions(
            [this](io_uring_cqe *cqe) { process_cqe_(cqe); });
        if (r < 0) {
            panic_on(std::format("io_uring_peek_cqe: {}",
                                 std::strerror(static_cast<int>(-r))));
        }
    }
 
    void flush_ring_wait_() noexcept {
        auto r = ring_.reap_completions_wait(
            [this](io_uring_cqe *cqe) { process_cqe_(cqe); });
        if (r < 0) {
            panic_on(std::format("io_uring_submit_and_wait: {}",
                                 std::strerror(static_cast<int>(-r))));
        }
    }
 
    void process_cqe_(io_uring_cqe *cqe) noexcept {
        auto [data, type] = decode_work(io_uring_cqe_get_data64(cqe));
 
        if (type == WorkType::Ignore) {
            // No-op
            assert(cqe->res != -EINVAL); // If EINVAL, something is wrong
        } else if (type == WorkType::Schedule) {
            if (data == nullptr) {
                if (cqe->res < 0) {
                    panic_on(std::format("io_uring_prep_msg_ring: {}",
                                         std::strerror(-cqe->res)));
                }
                pending_works_--;
            } else {
                auto *work = static_cast<WorkInvoker *>(data);
                tsan_acquire(work);
                (*work)();
            }
        } else if (type == WorkType::Cancel) {
            detail::CancelRequest *request =
                static_cast<detail::CancelRequest *>(data);
            tsan_acquire(request);
            io_uring_sqe *sqe = ring_.get_sqe();
            prep_cancel_(sqe, request->data());
            request->notify();
        } else if (type == WorkType::Common) {
            auto *handle = static_cast<OpFinishHandleBase *>(data);
            auto op_finish = handle->handle(cqe);
            if (op_finish) {
                pending_works_--;
            }
        } else {
            unreachable();
        }
    }
 
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() noexcept {
    return *detail::Context::current().runtime();
}
 
} // namespace condy